<commit_before>package pebbleclient

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

const maxPartialBody = 64 * 1024

type options Options

func (o options) Merge(other *options) options {
	op := Options(*other)
	return options(Options(o).Merge(&op))
}

// HTTPClient is a client for the Central API.
type HTTPClient struct {
	options
	hc *http.Client
}

// New constructs a new client.
func NewHTTPClient(opts Options) (*HTTPClient, error) {
	if err := opts.validate(); err != nil {
		return nil, err
	}
	var newOpts options = *(*options)(opts.applyDefaults())
	return &HTTPClient{
		options: newOpts,
		hc:      newOpts.HTTPClient,
	}, nil
}

// NewFromHTTPRequest constructs a new client that inherits the host name, protocol,
// session and request ID from an HTTP request. Any options specified will override inferred
// from the request.
func (client *HTTPClient) FromHTTPRequest(req *http.Request) (*HTTPClient, error) {
	opts := Options(client.options)

	var host string
	if hosts, ok := req.Header["X-Forwarded-Host"]; ok && len(hosts) > 0 {
		host = hosts[len(hosts)-1]
	} else {
		host = strings.SplitN(req.Host, ":", 2)[0]
	}
	opts.Host = host

	opts.Protocol = req.URL.Scheme

	if session := req.URL.Query().Get("session"); session != "" {
		opts.Session = session
	} else if cookie, err := req.Cookie("checkpoint.session"); err == nil {
		opts.Session = cookie.Value
	}

	if id := req.Header.Get("Request-Id"); id != "" {
		opts.RequestId = id
	}

	opts.HTTPClient = client.hc

	return NewHTTPClient(opts)
}

// GetOptions returns a copy of this client's options.
func (client *HTTPClient) GetOptions() Options {
	return Options(client.options)
}

func (client *HTTPClient) Options(opts Options) Client {
	newOpts := client.options.Merge((*options)(&opts))
	return &HTTPClient{
		options: newOpts,
		hc:      client.hc,
	}
}

func (client *HTTPClient) Get(path string, opts *RequestOptions, result interface{}) error {
	return client.do(opts, http.MethodGet, path, nil, result)
}

func (client *HTTPClient) Head(path string, opts *RequestOptions) error {
	return client.do(opts, http.MethodHead, path, nil, nil)
}

func (client *HTTPClient) Post(path string, opts *RequestOptions, body io.Reader, result interface{}) error {
	return client.do(opts, http.MethodPost, path, body, result)
}

func (client *HTTPClient) Put(path string, opts *RequestOptions, body io.Reader, result interface{}) error {
	return client.do(opts, http.MethodPut, path, body, result)
}

func (client *HTTPClient) Delete(path string, opts *RequestOptions, result interface{}) error {
	return client.do(opts, http.MethodDelete, path, nil, result)
}

func (client *HTTPClient) do(
	opts *RequestOptions,
	method string,
	path string,
	bodyIn io.Reader,
	result interface{}) error {
	if opts == nil {
		opts = &RequestOptions{}
	}

	req, err := http.NewRequest(method, client.formatEndpointUrl(path, opts.Params), bodyIn)
	if err != nil {
		return err
	}

	req.Header.Set("Content-Type", "application/json; charset=utf8")
	if client.options.RequestId != "" {
		req.Header.Set("Request-Id", client.options.RequestId)
	}

	ctx := client.Ctx
	if ctx == nil {
		ctx = context.Background()
	}

	resp, err := ctxhttp.Do(ctx, client.hc, req)
	if err != nil {
		return err
	}

	bodyOut := resp.Body
	defer func() {
		if bodyOut != nil {
			_ = bodyOut.Close()
		}
	}()

	if isNonSuccessStatus(resp.StatusCode) {
		return client.buildError(&RequestError{}, opts, resp)
	}

	if doesStatusCodeYieldBody(resp.StatusCode) && result != nil {
		return decodeResponseAsJSON(resp, bodyOut, result)
	}

	return nil
}

func (client *HTTPClient) buildError(
	err *RequestError,
	opts *RequestOptions,
	resp *http.Response) error {
	var buf bytes.Buffer
	b := make([]byte, 1024)
	for buf.Len() < maxPartialBody {
		count, err := resp.Body.Read(b[:])
		if count == 0 {
			break
		}
		if err != nil && err != io.EOF {
			break
		}
		_, wErr := buf.Write(b[0:count])
		if err != nil || wErr != nil {
			break
		}
	}
	err.PartialBody = buf.Bytes()
	err.client = client
	err.Resp = resp
	err.Options = opts
	return err
}

func (client *HTTPClient) formatEndpointUrl(path string, params Params) string {
	if path[0:1] == "/" {
		path = path[1:]
	}
	result := url.URL{
		Scheme: client.Protocol,
		Host:   client.Host,
		Path: fmt.Sprintf("/api/%s/v%d/%s",
			client.ServiceName, client.ApiVersion, escapedPath(path)),
	}

	query := result.Query()
	if params != nil {
		for key, value := range params {
			query.Set(key, fmt.Sprintf("%s", value))
		}
	}
	if client.Session != "" {
		query.Set("session", client.Session)
	}
	result.RawQuery = query.Encode()

	return result.String()
}
<commit_msg>Drain response body to work around Go bug.<commit_after>package pebbleclient

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

const maxPartialBody = 64 * 1024

type options Options

func (o options) Merge(other *options) options {
	op := Options(*other)
	return options(Options(o).Merge(&op))
}

// HTTPClient is a client for the Central API.
type HTTPClient struct {
	options
	hc *http.Client
}

// New constructs a new client.
func NewHTTPClient(opts Options) (*HTTPClient, error) {
	if err := opts.validate(); err != nil {
		return nil, err
	}
	var newOpts options = *(*options)(opts.applyDefaults())
	return &HTTPClient{
		options: newOpts,
		hc:      newOpts.HTTPClient,
	}, nil
}

// NewFromHTTPRequest constructs a new client that inherits the host name, protocol,
// session and request ID from an HTTP request. Any options specified will override inferred
// from the request.
func (client *HTTPClient) FromHTTPRequest(req *http.Request) (*HTTPClient, error) {
	opts := Options(client.options)

	var host string
	if hosts, ok := req.Header["X-Forwarded-Host"]; ok && len(hosts) > 0 {
		host = hosts[len(hosts)-1]
	} else {
		host = strings.SplitN(req.Host, ":", 2)[0]
	}
	opts.Host = host

	opts.Protocol = req.URL.Scheme

	if session := req.URL.Query().Get("session"); session != "" {
		opts.Session = session
	} else if cookie, err := req.Cookie("checkpoint.session"); err == nil {
		opts.Session = cookie.Value
	}

	if id := req.Header.Get("Request-Id"); id != "" {
		opts.RequestId = id
	}

	opts.HTTPClient = client.hc

	return NewHTTPClient(opts)
}

// GetOptions returns a copy of this client's options.
func (client *HTTPClient) GetOptions() Options {
	return Options(client.options)
}

func (client *HTTPClient) Options(opts Options) Client {
	newOpts := client.options.Merge((*options)(&opts))
	return &HTTPClient{
		options: newOpts,
		hc:      client.hc,
	}
}

func (client *HTTPClient) Get(path string, opts *RequestOptions, result interface{}) error {
	return client.do(opts, http.MethodGet, path, nil, result)
}

func (client *HTTPClient) Head(path string, opts *RequestOptions) error {
	return client.do(opts, http.MethodHead, path, nil, nil)
}

func (client *HTTPClient) Post(path string, opts *RequestOptions, body io.Reader, result interface{}) error {
	return client.do(opts, http.MethodPost, path, body, result)
}

func (client *HTTPClient) Put(path string, opts *RequestOptions, body io.Reader, result interface{}) error {
	return client.do(opts, http.MethodPut, path, body, result)
}

func (client *HTTPClient) Delete(path string, opts *RequestOptions, result interface{}) error {
	return client.do(opts, http.MethodDelete, path, nil, result)
}

func (client *HTTPClient) do(
	opts *RequestOptions,
	method string,
	path string,
	bodyIn io.Reader,
	result interface{}) error {
	if opts == nil {
		opts = &RequestOptions{}
	}

	req, err := http.NewRequest(method, client.formatEndpointUrl(path, opts.Params), bodyIn)
	if err != nil {
		return err
	}

	req.Header.Set("Content-Type", "application/json; charset=utf8")
	if client.options.RequestId != "" {
		req.Header.Set("Request-Id", client.options.RequestId)
	}

	ctx := client.Ctx
	if ctx == nil {
		ctx = context.Background()
	}

	resp, err := ctxhttp.Do(ctx, client.hc, req)
	if err != nil {
		return err
	}

	respBody := resp.Body
	defer func() {
		if respBody != nil {
			// Drain remaining body to work around bug in Go < 1.7
			_, _ = io.Copy(ioutil.Discard, respBody)

			_ = respBody.Close()
		}
	}()

	if isNonSuccessStatus(resp.StatusCode) {
		return client.buildError(&RequestError{}, opts, resp)
	}

	if doesStatusCodeYieldBody(resp.StatusCode) && result != nil {
		return decodeResponseAsJSON(resp, respBody, result)
	}

	return nil
}

func (client *HTTPClient) buildError(
	err *RequestError,
	opts *RequestOptions,
	resp *http.Response) error {
	var buf bytes.Buffer
	b := make([]byte, 1024)
	for buf.Len() < maxPartialBody {
		count, err := resp.Body.Read(b[:])
		if count == 0 {
			break
		}
		if err != nil && err != io.EOF {
			break
		}
		_, wErr := buf.Write(b[0:count])
		if err != nil || wErr != nil {
			break
		}
	}
	err.PartialBody = buf.Bytes()
	err.client = client
	err.Resp = resp
	err.Options = opts
	return err
}

func (client *HTTPClient) formatEndpointUrl(path string, params Params) string {
	if path[0:1] == "/" {
		path = path[1:]
	}
	result := url.URL{
		Scheme: client.Protocol,
		Host:   client.Host,
		Path: fmt.Sprintf("/api/%s/v%d/%s",
			client.ServiceName, client.ApiVersion, escapedPath(path)),
	}

	query := result.Query()
	if params != nil {
		for key, value := range params {
			query.Set(key, fmt.Sprintf("%s", value))
		}
	}
	if client.Session != "" {
		query.Set("session", client.Session)
	}
	result.RawQuery = query.Encode()

	return result.String()
}
<|endoftext|>


<commit_before>package tcp

import (
	"context"
	"crypto/tls"
	"net"
	"sync"
	"time"
)

// Handler responds to a TCP request.
//
// ServeTCP should write responses to the in connection and close
// it on return.
type Handler interface {
	ServeTCP(in net.Conn) error
}

type HandlerFunc func(in net.Conn) error

func (f HandlerFunc) ServeTCP(in net.Conn) error {
	return f(in)
}

// Server implements a generic TCP server.
type Server struct {
	Addr         string
	Handler      Handler
	ReadTimeout  time.Duration
	WriteTimeout time.Duration

	mu        sync.Mutex
	listeners []net.Listener
	conns     map[net.Conn]bool
}

func (s *Server) ListenAndServe() error {
	l, err := net.Listen("tcp", s.Addr)
	if err != nil {
		return err
	}
	defer l.Close()
	return s.Serve(l)
}

func (s *Server) ListenAndServeTLS(certFile, keyFile string) error {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}
	cfg := &tls.Config{Certificates: []tls.Certificate{cert}}
	l, err := tls.Listen("tcp", s.Addr, cfg)
	if err != nil {
		return err
	}
	defer l.Close()
	return s.Serve(l)
}

func (s *Server) Serve(l net.Listener) error {
	defer l.Close()

	s.mu.Lock()
	s.listeners = append(s.listeners, l)
	s.mu.Unlock()

	for {
		c, err := l.Accept()
		if err != nil {
			return err
		}
		c = &conn{
			c:            c,
			ReadTimeout:  s.ReadTimeout,
			WriteTimeout: s.WriteTimeout,
		}
		s.mu.Lock()
		if s.conns == nil {
			s.conns = map[net.Conn]bool{}
		}
		s.conns[c] = true
		s.mu.Unlock()
		go s.Handler.ServeTCP(c)
	}
}

func (s *Server) closeListeners() error {
	s.mu.Lock()
	for _, l := range s.listeners {
		l.Close()
	}
	s.listeners = nil
	s.mu.Unlock()
	return nil
}

func (s *Server) closeConns() error {
	s.mu.Lock()
	for c := range s.conns {
		c.Close()
	}
	s.conns = nil
	s.mu.Unlock()
	return nil
}

func (s *Server) Close() error {
	s.closeListeners()
	return s.closeConns()
}

func (s *Server) Shutdown(ctx context.Context) error {
	s.closeListeners()
	if ctx != nil {
		<-ctx.Done()
	}
	return s.closeConns()
}

