<commit_before>package store

import (
	"bytes"
	"crypto/x509"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/adamdecaf/cert-manage/tools/file"
	"github.com/adamdecaf/cert-manage/tools/pem"
	"github.com/adamdecaf/cert-manage/whitelist"
)

// Docs:
// - https://developer.mozilla.org/en-US/docs/Mozilla/Projects/NSS
// - https://developer.mozilla.org/en-US/docs/Mozilla/Projects/NSS/Tools/certutil
// - https://wiki.mozilla.org/NSS_Shared_DB
// - https://www.chromium.org/Home/chromium-security/root-ca-policy
//   - https://chromium.googlesource.com/chromium/src/+/master/docs/linux_cert_management.md
//   - https://wiki.mozilla.org/NSS_Shared_DB_And_LINUX

var (
	cutil = certutil{
		execPaths: []string{
			"/usr/local/opt/nss/bin/certutil", // Darwin
			"/usr/bin/certutil",               // Linux
		},
	}

	// We're only going to support the current version (cert8.db)
	// https://wiki.mozilla.org/NSS_Shared_DB#Where_we_are_today
	cert8Filename = "cert8.db"

	// Trust Attributes which signify "No Trust"
	// Run `certutil -A -H` for the full list
	trustAttrsNoTrust = "p,p,p"
)

type nssStore struct {
	// nssType refers to the application using this nss cert8.db instance
	// often multiple applications can share one db, or there are multiple
	// cert8.db files on a system.
	// This allows backups and restores to operate on the correct cert8.db file
	// for a given application.
	nssType string

	// paths represent the fs locations where cert8.db are stored
	// In the case of an app like Firefox this would be looking in the following locations:
	//  ~/Library/Application Support/Firefox/Profiles/*  (Darwin)
	// The expected result is a slice of directories that can be passed into certutil's -d option
	paths []cert8db

	// Holds a trigger if we've made modifications (useful for triggering a "Restart app" message
	// after we're done with modifications.
	notify *sync.Once
}

func NssStore(nssType string, paths []cert8db) nssStore {
	return nssStore{
		nssType: nssType,
		paths:   paths,
		notify:  &sync.Once{},
	}
}

func collectNssSuggestions(sugs []string) []cert8db {
	if debug {
		if len(sugs) == 0 {
			fmt.Println("store/nss: no cert8.db paths suggested")
			return nil
		}
		fmt.Printf("store/nss: suggestions: %s\n", strings.Join(sugs, ", "))
	}

	kept := make([]cert8db, 0)
	for i := range sugs {
		// Glob and find a cert8.db file
		matches, err := filepath.Glob(sugs[i])
		if err != nil {
			if debug {
				fmt.Printf("store/nss: %v\n", err)
			}
			return nil
		}

		// Accumulate dirs with a cert8.db file
		for j := range matches {
			if containsCert8db(matches[j]) {
				kept = append(kept, cert8db(matches[j]))
			}
		}
	}
	return kept
}

func containsCert8db(p string) bool {
	where := filepath.Join(p, cert8Filename)
	if debug {
		fmt.Printf("store/nss: guessing cert8.db location: %s\n", where)
	}
	s, err := os.Stat(where)
	if err != nil {
		return false
	}
	return s.Size() > 0
}

// NSS apps often require being restarted to get the updated set of trustAttrs for each certificate
func (s nssStore) notifyToRestart() {
	s.notify.Do(func() {
		fmt.Printf("Restart %s to refresh certificate trust\n", strings.Title(s.nssType))
	})
	return
}

// we should be able to backup a cert8.db file directly
func (s nssStore) Backup() error {
	dir, err := getCertManageDir(s.nssType)
	if err != nil {
		return err
	}

	// Only backup the first nss cert8.db path for now
	if len(s.paths) == 0 {
		return errors.New("No NSS cert db paths found")
	}

	// Copy file into backup location
	src := filepath.Join(string(s.paths[0]), cert8Filename)
	dst := filepath.Join(dir, fmt.Sprintf("cert8.db-%d", time.Now().Unix()))
	return file.CopyFile(src, dst)
}

// List returns the installed (and trusted) certificates contained in a NSS cert8.db file
//
// To list certificates with the NSS `certutil` tool the following would be ran
// $ /usr/local/opt/nss/bin/certutil -L -d "$dir"
//
// Note: `dir` represents a directory path which contains a cert8.db file
func (s nssStore) List() ([]*x509.Certificate, error) {
	if len(s.paths) == 0 {
		return nil, errors.New("unable to find NSS db directory")
	}

	items, err := cutil.listCertsFromDB(s.paths[0])
	if err != nil {
		return nil, err
	}

	kept := make([]*x509.Certificate, 0)
	for i := range items {
		if items[i].trustedForSSL() {
			kept = append(kept, items[i].certs...)
		}
	}
	return kept, nil
}

