<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package codegen

import (
	"path/filepath"
	"strings"

	"encoding/json"

	"github.com/pkg/errors"
	"github.com/uber/zanzibar/module"
)

// NewDefaultModuleSystem creates a fresh instance of the default zanzibar
// module system (clients, endpoints)
func NewDefaultModuleSystem(h *PackageHelper) (*module.System, error) {
	system := module.NewSystem()
	tmpl, err := NewTemplate()

	if err != nil {
		return nil, err
	}

	// Register client module class and type generators
	if err := system.RegisterClass("client", module.Class{
		Directory:         "clients",
		ClassType:         module.MultiModule,
		ClassDependencies: []string{},
	}); err != nil {
		return nil, errors.Wrapf(err, "Error registering client class")
	}

	if err := system.RegisterClassType("client", "http", &HTTPClientGenerator{
		templates:     tmpl,
		packageHelper: h,
	}); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error registering http client class type",
		)
	}

	if err := system.RegisterClassType("client", "tchannel", &TCahnnelClientGenerator{
		templates:     tmpl,
		packageHelper: h,
	}); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error registering TChannel client class type",
		)
	}

	// TODO: Register endpoint module class and type generators
	return system, nil
}

/*
 * HTTP Client Generator
 */

// HTTPClientGenerator generates an instance of a zanzibar http client
type HTTPClientGenerator struct {
	templates     *Template
	packageHelper *PackageHelper
}

// Generate returns the HTTP client generated files as a map of relative file
// path (relative to the target buid directory) to file bytes.
func (generator *HTTPClientGenerator) Generate(
	instance *module.Instance,
) (map[string][]byte, error) {
	// Parse the client config from the endpoint JSON file
	clientConfig, err := readClientConfig(instance.JSONFileRaw)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading HTTP client %s JSON config",
			instance.InstanceName,
		)
	}

	clientSpec, err := NewHTTPClientSpec(
		filepath.Join(
			instance.BaseDirectory,
			instance.Directory,
			instance.JSONFileName,
		),
		clientConfig,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error initializing HTTPClientSpec for %s",
			instance.InstanceName,
		)
	}

	clientMeta := &ClientMeta{
		PackageName:      clientSpec.ModuleSpec.PackageName,
		Services:         clientSpec.ModuleSpec.Services,
		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
		ClientID:         clientSpec.ClientID,
	}

	client, err := generator.templates.execTemplate(
		"http_client.tmpl",
		clientMeta,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing HTTP client template for %s",
			instance.InstanceName,
		)
	}

	structs, err := generator.templates.execTemplate(
		"structs.tmpl",
		clientMeta,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing HTTP client structs template for %s",
			instance.InstanceName,
		)
	}

	clientDirectory := filepath.Join(
		generator.packageHelper.CodeGenTargetPath(),
		instance.Directory,
	)

	clientFilePath, err := filepath.Rel(clientDirectory, clientSpec.GoFileName)
	if err != nil {
		clientFilePath = clientSpec.GoFileName
	}

	structFilePath, err := filepath.Rel(
		clientDirectory,
		clientSpec.GoStructsFileName,
	)
	if err != nil {
		structFilePath = clientSpec.GoStructsFileName
	}

	// Return the client files
	return map[string][]byte{
		clientFilePath: client,
		structFilePath: structs,
	}, nil
}

/*
 * TChannel Client Generator
 */

// TCahnnelClientGenerator generates an instance of a zanzibar TChannel client
type TCahnnelClientGenerator struct {
	templates     *Template
	packageHelper *PackageHelper
}

// Generate returns the TChannel client generated files as a map of relative file
// path (relative to the target build directory) to file bytes.
func (generator *TCahnnelClientGenerator) Generate(
	instance *module.Instance,
) (map[string][]byte, error) {
	// Parse the client config from the endpoint JSON file
	clientConfig, err := readClientConfig(instance.JSONFileRaw)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading TChannel client %s JSON config",
			instance.InstanceName,
		)
	}

	clientSpec, err := NewTChannelClientSpec(
		filepath.Join(
			instance.BaseDirectory,
			instance.Directory,
			instance.JSONFileName,
		),
		clientConfig,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error initializing TChannelClientSpec for %s",
			instance.InstanceName,
		)
	}

	clientMeta := &ClientMeta{
		PackageName:      clientSpec.ModuleSpec.PackageName,
		Services:         clientSpec.ModuleSpec.Services,
		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
		ClientID:         clientSpec.ClientID,
	}

	client, err := generator.templates.execTemplate(
		"tchannel_client.tmpl",
		clientMeta,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel client template for %s",
			instance.InstanceName,
		)
	}

	server, err := generator.templates.execTemplate(
		"tchannel_server.tmpl",
		clientMeta,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel server template for %s",
			instance.InstanceName,
		)
	}

	handler, err := generator.templates.execTemplate(
		"tchannel_handler.tmpl",
		clientMeta,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel handler template for %s",
			instance.InstanceName,
		)
	}

	clientDirectory := filepath.Join(
		generator.packageHelper.CodeGenTargetPath(),
		instance.Directory,
	)

	clientFilePath, err := filepath.Rel(clientDirectory, clientSpec.GoFileName)
	if err != nil {
		clientFilePath = clientSpec.GoFileName
	}

	// TODO:(lu) the locations of tchannel server and handler files are tentative
	serverFilePath := strings.TrimRight(clientFilePath, ".go") + "_server.go"
	handlerFilePath := strings.TrimRight(clientFilePath, ".go") + "_handler.go"

	// Return the client files
	return map[string][]byte{
		clientFilePath:  client,
		serverFilePath:  server,
		handlerFilePath: handler,
	}, nil
}

func readClientConfig(rawConfig []byte) (*clientClassConfig, error) {
	var clientConfig clientClassConfig
	if err := json.Unmarshal(rawConfig, &clientConfig); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading config for HTTP client instance",
		)
	}
	clientConfig.Config["clientId"] = clientConfig.Name
	clientConfig.Config["clientType"] = clientConfig.Type
	return &clientConfig, nil
}
<commit_msg>Fix typo<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package codegen

import (
	"path/filepath"
	"strings"

	"encoding/json"

	"github.com/pkg/errors"
	"github.com/uber/zanzibar/module"
)

// NewDefaultModuleSystem creates a fresh instance of the default zanzibar
// module system (clients, endpoints)
func NewDefaultModuleSystem(h *PackageHelper) (*module.System, error) {
	system := module.NewSystem()
	tmpl, err := NewTemplate()

	if err != nil {
		return nil, err
	}

	// Register client module class and type generators
	if err := system.RegisterClass("client", module.Class{
		Directory:         "clients",
		ClassType:         module.MultiModule,
		ClassDependencies: []string{},
	}); err != nil {
		return nil, errors.Wrapf(err, "Error registering client class")
	}

	if err := system.RegisterClassType("client", "http", &HTTPClientGenerator{
		templates:     tmpl,
		packageHelper: h,
	}); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error registering http client class type",
		)
	}

	if err := system.RegisterClassType("client", "tchannel", &TChannelClientGenerator{
		templates:     tmpl,
		packageHelper: h,
	}); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error registering TChannel client class type",
		)
	}

	// TODO: Register endpoint module class and type generators
	return system, nil
}

/*
 * HTTP Client Generator
 */

// HTTPClientGenerator generates an instance of a zanzibar http client
type HTTPClientGenerator struct {
	templates     *Template
	packageHelper *PackageHelper
}

// Generate returns the HTTP client generated files as a map of relative file
// path (relative to the target buid directory) to file bytes.
func (generator *HTTPClientGenerator) Generate(
	instance *module.Instance,
) (map[string][]byte, error) {
	// Parse the client config from the endpoint JSON file
	clientConfig, err := readClientConfig(instance.JSONFileRaw)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading HTTP client %s JSON config",
			instance.InstanceName,
		)
	}

	clientSpec, err := NewHTTPClientSpec(
		filepath.Join(
			instance.BaseDirectory,
			instance.Directory,
			instance.JSONFileName,
		),
		clientConfig,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error initializing HTTPClientSpec for %s",
			instance.InstanceName,
		)
	}

	clientMeta := &ClientMeta{
		PackageName:      clientSpec.ModuleSpec.PackageName,
		Services:         clientSpec.ModuleSpec.Services,
		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
		ClientID:         clientSpec.ClientID,
	}

	client, err := generator.templates.execTemplate(
		"http_client.tmpl",
		clientMeta,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing HTTP client template for %s",
			instance.InstanceName,
		)
	}

	structs, err := generator.templates.execTemplate(
		"structs.tmpl",
		clientMeta,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing HTTP client structs template for %s",
			instance.InstanceName,
		)
	}

	clientDirectory := filepath.Join(
		generator.packageHelper.CodeGenTargetPath(),
		instance.Directory,
	)

	clientFilePath, err := filepath.Rel(clientDirectory, clientSpec.GoFileName)
	if err != nil {
		clientFilePath = clientSpec.GoFileName
	}

	structFilePath, err := filepath.Rel(
		clientDirectory,
		clientSpec.GoStructsFileName,
	)
	if err != nil {
		structFilePath = clientSpec.GoStructsFileName
	}

	// Return the client files
	return map[string][]byte{
		clientFilePath: client,
		structFilePath: structs,
	}, nil
}

/*
 * TChannel Client Generator
 */

// TChannelClientGenerator generates an instance of a zanzibar TChannel client
type TChannelClientGenerator struct {
	templates     *Template
	packageHelper *PackageHelper
}

// Generate returns the TChannel client generated files as a map of relative file
// path (relative to the target build directory) to file bytes.
func (generator *TChannelClientGenerator) Generate(
	instance *module.Instance,
) (map[string][]byte, error) {
	// Parse the client config from the endpoint JSON file
	clientConfig, err := readClientConfig(instance.JSONFileRaw)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading TChannel client %s JSON config",
			instance.InstanceName,
		)
	}

	clientSpec, err := NewTChannelClientSpec(
		filepath.Join(
			instance.BaseDirectory,
			instance.Directory,
			instance.JSONFileName,
		),
		clientConfig,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error initializing TChannelClientSpec for %s",
			instance.InstanceName,
		)
	}

	clientMeta := &ClientMeta{
		PackageName:      clientSpec.ModuleSpec.PackageName,
		Services:         clientSpec.ModuleSpec.Services,
		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
		ClientID:         clientSpec.ClientID,
	}

	client, err := generator.templates.execTemplate(
		"tchannel_client.tmpl",
		clientMeta,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel client template for %s",
			instance.InstanceName,
		)
	}

	server, err := generator.templates.execTemplate(
		"tchannel_server.tmpl",
		clientMeta,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel server template for %s",
			instance.InstanceName,
		)
	}

	handler, err := generator.templates.execTemplate(
		"tchannel_handler.tmpl",
		clientMeta,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel handler template for %s",
			instance.InstanceName,
		)
	}

	clientDirectory := filepath.Join(
		generator.packageHelper.CodeGenTargetPath(),
		instance.Directory,
	)

	clientFilePath, err := filepath.Rel(clientDirectory, clientSpec.GoFileName)
	if err != nil {
		clientFilePath = clientSpec.GoFileName
	}

	// TODO:(lu) the locations of tchannel server and handler files are tentative
	serverFilePath := strings.TrimRight(clientFilePath, ".go") + "_server.go"
	handlerFilePath := strings.TrimRight(clientFilePath, ".go") + "_handler.go"

	// Return the client files
	return map[string][]byte{
		clientFilePath:  client,
		serverFilePath:  server,
		handlerFilePath: handler,
	}, nil
}