// conn implements a connection which honors read and write timeouts.
type conn struct {
	c            net.Conn
	ReadTimeout  time.Duration
	WriteTimeout time.Duration
}

func (c *conn) Read(b []byte) (int, error) {
	if c.ReadTimeout > 0 {
		c.c.SetReadDeadline(time.Now().Add(c.ReadTimeout))
	}
	return c.c.Read(b)
}

func (c *conn) Write(b []byte) (int, error) {
	if c.WriteTimeout > 0 {
		c.c.SetWriteDeadline(time.Now().Add(c.WriteTimeout))
	}
	return c.c.Write(b)
}

func (c *conn) Close() error {
	return c.c.Close()
}

func (c *conn) LocalAddr() net.Addr {
	return c.c.LocalAddr()
}

func (c *conn) RemoteAddr() net.Addr {
	return c.c.RemoteAddr()
}

func (c *conn) SetDeadline(t time.Time) error {
	return c.c.SetDeadline(t)
}

func (c *conn) SetReadDeadline(t time.Time) error {
	return c.c.SetReadDeadline(t)
}

func (c *conn) SetWriteDeadline(t time.Time) error {
	return c.c.SetWriteDeadline(t)
}
<commit_msg>PR #321: Cleanup TCP proxy connection<commit_after>package tcp

import (
	"context"
	"crypto/tls"
	"net"
	"sync"
	"time"
)

// Handler responds to a TCP request.
//
// ServeTCP should write responses to the in connection and close
// it on return.
type Handler interface {
	ServeTCP(in net.Conn) error
}

type HandlerFunc func(in net.Conn) error

func (f HandlerFunc) ServeTCP(in net.Conn) error {
	return f(in)
}

// Server implements a generic TCP server.
type Server struct {
	Addr         string
	Handler      Handler
	ReadTimeout  time.Duration
	WriteTimeout time.Duration

	mu        sync.Mutex
	listeners []net.Listener
	conns     map[net.Conn]bool
}

func (s *Server) ListenAndServe() error {
	l, err := net.Listen("tcp", s.Addr)
	if err != nil {
		return err
	}
	defer l.Close()
	return s.Serve(l)
}

func (s *Server) ListenAndServeTLS(certFile, keyFile string) error {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}
	cfg := &tls.Config{Certificates: []tls.Certificate{cert}}
	l, err := tls.Listen("tcp", s.Addr, cfg)
	if err != nil {
		return err
	}
	defer l.Close()
	return s.Serve(l)
}

func (s *Server) Serve(l net.Listener) error {
	defer l.Close()

	s.mu.Lock()
	s.listeners = append(s.listeners, l)
	s.mu.Unlock()

	for {
		c, err := l.Accept()
		if err != nil {
			return err
		}
		c = &conn{
			c:            c,
			ReadTimeout:  s.ReadTimeout,
			WriteTimeout: s.WriteTimeout,
		}
		s.mu.Lock()
		if s.conns == nil {
			s.conns = map[net.Conn]bool{}
		}
		s.conns[c] = true
		s.mu.Unlock()

		go func() {
			defer func() {
				// Explicitly close the connection
				c.Close()

				// Delete the connection from the conns map
				s.mu.Lock()
				delete(s.conns, c)
				s.mu.Unlock()
			}()

			s.Handler.ServeTCP(c)
		}()
	}
}

func (s *Server) closeListeners() error {
	s.mu.Lock()
	for _, l := range s.listeners {
		l.Close()
	}
	s.listeners = nil
	s.mu.Unlock()
	return nil
}

func (s *Server) closeConns() error {
	s.mu.Lock()
	for c := range s.conns {
		c.Close()
	}
	s.conns = nil
	s.mu.Unlock()
	return nil
}

func (s *Server) Close() error {
	s.closeListeners()
	return s.closeConns()
}

func (s *Server) Shutdown(ctx context.Context) error {
	s.closeListeners()
	if ctx != nil {
		<-ctx.Done()
	}
	return s.closeConns()
}

// conn implements a connection which honors read and write timeouts.
type conn struct {
	c            net.Conn
	ReadTimeout  time.Duration
	WriteTimeout time.Duration
}

func (c *conn) Read(b []byte) (int, error) {
	if c.ReadTimeout > 0 {
		c.c.SetReadDeadline(time.Now().Add(c.ReadTimeout))
	}
	return c.c.Read(b)
}

func (c *conn) Write(b []byte) (int, error) {
	if c.WriteTimeout > 0 {
		c.c.SetWriteDeadline(time.Now().Add(c.WriteTimeout))
	}
	return c.c.Write(b)
}

func (c *conn) Close() error {
	return c.c.Close()
}

func (c *conn) LocalAddr() net.Addr {
	return c.c.LocalAddr()
}

func (c *conn) RemoteAddr() net.Addr {
	return c.c.RemoteAddr()
}

func (c *conn) SetDeadline(t time.Time) error {
	return c.c.SetDeadline(t)
}

func (c *conn) SetReadDeadline(t time.Time) error {
	return c.c.SetReadDeadline(t)
}

func (c *conn) SetWriteDeadline(t time.Time) error {
	return c.c.SetWriteDeadline(t)
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build android

package libkb

import "sync"

// TODO: Make this implementation use GetStoredSecretServiceName(), as
// otherwise tests will clobber each other's passwords. See
// https://keybase.atlassian.net/browse/CORE-1934 .

// UnsafeExternalKeyStore is a simple interface that external clients can implement.
// It is unsafe because it returns raw bytes instead of the typed LKSecFullSecret
// Use with TypeSafeExternalKeyStoreProxy
type UnsafeExternalKeyStore interface {
	RetrieveSecret(serviceName string, key string) ([]byte, error)
	StoreSecret(serviceName string, key string, secret []byte) error
	ClearSecret(serviceName string, key string) error
	GetUsersWithStoredSecretsMsgPack(serviceName string) ([]byte, error)
	SetupKeyStore(serviceName string, key string) error
}

// ExternalKeyStore is the interface for the actual (external) keystore.
type ExternalKeyStore interface {
	RetrieveSecret(serviceName string, key string) (LKSecFullSecret, error)
	StoreSecret(serviceName string, key string, secret LKSecFullSecret) error
	ClearSecret(serviceName string, key string) error
	GetUsersWithStoredSecretsMsgPack(serviceName string) ([]byte, error)
	SetupKeyStore(serviceName string, key string) error
}

// TypeSafeExternalKeyStoreProxy wraps the UnsafeExternalKeyStore to provide
// the type-safe ExternalKeyStore interface to the rest of the code
type TypeSafeExternalKeyStoreProxy struct {
	UnsafeExternalKeyStore UnsafeExternalKeyStore
}

func (w TypeSafeExternalKeyStoreProxy) RetrieveSecret(serviceName string, key string) (LKSecFullSecret, error) {
	bytes, err := w.UnsafeExternalKeyStore.RetrieveSecret(serviceName, key)
	if err != nil {
		return LKSecFullSecret{}, err
	}

	return newLKSecFullSecretFromBytes(bytes)
}

func (w TypeSafeExternalKeyStoreProxy) StoreSecret(serviceName string, key string, secret LKSecFullSecret) error {
	return w.UnsafeExternalKeyStore.StoreSecret(serviceName, key, secret.Bytes())
}

func (w TypeSafeExternalKeyStoreProxy) ClearSecret(serviceName string, key string) error {
	return w.UnsafeExternalKeyStore.ClearSecret(serviceName, key)
}

func (w TypeSafeExternalKeyStoreProxy) GetUsersWithStoredSecretsMsgPack(serviceName string) ([]byte, error) {
	return w.UnsafeExternalKeyStore.GetUsersWithStoredSecretsMsgPack(serviceName)
}

func (w TypeSafeExternalKeyStoreProxy) SetupKeyStore(serviceName string, key string) error {
	return w.UnsafeExternalKeyStore.SetupKeyStore(serviceName, key)
}

// externalKeyStore is the reference to some external key store
var externalKeyStore ExternalKeyStore

var externalKeyStoreMu sync.Mutex

func (s *secretStoreAccountName) serviceName() string {
	return s.context.GetStoredSecretServiceName()
}

// SetGlobalExternalKeyStore is called by Android to register Android's KeyStore with Go
func SetGlobalExternalKeyStore(s UnsafeExternalKeyStore) {
	externalKeyStoreMu.Lock()
	defer externalKeyStoreMu.Unlock()
	externalKeyStore = TypeSafeExternalKeyStoreProxy{s}
}

func getGlobalExternalKeyStore() ExternalKeyStore {
	externalKeyStoreMu.Lock()
	defer externalKeyStoreMu.Unlock()
	return externalKeyStore
}

type secretStoreAccountName struct {
	externalKeyStore ExternalKeyStore
	context          SecretStoreContext
	setupOnce        sync.Once
}

var _ SecretStoreAll = &secretStoreAccountName{}

func NewSecretStoreAll(g *GlobalContext) SecretStoreAll {
	externalKeyStore := getGlobalExternalKeyStore()
	if externalKeyStore == nil {
		return nil
	}
	s := &secretStoreAccountName{
		externalKeyStore: externalKeyStore,
		context:          g,
	}
	go s.setup()
	return s
}

func (s *secretStoreAccountName) StoreSecret(username NormalizedUsername, secret LKSecFullSecret) (err error) {
	defer TraceTimed(s.context.GetLog(), "secret_store_external StoreSecret", func() error { return err })()
	s.setup()
	return s.externalKeyStore.StoreSecret(s.serviceName(), string(username), secret)
}

func (s *secretStoreAccountName) RetrieveSecret(username NormalizedUsername) (sec LKSecFullSecret, err error) {
	defer TraceTimed(s.context.GetLog(), "secret_store_external RetrieveSecret", func() error { return err })()
	s.setup()
	return s.externalKeyStore.RetrieveSecret(s.serviceName(), string(username))
}

func (s *secretStoreAccountName) ClearSecret(username NormalizedUsername) (err error) {
	defer TraceTimed(s.context.GetLog(), "secret_store_external ClearSecret", func() error { return err })()
	return s.externalKeyStore.ClearSecret(s.serviceName(), string(username))
}

func (s *secretStoreAccountName) GetUsersWithStoredSecrets() ([]string, error) {
	if s.externalKeyStore == nil {
		return nil, nil
	}
	usersMsgPack, err := s.externalKeyStore.GetUsersWithStoredSecretsMsgPack(s.serviceName())
	if err != nil {
		return nil, err
	}
	var users []string
	ch := codecHandle()
	err = MsgpackDecodeAll(usersMsgPack, ch, &users)
	return users, err
}

func (s *secretStoreAccountName) GetTerminalPrompt() string {
	return "Store secret in Android's KeyStore?"
}

func (s *secretStoreAccountName) GetApprovalPrompt() string {
	return "Store secret in Android's KeyStore?"
}

func (s *secretStoreAccountName) setup() {
	s.setupOnce.Do(func() {
		s.context.GetLog().Debug("+ secret_store_external:setup")
		// username not required
		s.externalKeyStore.SetupKeyStore(s.serviceName(), "")
		s.context.GetLog().Debug("- secret_store_external:setup")
	})
}

func NewTestSecretStoreAll(c SecretStoreContext, g *GlobalContext) SecretStoreAll {
	return nil
}
<commit_msg>Log the error from SetupKeyStore<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build android

package libkb

import "sync"

// TODO: Make this implementation use GetStoredSecretServiceName(), as
// otherwise tests will clobber each other's passwords. See
// https://keybase.atlassian.net/browse/CORE-1934 .

// UnsafeExternalKeyStore is a simple interface that external clients can implement.
// It is unsafe because it returns raw bytes instead of the typed LKSecFullSecret
// Use with TypeSafeExternalKeyStoreProxy
type UnsafeExternalKeyStore interface {
	RetrieveSecret(serviceName string, key string) ([]byte, error)
	StoreSecret(serviceName string, key string, secret []byte) error
	ClearSecret(serviceName string, key string) error
	GetUsersWithStoredSecretsMsgPack(serviceName string) ([]byte, error)
	SetupKeyStore(serviceName string, key string) error
}

// ExternalKeyStore is the interface for the actual (external) keystore.
type ExternalKeyStore interface {
	RetrieveSecret(serviceName string, key string) (LKSecFullSecret, error)
	StoreSecret(serviceName string, key string, secret LKSecFullSecret) error
	ClearSecret(serviceName string, key string) error
	GetUsersWithStoredSecretsMsgPack(serviceName string) ([]byte, error)
	SetupKeyStore(serviceName string, key string) error
}

// TypeSafeExternalKeyStoreProxy wraps the UnsafeExternalKeyStore to provide
// the type-safe ExternalKeyStore interface to the rest of the code
type TypeSafeExternalKeyStoreProxy struct {
	UnsafeExternalKeyStore UnsafeExternalKeyStore
}

func (w TypeSafeExternalKeyStoreProxy) RetrieveSecret(serviceName string, key string) (LKSecFullSecret, error) {
	bytes, err := w.UnsafeExternalKeyStore.RetrieveSecret(serviceName, key)
	if err != nil {
		return LKSecFullSecret{}, err
	}

	return newLKSecFullSecretFromBytes(bytes)
}

func (w TypeSafeExternalKeyStoreProxy) StoreSecret(serviceName string, key string, secret LKSecFullSecret) error {
	return w.UnsafeExternalKeyStore.StoreSecret(serviceName, key, secret.Bytes())
}

func (w TypeSafeExternalKeyStoreProxy) ClearSecret(serviceName string, key string) error {
	return w.UnsafeExternalKeyStore.ClearSecret(serviceName, key)
}

func (w TypeSafeExternalKeyStoreProxy) GetUsersWithStoredSecretsMsgPack(serviceName string) ([]byte, error) {
	return w.UnsafeExternalKeyStore.GetUsersWithStoredSecretsMsgPack(serviceName)
}

func (w TypeSafeExternalKeyStoreProxy) SetupKeyStore(serviceName string, key string) error {
	return w.UnsafeExternalKeyStore.SetupKeyStore(serviceName, key)
}

// externalKeyStore is the reference to some external key store
var externalKeyStore ExternalKeyStore

var externalKeyStoreMu sync.Mutex

func (s *secretStoreAccountName) serviceName() string {
	return s.context.GetStoredSecretServiceName()
}

// SetGlobalExternalKeyStore is called by Android to register Android's KeyStore with Go
func SetGlobalExternalKeyStore(s UnsafeExternalKeyStore) {
	externalKeyStoreMu.Lock()
	defer externalKeyStoreMu.Unlock()
	externalKeyStore = TypeSafeExternalKeyStoreProxy{s}
}

func getGlobalExternalKeyStore() ExternalKeyStore {
	externalKeyStoreMu.Lock()
	defer externalKeyStoreMu.Unlock()
	return externalKeyStore
}

type secretStoreAccountName struct {
	externalKeyStore ExternalKeyStore
	context          SecretStoreContext
	setupOnce        sync.Once
}

var _ SecretStoreAll = &secretStoreAccountName{}

func NewSecretStoreAll(g *GlobalContext) SecretStoreAll {
	externalKeyStore := getGlobalExternalKeyStore()
	if externalKeyStore == nil {
		return nil
	}
	s := &secretStoreAccountName{
		externalKeyStore: externalKeyStore,
		context:          g,
	}
	go s.setup()
	return s
}

func (s *secretStoreAccountName) StoreSecret(username NormalizedUsername, secret LKSecFullSecret) (err error) {
	defer TraceTimed(s.context.GetLog(), "secret_store_external StoreSecret", func() error { return err })()
	s.setup()
	return s.externalKeyStore.StoreSecret(s.serviceName(), string(username), secret)
}

func (s *secretStoreAccountName) RetrieveSecret(username NormalizedUsername) (sec LKSecFullSecret, err error) {
	defer TraceTimed(s.context.GetLog(), "secret_store_external RetrieveSecret", func() error { return err })()
	s.setup()
	return s.externalKeyStore.RetrieveSecret(s.serviceName(), string(username))
}

func (s *secretStoreAccountName) ClearSecret(username NormalizedUsername) (err error) {
	defer TraceTimed(s.context.GetLog(), "secret_store_external ClearSecret", func() error { return err })()
	return s.externalKeyStore.ClearSecret(s.serviceName(), string(username))
}

func (s *secretStoreAccountName) GetUsersWithStoredSecrets() ([]string, error) {
	if s.externalKeyStore == nil {
		return nil, nil
	}
	usersMsgPack, err := s.externalKeyStore.GetUsersWithStoredSecretsMsgPack(s.serviceName())
	if err != nil {
		return nil, err
	}
	var users []string
	ch := codecHandle()
	err = MsgpackDecodeAll(usersMsgPack, ch, &users)
	return users, err
}

func (s *secretStoreAccountName) GetTerminalPrompt() string {
	return "Store secret in Android's KeyStore?"
}

func (s *secretStoreAccountName) GetApprovalPrompt() string {
	return "Store secret in Android's KeyStore?"
}

func (s *secretStoreAccountName) setup() {
	s.setupOnce.Do(func() {
		s.context.GetLog().Debug("+ secret_store_external:setup")

		// username not required
		err := s.externalKeyStore.SetupKeyStore(s.serviceName(), "")
		if err != nil {
			s.context.GetLog().Debug("externalKeyStore.SetupKeyStore(%s) error: %s (%T)", s.serviceName(), err, err)
		} else {
			s.context.GetLog().Debug("externalKeyStore.SetupKeyStore(%s) success", s.serviceName())
		}

		s.context.GetLog().Debug("- secret_store_external:setup")
	})
}

func NewTestSecretStoreAll(c SecretStoreContext, g *GlobalContext) SecretStoreAll {
	return nil
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/albertyw/reaction-pics/tumblr"
	"github.com/ikeikeikeike/go-sitemap-generator/stm"
	// Used for getting tumblr env vars
	_ "github.com/joho/godotenv/autoload"
	newrelic "github.com/newrelic/go-agent"
)

const (
	maxResults = 20
)

var serverDir = filepath.Join(os.Getenv("ROOT_DIR"), "server")
var staticPath = fmt.Sprintf("%s/static/", serverDir)
var uRLFilePaths = map[string]func() (string, error){}
var board = tumblr.NewBoard([]tumblr.Post{})

// logURL is a closure that logs (to stdout) the url and query of requests
func logURL(
	targetFunc func(http.ResponseWriter, *http.Request),
) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		url := r.URL.String()
		fmt.Println(url)
		targetFunc(w, r)
	}
}