func (s nssStore) Remove(wh whitelist.Whitelist) error {
	if len(s.paths) == 0 {
		return errors.New("unable to find NSS db directory")
	}

	items, err := cutil.listCertsFromDB(s.paths[0])
	if err != nil {
		return err
	}

	// Remove trust from each cert if needed.
	for i := range items {
		if wh.MatchesAll(items[i].certs) {
			continue
		}

		// whitelist didn't match, blacklist cert
		defer s.notifyToRestart()
		err = cutil.modifyTrustAttributes(s.paths[0], items[i].nick, trustAttrsNoTrust)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s nssStore) Restore(where string) error {
	dir, err := getCertManageDir(s.nssType)
	if err != nil {
		return err
	}
	src, err := getLatestBackupFile(dir)
	if err != nil {
		return err
	}

	// Check we've got a restore point
	if len(s.paths) == 0 {
		return errors.New("No directory to restore NSS cert db into")
	}

	// Queue notification to restart app
	defer s.notifyToRestart()

	// Restore the latest backup file
	dst := filepath.Join(string(s.paths[0]), cert8Filename)
	return file.CopyFile(src, dst)
}

// cert8db represents a fs path for a cert8.db file
type cert8db string

// cert8Item represents an x509 Certificate with the NSS trust attributes
type cert8Item struct {
	nick  string
	certs []*x509.Certificate

	// The 'Trust Attributes' header from `certutil -L`, this represents three usecases for
	// x509 Certificates: SSL,S/MIME,JAR/XPI
	trustAttrs string
}

func (c cert8Item) trustedForSSL() bool {
	parts := strings.SplitN(c.trustAttrs, ",", 2) // We only care about the first C,.,. attribute
	if len(parts) != 2 {
		if debug {
			fmt.Printf("store/nss: after trustAttrs split (in %d parts): %s\n", len(parts), parts)
		}
		return false
	}

	// The only attribute for explicit distrust is 'p', which per the docs:
	// 'p 	 prohibited (explicitly distrusted)'
	//
	// The other flags refer to sending warnings (but still trusted), user certs,
	// or other attributes which may limit, but not explicitly remove trust
	// in regards to SSL/TLS communication.
	return !strings.Contains(parts[0], "p")
}

// certutil represents the NSS cli tool by the same name
type certutil struct {
	execPaths []string
}

func (c certutil) getExecPath() (string, error) {
	for i := range c.execPaths {
		if file.IsExecutable(c.execPaths[i]) {
			return c.execPaths[i], nil
		}
	}
	return "", errors.New("No executable for `certutil` found")
}

// Emulates the following
// /usr/local/opt/nss/bin/certutil -L -d '~/Library/Application Support/Firefox/Profiles/rrdlhe7o.default'
func (c certutil) listCertsFromDB(path cert8db) ([]cert8Item, error) {
	expath, err := c.getExecPath()
	if err != nil {
		return nil, err
	}

	args := []string{
		"-L",
		"-d", string(path),
	}
	cmd := exec.Command(expath, args...)

	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err = cmd.Run()
	if err != nil {
		if debug {
			fmt.Printf("Command was:\n %s", strings.Join(cmd.Args, " "))
		}
		return nil, err
	}

	// parse out each 'Certificate Nickname' field
	// there are some header lines that are either whitespace or contain the following:
	//
	// Certificate Nickname                                         Trust Attributes
	// SSL,S/MIME,JAR/XPI
	//
	items := make([]cert8Item, 0)
	for {
		// read next line
		line, err := stdout.ReadString(byte('\n'))
		if err != nil && err != io.EOF {
			return nil, err
		}

		// convert the line into a cert8Item
		nick, trust := c.parseCert8Item(line)
		if nick == "" || trust == "" {
			if err == io.EOF {
				break
			}
			continue // skip blank lines (headers are returned here as blank)
		}
		certs, err := c.readCertificatesForNick(path, nick)
		if err != nil {
			return nil, err
		}

		items = append(items, cert8Item{
			nick:       nick,
			certs:      certs,
			trustAttrs: trust,
		})

		// quit if we hit the end
		if err == io.EOF {
			break
		}
	}

	return items, nil
}

func (c certutil) parseCert8Item(line string) (nick string, trust string) {
	line = strings.TrimSpace(line)

	if len(line) == 0 ||
		strings.Contains(line, "Certificate Nickname") ||
		strings.Contains(line, "SSL,S/MIME,JAR/XPI") {
		return
	}

	// Pull out the 'Certificate Nickname'
	// Examples:
	// DigiCert SHA2 Extended Validation Server CA                  ,,
	// Symantec Class 3 Extended Validation SHA256 SSL CA           CT,C,C
	split := strings.Split(line, " ")

	nick = strings.TrimSpace(strings.Join(split[:len(split)-1], " "))
	trust = strings.TrimSpace(split[len(split)-1:][0]) // last

	return
}

// /usr/local/opt/nss/bin/certutil -L -d "$dir" -n 'Microsoft IT SSL SHA2' -a
func (c certutil) readCertificatesForNick(path cert8db, nick string) ([]*x509.Certificate, error) {
	expath, err := c.getExecPath()
	if err != nil {
		return nil, err
	}

	args := []string{
		"-L",
		"-a",
		"-d", string(path),
		"-n", nick,
	}
	cmd := exec.Command(expath, args...)
	var stdout bytes.Buffer
	cmd.Stdout = &stdout

	err = cmd.Run()
	if err != nil {
		if debug {
			fmt.Printf("Command was:\n %s", strings.Join(cmd.Args, " "))
		}
		return nil, err
	}

	certs, err := pem.Parse(stdout.Bytes())
	if err != nil {
		return nil, err
	}

	return certs, nil
}

// $ /usr/local/opt/nss/bin/certutil -M -n <nick> -t <trust-args> -d <dir>
func (c certutil) modifyTrustAttributes(path cert8db, nick, trustAttrs string) error {
	expath, err := c.getExecPath()
	if err != nil {
		return err
	}

	args := []string{
		"-M",
		"-n", nick,
		"-t", trustAttrs,
		"-d", string(path),
	}
	cmd := exec.Command(expath, args...)
	var stdout bytes.Buffer
	cmd.Stdout = &stdout

	err = cmd.Run()
	if debug {
		fmt.Printf("Command was: \n%s\nOutput was: \n%s\n", strings.Join(cmd.Args, " "), stdout.String())
	}
	return err
}
<commit_msg>store/nss: newline after cmd debug output<commit_after>package store

import (
	"bytes"
	"crypto/x509"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/adamdecaf/cert-manage/tools/file"
	"github.com/adamdecaf/cert-manage/tools/pem"
	"github.com/adamdecaf/cert-manage/whitelist"
)

// Docs:
// - https://developer.mozilla.org/en-US/docs/Mozilla/Projects/NSS
// - https://developer.mozilla.org/en-US/docs/Mozilla/Projects/NSS/Tools/certutil
// - https://wiki.mozilla.org/NSS_Shared_DB
// - https://www.chromium.org/Home/chromium-security/root-ca-policy
//   - https://chromium.googlesource.com/chromium/src/+/master/docs/linux_cert_management.md
//   - https://wiki.mozilla.org/NSS_Shared_DB_And_LINUX

var (
	cutil = certutil{
		execPaths: []string{
			"/usr/local/opt/nss/bin/certutil", // Darwin
			"/usr/bin/certutil",               // Linux
		},
	}

	// We're only going to support the current version (cert8.db)
	// https://wiki.mozilla.org/NSS_Shared_DB#Where_we_are_today
	cert8Filename = "cert8.db"

	// Trust Attributes which signify "No Trust"
	// Run `certutil -A -H` for the full list
	trustAttrsNoTrust = "p,p,p"
)

type nssStore struct {
	// nssType refers to the application using this nss cert8.db instance
	// often multiple applications can share one db, or there are multiple
	// cert8.db files on a system.
	// This allows backups and restores to operate on the correct cert8.db file
	// for a given application.
	nssType string

	// paths represent the fs locations where cert8.db are stored
	// In the case of an app like Firefox this would be looking in the following locations:
	//  ~/Library/Application Support/Firefox/Profiles/*  (Darwin)
	// The expected result is a slice of directories that can be passed into certutil's -d option
	paths []cert8db

	// Holds a trigger if we've made modifications (useful for triggering a "Restart app" message
	// after we're done with modifications.
	notify *sync.Once
}

func NssStore(nssType string, paths []cert8db) nssStore {
	return nssStore{
		nssType: nssType,
		paths:   paths,
		notify:  &sync.Once{},
	}
}

func collectNssSuggestions(sugs []string) []cert8db {
	if debug {
		if len(sugs) == 0 {
			fmt.Println("store/nss: no cert8.db paths suggested")
			return nil
		}
		fmt.Printf("store/nss: suggestions: %s\n", strings.Join(sugs, ", "))
	}

	kept := make([]cert8db, 0)
	for i := range sugs {
		// Glob and find a cert8.db file
		matches, err := filepath.Glob(sugs[i])
		if err != nil {
			if debug {
				fmt.Printf("store/nss: %v\n", err)
			}
			return nil
		}

		// Accumulate dirs with a cert8.db file
		for j := range matches {
			if containsCert8db(matches[j]) {
				kept = append(kept, cert8db(matches[j]))
			}
		}
	}
	return kept
}

func containsCert8db(p string) bool {
	where := filepath.Join(p, cert8Filename)
	if debug {
		fmt.Printf("store/nss: guessing cert8.db location: %s\n", where)
	}
	s, err := os.Stat(where)
	if err != nil {
		return false
	}
	return s.Size() > 0
}

// NSS apps often require being restarted to get the updated set of trustAttrs for each certificate
func (s nssStore) notifyToRestart() {
	s.notify.Do(func() {
		fmt.Printf("Restart %s to refresh certificate trust\n", strings.Title(s.nssType))
	})
	return
}

// we should be able to backup a cert8.db file directly
func (s nssStore) Backup() error {
	dir, err := getCertManageDir(s.nssType)
	if err != nil {
		return err
	}

	// Only backup the first nss cert8.db path for now
	if len(s.paths) == 0 {
		return errors.New("No NSS cert db paths found")
	}

	// Copy file into backup location
	src := filepath.Join(string(s.paths[0]), cert8Filename)
	dst := filepath.Join(dir, fmt.Sprintf("cert8.db-%d", time.Now().Unix()))
	return file.CopyFile(src, dst)
}

// List returns the installed (and trusted) certificates contained in a NSS cert8.db file
//
// To list certificates with the NSS `certutil` tool the following would be ran
// $ /usr/local/opt/nss/bin/certutil -L -d "$dir"
//
// Note: `dir` represents a directory path which contains a cert8.db file
func (s nssStore) List() ([]*x509.Certificate, error) {
	if len(s.paths) == 0 {
		return nil, errors.New("unable to find NSS db directory")
	}

	items, err := cutil.listCertsFromDB(s.paths[0])
	if err != nil {
		return nil, err
	}

	kept := make([]*x509.Certificate, 0)
	for i := range items {
		if items[i].trustedForSSL() {
			kept = append(kept, items[i].certs...)
		}
	}
	return kept, nil
}

func (s nssStore) Remove(wh whitelist.Whitelist) error {
	if len(s.paths) == 0 {
		return errors.New("unable to find NSS db directory")
	}

	items, err := cutil.listCertsFromDB(s.paths[0])
	if err != nil {
		return err
	}

	// Remove trust from each cert if needed.
	for i := range items {
		if wh.MatchesAll(items[i].certs) {
			continue
		}

		// whitelist didn't match, blacklist cert
		defer s.notifyToRestart()
		err = cutil.modifyTrustAttributes(s.paths[0], items[i].nick, trustAttrsNoTrust)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s nssStore) Restore(where string) error {
	dir, err := getCertManageDir(s.nssType)
	if err != nil {
		return err
	}
	src, err := getLatestBackupFile(dir)
	if err != nil {
		return err
	}

	// Check we've got a restore point
	if len(s.paths) == 0 {
		return errors.New("No directory to restore NSS cert db into")
	}

	// Queue notification to restart app
	defer s.notifyToRestart()

	// Restore the latest backup file
	dst := filepath.Join(string(s.paths[0]), cert8Filename)
	return file.CopyFile(src, dst)
}

// cert8db represents a fs path for a cert8.db file
type cert8db string

// cert8Item represents an x509 Certificate with the NSS trust attributes
type cert8Item struct {
	nick  string
	certs []*x509.Certificate

	// The 'Trust Attributes' header from `certutil -L`, this represents three usecases for
	// x509 Certificates: SSL,S/MIME,JAR/XPI
	trustAttrs string
}

func (c cert8Item) trustedForSSL() bool {
	parts := strings.SplitN(c.trustAttrs, ",", 2) // We only care about the first C,.,. attribute
	if len(parts) != 2 {
		if debug {
			fmt.Printf("store/nss: after trustAttrs split (in %d parts): %s\n", len(parts), parts)
		}
		return false
	}

	// The only attribute for explicit distrust is 'p', which per the docs:
	// 'p 	 prohibited (explicitly distrusted)'
	//
	// The other flags refer to sending warnings (but still trusted), user certs,
	// or other attributes which may limit, but not explicitly remove trust
	// in regards to SSL/TLS communication.
	return !strings.Contains(parts[0], "p")
}

// certutil represents the NSS cli tool by the same name
type certutil struct {
	execPaths []string
}

func (c certutil) getExecPath() (string, error) {
	for i := range c.execPaths {
		if file.IsExecutable(c.execPaths[i]) {
			return c.execPaths[i], nil
		}
	}
	return "", errors.New("No executable for `certutil` found")
}

// Emulates the following
// /usr/local/opt/nss/bin/certutil -L -d '~/Library/Application Support/Firefox/Profiles/rrdlhe7o.default'
func (c certutil) listCertsFromDB(path cert8db) ([]cert8Item, error) {
	expath, err := c.getExecPath()
	if err != nil {
		return nil, err
	}

	args := []string{
		"-L",
		"-d", string(path),
	}
	cmd := exec.Command(expath, args...)

	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err = cmd.Run()
	if err != nil {
		if debug {
			fmt.Printf("Command was:\n %s\n", strings.Join(cmd.Args, " "))
		}
		return nil, err
	}

	// parse out each 'Certificate Nickname' field
	// there are some header lines that are either whitespace or contain the following:
	//
	// Certificate Nickname                                         Trust Attributes
	// SSL,S/MIME,JAR/XPI
	//
	items := make([]cert8Item, 0)
	for {
		// read next line
		line, err := stdout.ReadString(byte('\n'))
		if err != nil && err != io.EOF {
			return nil, err
		}

		// convert the line into a cert8Item
		nick, trust := c.parseCert8Item(line)
		if nick == "" || trust == "" {
			if err == io.EOF {
				break
			}
			continue // skip blank lines (headers are returned here as blank)
		}
		certs, err := c.readCertificatesForNick(path, nick)
		if err != nil {
			return nil, err
		}

		items = append(items, cert8Item{
			nick:       nick,
			certs:      certs,
			trustAttrs: trust,
		})

		// quit if we hit the end
		if err == io.EOF {
			break
		}
	}

	return items, nil
}

func (c certutil) parseCert8Item(line string) (nick string, trust string) {
	line = strings.TrimSpace(line)

	if len(line) == 0 ||
		strings.Contains(line, "Certificate Nickname") ||
		strings.Contains(line, "SSL,S/MIME,JAR/XPI") {
		return
	}

	// Pull out the 'Certificate Nickname'
	// Examples:
	// DigiCert SHA2 Extended Validation Server CA                  ,,
	// Symantec Class 3 Extended Validation SHA256 SSL CA           CT,C,C
	split := strings.Split(line, " ")

	nick = strings.TrimSpace(strings.Join(split[:len(split)-1], " "))
	trust = strings.TrimSpace(split[len(split)-1:][0]) // last

	return
}

// /usr/local/opt/nss/bin/certutil -L -d "$dir" -n 'Microsoft IT SSL SHA2' -a
func (c certutil) readCertificatesForNick(path cert8db, nick string) ([]*x509.Certificate, error) {
	expath, err := c.getExecPath()
	if err != nil {
		return nil, err
	}

	args := []string{
		"-L",
		"-a",
		"-d", string(path),
		"-n", nick,
	}
	cmd := exec.Command(expath, args...)
	var stdout bytes.Buffer
	cmd.Stdout = &stdout

	err = cmd.Run()
	if err != nil {
		if debug {
			fmt.Printf("Command was:\n %s\n", strings.Join(cmd.Args, " "))
		}
		return nil, err
	}

	certs, err := pem.Parse(stdout.Bytes())
	if err != nil {
		return nil, err
	}

	return certs, nil
}

// $ /usr/local/opt/nss/bin/certutil -M -n <nick> -t <trust-args> -d <dir>
func (c certutil) modifyTrustAttributes(path cert8db, nick, trustAttrs string) error {
	expath, err := c.getExecPath()
	if err != nil {
		return err
	}

	args := []string{
		"-M",
		"-n", nick,
		"-t", trustAttrs,
		"-d", string(path),
	}
	cmd := exec.Command(expath, args...)
	var stdout bytes.Buffer
	cmd.Stdout = &stdout

	err = cmd.Run()
	if debug {
		fmt.Printf("Command was: \n%s\nOutput was: \n%s\n", strings.Join(cmd.Args, " "), stdout.String())
	}
	return err
}
<|endoftext|>


<commit_before>package couchbase

import (
	"encoding/json"
	"fmt"
	"github.com/couchbase/goutils/logging"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"time"
	"unsafe"
)

// Bucket auto-updater gets the latest version of the bucket config from
// the server. If the configuration has changed then updated the local
// bucket information. If the bucket has been deleted then notify anyone
// who is holding a reference to this bucket

const MAX_RETRY_COUNT = 5
const DISCONNECT_PERIOD = 120 * time.Second

type NotifyFn func(bucket string, err error)
type StreamingFn func(bucket *Bucket)

// Use TCP keepalive to detect half close sockets
var updaterTransport http.RoundTripper = &http.Transport{
	Proxy: http.ProxyFromEnvironment,
	Dial: (&net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}).Dial,
}

var updaterHTTPClient = &http.Client{Transport: updaterTransport}

func doHTTPRequestForUpdate(req *http.Request) (*http.Response, error) {

	var err error
	var res *http.Response

	for i := 0; i < HTTP_MAX_RETRY; i++ {
		res, err = updaterHTTPClient.Do(req)
		if err != nil && isHttpConnError(err) {
			continue
		}
		break
	}

	if err != nil {
		return nil, err
	}

	return res, err
}

func (b *Bucket) RunBucketUpdater(notify NotifyFn) {
	b.RunBucketUpdater2(nil, notify)
}

func (b *Bucket) RunBucketUpdater2(streamingFn StreamingFn, notify NotifyFn) {
	go func() {
		err := b.UpdateBucket2(streamingFn)
		if err != nil {
			if notify != nil {
				notify(b.GetName(), err)
			}
			logging.Errorf(" Bucket Updater exited with err %v", err)
		}
	}()
}

func (b *Bucket) replaceConnPools2(with []*connectionPool, bucketLocked bool) {
	if !bucketLocked {
		b.Lock()
		defer b.Unlock()
	}
	old := b.connPools
	b.connPools = unsafe.Pointer(&with)
	if old != nil {
		for _, pool := range *(*[]*connectionPool)(old) {
			if pool != nil && pool.inUse == false {
				pool.Close()
			}
		}
	}
	return
}

func (b *Bucket) UpdateBucket() error {
	return b.UpdateBucket2(nil)
}

func (b *Bucket) UpdateBucket2(streamingFn StreamingFn) error {
	var failures int
	var returnErr error
	var poolServices PoolServices

	for {

		if failures == MAX_RETRY_COUNT {
			logging.Errorf(" Maximum failures reached. Exiting loop...")
			return fmt.Errorf("Max failures reached. Last Error %v", returnErr)
		}

		nodes := b.Nodes()
		if len(nodes) < 1 {
			return fmt.Errorf("No healthy nodes found")
		}

		startNode := rand.Intn(len(nodes))
		node := nodes[(startNode)%len(nodes)]

		streamUrl := fmt.Sprintf("http://%s/pools/default/bucketsStreaming/%s", node.Hostname, uriAdj(b.GetName()))
		logging.Infof(" Trying with %s", streamUrl)
		req, err := http.NewRequest("GET", streamUrl, nil)
		if err != nil {
			return err
		}

		// Lock here to avoid having pool closed under us.
		b.RLock()
		err = maybeAddAuth(req, b.pool.client.ah)
		b.RUnlock()
		if err != nil {
			return err
		}

		res, err := doHTTPRequestForUpdate(req)
		if err != nil {
			return err
		}

		if res.StatusCode != 200 {
			bod, _ := ioutil.ReadAll(io.LimitReader(res.Body, 512))
			logging.Errorf("Failed to connect to host, unexpected status code: %v. Body %s", res.StatusCode, bod)
			res.Body.Close()
			returnErr = fmt.Errorf("Failed to connect to host. Status %v Body %s", res.StatusCode, bod)
			failures++
			continue
		}

		dec := json.NewDecoder(res.Body)

		tmpb := &Bucket{}
		for {

			err := dec.Decode(&tmpb)
			if err != nil {
				returnErr = err
				res.Body.Close()
				break
			}

			// if we got here, reset failure count
			failures = 0

			if b.pool.client.tlsConfig != nil {
				poolServices, err = b.pool.client.GetPoolServices("default")
				if err != nil {
					returnErr = err
					res.Body.Close()
					break
				}
			}

			b.Lock()

			// mark all the old connection pools for deletion
			pools := b.getConnPools(true /* already locked */)
			for _, pool := range pools {
				if pool != nil {
					pool.inUse = false
				}
			}

			newcps := make([]*connectionPool, len(tmpb.VBSMJson.ServerList))
			for i := range newcps {
				// get the old connection pool and check if it is still valid
				pool := b.getConnPoolByHost(tmpb.VBSMJson.ServerList[i], true /* bucket already locked */)
				if pool != nil && pool.inUse == false && pool.tlsConfig == b.pool.client.tlsConfig {
					// if the hostname and index is unchanged then reuse this pool
					newcps[i] = pool
					pool.inUse = true
					continue
				}
				// else create a new pool
				var encrypted bool
				hostport := tmpb.VBSMJson.ServerList[i]
				if b.pool.client.tlsConfig != nil {
					hostport, encrypted, err = MapKVtoSSL(hostport, &poolServices)
					if err != nil {
						b.Unlock()
						return err
					}
				}
				if b.ah != nil {
					newcps[i] = newConnectionPool(hostport,
						b.ah, false, PoolSize, PoolOverflow, b.pool.client.tlsConfig, b.Name, encrypted)

				} else {
					newcps[i] = newConnectionPool(hostport,
						b.authHandler(true /* bucket already locked */),
						false, PoolSize, PoolOverflow, b.pool.client.tlsConfig, b.Name, encrypted)
				}
			}

			b.replaceConnPools2(newcps, true /* bucket already locked */)

			tmpb.ah = b.ah
			b.vBucketServerMap = unsafe.Pointer(&tmpb.VBSMJson)
			b.nodeList = unsafe.Pointer(&tmpb.NodesJSON)
			b.Unlock()

			if streamingFn != nil {
				streamingFn(tmpb)
			}
			logging.Infof("Got new configuration for bucket %s", b.GetName())

		}
		// we are here because of an error
		failures++
		continue

	}
	return nil
}
<commit_msg>MB-41695 - Move logging to debug. - Query.log got bombarded with log message<commit_after>package couchbase

import (
	"encoding/json"
	"fmt"
	"github.com/couchbase/goutils/logging"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"time"
	"unsafe"
)

// Bucket auto-updater gets the latest version of the bucket config from
// the server. If the configuration has changed then updated the local
// bucket information. If the bucket has been deleted then notify anyone
// who is holding a reference to this bucket

const MAX_RETRY_COUNT = 5
const DISCONNECT_PERIOD = 120 * time.Second

type NotifyFn func(bucket string, err error)
type StreamingFn func(bucket *Bucket)

// Use TCP keepalive to detect half close sockets
var updaterTransport http.RoundTripper = &http.Transport{
	Proxy: http.ProxyFromEnvironment,
	Dial: (&net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}).Dial,
}