func readClientConfig(rawConfig []byte) (*clientClassConfig, error) {
	var clientConfig clientClassConfig
	if err := json.Unmarshal(rawConfig, &clientConfig); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading config for HTTP client instance",
		)
	}
	clientConfig.Config["clientId"] = clientConfig.Name
	clientConfig.Config["clientType"] = clientConfig.Type
	return &clientConfig, nil
}
<|endoftext|>


<commit_before>package grandcentral

import (
	"fmt"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	proto "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/goprotobuf/proto"
	datastore "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-datastore"
	peer "github.com/jbenet/go-ipfs/p2p/peer"
	dhtpb "github.com/jbenet/go-ipfs/routing/dht/pb"
	proxy "github.com/jbenet/go-ipfs/routing/grandcentral/proxy"
	util "github.com/jbenet/go-ipfs/util"
	errors "github.com/jbenet/go-ipfs/util/debugerror"
)

// Server handles routing queries using a database backend
type Server struct {
	local           peer.ID
	routingBackend  datastore.ThreadSafeDatastore
	peerstore       peer.Peerstore
	*proxy.Loopback // so server can be injected into client
}

// NewServer creates a new GrandCentral routing Server
func NewServer(ds datastore.ThreadSafeDatastore, ps peer.Peerstore, local peer.ID) (*Server, error) {
	s := &Server{local, ds, ps, nil}
	s.Loopback = &proxy.Loopback{
		Handler: s,
		Local:   local,
	}
	return s, nil
}

// HandleLocalRequest implements the proxy.RequestHandler interface. This is
// where requests are received from the outside world.
func (s *Server) HandleRequest(ctx context.Context, p peer.ID, req *dhtpb.Message) *dhtpb.Message {
	_, response := s.handleMessage(ctx, p, req) // ignore response peer. it's local.
	return response
}

// TODO extract backend. backend can be implemented with whatever database we desire
func (s *Server) handleMessage(
	ctx context.Context, p peer.ID, req *dhtpb.Message) (peer.ID, *dhtpb.Message) {

	defer log.EventBegin(ctx, "routingMessageReceived", req, p, s.local).Done() // TODO may need to differentiate between local and remote

	var response = dhtpb.NewMessage(req.GetType(), req.GetKey(), req.GetClusterLevel())
	switch req.GetType() {

	case dhtpb.Message_GET_VALUE:
		rawRecord, err := getRoutingRecord(s.routingBackend, util.Key(req.GetKey()))
		if err != nil {
			return "", nil
		}
		response.Record = rawRecord
		// TODO before merging: if we know any providers for the requested value, return those.
		return p, response

	case dhtpb.Message_PUT_VALUE:
		// TODO before merging: verifyRecord(req.GetRecord())
		putRoutingRecord(s.routingBackend, util.Key(req.GetKey()), req.GetRecord())
		return p, req // TODO before merging: verify that we should return record

	case dhtpb.Message_FIND_NODE:
		p := s.peerstore.PeerInfo(peer.ID(req.GetKey()))
		pri := []dhtpb.PeerRoutingInfo{
			dhtpb.PeerRoutingInfo{
				PeerInfo: p,
				// Connectedness: TODO
			},
		}
		response.CloserPeers = dhtpb.PeerRoutingInfosToPBPeers(pri)
		return p.ID, response

	case dhtpb.Message_ADD_PROVIDER:
		// FIXME(btc): do we want to store these locally? I think the
		// storeProvidersToPeerstore behavior is straight from the DHT message
		// handler.
		storeProvidersToPeerstore(s.peerstore, p, req.GetProviderPeers())

		if err := putRoutingProviders(s.routingBackend, util.Key(req.GetKey()), req.GetProviderPeers()); err != nil {
			return "", nil
		}
		return "", nil

	case dhtpb.Message_GET_PROVIDERS:
		providers, err := getRoutingProviders(s.routingBackend, util.Key(req.GetKey()))
		if err != nil {
			return "", nil
		}
		response.ProviderPeers = providers
		return p, response

	case dhtpb.Message_PING:
		return p, req
	default:
	}
	return "", nil
}

var _ proxy.RequestHandler = &Server{}
var _ proxy.Proxy = &Server{}

func getRoutingRecord(ds datastore.Datastore, k util.Key) (*dhtpb.Record, error) {
	dskey := k.DsKey()
	val, err := ds.Get(dskey)
	if err != nil {
		return nil, errors.Wrap(err)
	}
	recordBytes, ok := val.([]byte)
	if !ok {
		return nil, fmt.Errorf("datastore had non byte-slice value for %v", dskey)
	}
	var record dhtpb.Record
	if err := proto.Unmarshal(recordBytes, &record); err != nil {
		return nil, errors.New("failed to unmarshal dht record from datastore")
	}
	return &record, nil
}

func putRoutingRecord(ds datastore.Datastore, k util.Key, value *dhtpb.Record) error {
	data, err := proto.Marshal(value)
	if err != nil {
		return err
	}
	dskey := k.DsKey()
	// TODO namespace
	if err := ds.Put(dskey, data); err != nil {
		return err
	}
	return nil
}

func putRoutingProviders(ds datastore.Datastore, k util.Key, newRecords []*dhtpb.Message_Peer) error {
	log.Event(context.Background(), "putRoutingProviders", &k)
	oldRecords, err := getRoutingProviders(ds, k)
	if err != nil {
		return err
	}
	mergedRecords := make(map[string]*dhtpb.Message_Peer)
	for _, provider := range oldRecords {
		mergedRecords[provider.GetId()] = provider // add original records
	}
	for _, provider := range newRecords {
		mergedRecords[provider.GetId()] = provider // overwrite old record if new exists
	}
	var protomsg dhtpb.Message
	protomsg.ProviderPeers = make([]*dhtpb.Message_Peer, 0, len(mergedRecords))
	for _, provider := range mergedRecords {
		protomsg.ProviderPeers = append(protomsg.ProviderPeers, provider)
	}
	data, err := proto.Marshal(&protomsg)
	if err != nil {
		return err
	}
	return ds.Put(providerKey(k), data)
}

func storeProvidersToPeerstore(ps peer.Peerstore, p peer.ID, providers []*dhtpb.Message_Peer) {
	for _, provider := range providers {
		providerID := peer.ID(provider.GetId())
		if providerID != p {
			log.Errorf("provider message came from third-party %s", p)
			continue
		}
		for _, maddr := range provider.Addresses() {
			// as a router, we want to store addresses for peers who have provided
			ps.AddAddr(p, maddr, peer.AddressTTL)
		}
	}
}

func getRoutingProviders(ds datastore.Datastore, k util.Key) ([]*dhtpb.Message_Peer, error) {
	e := log.EventBegin(context.Background(), "getProviders", &k)
	defer e.Done()
	var providers []*dhtpb.Message_Peer
	if v, err := ds.Get(providerKey(k)); err == nil {
		if data, ok := v.([]byte); ok {
			var msg dhtpb.Message
			if err := proto.Unmarshal(data, &msg); err != nil {
				return nil, err
			}
			providers = append(providers, msg.GetProviderPeers()...)
		}
	}
	return providers, nil
}

func providerKey(k util.Key) datastore.Key {
	return datastore.KeyWithNamespaces([]string{"routing", "providers", k.String()})
}
<commit_msg>misc(gcr/s) rm TODO<commit_after>package grandcentral

import (
	"fmt"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	proto "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/goprotobuf/proto"
	datastore "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-datastore"
	peer "github.com/jbenet/go-ipfs/p2p/peer"
	dhtpb "github.com/jbenet/go-ipfs/routing/dht/pb"
	proxy "github.com/jbenet/go-ipfs/routing/grandcentral/proxy"
	util "github.com/jbenet/go-ipfs/util"
	errors "github.com/jbenet/go-ipfs/util/debugerror"
)

// Server handles routing queries using a database backend
type Server struct {
	local           peer.ID
	routingBackend  datastore.ThreadSafeDatastore
	peerstore       peer.Peerstore
	*proxy.Loopback // so server can be injected into client
}

// NewServer creates a new GrandCentral routing Server
func NewServer(ds datastore.ThreadSafeDatastore, ps peer.Peerstore, local peer.ID) (*Server, error) {
	s := &Server{local, ds, ps, nil}
	s.Loopback = &proxy.Loopback{
		Handler: s,
		Local:   local,
	}
	return s, nil
}

// HandleLocalRequest implements the proxy.RequestHandler interface. This is
// where requests are received from the outside world.
func (s *Server) HandleRequest(ctx context.Context, p peer.ID, req *dhtpb.Message) *dhtpb.Message {
	_, response := s.handleMessage(ctx, p, req) // ignore response peer. it's local.
	return response
}

func (s *Server) handleMessage(
	ctx context.Context, p peer.ID, req *dhtpb.Message) (peer.ID, *dhtpb.Message) {

	defer log.EventBegin(ctx, "routingMessageReceived", req, p, s.local).Done() // TODO may need to differentiate between local and remote

	var response = dhtpb.NewMessage(req.GetType(), req.GetKey(), req.GetClusterLevel())
	switch req.GetType() {

	case dhtpb.Message_GET_VALUE:
		rawRecord, err := getRoutingRecord(s.routingBackend, util.Key(req.GetKey()))
		if err != nil {
			return "", nil
		}
		response.Record = rawRecord
		// TODO before merging: if we know any providers for the requested value, return those.
		return p, response

	case dhtpb.Message_PUT_VALUE:
		// TODO before merging: verifyRecord(req.GetRecord())
		putRoutingRecord(s.routingBackend, util.Key(req.GetKey()), req.GetRecord())
		return p, req // TODO before merging: verify that we should return record

	case dhtpb.Message_FIND_NODE:
		p := s.peerstore.PeerInfo(peer.ID(req.GetKey()))
		pri := []dhtpb.PeerRoutingInfo{
			dhtpb.PeerRoutingInfo{
				PeerInfo: p,
				// Connectedness: TODO
			},
		}
		response.CloserPeers = dhtpb.PeerRoutingInfosToPBPeers(pri)
		return p.ID, response

	case dhtpb.Message_ADD_PROVIDER:
		// FIXME(btc): do we want to store these locally? I think the
		// storeProvidersToPeerstore behavior is straight from the DHT message
		// handler.
		storeProvidersToPeerstore(s.peerstore, p, req.GetProviderPeers())

		if err := putRoutingProviders(s.routingBackend, util.Key(req.GetKey()), req.GetProviderPeers()); err != nil {
			return "", nil
		}
		return "", nil

	case dhtpb.Message_GET_PROVIDERS:
		providers, err := getRoutingProviders(s.routingBackend, util.Key(req.GetKey()))
		if err != nil {
			return "", nil
		}
		response.ProviderPeers = providers
		return p, response

	case dhtpb.Message_PING:
		return p, req
	default:
	}
	return "", nil
}