// rewriteFS wraps a static file handler so to rewrite to the static directory
// and the root path is rewritten to index.htm
func rewriteFS(targetFunc func(http.ResponseWriter, *http.Request),
) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		path := strings.TrimPrefix(r.URL.Path, "/static/")
		if path == "/" {
			path = "index.htm"
		}
		r.URL.Path = path
		targetFunc(w, r)
	}
}

// searchHandler is an http handler to search data for keywords
// It matches the query against post titles and then ranks posts by number of likes
func searchHandler(w http.ResponseWriter, r *http.Request) {
	query := r.URL.Query().Get("query")
	query = strings.ToLower(query)
	queriedBoard := board.FilterBoard(query)
	offsetString := r.URL.Query().Get("offset")
	offset, err := strconv.Atoi(offsetString)
	if err != nil {
		offset = 0
	}
	data := map[string]interface{}{
		"offset":       offset,
		"totalResults": len(queriedBoard.Posts),
	}
	queriedBoard.LimitBoard(offset, maxResults)
	data["data"] = queriedBoard.PostsToJSON()
	dataBytes, _ := json.Marshal(data)
	fmt.Fprintf(w, string(dataBytes))
}

// postDataHandler is an http handler to return post data by ID
func postDataHandler(w http.ResponseWriter, r *http.Request) {
	pathStrings := strings.Split(r.URL.Path, "/")
	postIDString := pathStrings[2]
	postID, err := strconv.ParseInt(postIDString, 10, 64)
	if err != nil {
		fmt.Println("Cannot parse post id")
		http.NotFound(w, r)
		return
	}
	post := board.GetPostByID(postID)
	if post == nil {
		fmt.Println("Cannot find post")
		http.NotFound(w, r)
		return
	}
	data := map[string]interface{}{
		"offset":       0,
		"totalResults": 1,
		"data":         []tumblr.PostJSON{post.ToJSONStruct()},
	}
	marshalledPost, _ := json.Marshal(data)
	fmt.Fprintf(w, string(marshalledPost))
}

func postHandler(w http.ResponseWriter, r *http.Request) {
	pathStrings := strings.Split(r.URL.Path, "/")
	postIDString := pathStrings[2]
	postID, err := strconv.ParseInt(postIDString, 10, 64)
	if err != nil {
		fmt.Println("Cannot parse post id")
		http.NotFound(w, r)
		return
	}
	var post tumblr.Post
	for _, p := range board.Posts {
		if p.ID == postID {
			post = p
		}
	}
	if post == (tumblr.Post{}) {
		fmt.Println("Cannot find post")
		http.NotFound(w, r)
		return
	}
	templateData, err := ioutil.ReadFile(staticPath + "post.htm")
	if err != nil {
		http.Error(w, "Cannot read post template", 500)
		return
	}
	htmlTemplate, err := template.New("page").Parse(string(templateData))
	if err != nil {
		http.Error(w, "Cannot parse post template", 500)
		return
	}
	data := struct {
		Post *tumblr.Post
	}{
		Post: &post,
	}
	err = htmlTemplate.Execute(w, data)
	if err != nil {
		http.Error(w, "Cannot execute post template", 500)
		return
	}
}

func statsHandler(w http.ResponseWriter, r *http.Request) {
	postCount := strconv.Itoa(len(board.Posts))
	data := map[string]interface{}{
		"postCount": postCount,
		"keywords":  board.Keywords(),
	}
	stats, _ := json.Marshal(data)
	fmt.Fprintf(w, string(stats))
}

func sitemapHandler(w http.ResponseWriter, r *http.Request) {
	sm := stm.NewSitemap()
	sm.SetDefaultHost(os.Getenv("HOST"))

	sm.Create()
	sm.Add(stm.URL{"loc": "/"})
	for _, url := range board.URLs() {
		sm.Add(stm.URL{"loc": url})
	}
	w.Write(sm.XMLContent())
}

// Run starts up the HTTP server
func Run(postChan <-chan tumblr.Post, newrelicApp newrelic.Application) {
	go loadPosts(postChan)
	address := ":" + os.Getenv("PORT")
	fmt.Println("server listening on", address)
	staticFS := rewriteFS(http.FileServer(http.Dir(staticPath)).ServeHTTP)
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/", logURL(staticFS)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/search", logURL(searchHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/postdata/", logURL(postDataHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/post/", logURL(postHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/stats.json", logURL(statsHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/sitemap.xml", logURL(sitemapHandler)))
	http.ListenAndServe(address, nil)
}

func loadPosts(postChan <-chan tumblr.Post) {
	for p := range postChan {
		board.AddPost(p)
		board.SortPostsByLikes()
	}
}
<commit_msg>Simplify postHandler<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/albertyw/reaction-pics/tumblr"
	"github.com/ikeikeikeike/go-sitemap-generator/stm"
	// Used for getting tumblr env vars
	_ "github.com/joho/godotenv/autoload"
	newrelic "github.com/newrelic/go-agent"
)

const (
	maxResults = 20
)

var serverDir = filepath.Join(os.Getenv("ROOT_DIR"), "server")
var staticPath = fmt.Sprintf("%s/static/", serverDir)
var uRLFilePaths = map[string]func() (string, error){}
var board = tumblr.NewBoard([]tumblr.Post{})

// logURL is a closure that logs (to stdout) the url and query of requests
func logURL(
	targetFunc func(http.ResponseWriter, *http.Request),
) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		url := r.URL.String()
		fmt.Println(url)
		targetFunc(w, r)
	}
}

// rewriteFS wraps a static file handler so to rewrite to the static directory
// and the root path is rewritten to index.htm
func rewriteFS(targetFunc func(http.ResponseWriter, *http.Request),
) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		path := strings.TrimPrefix(r.URL.Path, "/static/")
		if path == "/" {
			path = "index.htm"
		}
		r.URL.Path = path
		targetFunc(w, r)
	}
}

// searchHandler is an http handler to search data for keywords
// It matches the query against post titles and then ranks posts by number of likes
func searchHandler(w http.ResponseWriter, r *http.Request) {
	query := r.URL.Query().Get("query")
	query = strings.ToLower(query)
	queriedBoard := board.FilterBoard(query)
	offsetString := r.URL.Query().Get("offset")
	offset, err := strconv.Atoi(offsetString)
	if err != nil {
		offset = 0
	}
	data := map[string]interface{}{
		"offset":       offset,
		"totalResults": len(queriedBoard.Posts),
	}
	queriedBoard.LimitBoard(offset, maxResults)
	data["data"] = queriedBoard.PostsToJSON()
	dataBytes, _ := json.Marshal(data)
	fmt.Fprintf(w, string(dataBytes))
}