var updaterHTTPClient = &http.Client{Transport: updaterTransport}

func doHTTPRequestForUpdate(req *http.Request) (*http.Response, error) {

	var err error
	var res *http.Response

	for i := 0; i < HTTP_MAX_RETRY; i++ {
		res, err = updaterHTTPClient.Do(req)
		if err != nil && isHttpConnError(err) {
			continue
		}
		break
	}

	if err != nil {
		return nil, err
	}

	return res, err
}

func (b *Bucket) RunBucketUpdater(notify NotifyFn) {
	b.RunBucketUpdater2(nil, notify)
}

func (b *Bucket) RunBucketUpdater2(streamingFn StreamingFn, notify NotifyFn) {
	go func() {
		err := b.UpdateBucket2(streamingFn)
		if err != nil {
			if notify != nil {
				notify(b.GetName(), err)
			}
			logging.Errorf(" Bucket Updater exited with err %v", err)
		}
	}()
}

func (b *Bucket) replaceConnPools2(with []*connectionPool, bucketLocked bool) {
	if !bucketLocked {
		b.Lock()
		defer b.Unlock()
	}
	old := b.connPools
	b.connPools = unsafe.Pointer(&with)
	if old != nil {
		for _, pool := range *(*[]*connectionPool)(old) {
			if pool != nil && pool.inUse == false {
				pool.Close()
			}
		}
	}
	return
}

func (b *Bucket) UpdateBucket() error {
	return b.UpdateBucket2(nil)
}