var _ proxy.RequestHandler = &Server{}
var _ proxy.Proxy = &Server{}

func getRoutingRecord(ds datastore.Datastore, k util.Key) (*dhtpb.Record, error) {
	dskey := k.DsKey()
	val, err := ds.Get(dskey)
	if err != nil {
		return nil, errors.Wrap(err)
	}
	recordBytes, ok := val.([]byte)
	if !ok {
		return nil, fmt.Errorf("datastore had non byte-slice value for %v", dskey)
	}
	var record dhtpb.Record
	if err := proto.Unmarshal(recordBytes, &record); err != nil {
		return nil, errors.New("failed to unmarshal dht record from datastore")
	}
	return &record, nil
}

func putRoutingRecord(ds datastore.Datastore, k util.Key, value *dhtpb.Record) error {
	data, err := proto.Marshal(value)
	if err != nil {
		return err
	}
	dskey := k.DsKey()
	// TODO namespace
	if err := ds.Put(dskey, data); err != nil {
		return err
	}
	return nil
}

func putRoutingProviders(ds datastore.Datastore, k util.Key, newRecords []*dhtpb.Message_Peer) error {
	log.Event(context.Background(), "putRoutingProviders", &k)
	oldRecords, err := getRoutingProviders(ds, k)
	if err != nil {
		return err
	}
	mergedRecords := make(map[string]*dhtpb.Message_Peer)
	for _, provider := range oldRecords {
		mergedRecords[provider.GetId()] = provider // add original records
	}
	for _, provider := range newRecords {
		mergedRecords[provider.GetId()] = provider // overwrite old record if new exists
	}
	var protomsg dhtpb.Message
	protomsg.ProviderPeers = make([]*dhtpb.Message_Peer, 0, len(mergedRecords))
	for _, provider := range mergedRecords {
		protomsg.ProviderPeers = append(protomsg.ProviderPeers, provider)
	}
	data, err := proto.Marshal(&protomsg)
	if err != nil {
		return err
	}
	return ds.Put(providerKey(k), data)
}

func storeProvidersToPeerstore(ps peer.Peerstore, p peer.ID, providers []*dhtpb.Message_Peer) {
	for _, provider := range providers {
		providerID := peer.ID(provider.GetId())
		if providerID != p {
			log.Errorf("provider message came from third-party %s", p)
			continue
		}
		for _, maddr := range provider.Addresses() {
			// as a router, we want to store addresses for peers who have provided
			ps.AddAddr(p, maddr, peer.AddressTTL)
		}
	}
}

func getRoutingProviders(ds datastore.Datastore, k util.Key) ([]*dhtpb.Message_Peer, error) {
	e := log.EventBegin(context.Background(), "getProviders", &k)
	defer e.Done()
	var providers []*dhtpb.Message_Peer
	if v, err := ds.Get(providerKey(k)); err == nil {
		if data, ok := v.([]byte); ok {
			var msg dhtpb.Message
			if err := proto.Unmarshal(data, &msg); err != nil {
				return nil, err
			}
			providers = append(providers, msg.GetProviderPeers()...)
		}
	}
	return providers, nil
}

func providerKey(k util.Key) datastore.Key {
	return datastore.KeyWithNamespaces([]string{"routing", "providers", k.String()})
}
<|endoftext|>


<commit_before>// Package testdata provides testing helpers and data for rules.
package testdata

import (
	"bytes"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"os"
	"os/exec"

	"github.com/golang/protobuf/v2/proto"
	descriptorpb "github.com/golang/protobuf/v2/types/descriptor"
	"github.com/jgeewax/api-linter/lint"
)

var protoc = "protoc"

func descriptorProtoFromSource(source []byte) (*descriptorpb.FileDescriptorProto, error) {
	tmpDir := os.TempDir()

	f, err := ioutil.TempFile(tmpDir, "proto*")

	if err != nil {
		return nil, err
	}
	defer mustCloseAndRemoveFile(f)

	if _, err = f.Write(source); err != nil {
		return nil, err
	}

	descSetF, err := ioutil.TempFile(tmpDir, "descset*")

	if err != nil {
		return nil, err
	}
	defer mustCloseAndRemoveFile(descSetF)

	cmd := exec.Command(
		protoc,
		"--include_source_info",
		fmt.Sprintf("--proto_path=%s", tmpDir),
		fmt.Sprintf("--descriptor_set_out=%s", descSetF.Name()),
		f.Name(),
	)

	var stdErrBuf bytes.Buffer

	cmd.Stderr = &stdErrBuf

	if err = cmd.Run(); err != nil {
		return nil, fmt.Errorf("protoc failed with %v and Stderr %q", err, stdErrBuf.String())
	}

	descSet, err := ioutil.ReadFile(descSetF.Name())

	if err != nil {
		return nil, err
	}

	protoset := &descriptorpb.FileDescriptorSet{}
	if err := proto.Unmarshal(descSet, protoset); err != nil {
		return nil, err
	}

	if len(protoset.GetFile()) == 0 {
		return nil, fmt.Errorf("protoset file list was empty")
	}

	return protoset.GetFile()[0], nil
}

func mustCloseAndRemoveFile(f *os.File) {
	if err := f.Close(); err != nil {
		log.Fatalf("Error closing proto file: %v", err)
	}

	if err := os.Remove(f.Name()); err != nil {
		log.Fatalf("Error removing proto file: %v", err)
	}
}

// MustCreateTemplate creates a template with name "test" from
// the provided template string.
func MustCreateTemplate(tmpl string) *template.Template {
	return template.Must(template.New("test").Parse(tmpl))
}

// MustCreateRequestFromTemplate creates a lint.Request from the provided template and test data.
func MustCreateRequestFromTemplate(tmpl *template.Template, testData interface{}) lint.Request {
	var b bytes.Buffer
	if err := tmpl.Execute(&b, testData); err != nil {
		log.Fatalf("Error executing template %v", err)
	}
	pd, err := descriptorProtoFromSource(b.Bytes())
	if err != nil {
		log.Fatalf("Error generating proto descriptor: %v", err)
	}
	req, err := lint.NewProtoFileRequest(pd)
	if err != nil {
		log.Fatalf("Error creating proto file request: %v", err)
	}
	return req
}
<commit_msg>Refactor.<commit_after>// Package testdata provides testing helpers and data for rules.
package testdata

import (
	"bytes"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"

	"github.com/golang/protobuf/v2/proto"
	descriptorpb "github.com/golang/protobuf/v2/types/descriptor"
	"github.com/jgeewax/api-linter/lint"
)

var protoc = "protoc"

func descriptorProtoFromSource(source io.Reader) (*descriptorpb.FileDescriptorProto, error) {
	tmpDir := os.TempDir()

	f, err := ioutil.TempFile(tmpDir, "proto*")
	if err != nil {
		return nil, err
	}
	defer mustCloseAndRemoveFile(f)

	if _, err = io.Copy(f, source); err != nil {
		return nil, err
	}

	descSetF, err := ioutil.TempFile(tmpDir, "descset*")
	if err != nil {
		return nil, err
	}
	defer mustCloseAndRemoveFile(descSetF)

	cmd := exec.Command(
		protoc,
		"--include_source_info",
		fmt.Sprintf("--proto_path=%s", tmpDir),
		fmt.Sprintf("--descriptor_set_out=%s", descSetF.Name()),
		f.Name(),
	)

	stderr := new(bytes.Buffer)
	cmd.Stderr = stderr

	if err = cmd.Run(); err != nil {
		return nil, fmt.Errorf("protoc failed with %v and Stderr %q", err, stderr.String())
	}

	descSet, err := ioutil.ReadFile(descSetF.Name())
	if err != nil {
		return nil, err
	}

	protoset := &descriptorpb.FileDescriptorSet{}
	if err := proto.Unmarshal(descSet, protoset); err != nil {
		return nil, err
	}

	if len(protoset.GetFile()) == 0 {
		return nil, fmt.Errorf("protoset.GetFile() returns empty list")
	}

	return protoset.GetFile()[0], nil
}

func mustCloseAndRemoveFile(f *os.File) {
	if err := f.Close(); err != nil {
		log.Fatalf("Error closing proto file: %v", err)
	}

	if err := os.Remove(f.Name()); err != nil {
		log.Fatalf("Error removing proto file: %v", err)
	}
}

// MustCreateTemplate creates a template with name "test" from
// the provided template string.
func MustCreateTemplate(tmpl string) *template.Template {
	return template.Must(template.New("test").Parse(tmpl))
}

// MustCreateRequestFromTemplate creates a lint.Request from the provided template and test data.
func MustCreateRequestFromTemplate(tmpl *template.Template, testData interface{}) lint.Request {
	b := new(bytes.Buffer)
	if err := tmpl.Execute(b, testData); err != nil {
		log.Fatalf("Error executing template %v", err)
	}

	pd, err := descriptorProtoFromSource(b)
	if err != nil {
		log.Fatalf("Error generating proto descriptor: %v", err)
	}

	req, err := lint.NewProtoFileRequest(pd)
	if err != nil {
		log.Fatalf("Error creating proto file request: %v", err)
	}
	return req
}
<|endoftext|>


<commit_before>package rest

import (
	"fmt"
	"sort"
	"strings"

	"github.com/lfq7413/tomato/config"
	"github.com/lfq7413/tomato/orm"
)

// Query ...
type Query struct {
	auth        *Auth
	className   string
	where       map[string]interface{}
	findOptions map[string]interface{}
	response    map[string]interface{}
	doCount     bool
	include     [][]string
	keys        []string
}