// postDataHandler is an http handler to return post data by ID
func postDataHandler(w http.ResponseWriter, r *http.Request) {
	pathStrings := strings.Split(r.URL.Path, "/")
	postIDString := pathStrings[2]
	postID, err := strconv.ParseInt(postIDString, 10, 64)
	if err != nil {
		fmt.Println("Cannot parse post id")
		http.NotFound(w, r)
		return
	}
	post := board.GetPostByID(postID)
	if post == nil {
		fmt.Println("Cannot find post")
		http.NotFound(w, r)
		return
	}
	data := map[string]interface{}{
		"offset":       0,
		"totalResults": 1,
		"data":         []tumblr.PostJSON{post.ToJSONStruct()},
	}
	marshalledPost, _ := json.Marshal(data)
	fmt.Fprintf(w, string(marshalledPost))
}

func postHandler(w http.ResponseWriter, r *http.Request) {
	pathStrings := strings.Split(r.URL.Path, "/")
	postIDString := pathStrings[2]
	postID, err := strconv.ParseInt(postIDString, 10, 64)
	if err != nil {
		fmt.Println("Cannot parse post id")
		http.NotFound(w, r)
		return
	}
	foundPost := false
	for _, p := range board.Posts {
		if p.ID == postID {
			foundPost = true
		}
	}
	if !foundPost {
		fmt.Println("Cannot find post")
		http.NotFound(w, r)
		return
	}
	templateData, err := ioutil.ReadFile(staticPath + "post.htm")
	if err != nil {
		http.Error(w, "Cannot read post template", 500)
		return
	}
	htmlTemplate, err := template.New("page").Parse(string(templateData))
	if err != nil {
		http.Error(w, "Cannot parse post template", 500)
		return
	}
	err = htmlTemplate.Execute(w, struct{}{})
	if err != nil {
		http.Error(w, "Cannot execute post template", 500)
		return
	}
}

func statsHandler(w http.ResponseWriter, r *http.Request) {
	postCount := strconv.Itoa(len(board.Posts))
	data := map[string]interface{}{
		"postCount": postCount,
		"keywords":  board.Keywords(),
	}
	stats, _ := json.Marshal(data)
	fmt.Fprintf(w, string(stats))
}

func sitemapHandler(w http.ResponseWriter, r *http.Request) {
	sm := stm.NewSitemap()
	sm.SetDefaultHost(os.Getenv("HOST"))

	sm.Create()
	sm.Add(stm.URL{"loc": "/"})
	for _, url := range board.URLs() {
		sm.Add(stm.URL{"loc": url})
	}
	w.Write(sm.XMLContent())
}

