<commit_before>// Package grpc transparently forwards the grpc protocol using a go-micro client.
package grpc

import (
	"context"
	"io"
	"strings"

	"github.com/micro/go-micro/v3/client"
	"github.com/micro/go-micro/v3/client/grpc"
	"github.com/micro/go-micro/v3/codec"
	"github.com/micro/go-micro/v3/proxy"
	"github.com/micro/go-micro/v3/server"
)

// Proxy will transparently proxy requests to the backend.
// If no backend is specified it will call a service using the client.
// If the service matches the Name it will use the server.DefaultRouter.
type Proxy struct {
	// The proxy options
	options proxy.Options

	// Endpoint specified the fixed endpoint to call.
	Endpoint string

	// The client to use for outbound requests
	Client client.Client
}

// read client request and write to server
func readLoop(r server.Request, s client.Stream) error {
	// request to backend server
	req := s.Request()

	for {
		// get data from client
		//  no need to decode it
		body, err := r.Read()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}

		// get the header from client
		hdr := r.Header()
		msg := &codec.Message{
			Type:   codec.Request,
			Header: hdr,
			Body:   body,
		}
		// write the raw request
		err = req.Codec().Write(msg, nil)
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

// ProcessMessage acts as a message exchange and forwards messages to ongoing topics
// TODO: should we look at p.Endpoint and only send to the local endpoint? probably
func (p *Proxy) ProcessMessage(ctx context.Context, msg server.Message) error {
	// TODO: check that we're not broadcast storming by sending to the same topic
	// that we're actually subscribed to

	// directly publish to the local client
	return p.Client.Publish(ctx, msg)
}

// ServeRequest honours the server.Proxy interface
func (p *Proxy) ServeRequest(ctx context.Context, req server.Request, rsp server.Response) error {
	// set default client
	if p.Client == nil {
		p.Client = grpc.NewClient()
	}

	opts := []client.CallOption{}

	// service name
	service := req.Service()
	endpoint := req.Endpoint()

	// call a specific backend
	if len(p.Endpoint) > 0 {
		// address:port
		if parts := strings.Split(p.Endpoint, ":"); len(parts) > 1 {
			opts = append(opts, client.WithAddress(p.Endpoint))
			// use as service name
		} else {
			service = p.Endpoint
		}
	}

	// create new request with raw bytes body
	creq := p.Client.NewRequest(service, endpoint, nil, client.WithContentType(req.ContentType()))

	// create new stream
	stream, err := p.Client.Stream(ctx, creq, opts...)
	if err != nil {
		return err
	}
	defer stream.Close()

	// create client request read loop
	go readLoop(req, stream)

	// get raw response
	resp := stream.Response()

	// create server response write loop
	for {
		// read backend response body
		body, err := resp.Read()
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}

		// read backend response header
		hdr := resp.Header()

		// write raw response header to client
		rsp.WriteHeader(hdr)

		// write raw response body to client
		err = rsp.Write(body)
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

func (p *Proxy) String() string {
	return "grpc"
}

// NewProxy returns a new grpc proxy server
func NewProxy(opts ...proxy.Option) proxy.Proxy {
	var options proxy.Options
	for _, o := range opts {
		o(&options)
	}

	p := new(Proxy)
	p.Endpoint = options.Endpoint
	p.Client = options.Client

	return p
}

// NewSingleHostProxy returns a router which sends requests to a single backend
func NewSingleHostProxy(url string) *Proxy {
	return &Proxy{
		Endpoint: url,
	}
}
<commit_msg>strip back the grpc proxy<commit_after>// Package grpc is a grpc proxy built for the go-micro/server
package grpc

import (
	"context"
	"io"

	"github.com/micro/go-micro/v3/client"
	grpcc "github.com/micro/go-micro/v3/client/grpc"
	"github.com/micro/go-micro/v3/codec"
	"github.com/micro/go-micro/v3/codec/bytes"
	"github.com/micro/go-micro/v3/errors"
	"github.com/micro/go-micro/v3/logger"
	"github.com/micro/go-micro/v3/proxy"
	"github.com/micro/go-micro/v3/server"
	"google.golang.org/grpc"
)

// Proxy will transparently proxy requests to an endpoint.
// If no endpoint is specified it will call a service using the client.
type Proxy struct {
	// embed options
	options proxy.Options

	// The client to use for outbound requests in the local network
	Client client.Client
}

// read client request and write to server
func readLoop(r server.Request, s client.Stream) error {
	// request to backend server
	req := s.Request()

	for {
		// get data from client
		//  no need to decode it
		body, err := r.Read()
		if err == io.EOF {
			return nil
		}

		if err != nil {
			return err
		}

		// get the header from client
		hdr := r.Header()
		msg := &codec.Message{
			Type:   codec.Request,
			Header: hdr,
			Body:   body,
		}

		// write the raw request
		err = req.Codec().Write(msg, nil)
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

// ProcessMessage acts as a message exchange and forwards messages to ongoing topics
// TODO: should we look at p.Endpoint and only send to the local endpoint? probably
func (p *Proxy) ProcessMessage(ctx context.Context, msg server.Message) error {
	// TODO: check that we're not broadcast storming by sending to the same topic
	// that we're actually subscribed to

	if logger.V(logger.TraceLevel, logger.DefaultLogger) {
		logger.Tracef("Proxy received message for %s", msg.Topic())
	}

	// directly publish to the local client
	return p.Client.Publish(ctx, msg)
}

// ServeRequest honours the server.Router interface
func (p *Proxy) ServeRequest(ctx context.Context, req server.Request, rsp server.Response) error {
	// service name to call
	service := req.Service()
	// endpoint to call
	endpoint := req.Endpoint()

	if len(service) == 0 {
		return errors.BadRequest("go.micro.proxy", "service name is blank")
	}

	if logger.V(logger.TraceLevel, logger.DefaultLogger) {
		logger.Tracef("Proxy received request for %s %s", service, endpoint)
	}

	opts := []client.CallOption{
		client.WithRetries(0),
	}

	// serve the normal way
	return p.serveRequest(ctx, p.Client, service, endpoint, req, rsp, opts...)
}

func (p *Proxy) serveRequest(ctx context.Context, link client.Client, service, endpoint string, req server.Request, rsp server.Response, opts ...client.CallOption) error {
	// read initial request
	body, err := req.Read()
	if err != nil {
		return err
	}

	// create new request with raw bytes body
	creq := link.NewRequest(service, endpoint, &bytes.Frame{Data: body}, client.WithContentType(req.ContentType()))

	// not a stream so make a client.Call request
	if !req.Stream() {
		crsp := new(bytes.Frame)

		// make a call to the backend
		if err := link.Call(ctx, creq, crsp, opts...); err != nil {
			return err
		}

		// write the response
		if err := rsp.Write(crsp.Data); err != nil {
			return err
		}

		return nil
	}

	// new context with cancel
	ctx, cancel := context.WithCancel(ctx)

	// create new stream
	stream, err := link.Stream(ctx, creq, opts...)
	if err != nil {
		return err
	}
	defer stream.Close()

	// with a grpc stream we have to refire the initial request
	// client request to start the server side

	// get the header from client
	msg := &codec.Message{
		Type:   codec.Request,
		Header: req.Header(),
		Body:   body,
	}

	// write the raw request
	err = stream.Request().Codec().Write(msg, nil)
	if err == io.EOF {
		return nil
	} else if err != nil {
		return err
	}

	// create client request read loop if streaming
	go func() {
		err := readLoop(req, stream)
		if err != nil && err != io.EOF {
			// cancel the context
			cancel()
		}
	}()

	// get raw response
	resp := stream.Response()

	// create server response write loop
	for {
		// read backend response body
		body, err := resp.Read()
		if err != nil {
			// when we're done if its a grpc stream we have to set the trailer
			if cc, ok := stream.(grpc.ClientStream); ok {
				if ss, ok := resp.Codec().(grpc.ServerStream); ok {
					ss.SetTrailer(cc.Trailer())
				}
			}
		}

		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}

		// read backend response header
		hdr := resp.Header()

		// write raw response header to client
		rsp.WriteHeader(hdr)

		// write raw response body to client
		err = rsp.Write(body)
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

func (p *Proxy) String() string {
	return "grpc"
}

// NewProxy returns a new proxy which will route based on mucp headers
func NewProxy(opts ...proxy.Option) proxy.Proxy {
	var options proxy.Options

	for _, o := range opts {
		o(&options)
	}

	// create a new grpc proxy
	p := new(Proxy)
	p.options = options

	// set the client
	p.Client = options.Client

	// set the default client
	if p.Client == nil {
		p.Client = grpcc.NewClient()
	}

	return p
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/coreos/go-etcd/etcd"

	"github.com/mozilla-services/pushgo/id"
)

const (
	minTTL = 2 * time.Second
)

var (
	ErrMinTTL       = fmt.Errorf("Default TTL too short; want at least %s", minTTL)
	ErrEtcdStatus   = fmt.Errorf("etcd returned unexpected health check result")
	ErrRouterClosed = fmt.Errorf("Router closed")
)

// IsEtcdKeyExist indicates whether the given error reports that an etcd key
// already exists.
func IsEtcdKeyExist(err error) bool {
	clientErr, ok := err.(*etcd.EtcdError)
	return ok && clientErr.ErrorCode == 105
}

// IsEtcdTemporary indicates whether the given error is a temporary
// etcd error.
func IsEtcdTemporary(err error) bool {
	clientErr, ok := err.(*etcd.EtcdError)
	// Raft (300-class) and internal (400-class) errors are temporary.
	return !ok || clientErr.ErrorCode >= 300 && clientErr.ErrorCode < 500
}

type EtcdLocatorConf struct {
	// Dir is the etcd key prefix for storing contacts. Defaults to
	// "push_hosts".
	Dir string

	// Servers is a list of etcd servers.
	Servers []string

	// DefaultTTL is the maximum amount of time that registered contacts will be
	// considered valid. Defaults to "24h".
	DefaultTTL string `env:"ttl"`

	// RefreshInterval is the maximum amount of time that a cached contact list
	// will be considered valid. Defaults to "5m".
	RefreshInterval string `toml:"refresh_interval" env:"refresh_interval"`

	// MaxRetries is the number of times to retry failed requests. Defaults to 5.
	MaxRetries int `toml:"max_retries" env:"max_retries"`

	// RetryDelay is the amount of time to wait before retrying requests.
	// Defaults to "200ms".
	RetryDelay string `toml:"retry_delay" env:"retry_delay"`

	// MaxJitter is the maximum per-retry randomized delay. Defaults to "400ms".
	MaxJitter string `toml:"max_jitter" env:"max_jitter"`

	// MaxDelay is the maximum amount of time to wait before retrying failed
	// requests with exponential backoff. Defaults to "5s".
	MaxDelay string `toml:"max_delay" env:"max_delay"`
}

// EtcdLocator stores routing endpoints in etcd and polls for new contacts.
type EtcdLocator struct {
	logger          *SimpleLogger
	metrics         *Metrics
	refreshInterval time.Duration
	defaultTTL      time.Duration
	maxRetries      int
	retryDelay      time.Duration
	maxJitter       time.Duration
	maxDelay        time.Duration
	serverList      []string
	dir             string
	url             string
	key             string
	client          *etcd.Client
	contactsLock    sync.RWMutex
	contacts        []string
	contactsErr     error
	lastFetch       time.Time
	isClosing       bool
	closeSignal     chan bool
	closeWait       sync.WaitGroup
	closeLock       sync.Mutex
	lastErr         error
}

func NewEtcdLocator() *EtcdLocator {
	return &EtcdLocator{
		closeSignal: make(chan bool),
	}
}

func (*EtcdLocator) ConfigStruct() interface{} {
	return &EtcdLocatorConf{
		Dir:             "push_hosts",
		Servers:         []string{"http://localhost:4001"},
		DefaultTTL:      "24h",
		RefreshInterval: "5m",
		MaxRetries:      5,
		RetryDelay:      "200ms",
		MaxJitter:       "400ms",
		MaxDelay:        "5s",
	}
}

func (l *EtcdLocator) Init(app *Application, config interface{}) (err error) {
	conf := config.(*EtcdLocatorConf)
	l.logger = app.Logger()
	l.metrics = app.Metrics()

	if l.refreshInterval, err = time.ParseDuration(conf.RefreshInterval); err != nil {
		l.logger.Alert("etcd", "Could not parse refreshInterval",
			LogFields{"error": err.Error(),
				"refreshInterval": conf.RefreshInterval})
		return err
	}
	// default time for the server to be "live"
	if l.defaultTTL, err = time.ParseDuration(conf.DefaultTTL); err != nil {
		l.logger.Alert("etcd",
			"Could not parse etcd default TTL",
			LogFields{"value": conf.DefaultTTL, "error": err.Error()})
		return err
	}
	if l.defaultTTL < minTTL {
		l.logger.Alert("etcd",
			"default TTL too short",
			LogFields{"value": conf.DefaultTTL})
		return ErrMinTTL
	}
	if l.retryDelay, err = time.ParseDuration(conf.RetryDelay); err != nil {
		l.logger.Alert("etcd",
			"Could not parse etcd 'retryDelay'",
			LogFields{"value": conf.RetryDelay, "error": err.Error()})
		return err
	}
	if l.maxJitter, err = time.ParseDuration(conf.MaxJitter); err != nil {
		l.logger.Alert("etcd",
			"Could not parse etcd 'maxJitter'",
			LogFields{"value": conf.MaxJitter, "error": err.Error()})
		return err
	}
	if l.maxDelay, err = time.ParseDuration(conf.MaxDelay); err != nil {
		l.logger.Alert("etcd",
			"Could not parse etcd 'maxDelay'",
			LogFields{"value": conf.MaxDelay, "error": err.Error()})
		return err
	}
	l.maxRetries = conf.MaxRetries

	l.serverList = conf.Servers
	l.dir = path.Clean(conf.Dir)

	// Use the hostname and port of the current server as the etcd key.
	l.url = app.Router().URL()
	uri, err := url.ParseRequestURI(l.url)
	if err != nil {
		l.logger.Alert("etcd", "Error parsing router URL", LogFields{
			"error": err.Error(), "url": l.url})
		return err
	}
	if len(uri.Host) > 0 {
		l.key = path.Join(l.dir, uri.Host)
	}

	if l.logger.ShouldLog(INFO) {
		l.logger.Info("etcd", "connecting to etcd servers",
			LogFields{"list": strings.Join(l.serverList, ";")})
	}
	etcd.SetLogger(log.New(&LogWriter{l.logger, "etcd", DEBUG}, "", 0))
	l.client = etcd.NewClient(l.serverList)
	l.client.CheckRetry = l.checkRetry

	// create the push hosts directory (if not already there)
	if _, err = l.client.CreateDir(l.dir, 0); err != nil {
		if !IsEtcdKeyExist(err) {
			l.logger.Alert("etcd", "etcd createDir error", LogFields{
				"error": err.Error()})
			return err
		}
	}
	if err = l.Register(); err != nil {
		l.logger.Alert("etcd", "Could not register with etcd",
			LogFields{"error": err.Error()})
		return err
	}
	if l.contacts, err = l.getServers(); err != nil {
		l.logger.Alert("etcd", "Could not fetch contact list",
			LogFields{"error": err.Error()})
		return err
	}

	l.closeWait.Add(2)
	go l.registerLoop()
	go l.fetchLoop()
	return nil
}

func (l *EtcdLocator) checkRetry(cluster *etcd.Cluster, retries int, lastResp http.Response, err error) error {
	if l.logger.ShouldLog(ERROR) {
		l.logger.Error("etcd", "etcd request error", LogFields{
			"error":   err.Error(),
			"retries": strconv.Itoa(retries),
			"status":  strconv.Itoa(lastResp.StatusCode)})
	}
	if retries >= l.maxRetries*len(cluster.Machines) {
		l.metrics.Increment("locator.etcd.error")
		return &etcd.EtcdError{
			ErrorCode: etcd.ErrCodeEtcdNotReachable,
			Message:   fmt.Sprintf("Error connecting to etcd after %d retries", retries),
		}
	}
	l.metrics.Increment("locator.etcd.retry.request")
	if lastResp.StatusCode >= 500 {
		retryDelay := time.Duration(int64(l.retryDelay) * (1 << uint(retries-1)))
		if retryDelay > l.maxDelay {
			retryDelay = l.maxDelay
		}
		delay := time.Duration(int64(retryDelay) + rand.Int63n(int64(l.maxJitter)))
		select {
		case <-l.closeSignal:
			return ErrRouterClosed
		case <-time.After(delay):
		}
	}
	return nil
}

// Close stops the locator and closes the etcd client connection. Implements
// Locator.Close().
func (l *EtcdLocator) Close() (err error) {
	defer l.closeLock.Unlock()
	l.closeLock.Lock()
	if l.isClosing {
		return l.lastErr
	}
	close(l.closeSignal)
	l.closeWait.Wait()
	if l.key != "" {
		_, err = l.client.Delete(l.key, false)
	}
	l.isClosing = true
	l.lastErr = err
	return err
}

// Contacts returns a shuffled list of all nodes in the Simple Push cluster.
// Implements Locator.Contacts().
func (l *EtcdLocator) Contacts(string) (contacts []string, err error) {
	l.contactsLock.RLock()
	contacts = make([]string, len(l.contacts))
	copy(contacts, l.contacts)
	if l.contactsErr != nil && time.Since(l.lastFetch) > l.defaultTTL {
		err = l.contactsErr
	}
	l.contactsLock.RUnlock()
	return
}

// Status determines whether etcd can respond to requests. Implements
// Locator.Status().
func (l *EtcdLocator) Status() (ok bool, err error) {
	fakeID, err := id.Generate()
	if err != nil {
		return false, err
	}
	key, expected := "status_"+fakeID, "test"
	if _, err = l.client.Set(key, expected, uint64(6*time.Second)); err != nil {
		return false, err
	}
	resp, err := l.client.Get(key, false, false)
	if err != nil {
		return false, err
	}
	if resp.Node.Value != expected {
		l.logger.Error("etcd", "Unexpected health check result",
			LogFields{"expected": expected, "actual": resp.Node.Value})
		return false, ErrEtcdStatus
	}
	l.client.Delete(key, false)
	return true, nil
}

// Register registers the server to the etcd cluster.
func (l *EtcdLocator) Register() (err error) {
	if l.logger.ShouldLog(INFO) {
		l.logger.Info("etcd", "Registering host", LogFields{
			"key": l.key, "url": l.url})
	}
	retries := 0
	retryDelay := l.retryDelay
	for ok := true; ok && retries < l.maxRetries; retries++ {
		if _, err = l.client.Set(l.key, l.url,
			uint64(l.defaultTTL/time.Second)); err != nil {

			if !IsEtcdTemporary(err) {
				break
			}
			if retryDelay > l.maxDelay {
				retryDelay = l.maxDelay
			}
			delay := time.Duration(int64(retryDelay) + rand.Int63n(int64(l.maxJitter)))
			select {
			case ok = <-l.closeSignal:
			case <-time.After(delay):
				retryDelay *= 2
			}
			continue
		}
		break
	}
	l.metrics.IncrementBy("locator.etcd.retry.register", int64(retries))
	if err != nil {
		if l.logger.ShouldLog(ERROR) {
			l.logger.Error("etcd", "Failed to register", LogFields{
				"error": err.Error(), "key": l.key, "url": l.url})
		}
		return err
	}
	return nil
}

// getServers gets the current contact list from etcd.
func (l *EtcdLocator) getServers() (servers []string, err error) {
	var nodeList *etcd.Response
	retries := 0
	retryDelay := l.retryDelay
	for ok := true; ok && retries < l.maxRetries; retries++ {
		if nodeList, err = l.client.Get(l.dir, false, false); err != nil {
			if !IsEtcdTemporary(err) {
				break
			}
			if retryDelay > l.maxDelay {
				retryDelay = l.maxDelay
			}
			delay := time.Duration(int64(retryDelay) + rand.Int63n(int64(l.maxJitter)))
			select {
			case ok = <-l.closeSignal:
			case <-time.After(delay):
				retryDelay *= 2
			}
			continue
		}
		break
	}
	l.metrics.IncrementBy("locator.etcd.retry.fetch", int64(retries))
	if err != nil {
		if l.logger.ShouldLog(ERROR) {
			l.logger.Error("etcd", "Could not get server list",
				LogFields{"error": err.Error()})
		}
		return nil, err
	}
	servers = make([]string, 0, len(nodeList.Node.Nodes))
	for _, node := range nodeList.Node.Nodes {
		if node.Value == l.url || node.Value == "" {
			continue
		}
		servers = append(servers, node.Value)
	}
	for length := len(servers); length > 0; {
		i := rand.Intn(length)
		length--
		servers[i], servers[length] = servers[length], servers[i]
	}
	return servers, nil
}

// refreshLoop periodically re-registers the current node with etcd.
func (l *EtcdLocator) registerLoop() {
	defer l.closeWait.Done()
	// auto refresh slightly more often than the TTL
	timeout := 0.75 * l.defaultTTL.Seconds()
	ticker := time.NewTicker(time.Duration(timeout) * time.Second)
	for ok := true; ok; {
		select {
		case ok = <-l.closeSignal:
		case <-ticker.C:
			l.Register()
		}
	}
	ticker.Stop()
}

// fetchLoop polls etcd for new nodes.
func (l *EtcdLocator) fetchLoop() {
	defer l.closeWait.Done()
	fetchTick := time.NewTicker(l.refreshInterval)
	for ok := true; ok; {
		select {
		case ok = <-l.closeSignal:
		case t := <-fetchTick.C:
			contacts, err := l.getServers()
			l.contactsLock.Lock()
			l.contacts, l.contactsErr = contacts, err
			l.lastFetch = t
			l.contactsLock.Unlock()
		}
	}
	fetchTick.Stop()
}

func init() {
	rand.Seed(time.Now().UnixNano())
	AvailableLocators["etcd"] = func() HasConfigStruct { return NewEtcdLocator() }
}
<commit_msg>Keep the original contact list for fetch fails.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/coreos/go-etcd/etcd"

	"github.com/mozilla-services/pushgo/id"
)

const (
	minTTL = 2 * time.Second
)

var (
	ErrMinTTL       = fmt.Errorf("Default TTL too short; want at least %s", minTTL)
	ErrEtcdStatus   = fmt.Errorf("etcd returned unexpected health check result")
	ErrRouterClosed = fmt.Errorf("Router closed")
)

// IsEtcdKeyExist indicates whether the given error reports that an etcd key
// already exists.
func IsEtcdKeyExist(err error) bool {
	clientErr, ok := err.(*etcd.EtcdError)
	return ok && clientErr.ErrorCode == 105
}

// IsEtcdTemporary indicates whether the given error is a temporary
// etcd error.
func IsEtcdTemporary(err error) bool {
	clientErr, ok := err.(*etcd.EtcdError)
	// Raft (300-class) and internal (400-class) errors are temporary.
	return !ok || clientErr.ErrorCode >= 300 && clientErr.ErrorCode < 500
}

type EtcdLocatorConf struct {
	// Dir is the etcd key prefix for storing contacts. Defaults to
	// "push_hosts".
	Dir string

	// Servers is a list of etcd servers.
	Servers []string

	// DefaultTTL is the maximum amount of time that registered contacts will be
	// considered valid. Defaults to "24h".
	DefaultTTL string `env:"ttl"`

	// RefreshInterval is the maximum amount of time that a cached contact list
	// will be considered valid. Defaults to "5m".
	RefreshInterval string `toml:"refresh_interval" env:"refresh_interval"`

	// MaxRetries is the number of times to retry failed requests. Defaults to 5.
	MaxRetries int `toml:"max_retries" env:"max_retries"`

	// RetryDelay is the amount of time to wait before retrying requests.
	// Defaults to "200ms".
	RetryDelay string `toml:"retry_delay" env:"retry_delay"`

	// MaxJitter is the maximum per-retry randomized delay. Defaults to "400ms".
	MaxJitter string `toml:"max_jitter" env:"max_jitter"`

	// MaxDelay is the maximum amount of time to wait before retrying failed
	// requests with exponential backoff. Defaults to "5s".
	MaxDelay string `toml:"max_delay" env:"max_delay"`
}

// EtcdLocator stores routing endpoints in etcd and polls for new contacts.
type EtcdLocator struct {
	logger          *SimpleLogger
	metrics         *Metrics
	refreshInterval time.Duration
	defaultTTL      time.Duration
	maxRetries      int
	retryDelay      time.Duration
	maxJitter       time.Duration
	maxDelay        time.Duration
	serverList      []string
	dir             string
	url             string
	key             string
	client          *etcd.Client
	contactsLock    sync.RWMutex
	contacts        []string
	contactsErr     error
	lastFetch       time.Time
	isClosing       bool
	closeSignal     chan bool
	closeWait       sync.WaitGroup
	closeLock       sync.Mutex
	lastErr         error
}

func NewEtcdLocator() *EtcdLocator {
	return &EtcdLocator{
		closeSignal: make(chan bool),
	}
}

func (*EtcdLocator) ConfigStruct() interface{} {
	return &EtcdLocatorConf{
		Dir:             "push_hosts",
		Servers:         []string{"http://localhost:4001"},
		DefaultTTL:      "24h",
		RefreshInterval: "5m",
		MaxRetries:      5,
		RetryDelay:      "200ms",
		MaxJitter:       "400ms",
		MaxDelay:        "5s",
	}
}

func (l *EtcdLocator) Init(app *Application, config interface{}) (err error) {
	conf := config.(*EtcdLocatorConf)
	l.logger = app.Logger()
	l.metrics = app.Metrics()

	if l.refreshInterval, err = time.ParseDuration(conf.RefreshInterval); err != nil {
		l.logger.Alert("etcd", "Could not parse refreshInterval",
			LogFields{"error": err.Error(),
				"refreshInterval": conf.RefreshInterval})
		return err
	}
	// default time for the server to be "live"
	if l.defaultTTL, err = time.ParseDuration(conf.DefaultTTL); err != nil {
		l.logger.Alert("etcd",
			"Could not parse etcd default TTL",
			LogFields{"value": conf.DefaultTTL, "error": err.Error()})
		return err
	}
	if l.defaultTTL < minTTL {
		l.logger.Alert("etcd",
			"default TTL too short",
			LogFields{"value": conf.DefaultTTL})
		return ErrMinTTL
	}
	if l.retryDelay, err = time.ParseDuration(conf.RetryDelay); err != nil {
		l.logger.Alert("etcd",
			"Could not parse etcd 'retryDelay'",
			LogFields{"value": conf.RetryDelay, "error": err.Error()})
		return err
	}
	if l.maxJitter, err = time.ParseDuration(conf.MaxJitter); err != nil {
		l.logger.Alert("etcd",
			"Could not parse etcd 'maxJitter'",
			LogFields{"value": conf.MaxJitter, "error": err.Error()})
		return err
	}
	if l.maxDelay, err = time.ParseDuration(conf.MaxDelay); err != nil {
		l.logger.Alert("etcd",
			"Could not parse etcd 'maxDelay'",
			LogFields{"value": conf.MaxDelay, "error": err.Error()})
		return err
	}
	l.maxRetries = conf.MaxRetries

	l.serverList = conf.Servers
	l.dir = path.Clean(conf.Dir)

	// Use the hostname and port of the current server as the etcd key.
	l.url = app.Router().URL()
	uri, err := url.ParseRequestURI(l.url)
	if err != nil {
		l.logger.Alert("etcd", "Error parsing router URL", LogFields{
			"error": err.Error(), "url": l.url})
		return err
	}
	if len(uri.Host) > 0 {
		l.key = path.Join(l.dir, uri.Host)
	}

	if l.logger.ShouldLog(INFO) {
		l.logger.Info("etcd", "connecting to etcd servers",
			LogFields{"list": strings.Join(l.serverList, ";")})
	}
	etcd.SetLogger(log.New(&LogWriter{l.logger, "etcd", DEBUG}, "", 0))
	l.client = etcd.NewClient(l.serverList)
	l.client.CheckRetry = l.checkRetry

	// create the push hosts directory (if not already there)
	if _, err = l.client.CreateDir(l.dir, 0); err != nil {
		if !IsEtcdKeyExist(err) {
			l.logger.Alert("etcd", "etcd createDir error", LogFields{
				"error": err.Error()})
			return err
		}
	}
	if err = l.Register(); err != nil {
		l.logger.Alert("etcd", "Could not register with etcd",
			LogFields{"error": err.Error()})
		return err
	}
	if l.contacts, err = l.getServers(); err != nil {
		l.logger.Alert("etcd", "Could not fetch contact list",
			LogFields{"error": err.Error()})
		return err
	}

	l.closeWait.Add(2)
	go l.registerLoop()
	go l.fetchLoop()
	return nil
}

func (l *EtcdLocator) checkRetry(cluster *etcd.Cluster, retries int, lastResp http.Response, err error) error {
	if l.logger.ShouldLog(ERROR) {
		l.logger.Error("etcd", "etcd request error", LogFields{
			"error":   err.Error(),
			"retries": strconv.Itoa(retries),
			"status":  strconv.Itoa(lastResp.StatusCode)})
	}
	if retries >= l.maxRetries*len(cluster.Machines) {
		l.metrics.Increment("locator.etcd.error")
		return &etcd.EtcdError{
			ErrorCode: etcd.ErrCodeEtcdNotReachable,
			Message:   fmt.Sprintf("Error connecting to etcd after %d retries", retries),
		}
	}
	l.metrics.Increment("locator.etcd.retry.request")
	if lastResp.StatusCode >= 500 {
		retryDelay := time.Duration(int64(l.retryDelay) * (1 << uint(retries-1)))
		if retryDelay > l.maxDelay {
			retryDelay = l.maxDelay
		}
		delay := time.Duration(int64(retryDelay) + rand.Int63n(int64(l.maxJitter)))
		select {
		case <-l.closeSignal:
			return ErrRouterClosed
		case <-time.After(delay):
		}
	}
	return nil
}

// Close stops the locator and closes the etcd client connection. Implements
// Locator.Close().
func (l *EtcdLocator) Close() (err error) {
	defer l.closeLock.Unlock()
	l.closeLock.Lock()
	if l.isClosing {
		return l.lastErr
	}
	close(l.closeSignal)
	l.closeWait.Wait()
	if l.key != "" {
		_, err = l.client.Delete(l.key, false)
	}
	l.isClosing = true
	l.lastErr = err
	return err
}

// Contacts returns a shuffled list of all nodes in the Simple Push cluster.
// Implements Locator.Contacts().
func (l *EtcdLocator) Contacts(string) (contacts []string, err error) {
	l.contactsLock.RLock()
	contacts = make([]string, len(l.contacts))
	copy(contacts, l.contacts)
	if l.contactsErr != nil && time.Since(l.lastFetch) > l.defaultTTL {
		err = l.contactsErr
	}
	l.contactsLock.RUnlock()
	return
}

// Status determines whether etcd can respond to requests. Implements
// Locator.Status().
func (l *EtcdLocator) Status() (ok bool, err error) {
	fakeID, err := id.Generate()
	if err != nil {
		return false, err
	}
	key, expected := "status_"+fakeID, "test"
	if _, err = l.client.Set(key, expected, uint64(6*time.Second)); err != nil {
		return false, err
	}
	resp, err := l.client.Get(key, false, false)
	if err != nil {
		return false, err
	}
	if resp.Node.Value != expected {
		l.logger.Error("etcd", "Unexpected health check result",
			LogFields{"expected": expected, "actual": resp.Node.Value})
		return false, ErrEtcdStatus
	}
	l.client.Delete(key, false)
	return true, nil
}

// Register registers the server to the etcd cluster.
func (l *EtcdLocator) Register() (err error) {
	if l.logger.ShouldLog(INFO) {
		l.logger.Info("etcd", "Registering host", LogFields{
			"key": l.key, "url": l.url})
	}
	retries := 0
	retryDelay := l.retryDelay
	for ok := true; ok && retries < l.maxRetries; retries++ {
		if _, err = l.client.Set(l.key, l.url,
			uint64(l.defaultTTL/time.Second)); err != nil {

			if !IsEtcdTemporary(err) {
				break
			}
			if retryDelay > l.maxDelay {
				retryDelay = l.maxDelay
			}
			delay := time.Duration(int64(retryDelay) + rand.Int63n(int64(l.maxJitter)))
			select {
			case ok = <-l.closeSignal:
			case <-time.After(delay):
				retryDelay *= 2
			}
			continue
		}
		break
	}
	l.metrics.IncrementBy("locator.etcd.retry.register", int64(retries))
	if err != nil {
		if l.logger.ShouldLog(ERROR) {
			l.logger.Error("etcd", "Failed to register", LogFields{
				"error": err.Error(), "key": l.key, "url": l.url})
		}
		return err
	}
	return nil
}

// getServers gets the current contact list from etcd.
func (l *EtcdLocator) getServers() (servers []string, err error) {
	var nodeList *etcd.Response
	retries := 0
	retryDelay := l.retryDelay
	for ok := true; ok && retries < l.maxRetries; retries++ {
		if nodeList, err = l.client.Get(l.dir, false, false); err != nil {
			if !IsEtcdTemporary(err) {
				break
			}
			if retryDelay > l.maxDelay {
				retryDelay = l.maxDelay
			}
			delay := time.Duration(int64(retryDelay) + rand.Int63n(int64(l.maxJitter)))
			select {
			case ok = <-l.closeSignal:
			case <-time.After(delay):
				retryDelay *= 2
			}
			continue
		}
		break
	}
	l.metrics.IncrementBy("locator.etcd.retry.fetch", int64(retries))
	if err != nil {
		if l.logger.ShouldLog(ERROR) {
			l.logger.Error("etcd", "Could not get server list",
				LogFields{"error": err.Error()})
		}
		return nil, err
	}
	servers = make([]string, 0, len(nodeList.Node.Nodes))
	for _, node := range nodeList.Node.Nodes {
		if node.Value == l.url || node.Value == "" {
			continue
		}
		servers = append(servers, node.Value)
	}
	for length := len(servers); length > 0; {
		i := rand.Intn(length)
		length--
		servers[i], servers[length] = servers[length], servers[i]
	}
	return servers, nil
}

// refreshLoop periodically re-registers the current node with etcd.
func (l *EtcdLocator) registerLoop() {
	defer l.closeWait.Done()
	// auto refresh slightly more often than the TTL
	timeout := 0.75 * l.defaultTTL.Seconds()
	ticker := time.NewTicker(time.Duration(timeout) * time.Second)
	for ok := true; ok; {
		select {
		case ok = <-l.closeSignal:
		case <-ticker.C:
			l.Register()
		}
	}
	ticker.Stop()
}

// fetchLoop polls etcd for new nodes.
func (l *EtcdLocator) fetchLoop() {
	defer l.closeWait.Done()
	fetchTick := time.NewTicker(l.refreshInterval)
	for ok := true; ok; {
		select {
		case ok = <-l.closeSignal:
		case t := <-fetchTick.C:
			contacts, err := l.getServers()
			l.contactsLock.Lock()
			if err != nil {
				l.contactsErr = err
			} else {
				l.contacts = contacts
				l.contactsErr = nil
			}
			l.lastFetch = t
			l.contactsLock.Unlock()
		}
	}
	fetchTick.Stop()
}

func init() {
	rand.Seed(time.Now().UnixNano())
	AvailableLocators["etcd"] = func() HasConfigStruct { return NewEtcdLocator() }
}
<|endoftext|>


<commit_before>package eth

import (
	"bytes"
	"io"
	"log"
	"math/big"
	"os"
	"testing"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/errs"
	ethlogger "github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/p2p"
	"github.com/ethereum/go-ethereum/p2p/discover"
)

var logsys = ethlogger.NewStdLogSystem(os.Stdout, log.LstdFlags, ethlogger.LogLevel(ethlogger.DebugDetailLevel))

var ini = false

func logInit() {
	if !ini {
		ethlogger.AddLogSystem(logsys)
		ini = true
	}
}

type testMsgReadWriter struct {
	in  chan p2p.Msg
	out []p2p.Msg
}

func (self *testMsgReadWriter) In(msg p2p.Msg) {
	self.in <- msg
}

func (self *testMsgReadWriter) Out() (msg p2p.Msg, ok bool) {
	if len(self.out) > 0 {
		msg = self.out[0]
		self.out = self.out[1:]
		ok = true
	}
	return
}

func (self *testMsgReadWriter) WriteMsg(msg p2p.Msg) error {
	self.out = append(self.out, msg)
	return nil
}

func (self *testMsgReadWriter) ReadMsg() (p2p.Msg, error) {
	msg, ok := <-self.in
	if !ok {
		return msg, io.EOF
	}
	return msg, nil
}

type testTxPool struct {
	getTransactions func() []*types.Transaction
	addTransactions func(txs []*types.Transaction)
}

type testChainManager struct {
	getBlockHashes func(hash []byte, amount uint64) (hashes [][]byte)
	getBlock       func(hash []byte) *types.Block
	status         func() (td *big.Int, currentBlock []byte, genesisBlock []byte)
}

type testBlockPool struct {
	addBlockHashes func(next func() ([]byte, bool), peerId string)
	addBlock       func(block *types.Block, peerId string) (err error)
	addPeer        func(td *big.Int, currentBlock []byte, peerId string, requestHashes func([]byte) error, requestBlocks func([][]byte) error, peerError func(*errs.Error)) (best bool)
	removePeer     func(peerId string)
}

// func (self *testTxPool) GetTransactions() (txs []*types.Transaction) {
// 	if self.getTransactions != nil {
// 		txs = self.getTransactions()
// 	}
// 	return
// }

func (self *testTxPool) AddTransactions(txs []*types.Transaction) {
	if self.addTransactions != nil {
		self.addTransactions(txs)
	}
}

func (self *testTxPool) GetTransactions() types.Transactions { return nil }

func (self *testChainManager) GetBlockHashesFromHash(hash []byte, amount uint64) (hashes [][]byte) {
	if self.getBlockHashes != nil {
		hashes = self.getBlockHashes(hash, amount)
	}
	return
}

func (self *testChainManager) Status() (td *big.Int, currentBlock []byte, genesisBlock []byte) {
	if self.status != nil {
		td, currentBlock, genesisBlock = self.status()
	}
	return
}

func (self *testChainManager) GetBlock(hash []byte) (block *types.Block) {
	if self.getBlock != nil {
		block = self.getBlock(hash)
	}
	return
}

func (self *testBlockPool) AddBlockHashes(next func() ([]byte, bool), peerId string) {
	if self.addBlockHashes != nil {
		self.addBlockHashes(next, peerId)
	}
}

func (self *testBlockPool) AddBlock(block *types.Block, peerId string) {
	if self.addBlock != nil {
		self.addBlock(block, peerId)
	}
}

func (self *testBlockPool) AddPeer(td *big.Int, currentBlock []byte, peerId string, requestBlockHashes func([]byte) error, requestBlocks func([][]byte) error, peerError func(*errs.Error)) (best bool) {
	if self.addPeer != nil {
		best = self.addPeer(td, currentBlock, peerId, requestBlockHashes, requestBlocks, peerError)
	}
	return
}

func (self *testBlockPool) RemovePeer(peerId string) {
	if self.removePeer != nil {
		self.removePeer(peerId)
	}
}

func testPeer() *p2p.Peer {
	var id discover.NodeID
	pk := crypto.GenerateNewKeyPair().PublicKey
	copy(id[:], pk)
	return p2p.NewPeer(id, "test peer", []p2p.Cap{})
}

type ethProtocolTester struct {
	quit         chan error
	rw           *testMsgReadWriter // p2p.MsgReadWriter
	txPool       *testTxPool        // txPool
	chainManager *testChainManager  // chainManager
	blockPool    *testBlockPool     // blockPool
	t            *testing.T
}

func newEth(t *testing.T) *ethProtocolTester {
	return &ethProtocolTester{
		quit:         make(chan error),
		rw:           &testMsgReadWriter{in: make(chan p2p.Msg, 10)},
		txPool:       &testTxPool{},
		chainManager: &testChainManager{},
		blockPool:    &testBlockPool{},
		t:            t,
	}
}

func (self *ethProtocolTester) reset() {
	self.rw = &testMsgReadWriter{in: make(chan p2p.Msg, 10)}
	self.quit = make(chan error)
}

func (self *ethProtocolTester) checkError(expCode int, delay time.Duration) (err error) {
	var timer = time.After(delay)
	select {
	case err = <-self.quit:
	case <-timer:
		self.t.Errorf("no error after %v, expected %v", delay, expCode)
		return
	}
	perr, ok := err.(*errs.Error)
	if ok && perr != nil {
		if code := perr.Code; code != expCode {
			self.t.Errorf("expected protocol error (code %v), got %v (%v)", expCode, code, err)
		}
	} else {
		self.t.Errorf("expected protocol error (code %v), got %v", expCode, err)
	}
	return
}

func (self *ethProtocolTester) In(msg p2p.Msg) {
	self.rw.In(msg)
}

func (self *ethProtocolTester) Out() (p2p.Msg, bool) {
	return self.rw.Out()
}

func (self *ethProtocolTester) checkMsg(i int, code uint64, val interface{}) (msg p2p.Msg) {
	if i >= len(self.rw.out) {
		self.t.Errorf("expected at least %v msgs, got %v", i, len(self.rw.out))
		return
	}
	msg = self.rw.out[i]
	if msg.Code != code {
		self.t.Errorf("expected msg code %v, got %v", code, msg.Code)
	}
	if val != nil {
		if err := msg.Decode(val); err != nil {
			self.t.Errorf("rlp encoding error: %v", err)
		}
	}
	return
}

func (self *ethProtocolTester) run() {
	err := runEthProtocol(self.txPool, self.chainManager, self.blockPool, testPeer(), self.rw)
	self.quit <- err
}

func TestStatusMsgErrors(t *testing.T) {
	logInit()
	eth := newEth(t)
	td := common.Big1
	currentBlock := []byte{1}
	genesis := []byte{2}
	eth.chainManager.status = func() (*big.Int, []byte, []byte) { return td, currentBlock, genesis }
	go eth.run()
	statusMsg := p2p.NewMsg(4)
	eth.In(statusMsg)
	delay := 1 * time.Second
	eth.checkError(ErrNoStatusMsg, delay)
	var status statusMsgData
	eth.checkMsg(0, StatusMsg, &status) // first outgoing msg should be StatusMsg
	if status.TD.Cmp(td) != 0 ||
		status.ProtocolVersion != eth.ProtocolVersion ||
		status.NetworkId != eth.NetworkId ||
		status.TD.Cmp(td) != 0 ||
		bytes.Compare(status.CurrentBlock, currentBlock) != 0 ||
		bytes.Compare(status.GenesisBlock, genesis) != 0 {
		t.Errorf("incorrect outgoing status")
	}

	eth.reset()
	go eth.run()
	statusMsg = p2p.NewMsg(0, uint32(48), uint32(0), td, currentBlock, genesis)
	eth.In(statusMsg)
	eth.checkError(ErrProtocolVersionMismatch, delay)

	eth.reset()
	go eth.run()
	statusMsg = p2p.NewMsg(0, uint32(49), uint32(1), td, currentBlock, genesis)
	eth.In(statusMsg)
	eth.checkError(ErrNetworkIdMismatch, delay)

	eth.reset()
	go eth.run()
	statusMsg = p2p.NewMsg(0, uint32(49), uint32(0), td, currentBlock, []byte{3})
	eth.In(statusMsg)
	eth.checkError(ErrGenesisBlockMismatch, delay)

}
<commit_msg>fix eth tests<commit_after>package eth

import (
	"bytes"
	"io"
	"log"
	"math/big"
	"os"
	"testing"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/errs"
	ethlogger "github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/p2p"
	"github.com/ethereum/go-ethereum/p2p/discover"
)

var logsys = ethlogger.NewStdLogSystem(os.Stdout, log.LstdFlags, ethlogger.LogLevel(ethlogger.DebugDetailLevel))

var ini = false

func logInit() {
	if !ini {
		ethlogger.AddLogSystem(logsys)
		ini = true
	}
}

type testMsgReadWriter struct {
	in  chan p2p.Msg
	out []p2p.Msg
}

func (self *testMsgReadWriter) In(msg p2p.Msg) {
	self.in <- msg
}

func (self *testMsgReadWriter) Out() (msg p2p.Msg, ok bool) {
	if len(self.out) > 0 {
		msg = self.out[0]
		self.out = self.out[1:]
		ok = true
	}
	return
}

func (self *testMsgReadWriter) WriteMsg(msg p2p.Msg) error {
	self.out = append(self.out, msg)
	return nil
}

func (self *testMsgReadWriter) ReadMsg() (p2p.Msg, error) {
	msg, ok := <-self.in
	if !ok {
		return msg, io.EOF
	}
	return msg, nil
}

type testTxPool struct {
	getTransactions func() []*types.Transaction
	addTransactions func(txs []*types.Transaction)
}

type testChainManager struct {
	getBlockHashes func(hash []byte, amount uint64) (hashes [][]byte)
	getBlock       func(hash []byte) *types.Block
	status         func() (td *big.Int, currentBlock []byte, genesisBlock []byte)
}

type testBlockPool struct {
	addBlockHashes func(next func() ([]byte, bool), peerId string)
	addBlock       func(block *types.Block, peerId string) (err error)
	addPeer        func(td *big.Int, currentBlock []byte, peerId string, requestHashes func([]byte) error, requestBlocks func([][]byte) error, peerError func(*errs.Error)) (best bool)
	removePeer     func(peerId string)
}

// func (self *testTxPool) GetTransactions() (txs []*types.Transaction) {
// 	if self.getTransactions != nil {
// 		txs = self.getTransactions()
// 	}
// 	return
// }

func (self *testTxPool) AddTransactions(txs []*types.Transaction) {
	if self.addTransactions != nil {
		self.addTransactions(txs)
	}
}

func (self *testTxPool) GetTransactions() types.Transactions { return nil }

func (self *testChainManager) GetBlockHashesFromHash(hash []byte, amount uint64) (hashes [][]byte) {
	if self.getBlockHashes != nil {
		hashes = self.getBlockHashes(hash, amount)
	}
	return
}

func (self *testChainManager) Status() (td *big.Int, currentBlock []byte, genesisBlock []byte) {
	if self.status != nil {
		td, currentBlock, genesisBlock = self.status()
	}
	return
}

func (self *testChainManager) GetBlock(hash []byte) (block *types.Block) {
	if self.getBlock != nil {
		block = self.getBlock(hash)
	}
	return
}

func (self *testBlockPool) AddBlockHashes(next func() ([]byte, bool), peerId string) {
	if self.addBlockHashes != nil {
		self.addBlockHashes(next, peerId)
	}
}

func (self *testBlockPool) AddBlock(block *types.Block, peerId string) {
	if self.addBlock != nil {
		self.addBlock(block, peerId)
	}
}

func (self *testBlockPool) AddPeer(td *big.Int, currentBlock []byte, peerId string, requestBlockHashes func([]byte) error, requestBlocks func([][]byte) error, peerError func(*errs.Error)) (best bool) {
	if self.addPeer != nil {
		best = self.addPeer(td, currentBlock, peerId, requestBlockHashes, requestBlocks, peerError)
	}
	return
}

func (self *testBlockPool) RemovePeer(peerId string) {
	if self.removePeer != nil {
		self.removePeer(peerId)
	}
}

func testPeer() *p2p.Peer {
	var id discover.NodeID
	pk := crypto.GenerateNewKeyPair().PublicKey
	copy(id[:], pk)
	return p2p.NewPeer(id, "test peer", []p2p.Cap{})
}

type ethProtocolTester struct {
	quit         chan error
	rw           *testMsgReadWriter // p2p.MsgReadWriter
	txPool       *testTxPool        // txPool
	chainManager *testChainManager  // chainManager
	blockPool    *testBlockPool     // blockPool
	t            *testing.T
}

func newEth(t *testing.T) *ethProtocolTester {
	return &ethProtocolTester{
		quit:         make(chan error),
		rw:           &testMsgReadWriter{in: make(chan p2p.Msg, 10)},
		txPool:       &testTxPool{},
		chainManager: &testChainManager{},
		blockPool:    &testBlockPool{},
		t:            t,
	}
}

func (self *ethProtocolTester) reset() {
	self.rw = &testMsgReadWriter{in: make(chan p2p.Msg, 10)}
	self.quit = make(chan error)
}

func (self *ethProtocolTester) checkError(expCode int, delay time.Duration) (err error) {
	var timer = time.After(delay)
	select {
	case err = <-self.quit:
	case <-timer:
		self.t.Errorf("no error after %v, expected %v", delay, expCode)
		return
	}
	perr, ok := err.(*errs.Error)
	if ok && perr != nil {
		if code := perr.Code; code != expCode {
			self.t.Errorf("expected protocol error (code %v), got %v (%v)", expCode, code, err)
		}
	} else {
		self.t.Errorf("expected protocol error (code %v), got %v", expCode, err)
	}
	return
}

func (self *ethProtocolTester) In(msg p2p.Msg) {
	self.rw.In(msg)
}

func (self *ethProtocolTester) Out() (p2p.Msg, bool) {
	return self.rw.Out()
}

func (self *ethProtocolTester) checkMsg(i int, code uint64, val interface{}) (msg p2p.Msg) {
	if i >= len(self.rw.out) {
		self.t.Errorf("expected at least %v msgs, got %v", i, len(self.rw.out))
		return
	}
	msg = self.rw.out[i]
	if msg.Code != code {
		self.t.Errorf("expected msg code %v, got %v", code, msg.Code)
	}
	if val != nil {
		if err := msg.Decode(val); err != nil {
			self.t.Errorf("rlp encoding error: %v", err)
		}
	}
	return
}

func (self *ethProtocolTester) run() {
	err := runEthProtocol(ProtocolVersion, NetworkId, self.txPool, self.chainManager, self.blockPool, testPeer(), self.rw)
	self.quit <- err
}

func TestStatusMsgErrors(t *testing.T) {
	logInit()
	eth := newEth(t)
	td := common.Big1
	currentBlock := []byte{1}
	genesis := []byte{2}
	eth.chainManager.status = func() (*big.Int, []byte, []byte) { return td, currentBlock, genesis }
	go eth.run()
	statusMsg := p2p.NewMsg(4)
	eth.In(statusMsg)
	delay := 1 * time.Second
	eth.checkError(ErrNoStatusMsg, delay)
	var status statusMsgData
	eth.checkMsg(0, StatusMsg, &status) // first outgoing msg should be StatusMsg
	if status.TD.Cmp(td) != 0 ||
		status.ProtocolVersion != ProtocolVersion ||
		status.NetworkId != NetworkId ||
		status.TD.Cmp(td) != 0 ||
		bytes.Compare(status.CurrentBlock, currentBlock) != 0 ||
		bytes.Compare(status.GenesisBlock, genesis) != 0 {
		t.Errorf("incorrect outgoing status")
	}

	eth.reset()
	go eth.run()
	statusMsg = p2p.NewMsg(0, uint32(48), uint32(0), td, currentBlock, genesis)
	eth.In(statusMsg)
	eth.checkError(ErrProtocolVersionMismatch, delay)

	eth.reset()
	go eth.run()
	statusMsg = p2p.NewMsg(0, uint32(49), uint32(1), td, currentBlock, genesis)
	eth.In(statusMsg)
	eth.checkError(ErrNetworkIdMismatch, delay)

	eth.reset()
	go eth.run()
	statusMsg = p2p.NewMsg(0, uint32(49), uint32(0), td, currentBlock, []byte{3})
	eth.In(statusMsg)
	eth.checkError(ErrGenesisBlockMismatch, delay)

}
<|endoftext|>


<commit_before>package unit

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"syscall"

	systemdDbus "github.com/coreos/go-systemd/dbus"
	"github.com/guelfey/go.dbus"
	log "github.com/golang/glog"

	"github.com/coreos/coreinit/job"
	"github.com/coreos/coreinit/machine"
)

const (
	defaultSystemdRuntimePath = "/run/systemd/system/"
	defaultSystemdDbusPath    = "/org/freedesktop/systemd1/unit/"
)

type SystemdManager struct {
	Systemd  *systemdDbus.Conn
	Target   *SystemdTarget
	Machine  *machine.Machine
	unitPath string
	dbusPath string
}

func NewSystemdManager(machine *machine.Machine) *SystemdManager {
	systemd := systemdDbus.New()

	name := "coreinit-" + machine.BootId + ".target"
	target := NewSystemdTarget(name)

	mgr := &SystemdManager{systemd, target, machine, defaultSystemdRuntimePath, defaultSystemdDbusPath}

	mgr.writeUnit(target.Name(), "")

	return mgr
}

func (m *SystemdManager) getUnitByName(name string) (*SystemdUnit, error) {
	var unit SystemdUnit
	if strings.HasSuffix(name, ".service") {
		unit = NewSystemdService(m, name)
	} else if strings.HasSuffix(name, ".socket") {
		unit = NewSystemdSocket(m, name)
	} else {
		panic("WAT")
	}

	return &unit, nil
}

func (m *SystemdManager) getUnitsByTarget(target *SystemdTarget) []SystemdUnit {
	object := m.getDbusPath(target.Name())
	info, err := m.Systemd.GetUnitInfo(object)

	if err != nil {
		panic(err)
	}

	names := info["Wants"].Value().([]string)

	var units []SystemdUnit
	for _, name := range names {
		unit, err := m.getUnitByName(name)
		if err == nil {
			units = append(units, *unit)
		} else {
			log.V(1).Infof("Unit %s seems to exist, yet unable to get corresponding SystemdUnit object", name)
		}
	}

	return units
}

func (m *SystemdManager) GetJobs() map[string]job.Job {
	units := m.getUnitsByTarget(m.Target)
	jobs := make(map[string]job.Job, len(units))
	for _, u := range units {
		state := m.getJobStateFromUnit(&u)
		j, _ := job.NewJob(u.Name(), state, nil)
		jobs[j.Name] = *j
	}

	return jobs
}

func (m *SystemdManager) getJobStateFromUnit(u *SystemdUnit) *job.JobState {
	loadState, activeState, subState, sockets, err := (*u).State()
	if err != nil {
		log.V(1).Infof("Failed to get state for unit %s", (*u).Name())
		return nil
	} else {
		return job.NewJobState(loadState, activeState, subState, sockets, m.Machine)
	}
}

func (m *SystemdManager) GetJobState(j *job.Job) *job.JobState {
	unit, err := m.getUnitByName(j.Name)
	if err != nil {
		log.V(1).Infof("No local unit corresponding to job %s", j.Name)
		return nil
	}

	return m.getJobStateFromUnit(unit)
}

func (m *SystemdManager) StartJob(job *job.Job) {
	//This is probably not the right place to force the service to be
	// WantedBy our systemd target
	job.Payload.Value += "\r\n\r\n[Install]\r\nWantedBy=" + m.Target.Name()

	m.writeUnit(job.Name, job.Payload.Value)
	m.startUnit(job.Name)
}

func (m *SystemdManager) StopJob(job *job.Job) {
	m.stopUnit(job.Name)
	m.removeUnit(job.Name)
}

func (m *SystemdManager) getUnitStates(name string) (string, string, string, error) {
	info, err := m.Systemd.GetUnitInfo(m.getDbusPath(name))

	if err != nil {
		return "", "", "", err
	} else {
		loadState := info["LoadState"].Value().(string)
		activeState := info["ActiveState"].Value().(string)
		subState := info["SubState"].Value().(string)
		return loadState, activeState, subState, nil
	}
}

func (m *SystemdManager) startUnit(name string) {
	log.Infof("Starting systemd unit %s", name)

	files := []string{name}
	m.Systemd.EnableUnitFiles(files, true, false)

	m.Systemd.StartUnit(name, "replace")
}

func (m *SystemdManager) stopUnit(name string) {
	log.Infof("Stopping systemd unit %s", name)

	m.Systemd.StopUnit(name, "replace")

	// go-systemd does not yet have this implemented
	//files := []string{name}
	//Systemd.DisableUnitFiles(files, true, false)
}

func (m *SystemdManager) removeUnit(name string) {
	log.Infof("Unlinking systemd unit %s from target %s", name, m.Target.Name())
	link := m.getLocalPath(path.Join(m.Target.Name()+".wants", name))
	syscall.Unlink(link)

	file := m.getLocalPath(name)
	log.Infof("Removing systemd unit file %s", file)
	syscall.Unlink(file)
}

func (m *SystemdManager) readUnit(name string) (string, error) {
	path := m.getLocalPath(name)
	contents, err := ioutil.ReadFile(path)
	if err == nil {
		return string(contents), nil
	} else {
		return "", errors.New(fmt.Sprintf("No unit file at local path %s", path))
	}
}

func (m *SystemdManager) writeUnit(name string, contents string) error {
	log.Infof("Writing systemd unit file %s", name)

	path := path.Join(m.unitPath, name)
	file, err := os.Create(path)
	defer file.Close()

	if err != nil {
		return err
	}

	file.Write([]byte(contents))
	return nil
}

func (m *SystemdManager) getDbusPath(name string) dbus.ObjectPath {
	path := path.Join(m.dbusPath, name)
	path = serializeDbusPath(path)
	return dbus.ObjectPath(path)
}

func (m *SystemdManager) getLocalPath(name string) string {
	return path.Join(m.unitPath, name)
}
<commit_msg>refacotr(unit): Prefix unit files with machine name<commit_after>package unit

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"syscall"

	systemdDbus "github.com/coreos/go-systemd/dbus"
	"github.com/guelfey/go.dbus"
	log "github.com/golang/glog"

	"github.com/coreos/coreinit/job"
	"github.com/coreos/coreinit/machine"
)

const (
	defaultSystemdRuntimePath = "/run/systemd/system/"
	defaultSystemdDbusPath    = "/org/freedesktop/systemd1/unit/"
)

type SystemdManager struct {
	Systemd  *systemdDbus.Conn
	Target   *SystemdTarget
	Machine  *machine.Machine
	unitPath string
	dbusPath string
}

func NewSystemdManager(machine *machine.Machine) *SystemdManager {
	systemd := systemdDbus.New()

	name := "coreinit-" + machine.BootId + ".target"
	target := NewSystemdTarget(name)

	mgr := &SystemdManager{systemd, target, machine, defaultSystemdRuntimePath, defaultSystemdDbusPath}

	mgr.writeUnit(target.Name(), "")

	return mgr
}

func (m *SystemdManager) getUnitByName(name string) (*SystemdUnit, error) {
	var unit SystemdUnit
	if strings.HasSuffix(name, ".service") {
		unit = NewSystemdService(m, name)
	} else if strings.HasSuffix(name, ".socket") {
		unit = NewSystemdSocket(m, name)
	} else {
		panic("WAT")
	}

	return &unit, nil
}

func (m *SystemdManager) getUnitsByTarget(target *SystemdTarget) []SystemdUnit {
	object := m.getDbusPath(target.Name())
	info, err := m.Systemd.GetUnitInfo(object)

	if err != nil {
		panic(err)
	}

	names := info["Wants"].Value().([]string)

	var units []SystemdUnit
	for _, name := range names {
		unit, err := m.getUnitByName(name)
		if err == nil {
			units = append(units, *unit)
		} else {
			log.V(1).Infof("Unit %s seems to exist, yet unable to get corresponding SystemdUnit object", name)
		}
	}

	return units
}

func (m *SystemdManager) GetJobs() map[string]job.Job {
	units := m.getUnitsByTarget(m.Target)
	jobs := make(map[string]job.Job, len(units))
	for _, u := range units {
		state := m.getJobStateFromUnit(&u)
		name := m.stripUnitNamePrefix(u.Name())
		j, _ := job.NewJob(name, state, nil)
		jobs[j.Name] = *j
	}

	return jobs
}

func (m *SystemdManager) getJobStateFromUnit(u *SystemdUnit) *job.JobState {
	loadState, activeState, subState, sockets, err := (*u).State()
	if err != nil {
		log.V(1).Infof("Failed to get state for unit %s", (*u).Name())
		return nil
	} else {
		return job.NewJobState(loadState, activeState, subState, sockets, m.Machine)
	}
}

func (m *SystemdManager) GetJobState(j *job.Job) *job.JobState {
	name := m.addUnitNamePrefix(j.Name)
	unit, err := m.getUnitByName(name)
	if err != nil {
		log.V(1).Infof("No local unit corresponding to job %s", j.Name)
		return nil
	}

	return m.getJobStateFromUnit(unit)
}

func (m *SystemdManager) StartJob(job *job.Job) {
	//This is probably not the right place to force the service to be
	// WantedBy our systemd target
	job.Payload.Value += "\r\n\r\n[Install]\r\nWantedBy=" + m.Target.Name()

	name := m.addUnitNamePrefix(job.Name)
	m.writeUnit(name, job.Payload.Value)
	m.startUnit(name)
}

func (m *SystemdManager) StopJob(job *job.Job) {
	name := m.addUnitNamePrefix(job.Name)
	m.stopUnit(name)
	m.removeUnit(name)
}

func (m *SystemdManager) getUnitStates(name string) (string, string, string, error) {
	info, err := m.Systemd.GetUnitInfo(m.getDbusPath(name))

	if err != nil {
		return "", "", "", err
	} else {
		loadState := info["LoadState"].Value().(string)
		activeState := info["ActiveState"].Value().(string)
		subState := info["SubState"].Value().(string)
		return loadState, activeState, subState, nil
	}
}

func (m *SystemdManager) startUnit(name string) {
	log.Infof("Starting systemd unit %s", name)

	files := []string{name}
	m.Systemd.EnableUnitFiles(files, true, false)

	m.Systemd.StartUnit(name, "replace")
}

func (m *SystemdManager) stopUnit(name string) {
	log.Infof("Stopping systemd unit %s", name)

	m.Systemd.StopUnit(name, "replace")

	// go-systemd does not yet have this implemented
	//files := []string{name}
	//Systemd.DisableUnitFiles(files, true, false)
}

func (m *SystemdManager) removeUnit(name string) {
	log.Infof("Unlinking systemd unit %s from target %s", name, m.Target.Name())
	link := m.getLocalPath(path.Join(m.Target.Name()+".wants", name))
	syscall.Unlink(link)

	file := m.getLocalPath(name)
	log.Infof("Removing systemd unit file %s", file)
	syscall.Unlink(file)
}

func (m *SystemdManager) readUnit(name string) (string, error) {
	path := m.getLocalPath(name)
	contents, err := ioutil.ReadFile(path)
	if err == nil {
		return string(contents), nil
	} else {
		return "", errors.New(fmt.Sprintf("No unit file at local path %s", path))
	}
}

func (m *SystemdManager) writeUnit(name string, contents string) error {
	log.Infof("Writing systemd unit file %s", name)

	path := path.Join(m.unitPath, name)
	file, err := os.Create(path)
	defer file.Close()

	if err != nil {
		return err
	}

	file.Write([]byte(contents))
	return nil
}

func (m *SystemdManager) getDbusPath(name string) dbus.ObjectPath {
	path := path.Join(m.dbusPath, name)
	path = serializeDbusPath(path)
	return dbus.ObjectPath(path)
}

func (m *SystemdManager) getLocalPath(name string) string {
	return path.Join(m.unitPath, name)
}

func (m *SystemdManager) addUnitNamePrefix(name string) string {
	return fmt.Sprintf("%s.%s", m.Machine.BootId, name)
}

func (m *SystemdManager) stripUnitNamePrefix(name string) string {
	return strings.TrimPrefix(name, fmt.Sprintf("%s.", m.Machine.BootId))
}
<|endoftext|>


<commit_before>package docker

import (
	"errors"
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"
	"time"

	authProvider "github.com/ViBiOh/auth/provider"
	"github.com/ViBiOh/httputils/request"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
)

func TestGetConfig(t *testing.T) {
	var cases = []struct {
		service *dockerComposeService
		user    *authProvider.User
		appName string
		want    *container.Config
		wantErr error
	}{
		{
			&dockerComposeService{
				Healthcheck: &dockerComposeHealthcheck{
					Interval: `abcd`,
				},
			},
			authProvider.NewUser(0, `admin`, `admin`),
			`test`,
			nil,
			errors.New(`Error while parsing healthcheck interval: time: invalid duration abcd`),
		},
		{
			&dockerComposeService{
				Healthcheck: &dockerComposeHealthcheck{
					Interval: `30s`,
					Timeout:  `abcd`,
				},
			},
			authProvider.NewUser(0, `admin`, `admin`),
			`test`,
			nil,
			errors.New(`Error while parsing healthcheck timeout: time: invalid duration abcd`),
		},
		{
			&dockerComposeService{},
			authProvider.NewUser(0, `admin`, `admin`),
			`test`,
			&container.Config{
				Labels: map[string]string{`owner`: `admin`, `app`: `test`},
				Env:    []string{},
			},
			nil,
		},
		{
			&dockerComposeService{
				Image:       `vibioh/dashboard`,
				Environment: map[string]string{`PATH`: `/usr/bin`},
				Labels:      map[string]string{`CUSTOM_LABEL`: `testing`},
				Command:     []string{`entrypoint.sh`, `start`},
				Healthcheck: &dockerComposeHealthcheck{
					Test:     []string{`CMD`, `alcotest`},
					Retries:  10,
					Interval: `30s`,
					Timeout:  `10s`,
				},
			},
			authProvider.NewUser(0, `admin`, `admin`),
			`test`,
			&container.Config{
				Image:  `vibioh/dashboard`,
				Labels: map[string]string{`CUSTOM_LABEL`: `testing`, `owner`: `admin`, `app`: `test`},
				Env:    []string{`PATH=/usr/bin`},
				Cmd:    []string{`entrypoint.sh`, `start`},
				Healthcheck: &container.HealthConfig{
					Test:     []string{`CMD`, `alcotest`},
					Retries:  10,
					Interval: time.Second * 30,
					Timeout:  time.Second * 10,
				},
			},
			nil,
		},
	}

	var failed bool

	for _, testCase := range cases {
		result, err := getConfig(testCase.service, testCase.user, testCase.appName)

		failed = false

		if err == nil && testCase.wantErr != nil {
			failed = true
		} else if err != nil && testCase.wantErr == nil {
			failed = true
		} else if err != nil && err.Error() != testCase.wantErr.Error() {
			failed = true
		} else if !reflect.DeepEqual(result, testCase.want) {
			failed = true
		}

		if failed {
			t.Errorf(`getConfig(%v, %v, %v) = (%v, %v), want (%v, %v)`, testCase.service, testCase.user, testCase.appName, result, err, testCase.want, testCase.wantErr)
		}
	}
}

func TestGetHostConfig(t *testing.T) {
	var cases = []struct {
		service *dockerComposeService
		want    *container.HostConfig
	}{
		{
			&dockerComposeService{},
			&container.HostConfig{
				LogConfig: container.LogConfig{Type: `json-file`, Config: map[string]string{
					`max-size`: `10m`,
				}},
				NetworkMode:   container.NetworkMode(*dockerNetwork),
				RestartPolicy: container.RestartPolicy{Name: `on-failure`, MaximumRetryCount: 5},
				Resources: container.Resources{
					CPUShares: defaultCPUShares,
					Memory:    minMemory,
				},
				SecurityOpt: []string{`no-new-privileges`},
			},
		},
		{
			&dockerComposeService{
				ReadOnly:    true,
				CPUShares:   512,
				MemoryLimit: 33554432,
			},
			&container.HostConfig{
				LogConfig: container.LogConfig{Type: `json-file`, Config: map[string]string{
					`max-size`: `10m`,
				}},
				NetworkMode:    container.NetworkMode(*dockerNetwork),
				RestartPolicy:  container.RestartPolicy{Name: `on-failure`, MaximumRetryCount: 5},
				ReadonlyRootfs: true,
				Resources: container.Resources{
					CPUShares: 512,
					Memory:    33554432,
				},
				SecurityOpt: []string{`no-new-privileges`},
			},
		},
		{
			&dockerComposeService{
				ReadOnly:    true,
				CPUShares:   512,
				MemoryLimit: 20973619200,
			},
			&container.HostConfig{
				LogConfig: container.LogConfig{Type: `json-file`, Config: map[string]string{
					`max-size`: `10m`,
				}},
				NetworkMode:    container.NetworkMode(*dockerNetwork),
				RestartPolicy:  container.RestartPolicy{Name: `on-failure`, MaximumRetryCount: 5},
				ReadonlyRootfs: true,
				Resources: container.Resources{
					CPUShares: 512,
					Memory:    maxMemory,
				},
				SecurityOpt: []string{`no-new-privileges`},
			},
		},
	}

	for _, testCase := range cases {
		if result := getHostConfig(testCase.service, nil); !reflect.DeepEqual(result, testCase.want) {
			t.Errorf(`getHostConfig(%v) = %v, want %v`, testCase.service, result, testCase.want)
		}
	}
}

func TestGetNetworkConfig(t *testing.T) {
	var cases = []struct {
		service *dockerComposeService
		want    *network.NetworkingConfig
	}{
		{
			&dockerComposeService{},
			&network.NetworkingConfig{
				EndpointsConfig: map[string]*network.EndpointSettings{
					*dockerNetwork: {
						Aliases: []string{`service`},
					},
				},
			},
		},
		{
			&dockerComposeService{
				Links: []string{`db`},
			},
			&network.NetworkingConfig{
				EndpointsConfig: map[string]*network.EndpointSettings{
					*dockerNetwork: {
						Aliases: []string{`service`},
						Links:   []string{`db:db`},
					},
				},
			},
		},
		{
			&dockerComposeService{
				Links: []string{`postgres:db`},
			},
			&network.NetworkingConfig{
				EndpointsConfig: map[string]*network.EndpointSettings{
					*dockerNetwork: {
						Aliases: []string{`service`},
						Links:   []string{`postgres:db`},
					},
				},
			},
		},
		{
			&dockerComposeService{
				Links: []string{`db:postgres`},
			},
			&network.NetworkingConfig{
				EndpointsConfig: map[string]*network.EndpointSettings{
					*dockerNetwork: {
						Aliases: []string{`service`},
						Links:   []string{`db:postgres`},
					},
				},
			},
		},
	}

	for _, testCase := range cases {
		if result := getNetworkConfig(`service`, testCase.service); !reflect.DeepEqual(result, testCase.want) {
			t.Errorf(`getNetworkConfig(%+v) = %+v, want %+v`, testCase.service, result.EndpointsConfig[`traefik`], testCase.want.EndpointsConfig[`traefik`])
		}
	}
}

func TestGetServiceFullName(t *testing.T) {
	var cases = []struct {
		app     string
		service string
		want    string
	}{
		{
			`dashboard`,
			`api`,
			`dashboard_api_deploy`,
		},
	}

	for _, testCase := range cases {
		if result := getServiceFullName(testCase.app, testCase.service); result != testCase.want {
			t.Errorf(`getServiceFullName(%v, %v) = %v, want %v`, testCase.app, testCase.service, result, testCase.want)
		}
	}
}

func TestGetFinalName(t *testing.T) {
	var cases = []struct {
		serviceFullName string
		want            string
	}{
		{
			`dashboard_deploy`,
			`dashboard`,
		},
		{
			`dashboard`,
			`dashboard`,
		},
	}

	for _, testCase := range cases {
		if result := getFinalName(testCase.serviceFullName); result != testCase.want {
			t.Errorf(`getFinalName(%v) = %v, want %v`, testCase.serviceFullName, result, testCase.want)
		}
	}
}

func TestComposeFailed(t *testing.T) {
	var cases = []struct {
		user       *authProvider.User
		appName    string
		err        error
		want       string
		wantStatus int
	}{
		{
			authProvider.NewUser(0, `admin`, `admin`),
			`test`,
			errors.New(`test unit error`),
			`[admin] [test] Failed to deploy: test unit error
`,
			http.StatusInternalServerError,
		},
	}

	for _, testCase := range cases {
		writer := httptest.NewRecorder()

		composeFailed(writer, testCase.user, testCase.appName, testCase.err)

		if result := writer.Code; result != testCase.wantStatus {
			t.Errorf(`composeFailed(%v, %v, %v) = %v, want %v`, testCase.user, testCase.appName, testCase.err, result, testCase.wantStatus)
		}

		if result, _ := request.ReadBody(writer.Result().Body); string(result) != testCase.want {
			t.Errorf(`composeFailed(%v, %v, %v) = %v, want %v`, testCase.user, testCase.appName, testCase.err, string(result), testCase.want)
		}
	}
}
<commit_msg>Fixing test<commit_after>package docker

import (
	"errors"
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"
	"time"

	authProvider "github.com/ViBiOh/auth/provider"
	"github.com/ViBiOh/httputils/request"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
)

func TestGetConfig(t *testing.T) {
	var cases = []struct {
		service *dockerComposeService
		user    *authProvider.User
		appName string
		want    *container.Config
		wantErr error
	}{
		{
			&dockerComposeService{
				Healthcheck: &dockerComposeHealthcheck{
					Interval: `abcd`,
				},
			},
			authProvider.NewUser(0, `admin`, `admin`),
			`test`,
			nil,
			errors.New(`Error while parsing healthcheck interval: time: invalid duration abcd`),
		},
		{
			&dockerComposeService{
				Healthcheck: &dockerComposeHealthcheck{
					Interval: `30s`,
					Timeout:  `abcd`,
				},
			},
			authProvider.NewUser(0, `admin`, `admin`),
			`test`,
			nil,
			errors.New(`Error while parsing healthcheck timeout: time: invalid duration abcd`),
		},
		{
			&dockerComposeService{},
			authProvider.NewUser(0, `admin`, `admin`),
			`test`,
			&container.Config{
				Labels: map[string]string{`owner`: `admin`, `app`: `test`},
				Env:    []string{},
				User:   `1000`,
			},
			nil,
		},
		{
			&dockerComposeService{
				Image:       `vibioh/dashboard`,
				Environment: map[string]string{`PATH`: `/usr/bin`},
				Labels:      map[string]string{`CUSTOM_LABEL`: `testing`},
				Command:     []string{`entrypoint.sh`, `start`},
				User:        `1000`,
				Healthcheck: &dockerComposeHealthcheck{
					Test:     []string{`CMD`, `alcotest`},
					Retries:  10,
					Interval: `30s`,
					Timeout:  `10s`,
				},
			},
			authProvider.NewUser(0, `admin`, `admin`),
			`test`,
			&container.Config{
				Image:  `vibioh/dashboard`,
				Labels: map[string]string{`CUSTOM_LABEL`: `testing`, `owner`: `admin`, `app`: `test`},
				Env:    []string{`PATH=/usr/bin`},
				Cmd:    []string{`entrypoint.sh`, `start`},
				User:   `1000`,
				Healthcheck: &container.HealthConfig{
					Test:     []string{`CMD`, `alcotest`},
					Retries:  10,
					Interval: time.Second * 30,
					Timeout:  time.Second * 10,
				},
			},
			nil,
		},
	}

	var failed bool

	for _, testCase := range cases {
		result, err := getConfig(testCase.service, testCase.user, testCase.appName)

		failed = false

		if err == nil && testCase.wantErr != nil {
			failed = true
		} else if err != nil && testCase.wantErr == nil {
			failed = true
		} else if err != nil && err.Error() != testCase.wantErr.Error() {
			failed = true
		} else if !reflect.DeepEqual(result, testCase.want) {
			failed = true
		}

		if failed {
			t.Errorf(`getConfig(%+v, %+v, %+v) = (%+v, %+v), want (%+v, %+v)`, testCase.service, testCase.user, testCase.appName, result, err, testCase.want, testCase.wantErr)
		}
	}
}

func TestGetHostConfig(t *testing.T) {
	var cases = []struct {
		service *dockerComposeService
		want    *container.HostConfig
	}{
		{
			&dockerComposeService{},
			&container.HostConfig{
				LogConfig: container.LogConfig{Type: `json-file`, Config: map[string]string{
					`max-size`: `10m`,
				}},
				NetworkMode:   container.NetworkMode(*dockerNetwork),
				RestartPolicy: container.RestartPolicy{Name: `on-failure`, MaximumRetryCount: 5},
				Resources: container.Resources{
					CPUShares: defaultCPUShares,
					Memory:    minMemory,
				},
				SecurityOpt: []string{`no-new-privileges`},
			},
		},
		{
			&dockerComposeService{
				ReadOnly:    true,
				CPUShares:   512,
				MemoryLimit: 33554432,
			},
			&container.HostConfig{
				LogConfig: container.LogConfig{Type: `json-file`, Config: map[string]string{
					`max-size`: `10m`,
				}},
				NetworkMode:    container.NetworkMode(*dockerNetwork),
				RestartPolicy:  container.RestartPolicy{Name: `on-failure`, MaximumRetryCount: 5},
				ReadonlyRootfs: true,
				Resources: container.Resources{
					CPUShares: 512,
					Memory:    33554432,
				},
				SecurityOpt: []string{`no-new-privileges`},
			},
		},
		{
			&dockerComposeService{
				ReadOnly:    true,
				CPUShares:   512,
				MemoryLimit: 20973619200,
			},
			&container.HostConfig{
				LogConfig: container.LogConfig{Type: `json-file`, Config: map[string]string{
					`max-size`: `10m`,
				}},
				NetworkMode:    container.NetworkMode(*dockerNetwork),
				RestartPolicy:  container.RestartPolicy{Name: `on-failure`, MaximumRetryCount: 5},
				ReadonlyRootfs: true,
				Resources: container.Resources{
					CPUShares: 512,
					Memory:    maxMemory,
				},
				SecurityOpt: []string{`no-new-privileges`},
			},
		},
	}

	for _, testCase := range cases {
		if result := getHostConfig(testCase.service, nil); !reflect.DeepEqual(result, testCase.want) {
			t.Errorf(`getHostConfig(%+v) = %+v, want %+v`, testCase.service, result, testCase.want)
		}
	}
}

func TestGetNetworkConfig(t *testing.T) {
	var cases = []struct {
		service *dockerComposeService
		want    *network.NetworkingConfig
	}{
		{
			&dockerComposeService{},
			&network.NetworkingConfig{
				EndpointsConfig: map[string]*network.EndpointSettings{
					*dockerNetwork: {
						Aliases: []string{`service`},
					},
				},
			},
		},
		{
			&dockerComposeService{
				Links: []string{`db`},
			},
			&network.NetworkingConfig{
				EndpointsConfig: map[string]*network.EndpointSettings{
					*dockerNetwork: {
						Aliases: []string{`service`},
						Links:   []string{`db:db`},
					},
				},
			},
		},
		{
			&dockerComposeService{
				Links: []string{`postgres:db`},
			},
			&network.NetworkingConfig{
				EndpointsConfig: map[string]*network.EndpointSettings{
					*dockerNetwork: {
						Aliases: []string{`service`},
						Links:   []string{`postgres:db`},
					},
				},
			},
		},
		{
			&dockerComposeService{
				Links: []string{`db:postgres`},
			},
			&network.NetworkingConfig{
				EndpointsConfig: map[string]*network.EndpointSettings{
					*dockerNetwork: {
						Aliases: []string{`service`},
						Links:   []string{`db:postgres`},
					},
				},
			},
		},
	}

	for _, testCase := range cases {
		if result := getNetworkConfig(`service`, testCase.service); !reflect.DeepEqual(result, testCase.want) {
			t.Errorf(`getNetworkConfig(%+v) = %+v, want %+v`, testCase.service, result.EndpointsConfig[`traefik`], testCase.want.EndpointsConfig[`traefik`])
		}
	}
}

func TestGetServiceFullName(t *testing.T) {
	var cases = []struct {
		app     string
		service string
		want    string
	}{
		{
			`dashboard`,
			`api`,
			`dashboard_api_deploy`,
		},
	}

	for _, testCase := range cases {
		if result := getServiceFullName(testCase.app, testCase.service); result != testCase.want {
			t.Errorf(`getServiceFullName(%+v, %+v) = %+v, want %+v`, testCase.app, testCase.service, result, testCase.want)
		}
	}
}

func TestGetFinalName(t *testing.T) {
	var cases = []struct {
		serviceFullName string
		want            string
	}{
		{
			`dashboard_deploy`,
			`dashboard`,
		},
		{
			`dashboard`,
			`dashboard`,
		},
	}

	for _, testCase := range cases {
		if result := getFinalName(testCase.serviceFullName); result != testCase.want {
			t.Errorf(`getFinalName(%+v) = %+v, want %+v`, testCase.serviceFullName, result, testCase.want)
		}
	}
}

func TestComposeFailed(t *testing.T) {
	var cases = []struct {
		user       *authProvider.User
		appName    string
		err        error
		want       string
		wantStatus int
	}{
		{
			authProvider.NewUser(0, `admin`, `admin`),
			`test`,
			errors.New(`test unit error`),
			`[admin] [test] Failed to deploy: test unit error
`,
			http.StatusInternalServerError,
		},
	}

	for _, testCase := range cases {
		writer := httptest.NewRecorder()

		composeFailed(writer, testCase.user, testCase.appName, testCase.err)

		if result := writer.Code; result != testCase.wantStatus {
			t.Errorf(`composeFailed(%+v, %+v, %+v) = %+v, want %+v`, testCase.user, testCase.appName, testCase.err, result, testCase.wantStatus)
		}

		if result, _ := request.ReadBody(writer.Result().Body); string(result) != testCase.want {
			t.Errorf(`composeFailed(%+v, %+v, %+v) = %+v, want %+v`, testCase.user, testCase.appName, testCase.err, string(result), testCase.want)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

func main() {
	log.Fatal(http.ListenAndServe(":8080", mux("/static")))
}

const jsonParam = "json"

func mux(dir string) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request){
		if r.Method == http.MethodDelete {
			deleteFileIfExists(w, r, dir)
			return
		}
		if _, ok := r.URL.Query()[jsonParam]; ok {
			listFilesAsJson(w, dir)
			return
		}
		http.FileServer(http.Dir(dir)).ServeHTTP(w, r)
	})
	return mux
}

func listFilesAsJson(w http.ResponseWriter, dir string) {
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	ret := []string{}
	for _, f := range files {
		ret = append(ret, f.Name())
	}
	w.Header().Add("Content-Type", "application/json")
	json.NewEncoder(w).Encode(ret)
}

func deleteFileIfExists(w http.ResponseWriter, r *http.Request, dir string) {
	fileName := strings.TrimPrefix(r.URL.Path, "/")
	filePath := filepath.Join(dir, fileName)
	_, err := os.Stat(filePath)
	if err != nil {
		http.Error(w, fmt.Sprintf("Unknown file %s", fileName), http.StatusNotFound)
		return
	}
	err = os.Remove(filePath)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to delete file %s: %v", fileName, err), http.StatusInternalServerError)
		return
	}
}
<commit_msg>Sorting downloaded files by their modification time in descending order (fixes #211)<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

func main() {
	log.Fatal(http.ListenAndServe(":8080", mux("/static")))
}

const jsonParam = "json"

func mux(dir string) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request){
		if r.Method == http.MethodDelete {
			deleteFileIfExists(w, r, dir)
			return
		}
		if _, ok := r.URL.Query()[jsonParam]; ok {
			listFilesAsJson(w, dir)
			return
		}
		http.FileServer(http.Dir(dir)).ServeHTTP(w, r)
	})
	return mux
}

func listFilesAsJson(w http.ResponseWriter, dir string) {
	files, err := ioutil.ReadDir(dir)
	sort.Slice(files, func(i, j int) bool {
		return files[i].ModTime().After(files[j].ModTime())
	})
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	ret := []string{}
	for _, f := range files {
		ret = append(ret, f.Name())
	}
	w.Header().Add("Content-Type", "application/json")
	json.NewEncoder(w).Encode(ret)
}

func deleteFileIfExists(w http.ResponseWriter, r *http.Request, dir string) {
	fileName := strings.TrimPrefix(r.URL.Path, "/")
	filePath := filepath.Join(dir, fileName)
	_, err := os.Stat(filePath)
	if err != nil {
		http.Error(w, fmt.Sprintf("Unknown file %s", fileName), http.StatusNotFound)
		return
	}
	err = os.Remove(filePath)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to delete file %s: %v", fileName, err), http.StatusInternalServerError)
		return
	}
}
<|endoftext|>