// NewQuery ...
func NewQuery(
	auth *Auth,
	className string,
	where map[string]interface{},
	options map[string]interface{},
) *Query {
	query := &Query{
		auth:        auth,
		className:   className,
		where:       where,
		findOptions: map[string]interface{}{},
		response:    nil,
		doCount:     false,
		include:     [][]string{},
		keys:        []string{},
	}

	if auth.IsMaster == false {
		if auth.User != nil {
			query.findOptions["acl"] = []string{auth.User.ID}
		} else {
			query.findOptions["acl"] = nil
		}
		if className == "_Session" {
			if query.findOptions["acl"] == nil {
				// TODO session 无效
			}
			user := map[string]string{"__type": "Pointer", "className": "_User", "objectId": auth.User.ID}
			and := []interface{}{where, user}
			query.where = map[string]interface{}{"$and": and}
		}
	}

	for k, v := range options {
		switch k {
		case "keys":
			if s, ok := v.(string); ok {
				query.keys = strings.Split(s, ",")
				query.keys = append(query.keys, "objectId", "createdAt", "updatedAt")
			}
		case "count":
			query.doCount = true
		case "skip":
			query.findOptions["skip"] = v
		case "limit":
			query.findOptions["limit"] = v
		case "order":
			if s, ok := v.(string); ok {
				fields := strings.Split(s, ",")
				sortMap := map[string]int{}
				for _, v := range fields {
					if strings.HasPrefix(v, "-") {
						sortMap[v[1:]] = -1
					} else {
						sortMap[v] = 1
					}
				}
				query.findOptions["sort"] = sortMap
			}
		case "include":
			if s, ok := v.(string); ok { // v = "user.session,name.friend"
				paths := strings.Split(s, ",") // paths = ["user.session","name.friend"]
				pathSet := []string{}
				for _, path := range paths {
					parts := strings.Split(path, ".") // parts = ["user","session"]
					for lenght := 1; lenght <= len(parts); lenght++ {
						pathSet = append(pathSet, strings.Join(parts[0:lenght], "."))
					} // pathSet = ["user","user.session"]
				} // pathSet = ["user","user.session","name","name.friend"]
				sort.Strings(pathSet) // pathSet = ["name","name.friend","user","user.session"]
				for _, set := range pathSet {
					query.include = append(query.include, strings.Split(set, "."))
				} // query.include = [["name"],["name","friend"],["user"],["user","seeeion"]]
			}
		default:
		}
	}

	return query
}

// Execute ...
func (q *Query) Execute() map[string]interface{} {

	fmt.Println("keys       ", q.keys)
	fmt.Println("doCount    ", q.doCount)
	fmt.Println("findOptions", q.findOptions)
	fmt.Println("include    ", q.include)

	q.buildRestWhere()
	q.runFind()
	q.runCount()
	q.handleInclude()
	return q.response
}

func (q *Query) buildRestWhere() error {
	q.getUserAndRoleACL()
	q.validateClientClassCreation()
	q.replaceSelect()
	q.replaceDontSelect()
	q.replaceInQuery()
	q.replaceNotInQuery()
	return nil
}

func (q *Query) getUserAndRoleACL() error {
	if q.auth.IsMaster || q.auth.User == nil {
		return nil
	}
	roles := q.auth.GetUserRoles()
	roles = append(roles, q.auth.User.ID)
	q.findOptions["acl"] = roles
	return nil
}

func (q *Query) validateClientClassCreation() error {
	sysClass := []string{"_User", "_Installation", "_Role", "_Session", "_Product"}
	if config.TConfig.AllowClientClassCreation {
		return nil
	}
	if q.auth.IsMaster {
		return nil
	}
	for _, v := range sysClass {
		if v == q.className {
			return nil
		}
	}
	if orm.CollectionExists(q.className) {
		return nil
	}
	// TODO 无法操作不存在的表
	return nil
}

func (q *Query) replaceSelect() error {
	return nil
}

func (q *Query) replaceDontSelect() error {
	return nil
}

func (q *Query) replaceInQuery() error {
	return nil
}

func (q *Query) replaceNotInQuery() error {
	return nil
}

func (q *Query) runFind() error {
	return nil
}

func (q *Query) runCount() error {
	return nil
}

func (q *Query) handleInclude() error {
	return nil
}
<commit_msg>完成findObjectWithKey<commit_after>package rest

import (
	"fmt"
	"sort"
	"strings"

	"github.com/lfq7413/tomato/config"
	"github.com/lfq7413/tomato/orm"
	"github.com/lfq7413/tomato/utils"
)

// Query ...
type Query struct {
	auth        *Auth
	className   string
	where       map[string]interface{}
	findOptions map[string]interface{}
	response    map[string]interface{}
	doCount     bool
	include     [][]string
	keys        []string
}

// NewQuery ...
func NewQuery(
	auth *Auth,
	className string,
	where map[string]interface{},
	options map[string]interface{},
) *Query {
	query := &Query{
		auth:        auth,
		className:   className,
		where:       where,
		findOptions: map[string]interface{}{},
		response:    nil,
		doCount:     false,
		include:     [][]string{},
		keys:        []string{},
	}

	if auth.IsMaster == false {
		if auth.User != nil {
			query.findOptions["acl"] = []string{auth.User.ID}
		} else {
			query.findOptions["acl"] = nil
		}
		if className == "_Session" {
			if query.findOptions["acl"] == nil {
				// TODO session 无效
			}
			user := map[string]string{"__type": "Pointer", "className": "_User", "objectId": auth.User.ID}
			and := []interface{}{where, user}
			query.where = map[string]interface{}{"$and": and}
		}
	}

	for k, v := range options {
		switch k {
		case "keys":
			if s, ok := v.(string); ok {
				query.keys = strings.Split(s, ",")
				query.keys = append(query.keys, "objectId", "createdAt", "updatedAt")
			}
		case "count":
			query.doCount = true
		case "skip":
			query.findOptions["skip"] = v
		case "limit":
			query.findOptions["limit"] = v
		case "order":
			if s, ok := v.(string); ok {
				fields := strings.Split(s, ",")
				sortMap := map[string]int{}
				for _, v := range fields {
					if strings.HasPrefix(v, "-") {
						sortMap[v[1:]] = -1
					} else {
						sortMap[v] = 1
					}
				}
				query.findOptions["sort"] = sortMap
			}
		case "include":
			if s, ok := v.(string); ok { // v = "user.session,name.friend"
				paths := strings.Split(s, ",") // paths = ["user.session","name.friend"]
				pathSet := []string{}
				for _, path := range paths {
					parts := strings.Split(path, ".") // parts = ["user","session"]
					for lenght := 1; lenght <= len(parts); lenght++ {
						pathSet = append(pathSet, strings.Join(parts[0:lenght], "."))
					} // pathSet = ["user","user.session"]
				} // pathSet = ["user","user.session","name","name.friend"]
				sort.Strings(pathSet) // pathSet = ["name","name.friend","user","user.session"]
				for _, set := range pathSet {
					query.include = append(query.include, strings.Split(set, "."))
				} // query.include = [["name"],["name","friend"],["user"],["user","seeeion"]]
			}
		default:
		}
	}

	return query
}

// Execute ...
func (q *Query) Execute() map[string]interface{} {

	fmt.Println("keys       ", q.keys)
	fmt.Println("doCount    ", q.doCount)
	fmt.Println("findOptions", q.findOptions)
	fmt.Println("include    ", q.include)

	q.buildRestWhere()
	q.runFind()
	q.runCount()
	q.handleInclude()
	return q.response
}

func (q *Query) buildRestWhere() error {
	q.getUserAndRoleACL()
	q.validateClientClassCreation()
	q.replaceSelect()
	q.replaceDontSelect()
	q.replaceInQuery()
	q.replaceNotInQuery()
	return nil
}

func (q *Query) getUserAndRoleACL() error {
	if q.auth.IsMaster || q.auth.User == nil {
		return nil
	}
	roles := q.auth.GetUserRoles()
	roles = append(roles, q.auth.User.ID)
	q.findOptions["acl"] = roles
	return nil
}

func (q *Query) validateClientClassCreation() error {
	sysClass := []string{"_User", "_Installation", "_Role", "_Session", "_Product"}
	if config.TConfig.AllowClientClassCreation {
		return nil
	}
	if q.auth.IsMaster {
		return nil
	}
	for _, v := range sysClass {
		if v == q.className {
			return nil
		}
	}
	if orm.CollectionExists(q.className) {
		return nil
	}
	// TODO 无法操作不存在的表
	return nil
}

func (q *Query) replaceSelect() error {
	return nil
}

func (q *Query) replaceDontSelect() error {
	return nil
}

func (q *Query) replaceInQuery() error {
	return nil
}

func (q *Query) replaceNotInQuery() error {
	return nil
}

func (q *Query) runFind() error {
	return nil
}

func (q *Query) runCount() error {
	return nil
}

func (q *Query) handleInclude() error {
	return nil
}