// Run starts up the HTTP server
func Run(postChan <-chan tumblr.Post, newrelicApp newrelic.Application) {
	go loadPosts(postChan)
	address := ":" + os.Getenv("PORT")
	fmt.Println("server listening on", address)
	staticFS := rewriteFS(http.FileServer(http.Dir(staticPath)).ServeHTTP)
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/", logURL(staticFS)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/search", logURL(searchHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/postdata/", logURL(postDataHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/post/", logURL(postHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/stats.json", logURL(statsHandler)))
	http.HandleFunc(newrelic.WrapHandleFunc(newrelicApp, "/sitemap.xml", logURL(sitemapHandler)))
	http.ListenAndServe(address, nil)
}

func loadPosts(postChan <-chan tumblr.Post) {
	for p := range postChan {
		board.AddPost(p)
		board.SortPostsByLikes()
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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.

// A command that reads all blobs necessary for verifying the directory
// structure rooted at a set of backup root scores, ensuring that the entire
// directory structure is intact in GCS.
//
// Optionally, all file content is also read and verified. This is less
// important than verifying directory connectedness if we trust that GCS does
// not corrupt object metadata (where we store expected CRC32C and MD5) and
// does correctly report the object's CRC32C and MD5 sums in listings,
// verifying them periodically.
//
// Output of the following form is written to stdout and a file in the user's
// home directory:
//
//     <timestamp> <node> [<child node> ...]
//
// where:
//
//  *  Timestamps are formatted according to time.RFC3339.
//
//  *  Node names have one of two forms:
//
//     *   Nodes of the form "d:<hex score>" represent the directory listing
//         contained within the blob of the given score.
//
//     *   Nodes of the form "f:<hex score>" represent a piece of a file,
//         contained within the blob of the given score.
//
// An output line for a directory node means that at the given timestamp we
// certified that a piece of content with the given score was parseable as a
// directory listing that referred to the given scores for its direct children.
//
// An output line for a file node means that at the given timestamp we
// certified that a piece of content with the given score was parseable as a
// piece of a file. File nodes never have children.

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"os/user"
	"path"
	"runtime"
	"strings"
	"time"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/graph"
	"github.com/jacobsa/comeback/util"
	"github.com/jacobsa/comeback/verify"
	"github.com/jacobsa/comeback/wiring"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/syncutil"
	"golang.org/x/net/context"
)

var cmdVerify = &Command{
	Name: "verify",
}

// TODO(jacobsa): Get these automatically from the registry.
var fRoots = cmdVerify.Flags.String(
	"roots",
	"",
	"Comma-separated list of backup root scores to verify.")

var fFast = cmdVerify.Flags.Bool(
	"fast",
	false,
	"When set, don't verify file content.")

func init() {
	cmdVerify.Run = runVerify // Break flag-related dependency loop.
}

////////////////////////////////////////////////////////////////////////
// Visitor types
////////////////////////////////////////////////////////////////////////

// A record certifying something we confirmed about a node at a certain time.
// See the notes at the top of verify.go for details.
type verifyRecord struct {
	t        time.Time
	node     verify.Node
	adjacent []verify.Node
}

// A visitor that writes the information it gleans from the wrapped visitor to
// a channel.
type snoopingVisitor struct {
	records chan<- verifyRecord
	wrapped graph.Visitor
}

func (v *snoopingVisitor) Visit(
	ctx context.Context,
	node string) (adjacent []string, err error) {
	// Call through.
	adjacent, err = v.wrapped.Visit(ctx, node)
	if err != nil {
		return
	}

	// Build a record.
	r := verifyRecord{
		t: time.Now(),
	}

	r.node, err = verify.ParseNode(node)
	if err != nil {
		err = fmt.Errorf("ParseNode(%q): %v", node, err)
		return
	}

	for _, a := range adjacent {
		var adjacentNode verify.Node
		adjacentNode, err = verify.ParseNode(a)
		if err != nil {
			err = fmt.Errorf("ParseNode(%q): %v", a, err)
			return
		}

		r.adjacent = append(r.adjacent, adjacentNode)
	}

	// Write out the record.
	select {
	case <-ctx.Done():
		err = ctx.Err()
		return

	case v.records <- r:
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Output
////////////////////////////////////////////////////////////////////////

// Print output based on the visitor results arriving on the supplied channel.
func formatVerifyOutput(r verifyRecord) (s string) {
	var extra string
	if len(r.adjacent) != 0 {
		extra = fmt.Sprintf(" %s", strings.Join(r.adjacent, " "))
	}

	s = fmt.Sprintf(
		"%s %s%s",
		r.t.Format(time.RFC3339),
		r.node,
		extra)

	return
}

// Parse the supplied line (without line break) previously output by the verify
// command.
func parseVerifyRecord(line []byte) (r verifyRecord, err error) {
	// We expect space-separate components.
	components := bytes.Split(line, []byte{' '})
	if len(components) < 2 {
		err = fmt.Errorf(
			"Expected at least two components, got %d.",
			len(components))

		return
	}

	// The first should be the timestmap.
	r.t, err = time.Parse(time.RFC3339, string(components[0]))
	if err != nil {
		err = fmt.Errorf("time.Parse(%q): %v", components[0], err)
		return
	}

	// The next should be the node name.
	r.node = string(components[1])

	// The rest should be adjacent node names.
	for i := 2; i < len(components); i++ {
		r.adjacent = append(r.adjacent, string(components[i]))
	}

	// Make sure the record is legal.
	err = r.Check()
	if err != nil {
		err = fmt.Errorf("Check: %v", err)
		return
	}

	return
}

// Parse the supplied output from a run of the verify command, writing records
// into the supplied channel.
func parseVerifyOutput(
	ctx context.Context,
	in io.Reader,
	records chan<- verifyRecord) (err error) {
	reader := bufio.NewReader(in)

	for {
		// Find the next line. EOF with no data means we are done; otherwise ignore
		// EOF in case the file doesn't end with a newline.
		var line []byte
		line, err = reader.ReadBytes('\n')
		if err == io.EOF {
			err = nil
			if len(line) == 0 {
				break
			}
		}

		// Propagate other errors.
		if err != nil {
			err = fmt.Errorf("ReadBytes: %v", err)
			return
		}

		// Trim the delimiter.
		line = line[:len(line)-1]

		// Parse the line.
		var r verifyRecord
		r, err = parseVerifyRecord(line)
		if err != nil {
			err = fmt.Errorf("parseVerifyRecord(%q): %v", line, err)
			return
		}

		// Write out the record.
		select {
		case <-ctx.Done():
			err = ctx.Err()
			return

		case records <- r:
		}
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

// List blob.ListScores, but returns a slice instead of writing into a channel.
func listAllScores(
	ctx context.Context,
	bucket gcs.Bucket,
	namePrefix string) (scores []blob.Score, err error) {
	b := syncutil.NewBundle(context.Background())
	defer func() { err = b.Join() }()

	// List scores into a channel.
	scoreChan := make(chan blob.Score, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(scoreChan)
		err = blob.ListScores(ctx, bucket, wiring.BlobObjectNamePrefix, scoreChan)
		if err != nil {
			err = fmt.Errorf("ListScores: %v", err)
			return
		}

		return
	})

	// Accumulate into the slice.
	b.Add(func(ctx context.Context) (err error) {
		for score := range scoreChan {
			scores = append(scores, score)
		}

		return
	})

	return
}

////////////////////////////////////////////////////////////////////////
// Verify
////////////////////////////////////////////////////////////////////////

// Run the verification pipeline. Return a count of the number of scores
// verified and the number skipped due to readFiles being false.
func verifyImpl(
	ctx context.Context,
	readFiles bool,
	rootScores []blob.Score,
	knownScores []blob.Score,
	blobStore blob.Store,
	output io.Writer) (nodesVerified uint64, nodesSkipped uint64, err error) {
	b := syncutil.NewBundle(ctx)

	// Visit every node in the graph, snooping on the graph structure into a
	// channel.
	visitorRecords := make(chan verifyRecord, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(visitorRecords)

		visitor := verify.NewVisitor(
			readFiles,
			knownScores,
			blobStore)

		visitor = &snoopingVisitor{
			wrapped: visitor,
			records: visitorRecords,
		}

		// Format root node names.
		var roots []string
		for _, score := range rootScores {
			roots = append(roots, verify.FormatNodeName(true, score))
		}

		// Traverse starting at the specified roots. Use an "experimentally
		// determined" parallelism, which in theory should depend on bandwidth-delay
		// products but in practice comes down to when the OS gets cranky about open
		// files.
		const parallelism = 128

		err = graph.Traverse(
			context.Background(),
			parallelism,
			roots,
			visitor)

		if err != nil {
			err = fmt.Errorf("Traverse: %v", err)
			return
		}

		return
	})

	// Count and output the nodes visited, filtering out file nodes if we're not
	// actually reading and verifying them.
	b.Add(func(ctx context.Context) (err error) {
		for r := range visitorRecords {
			var dir bool
			dir, _, err = verify.ParseNodeName(r.node)
			if err != nil {
				err = fmt.Errorf("ParseNodeName(%q): %v", r.node, err)
				return
			}

			// Skip files if appropriate.
			if !readFiles && !dir {
				nodesSkipped++
				continue
			}

			// Increment the count.
			nodesVerified++

			// Output the information.
			_, err = fmt.Fprintf(output, "%s\n", formatVerifyOutput(r))
			if err != nil {
				err = fmt.Errorf("Fprintf: %v", err)
				return
			}
		}

		return
	})

	err = b.Join()
	return
}

// Open the file to which we log verify output.
func openVerifyLog() (w io.WriteCloser, err error) {
	// Find the current user.
	u, err := user.Current()
	if err != nil {
		err = fmt.Errorf("user.Current: %v", err)
		return
	}

	// Put the file in her home directory. Append to whatever is already there.
	w, err = os.OpenFile(
		path.Join(u.HomeDir, ".comeback.verify.log"),
		os.O_WRONLY|os.O_APPEND|os.O_CREATE,
		0600)

	if err != nil {
		err = fmt.Errorf("OpenFile: %v", err)
		return
	}

	return
}

func runVerify(args []string) {
	// Allow parallelism.
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Die on error.
	var err error
	defer func() {
		if err != nil {
			log.Fatalln(err)
		}
	}()

	// Read flags.
	readFiles := !*fFast

	if *fRoots == "" {
		err = fmt.Errorf("You must set --roots.")
		return
	}

	rootHexScores := strings.Split(*fRoots, ",")
	var rootScores []blob.Score
	for _, hexScore := range rootHexScores {
		var score blob.Score
		score, err = blob.ParseHexScore(hexScore)
		if err != nil {
			err = fmt.Errorf("Invalid root %q: %v", hexScore, err)
			return
		}

		rootScores = append(rootScores, score)
	}

	// Grab dependencies.
	bucket := getBucket()
	crypter := getCrypter()

	// Open the log file.
	logFile, err := openVerifyLog()
	if err != nil {
		err = fmt.Errorf("openVerifyLog: %v", err)
		return
	}

	defer logFile.Close()

	// Create a blob store.
	blobStore, err := wiring.MakeBlobStore(
		bucket,
		crypter,
		util.NewStringSet())

	if err != nil {
		err = fmt.Errorf("MakeBlobStore: %v", err)
		return
	}

	// List all scores in the bucket, verifying the object record metadata in the
	// process.
	log.Println("Listing scores...")
	knownScores, err := listAllScores(
		context.Background(),
		bucket,
		wiring.BlobObjectNamePrefix)

	if err != nil {
		err = fmt.Errorf("listAllScores: %v", err)
		return
	}

	log.Printf("Listed %d scores.", len(knownScores))

	// Run the rest of the pipeline.
	nodesVerified, nodesSkipped, err := verifyImpl(
		context.Background(),
		readFiles,
		rootScores,
		knownScores,
		blobStore,
		io.MultiWriter(logFile, os.Stdout))

	if err != nil {
		err = fmt.Errorf("verifyImpl: %v", err)
		return
	}

	log.Printf(
		"Successfully verified %d nodes (%d skipped due to fast mode).",
		nodesVerified,
		nodesSkipped)

	return
}
<commit_msg>formatVerifyOutput<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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.

// A command that reads all blobs necessary for verifying the directory
// structure rooted at a set of backup root scores, ensuring that the entire
// directory structure is intact in GCS.
//
// Optionally, all file content is also read and verified. This is less
// important than verifying directory connectedness if we trust that GCS does
// not corrupt object metadata (where we store expected CRC32C and MD5) and
// does correctly report the object's CRC32C and MD5 sums in listings,
// verifying them periodically.
//
// Output of the following form is written to stdout and a file in the user's
// home directory:
//
//     <timestamp> <node> [<child node> ...]
//
// where:
//
//  *  Timestamps are formatted according to time.RFC3339.
//
//  *  Node names have one of two forms:
//
//     *   Nodes of the form "d:<hex score>" represent the directory listing
//         contained within the blob of the given score.
//
//     *   Nodes of the form "f:<hex score>" represent a piece of a file,
//         contained within the blob of the given score.
//
// An output line for a directory node means that at the given timestamp we
// certified that a piece of content with the given score was parseable as a
// directory listing that referred to the given scores for its direct children.
//
// An output line for a file node means that at the given timestamp we
// certified that a piece of content with the given score was parseable as a
// piece of a file. File nodes never have children.

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"os/user"
	"path"
	"runtime"
	"strings"
	"time"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/graph"
	"github.com/jacobsa/comeback/util"
	"github.com/jacobsa/comeback/verify"
	"github.com/jacobsa/comeback/wiring"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/syncutil"
	"golang.org/x/net/context"
)

var cmdVerify = &Command{
	Name: "verify",
}

// TODO(jacobsa): Get these automatically from the registry.
var fRoots = cmdVerify.Flags.String(
	"roots",
	"",
	"Comma-separated list of backup root scores to verify.")

var fFast = cmdVerify.Flags.Bool(
	"fast",
	false,
	"When set, don't verify file content.")

func init() {
	cmdVerify.Run = runVerify // Break flag-related dependency loop.
}

////////////////////////////////////////////////////////////////////////
// Visitor types
////////////////////////////////////////////////////////////////////////

// A record certifying something we confirmed about a node at a certain time.
// See the notes at the top of verify.go for details.
type verifyRecord struct {
	t        time.Time
	node     verify.Node
	adjacent []verify.Node
}

// A visitor that writes the information it gleans from the wrapped visitor to
// a channel.
type snoopingVisitor struct {
	records chan<- verifyRecord
	wrapped graph.Visitor
}

func (v *snoopingVisitor) Visit(
	ctx context.Context,
	node string) (adjacent []string, err error) {
	// Call through.
	adjacent, err = v.wrapped.Visit(ctx, node)
	if err != nil {
		return
	}

	// Build a record.
	r := verifyRecord{
		t: time.Now(),
	}

	r.node, err = verify.ParseNode(node)
	if err != nil {
		err = fmt.Errorf("ParseNode(%q): %v", node, err)
		return
	}

	for _, a := range adjacent {
		var adjacentNode verify.Node
		adjacentNode, err = verify.ParseNode(a)
		if err != nil {
			err = fmt.Errorf("ParseNode(%q): %v", a, err)
			return
		}

		r.adjacent = append(r.adjacent, adjacentNode)
	}

	// Write out the record.
	select {
	case <-ctx.Done():
		err = ctx.Err()
		return

	case v.records <- r:
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Output
////////////////////////////////////////////////////////////////////////

// Print output based on the visitor results arriving on the supplied channel.
func formatVerifyOutput(r verifyRecord) (s string) {
	s = fmt.Sprintf(
		"%s %s",
		r.t.Format(time.RFC3339),
		r.node)

	for _, a := range r.adjacent {
		s += fmt.Sprintf(" %s", a.String())
	}

	return
}

// Parse the supplied line (without line break) previously output by the verify
// command.
func parseVerifyRecord(line []byte) (r verifyRecord, err error) {
	// We expect space-separate components.
	components := bytes.Split(line, []byte{' '})
	if len(components) < 2 {
		err = fmt.Errorf(
			"Expected at least two components, got %d.",
			len(components))

		return
	}

	// The first should be the timestmap.
	r.t, err = time.Parse(time.RFC3339, string(components[0]))
	if err != nil {
		err = fmt.Errorf("time.Parse(%q): %v", components[0], err)
		return
	}

	// The next should be the node name.
	r.node = string(components[1])

	// The rest should be adjacent node names.
	for i := 2; i < len(components); i++ {
		r.adjacent = append(r.adjacent, string(components[i]))
	}

	// Make sure the record is legal.
	err = r.Check()
	if err != nil {
		err = fmt.Errorf("Check: %v", err)
		return
	}

	return
}

// Parse the supplied output from a run of the verify command, writing records
// into the supplied channel.
func parseVerifyOutput(
	ctx context.Context,
	in io.Reader,
	records chan<- verifyRecord) (err error) {
	reader := bufio.NewReader(in)

	for {
		// Find the next line. EOF with no data means we are done; otherwise ignore
		// EOF in case the file doesn't end with a newline.
		var line []byte
		line, err = reader.ReadBytes('\n')
		if err == io.EOF {
			err = nil
			if len(line) == 0 {
				break
			}
		}

		// Propagate other errors.
		if err != nil {
			err = fmt.Errorf("ReadBytes: %v", err)
			return
		}

		// Trim the delimiter.
		line = line[:len(line)-1]

		// Parse the line.
		var r verifyRecord
		r, err = parseVerifyRecord(line)
		if err != nil {
			err = fmt.Errorf("parseVerifyRecord(%q): %v", line, err)
			return
		}

		// Write out the record.
		select {
		case <-ctx.Done():
			err = ctx.Err()
			return

		case records <- r:
		}
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

// List blob.ListScores, but returns a slice instead of writing into a channel.
func listAllScores(
	ctx context.Context,
	bucket gcs.Bucket,
	namePrefix string) (scores []blob.Score, err error) {
	b := syncutil.NewBundle(context.Background())
	defer func() { err = b.Join() }()

	// List scores into a channel.
	scoreChan := make(chan blob.Score, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(scoreChan)
		err = blob.ListScores(ctx, bucket, wiring.BlobObjectNamePrefix, scoreChan)
		if err != nil {
			err = fmt.Errorf("ListScores: %v", err)
			return
		}

		return
	})

	// Accumulate into the slice.
	b.Add(func(ctx context.Context) (err error) {
		for score := range scoreChan {
			scores = append(scores, score)
		}

		return
	})

	return
}

////////////////////////////////////////////////////////////////////////
// Verify
////////////////////////////////////////////////////////////////////////

// Run the verification pipeline. Return a count of the number of scores
// verified and the number skipped due to readFiles being false.
func verifyImpl(
	ctx context.Context,
	readFiles bool,
	rootScores []blob.Score,
	knownScores []blob.Score,
	blobStore blob.Store,
	output io.Writer) (nodesVerified uint64, nodesSkipped uint64, err error) {
	b := syncutil.NewBundle(ctx)

	// Visit every node in the graph, snooping on the graph structure into a
	// channel.
	visitorRecords := make(chan verifyRecord, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(visitorRecords)

		visitor := verify.NewVisitor(
			readFiles,
			knownScores,
			blobStore)

		visitor = &snoopingVisitor{
			wrapped: visitor,
			records: visitorRecords,
		}

		// Format root node names.
		var roots []string
		for _, score := range rootScores {
			roots = append(roots, verify.FormatNodeName(true, score))
		}

		// Traverse starting at the specified roots. Use an "experimentally
		// determined" parallelism, which in theory should depend on bandwidth-delay
		// products but in practice comes down to when the OS gets cranky about open
		// files.
		const parallelism = 128

		err = graph.Traverse(
			context.Background(),
			parallelism,
			roots,
			visitor)

		if err != nil {
			err = fmt.Errorf("Traverse: %v", err)
			return
		}

		return
	})

	// Count and output the nodes visited, filtering out file nodes if we're not
	// actually reading and verifying them.
	b.Add(func(ctx context.Context) (err error) {
		for r := range visitorRecords {
			var dir bool
			dir, _, err = verify.ParseNodeName(r.node)
			if err != nil {
				err = fmt.Errorf("ParseNodeName(%q): %v", r.node, err)
				return
			}

			// Skip files if appropriate.
			if !readFiles && !dir {
				nodesSkipped++
				continue
			}

			// Increment the count.
			nodesVerified++

			// Output the information.
			_, err = fmt.Fprintf(output, "%s\n", formatVerifyOutput(r))
			if err != nil {
				err = fmt.Errorf("Fprintf: %v", err)
				return
			}
		}

		return
	})

	err = b.Join()
	return
}

// Open the file to which we log verify output.
func openVerifyLog() (w io.WriteCloser, err error) {
	// Find the current user.
	u, err := user.Current()
	if err != nil {
		err = fmt.Errorf("user.Current: %v", err)
		return
	}

	// Put the file in her home directory. Append to whatever is already there.
	w, err = os.OpenFile(
		path.Join(u.HomeDir, ".comeback.verify.log"),
		os.O_WRONLY|os.O_APPEND|os.O_CREATE,
		0600)

	if err != nil {
		err = fmt.Errorf("OpenFile: %v", err)
		return
	}

	return
}

func runVerify(args []string) {
	// Allow parallelism.
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Die on error.
	var err error
	defer func() {
		if err != nil {
			log.Fatalln(err)
		}
	}()

	// Read flags.
	readFiles := !*fFast

	if *fRoots == "" {
		err = fmt.Errorf("You must set --roots.")
		return
	}

	rootHexScores := strings.Split(*fRoots, ",")
	var rootScores []blob.Score
	for _, hexScore := range rootHexScores {
		var score blob.Score
		score, err = blob.ParseHexScore(hexScore)
		if err != nil {
			err = fmt.Errorf("Invalid root %q: %v", hexScore, err)
			return
		}

		rootScores = append(rootScores, score)
	}

	// Grab dependencies.
	bucket := getBucket()
	crypter := getCrypter()

	// Open the log file.
	logFile, err := openVerifyLog()
	if err != nil {
		err = fmt.Errorf("openVerifyLog: %v", err)
		return
	}

	defer logFile.Close()

	// Create a blob store.
	blobStore, err := wiring.MakeBlobStore(
		bucket,
		crypter,
		util.NewStringSet())

	if err != nil {
		err = fmt.Errorf("MakeBlobStore: %v", err)
		return
	}

	// List all scores in the bucket, verifying the object record metadata in the
	// process.
	log.Println("Listing scores...")
	knownScores, err := listAllScores(
		context.Background(),
		bucket,
		wiring.BlobObjectNamePrefix)

	if err != nil {
		err = fmt.Errorf("listAllScores: %v", err)
		return
	}

	log.Printf("Listed %d scores.", len(knownScores))

	// Run the rest of the pipeline.
	nodesVerified, nodesSkipped, err := verifyImpl(
		context.Background(),
		readFiles,
		rootScores,
		knownScores,
		blobStore,
		io.MultiWriter(logFile, os.Stdout))

	if err != nil {
		err = fmt.Errorf("verifyImpl: %v", err)
		return
	}

	log.Printf(
		"Successfully verified %d nodes (%d skipped due to fast mode).",
		nodesVerified,
		nodesSkipped)

	return
}
<|endoftext|>


<commit_before>// Copyright 2022 The etcd Authors
//
// 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 common

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	clientv3 "go.etcd.io/etcd/client/v3"
	"go.etcd.io/etcd/tests/v3/framework/config"
	"go.etcd.io/etcd/tests/v3/framework/testutils"
)

func TestLeaseGrantTimeToLive(t *testing.T) {
	testRunner.BeforeTest(t)

	tcs := []struct {
		name   string
		config config.ClusterConfig
	}{
		{
			name:   "NoTLS",
			config: config.ClusterConfig{ClusterSize: 1},
		},
		{
			name:   "PeerTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
		},
		{
			name:   "PeerAutoTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
		},
		{
			name:   "ClientTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
		},
		{
			name:   "ClientAutoTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			clus := testRunner.NewCluster(t, tc.config)
			defer clus.Close()
			cc := clus.Client()

			testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
				ttl := int64(10)
				leaseResp, err := cc.Grant(ttl)
				require.NoError(t, err)

				ttlResp, err := cc.TimeToLive(leaseResp.ID, config.LeaseOption{})
				require.NoError(t, err)
				require.Equal(t, ttl, ttlResp.GrantedTTL)
			})
		})
	}
}

func TestLeaseGrantAndList(t *testing.T) {
	testRunner.BeforeTest(t)

	tcs := []struct {
		name   string
		config config.ClusterConfig
	}{
		{
			name:   "NoTLS",
			config: config.ClusterConfig{ClusterSize: 1},
		},
		{
			name:   "PeerTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
		},
		{
			name:   "PeerAutoTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
		},
		{
			name:   "ClientTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
		},
		{
			name:   "ClientAutoTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
		},
	}
	for _, tc := range tcs {
		nestedCases := []struct {
			name       string
			leaseCount int
		}{
			{
				name:       "no_leases",
				leaseCount: 0,
			},
			{
				name:       "one_lease",
				leaseCount: 1,
			},
			{
				name:       "many_leases",
				leaseCount: 3,
			},
		}

		for _, nc := range nestedCases {
			t.Run(tc.name+"/"+nc.name, func(t *testing.T) {
				t.Logf("Creating cluster...")
				clus := testRunner.NewCluster(t, tc.config)
				defer clus.Close()
				cc := clus.Client()
				t.Logf("Created cluster and client")
				testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
					createdLeases := []clientv3.LeaseID{}
					lastRev := int64(0)
					for i := 0; i < nc.leaseCount; i++ {
						leaseResp, err := cc.Grant(10)
						t.Logf("Grant returned: resp:%s err:%v", leaseResp.String(), err)
						require.NoError(t, err)
						createdLeases = append(createdLeases, leaseResp.ID)
						lastRev = leaseResp.GetRevision()
					}

					// Because we're not guarunteed to talk to the same member, wait for
					// listing to eventually return true, either by the result propagaing
					// or by hitting an up to date member.
					leases := []clientv3.LeaseStatus{}
					require.Eventually(t, func() bool {
						resp, err := cc.LeaseList()
						if err != nil {
							return false
						}
						leases = resp.Leases
						return resp.GetRevision() >= lastRev
					}, 2*time.Second, 10*time.Millisecond)

					returnedLeases := make([]clientv3.LeaseID, 0, nc.leaseCount)
					for _, status := range leases {
						returnedLeases = append(returnedLeases, status.ID)
					}

					require.ElementsMatch(t, createdLeases, returnedLeases)
				})
			})
		}
	}
}