func (b *Bucket) UpdateBucket2(streamingFn StreamingFn) error {
	var failures int
	var returnErr error
	var poolServices PoolServices

	for {

		if failures == MAX_RETRY_COUNT {
			logging.Errorf(" Maximum failures reached. Exiting loop...")
			return fmt.Errorf("Max failures reached. Last Error %v", returnErr)
		}

		nodes := b.Nodes()
		if len(nodes) < 1 {
			return fmt.Errorf("No healthy nodes found")
		}

		startNode := rand.Intn(len(nodes))
		node := nodes[(startNode)%len(nodes)]

		streamUrl := fmt.Sprintf("http://%s/pools/default/bucketsStreaming/%s", node.Hostname, uriAdj(b.GetName()))
		logging.Infof(" Trying with %s", streamUrl)
		req, err := http.NewRequest("GET", streamUrl, nil)
		if err != nil {
			return err
		}

		// Lock here to avoid having pool closed under us.
		b.RLock()
		err = maybeAddAuth(req, b.pool.client.ah)
		b.RUnlock()
		if err != nil {
			return err
		}

		res, err := doHTTPRequestForUpdate(req)
		if err != nil {
			return err
		}

		if res.StatusCode != 200 {
			bod, _ := ioutil.ReadAll(io.LimitReader(res.Body, 512))
			logging.Errorf("Failed to connect to host, unexpected status code: %v. Body %s", res.StatusCode, bod)
			res.Body.Close()
			returnErr = fmt.Errorf("Failed to connect to host. Status %v Body %s", res.StatusCode, bod)
			failures++
			continue
		}

		dec := json.NewDecoder(res.Body)

		tmpb := &Bucket{}
		for {

			err := dec.Decode(&tmpb)
			if err != nil {
				returnErr = err
				res.Body.Close()
				break
			}

			// if we got here, reset failure count
			failures = 0

			if b.pool.client.tlsConfig != nil {
				poolServices, err = b.pool.client.GetPoolServices("default")
				if err != nil {
					returnErr = err
					res.Body.Close()
					break
				}
			}

			b.Lock()

			// mark all the old connection pools for deletion
			pools := b.getConnPools(true /* already locked */)
			for _, pool := range pools {
				if pool != nil {
					pool.inUse = false
				}
			}

			newcps := make([]*connectionPool, len(tmpb.VBSMJson.ServerList))
			for i := range newcps {
				// get the old connection pool and check if it is still valid
				pool := b.getConnPoolByHost(tmpb.VBSMJson.ServerList[i], true /* bucket already locked */)
				if pool != nil && pool.inUse == false && pool.tlsConfig == b.pool.client.tlsConfig {
					// if the hostname and index is unchanged then reuse this pool
					newcps[i] = pool
					pool.inUse = true
					continue
				}
				// else create a new pool
				var encrypted bool
				hostport := tmpb.VBSMJson.ServerList[i]
				if b.pool.client.tlsConfig != nil {
					hostport, encrypted, err = MapKVtoSSL(hostport, &poolServices)
					if err != nil {
						b.Unlock()
						return err
					}
				}
				if b.ah != nil {
					newcps[i] = newConnectionPool(hostport,
						b.ah, false, PoolSize, PoolOverflow, b.pool.client.tlsConfig, b.Name, encrypted)

				} else {
					newcps[i] = newConnectionPool(hostport,
						b.authHandler(true /* bucket already locked */),
						false, PoolSize, PoolOverflow, b.pool.client.tlsConfig, b.Name, encrypted)
				}
			}

			b.replaceConnPools2(newcps, true /* bucket already locked */)

			tmpb.ah = b.ah
			b.vBucketServerMap = unsafe.Pointer(&tmpb.VBSMJson)
			b.nodeList = unsafe.Pointer(&tmpb.NodesJSON)
			b.Unlock()

			if streamingFn != nil {
				streamingFn(tmpb)
			}
			logging.Debugf("Got new configuration for bucket %s", b.GetName())

		}
		// we are here because of an error
		failures++
		continue

	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"go/build"
	"sync"

	"github.com/bradfitz/iter"
	"github.com/shurcooL/vcsstate"
	"golang.org/x/tools/go/vcs"
)

// workspace is a Go workspace environment; each repo has local and remote components.
type workspace struct {
	ImportPaths       chan string // ImportPaths is the input for Go packages to be processed.
	unique            chan *Repo  // unique repos.
	processedFiltered chan *Repo  // processed repos, populated with local and remote info, filtered with shouldShow.
	Statuses          chan string // Statuses has results of running presenter on processed repos.
	Errors            chan error  // Errors contains errors that were encountered during processing of repos.

	shouldShow RepoFilter
	presenter  RepoPresenter

	reposMu sync.Mutex
	repos   map[string]*Repo // Map key is repoRoot.
}

func NewWorkspace(shouldShow RepoFilter, presenter RepoPresenter) *workspace {
	w := &workspace{
		ImportPaths:       make(chan string, 64),
		unique:            make(chan *Repo, 64),
		processedFiltered: make(chan *Repo, 64),
		Statuses:          make(chan string, 64),
		Errors:            make(chan error, 64),

		shouldShow: shouldShow,
		presenter:  presenter,

		repos: make(map[string]*Repo),
	}

	{
		var wg sync.WaitGroup
		for range iter.N(parallelism) {
			wg.Add(1)
			go w.uniqueWorker(&wg)
		}
		go func() {
			wg.Wait()
			close(w.unique)
		}()
	}
	{
		var wg sync.WaitGroup
		for range iter.N(parallelism) {
			wg.Add(1)
			go w.processFilterWorker(&wg)
		}
		go func() {
			wg.Wait()
			close(w.processedFiltered)
		}()
	}
	{
		var wg sync.WaitGroup
		for range iter.N(parallelism) {
			wg.Add(1)
			go w.presenterWorker(&wg)
		}
		go func() {
			wg.Wait()
			close(w.Statuses)
			close(w.Errors)
		}()
	}

	return w
}

// uniqueWorker finds unique repos out of all input Go packages.
func (w *workspace) uniqueWorker(wg *sync.WaitGroup) {
	defer wg.Done()
	for importPath := range w.ImportPaths {
		// Determine repo root.
		// This is potentially somewhat slow.
		bpkg, err := build.Import(importPath, wd, build.FindOnly)
		if err != nil {
			w.Errors <- err
			continue
		}
		if bpkg.Goroot {
			// gostatus has no support for printing status of packages in GOROOT, so skip those.
			continue
		}
		vcsCmd, root, err := vcs.FromDir(bpkg.Dir, bpkg.SrcRoot)
		if err != nil {
			// Go package not under VCS.
			var pkg *Repo
			w.reposMu.Lock()
			if _, ok := w.repos[bpkg.ImportPath]; !ok {
				pkg = &Repo{
					Path: bpkg.Dir,
					Root: bpkg.ImportPath,
					VCS:  nil,
				}
				w.repos[bpkg.ImportPath] = pkg
			}
			w.reposMu.Unlock()

			// If new package, send off to next stage.
			if pkg != nil {
				w.unique <- pkg
			}
			continue
		}
		vcs, err := vcsstate.NewVCS(vcsCmd)
		if err != nil {
			w.Errors <- fmt.Errorf("repo %v not supported by vcsstate: %v", root, err)
			continue
		}

		var repo *Repo
		w.reposMu.Lock()
		if _, ok := w.repos[root]; !ok {
			repo = &Repo{
				Path: bpkg.Dir,
				Root: root,
				VCS:  vcs,
			}
			w.repos[root] = repo
		}
		w.reposMu.Unlock()

		// If new repo, send off to next stage.
		if repo != nil {
			w.unique <- repo
		}
	}
}

// processFilterWorker figures out repo local and remote info, and filters with shouldShow.
func (w *workspace) processFilterWorker(wg *sync.WaitGroup) {
	defer wg.Done()
	for repo := range w.unique {
		w.computeVCSState(repo)

		if !w.shouldShow(repo) {
			continue
		}

		w.processedFiltered <- repo
	}
}

func (*workspace) computeVCSState(r *Repo) {
	if r.VCS == nil {
		// Go package not under VCS.
		return
	}

	if s, err := r.VCS.Status(r.Path); err == nil {
		r.Local.Status = s
	}
	if b, err := r.VCS.Branch(r.Path); err == nil {
		r.Local.Branch = b
	}
	if rev, err := r.VCS.LocalRevision(r.Path); err == nil {
		r.Local.Revision = rev
	}
	if s, err := r.VCS.Stash(r.Path); err == nil {
		r.Local.Stash = s
	}
	if remote, err := r.VCS.RemoteURL(r.Path); err == nil {
		r.Local.RemoteURL = remote
	}
	if rev, err := r.VCS.RemoteRevision(r.Path); err == nil {
		r.Remote.Revision = rev
	}
	if r.Remote.Revision != "" {
		if c, err := r.VCS.Contains(r.Path, r.Remote.Revision); err == nil {
			r.LocalContainsRemoteRevision = c
		}
	}
	if rr, err := vcs.RepoRootForImportPath(r.Root, false); err == nil {
		r.Remote.RepoURL = rr.Repo
	}
}

// presenterWorker runs presenter on processed and filtered repos.
func (w *workspace) presenterWorker(wg *sync.WaitGroup) {
	defer wg.Done()
	for repo := range w.processedFiltered {
		w.Statuses <- w.presenter(repo)
	}
}
<commit_msg>Minor comment improvements.<commit_after>package main

import (
	"fmt"
	"go/build"
	"sync"

	"github.com/bradfitz/iter"
	"github.com/shurcooL/vcsstate"
	"golang.org/x/tools/go/vcs"
)

// workspace is a Go workspace environment; each repo has local and remote components.
type workspace struct {
	ImportPaths       chan string // ImportPaths is the input for Go packages to be processed.
	unique            chan *Repo  // Unique repos.
	processedFiltered chan *Repo  // Processed repos, populated with local and remote state, filtered with shouldShow.
	Statuses          chan string // Statuses has results of running presenter on processed repos.
	Errors            chan error  // Errors contains errors that were encountered during processing of repos.

	shouldShow RepoFilter
	presenter  RepoPresenter

	reposMu sync.Mutex
	repos   map[string]*Repo // Map key is the import path corresponding to the root of the repository or Go package.
}

func NewWorkspace(shouldShow RepoFilter, presenter RepoPresenter) *workspace {
	w := &workspace{
		ImportPaths:       make(chan string, 64),
		unique:            make(chan *Repo, 64),
		processedFiltered: make(chan *Repo, 64),
		Statuses:          make(chan string, 64),
		Errors:            make(chan error, 64),

		shouldShow: shouldShow,
		presenter:  presenter,

		repos: make(map[string]*Repo),
	}

	{
		var wg sync.WaitGroup
		for range iter.N(parallelism) {
			wg.Add(1)
			go w.uniqueWorker(&wg)
		}
		go func() {
			wg.Wait()
			close(w.unique)
		}()
	}
	{
		var wg sync.WaitGroup
		for range iter.N(parallelism) {
			wg.Add(1)
			go w.processFilterWorker(&wg)
		}
		go func() {
			wg.Wait()
			close(w.processedFiltered)
		}()
	}
	{
		var wg sync.WaitGroup
		for range iter.N(parallelism) {
			wg.Add(1)
			go w.presenterWorker(&wg)
		}
		go func() {
			wg.Wait()
			close(w.Statuses)
			close(w.Errors)
		}()
	}

	return w
}

// uniqueWorker finds unique repos out of all input Go packages.
func (w *workspace) uniqueWorker(wg *sync.WaitGroup) {
	defer wg.Done()
	for importPath := range w.ImportPaths {
		// Determine repo root.
		// This is potentially somewhat slow.
		bpkg, err := build.Import(importPath, wd, build.FindOnly)
		if err != nil {
			w.Errors <- err
			continue
		}
		if bpkg.Goroot {
			// gostatus has no support for printing status of packages in GOROOT, so skip those.
			continue
		}
		vcsCmd, root, err := vcs.FromDir(bpkg.Dir, bpkg.SrcRoot)
		if err != nil {
			// Go package not under VCS.
			var pkg *Repo
			w.reposMu.Lock()
			if _, ok := w.repos[bpkg.ImportPath]; !ok {
				pkg = &Repo{
					Path: bpkg.Dir,
					Root: bpkg.ImportPath,
					VCS:  nil,
				}
				w.repos[bpkg.ImportPath] = pkg
			}
			w.reposMu.Unlock()

			// If new package, send off to next stage.
			if pkg != nil {
				w.unique <- pkg
			}
			continue
		}
		vcs, err := vcsstate.NewVCS(vcsCmd)
		if err != nil {
			w.Errors <- fmt.Errorf("repo %v not supported by vcsstate: %v", root, err)
			continue
		}

		var repo *Repo
		w.reposMu.Lock()
		if _, ok := w.repos[root]; !ok {
			repo = &Repo{
				Path: bpkg.Dir,
				Root: root,
				VCS:  vcs,
			}
			w.repos[root] = repo
		}
		w.reposMu.Unlock()

		// If new repo, send off to next stage.
		if repo != nil {
			w.unique <- repo
		}
	}
}

// processFilterWorker computes repository local and remote state, and filters with shouldShow.
func (w *workspace) processFilterWorker(wg *sync.WaitGroup) {
	defer wg.Done()
	for repo := range w.unique {
		w.computeVCSState(repo)

		if !w.shouldShow(repo) {
			continue
		}

		w.processedFiltered <- repo
	}
}

func (*workspace) computeVCSState(r *Repo) {
	if r.VCS == nil {
		// Go package not under VCS.
		return
	}

	if s, err := r.VCS.Status(r.Path); err == nil {
		r.Local.Status = s
	}
	if b, err := r.VCS.Branch(r.Path); err == nil {
		r.Local.Branch = b
	}
	if rev, err := r.VCS.LocalRevision(r.Path); err == nil {
		r.Local.Revision = rev
	}
	if s, err := r.VCS.Stash(r.Path); err == nil {
		r.Local.Stash = s
	}
	if remote, err := r.VCS.RemoteURL(r.Path); err == nil {
		r.Local.RemoteURL = remote
	}
	if rev, err := r.VCS.RemoteRevision(r.Path); err == nil {
		r.Remote.Revision = rev
	}
	if r.Remote.Revision != "" {
		if c, err := r.VCS.Contains(r.Path, r.Remote.Revision); err == nil {
			r.LocalContainsRemoteRevision = c
		}
	}
	if rr, err := vcs.RepoRootForImportPath(r.Root, false); err == nil {
		r.Remote.RepoURL = rr.Repo
	}
}

// presenterWorker runs presenter on processed and filtered repos.
func (w *workspace) presenterWorker(wg *sync.WaitGroup) {
	defer wg.Done()
	for repo := range w.processedFiltered {
		w.Statuses <- w.presenter(repo)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math"
	"os"

	"github.com/dustin/go-humanize"
	"github.com/go-errors/errors"
	"github.com/itchio/butler/comm"
	"github.com/itchio/wharf/pwr"
	"github.com/itchio/wharf/tlc"
	"github.com/itchio/wharf/wire"
)

func ranges(patch string) {
	must(doRanges(patch))
}

func doRanges(patch string) error {
	patchStats, err := os.Lstat(patch)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	patchReader, err := os.Open(patch)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	rawPatchWire := wire.NewReadContext(patchReader)
	err = rawPatchWire.ExpectMagic(pwr.PatchMagic)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	header := &pwr.PatchHeader{}
	err = rawPatchWire.ReadMessage(header)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	patchWire, err := pwr.DecompressWire(rawPatchWire, header.Compression)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	targetContainer := &tlc.Container{}
	err = patchWire.ReadMessage(targetContainer)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	sourceContainer := &tlc.Container{}
	err = patchWire.ReadMessage(sourceContainer)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	comm.Opf("Showing ranges for %s patch", humanize.IBytes(uint64(patchStats.Size())))
	comm.Statf("Old version: %s in %s", humanize.IBytes(uint64(targetContainer.Size)), targetContainer.Stats())
	comm.Statf("New version: %s in %s", humanize.IBytes(uint64(sourceContainer.Size)), sourceContainer.Stats())
	deltaOp := "+"
	if sourceContainer.Size < targetContainer.Size {
		deltaOp = "-"
	}
	delta := math.Abs(float64(sourceContainer.Size - targetContainer.Size))
	comm.Statf("Delta: %s%s (%s%.2f%%)", deltaOp, humanize.IBytes(uint64(delta)), deltaOp, delta/float64(targetContainer.Size)*100.0)
	comm.Log("")

	numDatas := 0
	numBlockRanges := 0
	blockSize := int64(pwr.BlockSize)
	unchangedBytes := int64(0)
	movedBytes := int64(0)
	freshBytes := int64(0)

	sh := &pwr.SyncHeader{}
	for fileIndex, sourceFile := range sourceContainer.Files {
		sh.Reset()
		err := patchWire.ReadMessage(sh)
		if err != nil {
			return errors.Wrap(err, 1)
		}

		if sh.FileIndex != int64(fileIndex) {
			fmt.Printf("expected fileIndex = %d, got fileIndex %d\n", fileIndex, sh.FileIndex)
			return errors.Wrap(pwr.ErrMalformedPatch, 1)
		}

		rop := &pwr.SyncOp{}

		err = (func() error {
			sourceOffset := int64(0)

			for {
				rop.Reset()
				pErr := patchWire.ReadMessage(rop)
				if pErr != nil {
					return errors.Wrap(pErr, 1)
				}

				switch rop.Type {
				case pwr.SyncOp_BLOCK_RANGE:
					targetOffset := blockSize * rop.BlockIndex
					targetFile := targetContainer.Files[rop.FileIndex]

					size := blockSize * rop.BlockSpan

					alignedSize := blockSize * (rop.BlockIndex + rop.BlockSpan)
					if alignedSize > targetFile.Size {
						size -= blockSize
						size += targetFile.Size % blockSize
					}

					if targetFile.Path == sourceFile.Path && targetOffset == sourceOffset {
						// comm.Statf("%d unchanged blocks %d bytes into %s", rop.BlockSpan, sourceOffset, targetFile.Path)
						unchangedBytes += size
					} else {
						movedBytes += size
					}

					numBlockRanges++
					sourceOffset += size
				case pwr.SyncOp_DATA:
					size := int64(len(rop.Data))
					sourceOffset += size
					freshBytes += size
					numDatas++
				case pwr.SyncOp_HEY_YOU_DID_IT:
					return nil
				}
			}
		})()
		if err != nil {
			return errors.Wrap(err, 1)
		}
	}

	comm.Statf("%d BlockRange ops, %d Data ops", numBlockRanges, numDatas)
	comm.Statf("Unchanged bytes: %s", humanize.IBytes(uint64(unchangedBytes)))
	comm.Statf("Moved bytes    : %s", humanize.IBytes(uint64(movedBytes)))
	comm.Statf("Fresh bytes    : %s", humanize.IBytes(uint64(freshBytes)))

	return nil
}
<commit_msg>Try out FakePool<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"math"
	"os"

	"github.com/dustin/go-humanize"
	"github.com/go-errors/errors"
	"github.com/itchio/butler/comm"
	"github.com/itchio/wharf/pwr"
	"github.com/itchio/wharf/sync"
	"github.com/itchio/wharf/tlc"
	"github.com/itchio/wharf/wire"
)

func ranges(patch string) {
	must(doRanges(patch))
}

type FakePool struct {
	container *tlc.Container
}

func (fp *FakePool) GetReader(fileIndex int64) (io.Reader, error) {
	return fp.GetReadSeeker(fileIndex)
}

func (fp *FakePool) GetReadSeeker(fileIndex int64) (io.ReadSeeker, error) {
	buf := make([]byte, fp.container.Files[fileIndex].Size)
	return bytes.NewReader(buf), nil
}

func (fp *FakePool) Close() error {
	return nil
}

var _ sync.FilePool = (*FakePool)(nil)

func doRanges(patch string) error {
	patchStats, err := os.Lstat(patch)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	patchReader, err := os.Open(patch)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	rawPatchWire := wire.NewReadContext(patchReader)
	err = rawPatchWire.ExpectMagic(pwr.PatchMagic)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	header := &pwr.PatchHeader{}
	err = rawPatchWire.ReadMessage(header)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	patchWire, err := pwr.DecompressWire(rawPatchWire, header.Compression)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	targetContainer := &tlc.Container{}
	err = patchWire.ReadMessage(targetContainer)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	sourceContainer := &tlc.Container{}
	err = patchWire.ReadMessage(sourceContainer)
	if err != nil {
		return errors.Wrap(err, 1)
	}

	comm.Opf("Showing ranges for %s patch", humanize.IBytes(uint64(patchStats.Size())))
	comm.Statf("Old version: %s in %s", humanize.IBytes(uint64(targetContainer.Size)), targetContainer.Stats())
	comm.Statf("New version: %s in %s", humanize.IBytes(uint64(sourceContainer.Size)), sourceContainer.Stats())
	deltaOp := "+"
	if sourceContainer.Size < targetContainer.Size {
		deltaOp = "-"
	}
	delta := math.Abs(float64(sourceContainer.Size - targetContainer.Size))
	comm.Statf("Delta: %s%s (%s%.2f%%)", deltaOp, humanize.IBytes(uint64(delta)), deltaOp, delta/float64(targetContainer.Size)*100.0)
	comm.Log("")

	numDatas := 0
	numBlockRanges := 0
	blockSize := int64(pwr.BlockSize)
	bigBlockSize := int64(4 * 1024 * 1024) // 4MB

	unchangedBytes := int64(0)
	movedBytes := int64(0)
	freshBytes := int64(0)

	requiredOldBlocks := make([]map[int64]bool, len(targetContainer.Files))
	for i := 0; i < len(targetContainer.Files); i++ {
		requiredOldBlocks[i] = make(map[int64]bool)
	}

	sh := &pwr.SyncHeader{}
	for fileIndex, sourceFile := range sourceContainer.Files {
		sh.Reset()
		err := patchWire.ReadMessage(sh)
		if err != nil {
			return errors.Wrap(err, 1)
		}

		if sh.FileIndex != int64(fileIndex) {
			fmt.Printf("expected fileIndex = %d, got fileIndex %d\n", fileIndex, sh.FileIndex)
			return errors.Wrap(pwr.ErrMalformedPatch, 1)
		}

		rop := &pwr.SyncOp{}

		err = (func() error {
			sourceOffset := int64(0)
			numMoved := 0
			numUnchanged := 0
			numFresh := 0

			for {
				rop.Reset()
				pErr := patchWire.ReadMessage(rop)
				if pErr != nil {
					return errors.Wrap(pErr, 1)
				}

				switch rop.Type {
				case pwr.SyncOp_BLOCK_RANGE:
					targetOffset := blockSize * rop.BlockIndex
					targetFile := targetContainer.Files[rop.FileIndex]

					size := blockSize * rop.BlockSpan

					alignedSize := blockSize * (rop.BlockIndex + rop.BlockSpan)
					// op includes last block which is smaller than blockSize
					if alignedSize > targetFile.Size {
						size -= blockSize
						size += targetFile.Size % blockSize
					}

					if targetFile.Path == sourceFile.Path && targetOffset == sourceOffset {
						// comm.Statf("%d unchanged blocks %d bytes into %s", rop.BlockSpan, sourceOffset, targetFile.Path)
						unchangedBytes += size
						numUnchanged++
					} else {
						movedBytes += size
						numMoved++

						bigBlockStart := int64(math.Floor(float64(rop.BlockIndex*blockSize) / float64(bigBlockSize)))
						bigBlockEnd := int64(math.Ceil(float64(rop.BlockIndex*blockSize+size) / float64(bigBlockSize)))

						for i := bigBlockStart; i < bigBlockEnd; i++ {
							requiredOldBlocks[rop.FileIndex][i] = true
						}
					}

					numBlockRanges++
					sourceOffset += size
				case pwr.SyncOp_DATA:
					size := int64(len(rop.Data))
					sourceOffset += size
					freshBytes += size
					numDatas++
					numFresh++
				case pwr.SyncOp_HEY_YOU_DID_IT:
					if numFresh == 0 && numUnchanged == 0 && numMoved == 1 {
						comm.Statf("Found rename!")
					} else {
						// comm.Statf("Fresh %d, unchanged %d, moved %d", numFresh, numUnchanged, numMoved)
					}
					return nil
				}
			}
		})()
		if err != nil {
			return errors.Wrap(err, 1)
		}
	}

	comm.Statf("%d BlockRange ops, %d Data ops", numBlockRanges, numDatas)
	comm.Statf("Unchanged bytes: %s", humanize.IBytes(uint64(unchangedBytes)))
	comm.Statf("Moved bytes    : %s", humanize.IBytes(uint64(movedBytes)))
	comm.Statf("Fresh bytes    : %s", humanize.IBytes(uint64(freshBytes)))

	comm.Log("")

	totalBlocks := 0
	partialBlocks := 0
	neededBlocks := 0
	neededBlockSize := int64(0)

	for i, blockMap := range requiredOldBlocks {
		f := targetContainer.Files[i]
		fileNumBlocks := int64(math.Ceil(float64(f.Size) / float64(bigBlockSize)))
		for j := int64(0); j < fileNumBlocks; j++ {
			totalBlocks++
			if blockMap[j] {
				size := bigBlockSize
				if (j+1)*bigBlockSize > f.Size {
					partialBlocks++
					size = f.Size % bigBlockSize
				}
				neededBlockSize += size
				neededBlocks++
			}
		}
	}
	comm.Statf("Total old blocks: %d, needed: %d (of which %d are smaller than %s)", totalBlocks, neededBlocks, partialBlocks, humanize.IBytes(uint64(bigBlockSize)))
	comm.Statf("Needed block size: %s (%.2f%% of full old build size)", humanize.IBytes(uint64(neededBlockSize)), float64(neededBlockSize)/float64(targetContainer.Size)*100.0)

	outputPath := "out"

	fakePool := &FakePool{
		container: targetContainer,
	}

	actx := &pwr.ApplyContext{
		Consumer:   comm.NewStateConsumer(),
		OutputPath: outputPath,
		TargetPool: fakePool,
	}

	_, err = patchReader.Seek(0, os.SEEK_SET)
	if err != nil {
		return err
	}

	err = actx.ApplyPatch(patchReader)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package run

import (
	"flag"

	clibazil "bazil.org/bazil/cli"
	"bazil.org/bazil/cliutil/subcommands"
	"bazil.org/bazil/control"
	"bazil.org/bazil/server"
)

type runCommand struct {
	subcommands.Description
	flag.FlagSet
}

func (cmd *runCommand) Run() error {
	app, err := server.New(clibazil.Bazil.Config.DataDir.String())
	if err != nil {
		return err
	}
	defer app.Close()
	c := control.New(app)
	return c.ListenAndServe()
}

var run = runCommand{
	Description: "run bazil server",
}

func init() {
	subcommands.Register(&run)
}
<commit_msg>cli/server/run: Use all CPU cores by default<commit_after>package run

import (
	"flag"

	clibazil "bazil.org/bazil/cli"
	"bazil.org/bazil/cliutil/subcommands"
	"bazil.org/bazil/control"
	"bazil.org/bazil/server"
	"github.com/cespare/gomaxprocs"
)

type runCommand struct {
	subcommands.Description
	flag.FlagSet
}

func (cmd *runCommand) Run() error {
	gomaxprocs.SetToNumCPU()
	app, err := server.New(clibazil.Bazil.Config.DataDir.String())
	if err != nil {
		return err
	}
	defer app.Close()
	c := control.New(app)
	return c.ListenAndServe()
}

var run = runCommand{
	Description: "run bazil server",
}

func init() {
	subcommands.Register(&run)
}
<|endoftext|>


<commit_before>package rollbar

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"runtime"
	"time"
)

type Client struct {
	AppName    string
	AppVersion string
	Endpoint   string
	Token      string
}

func (c *Client) Report(message string) error {
	body := c.buildBody()
	data := body["data"].(map[string]interface{})
	data["body"] = c.errorBody(message, 2)

	jsonBody, err := json.Marshal(body)
	if err != nil {
		return err
	}

	resp, err := http.Post(c.Endpoint, "application/json", bytes.NewReader(jsonBody))
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode/100 != 2 { // 200, 201, 202, etc
		return fmt.Errorf("unexpected status code %d", resp.StatusCode)
	}
	return nil
}

func (c *Client) errorBody(message string, skip int) map[string]interface{} {
	return map[string]interface{}{
		"trace": map[string]interface{}{
			"frames": c.stacktraceFrames(3 + skip),
			"exception": map[string]interface{}{
				"class":   "",
				"message": message,
			},
		},
	}
}

func (c *Client) stacktraceFrames(skip int) []map[string]interface{} {
	frames := []map[string]interface{}{}

	for i := skip; ; i++ {
		pc, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		f := runtime.FuncForPC(pc)
		fname := "unknown"
		if f != nil {
			fname = f.Name()
		}
		frames = append(frames, map[string]interface{}{
			"filename": file,
			"lineno":   line,
			"method:":  fname,
		})
	}
	return frames
}

func (c *Client) buildBody() map[string]interface{} {
	timestamp := time.Now().UTC().Unix()

	return map[string]interface{}{
		"access_token": c.Token,
		"data": map[string]interface{}{
			"environment": "production",
			"timestamp":   timestamp,
			"platform":    "client",
			"language":    "go",
			"notifier": map[string]interface{}{
				"name":    c.AppName,
				"version": c.AppVersion,
			},
			"client": map[string]interface{}{
				"arch":     runtime.GOARCH,
				"platform": runtime.GOOS,
			},
		},
	}
}
<commit_msg>Remove rollbar.<commit_after><|endoftext|>


<commit_before>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package client

import (
	"bytes"
	"context"
	"fmt"
	"testing"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"

	"github.com/google/trillian"
	"github.com/google/trillian/testonly"
	"github.com/google/trillian/testonly/integration"
)

const timeout = 100 * time.Millisecond

func TestAddGetLeaf(t *testing.T) {
	// TODO: Build a GetLeaf method and test a full get/set cycle.
}

// addSequencedLeaves is a temporary stand-in function for tests until the real API gets built.
func addSequencedLeaves(ctx context.Context, env *integration.LogEnv, client VerifyingLogClient, leaves [][]byte) error {
	// TODO(gdbelvin): Replace with batch API.
	// TODO(gdbelvin): Replace with AddSequencedLeaves API.
	for _, l := range leaves {
		ctx, cancel := context.WithDeadline(ctx, time.Now().Add(timeout))
		defer cancel()
		if err, want := client.AddLeaf(ctx, l), codes.DeadlineExceeded; grpc.Code(err) != want {
			return fmt.Errorf("AddLeaf(%v): %v, want, %v", l, err, want)
		}
	}
	return nil
}

func TestGetByIndex(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestGetByIndex")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}

	cli := trillian.NewTrillianLogClient(env.ClientConn)
	client := New(logID, cli, testonly.Hasher, env.PublicKey)
	// Add a few test leaves.
	leafData := [][]byte{
		[]byte("A"),
		[]byte("B"),
		[]byte("C"),
	}

	if err := addSequencedLeaves(ctx, env, client, leafData); err != nil {
		t.Errorf("Failed to add leaves: %v", err)
	}

	for i, l := range leafData {
		leaf, err := client.GetByIndex(ctx, int64(i))
		if err != nil {
			t.Errorf("Failed to GetByIndex(%v): %v", i, err)
			continue
		}
		if got, want := leaf.LeafValue, l; !bytes.Equal(got, want) {
			t.Errorf("GetByIndex(%v) = %x, want %x", i, got, want)
		}
	}
}

func TestListByIndex(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestGetByIndex")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}

	cli := trillian.NewTrillianLogClient(env.ClientConn)
	client := New(logID, cli, testonly.Hasher, env.PublicKey)
	// Add a few test leaves.
	leafData := [][]byte{
		[]byte("A"),
		[]byte("B"),
		[]byte("C"),
	}

	if err := addSequencedLeaves(ctx, env, client, leafData); err != nil {
		t.Errorf("Failed to add leaves: %v", err)
	}

	// Fetch leaves.
	leaves, err := client.ListByIndex(ctx, 0, 3)
	if err != nil {
		t.Errorf("Failed to ListByIndex: %v", err)
	}
	for i, l := range leaves {
		if got, want := l.LeafValue, leafData[i]; !bytes.Equal(got, want) {
			t.Errorf("ListIndex()[%v] = %v, want %v", i, got, want)
		}
	}
}

func TestVerifyInclusion(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestVerifyInclusion")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}

	cli := trillian.NewTrillianLogClient(env.ClientConn)
	client := New(logID, cli, testonly.Hasher, env.PublicKey)
	// Add a few test leaves.
	leafData := [][]byte{
		[]byte("A"),
		[]byte("B"),
	}

	if err := addSequencedLeaves(ctx, env, client, leafData); err != nil {
		t.Errorf("Failed to add leaves: %v", err)
	}

	for _, l := range leafData {
		if err := client.VerifyInclusion(ctx, l); err != nil {
			t.Errorf("VerifyInclusion(%s) = %v, want nil", l, err)
		}
	}
}

func TestVerifyInclusionAtIndex(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestVerifyInclusionAtIndex")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}

	cli := trillian.NewTrillianLogClient(env.ClientConn)
	client := New(logID, cli, testonly.Hasher, env.PublicKey)
	// Add a few test leaves.
	leafData := [][]byte{
		[]byte("A"),
		[]byte("B"),
	}

	if err := addSequencedLeaves(ctx, env, client, leafData); err != nil {
		t.Errorf("Failed to add leaves: %v", err)
	}

	for i, l := range leafData {
		if err := client.VerifyInclusionAtIndex(ctx, l, int64(i)); err != nil {
			t.Errorf("VerifyInclusion(%s) = %v, want nil", l, err)
		}
	}
}

func TestAddLeaf(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestAddLeaf")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}

	cli := trillian.NewTrillianLogClient(env.ClientConn)
	for _, test := range []struct {
		desc    string
		client  trillian.TrillianLogClient
		wantErr bool
	}{
		{
			desc:   "success 1",
			client: cli,
		},
		{
			desc:   "success 2",
			client: cli,
		},
		{
			desc:    "invalid inclusion proof",
			client:  &MockLogClient{c: cli, mGetInclusionProof: true},
			wantErr: true,
		},
		{
			desc:    "invalid consistency proof",
			client:  &MockLogClient{c: cli, mGetConsistencyProof: true},
			wantErr: true,
		},
	} {
		client := New(logID, test.client, testonly.Hasher, env.PublicKey)
		{
			ctx, cancel := context.WithDeadline(ctx, time.Now().Add(timeout))
			defer cancel()
			if err, want := client.AddLeaf(ctx, []byte(test.desc)), codes.DeadlineExceeded; grpc.Code(err) != want {
				t.Errorf("AddLeaf(%v): %v, want, %v", test.desc, err, want)
				continue
			}
		}
		err := client.AddLeaf(ctx, []byte(test.desc))
		if got := err != nil; got != test.wantErr {
			t.Errorf("AddLeaf(%v): %v, want error: %v", test.desc, err, test.wantErr)
		}
	}
}

func TestUpdateRoot(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestUpdateRoot")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}
	cli := trillian.NewTrillianLogClient(env.ClientConn)
	client := New(logID, cli, testonly.Hasher, env.PublicKey)

	before := client.Root().TreeSize

	{
		ctx, cancel := context.WithDeadline(ctx, time.Now().Add(timeout))
		defer cancel()
		if err, want := client.AddLeaf(ctx, []byte("foo")), codes.DeadlineExceeded; grpc.Code(err) != want {
			t.Errorf("AddLeaf(): %v, want, %v", err, want)
		}
	}

	if err := client.UpdateRoot(ctx); err != nil {
		t.Error(err)
	}
	if got, want := client.Root().TreeSize, before; got <= want {
		t.Errorf("Tree size after add Leaf: %v, want > %v", got, want)
	}
}
<commit_msg>client: drop unused parameter from test fn<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package client

import (
	"bytes"
	"context"
	"fmt"
	"testing"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"

	"github.com/google/trillian"
	"github.com/google/trillian/testonly"
	"github.com/google/trillian/testonly/integration"
)

const timeout = 100 * time.Millisecond

func TestAddGetLeaf(t *testing.T) {
	// TODO: Build a GetLeaf method and test a full get/set cycle.
}

// addSequencedLeaves is a temporary stand-in function for tests until the real API gets built.
func addSequencedLeaves(ctx context.Context, client VerifyingLogClient, leaves [][]byte) error {
	// TODO(gdbelvin): Replace with batch API.
	// TODO(gdbelvin): Replace with AddSequencedLeaves API.
	for _, l := range leaves {
		ctx, cancel := context.WithDeadline(ctx, time.Now().Add(timeout))
		defer cancel()
		if err, want := client.AddLeaf(ctx, l), codes.DeadlineExceeded; grpc.Code(err) != want {
			return fmt.Errorf("AddLeaf(%v): %v, want, %v", l, err, want)
		}
	}
	return nil
}

func TestGetByIndex(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestGetByIndex")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}

	cli := trillian.NewTrillianLogClient(env.ClientConn)
	client := New(logID, cli, testonly.Hasher, env.PublicKey)
	// Add a few test leaves.
	leafData := [][]byte{
		[]byte("A"),
		[]byte("B"),
		[]byte("C"),
	}

	if err := addSequencedLeaves(ctx, client, leafData); err != nil {
		t.Errorf("Failed to add leaves: %v", err)
	}

	for i, l := range leafData {
		leaf, err := client.GetByIndex(ctx, int64(i))
		if err != nil {
			t.Errorf("Failed to GetByIndex(%v): %v", i, err)
			continue
		}
		if got, want := leaf.LeafValue, l; !bytes.Equal(got, want) {
			t.Errorf("GetByIndex(%v) = %x, want %x", i, got, want)
		}
	}
}

func TestListByIndex(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestGetByIndex")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}

	cli := trillian.NewTrillianLogClient(env.ClientConn)
	client := New(logID, cli, testonly.Hasher, env.PublicKey)
	// Add a few test leaves.
	leafData := [][]byte{
		[]byte("A"),
		[]byte("B"),
		[]byte("C"),
	}

	if err := addSequencedLeaves(ctx, client, leafData); err != nil {
		t.Errorf("Failed to add leaves: %v", err)
	}

	// Fetch leaves.
	leaves, err := client.ListByIndex(ctx, 0, 3)
	if err != nil {
		t.Errorf("Failed to ListByIndex: %v", err)
	}
	for i, l := range leaves {
		if got, want := l.LeafValue, leafData[i]; !bytes.Equal(got, want) {
			t.Errorf("ListIndex()[%v] = %v, want %v", i, got, want)
		}
	}
}

func TestVerifyInclusion(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestVerifyInclusion")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}

	cli := trillian.NewTrillianLogClient(env.ClientConn)
	client := New(logID, cli, testonly.Hasher, env.PublicKey)
	// Add a few test leaves.
	leafData := [][]byte{
		[]byte("A"),
		[]byte("B"),
	}

	if err := addSequencedLeaves(ctx, client, leafData); err != nil {
		t.Errorf("Failed to add leaves: %v", err)
	}

	for _, l := range leafData {
		if err := client.VerifyInclusion(ctx, l); err != nil {
			t.Errorf("VerifyInclusion(%s) = %v, want nil", l, err)
		}
	}
}

func TestVerifyInclusionAtIndex(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestVerifyInclusionAtIndex")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}

	cli := trillian.NewTrillianLogClient(env.ClientConn)
	client := New(logID, cli, testonly.Hasher, env.PublicKey)
	// Add a few test leaves.
	leafData := [][]byte{
		[]byte("A"),
		[]byte("B"),
	}

	if err := addSequencedLeaves(ctx, client, leafData); err != nil {
		t.Errorf("Failed to add leaves: %v", err)
	}

	for i, l := range leafData {
		if err := client.VerifyInclusionAtIndex(ctx, l, int64(i)); err != nil {
			t.Errorf("VerifyInclusion(%s) = %v, want nil", l, err)
		}
	}
}

func TestAddLeaf(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestAddLeaf")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}

	cli := trillian.NewTrillianLogClient(env.ClientConn)
	for _, test := range []struct {
		desc    string
		client  trillian.TrillianLogClient
		wantErr bool
	}{
		{
			desc:   "success 1",
			client: cli,
		},
		{
			desc:   "success 2",
			client: cli,
		},
		{
			desc:    "invalid inclusion proof",
			client:  &MockLogClient{c: cli, mGetInclusionProof: true},
			wantErr: true,
		},
		{
			desc:    "invalid consistency proof",
			client:  &MockLogClient{c: cli, mGetConsistencyProof: true},
			wantErr: true,
		},
	} {
		client := New(logID, test.client, testonly.Hasher, env.PublicKey)
		{
			ctx, cancel := context.WithDeadline(ctx, time.Now().Add(timeout))
			defer cancel()
			if err, want := client.AddLeaf(ctx, []byte(test.desc)), codes.DeadlineExceeded; grpc.Code(err) != want {
				t.Errorf("AddLeaf(%v): %v, want, %v", test.desc, err, want)
				continue
			}
		}
		err := client.AddLeaf(ctx, []byte(test.desc))
		if got := err != nil; got != test.wantErr {
			t.Errorf("AddLeaf(%v): %v, want error: %v", test.desc, err, test.wantErr)
		}
	}
}

func TestUpdateRoot(t *testing.T) {
	ctx := context.Background()
	env, err := integration.NewLogEnv(ctx, 1, "TestUpdateRoot")
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()
	logID, err := env.CreateLog()
	if err != nil {
		t.Fatalf("Failed to create log: %v", err)
	}
	cli := trillian.NewTrillianLogClient(env.ClientConn)
	client := New(logID, cli, testonly.Hasher, env.PublicKey)

	before := client.Root().TreeSize

	{
		ctx, cancel := context.WithDeadline(ctx, time.Now().Add(timeout))
		defer cancel()
		if err, want := client.AddLeaf(ctx, []byte("foo")), codes.DeadlineExceeded; grpc.Code(err) != want {
			t.Errorf("AddLeaf(): %v, want, %v", err, want)
		}
	}

	if err := client.UpdateRoot(ctx); err != nil {
		t.Error(err)
	}
	if got, want := client.Root().TreeSize, before; got <= want {
		t.Errorf("Tree size after add Leaf: %v, want > %v", got, want)
	}
}
<|endoftext|>


<commit_before>package client_test

import (
	"encoding/json"
	"github.com/nelhage/livegrep/client"
	"io"
	. "launchpad.net/gocheck"
	"net"
	"testing"
)

// We assume a codesearch running on localhost:9999. This could be
// improved.

func Test(t *testing.T) { TestingT(t) }

type ClientSuite struct {
	client client.Client
}

func (c *ClientSuite) TearDownTest(*C) {
	if c.client != nil {
		c.client.Close()
	}
}

var _ = Suite(&ClientSuite{})

type MockServer struct {
	Info    *client.ServerInfo
	Results []*client.Result
}

func (m *MockServer) handle(conn net.Conn) {
	defer conn.Close()

	encoder := json.NewEncoder(conn)
	reader := json.NewDecoder(conn)
	for {
		io.WriteString(conn, "READY ")
		encoder.Encode(m.Info)

		var q client.Query
		if err := reader.Decode(&q); err != nil {
			if err == io.EOF {
				return
			}
			panic(err.Error())
		}

		for _, r := range m.Results {
			encoder.Encode(r)
		}
		io.WriteString(conn, "DONE ")
		encoder.Encode(&client.Stats{})
	}
}

func runMockServer(handle func(net.Conn)) <-chan string {
	ready := make(chan string, 1)
	go func() {
		ln, err := net.Listen("tcp", ":0")
		if err != nil {
			panic(err.Error())
		}
		defer ln.Close()
		ready <- ln.Addr().String()
		conn, err := ln.Accept()
		if err != nil {
			panic(err.Error())
		}
		handle(conn)
	}()
	return ready
}

func (s *ClientSuite) connect(c *C, addr string) {
	var err error
	s.client, err = client.Dial("tcp", addr)
	if err != nil {
		c.Fatalf("connecting to %s: %s", addr, err.Error())
	}
}

func (s *ClientSuite) TestQuery(c *C) {
	s.connect(c, <-runMockServer((&MockServer{
		Results: []*client.Result{
			{Line: "match line 1"},
		},
	}).handle))
	search, err := s.client.Query(&client.Query{".", "", ""})
	c.Assert(err, IsNil)
	var n int
	for r := range search.Results() {
		n++
		c.Assert(r.Line, Not(Equals), "")
	}
	c.Assert(n, Equals, 1)
	st, e := search.Close()
	c.Assert(e, IsNil)
	c.Assert(st, Not(IsNil))
}

func (s *ClientSuite) TestTwoQueries(c *C) {
	s.connect(c, <-runMockServer((&MockServer{
		Results: []*client.Result{
			{Line: "match line 1"},
		},
	}).handle))

	search, err := s.client.Query(&client.Query{".", "", ""})
	c.Assert(err, IsNil)
	_, err = search.Close()
	c.Assert(err, IsNil)

	search, err = s.client.Query(&client.Query{".", "", ""})
	c.Assert(err, IsNil)
	n := 0
	for _ = range search.Results() {
		n++
	}
	_, err = search.Close()
	if err != nil {
		c.Fatalf("Unexpected error: %s", err.Error())
	}
	c.Assert(n, Not(Equals), 0)
}

type MockServerQueryError struct {
	Info *client.ServerInfo
	Err  string
}

func (m *MockServerQueryError) handle(conn net.Conn) {
	defer conn.Close()
	encoder := json.NewEncoder(conn)
	reader := json.NewDecoder(conn)
	for {
		io.WriteString(conn, "READY ")
		encoder.Encode(m.Info)

		var q client.Query
		if err := reader.Decode(&q); err != nil {
			if err == io.EOF {
				return
			}
			panic(err.Error())
		}

		io.WriteString(conn, "FATAL ")
		io.WriteString(conn, m.Err)
		io.WriteString(conn, "\n")
	}
}

func (s *ClientSuite) TestBadRegex(c *C) {
	errStr := "Invalid query: ("
	s.connect(c, <-runMockServer((&MockServerQueryError{
		Err: errStr,
	}).handle))

	search, err := s.client.Query(&client.Query{"(", "", ""})
	c.Assert(err, IsNil)
	for _ = range search.Results() {
		c.Fatal("Got back a result from an erroneous query!")
	}
	st, e := search.Close()
	c.Assert(st, IsNil)
	if e == nil {
		c.Fatal("Didn't get back an error")
	}
	if q, ok := e.(client.QueryError); ok {
		c.Assert(q.Query.Line, Equals, "(")
		c.Assert(q.Err, Equals, errStr)
	} else {
		c.Fatalf("Error %v wasn't a QueryError", e)
	}
}

func mockServerShutdown() <-chan string {
	return runMockServer(func(conn net.Conn) {
		conn.Write([]byte("READY {}\n"))
		conn.Close()
	})
}

func (s *ClientSuite) TestShutdown(c *C) {
	addr := <-mockServerShutdown()

	s.connect(c, addr)

	search, err := s.client.Query(&client.Query{Line: "l"})
	c.Assert(err, IsNil)
	c.Assert(search, Not(IsNil))

	results := search.Results()
	c.Assert(results, Not(IsNil))
	for r := range results {
		c.Errorf("Got a result back: %+v", r)
	}
	st, err := search.Close()
	c.Assert(st, IsNil)
	c.Assert(err, Not(IsNil))

	search, err = s.client.Query(&client.Query{Line: "l"})
	c.Assert(err, Not(IsNil))
	c.Assert(search, IsNil)
}
<commit_msg>(failing) test for handling really long lines<commit_after>package client_test

import (
	"encoding/json"
	"github.com/nelhage/livegrep/client"
	"io"
	. "launchpad.net/gocheck"
	"net"
	"strings"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type ClientSuite struct {
	client client.Client
}

func (c *ClientSuite) TearDownTest(*C) {
	if c.client != nil {
		c.client.Close()
	}
}

var _ = Suite(&ClientSuite{})

type MockServer struct {
	Info    *client.ServerInfo
	Results []*client.Result
}

func (m *MockServer) handle(conn net.Conn) {
	defer conn.Close()

	encoder := json.NewEncoder(conn)
	reader := json.NewDecoder(conn)
	for {
		io.WriteString(conn, "READY ")
		encoder.Encode(m.Info)

		var q client.Query
		if err := reader.Decode(&q); err != nil {
			if err == io.EOF {
				return
			}
			panic(err.Error())
		}

		for _, r := range m.Results {
			encoder.Encode(r)
		}
		io.WriteString(conn, "DONE ")
		encoder.Encode(&client.Stats{})
	}
}

func runMockServer(handle func(net.Conn)) <-chan string {
	ready := make(chan string, 1)
	go func() {
		ln, err := net.Listen("tcp", ":0")
		if err != nil {
			panic(err.Error())
		}
		defer ln.Close()
		ready <- ln.Addr().String()
		conn, err := ln.Accept()
		if err != nil {
			panic(err.Error())
		}
		handle(conn)
	}()
	return ready
}

func (s *ClientSuite) connect(c *C, addr string) {
	var err error
	s.client, err = client.Dial("tcp", addr)
	if err != nil {
		c.Fatalf("connecting to %s: %s", addr, err.Error())
	}
}

func (s *ClientSuite) TestQuery(c *C) {
	s.connect(c, <-runMockServer((&MockServer{
		Results: []*client.Result{
			{Line: "match line 1"},
		},
	}).handle))
	search, err := s.client.Query(&client.Query{".", "", ""})
	c.Assert(err, IsNil)
	var n int
	for r := range search.Results() {
		n++
		c.Assert(r.Line, Not(Equals), "")
	}
	c.Assert(n, Equals, 1)
	st, e := search.Close()
	c.Assert(e, IsNil)
	c.Assert(st, Not(IsNil))
}

func (s *ClientSuite) TestTwoQueries(c *C) {
	s.connect(c, <-runMockServer((&MockServer{
		Results: []*client.Result{
			{Line: "match line 1"},
		},
	}).handle))

	search, err := s.client.Query(&client.Query{".", "", ""})
	c.Assert(err, IsNil)
	_, err = search.Close()
	c.Assert(err, IsNil)

	search, err = s.client.Query(&client.Query{".", "", ""})
	c.Assert(err, IsNil)
	n := 0
	for _ = range search.Results() {
		n++
	}
	_, err = search.Close()
	if err != nil {
		c.Fatalf("Unexpected error: %s", err.Error())
	}
	c.Assert(n, Not(Equals), 0)
}

func (s *ClientSuite) TestLongLine(c *C) {
	longLine := strings.Repeat("X", 1<<16)
	s.connect(c, <-runMockServer((&MockServer{
		Results: []*client.Result{
			{Line: longLine},
		},
	}).handle))
	search, err := s.client.Query(&client.Query{".", "", ""})
	c.Assert(err, IsNil)
	var rs []*client.Result
	for r := range search.Results() {
		rs = append(rs, r)
	}

	c.Assert(len(rs), Equals, 1)
	c.Assert(rs[0].Line, Equals, longLine)
}

type MockServerQueryError struct {
	Info *client.ServerInfo
	Err  string
}

func (m *MockServerQueryError) handle(conn net.Conn) {
	defer conn.Close()
	encoder := json.NewEncoder(conn)
	reader := json.NewDecoder(conn)
	for {
		io.WriteString(conn, "READY ")
		encoder.Encode(m.Info)

		var q client.Query
		if err := reader.Decode(&q); err != nil {
			if err == io.EOF {
				return
			}
			panic(err.Error())
		}

		io.WriteString(conn, "FATAL ")
		io.WriteString(conn, m.Err)
		io.WriteString(conn, "\n")
	}
}

func (s *ClientSuite) TestBadRegex(c *C) {
	errStr := "Invalid query: ("
	s.connect(c, <-runMockServer((&MockServerQueryError{
		Err: errStr,
	}).handle))

	search, err := s.client.Query(&client.Query{"(", "", ""})
	c.Assert(err, IsNil)
	for _ = range search.Results() {
		c.Fatal("Got back a result from an erroneous query!")
	}
	st, e := search.Close()
	c.Assert(st, IsNil)
	if e == nil {
		c.Fatal("Didn't get back an error")
	}
	if q, ok := e.(client.QueryError); ok {
		c.Assert(q.Query.Line, Equals, "(")
		c.Assert(q.Err, Equals, errStr)
	} else {
		c.Fatalf("Error %v wasn't a QueryError", e)
	}
}

func mockServerShutdown() <-chan string {
	return runMockServer(func(conn net.Conn) {
		conn.Write([]byte("READY {}\n"))
		conn.Close()
	})
}

func (s *ClientSuite) TestShutdown(c *C) {
	addr := <-mockServerShutdown()

	s.connect(c, addr)

	search, err := s.client.Query(&client.Query{Line: "l"})
	c.Assert(err, IsNil)
	c.Assert(search, Not(IsNil))

	results := search.Results()
	c.Assert(results, Not(IsNil))
	for r := range results {
		c.Errorf("Got a result back: %+v", r)
	}
	st, err := search.Close()
	c.Assert(st, IsNil)
	c.Assert(err, Not(IsNil))

	search, err = s.client.Query(&client.Query{Line: "l"})
	c.Assert(err, Not(IsNil))
	c.Assert(search, IsNil)
}
<|endoftext|>


<commit_before>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"

	"github.com/docker/docker-registry/common/testutil"
	"github.com/docker/docker-registry/digest"
	"github.com/docker/docker-registry/storage"
)

type testBlob struct {
	digest   digest.Digest
	contents []byte
}

func TestPush(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	uploadLocations := make([]string, len(testBlobs))
	blobs := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, blob := range testBlobs {
		// TODO(bbland): this is returning the same location for all uploads,
		// because we can't know which blob will get which location.
		// It's sort of okay because we're using unique digests, but this needs
		// to change at some point.
		uploadLocations[i] = fmt.Sprintf("/v2/%s/blob/test-uuid", name)
		blobs[i] = storage.FSLayer{BlobSum: blob.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: blob.digest.String()}
	}

	manifest := &storage.SignedManifest{
		Manifest: storage.Manifest{
			Name:         name,
			Tag:          tag,
			Architecture: "x86",
			FSLayers:     blobs,
			History:      history,
			Versioned: storage.Versioned{
				SchemaVersion: 1,
			},
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	blobRequestResponseMappings := make([]testutil.RequestResponseMapping, 2*len(testBlobs))
	for i, blob := range testBlobs {
		blobRequestResponseMappings[2*i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "POST",
				Route:  "/v2/" + name + "/blob/upload/",
			},
			Response: testutil.Response{
				StatusCode: http.StatusAccepted,
				Headers: http.Header(map[string][]string{
					"Location": {uploadLocations[i]},
				}),
			},
		}
		blobRequestResponseMappings[2*i+1] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "PUT",
				Route:  uploadLocations[i],
				QueryParams: map[string][]string{
					"length": {fmt.Sprint(len(blob.contents))},
					"digest": {blob.digest.String()},
				},
				Body: blob.contents,
			},
			Response: testutil.Response{
				StatusCode: http.StatusCreated,
			},
		}
	}

	handler := testutil.NewHandler(append(blobRequestResponseMappings, testutil.RequestResponseMap{
		testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "PUT",
				Route:  "/v2/" + name + "/manifest/" + tag,
				Body:   manifestBytes,
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
			},
		},
	}...))
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		writer, err := l.Writer()
		if err != nil {
			t.Fatal(err)
		}

		writer.SetSize(len(blob.contents))
		writer.Write(blob.contents)
		writer.Close()
	}

	objectStore.WriteManifest(name, tag, manifest)

	err = Push(client, objectStore, name, tag)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPull(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	blobs := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, blob := range testBlobs {
		blobs[i] = storage.FSLayer{BlobSum: blob.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: blob.digest.String()}
	}

	manifest := &storage.SignedManifest{
		Manifest: storage.Manifest{
			Name:         name,
			Tag:          tag,
			Architecture: "x86",
			FSLayers:     blobs,
			History:      history,
			Versioned: storage.Versioned{
				SchemaVersion: 1,
			},
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	blobRequestResponseMappings := make([]testutil.RequestResponseMapping, len(testBlobs))
	for i, blob := range testBlobs {
		blobRequestResponseMappings[i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents,
			},
		}
	}

	handler := testutil.NewHandler(append(blobRequestResponseMappings, testutil.RequestResponseMap{
		testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/manifest/" + tag,
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       manifestBytes,
			},
		},
	}...))
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	err = Pull(client, objectStore, name, tag)
	if err != nil {
		t.Fatal(err)
	}

	m, err := objectStore.Manifest(name, tag)
	if err != nil {
		t.Fatal(err)
	}

	mBytes, err := json.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}

	if string(mBytes) != string(manifestBytes) {
		t.Fatal("Incorrect manifest")
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		reader, err := l.Reader()
		if err != nil {
			t.Fatal(err)
		}
		defer reader.Close()

		blobBytes, err := ioutil.ReadAll(reader)
		if err != nil {
			t.Fatal(err)
		}

		if string(blobBytes) != string(blob.contents) {
			t.Fatal("Incorrect blob")
		}
	}
}