// 查找带有指定 key 的对象
func findObjectWithKey(root interface{}, key string) map[string]interface{} {
	if s := utils.SliceInterface(root); s != nil {
		for _, v := range s {
			answer := findObjectWithKey(v, key)
			if answer != nil {
				return answer
			}
		}
	}
	if m := utils.MapInterface(root); m != nil {
		if m[key] != nil {
			return m
		}
		for _, v := range m {
			answer := findObjectWithKey(v, key)
			if answer != nil {
				return answer
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package retry

import (
	"errors"
	"reflect"
	"strings"
	"testing"
	"time"
)

func TestDo(t *testing.T) {
	t.Parallel()

	err := Do(happy)
	if err != nil {
		t.Errorf("should have succeeded without an error, got %v", err)
	}

	err = Do(sad, Tries(5))
	if err == nil {
		t.Errorf("should have failed with an error")
	}
}

func TestDefaultNew(t *testing.T) {
	t.Parallel()

	r := New()
	if r.Tries != MaxRetries {
		t.Fatalf("bad tries config, got %d want %d", r.Tries, MaxRetries)
	}

	if r.Tries != MaxRetries {
		t.Fatalf("bad tries config, got %d want %d", r.Tries, MaxRetries)
	}

	err := r.Do(happy)
	if err != nil {
		t.Errorf("should have succeeded without an error, got %v. Retryer state %#v", err, r)
	}

	err = r.Do(sad)
	if err == nil {
		t.Errorf("should have failed with an error. Retryer state %#v", r)
	}
}

func TestInfNumberOfTries(t *testing.T) {
	t.Parallel()

	r := New(Tries(0))
	if r.Tries != 0 {
		t.Fatalf("bad tries config, got %d want %d", r.Tries, 0)
	}

	for i := 0; i <= 1000; i++ {
		err := r.Do(sad)
		if err == nil {
			t.Errorf("should have failed with an error")
		}
	}
}

func TestSleep(t *testing.T) {
	t.Parallel()

	// using a ever failing function
	tcs := []struct {
		sleep int
		tries int
		wait  int
	}{
		{
			sleep: 50,
			tries: 1,
			wait:  100,
		},
		{
			sleep: 50,
			tries: 5,
			wait:  100,
		},
	}

	for i, tc := range tcs {
		r := New(Sleep(tc.sleep), Tries(tc.tries))
		ch := make(chan error)

		start := time.Now()
		go func() {
			ch <- r.Do(sad)
		}()

		select {
		case <-time.After(time.Duration((tc.sleep*tc.tries)+tc.wait) * time.Millisecond):
			t.Errorf("tc %d: should have slept only for %v took too long", i, time.Duration((tc.sleep*tc.tries)+tc.wait)*time.Millisecond)
		case err := <-ch:
			// have we finished sooner?
			if d := time.Since(start); d < time.Duration(tc.sleep*tc.tries)*time.Millisecond {
				t.Errorf("tc %d: retryer didn't sleep for the desired time, ended after %v", i, d)
			}
			if err == nil {
				t.Errorf("tc %d: should have failed with an error, Retryer state %#v", i, r)
			}
		}
	}
}

func TestSleepFn(t *testing.T) {
	t.Parallel()

	// we want to sleep for 100+200+300 ms = 600 ms -> linearly growing backoff
	sleepFn := func(attempts int) {
		sleep := time.Duration(100*attempts) * time.Millisecond
		time.Sleep(sleep)
	}

	r := New(SleepFn(sleepFn), Tries(3))
	ch := make(chan error)
	start := time.Now()
	go func() {
		ch <- r.Do(sad)
	}()

	expectedSleepDur := time.Duration(600 * time.Millisecond)
	select {
	case <-time.After(expectedSleepDur + 100*time.Millisecond):
		t.Errorf("should have slept only for %v took too long", expectedSleepDur)
	case err := <-ch:
		// have we finished sooner?
		if d := time.Since(start); d < expectedSleepDur {
			t.Errorf("retryer didn't sleep for the desired time, ended after %v", d)
		}
		if err == nil {
			t.Errorf("should have failed with an error, Retryer state %#v", r)
		}
	}
}

func TestPanicRecoveryEnabled(t *testing.T) {
	defer func() {
		if r := recover(); r != nil {
			t.Errorf("retryer with default panic recovery option - Recover - shouldn't have panicked: %v", r)
		}
	}()

	// testing retryer with panic recovery
	err := New(Recover()).Do(panicked)
	if err == nil {
		t.Error("expected an error containing panic stacktracke")
	}
	if !strings.HasPrefix(err.Error(), "retryer has recovered panic: explicit trigger of panic goroutine") ||
		!strings.Contains(err.Error(), "stack") {
		t.Errorf("unexpected error returned from panic recovery %v", err)
	}
}

func TestPanicRecoveryDisabled(t *testing.T) {
	t.Parallel()

	defer func() {
		if r := recover(); r == nil {
			t.Errorf("retryer should have panicked with default settings")
		}
	}()

	New().Do(panicked)
}
func TestEnsureFn(t *testing.T) {
	t.Parallel()

	touched := false
	toggler := func(error) { touched = true }

	r := New(Ensure(toggler))
	err := r.Do(sad)
	if err == nil {
		t.Errorf("should have succeeded without an error, got %v. Retryer state %#v", err, r)
	}
	if !touched {
		t.Error("ensure function wasn't called")
	}
}

func TestErrorFnOn(t *testing.T) {
	t.Parallel()

	r := New(On([]error{errorTypeA{}, errorTypeB{}}))
	err := r.Do(sad)
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}

	// errorTypeC is not amongst the ones to retry on, we try only once
	fn := func() error { return errorTypeC{S: "error c triggered"} }
	err = New(Tries(1), On([]error{errorTypeA{}, errorTypeC{}})).Do(fn)
	if err == nil {
		t.Errorf("expected errorType error")
	}

	// errorTypeC is not amongst the ones to retry on, as we supply empty slice of errors, we try only once
	fn = func() error { return errorTypeC{S: "error c triggered"} }
	err = New(Tries(1), On([]error{})).Do(fn)
	if err == nil {
		t.Errorf("expected errorType error")
	}
}

func TestErrorFnNot(t *testing.T) {
	t.Parallel()

	// errorTypeC is amongst the ones to ignore in the not slice, we expect a success after 1st run
	errMsg := "error c triggered"
	fn := func() error { return errorTypeC{S: errMsg} }

	r := New(Tries(3), Not([]error{errorTypeA{}, errorTypeC{}}))
	err := r.Do(fn)
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}
	if r.Attempts() != 1 {
		t.Errorf("incorrect attempts count, got %d want 1", r.attempts)
	}

	// errorTypeC is not amongst the ones to retry on, we try only once
	fn = func() error { return errorTypeC{S: "error c triggered"} }
	err = New(Tries(1), Not([]error{errorTypeA{}, errorTypeB{}})).Do(fn)
	if err == nil {
		t.Errorf("expected errorType error")
	}
	if !strings.Contains(err.Error(), errMsg) {
		t.Errorf("unexpected error returned, got: type:%v msg:'%v',  want it to contain: type:%v msg:'%v'", reflect.TypeOf(err), err.Error(), "errorTypeC", errMsg)
	}
}

func TestAfterEachFail(t *testing.T) {
	t.Parallel()

	ch := make(chan error)
	fn := func(err error) { ch <- err }

	bodyFn := func() error { return errors.New("generic error") }
	ab := attemptsBased{
		succeedOnNth: 5,
		fn:           bodyFn,
	}

	go func() {
		New(AfterEachFail(fn), Tries(5)).Do(ab.run)
	}()

	for i := 0; i < 5; i++ {
		select {
		case <-ch:
			continue
		case <-time.After(50 * time.Millisecond):
			t.Fatal("didn't receive an error from the fail callback function")
		}
	}
}

func TestCombinedOptions(t *testing.T) {
	t.Parallel()

	ab := attemptsBased{
		succeedOnNth: 3,
		fn:           sad,
	}
	touched := false
	toggler := func(error) { touched = true }

	r := New(Tries(5), Sleep(10), Ensure(toggler))
	err := r.Do(ab.run)
	if err != nil {
		t.Errorf("got unexpected error: %v", err)
	}
	if !touched {
		t.Error("ensure function wasn't called")
	}
}

func TestSleepFnPriorityOverSleep(t *testing.T) {
	t.Parallel()

	// we want to sleep for 100+200+300 ms = 600 ms -> linearly growing backoff
	sleepFn := func(attempts int) {
		sleep := time.Duration(100*attempts) * time.Millisecond
		time.Sleep(sleep)
	}

	// the Sleep(1000) won't be used, if yes, it will be caught by the timers below
	r := New(SleepFn(sleepFn), Sleep(1000), Tries(3))
	ch := make(chan error)
	start := time.Now()
	go func() {
		ch <- r.Do(sad)
	}()

	expectedSleepDur := time.Duration(600 * time.Millisecond)
	select {
	case <-time.After(expectedSleepDur + 100*time.Millisecond):
		t.Errorf("should have slept only for %v took too long", expectedSleepDur)
	case err := <-ch:
		// have we finished sooner?
		if d := time.Since(start); d < expectedSleepDur {
			t.Errorf("retryer didn't sleep for the desired time, ended after %v", d)
		}
		if err == nil {
			t.Errorf("should have failed with an error, Retryer state %#v", r)
		}
	}
}

type errorTypeA struct {
	s string
}

func (e errorTypeA) Error() string {
	return e.s
}

type errorTypeB struct {
	s string
}

func (e errorTypeB) Error() string {
	return e.s
}

type errorTypeC struct {
	S string
}

func (e errorTypeC) Error() string {
	return e.S
}

func happy() error {
	_ = 2 + 3
	return nil
}

func sad() error {
	_ = 2 + 3
	return errors.New("error on primitive addition")
}

func panicked() error {
	panic("explicit trigger of panic")
}

type attemptsBased struct {
	attempts     int
	succeedOnNth int
	fn           func() error
}

func (ab *attemptsBased) run() error {
	if ab.attempts == ab.succeedOnNth {
		return nil
	}

	ab.attempts++
	return ab.fn()
}
<commit_msg>Paralellize all the tests<commit_after>package retry

import (
	"errors"
	"reflect"
	"strings"
	"testing"
	"time"
)

func TestDo(t *testing.T) {
	t.Parallel()

	err := Do(happy)
	if err != nil {
		t.Errorf("should have succeeded without an error, got %v", err)
	}

	err = Do(sad, Tries(5))
	if err == nil {
		t.Errorf("should have failed with an error")
	}
}

func TestDefaultNew(t *testing.T) {
	t.Parallel()

	r := New()
	if r.Tries != MaxRetries {
		t.Fatalf("bad tries config, got %d want %d", r.Tries, MaxRetries)
	}

	err := r.Do(happy)
	if err != nil {
		t.Errorf("should have succeeded without an error, got %v. Retryer state %#v", err, r)
	}

	err = r.Do(sad)
	if err == nil {
		t.Errorf("should have failed with an error. Retryer state %#v", r)
	}
}

func TestInfNumberOfTries(t *testing.T) {
	t.Parallel()

	r := New(Tries(0))
	if r.Tries != 0 {
		t.Fatalf("bad tries config, got %d want %d", r.Tries, 0)
	}

	for i := 0; i <= 1000; i++ {
		err := r.Do(sad)
		if err == nil {
			t.Errorf("should have failed with an error")
		}
	}
}

func TestSleep(t *testing.T) {
	t.Parallel()

	// using a ever failing function
	tcs := []struct {
		sleep int
		tries int
		wait  int
	}{
		{
			sleep: 50,
			tries: 1,
			wait:  100,
		},
		{
			sleep: 50,
			tries: 5,
			wait:  100,
		},
	}

	for i, tc := range tcs {
		r := New(Sleep(tc.sleep), Tries(tc.tries))
		ch := make(chan error)

		start := time.Now()
		go func() {
			ch <- r.Do(sad)
		}()

		select {
		case <-time.After(time.Duration((tc.sleep*tc.tries)+tc.wait) * time.Millisecond):
			t.Errorf("tc %d: should have slept only for %v took too long", i, time.Duration((tc.sleep*tc.tries)+tc.wait)*time.Millisecond)
		case err := <-ch:
			// have we finished sooner?
			if d := time.Since(start); d < time.Duration(tc.sleep*tc.tries)*time.Millisecond {
				t.Errorf("tc %d: retryer didn't sleep for the desired time, ended after %v", i, d)
			}
			if err == nil {
				t.Errorf("tc %d: should have failed with an error, Retryer state %#v", i, r)
			}
		}
	}
}

func TestSleepFn(t *testing.T) {
	t.Parallel()

	// we want to sleep for 100+200+300 ms = 600 ms -> linearly growing backoff
	sleepFn := func(attempts int) {
		sleep := time.Duration(100*attempts) * time.Millisecond
		time.Sleep(sleep)
	}

	r := New(SleepFn(sleepFn), Tries(3))
	ch := make(chan error)
	start := time.Now()
	go func() {
		ch <- r.Do(sad)
	}()

	expectedSleepDur := time.Duration(600 * time.Millisecond)
	select {
	case <-time.After(expectedSleepDur + 100*time.Millisecond):
		t.Errorf("should have slept only for %v took too long", expectedSleepDur)
	case err := <-ch:
		// have we finished sooner?
		if d := time.Since(start); d < expectedSleepDur {
			t.Errorf("retryer didn't sleep for the desired time, ended after %v", d)
		}
		if err == nil {
			t.Errorf("should have failed with an error, Retryer state %#v", r)
		}
	}
}

func TestPanicRecoveryEnabled(t *testing.T) {
	t.Parallel()

	defer func() {
		if r := recover(); r != nil {
			t.Errorf("retryer with default panic recovery option - Recover - shouldn't have panicked: %v", r)
		}
	}()

	// testing retryer with panic recovery
	err := New(Recover()).Do(panicked)
	if err == nil {
		t.Error("expected an error containing panic stacktracke")
	}
	if !strings.HasPrefix(err.Error(), "retryer has recovered panic: explicit trigger of panic goroutine") ||
		!strings.Contains(err.Error(), "stack") {
		t.Errorf("unexpected error returned from panic recovery %v", err)
	}
}

func TestPanicRecoveryDisabled(t *testing.T) {
	t.Parallel()

	defer func() {
		if r := recover(); r == nil {
			t.Errorf("retryer should have panicked with default settings")
		}
	}()

	New().Do(panicked)
}
func TestEnsureFn(t *testing.T) {
	t.Parallel()

	touched := false
	toggler := func(error) { touched = true }

	r := New(Ensure(toggler))
	err := r.Do(sad)
	if err == nil {
		t.Errorf("should have succeeded without an error, got %v. Retryer state %#v", err, r)
	}
	if !touched {
		t.Error("ensure function wasn't called")
	}
}

func TestErrorFnOn(t *testing.T) {
	t.Parallel()

	r := New(On([]error{errorTypeA{}, errorTypeB{}}))
	err := r.Do(sad)
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}

	// errorTypeC is not amongst the ones to retry on, we try only once
	fn := func() error { return errorTypeC{S: "error c triggered"} }
	err = New(Tries(1), On([]error{errorTypeA{}, errorTypeC{}})).Do(fn)
	if err == nil {
		t.Errorf("expected errorType error")
	}

	// errorTypeC is not amongst the ones to retry on, as we supply empty slice of errors, we try only once
	fn = func() error { return errorTypeC{S: "error c triggered"} }
	err = New(Tries(1), On([]error{})).Do(fn)
	if err == nil {
		t.Errorf("expected errorType error")
	}
}

func TestErrorFnNot(t *testing.T) {
	t.Parallel()

	// errorTypeC is amongst the ones to ignore in the not slice, we expect a success after 1st run
	errMsg := "error c triggered"
	fn := func() error { return errorTypeC{S: errMsg} }

	r := New(Tries(3), Not([]error{errorTypeA{}, errorTypeC{}}))
	err := r.Do(fn)
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}
	if r.Attempts() != 1 {
		t.Errorf("incorrect attempts count, got %d want 1", r.attempts)
	}

	// errorTypeC is not amongst the ones to retry on, we try only once
	fn = func() error { return errorTypeC{S: "error c triggered"} }
	err = New(Tries(1), Not([]error{errorTypeA{}, errorTypeB{}})).Do(fn)
	if err == nil {
		t.Errorf("expected errorType error")
	}
	if !strings.Contains(err.Error(), errMsg) {
		t.Errorf("unexpected error returned, got: type:%v msg:'%v',  want it to contain: type:%v msg:'%v'", reflect.TypeOf(err), err.Error(), "errorTypeC", errMsg)
	}
}