func TestLeaseGrantTimeToLiveExpired(t *testing.T) {
	testRunner.BeforeTest(t)

	tcs := []struct {
		name   string
		config config.ClusterConfig
	}{
		{
			name:   "NoTLS",
			config: config.ClusterConfig{ClusterSize: 1},
		},
		{
			name:   "PeerTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
		},
		{
			name:   "PeerAutoTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
		},
		{
			name:   "ClientTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
		},
		{
			name:   "ClientAutoTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			clus := testRunner.NewCluster(t, tc.config)
			defer clus.Close()
			cc := clus.Client()

			testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
				leaseResp, err := cc.Grant(2)
				require.NoError(t, err)

				err = cc.Put("foo", "bar", config.PutOptions{LeaseID: leaseResp.ID})
				require.NoError(t, err)

				getResp, err := cc.Get("foo", config.GetOptions{})
				require.NoError(t, err)
				require.Equal(t, int64(1), getResp.Count)

				time.Sleep(3 * time.Second)

				ttlResp, err := cc.TimeToLive(leaseResp.ID, config.LeaseOption{})
				require.NoError(t, err)
				require.Equal(t, int64(-1), ttlResp.TTL)

				getResp, err = cc.Get("foo", config.GetOptions{})
				require.NoError(t, err)
				// Value should expire with the lease
				require.Equal(t, int64(0), getResp.Count)
			})
		})
	}
}

func TestLeaseGrantKeepAliveOnce(t *testing.T) {
	testRunner.BeforeTest(t)

	tcs := []struct {
		name   string
		config config.ClusterConfig
	}{
		{
			name:   "NoTLS",
			config: config.ClusterConfig{ClusterSize: 1},
		},
		{
			name:   "PeerTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
		},
		{
			name:   "PeerAutoTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
		},
		{
			name:   "ClientTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
		},
		{
			name:   "ClientAutoTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			clus := testRunner.NewCluster(t, tc.config)
			defer clus.Close()
			cc := clus.Client()

			testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
				leaseResp, err := cc.Grant(2)
				require.NoError(t, err)

				_, err = cc.LeaseKeepAliveOnce(leaseResp.ID)
				require.NoError(t, err)

				time.Sleep(2 * time.Second) // Wait for the original lease to expire

				ttlResp, err := cc.TimeToLive(leaseResp.ID, config.LeaseOption{})
				require.NoError(t, err)
				// We still have a lease!
				require.Greater(t, int64(2), ttlResp.TTL)
			})
		})
	}
}

func TestLeaseGrantRevoke(t *testing.T) {
	testRunner.BeforeTest(t)

	tcs := []struct {
		name   string
		config config.ClusterConfig
	}{
		{
			name:   "NoTLS",
			config: config.ClusterConfig{ClusterSize: 1},
		},
		{
			name:   "PeerTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
		},
		{
			name:   "PeerAutoTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
		},
		{
			name:   "ClientTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
		},
		{
			name:   "ClientAutoTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			clus := testRunner.NewCluster(t, tc.config)
			defer clus.Close()
			cc := clus.Client()

			testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
				leaseResp, err := cc.Grant(20)
				require.NoError(t, err)

				err = cc.Put("foo", "bar", config.PutOptions{LeaseID: leaseResp.ID})
				require.NoError(t, err)

				getResp, err := cc.Get("foo", config.GetOptions{})
				require.NoError(t, err)
				require.Equal(t, int64(1), getResp.Count)

				_, err = cc.LeaseRevoke(leaseResp.ID)
				require.NoError(t, err)

				ttlResp, err := cc.TimeToLive(leaseResp.ID, config.LeaseOption{})
				require.NoError(t, err)
				require.Equal(t, int64(-1), ttlResp.TTL)

				getResp, err = cc.Get("foo", config.GetOptions{})
				require.NoError(t, err)
				// Value should expire with the lease
				require.Equal(t, int64(0), getResp.Count)
			})
		})
	}
}
<commit_msg>tests/common/lease: don't use revision to wait for leases<commit_after>// Copyright 2022 The etcd Authors
//
// 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 common

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	clientv3 "go.etcd.io/etcd/client/v3"
	"go.etcd.io/etcd/tests/v3/framework/config"
	"go.etcd.io/etcd/tests/v3/framework/testutils"
)

func TestLeaseGrantTimeToLive(t *testing.T) {
	testRunner.BeforeTest(t)

	tcs := []struct {
		name   string
		config config.ClusterConfig
	}{
		{
			name:   "NoTLS",
			config: config.ClusterConfig{ClusterSize: 1},
		},
		{
			name:   "PeerTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
		},
		{
			name:   "PeerAutoTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
		},
		{
			name:   "ClientTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
		},
		{
			name:   "ClientAutoTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			clus := testRunner.NewCluster(t, tc.config)
			defer clus.Close()
			cc := clus.Client()

			testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
				ttl := int64(10)
				leaseResp, err := cc.Grant(ttl)
				require.NoError(t, err)

				ttlResp, err := cc.TimeToLive(leaseResp.ID, config.LeaseOption{})
				require.NoError(t, err)
				require.Equal(t, ttl, ttlResp.GrantedTTL)
			})
		})
	}
}

func TestLeaseGrantAndList(t *testing.T) {
	testRunner.BeforeTest(t)

	tcs := []struct {
		name   string
		config config.ClusterConfig
	}{
		{
			name:   "NoTLS",
			config: config.ClusterConfig{ClusterSize: 1},
		},
		{
			name:   "PeerTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
		},
		{
			name:   "PeerAutoTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
		},
		{
			name:   "ClientTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
		},
		{
			name:   "ClientAutoTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
		},
	}
	for _, tc := range tcs {
		nestedCases := []struct {
			name       string
			leaseCount int
		}{
			{
				name:       "no_leases",
				leaseCount: 0,
			},
			{
				name:       "one_lease",
				leaseCount: 1,
			},
			{
				name:       "many_leases",
				leaseCount: 3,
			},
		}

		for _, nc := range nestedCases {
			t.Run(tc.name+"/"+nc.name, func(t *testing.T) {
				t.Logf("Creating cluster...")
				clus := testRunner.NewCluster(t, tc.config)
				defer clus.Close()
				cc := clus.Client()
				t.Logf("Created cluster and client")
				testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
					createdLeases := []clientv3.LeaseID{}
					for i := 0; i < nc.leaseCount; i++ {
						leaseResp, err := cc.Grant(10)
						t.Logf("Grant returned: resp:%s err:%v", leaseResp.String(), err)
						require.NoError(t, err)
						createdLeases = append(createdLeases, leaseResp.ID)
					}

					// Because we're not guarunteed to talk to the same member, wait for
					// listing to eventually return true, either by the result propagaing
					// or by hitting an up to date member.
					leases := []clientv3.LeaseStatus{}
					require.Eventually(t, func() bool {
						resp, err := cc.LeaseList()
						if err != nil {
							return false
						}
						leases = resp.Leases
						// TODO: update this to use last Revision from leaseResp
						// after https://github.com/etcd-io/etcd/issues/13989 is fixed
						return len(leases) == len(createdLeases)
					}, 2*time.Second, 10*time.Millisecond)

					returnedLeases := make([]clientv3.LeaseID, 0, nc.leaseCount)
					for _, status := range leases {
						returnedLeases = append(returnedLeases, status.ID)
					}

					require.ElementsMatch(t, createdLeases, returnedLeases)
				})
			})
		}
	}
}

func TestLeaseGrantTimeToLiveExpired(t *testing.T) {
	testRunner.BeforeTest(t)

	tcs := []struct {
		name   string
		config config.ClusterConfig
	}{
		{
			name:   "NoTLS",
			config: config.ClusterConfig{ClusterSize: 1},
		},
		{
			name:   "PeerTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
		},
		{
			name:   "PeerAutoTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
		},
		{
			name:   "ClientTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
		},
		{
			name:   "ClientAutoTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			clus := testRunner.NewCluster(t, tc.config)
			defer clus.Close()
			cc := clus.Client()

			testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
				leaseResp, err := cc.Grant(2)
				require.NoError(t, err)

				err = cc.Put("foo", "bar", config.PutOptions{LeaseID: leaseResp.ID})
				require.NoError(t, err)

				getResp, err := cc.Get("foo", config.GetOptions{})
				require.NoError(t, err)
				require.Equal(t, int64(1), getResp.Count)

				time.Sleep(3 * time.Second)

				ttlResp, err := cc.TimeToLive(leaseResp.ID, config.LeaseOption{})
				require.NoError(t, err)
				require.Equal(t, int64(-1), ttlResp.TTL)

				getResp, err = cc.Get("foo", config.GetOptions{})
				require.NoError(t, err)
				// Value should expire with the lease
				require.Equal(t, int64(0), getResp.Count)
			})
		})
	}
}