func TestPullResume(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	layers := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, layer := range testBlobs {
		layers[i] = storage.FSLayer{BlobSum: layer.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: layer.digest.String()}
	}

	manifest := &storage.Manifest{
		Name:         name,
		Tag:          tag,
		Architecture: "x86",
		FSLayers:     layers,
		History:      history,
		Versioned: storage.Versioned{
			SchemaVersion: 1,
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	layerRequestResponseMappings := make([]testutil.RequestResponseMapping, 2*len(testBlobs))
	for i, blob := range testBlobs {
		layerRequestResponseMappings[2*i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents[:len(blob.contents)/2],
				Headers: http.Header(map[string][]string{
					"Content-Length": {fmt.Sprint(len(blob.contents))},
				}),
			},
		}
		layerRequestResponseMappings[2*i+1] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents[len(blob.contents)/2:],
			},
		}
	}

	for i := 0; i < 3; i++ {
		layerRequestResponseMappings = append(layerRequestResponseMappings, testutil.RequestResponseMap{
			testutil.RequestResponseMapping{
				Request: testutil.Request{
					Method: "GET",
					Route:  "/v2/" + name + "/manifest/" + tag,
				},
				Response: testutil.Response{
					StatusCode: http.StatusOK,
					Body:       manifestBytes,
				},
			},
		}...)
	}

	handler := testutil.NewHandler(layerRequestResponseMappings)
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	for attempts := 0; attempts < 3; attempts++ {
		err = Pull(client, objectStore, name, tag)
		if err == nil {
			break
		}
	}

	if err != nil {
		t.Fatal(err)
	}

	m, err := objectStore.Manifest(name, tag)
	if err != nil {
		t.Fatal(err)
	}

	mBytes, err := json.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}

	if string(mBytes) != string(manifestBytes) {
		t.Fatal("Incorrect manifest")
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		reader, err := l.Reader()
		if err != nil {
			t.Fatal(err)
		}
		defer reader.Close()

		layerBytes, err := ioutil.ReadAll(reader)
		if err != nil {
			t.Fatal(err)
		}

		if string(layerBytes) != string(blob.contents) {
			t.Fatal("Incorrect blob")
		}
	}
}
<commit_msg>Address go vet declaration issue in tests<commit_after>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"

	"github.com/docker/docker-registry/common/testutil"
	"github.com/docker/docker-registry/digest"
	"github.com/docker/docker-registry/storage"
)