func TestAfterEachFail(t *testing.T) {
	t.Parallel()

	ch := make(chan error)
	fn := func(err error) { ch <- err }

	bodyFn := func() error { return errors.New("generic error") }
	ab := attemptsBased{
		succeedOnNth: 5,
		fn:           bodyFn,
	}

	go func() {
		New(AfterEachFail(fn), Tries(5)).Do(ab.run)
	}()

	for i := 0; i < 5; i++ {
		select {
		case <-ch:
			continue
		case <-time.After(50 * time.Millisecond):
			t.Fatal("didn't receive an error from the fail callback function")
		}
	}
}

func TestCombinedOptions(t *testing.T) {
	t.Parallel()

	ab := attemptsBased{
		succeedOnNth: 3,
		fn:           sad,
	}
	touched := false
	toggler := func(error) { touched = true }

	r := New(Tries(5), Sleep(10), Ensure(toggler))
	err := r.Do(ab.run)
	if err != nil {
		t.Errorf("got unexpected error: %v", err)
	}
	if !touched {
		t.Error("ensure function wasn't called")
	}
}

func TestSleepFnPriorityOverSleep(t *testing.T) {
	t.Parallel()

	// we want to sleep for 100+200+300 ms = 600 ms -> linearly growing backoff
	sleepFn := func(attempts int) {
		sleep := time.Duration(100*attempts) * time.Millisecond
		time.Sleep(sleep)
	}

	// the Sleep(1000) won't be used, if yes, it will be caught by the timers below
	r := New(SleepFn(sleepFn), Sleep(1000), Tries(3))
	ch := make(chan error)
	start := time.Now()
	go func() {
		ch <- r.Do(sad)
	}()

	expectedSleepDur := time.Duration(600 * time.Millisecond)
	select {
	case <-time.After(expectedSleepDur + 100*time.Millisecond):
		t.Errorf("should have slept only for %v took too long", expectedSleepDur)
	case err := <-ch:
		// have we finished sooner?
		if d := time.Since(start); d < expectedSleepDur {
			t.Errorf("retryer didn't sleep for the desired time, ended after %v", d)
		}
		if err == nil {
			t.Errorf("should have failed with an error, Retryer state %#v", r)
		}
	}
}

type errorTypeA struct {
	s string
}

func (e errorTypeA) Error() string {
	return e.s
}

type errorTypeB struct {
	s string
}

func (e errorTypeB) Error() string {
	return e.s
}

type errorTypeC struct {
	S string
}

func (e errorTypeC) Error() string {
	return e.S
}

func happy() error {
	_ = 2 + 3
	return nil
}

func sad() error {
	_ = 2 + 3
	return errors.New("error on primitive addition")
}

func panicked() error {
	panic("explicit trigger of panic")
}

type attemptsBased struct {
	attempts     int
	succeedOnNth int
	fn           func() error
}

func (ab *attemptsBased) run() error {
	if ab.attempts == ab.succeedOnNth {
		return nil
	}

	ab.attempts++
	return ab.fn()
}
<|endoftext|>


<commit_before>package protocol

//TODO: Handle errors

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"reflect"
)

// Represents server list ping response.
type ServerListPing struct {
	Ver  Version       `json:"version"`
	Pl   Players       `json:"players"`
	Desc ChatComponent `json:"description"`
	Fav  string        `json:"favicon,omitempty"`
}

type Version struct {
	Name     string `json:"name"`
	Protocol uint32 `json:"protocol"`
}

type Players struct {
	Max    uint `json:"max"`
	Online uint `json:"online"`
}

type ChatComponent struct {
	Text string `json:"text"`
}

type Response struct {
	data *bytes.Buffer
}

// NewResponse creates a new response.
func NewResponse() *Response {
	return &Response{data: new(bytes.Buffer)}
}

// WriteBoolean writes the given boolean to the current response.
func (r *Response) WriteBoolean(b bool) *Response {
	if b {
		return r.WriteByte(1)
	} else {
		return r.WriteByte(0)
	}
}

// WriteByte writes the given byte to the current response.
func (r *Response) WriteByte(b int8) *Response {
	r.data.Write([]byte{byte(b)})
	return r
}

// WriteUnsignedByte writes the given unsigned byte to the current response.
func (r *Response) WriteUnsignedByte(b uint8) *Response {
	binary.Write(r.data, ByteOrder, b)
	return r
}

// WriteUVarint writes the given UVarint to the current response.
func (r *Response) WriteUVarint(i uint32) *Response {
	_, err := r.data.Write(Uvarint(i))
	if err != nil {
		panic(err)
	}
	return r
}

// WriteVarint writes the given Varint to the current response.
func (r *Response) WriteVarint(i int32) *Response {
	_, err := r.data.Write(Varint(i))
	if err != nil {
		panic(err)
	}
	return r
}

// WriteInt writes the given integer to the current response.
func (r *Response) WriteInt(i int) *Response {
	binary.Write(r.data, ByteOrder, int32(i))
	return r
}

// WriteLong writes the given long to the current response.
func (r *Response) WriteLong(l int64) *Response {
	binary.Write(r.data, ByteOrder, l)
	return r
}

// WriteFloat writes the given float to the current response.
func (r *Response) WriteFloat(f float32) *Response {
	binary.Write(r.data, ByteOrder, f)
	return r
}

// WriteDouble writes the given double to the current response.
func (r *Response) WriteDouble(d float64) *Response {
	binary.Write(r.data, ByteOrder, d)
	return r
}

// WriteByteArray writes the given byte array to the current response.
func (r *Response) WriteByteArray(b []byte) *Response {
	r.WriteUVarint(uint32(len(b)))
	r.data.Write(b)
	return r
}

// WriteArray writes the length of the given array followed by each of
// its components. Ignored if array's length is equal to 0.
func (r *Response) WriteArray(arr []interface{}) *Response {
	if len(arr) == 0 {
		return r
	}

	r.WriteUVarint(uint32(len(arr)))
	for _, elem := range arr {
		r.WriteStructure(elem)
	}

	return r
}

// WriteJSON writes the given interface as a JSON string to the current response.
// (Currently ignores failure.)
func (r *Response) WriteJSON(obj interface{}) *Response {
	j, err := json.Marshal(obj)
	if err != nil {
		panic(err)
	}
	return r.WriteByteArray(j)
}

// WriteStructure explores the field of the given interface, and, if
// their serialization is supported, writes them to the response;
// otherwise, tries to explore the field as an interface too--if
// it fails, skips the field.
func (r *Response) WriteStructure(object interface{}) *Response {
	switch object.(type) {
	case int8:
		r.WriteByte(object.(int8))
	case uint8:
		r.WriteUnsignedByte(object.(byte))
	case bool:
		r.WriteBoolean(object.(bool))
	case int:
		r.WriteInt(object.(int))
	case int32:
		r.WriteVarint(object.(int32))
	case int64:
		r.WriteLong(object.(int64))
	case uint32:
		r.WriteUVarint(object.(uint32))
	case float32:
		r.WriteFloat(object.(float32))
	case float64:
		r.WriteDouble(object.(float64))
	case string:
		r.WriteString(object.(string))
	case []byte:
		r.WriteByteArray(object.([]byte))
	default:
		t := reflect.ValueOf(object)
		if t.CanInterface() {
			for i := 0; i < t.NumField(); i++ {
				r.WriteStructure(t.Field(i).Interface())
			}
		}
	}
	return r
}