func TestLeaseGrantKeepAliveOnce(t *testing.T) {
	testRunner.BeforeTest(t)

	tcs := []struct {
		name   string
		config config.ClusterConfig
	}{
		{
			name:   "NoTLS",
			config: config.ClusterConfig{ClusterSize: 1},
		},
		{
			name:   "PeerTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
		},
		{
			name:   "PeerAutoTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
		},
		{
			name:   "ClientTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
		},
		{
			name:   "ClientAutoTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			clus := testRunner.NewCluster(t, tc.config)
			defer clus.Close()
			cc := clus.Client()

			testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
				leaseResp, err := cc.Grant(2)
				require.NoError(t, err)

				_, err = cc.LeaseKeepAliveOnce(leaseResp.ID)
				require.NoError(t, err)

				time.Sleep(2 * time.Second) // Wait for the original lease to expire

				ttlResp, err := cc.TimeToLive(leaseResp.ID, config.LeaseOption{})
				require.NoError(t, err)
				// We still have a lease!
				require.Greater(t, int64(2), ttlResp.TTL)
			})
		})
	}
}

func TestLeaseGrantRevoke(t *testing.T) {
	testRunner.BeforeTest(t)

	tcs := []struct {
		name   string
		config config.ClusterConfig
	}{
		{
			name:   "NoTLS",
			config: config.ClusterConfig{ClusterSize: 1},
		},
		{
			name:   "PeerTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.ManualTLS},
		},
		{
			name:   "PeerAutoTLS",
			config: config.ClusterConfig{ClusterSize: 3, PeerTLS: config.AutoTLS},
		},
		{
			name:   "ClientTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.ManualTLS},
		},
		{
			name:   "ClientAutoTLS",
			config: config.ClusterConfig{ClusterSize: 1, ClientTLS: config.AutoTLS},
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			clus := testRunner.NewCluster(t, tc.config)
			defer clus.Close()
			cc := clus.Client()

			testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
				leaseResp, err := cc.Grant(20)
				require.NoError(t, err)

				err = cc.Put("foo", "bar", config.PutOptions{LeaseID: leaseResp.ID})
				require.NoError(t, err)

				getResp, err := cc.Get("foo", config.GetOptions{})
				require.NoError(t, err)
				require.Equal(t, int64(1), getResp.Count)

				_, err = cc.LeaseRevoke(leaseResp.ID)
				require.NoError(t, err)

				ttlResp, err := cc.TimeToLive(leaseResp.ID, config.LeaseOption{})
				require.NoError(t, err)
				require.Equal(t, int64(-1), ttlResp.TTL)

				getResp, err = cc.Get("foo", config.GetOptions{})
				require.NoError(t, err)
				// Value should expire with the lease
				require.Equal(t, int64(0), getResp.Count)
			})
		})
	}
}
<|endoftext|>


<commit_before>package checks

import (
	"testing"

	"github.com/mackerelio/mackerel-agent/config"
)

func TestChecker_Check(t *testing.T) {
	checkerOK := Checker{
		Config: config.PluginConfig{
			Command: "go run testdata/exit.go -code 0 -message OK",
		},
	}

	checkerWarning := Checker{
		Config: config.PluginConfig{
			Command: "go run testdata/exit.go -code 1 -message 'something is going wrong'",
		},
	}

	{
		report, err := checkerOK.Check()
		if err != nil {
			t.Errorf("err should be nil: %v", err)
		}
		if report.Status != StatusOK {
			t.Errorf("status should be OK: %v", report.Status)
		}
		if report.Message != "OK\n" {
			t.Errorf("wrong message: %q", report.Message)
		}
	}

	{
		report, err := checkerWarning.Check()
		if err != nil {
			t.Errorf("err should be nil: %v", err)
		}
		if report.Status != StatusWarning {
			t.Errorf("status should be WARNING: %v", report.Status)
		}
		if report.Message != "something is going wrong\n" {
			t.Errorf("wrong message: %q", report.Message)
		}
	}
}
<commit_msg>quotation issssue<commit_after>package checks

import (
	"testing"

	"github.com/mackerelio/mackerel-agent/config"
)