type testBlob struct {
	digest   digest.Digest
	contents []byte
}

func TestPush(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	uploadLocations := make([]string, len(testBlobs))
	blobs := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, blob := range testBlobs {
		// TODO(bbland): this is returning the same location for all uploads,
		// because we can't know which blob will get which location.
		// It's sort of okay because we're using unique digests, but this needs
		// to change at some point.
		uploadLocations[i] = fmt.Sprintf("/v2/%s/blob/test-uuid", name)
		blobs[i] = storage.FSLayer{BlobSum: blob.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: blob.digest.String()}
	}

	manifest := &storage.SignedManifest{
		Manifest: storage.Manifest{
			Name:         name,
			Tag:          tag,
			Architecture: "x86",
			FSLayers:     blobs,
			History:      history,
			Versioned: storage.Versioned{
				SchemaVersion: 1,
			},
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	blobRequestResponseMappings := make([]testutil.RequestResponseMapping, 2*len(testBlobs))
	for i, blob := range testBlobs {
		blobRequestResponseMappings[2*i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "POST",
				Route:  "/v2/" + name + "/blob/upload/",
			},
			Response: testutil.Response{
				StatusCode: http.StatusAccepted,
				Headers: http.Header(map[string][]string{
					"Location": {uploadLocations[i]},
				}),
			},
		}
		blobRequestResponseMappings[2*i+1] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "PUT",
				Route:  uploadLocations[i],
				QueryParams: map[string][]string{
					"length": {fmt.Sprint(len(blob.contents))},
					"digest": {blob.digest.String()},
				},
				Body: blob.contents,
			},
			Response: testutil.Response{
				StatusCode: http.StatusCreated,
			},
		}
	}

	handler := testutil.NewHandler(append(blobRequestResponseMappings, testutil.RequestResponseMap{
		{
			Request: testutil.Request{
				Method: "PUT",
				Route:  "/v2/" + name + "/manifest/" + tag,
				Body:   manifestBytes,
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
			},
		},
	}...))
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		writer, err := l.Writer()
		if err != nil {
			t.Fatal(err)
		}

		writer.SetSize(len(blob.contents))
		writer.Write(blob.contents)
		writer.Close()
	}

	objectStore.WriteManifest(name, tag, manifest)

	err = Push(client, objectStore, name, tag)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPull(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	blobs := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, blob := range testBlobs {
		blobs[i] = storage.FSLayer{BlobSum: blob.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: blob.digest.String()}
	}

	manifest := &storage.SignedManifest{
		Manifest: storage.Manifest{
			Name:         name,
			Tag:          tag,
			Architecture: "x86",
			FSLayers:     blobs,
			History:      history,
			Versioned: storage.Versioned{
				SchemaVersion: 1,
			},
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	blobRequestResponseMappings := make([]testutil.RequestResponseMapping, len(testBlobs))
	for i, blob := range testBlobs {
		blobRequestResponseMappings[i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents,
			},
		}
	}

	handler := testutil.NewHandler(append(blobRequestResponseMappings, testutil.RequestResponseMap{
		{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/manifest/" + tag,
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       manifestBytes,
			},
		},
	}...))
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	err = Pull(client, objectStore, name, tag)
	if err != nil {
		t.Fatal(err)
	}

	m, err := objectStore.Manifest(name, tag)
	if err != nil {
		t.Fatal(err)
	}

	mBytes, err := json.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}

	if string(mBytes) != string(manifestBytes) {
		t.Fatal("Incorrect manifest")
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		reader, err := l.Reader()
		if err != nil {
			t.Fatal(err)
		}
		defer reader.Close()

		blobBytes, err := ioutil.ReadAll(reader)
		if err != nil {
			t.Fatal(err)
		}

		if string(blobBytes) != string(blob.contents) {
			t.Fatal("Incorrect blob")
		}
	}
}