// WriteString writes the given string to the current response.
func (r *Response) WriteString(str string) *Response {
	return r.WriteByteArray([]byte(str))
}

// Returns the raw packet created from the written bytes and the provided id.
func (r *Response) ToRawPacket(id uint64) *RawPacket {
	return NewRawPacket(id, r.data.Bytes(), nil)
}

// Clear clears the data from the response's buffer.
func (r *Response) Clear() {
	r.data = new(bytes.Buffer)
}
<commit_msg>:gift: Added JsonString support to WriteStructure.<commit_after>package protocol

//TODO: Handle errors

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"reflect"
)

// Represents server list ping response.
type ServerListPing struct {
	Ver  Version       `json:"version"`
	Pl   Players       `json:"players"`
	Desc ChatComponent `json:"description"`
	Fav  string        `json:"favicon,omitempty"`
}

type Version struct {
	Name     string `json:"name"`
	Protocol uint32 `json:"protocol"`
}

type Players struct {
	Max    uint `json:"max"`
	Online uint `json:"online"`
}

type ChatComponent struct {
	Text string `json:"text"`
}

type Response struct {
	data *bytes.Buffer
}

// NewResponse creates a new response.
func NewResponse() *Response {
	return &Response{data: new(bytes.Buffer)}
}

// WriteBoolean writes the given boolean to the current response.
func (r *Response) WriteBoolean(b bool) *Response {
	if b {
		return r.WriteByte(1)
	} else {
		return r.WriteByte(0)
	}
}

// WriteByte writes the given byte to the current response.
func (r *Response) WriteByte(b int8) *Response {
	r.data.Write([]byte{byte(b)})
	return r
}

// WriteUnsignedByte writes the given unsigned byte to the current response.
func (r *Response) WriteUnsignedByte(b uint8) *Response {
	binary.Write(r.data, ByteOrder, b)
	return r
}

// WriteUVarint writes the given UVarint to the current response.
func (r *Response) WriteUVarint(i uint32) *Response {
	_, err := r.data.Write(Uvarint(i))
	if err != nil {
		panic(err)
	}
	return r
}

// WriteVarint writes the given Varint to the current response.
func (r *Response) WriteVarint(i int32) *Response {
	_, err := r.data.Write(Varint(i))
	if err != nil {
		panic(err)
	}
	return r
}

// WriteInt writes the given integer to the current response.
func (r *Response) WriteInt(i int) *Response {
	binary.Write(r.data, ByteOrder, int32(i))
	return r
}

// WriteLong writes the given long to the current response.
func (r *Response) WriteLong(l int64) *Response {
	binary.Write(r.data, ByteOrder, l)
	return r
}

// WriteFloat writes the given float to the current response.
func (r *Response) WriteFloat(f float32) *Response {
	binary.Write(r.data, ByteOrder, f)
	return r
}

// WriteDouble writes the given double to the current response.
func (r *Response) WriteDouble(d float64) *Response {
	binary.Write(r.data, ByteOrder, d)
	return r
}

// WriteByteArray writes the given byte array to the current response.
func (r *Response) WriteByteArray(b []byte) *Response {
	r.WriteUVarint(uint32(len(b)))
	r.data.Write(b)
	return r
}

// WriteArray writes the length of the given array followed by each of
// its components. Ignored if array's length is equal to 0.
func (r *Response) WriteArray(arr []interface{}) *Response {
	if len(arr) == 0 {
		return r
	}

	r.WriteUVarint(uint32(len(arr)))
	for _, elem := range arr {
		r.WriteStructure(elem)
	}

	return r
}

// WriteJSON writes the given interface as a JSON string to the current response.
// (Currently ignores failure.)
func (r *Response) WriteJSON(obj interface{}) *Response {
	j, err := json.Marshal(obj)
	if err != nil {
		panic(err)
	}
	return r.WriteByteArray(j)
}

// WriteStructure explores the field of the given interface, and, if
// their serialization is supported, writes them to the response;
// otherwise, tries to explore the field as an interface too--if
// it fails, skips the field.
func (r *Response) WriteStructure(object interface{}) *Response {
	switch object.(type) {
	case int8:
		r.WriteByte(object.(int8))
	case uint8:
		r.WriteUnsignedByte(object.(byte))
	case bool:
		r.WriteBoolean(object.(bool))
	case int:
		r.WriteInt(object.(int))
	case int32:
		r.WriteVarint(object.(int32))
	case int64:
		r.WriteLong(object.(int64))
	case uint32:
		r.WriteUVarint(object.(uint32))
	case float32:
		r.WriteFloat(object.(float32))
	case float64:
		r.WriteDouble(object.(float64))
	case string:
		r.WriteString(object.(string))
	case []byte:
		r.WriteByteArray(object.([]byte))
	case JsonString:
		r.WriteJSON(object)
	default:
		t := reflect.ValueOf(object)
		if t.CanInterface() {
			for i := 0; i < t.NumField(); i++ {
				r.WriteStructure(t.Field(i).Interface())
			}
		}
	}
	return r
}

// WriteString writes the given string to the current response.
func (r *Response) WriteString(str string) *Response {
	return r.WriteByteArray([]byte(str))
}

// Returns the raw packet created from the written bytes and the provided id.
func (r *Response) ToRawPacket(id uint64) *RawPacket {
	return NewRawPacket(id, r.data.Bytes(), nil)
}

// Clear clears the data from the response's buffer.
func (r *Response) Clear() {
	r.data = new(bytes.Buffer)
}
<|endoftext|>


<commit_before>package proxy

import (
	"context"
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/fabiolb/fabio/config"
	"github.com/fabiolb/fabio/metrics"
	"github.com/fabiolb/fabio/route"
	grpc_proxy "github.com/mwitkow/grpc-proxy/proxy"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/stats"
	"google.golang.org/grpc/status"
)

type gRPCServer struct {
	server *grpc.Server
}

func (s *gRPCServer) Close() error {
	s.server.Stop()
	return nil
}

func (s *gRPCServer) Shutdown(ctx context.Context) error {
	s.server.GracefulStop()
	return nil
}

func (s *gRPCServer) Serve(lis net.Listener) error {
	return s.server.Serve(lis)
}

func GetGRPCDirector(tlscfg *tls.Config) func(ctx context.Context, fullMethodName string) (context.Context, *grpc.ClientConn, error) {

	connectionPool := newGrpcConnectionPool(tlscfg)

	return func(ctx context.Context, fullMethodName string) (context.Context, *grpc.ClientConn, error) {
		md, ok := metadata.FromIncomingContext(ctx)

		if !ok {
			return ctx, nil, fmt.Errorf("error extracting metadata from request")
		}

		outCtx, _ := context.WithCancel(ctx)
		outCtx = metadata.NewOutgoingContext(outCtx, md.Copy())

		target, _ := ctx.Value(targetKey{}).(*route.Target)

		if target == nil {
			log.Println("[WARN] grpc: no route for ", fullMethodName)
			return outCtx, nil, fmt.Errorf("no route found")
		}

		conn, err := connectionPool.Get(outCtx, target)

		return outCtx, conn, err
	}

}

type GrpcProxyInterceptor struct {
	Config       *config.Config
	StatsHandler *GrpcStatsHandler
	GlobCache    *route.GlobCache
}

type targetKey struct{}

type proxyStream struct {
	grpc.ServerStream
	ctx context.Context
}

func (p proxyStream) Context() context.Context {
	return p.ctx
}

func (g GrpcProxyInterceptor) Stream(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	ctx := stream.Context()

	target, err := g.lookup(ctx, info.FullMethod)

	if err != nil {
		log.Println("[ERROR] grpc: error looking up route", err)
		return status.Error(codes.Internal, "internal error")
	}

	if target == nil {
		g.StatsHandler.NoRoute.Inc(1)
		log.Println("[WARN] grpc: no route found for", info.FullMethod)
		return status.Error(codes.NotFound, "no route found")
	}

	ctx = context.WithValue(ctx, targetKey{}, target)

	proxyStream := proxyStream{
		ServerStream: stream,
		ctx:          ctx,
	}

	start := time.Now()

	err = handler(srv, proxyStream)

	end := time.Now()
	dur := end.Sub(start)

	target.Timer.Update(dur)

	return err
}

func (g GrpcProxyInterceptor) lookup(ctx context.Context, fullMethodName string) (*route.Target, error) {
	pick := route.Picker[g.Config.Proxy.Strategy]
	match := route.Matcher[g.Config.Proxy.Matcher]

	md, ok := metadata.FromIncomingContext(ctx)

	if !ok {
		return nil, fmt.Errorf("error extracting metadata from request")
	}

	reqUrl, err := url.ParseRequestURI(fullMethodName)

	if err != nil {
		log.Print("[WARN] Error parsing grpc request url ", fullMethodName)
		return nil, fmt.Errorf("error parsing request url")
	}

	headers := http.Header{}

	for k, v := range md {
		for _, h := range v {
			headers.Add(k, h)
		}
	}

	req := &http.Request{
		Host:   "",
		URL:    reqUrl,
		Header: headers,
	}

	return route.GetTable().Lookup(req, req.Header.Get("trace"), pick, match, g.GlobCache, g.Config.GlobMatchingDisabled), nil
}

type GrpcStatsHandler struct {
	Connect metrics.Counter
	Request metrics.Timer
	NoRoute metrics.Counter
}

type connCtxKey struct{}
type rpcCtxKey struct{}

func (h *GrpcStatsHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
	return context.WithValue(ctx, connCtxKey{}, info)
}

func (h *GrpcStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
	return context.WithValue(ctx, rpcCtxKey{}, info)
}

func (h *GrpcStatsHandler) HandleRPC(ctx context.Context, rpc stats.RPCStats) {
	rpcStats, _ := rpc.(*stats.End)

	if rpcStats == nil {
		return
	}

	dur := rpcStats.EndTime.Sub(rpcStats.BeginTime)

	h.Request.Update(dur)

	s, _ := status.FromError(rpcStats.Error)
	metrics.DefaultRegistry.GetTimer(fmt.Sprintf("grpc.status.%s", strings.ToLower(s.Code().String())))
}

// HandleConn processes the Conn stats.
func (h *GrpcStatsHandler) HandleConn(ctx context.Context, conn stats.ConnStats) {
	connBegin, _ := conn.(*stats.ConnBegin)

	if connBegin != nil {
		h.Connect.Inc(1)
	}
}

type grpcConnectionPool struct {
	connections     map[*route.Target]*grpc.ClientConn
	lock            sync.RWMutex
	cleanupInterval time.Duration
	tlscfg          *tls.Config
}