func TestChecker_Check(t *testing.T) {
	checkerOK := Checker{
		Config: config.PluginConfig{
			Command: "go run testdata/exit.go -code 0 -message OK",
		},
	}

	checkerWarning := Checker{
		Config: config.PluginConfig{
			Command: "go run testdata/exit.go -code 1 -message something_is_going_wrong",
		},
	}

	{
		report, err := checkerOK.Check()
		if err != nil {
			t.Errorf("err should be nil: %v", err)
		}
		if report.Status != StatusOK {
			t.Errorf("status should be OK: %v", report.Status)
		}
		if report.Message != "OK\n" {
			t.Errorf("wrong message: %q", report.Message)
		}
	}

	{
		report, err := checkerWarning.Check()
		if err != nil {
			t.Errorf("err should be nil: %v", err)
		}
		if report.Status != StatusWarning {
			t.Errorf("status should be WARNING: %v", report.Status)
		}
		if report.Message != "something_is_going_wrong\n" {
			t.Errorf("wrong message: %q", report.Message)
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2013 The Camlistore Authors

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 storagetest tests blobserver.Storage implementations
package storagetest

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"testing"
	"time"

	"camlistore.org/pkg/blob"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/context"
	"camlistore.org/pkg/syncutil"
	"camlistore.org/pkg/test"
)

type Opts struct {
	// New is required and must return the storage server to test, along with a func to
	// clean it up. The cleanup may be nil.
	New func(*testing.T) (sto blobserver.Storage, cleanup func())

	// Retries specifies how long to wait to retry after each failure
	// that may be an eventual consistency issue (enumerate, stat), etc.
	Retries []time.Duration

	SkipEnum bool // for when EnumerateBlobs is not implemented
}

func Test(t *testing.T, fn func(*testing.T) (sto blobserver.Storage, cleanup func())) {
	TestOpt(t, Opts{New: fn})
}

type run struct {
	t   *testing.T
	opt Opts
	sto blobserver.Storage
}

func TestOpt(t *testing.T, opt Opts) {
	sto, cleanup := opt.New(t)
	defer func() {
		if t.Failed() {
			t.Logf("test %T FAILED, skipping cleanup!", sto)
		} else {
			if cleanup != nil {
				cleanup()
			}
		}
	}()
	r := &run{
		t:   t,
		opt: opt,
		sto: sto,
	}
	t.Logf("Testing blobserver storage %T", sto)

	t.Logf("Testing Enumerate for empty")
	r.testEnumerate(nil)

	var blobs []*test.Blob
	var blobRefs []blob.Ref
	var blobSizedRefs []blob.SizedRef

	contents := []string{"foo", "quux", "asdf", "qwerty", "0123456789"}
	if !testing.Short() {
		for i := 0; i < 95; i++ {
			contents = append(contents, "foo-"+strconv.Itoa(i))
		}
	}
	t.Logf("Testing receive")
	for i, x := range contents {
		b1 := &test.Blob{x}
		if testing.Short() {
			t.Logf("blob[%d] = %s: %q", i, b1.BlobRef(), x)
		}
		b1s, err := sto.ReceiveBlob(b1.BlobRef(), b1.Reader())
		if err != nil {
			t.Fatalf("ReceiveBlob of %s: %v", b1, err)
		}
		if b1s != b1.SizedRef() {
			t.Fatalf("Received %v; want %v", b1s, b1.SizedRef())
		}
		blobs = append(blobs, b1)
		blobRefs = append(blobRefs, b1.BlobRef())
		blobSizedRefs = append(blobSizedRefs, b1.SizedRef())

		switch len(blobSizedRefs) {
		case 1, 5, 100:
			t.Logf("Testing Enumerate for %d blobs", len(blobSizedRefs))
			r.testEnumerate(blobSizedRefs)
		}
	}
	b1 := blobs[0]

	t.Logf("Testing Fetch")
	for i, b2 := range blobs {
		rc, size, err := sto.Fetch(b2.BlobRef())
		if err != nil {
			t.Fatalf("error fetching %d. %s: %v", i, b2, err)
		}
		defer rc.Close()
		testSizedBlob(t, rc, b2.BlobRef(), int64(size))
	}

	t.Logf("Testing Stat")
	dest := make(chan blob.SizedRef)
	errc := make(chan error, 1)
	go func() {
		errc <- sto.StatBlobs(dest, blobRefs)
	}()
	testStat(t, dest, blobSizedRefs)
	if err := <-errc; err != nil {
		t.Fatalf("error stating blobs %s: %v", blobRefs, err)
	}

	// Enumerate tests.
	sort.Sort(blob.SizedByRef(blobSizedRefs))

	t.Logf("Testing Enumerate on all")
	r.testEnumerate(blobSizedRefs)

	t.Logf("Testing Enumerate 'limit' param")
	r.testEnumerate(blobSizedRefs[:3], 3)

	// Enumerate 'after'
	{
		after := blobSizedRefs[2].Ref.String()
		t.Logf("Testing Enumerate 'after' param; after %q", after)
		r.testEnumerate(blobSizedRefs[3:], after)
	}

	// Enumerate 'after' + limit
	{
		after := blobSizedRefs[2].Ref.String()
		t.Logf("Testing Enumerate 'after' + 'limit' param; after %q, limit 1", after)
		r.testEnumerate(blobSizedRefs[3:4], after, 1)
	}

	// Enumerate 'after' with prefix of a blobref + limit
	{
		after := "a"
		t.Logf("Testing Enumerate 'after' + 'limit' param; after %q, limit 1", after)
		r.testEnumerate(blobSizedRefs[:1], after, 1)
	}

	t.Logf("Testing Remove")
	if err := sto.RemoveBlobs(blobRefs); err != nil {
		if strings.Contains(err.Error(), "not implemented") {
			t.Logf("RemoveBlob %s: %v", b1, err)
		} else {
			t.Fatalf("RemoveBlob %s: %v", b1, err)
		}
	}

	testSubFetcher(t, sto)
}

func testSubFetcher(t *testing.T, sto blobserver.Storage) {
	sf, ok := sto.(blob.SubFetcher)
	if !ok {
		t.Logf("%T is not a SubFetcher", sto)
		return
	}
	t.Logf("Testing SubFetch")
	big := &test.Blob{"Some big blob"}
	if _, err := sto.ReceiveBlob(big.BlobRef(), big.Reader()); err != nil {
		t.Fatal(err)
	}
	regions := []struct {
		off, limit int64
		want       string
		errok      bool
	}{
		{5, 3, "big", false},
		{5, 8, "big blob", false},
		{5, 100, "big blob", true},
	}
	for _, tt := range regions {
		r, err := sf.SubFetch(big.BlobRef(), tt.off, tt.limit)
		if err != nil {
			t.Fatalf("Error fetching big blob for SubFetch: %v", err)
		}
		if r == nil {
			t.Fatal("SubFetch returned nil, nil")
		}
		all, err := ioutil.ReadAll(r)
		r.Close()
		if err != nil && !tt.errok {
			t.Errorf("Unexpected error reading SubFetch region %+v: %v", tt, err)
		}
		if string(all) != tt.want {
			t.Errorf("SubFetch region %+v got %q; want %q", tt, all, tt.want)
		}
	}
}

func testSizedBlob(t *testing.T, r io.Reader, b1 blob.Ref, size int64) {
	h := b1.Hash()
	n, err := io.Copy(h, r)
	if err != nil {
		t.Fatalf("error reading from %s: %v", r, err)
	}
	if n != size {
		t.Fatalf("read %d bytes from %s, metadata said %d!", n, r, size)
	}
	b2 := blob.RefFromHash(h)
	if b2 != b1 {
		t.Fatalf("content mismatch (awaited %s, got %s)", b1, b2)
	}
}

func CheckEnumerate(sto blobserver.Storage, wantUnsorted []blob.SizedRef, opts ...interface{}) error {
	var after string
	var n = 1000
	for _, opt := range opts {
		switch v := opt.(type) {
		case string:
			after = v
		case int:
			n = v
		default:
			panic("bad option of type " + fmt.Sprintf("%T", v))
		}
	}

	want := append([]blob.SizedRef(nil), wantUnsorted...)
	sort.Sort(blob.SizedByRef(want))

	sbc := make(chan blob.SizedRef, 10)

	var got []blob.SizedRef
	var grp syncutil.Group
	sawEnd := make(chan bool, 1)
	grp.Go(func() error {
		ctx := context.New()
		defer ctx.Cancel()
		if err := sto.EnumerateBlobs(ctx, sbc, after, n); err != nil {
			return fmt.Errorf("EnumerateBlobs(%q, %d): %v", after, n, err)
		}
		return nil
	})
	grp.Go(func() error {
		var lastRef blob.Ref
		for sb := range sbc {
			if !sb.Valid() {
				return fmt.Errorf("invalid blobref %#v received in enumerate", sb)
			}
			got = append(got, sb)
			if lastRef.Valid() && sb.Ref.Less(lastRef) {
				return fmt.Errorf("blobs appearing out of order")
			}
			lastRef = sb.Ref
		}
		sawEnd <- true
		return nil

	})
	grp.Go(func() error {
		select {
		case <-sawEnd:
			return nil
		case <-time.After(10 * time.Second):
			return errors.New("timeout waiting for EnumerateBlobs to close its channel")
		}

	})
	if err := grp.Err(); err != nil {
		return fmt.Errorf("Enumerate error: %v", err)
	}
	if len(got) == 0 && len(want) == 0 {
		return nil
	}
	var gotSet = map[blob.SizedRef]bool{}
	for _, sb := range got {
		if gotSet[sb] {
			return fmt.Errorf("duplicate blob %v returned in enumerate", sb)
		}
		gotSet[sb] = true
	}

	if !reflect.DeepEqual(got, want) {
		return fmt.Errorf("Enumerate mismatch. Got %d; want %d.\n Got: %v\nWant: %v\n",
			len(got), len(want), got, want)
	}
	return nil
}

func (r *run) testEnumerate(wantUnsorted []blob.SizedRef, opts ...interface{}) {
	if r.opt.SkipEnum {
		r.t.Log("Skipping enum test")
		return
	}
	if err := r.withRetries(func() error {
		return CheckEnumerate(r.sto, wantUnsorted, opts...)
	}); err != nil {
		r.t.Fatalf("%v", err)
	}
}

func (r *run) withRetries(fn func() error) error {
	delays := r.opt.Retries
	for {
		err := fn()
		if err == nil || len(delays) == 0 {
			return err
		}
		r.t.Logf("(operation failed; retrying after %v)", delays[0])
		time.Sleep(delays[0])
		delays = delays[1:]
	}
}

func testStat(t *testing.T, enum <-chan blob.SizedRef, want []blob.SizedRef) {
	// blobs may arrive in ANY order
	m := make(map[string]int, len(want))
	for i, sb := range want {
		m[sb.Ref.String()] = i
	}

	i := 0
	for sb := range enum {
		if !sb.Valid() {
			break
		}
		wanted := want[m[sb.Ref.String()]]
		if wanted.Size != sb.Size {
			t.Fatalf("received blob size is %d, wanted %d for &%d", sb.Size, wanted.Size, i)
		}
		if wanted.Ref != sb.Ref {
			t.Fatalf("received blob ref mismatch &%d: wanted %s, got %s", i, sb.Ref, wanted.Ref)
		}
		i++
		if i >= len(want) {
			break
		}
	}
}
<commit_msg>storagetest: clean up a bit<commit_after>/*
Copyright 2013 The Camlistore Authors

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 storagetest tests blobserver.Storage implementations
package storagetest

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"testing"
	"time"

	"camlistore.org/pkg/blob"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/context"
	"camlistore.org/pkg/syncutil"
	"camlistore.org/pkg/test"
)

type Opts struct {
	// New is required and must return the storage server to test, along with a func to
	// clean it up. The cleanup may be nil.
	New func(*testing.T) (sto blobserver.Storage, cleanup func())

	// Retries specifies how long to wait to retry after each failure
	// that may be an eventual consistency issue (enumerate, stat), etc.
	Retries []time.Duration

	SkipEnum bool // for when EnumerateBlobs is not implemented
}

func Test(t *testing.T, fn func(*testing.T) (sto blobserver.Storage, cleanup func())) {
	TestOpt(t, Opts{New: fn})
}

type run struct {
	t   *testing.T
	opt Opts
	sto blobserver.Storage
}

func TestOpt(t *testing.T, opt Opts) {
	sto, cleanup := opt.New(t)
	defer func() {
		if t.Failed() {
			t.Logf("test %T FAILED, skipping cleanup!", sto)
		} else {
			if cleanup != nil {
				cleanup()
			}
		}
	}()
	r := &run{
		t:   t,
		opt: opt,
		sto: sto,
	}
	t.Logf("Testing blobserver storage %T", sto)

	t.Logf("Testing Enumerate for empty")
	r.testEnumerate(nil)

	var blobs []*test.Blob
	var blobRefs []blob.Ref
	var blobSizedRefs []blob.SizedRef

	contents := []string{"foo", "quux", "asdf", "qwerty", "0123456789"}
	if !testing.Short() {
		for i := 0; i < 95; i++ {
			contents = append(contents, "foo-"+strconv.Itoa(i))
		}
	}
	t.Logf("Testing receive")
	for i, x := range contents {
		b1 := &test.Blob{x}
		if testing.Short() {
			t.Logf("blob[%d] = %s: %q", i, b1.BlobRef(), x)
		}
		b1s, err := sto.ReceiveBlob(b1.BlobRef(), b1.Reader())
		if err != nil {
			t.Fatalf("ReceiveBlob of %s: %v", b1, err)
		}
		if b1s != b1.SizedRef() {
			t.Fatalf("Received %v; want %v", b1s, b1.SizedRef())
		}
		blobs = append(blobs, b1)
		blobRefs = append(blobRefs, b1.BlobRef())
		blobSizedRefs = append(blobSizedRefs, b1.SizedRef())

		switch len(blobSizedRefs) {
		case 1, 5, 100:
			t.Logf("Testing Enumerate for %d blobs", len(blobSizedRefs))
			r.testEnumerate(blobSizedRefs)
		}
	}

	t.Logf("Testing Fetch")
	for i, b2 := range blobs {
		rc, size, err := sto.Fetch(b2.BlobRef())
		if err != nil {
			t.Fatalf("error fetching %d. %s: %v", i, b2, err)
		}
		defer rc.Close()
		testSizedBlob(t, rc, b2.BlobRef(), int64(size))
	}

	t.Logf("Testing Stat")
	dest := make(chan blob.SizedRef)
	errc := make(chan error, 1)
	go func() {
		errc <- sto.StatBlobs(dest, blobRefs)
	}()
	testStat(t, dest, blobSizedRefs)
	if err := <-errc; err != nil {
		t.Fatalf("error stating blobs %s: %v", blobRefs, err)
	}

	// Enumerate tests.
	sort.Sort(blob.SizedByRef(blobSizedRefs))

	t.Logf("Testing Enumerate on all")
	r.testEnumerate(blobSizedRefs)

	t.Logf("Testing Enumerate 'limit' param")
	r.testEnumerate(blobSizedRefs[:3], 3)

	// Enumerate 'after'
	{
		after := blobSizedRefs[2].Ref.String()
		t.Logf("Testing Enumerate 'after' param; after %q", after)
		r.testEnumerate(blobSizedRefs[3:], after)
	}

	// Enumerate 'after' + limit
	{
		after := blobSizedRefs[2].Ref.String()
		t.Logf("Testing Enumerate 'after' + 'limit' param; after %q, limit 1", after)
		r.testEnumerate(blobSizedRefs[3:4], after, 1)
	}

	// Enumerate 'after' with prefix of a blobref + limit
	{
		after := "a"
		t.Logf("Testing Enumerate 'after' + 'limit' param; after %q, limit 1", after)
		r.testEnumerate(blobSizedRefs[:1], after, 1)
	}

	r.testRemove(blobRefs)

	r.testSubFetcher()
}

func (r *run) testRemove(blobRefs []blob.Ref) {
	t, sto := r.t, r.sto
	t.Logf("Testing Remove")
	if err := sto.RemoveBlobs(blobRefs); err != nil {
		if strings.Contains(err.Error(), "not implemented") {
			t.Logf("RemoveBlobs: %v", err)
		} else {
			t.Fatalf("RemoveBlobs: %v", err)
		}
	}
	r.testEnumerate(nil) // verify they're all gone
}

func (r *run) testSubFetcher() {
	t, sto := r.t, r.sto
	sf, ok := sto.(blob.SubFetcher)
	if !ok {
		t.Logf("%T is not a SubFetcher", sto)
		return
	}
	t.Logf("Testing SubFetch")
	big := &test.Blob{"Some big blob"}
	if _, err := sto.ReceiveBlob(big.BlobRef(), big.Reader()); err != nil {
		t.Fatal(err)
	}
	regions := []struct {
		off, limit int64
		want       string
		errok      bool
	}{
		{5, 3, "big", false},
		{5, 8, "big blob", false},
		{5, 100, "big blob", true},
	}
	for _, tt := range regions {
		r, err := sf.SubFetch(big.BlobRef(), tt.off, tt.limit)
		if err != nil {
			t.Fatalf("Error fetching big blob for SubFetch: %v", err)
		}
		if r == nil {
			t.Fatal("SubFetch returned nil, nil")
		}
		all, err := ioutil.ReadAll(r)
		r.Close()
		if err != nil && !tt.errok {
			t.Errorf("Unexpected error reading SubFetch region %+v: %v", tt, err)
		}
		if string(all) != tt.want {
			t.Errorf("SubFetch region %+v got %q; want %q", tt, all, tt.want)
		}
	}
}

func testSizedBlob(t *testing.T, r io.Reader, b1 blob.Ref, size int64) {
	h := b1.Hash()
	n, err := io.Copy(h, r)
	if err != nil {
		t.Fatalf("error reading from %s: %v", r, err)
	}
	if n != size {
		t.Fatalf("read %d bytes from %s, metadata said %d!", n, r, size)
	}
	b2 := blob.RefFromHash(h)
	if b2 != b1 {
		t.Fatalf("content mismatch (awaited %s, got %s)", b1, b2)
	}
}

func CheckEnumerate(sto blobserver.Storage, wantUnsorted []blob.SizedRef, opts ...interface{}) error {
	var after string
	var n = 1000
	for _, opt := range opts {
		switch v := opt.(type) {
		case string:
			after = v
		case int:
			n = v
		default:
			panic("bad option of type " + fmt.Sprintf("%T", v))
		}
	}

	want := append([]blob.SizedRef(nil), wantUnsorted...)
	sort.Sort(blob.SizedByRef(want))

	sbc := make(chan blob.SizedRef, 10)

	var got []blob.SizedRef
	var grp syncutil.Group
	sawEnd := make(chan bool, 1)
	grp.Go(func() error {
		ctx := context.New()
		defer ctx.Cancel()
		if err := sto.EnumerateBlobs(ctx, sbc, after, n); err != nil {
			return fmt.Errorf("EnumerateBlobs(%q, %d): %v", after, n, err)
		}
		return nil
	})
	grp.Go(func() error {
		var lastRef blob.Ref
		for sb := range sbc {
			if !sb.Valid() {
				return fmt.Errorf("invalid blobref %#v received in enumerate", sb)
			}
			got = append(got, sb)
			if lastRef.Valid() && sb.Ref.Less(lastRef) {
				return fmt.Errorf("blobs appearing out of order")
			}
			lastRef = sb.Ref
		}
		sawEnd <- true
		return nil

	})
	grp.Go(func() error {
		select {
		case <-sawEnd:
			return nil
		case <-time.After(10 * time.Second):
			return errors.New("timeout waiting for EnumerateBlobs to close its channel")
		}

	})
	if err := grp.Err(); err != nil {
		return fmt.Errorf("Enumerate error: %v", err)
	}
	if len(got) == 0 && len(want) == 0 {
		return nil
	}
	var gotSet = map[blob.SizedRef]bool{}
	for _, sb := range got {
		if gotSet[sb] {
			return fmt.Errorf("duplicate blob %v returned in enumerate", sb)
		}
		gotSet[sb] = true
	}

	if !reflect.DeepEqual(got, want) {
		return fmt.Errorf("Enumerate mismatch. Got %d; want %d.\n Got: %v\nWant: %v\n",
			len(got), len(want), got, want)
	}
	return nil
}

func (r *run) testEnumerate(wantUnsorted []blob.SizedRef, opts ...interface{}) {
	if r.opt.SkipEnum {
		r.t.Log("Skipping enum test")
		return
	}
	if err := r.withRetries(func() error {
		return CheckEnumerate(r.sto, wantUnsorted, opts...)
	}); err != nil {
		r.t.Fatalf("%v", err)
	}
}

func (r *run) withRetries(fn func() error) error {
	delays := r.opt.Retries
	for {
		err := fn()
		if err == nil || len(delays) == 0 {
			return err
		}
		r.t.Logf("(operation failed; retrying after %v)", delays[0])
		time.Sleep(delays[0])
		delays = delays[1:]
	}
}

func testStat(t *testing.T, enum <-chan blob.SizedRef, want []blob.SizedRef) {
	// blobs may arrive in ANY order
	m := make(map[string]int, len(want))
	for i, sb := range want {
		m[sb.Ref.String()] = i
	}

	i := 0
	for sb := range enum {
		if !sb.Valid() {
			break
		}
		wanted := want[m[sb.Ref.String()]]
		if wanted.Size != sb.Size {
			t.Fatalf("received blob size is %d, wanted %d for &%d", sb.Size, wanted.Size, i)
		}
		if wanted.Ref != sb.Ref {
			t.Fatalf("received blob ref mismatch &%d: wanted %s, got %s", i, sb.Ref, wanted.Ref)
		}
		i++
		if i >= len(want) {
			break
		}
	}
}
<|endoftext|>