func TestPullResume(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	layers := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, layer := range testBlobs {
		layers[i] = storage.FSLayer{BlobSum: layer.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: layer.digest.String()}
	}

	manifest := &storage.Manifest{
		Name:         name,
		Tag:          tag,
		Architecture: "x86",
		FSLayers:     layers,
		History:      history,
		Versioned: storage.Versioned{
			SchemaVersion: 1,
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	layerRequestResponseMappings := make([]testutil.RequestResponseMapping, 2*len(testBlobs))
	for i, blob := range testBlobs {
		layerRequestResponseMappings[2*i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents[:len(blob.contents)/2],
				Headers: http.Header(map[string][]string{
					"Content-Length": {fmt.Sprint(len(blob.contents))},
				}),
			},
		}
		layerRequestResponseMappings[2*i+1] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents[len(blob.contents)/2:],
			},
		}
	}

	for i := 0; i < 3; i++ {
		layerRequestResponseMappings = append(layerRequestResponseMappings, testutil.RequestResponseMap{
			{
				Request: testutil.Request{
					Method: "GET",
					Route:  "/v2/" + name + "/manifest/" + tag,
				},
				Response: testutil.Response{
					StatusCode: http.StatusOK,
					Body:       manifestBytes,
				},
			},
		}...)
	}

	handler := testutil.NewHandler(layerRequestResponseMappings)
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	for attempts := 0; attempts < 3; attempts++ {
		err = Pull(client, objectStore, name, tag)
		if err == nil {
			break
		}
	}

	if err != nil {
		t.Fatal(err)
	}

	m, err := objectStore.Manifest(name, tag)
	if err != nil {
		t.Fatal(err)
	}

	mBytes, err := json.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}

	if string(mBytes) != string(manifestBytes) {
		t.Fatal("Incorrect manifest")
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		reader, err := l.Reader()
		if err != nil {
			t.Fatal(err)
		}
		defer reader.Close()

		layerBytes, err := ioutil.ReadAll(reader)
		if err != nil {
			t.Fatal(err)
		}

		if string(layerBytes) != string(blob.contents) {
			t.Fatal("Incorrect blob")
		}
	}
}
<|endoftext|>