func newGrpcConnectionPool(tlscfg *tls.Config) *grpcConnectionPool {
	cp := &grpcConnectionPool{
		connections:     make(map[*route.Target]*grpc.ClientConn),
		lock:            sync.RWMutex{},
		cleanupInterval: time.Second * 5,
		tlscfg:          tlscfg,
	}

	go cp.cleanup()

	return cp
}

func (p *grpcConnectionPool) Get(ctx context.Context, target *route.Target) (*grpc.ClientConn, error) {
	p.lock.RLock()
	conn := p.connections[target]
	p.lock.RUnlock()

	if conn != nil && conn.GetState() != connectivity.Shutdown {
		return conn, nil
	}

	return p.newConnection(ctx, target)
}

func (p *grpcConnectionPool) newConnection(ctx context.Context, target *route.Target) (*grpc.ClientConn, error) {
	opts := []grpc.DialOption{
		grpc.WithDefaultCallOptions(grpc.CallCustomCodec(grpc_proxy.Codec())),
	}

	if target.URL.Scheme == "grpcs" && p.tlscfg != nil {
		opts = append(opts, grpc.WithTransportCredentials(
			credentials.NewTLS(&tls.Config{
				ClientCAs:          p.tlscfg.ClientCAs,
				InsecureSkipVerify: target.TLSSkipVerify,
				// as per the http/2 spec, the host header isn't required, so if your
				// target service doesn't have IP SANs in it's certificate
				// then you will need to override the servername
				ServerName: target.Opts["grpcservername"],
			})))
	} else {
		opts = append(opts, grpc.WithInsecure())
	}

	conn, err := grpc.DialContext(ctx, target.URL.Host, opts...)

	if err == nil {
		p.Set(target, conn)
	}

	return conn, err
}

func (p *grpcConnectionPool) Set(target *route.Target, conn *grpc.ClientConn) {
	p.lock.Lock()
	defer p.lock.Unlock()
	p.connections[target] = conn
}

func (p *grpcConnectionPool) cleanup() {
	for {
		p.lock.Lock()
		table := route.GetTable()
		for target, cs := range p.connections {
			if cs.GetState() == connectivity.Shutdown {
				delete(p.connections, target)
				continue
			}

			if !hasTarget(target, table) {
				log.Println("[DEBUG] grpc: cleaning up connection to", target.URL.Host)
				cs.Close()
				delete(p.connections, target)
			}
		}
		p.lock.Unlock()
		time.Sleep(p.cleanupInterval)
	}
}

func hasTarget(target *route.Target, table route.Table) bool {
	for _, routes := range table {
		for _, r := range routes {
			for _, t := range r.Targets {
				if target == t {
					return true
				}
			}
		}
	}
	return false
}
<commit_msg>fixes #807 - changes map for grpc connections to be a string key<commit_after>package proxy

import (
	"context"
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/fabiolb/fabio/config"
	"github.com/fabiolb/fabio/metrics"
	"github.com/fabiolb/fabio/route"
	grpc_proxy "github.com/mwitkow/grpc-proxy/proxy"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/stats"
	"google.golang.org/grpc/status"
)

type gRPCServer struct {
	server *grpc.Server
}

func (s *gRPCServer) Close() error {
	s.server.Stop()
	return nil
}

func (s *gRPCServer) Shutdown(ctx context.Context) error {
	s.server.GracefulStop()
	return nil
}

func (s *gRPCServer) Serve(lis net.Listener) error {
	return s.server.Serve(lis)
}

func GetGRPCDirector(tlscfg *tls.Config) func(ctx context.Context, fullMethodName string) (context.Context, *grpc.ClientConn, error) {

	connectionPool := newGrpcConnectionPool(tlscfg)

	return func(ctx context.Context, fullMethodName string) (context.Context, *grpc.ClientConn, error) {
		md, ok := metadata.FromIncomingContext(ctx)

		if !ok {
			return ctx, nil, fmt.Errorf("error extracting metadata from request")
		}

		outCtx := metadata.NewOutgoingContext(ctx, md.Copy())

		target, _ := ctx.Value(targetKey{}).(*route.Target)

		if target == nil {
			log.Println("[WARN] grpc: no route for ", fullMethodName)
			return outCtx, nil, fmt.Errorf("no route found")
		}

		conn, err := connectionPool.Get(outCtx, target)

		return outCtx, conn, err
	}

}

type GrpcProxyInterceptor struct {
	Config       *config.Config
	StatsHandler *GrpcStatsHandler
	GlobCache    *route.GlobCache
}

type targetKey struct{}

type proxyStream struct {
	grpc.ServerStream
	ctx context.Context
}

func (p proxyStream) Context() context.Context {
	return p.ctx
}

func makeGRPCTargetKey(t *route.Target) string {
	return t.URL.String()
}

func (g GrpcProxyInterceptor) Stream(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	ctx := stream.Context()

	target, err := g.lookup(ctx, info.FullMethod)

	if err != nil {
		log.Println("[ERROR] grpc: error looking up route", err)
		return status.Error(codes.Internal, "internal error")
	}

	if target == nil {
		g.StatsHandler.NoRoute.Inc(1)
		log.Println("[WARN] grpc: no route found for", info.FullMethod)
		return status.Error(codes.NotFound, "no route found")
	}

	ctx = context.WithValue(ctx, targetKey{}, target)

	proxyStream := proxyStream{
		ServerStream: stream,
		ctx:          ctx,
	}

	start := time.Now()

	err = handler(srv, proxyStream)

	end := time.Now()
	dur := end.Sub(start)

	target.Timer.Update(dur)

	return err
}

func (g GrpcProxyInterceptor) lookup(ctx context.Context, fullMethodName string) (*route.Target, error) {
	pick := route.Picker[g.Config.Proxy.Strategy]
	match := route.Matcher[g.Config.Proxy.Matcher]

	md, ok := metadata.FromIncomingContext(ctx)

	if !ok {
		return nil, fmt.Errorf("error extracting metadata from request")
	}

	reqUrl, err := url.ParseRequestURI(fullMethodName)

	if err != nil {
		log.Print("[WARN] Error parsing grpc request url ", fullMethodName)
		return nil, fmt.Errorf("error parsing request url")
	}

	headers := http.Header{}

	for k, v := range md {
		for _, h := range v {
			headers.Add(k, h)
		}
	}

	req := &http.Request{
		Host:   "",
		URL:    reqUrl,
		Header: headers,
	}

	return route.GetTable().Lookup(req, req.Header.Get("trace"), pick, match, g.GlobCache, g.Config.GlobMatchingDisabled), nil
}

type GrpcStatsHandler struct {
	Connect metrics.Counter
	Request metrics.Timer
	NoRoute metrics.Counter
}

type connCtxKey struct{}
type rpcCtxKey struct{}

func (h *GrpcStatsHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
	return context.WithValue(ctx, connCtxKey{}, info)
}

func (h *GrpcStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
	return context.WithValue(ctx, rpcCtxKey{}, info)
}

func (h *GrpcStatsHandler) HandleRPC(ctx context.Context, rpc stats.RPCStats) {
	rpcStats, _ := rpc.(*stats.End)

	if rpcStats == nil {
		return
	}

	dur := rpcStats.EndTime.Sub(rpcStats.BeginTime)

	h.Request.Update(dur)

	s, _ := status.FromError(rpcStats.Error)
	metrics.DefaultRegistry.GetTimer(fmt.Sprintf("grpc.status.%s", strings.ToLower(s.Code().String())))
}

// HandleConn processes the Conn stats.
func (h *GrpcStatsHandler) HandleConn(ctx context.Context, conn stats.ConnStats) {
	connBegin, _ := conn.(*stats.ConnBegin)

	if connBegin != nil {
		h.Connect.Inc(1)
	}
}

type grpcConnectionPool struct {
	connections     map[string]*grpc.ClientConn
	lock            sync.RWMutex
	cleanupInterval time.Duration
	tlscfg          *tls.Config
}

func newGrpcConnectionPool(tlscfg *tls.Config) *grpcConnectionPool {
	cp := &grpcConnectionPool{
		connections:     make(map[string]*grpc.ClientConn),
		lock:            sync.RWMutex{},
		cleanupInterval: time.Second * 5,
		tlscfg:          tlscfg,
	}

	go cp.cleanup()

	return cp
}

func (p *grpcConnectionPool) Get(ctx context.Context, target *route.Target) (*grpc.ClientConn, error) {
	p.lock.RLock()
	conn := p.connections[makeGRPCTargetKey(target)]
	p.lock.RUnlock()

	if conn != nil && conn.GetState() != connectivity.Shutdown {
		return conn, nil
	}

	return p.newConnection(ctx, target)
}

func (p *grpcConnectionPool) newConnection(ctx context.Context, target *route.Target) (*grpc.ClientConn, error) {
	opts := []grpc.DialOption{
		grpc.WithDefaultCallOptions(grpc.CallCustomCodec(grpc_proxy.Codec())),
	}

	if target.URL.Scheme == "grpcs" && p.tlscfg != nil {
		opts = append(opts, grpc.WithTransportCredentials(
			credentials.NewTLS(&tls.Config{
				ClientCAs:          p.tlscfg.ClientCAs,
				InsecureSkipVerify: target.TLSSkipVerify,
				// as per the http/2 spec, the host header isn't required, so if your
				// target service doesn't have IP SANs in it's certificate
				// then you will need to override the servername
				ServerName: target.Opts["grpcservername"],
			})))
	} else {
		opts = append(opts, grpc.WithInsecure())
	}

	conn, err := grpc.DialContext(ctx, target.URL.Host, opts...)

	if err == nil {
		p.Set(target, conn)
	}

	return conn, err
}

func (p *grpcConnectionPool) Set(target *route.Target, conn *grpc.ClientConn) {
	p.lock.Lock()
	defer p.lock.Unlock()

	p.connections[makeGRPCTargetKey(target)] = conn
}

func (p *grpcConnectionPool) cleanup() {
	for {
		p.lock.Lock()
		table := route.GetTable()
		for tKey, cs := range p.connections {
			if cs.GetState() == connectivity.Shutdown {
				delete(p.connections, tKey)
				continue
			}

			if !hasTarget(tKey, table) {
				log.Println("[DEBUG] grpc: cleaning up connection to", tKey)
				cs.Close()
				delete(p.connections, tKey)
			}
		}
		p.lock.Unlock()
		time.Sleep(p.cleanupInterval)
	}
}

func hasTarget(tKey string, table route.Table) bool {
	for _, routes := range table {
		for _, r := range routes {
			for _, t := range r.Targets {
				if tKey == makeGRPCTargetKey(t) {
					return true
				}
			}
		}
	}
	return false
}
<|endoftext|>