<commit_before>package config

import (
	"fmt"
	"io"
	"path/filepath"
)

// configurable is an interface that must be implemented by any configuration
// formats of Terraform in order to return a *Config.
type configurable interface {
	Config() (*Config, error)
}

// importTree is the result of the first-pass load of the configuration
// files. It is a tree of raw configurables and then any children (their
// imports).
//
// An importTree can be turned into a configTree.
type importTree struct {
	Path     string
	Raw      configurable
	Children []*importTree
}

// This is the function type that must be implemented by the configuration
// file loader to turn a single file into a configurable and any additional
// imports.
type fileLoaderFunc func(path string) (configurable, []string, error)

// loadTree takes a single file and loads the entire importTree for that
// file. This function detects what kind of configuration file it is an
// executes the proper fileLoaderFunc.
func loadTree(root string) (*importTree, error) {
	var f fileLoaderFunc
	switch filepath.Ext(root) {
	case ".tf":
		f = loadFileLibucl
	default:
		return nil, fmt.Errorf(
			"%s: uknown configuration format. Use '.tf' or '.tf.rb' extension",
			root)
	}

	c, imps, err := f(root)
	if err != nil {
		return nil, err
	}

	children := make([]*importTree, len(imps))
	for i, imp := range imps {
		t, err := loadTree(imp)
		if err != nil {
			return nil, err
		}

		children[i] = t
	}

	return &importTree{
		Path:     root,
		Raw:      c,
		Children: children,
	}, nil
}

// Close releases any resources we might be holding open for the importTree.
//
// This can safely be called even while ConfigTree results are alive. The
// importTree is not bound to these.
func (t *importTree) Close() error {
	if c, ok := t.Raw.(io.Closer); ok {
		c.Close()
	}
	for _, ct := range t.Children {
		ct.Close()
	}

	return nil
}

// ConfigTree traverses the importTree and turns each node into a *Config
// object, ultimately returning a *configTree.
func (t *importTree) ConfigTree() (*configTree, error) {
	config, err := t.Raw.Config()
	if err != nil {
		return nil, fmt.Errorf(
			"Error loading %s: %s",
			t.Path,
			err)
	}

	// Build our result
	result := &configTree{
		Path:   t.Path,
		Config: config,
	}

	// Build the config trees for the children
	result.Children = make([]*configTree, len(t.Children))
	for i, ct := range t.Children {
		t, err := ct.ConfigTree()
		if err != nil {
			return nil, err
		}

		result.Children[i] = t
	}

	return result, nil
}
<commit_msg>config: reword for ".tf" extension force<commit_after>package config

import (
	"fmt"
	"io"
	"path/filepath"
)

// configurable is an interface that must be implemented by any configuration
// formats of Terraform in order to return a *Config.
type configurable interface {
	Config() (*Config, error)
}

// importTree is the result of the first-pass load of the configuration
// files. It is a tree of raw configurables and then any children (their
// imports).
//
// An importTree can be turned into a configTree.
type importTree struct {
	Path     string
	Raw      configurable
	Children []*importTree
}

// This is the function type that must be implemented by the configuration
// file loader to turn a single file into a configurable and any additional
// imports.
type fileLoaderFunc func(path string) (configurable, []string, error)

// loadTree takes a single file and loads the entire importTree for that
// file. This function detects what kind of configuration file it is an
// executes the proper fileLoaderFunc.
func loadTree(root string) (*importTree, error) {
	var f fileLoaderFunc
	switch filepath.Ext(root) {
	case ".tf":
		f = loadFileLibucl
	default:
		return nil, fmt.Errorf(
			"%s: unknown configuration format. Use '.tf' extension",
			root)
	}

	c, imps, err := f(root)
	if err != nil {
		return nil, err
	}

	children := make([]*importTree, len(imps))
	for i, imp := range imps {
		t, err := loadTree(imp)
		if err != nil {
			return nil, err
		}

		children[i] = t
	}

	return &importTree{
		Path:     root,
		Raw:      c,
		Children: children,
	}, nil
}

// Close releases any resources we might be holding open for the importTree.
//
// This can safely be called even while ConfigTree results are alive. The
// importTree is not bound to these.
func (t *importTree) Close() error {
	if c, ok := t.Raw.(io.Closer); ok {
		c.Close()
	}
	for _, ct := range t.Children {
		ct.Close()
	}

	return nil
}

// ConfigTree traverses the importTree and turns each node into a *Config
// object, ultimately returning a *configTree.
func (t *importTree) ConfigTree() (*configTree, error) {
	config, err := t.Raw.Config()
	if err != nil {
		return nil, fmt.Errorf(
			"Error loading %s: %s",
			t.Path,
			err)
	}

	// Build our result
	result := &configTree{
		Path:   t.Path,
		Config: config,
	}

	// Build the config trees for the children
	result.Children = make([]*configTree, len(t.Children))
	for i, ct := range t.Children {
		t, err := ct.ConfigTree()
		if err != nil {
			return nil, err
		}

		result.Children[i] = t
	}

	return result, nil
}
<|endoftext|>


<commit_before>package vault

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/vault/helper/forwarding"
	"golang.org/x/net/context"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/keepalive"
)

const (
	clusterListenerAcceptDeadline = 500 * time.Millisecond
	heartbeatInterval             = 30 * time.Second
)

// Starts the listeners and servers necessary to handle forwarded requests
func (c *Core) startForwarding() error {
	c.logger.Trace("core: cluster listener setup function")
	defer c.logger.Trace("core: leaving cluster listener setup function")

	// Clean up in case we have transitioned from a client to a server
	c.requestForwardingConnectionLock.Lock()
	c.clearForwardingClients()
	c.requestForwardingConnectionLock.Unlock()

	// Resolve locally to avoid races
	ha := c.ha != nil

	// Get our TLS config
	tlsConfig, err := c.ClusterTLSConfig()
	if err != nil {
		c.logger.Error("core: failed to get tls configuration when starting forwarding", "error", err)
		return err
	}

	// The server supports all of the possible protos
	tlsConfig.NextProtos = []string{"h2", "req_fw_sb-act_v1"}

	// Create our RPC server and register the request handler server
	c.clusterParamsLock.Lock()

	if c.rpcServer != nil {
		c.logger.Warn("core: forwarding rpc server already running")
		return nil
	}

	c.rpcServer = grpc.NewServer(
		grpc.KeepaliveParams(keepalive.ServerParameters{
			Time: 2 * heartbeatInterval,
		}),
	)

	if ha && c.clusterHandler != nil {
		RegisterRequestForwardingServer(c.rpcServer, &forwardedRequestRPCServer{
			core:    c,
			handler: c.clusterHandler,
		})
	}
	c.clusterParamsLock.Unlock()

	// Create the HTTP/2 server that will be shared by both RPC and regular
	// duties. Doing it this way instead of listening via the server and gRPC
	// allows us to re-use the same port via ALPN. We can just tell the server
	// to serve a given conn and which handler to use.
	fws := &http2.Server{}

	// Shutdown coordination logic
	var shutdown uint32
	shutdownWg := &sync.WaitGroup{}

	for _, addr := range c.clusterListenerAddrs {
		shutdownWg.Add(1)

		// Force a local resolution to avoid data races
		laddr := addr

		// Start our listening loop
		go func() {
			defer shutdownWg.Done()

			if c.logger.IsInfo() {
				c.logger.Info("core/startClusterListener: starting listener", "listener_address", laddr)
			}

			// Create a TCP listener. We do this separately and specifically
			// with TCP so that we can set deadlines.
			tcpLn, err := net.ListenTCP("tcp", laddr)
			if err != nil {
				c.logger.Error("core/startClusterListener: error starting listener", "error", err)
				return
			}

			// Wrap the listener with TLS
			tlsLn := tls.NewListener(tcpLn, tlsConfig)
			defer tlsLn.Close()

			if c.logger.IsInfo() {
				c.logger.Info("core/startClusterListener: serving cluster requests", "cluster_listen_address", tlsLn.Addr())
			}

			for {
				if atomic.LoadUint32(&shutdown) > 0 {
					return
				}

				// Set the deadline for the accept call. If it passes we'll get
				// an error, causing us to check the condition at the top
				// again.
				tcpLn.SetDeadline(time.Now().Add(clusterListenerAcceptDeadline))

				// Accept the connection
				conn, err := tlsLn.Accept()
				if conn != nil {
					// Always defer although it may be closed ahead of time
					defer conn.Close()
				}
				if err != nil {
					continue
				}

				// Type assert to TLS connection and handshake to populate the
				// connection state
				tlsConn := conn.(*tls.Conn)
				err = tlsConn.Handshake()
				if err != nil {
					if c.logger.IsDebug() {
						c.logger.Debug("core: error handshaking cluster connection", "error", err)
					}
					if conn != nil {
						conn.Close()
					}
					continue
				}

				switch tlsConn.ConnectionState().NegotiatedProtocol {
				case "req_fw_sb-act_v1":
					if !ha {
						conn.Close()
						continue
					}

					c.logger.Trace("core: got req_fw_sb-act_v1 connection")
					go fws.ServeConn(conn, &http2.ServeConnOpts{
						Handler: c.rpcServer,
					})

				default:
					c.logger.Debug("core: unknown negotiated protocol on cluster port")
					conn.Close()
					continue
				}
			}
		}()
	}

	// This is in its own goroutine so that we don't block the main thread, and
	// thus we use atomic and channels to coordinate
	// However, because you can't query the status of a channel, we set a bool
	// here while we have the state lock to know whether to actually send a
	// shutdown (e.g. whether the channel will block). See issue #2083.
	c.clusterListenersRunning = true
	go func() {
		// If we get told to shut down...
		<-c.clusterListenerShutdownCh

		// Stop the RPC server
		c.logger.Info("core: shutting down forwarding rpc listeners")
		c.clusterParamsLock.Lock()
		c.rpcServer.Stop()
		c.rpcServer = nil
		c.clusterParamsLock.Unlock()
		c.logger.Info("core: forwarding rpc listeners stopped")

		// Set the shutdown flag. This will cause the listeners to shut down
		// within the deadline in clusterListenerAcceptDeadline
		atomic.StoreUint32(&shutdown, 1)

		// Wait for them all to shut down
		shutdownWg.Wait()
		c.logger.Info("core: rpc listeners successfully shut down")

		// Tell the main thread that shutdown is done.
		c.clusterListenerShutdownSuccessCh <- struct{}{}
	}()

	return nil
}

// refreshRequestForwardingConnection ensures that the client/transport are
// alive and that the current active address value matches the most
// recently-known address.
func (c *Core) refreshRequestForwardingConnection(clusterAddr string) error {
	c.logger.Trace("core: refreshing forwarding connection")
	defer c.logger.Trace("core: done refreshing forwarding connection")

	c.requestForwardingConnectionLock.Lock()
	defer c.requestForwardingConnectionLock.Unlock()

	// Clean things up first
	c.clearForwardingClients()

	// If we don't have anything to connect to, just return
	if clusterAddr == "" {
		return nil
	}

	clusterURL, err := url.Parse(clusterAddr)
	if err != nil {
		c.logger.Error("core: error parsing cluster address attempting to refresh forwarding connection", "error", err)
		return err
	}

	// Set up grpc forwarding handling
	// It's not really insecure, but we have to dial manually to get the
	// ALPN header right. It's just "insecure" because GRPC isn't managing
	// the TLS state.
	ctx, cancelFunc := context.WithCancel(context.Background())
	c.rpcClientConn, err = grpc.DialContext(ctx, clusterURL.Host,
		grpc.WithDialer(c.getGRPCDialer("req_fw_sb-act_v1", "", nil)),
		grpc.WithInsecure(), // it's not, we handle it in the dialer
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time: 2 * heartbeatInterval,
		}))
	if err != nil {
		cancelFunc()
		c.logger.Error("core: err setting up forwarding rpc client", "error", err)
		return err
	}
	c.rpcClientConnContext = ctx
	c.rpcClientConnCancelFunc = cancelFunc
	c.rpcForwardingClient = &forwardingClient{
		RequestForwardingClient: NewRequestForwardingClient(c.rpcClientConn),
		core:        c,
		echoTicker:  time.NewTicker(heartbeatInterval),
		echoContext: ctx,
	}
	c.rpcForwardingClient.startHeartbeat()

	return nil
}

func (c *Core) clearForwardingClients() {
	c.logger.Trace("core: clearing forwarding clients")
	defer c.logger.Trace("core: done clearing forwarding clients")

	if c.rpcClientConnCancelFunc != nil {
		c.rpcClientConnCancelFunc()
		c.rpcClientConnCancelFunc = nil
	}
	if c.rpcClientConn != nil {
		c.rpcClientConn.Close()
		c.rpcClientConn = nil
	}

	c.rpcClientConnContext = nil
	c.rpcForwardingClient = nil
}

// ForwardRequest forwards a given request to the active node and returns the
// response.
func (c *Core) ForwardRequest(req *http.Request) (int, http.Header, []byte, error) {
	c.requestForwardingConnectionLock.RLock()
	defer c.requestForwardingConnectionLock.RUnlock()

	if c.rpcForwardingClient == nil {
		return 0, nil, nil, ErrCannotForward
	}

	freq, err := forwarding.GenerateForwardedRequest(req)
	if err != nil {
		c.logger.Error("core: error creating forwarding RPC request", "error", err)
		return 0, nil, nil, fmt.Errorf("error creating forwarding RPC request")
	}
	if freq == nil {
		c.logger.Error("core: got nil forwarding RPC request")
		return 0, nil, nil, fmt.Errorf("got nil forwarding RPC request")
	}
	resp, err := c.rpcForwardingClient.ForwardRequest(c.rpcClientConnContext, freq)
	if err != nil {
		c.logger.Error("core: error during forwarded RPC request", "error", err)
		return 0, nil, nil, fmt.Errorf("error during forwarding RPC request")
	}

	var header http.Header
	if resp.HeaderEntries != nil {
		header = make(http.Header)
		for k, v := range resp.HeaderEntries {
			for _, j := range v.Values {
				header.Add(k, j)
			}
		}
	}

	return int(resp.StatusCode), header, resp.Body, nil
}

// getGRPCDialer is used to return a dialer that has the correct TLS
// configuration. Otherwise gRPC tries to be helpful and stomps all over our
// NextProtos.
func (c *Core) getGRPCDialer(alpnProto, serverName string, caCert *x509.Certificate) func(string, time.Duration) (net.Conn, error) {
	return func(addr string, timeout time.Duration) (net.Conn, error) {
		tlsConfig, err := c.ClusterTLSConfig()
		if err != nil {
			c.logger.Error("core: failed to get tls configuration", "error", err)
			return nil, err
		}
		if serverName != "" {
			tlsConfig.ServerName = serverName
		}
		if caCert != nil {
			pool := x509.NewCertPool()
			pool.AddCert(caCert)
			tlsConfig.RootCAs = pool
			tlsConfig.ClientCAs = pool
		}
		c.logger.Trace("core: creating rpc dialer", "host", tlsConfig.ServerName)

		tlsConfig.NextProtos = []string{alpnProto}
		dialer := &net.Dialer{
			Timeout: timeout,
		}
		return tls.DialWithDialer(dialer, "tcp", addr, tlsConfig)
	}
}

type forwardedRequestRPCServer struct {
	core    *Core
	handler http.Handler
}

func (s *forwardedRequestRPCServer) ForwardRequest(ctx context.Context, freq *forwarding.Request) (*forwarding.Response, error) {
	//s.core.logger.Trace("forwarding: serving rpc forwarded request")

	// Parse an http.Request out of it
	req, err := forwarding.ParseForwardedRequest(freq)
	if err != nil {
		return nil, err
	}

	// A very dummy response writer that doesn't follow normal semantics, just
	// lets you write a status code (last written wins) and a body. But it
	// meets the interface requirements.
	w := forwarding.NewRPCResponseWriter()

	resp := &forwarding.Response{}
	var respSet bool

	runRequest := func() {
		defer func() {
			// Logic here comes mostly from the Go source code
			if err := recover(); err != nil {
				const size = 64 << 10
				buf := make([]byte, size)
				buf = buf[:runtime.Stack(buf, false)]
				resp.StatusCode = 500
				s.core.logger.Error("forwarding: panic serving request for %v: %v\n%s", req.URL.Path, err, buf)
				respSet = true
			}
		}()
		s.handler.ServeHTTP(w, req)
	}
	runRequest()
	if !respSet {
		resp.StatusCode = uint32(w.StatusCode())
		resp.Body = w.Body().Bytes()
	}

	header := w.Header()
	if header != nil {
		resp.HeaderEntries = make(map[string]*forwarding.HeaderEntry, len(header))
		for k, v := range header {
			resp.HeaderEntries[k] = &forwarding.HeaderEntry{
				Values: v,
			}
		}
	}

	return resp, nil
}

func (s *forwardedRequestRPCServer) Echo(ctx context.Context, in *EchoRequest) (*EchoReply, error) {
	if in.ClusterAddr != "" {
		s.core.clusterPeerClusterAddrsCache.Set(in.ClusterAddr, nil, 0)
	}
	return &EchoReply{
		Message: "pong",
	}, nil
}

type forwardingClient struct {
	RequestForwardingClient

	core *Core

	echoTicker  *time.Ticker
	echoContext context.Context
}

// NOTE: we also take advantage of gRPC's keepalive bits, but as we send data
// with these requests it's useful to keep this as well
func (c *forwardingClient) startHeartbeat() {
	go func() {
		tick := func() {
			c.core.stateLock.RLock()
			clusterAddr := c.core.clusterAddr
			c.core.stateLock.RUnlock()

			ctx, cancel := context.WithTimeout(c.echoContext, 2*time.Second)
			resp, err := c.RequestForwardingClient.Echo(ctx, &EchoRequest{
				Message:     "ping",
				ClusterAddr: clusterAddr,
			})
			cancel()
			if err != nil {
				c.core.logger.Debug("forwarding: error sending echo request to active node", "error", err)
				return
			}
			if resp == nil {
				c.core.logger.Debug("forwarding: empty echo response from active node")
				return
			}
			if resp.Message != "pong" {
				c.core.logger.Debug("forwarding: unexpected echo response from active node", "message", resp.Message)
				return
			}
			c.core.logger.Trace("forwarding: successful heartbeat")
		}

		tick()

		for {
			select {
			case <-c.echoContext.Done():
				c.echoTicker.Stop()
				c.core.logger.Trace("forwarding: stopping heartbeating")
				return
			case <-c.echoTicker.C:
				tick()
			}
		}
	}()
}
<commit_msg>Fix error message formatting and response body<commit_after>package vault

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/vault/helper/forwarding"
	"golang.org/x/net/context"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/keepalive"
)

const (
	clusterListenerAcceptDeadline = 500 * time.Millisecond
	heartbeatInterval             = 30 * time.Second
)

// Starts the listeners and servers necessary to handle forwarded requests
func (c *Core) startForwarding() error {
	c.logger.Trace("core: cluster listener setup function")
	defer c.logger.Trace("core: leaving cluster listener setup function")

	// Clean up in case we have transitioned from a client to a server
	c.requestForwardingConnectionLock.Lock()
	c.clearForwardingClients()
	c.requestForwardingConnectionLock.Unlock()

	// Resolve locally to avoid races
	ha := c.ha != nil

	// Get our TLS config
	tlsConfig, err := c.ClusterTLSConfig()
	if err != nil {
		c.logger.Error("core: failed to get tls configuration when starting forwarding", "error", err)
		return err
	}

	// The server supports all of the possible protos
	tlsConfig.NextProtos = []string{"h2", "req_fw_sb-act_v1"}

	// Create our RPC server and register the request handler server
	c.clusterParamsLock.Lock()

	if c.rpcServer != nil {
		c.logger.Warn("core: forwarding rpc server already running")
		return nil
	}

	c.rpcServer = grpc.NewServer(
		grpc.KeepaliveParams(keepalive.ServerParameters{
			Time: 2 * heartbeatInterval,
		}),
	)

	if ha && c.clusterHandler != nil {
		RegisterRequestForwardingServer(c.rpcServer, &forwardedRequestRPCServer{
			core:    c,
			handler: c.clusterHandler,
		})
	}
	c.clusterParamsLock.Unlock()

	// Create the HTTP/2 server that will be shared by both RPC and regular
	// duties. Doing it this way instead of listening via the server and gRPC
	// allows us to re-use the same port via ALPN. We can just tell the server
	// to serve a given conn and which handler to use.
	fws := &http2.Server{}

	// Shutdown coordination logic
	var shutdown uint32
	shutdownWg := &sync.WaitGroup{}

	for _, addr := range c.clusterListenerAddrs {
		shutdownWg.Add(1)

		// Force a local resolution to avoid data races
		laddr := addr

		// Start our listening loop
		go func() {
			defer shutdownWg.Done()

			if c.logger.IsInfo() {
				c.logger.Info("core/startClusterListener: starting listener", "listener_address", laddr)
			}

			// Create a TCP listener. We do this separately and specifically
			// with TCP so that we can set deadlines.
			tcpLn, err := net.ListenTCP("tcp", laddr)
			if err != nil {
				c.logger.Error("core/startClusterListener: error starting listener", "error", err)
				return
			}

			// Wrap the listener with TLS
			tlsLn := tls.NewListener(tcpLn, tlsConfig)
			defer tlsLn.Close()

			if c.logger.IsInfo() {
				c.logger.Info("core/startClusterListener: serving cluster requests", "cluster_listen_address", tlsLn.Addr())
			}

			for {
				if atomic.LoadUint32(&shutdown) > 0 {
					return
				}

				// Set the deadline for the accept call. If it passes we'll get
				// an error, causing us to check the condition at the top
				// again.
				tcpLn.SetDeadline(time.Now().Add(clusterListenerAcceptDeadline))

				// Accept the connection
				conn, err := tlsLn.Accept()
				if conn != nil {
					// Always defer although it may be closed ahead of time
					defer conn.Close()
				}
				if err != nil {
					continue
				}

				// Type assert to TLS connection and handshake to populate the
				// connection state
				tlsConn := conn.(*tls.Conn)
				err = tlsConn.Handshake()
				if err != nil {
					if c.logger.IsDebug() {
						c.logger.Debug("core: error handshaking cluster connection", "error", err)
					}
					if conn != nil {
						conn.Close()
					}
					continue
				}

				switch tlsConn.ConnectionState().NegotiatedProtocol {
				case "req_fw_sb-act_v1":
					if !ha {
						conn.Close()
						continue
					}

					c.logger.Trace("core: got req_fw_sb-act_v1 connection")
					go fws.ServeConn(conn, &http2.ServeConnOpts{
						Handler: c.rpcServer,
					})

				default:
					c.logger.Debug("core: unknown negotiated protocol on cluster port")
					conn.Close()
					continue
				}
			}
		}()
	}

	// This is in its own goroutine so that we don't block the main thread, and
	// thus we use atomic and channels to coordinate
	// However, because you can't query the status of a channel, we set a bool
	// here while we have the state lock to know whether to actually send a
	// shutdown (e.g. whether the channel will block). See issue #2083.
	c.clusterListenersRunning = true
	go func() {
		// If we get told to shut down...
		<-c.clusterListenerShutdownCh

		// Stop the RPC server
		c.logger.Info("core: shutting down forwarding rpc listeners")
		c.clusterParamsLock.Lock()
		c.rpcServer.Stop()
		c.rpcServer = nil
		c.clusterParamsLock.Unlock()
		c.logger.Info("core: forwarding rpc listeners stopped")

		// Set the shutdown flag. This will cause the listeners to shut down
		// within the deadline in clusterListenerAcceptDeadline
		atomic.StoreUint32(&shutdown, 1)

		// Wait for them all to shut down
		shutdownWg.Wait()
		c.logger.Info("core: rpc listeners successfully shut down")

		// Tell the main thread that shutdown is done.
		c.clusterListenerShutdownSuccessCh <- struct{}{}
	}()

	return nil
}

// refreshRequestForwardingConnection ensures that the client/transport are
// alive and that the current active address value matches the most
// recently-known address.
func (c *Core) refreshRequestForwardingConnection(clusterAddr string) error {
	c.logger.Trace("core: refreshing forwarding connection")
	defer c.logger.Trace("core: done refreshing forwarding connection")

	c.requestForwardingConnectionLock.Lock()
	defer c.requestForwardingConnectionLock.Unlock()

	// Clean things up first
	c.clearForwardingClients()

	// If we don't have anything to connect to, just return
	if clusterAddr == "" {
		return nil
	}

	clusterURL, err := url.Parse(clusterAddr)
	if err != nil {
		c.logger.Error("core: error parsing cluster address attempting to refresh forwarding connection", "error", err)
		return err
	}

	// Set up grpc forwarding handling
	// It's not really insecure, but we have to dial manually to get the
	// ALPN header right. It's just "insecure" because GRPC isn't managing
	// the TLS state.
	ctx, cancelFunc := context.WithCancel(context.Background())
	c.rpcClientConn, err = grpc.DialContext(ctx, clusterURL.Host,
		grpc.WithDialer(c.getGRPCDialer("req_fw_sb-act_v1", "", nil)),
		grpc.WithInsecure(), // it's not, we handle it in the dialer
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time: 2 * heartbeatInterval,
		}))
	if err != nil {
		cancelFunc()
		c.logger.Error("core: err setting up forwarding rpc client", "error", err)
		return err
	}
	c.rpcClientConnContext = ctx
	c.rpcClientConnCancelFunc = cancelFunc
	c.rpcForwardingClient = &forwardingClient{
		RequestForwardingClient: NewRequestForwardingClient(c.rpcClientConn),
		core:        c,
		echoTicker:  time.NewTicker(heartbeatInterval),
		echoContext: ctx,
	}
	c.rpcForwardingClient.startHeartbeat()

	return nil
}

func (c *Core) clearForwardingClients() {
	c.logger.Trace("core: clearing forwarding clients")
	defer c.logger.Trace("core: done clearing forwarding clients")

	if c.rpcClientConnCancelFunc != nil {
		c.rpcClientConnCancelFunc()
		c.rpcClientConnCancelFunc = nil
	}
	if c.rpcClientConn != nil {
		c.rpcClientConn.Close()
		c.rpcClientConn = nil
	}

	c.rpcClientConnContext = nil
	c.rpcForwardingClient = nil
}

// ForwardRequest forwards a given request to the active node and returns the
// response.
func (c *Core) ForwardRequest(req *http.Request) (int, http.Header, []byte, error) {
	c.requestForwardingConnectionLock.RLock()
	defer c.requestForwardingConnectionLock.RUnlock()

	if c.rpcForwardingClient == nil {
		return 0, nil, nil, ErrCannotForward
	}

	freq, err := forwarding.GenerateForwardedRequest(req)
	if err != nil {
		c.logger.Error("core: error creating forwarding RPC request", "error", err)
		return 0, nil, nil, fmt.Errorf("error creating forwarding RPC request")
	}
	if freq == nil {
		c.logger.Error("core: got nil forwarding RPC request")
		return 0, nil, nil, fmt.Errorf("got nil forwarding RPC request")
	}
	resp, err := c.rpcForwardingClient.ForwardRequest(c.rpcClientConnContext, freq)
	if err != nil {
		c.logger.Error("core: error during forwarded RPC request", "error", err)
		return 0, nil, nil, fmt.Errorf("error during forwarding RPC request")
	}

	var header http.Header
	if resp.HeaderEntries != nil {
		header = make(http.Header)
		for k, v := range resp.HeaderEntries {
			for _, j := range v.Values {
				header.Add(k, j)
			}
		}
	}

	return int(resp.StatusCode), header, resp.Body, nil
}

// getGRPCDialer is used to return a dialer that has the correct TLS
// configuration. Otherwise gRPC tries to be helpful and stomps all over our
// NextProtos.
func (c *Core) getGRPCDialer(alpnProto, serverName string, caCert *x509.Certificate) func(string, time.Duration) (net.Conn, error) {
	return func(addr string, timeout time.Duration) (net.Conn, error) {
		tlsConfig, err := c.ClusterTLSConfig()
		if err != nil {
			c.logger.Error("core: failed to get tls configuration", "error", err)
			return nil, err
		}
		if serverName != "" {
			tlsConfig.ServerName = serverName
		}
		if caCert != nil {
			pool := x509.NewCertPool()
			pool.AddCert(caCert)
			tlsConfig.RootCAs = pool
			tlsConfig.ClientCAs = pool
		}
		c.logger.Trace("core: creating rpc dialer", "host", tlsConfig.ServerName)

		tlsConfig.NextProtos = []string{alpnProto}
		dialer := &net.Dialer{
			Timeout: timeout,
		}
		return tls.DialWithDialer(dialer, "tcp", addr, tlsConfig)
	}
}

type forwardedRequestRPCServer struct {
	core    *Core
	handler http.Handler
}

func (s *forwardedRequestRPCServer) ForwardRequest(ctx context.Context, freq *forwarding.Request) (*forwarding.Response, error) {
	//s.core.logger.Trace("forwarding: serving rpc forwarded request")

	// Parse an http.Request out of it
	req, err := forwarding.ParseForwardedRequest(freq)
	if err != nil {
		return nil, err
	}

	// A very dummy response writer that doesn't follow normal semantics, just
	// lets you write a status code (last written wins) and a body. But it
	// meets the interface requirements.
	w := forwarding.NewRPCResponseWriter()

	resp := &forwarding.Response{}

	runRequest := func() {
		defer func() {
			// Logic here comes mostly from the Go source code
			if err := recover(); err != nil {
				const size = 64 << 10
				buf := make([]byte, size)
				buf = buf[:runtime.Stack(buf, false)]
				s.core.logger.Error("forwarding: panic serving request", "path", req.URL.Path, "error", err, "stacktrace", buf)
			}
		}()
		s.handler.ServeHTTP(w, req)
	}
	runRequest()
	resp.StatusCode = uint32(w.StatusCode())
	resp.Body = w.Body().Bytes()

	header := w.Header()
	if header != nil {
		resp.HeaderEntries = make(map[string]*forwarding.HeaderEntry, len(header))
		for k, v := range header {
			resp.HeaderEntries[k] = &forwarding.HeaderEntry{
				Values: v,
			}
		}
	}

	return resp, nil
}

func (s *forwardedRequestRPCServer) Echo(ctx context.Context, in *EchoRequest) (*EchoReply, error) {
	if in.ClusterAddr != "" {
		s.core.clusterPeerClusterAddrsCache.Set(in.ClusterAddr, nil, 0)
	}
	return &EchoReply{
		Message: "pong",
	}, nil
}

type forwardingClient struct {
	RequestForwardingClient

	core *Core

	echoTicker  *time.Ticker
	echoContext context.Context
}

// NOTE: we also take advantage of gRPC's keepalive bits, but as we send data
// with these requests it's useful to keep this as well
func (c *forwardingClient) startHeartbeat() {
	go func() {
		tick := func() {
			c.core.stateLock.RLock()
			clusterAddr := c.core.clusterAddr
			c.core.stateLock.RUnlock()

			ctx, cancel := context.WithTimeout(c.echoContext, 2*time.Second)
			resp, err := c.RequestForwardingClient.Echo(ctx, &EchoRequest{
				Message:     "ping",
				ClusterAddr: clusterAddr,
			})
			cancel()
			if err != nil {
				c.core.logger.Debug("forwarding: error sending echo request to active node", "error", err)
				return
			}
			if resp == nil {
				c.core.logger.Debug("forwarding: empty echo response from active node")
				return
			}
			if resp.Message != "pong" {
				c.core.logger.Debug("forwarding: unexpected echo response from active node", "message", resp.Message)
				return
			}
			c.core.logger.Trace("forwarding: successful heartbeat")
		}

		tick()

		for {
			select {
			case <-c.echoContext.Done():
				c.echoTicker.Stop()
				c.core.logger.Trace("forwarding: stopping heartbeating")
				return
			case <-c.echoTicker.C:
				tick()
			}
		}
	}()
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 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 serialize

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"

	"github.com/m3db/m3/src/x/checked"
	xerrors "github.com/m3db/m3/src/x/errors"
	"github.com/m3db/m3/src/x/ident"
)

/*
 * Serialization scheme to combat Thrift's allocation hell.
 *
 * Given Tags (i.e. key-values) this allows the bijective serialization to,
 * and from Tags <--> []byte.
 *
 * Consider example, Tags: {"abc": "defg", "x": "foo"}
 * this translates to:
 * []byte(
 *    MAGIC_MARKER + NUMBER_TAGS
 *                 + LENGTH([]byte("abc"))  + []byte("abc")
 *                 + LENGTH([]byte("defg")) + []byte("abc")
 *                 + LENGTH([]byte("x"))    + []byte("x")
 *                 + LENGTH([]byte("foo"))  + []byte("foo")
 * )
 *
 * Where MAGIC_MARKER/NUMBER_TAGS/LENGTH are maximum 2 bytes.
 */

var (
	// ByteOrder is the byte order used for encoding tags into a byte sequence.
	ByteOrder        binary.ByteOrder = binary.LittleEndian
	headerMagicBytes                  = make([]byte, 2)
)

func init() {
	encodeUInt16(HeaderMagicNumber, headerMagicBytes)
}

var (
	errTagEncoderInUse   = errors.New("encoder already in use")
	errTagLiteralTooLong = errors.New("literal is too long")
	// ErrEmptyTagNameLiteral is an error when encoded tag name is empty.
	ErrEmptyTagNameLiteral = xerrors.NewInvalidParamsError(errors.New("tag name cannot be empty"))
)

type newCheckedBytesFn func([]byte, checked.BytesOptions) checked.Bytes

var defaultNewCheckedBytesFn = checked.NewBytes

type encoder struct {
	buf               *bytes.Buffer
	checkedBytes      checked.Bytes
	staticBuffer      [2]byte
	staticBufferSlice []byte

	opts TagEncoderOptions
	pool TagEncoderPool
}

func newTagEncoder(
	newFn newCheckedBytesFn,
	opts TagEncoderOptions,
	pool TagEncoderPool,
) TagEncoder {
	b := make([]byte, 0, opts.InitialCapacity())
	cb := newFn(nil, nil)
	e := &encoder{
		buf:          bytes.NewBuffer(b),
		checkedBytes: cb,
		opts:         opts,
		pool:         pool,
	}
	e.staticBufferSlice = e.staticBuffer[:]
	return e
}

func (e *encoder) Encode(tags ident.TagIterator) error {
	if e.checkedBytes.NumRef() > 0 {
		return errTagEncoderInUse
	}

	tags.Rewind()
	defer tags.Rewind()

	numTags := tags.Remaining()
	max := int(e.opts.TagSerializationLimits().MaxNumberTags())
	if numTags > max {
		return fmt.Errorf("too many tags to encode (%d), limit is: %d", numTags, max)
	}

	if _, err := e.buf.Write(headerMagicBytes); err != nil {
		e.buf.Reset()
		return err
	}

	if _, err := e.buf.Write(e.encodeUInt16(uint16(numTags))); err != nil {
		e.buf.Reset()
		return err
	}

	for tags.Next() {
		tag := tags.Current()
		if err := e.encodeTag(tag); err != nil {
			e.buf.Reset()
			return err
		}
	}

	if err := tags.Err(); err != nil {
		e.buf.Reset()
		return err
	}

	e.checkedBytes.IncRef()
	e.checkedBytes.Reset(e.buf.Bytes())

	return nil
}

func (e *encoder) Data() (checked.Bytes, bool) {
	if e.checkedBytes.NumRef() == 0 {
		return nil, false
	}
	return e.checkedBytes, true
}

func (e *encoder) Reset() {
	if e.checkedBytes.NumRef() == 0 {
		return
	}
	e.buf.Reset()
	e.checkedBytes.Reset(nil)
	e.checkedBytes.DecRef()
}

func (e *encoder) Finalize() {
	e.Reset()
	p := e.pool
	if p == nil {
		return
	}
	p.Put(e)
}

func (e *encoder) encodeTag(t ident.Tag) error {
	if len(t.Name.Bytes()) == 0 {
		return ErrEmptyTagNameLiteral
	}

	if err := e.encodeID(t.Name); err != nil {
		return err
	}

	return e.encodeID(t.Value)
}

func (e *encoder) encodeID(i ident.ID) error {
	d := i.Bytes()

	max := int(e.opts.TagSerializationLimits().MaxTagLiteralLength())
	if len(d) >= max {
		return errTagLiteralTooLong
	}

	ld := uint16(len(d))
	if _, err := e.buf.Write(e.encodeUInt16(ld)); err != nil {
		return err
	}

	if _, err := e.buf.Write(d); err != nil {
		return err
	}

	return nil
}

func (e *encoder) encodeUInt16(v uint16) []byte {
	// NB(r): Use static buffer on the struct for encoding, otherwise if it's
	// statically defined inline in the function it will escape to heap.
	dest := e.staticBufferSlice[:2]
	return encodeUInt16(v, dest)
}

func encodeUInt16(v uint16, dest []byte) []byte {
	ByteOrder.PutUint16(dest, v)
	return dest
}

func decodeUInt16(b []byte) uint16 {
	return ByteOrder.Uint16(b)
}
<commit_msg>[x/serialize] Change "literal is too long" error to 400 status code (#3790)<commit_after>// Copyright (c) 2018 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 serialize

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"

	"github.com/m3db/m3/src/x/checked"
	xerrors "github.com/m3db/m3/src/x/errors"
	"github.com/m3db/m3/src/x/ident"
)

/*
 * Serialization scheme to combat Thrift's allocation hell.
 *
 * Given Tags (i.e. key-values) this allows the bijective serialization to,
 * and from Tags <--> []byte.
 *
 * Consider example, Tags: {"abc": "defg", "x": "foo"}
 * this translates to:
 * []byte(
 *    MAGIC_MARKER + NUMBER_TAGS
 *                 + LENGTH([]byte("abc"))  + []byte("abc")
 *                 + LENGTH([]byte("defg")) + []byte("abc")
 *                 + LENGTH([]byte("x"))    + []byte("x")
 *                 + LENGTH([]byte("foo"))  + []byte("foo")
 * )
 *
 * Where MAGIC_MARKER/NUMBER_TAGS/LENGTH are maximum 2 bytes.
 */

var (
	// ByteOrder is the byte order used for encoding tags into a byte sequence.
	ByteOrder        binary.ByteOrder = binary.LittleEndian
	headerMagicBytes                  = make([]byte, 2)
)

func init() {
	encodeUInt16(HeaderMagicNumber, headerMagicBytes)
}

var (
	errTagEncoderInUse   = errors.New("encoder already in use")
	errTagLiteralTooLong = xerrors.NewInvalidParamsError(errors.New("literal is too long"))
	// ErrEmptyTagNameLiteral is an error when encoded tag name is empty.
	ErrEmptyTagNameLiteral = xerrors.NewInvalidParamsError(errors.New("tag name cannot be empty"))
)

type newCheckedBytesFn func([]byte, checked.BytesOptions) checked.Bytes

var defaultNewCheckedBytesFn = checked.NewBytes

type encoder struct {
	buf               *bytes.Buffer
	checkedBytes      checked.Bytes
	staticBuffer      [2]byte
	staticBufferSlice []byte

	opts TagEncoderOptions
	pool TagEncoderPool
}

func newTagEncoder(
	newFn newCheckedBytesFn,
	opts TagEncoderOptions,
	pool TagEncoderPool,
) TagEncoder {
	b := make([]byte, 0, opts.InitialCapacity())
	cb := newFn(nil, nil)
	e := &encoder{
		buf:          bytes.NewBuffer(b),
		checkedBytes: cb,
		opts:         opts,
		pool:         pool,
	}
	e.staticBufferSlice = e.staticBuffer[:]
	return e
}

func (e *encoder) Encode(tags ident.TagIterator) error {
	if e.checkedBytes.NumRef() > 0 {
		return errTagEncoderInUse
	}

	tags.Rewind()
	defer tags.Rewind()

	numTags := tags.Remaining()
	max := int(e.opts.TagSerializationLimits().MaxNumberTags())
	if numTags > max {
		return fmt.Errorf("too many tags to encode (%d), limit is: %d", numTags, max)
	}

	if _, err := e.buf.Write(headerMagicBytes); err != nil {
		e.buf.Reset()
		return err
	}

	if _, err := e.buf.Write(e.encodeUInt16(uint16(numTags))); err != nil {
		e.buf.Reset()
		return err
	}

	for tags.Next() {
		tag := tags.Current()
		if err := e.encodeTag(tag); err != nil {
			e.buf.Reset()
			return err
		}
	}

	if err := tags.Err(); err != nil {
		e.buf.Reset()
		return err
	}

	e.checkedBytes.IncRef()
	e.checkedBytes.Reset(e.buf.Bytes())

	return nil
}

func (e *encoder) Data() (checked.Bytes, bool) {
	if e.checkedBytes.NumRef() == 0 {
		return nil, false
	}
	return e.checkedBytes, true
}

func (e *encoder) Reset() {
	if e.checkedBytes.NumRef() == 0 {
		return
	}
	e.buf.Reset()
	e.checkedBytes.Reset(nil)
	e.checkedBytes.DecRef()
}

func (e *encoder) Finalize() {
	e.Reset()
	p := e.pool
	if p == nil {
		return
	}
	p.Put(e)
}

func (e *encoder) encodeTag(t ident.Tag) error {
	if len(t.Name.Bytes()) == 0 {
		return ErrEmptyTagNameLiteral
	}

	if err := e.encodeID(t.Name); err != nil {
		return err
	}

	return e.encodeID(t.Value)
}

func (e *encoder) encodeID(i ident.ID) error {
	d := i.Bytes()

	max := int(e.opts.TagSerializationLimits().MaxTagLiteralLength())
	if len(d) >= max {
		return errTagLiteralTooLong
	}

	ld := uint16(len(d))
	if _, err := e.buf.Write(e.encodeUInt16(ld)); err != nil {
		return err
	}

	if _, err := e.buf.Write(d); err != nil {
		return err
	}

	return nil
}

func (e *encoder) encodeUInt16(v uint16) []byte {
	// NB(r): Use static buffer on the struct for encoding, otherwise if it's
	// statically defined inline in the function it will escape to heap.
	dest := e.staticBufferSlice[:2]
	return encodeUInt16(v, dest)
}

func encodeUInt16(v uint16, dest []byte) []byte {
	ByteOrder.PutUint16(dest, v)
	return dest
}

func decodeUInt16(b []byte) uint16 {
	return ByteOrder.Uint16(b)
}
<|endoftext|>


<commit_before>package tracker

import (
	"errors"
	"fmt"
	"github.com/zeebo/bencode"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
	"xd/lib/common"
	"xd/lib/log"
	"xd/lib/sync"
)

// http tracker
type HttpTracker struct {
	u *url.URL
	// last time we resolved the remote address
	lastResolved time.Time
	// cached network address of tracker
	addr net.Addr
	// how often to resolve network address
	resolveInterval time.Duration
	// currently resolving the address ?
	resolving sync.Mutex
}

// create new http tracker from url
func NewHttpTracker(u *url.URL) *HttpTracker {
	t := &HttpTracker{
		u:               u,
		resolveInterval: time.Hour,
		lastResolved:    time.Unix(0, 0),
	}

	return t
}

func (t *HttpTracker) shouldResolve() bool {
	return t.lastResolved.Add(t.resolveInterval).Before(time.Now())
}

// http compact response
type compactHttpAnnounceResponse struct {
	Peers    interface{} `bencode:"peers"`
	Interval int         `bencode:"interval"`
	Error    string      `bencode:"failure reason"`
}

func (t *HttpTracker) Name() string {
	return t.u.String()
}

// send announce via http request
func (t *HttpTracker) Announce(req *Request) (resp *Response, err error) {
	//if req == nil {
	//	return
	//}
	// http client
	var client http.Client

	client.Transport = &http.Transport{
		Dial: func(_, _ string) (c net.Conn, e error) {
			var a net.Addr
			t.resolving.Lock()
			if t.shouldResolve() {
				var h, p string
				// XXX: hack
				if strings.Index(t.u.Host, ":") == -1 {
					t.u.Host += ":80"
				}
				h, p, e = net.SplitHostPort(t.u.Host)
				if e == nil {
					a, e = req.GetNetwork().Lookup(h, p)
					if e == nil {
						t.addr = a
						t.lastResolved = time.Now()
					}
				}
			} else {
				a = t.addr
			}
			t.resolving.Unlock()
			if e == nil {
				c, e = req.GetNetwork().Dial(a.Network(), a.String())
			}
			return
		},
	}

	resp = new(Response)
	interval := 30
	// build query
	var u *url.URL
	u, err = url.Parse(t.u.String())
	if err == nil {
		v := u.Query()
		n := req.GetNetwork()
		a := n.Addr()
		addr := a.String() + ".i2p"
		v.Add("ip", addr)
		v.Add("info_hash", string(req.Infohash.Bytes()))
		v.Add("peer_id", string(req.PeerID.Bytes()))
		v.Add("port", fmt.Sprintf("%d", req.Port))
		v.Add("numwant", fmt.Sprintf("%d", req.NumWant))
		v.Add("left", fmt.Sprintf("%d", req.Left))
		if req.Event != Nop {
			v.Add("event", req.Event.String())
		}
		v.Add("downloaded", fmt.Sprintf("%d", req.Downloaded))
		v.Add("uploaded", fmt.Sprintf("%d", req.Uploaded))

		// compact response
		if req.Compact || u.Path != "/a" {
			req.Compact = true
			v.Add("compact", "1")
		}
		u.RawQuery = v.Encode()
		var r *http.Response
		log.Debugf("%s announcing", t.Name())
		r, err = client.Get(u.String())
		if err == nil {
			defer r.Body.Close()
			dec := bencode.NewDecoder(r.Body)
			if req.Compact {
				cresp := new(compactHttpAnnounceResponse)
				err = dec.Decode(cresp)
				if err == nil {
					interval = cresp.Interval
					var cpeers string

					_, ok := cresp.Peers.(string)
					if ok {
						cpeers = cresp.Peers.(string)
						l := len(cpeers) / 32
						for l > 0 {
							var p common.Peer
							// TODO: bounds check
							copy(p.Compact[:], cpeers[(l-1)*32:l*32])
							resp.Peers = append(resp.Peers, p)
							l--
						}
					} else {
						fullpeers, ok := cresp.Peers.([]interface{})
						if ok {
							for idx := range fullpeers {
								// XXX: this is horribad :DDDDDDDDD
								var peer map[string]interface{}
								peer, ok = fullpeers[idx].(map[string]interface{})
								if ok {
									var p common.Peer
									p.IP = fmt.Sprintf("%s", peer["ip"])
									resp.Peers = append(resp.Peers, p)
								}
							}
						}
					}

					if len(cresp.Error) > 0 {
						err = errors.New(cresp.Error)
					}
				}
			} else {
				// decode non compact response
				err = dec.Decode(resp)
				interval = resp.Interval
				if len(resp.Error) > 0 {
					err = errors.New(resp.Error)
				}
			}
		}
	}

	if err == nil {
		log.Infof("%s got %d peers for %s", t.Name(), len(resp.Peers), req.Infohash.Hex())
	} else {
		log.Warnf("%s got error while announcing: %s", t.Name(), err)
	}
	if interval == 0 {
		interval = 60
	}
	resp.NextAnnounce = time.Now().Add(time.Second * time.Duration(interval))
	return
}
<commit_msg>use correct value in announce<commit_after>package tracker

import (
	"errors"
	"fmt"
	"github.com/zeebo/bencode"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
	"xd/lib/common"
	"xd/lib/log"
	"xd/lib/sync"
)

// http tracker
type HttpTracker struct {
	u *url.URL
	// last time we resolved the remote address
	lastResolved time.Time
	// cached network address of tracker
	addr net.Addr
	// how often to resolve network address
	resolveInterval time.Duration
	// currently resolving the address ?
	resolving sync.Mutex
}

// create new http tracker from url
func NewHttpTracker(u *url.URL) *HttpTracker {
	t := &HttpTracker{
		u:               u,
		resolveInterval: time.Hour,
		lastResolved:    time.Unix(0, 0),
	}

	return t
}

func (t *HttpTracker) shouldResolve() bool {
	return t.lastResolved.Add(t.resolveInterval).Before(time.Now())
}

// http compact response
type compactHttpAnnounceResponse struct {
	Peers    interface{} `bencode:"peers"`
	Interval int         `bencode:"interval"`
	Error    string      `bencode:"failure reason"`
}

func (t *HttpTracker) Name() string {
	return t.u.String()
}

// send announce via http request
func (t *HttpTracker) Announce(req *Request) (resp *Response, err error) {
	//if req == nil {
	//	return
	//}
	// http client
	var client http.Client

	client.Transport = &http.Transport{
		Dial: func(_, _ string) (c net.Conn, e error) {
			var a net.Addr
			t.resolving.Lock()
			if t.shouldResolve() {
				var h, p string
				// XXX: hack
				if strings.Index(t.u.Host, ":") == -1 {
					t.u.Host += ":80"
				}
				h, p, e = net.SplitHostPort(t.u.Host)
				if e == nil {
					a, e = req.GetNetwork().Lookup(h, p)
					if e == nil {
						t.addr = a
						t.lastResolved = time.Now()
					}
				}
			} else {
				a = t.addr
			}
			t.resolving.Unlock()
			if e == nil {
				c, e = req.GetNetwork().Dial(a.Network(), a.String())
			}
			return
		},
	}

	resp = new(Response)
	interval := 30
	// build query
	var u *url.URL
	u, err = url.Parse(t.u.String())
	if err == nil {
		v := u.Query()
		n := req.GetNetwork()
		a := n.Addr()
		host, _, _ := net.SplitHostPort(a.String())
		addr := host + ".i2p"
		v.Add("ip", addr)
		v.Add("info_hash", string(req.Infohash.Bytes()))
		v.Add("peer_id", string(req.PeerID.Bytes()))
		v.Add("port", fmt.Sprintf("%d", req.Port))
		v.Add("numwant", fmt.Sprintf("%d", req.NumWant))
		v.Add("left", fmt.Sprintf("%d", req.Left))
		if req.Event != Nop {
			v.Add("event", req.Event.String())
		}
		v.Add("downloaded", fmt.Sprintf("%d", req.Downloaded))
		v.Add("uploaded", fmt.Sprintf("%d", req.Uploaded))

		// compact response
		if req.Compact || u.Path != "/a" {
			req.Compact = true
			v.Add("compact", "1")
		}
		u.RawQuery = v.Encode()
		var r *http.Response
		log.Debugf("%s announcing", t.Name())
		r, err = client.Get(u.String())
		if err == nil {
			defer r.Body.Close()
			dec := bencode.NewDecoder(r.Body)
			if req.Compact {
				cresp := new(compactHttpAnnounceResponse)
				err = dec.Decode(cresp)
				if err == nil {
					interval = cresp.Interval
					var cpeers string

					_, ok := cresp.Peers.(string)
					if ok {
						cpeers = cresp.Peers.(string)
						l := len(cpeers) / 32
						for l > 0 {
							var p common.Peer
							// TODO: bounds check
							copy(p.Compact[:], cpeers[(l-1)*32:l*32])
							resp.Peers = append(resp.Peers, p)
							l--
						}
					} else {
						fullpeers, ok := cresp.Peers.([]interface{})
						if ok {
							for idx := range fullpeers {
								// XXX: this is horribad :DDDDDDDDD
								var peer map[string]interface{}
								peer, ok = fullpeers[idx].(map[string]interface{})
								if ok {
									var p common.Peer
									p.IP = fmt.Sprintf("%s", peer["ip"])
									resp.Peers = append(resp.Peers, p)
								}
							}
						}
					}

					if len(cresp.Error) > 0 {
						err = errors.New(cresp.Error)
					}
				}
			} else {
				// decode non compact response
				err = dec.Decode(resp)
				interval = resp.Interval
				if len(resp.Error) > 0 {
					err = errors.New(resp.Error)
				}
			}
		}
	}

	if err == nil {
		log.Infof("%s got %d peers for %s", t.Name(), len(resp.Peers), req.Infohash.Hex())
	} else {
		log.Warnf("%s got error while announcing: %s", t.Name(), err)
	}
	if interval == 0 {
		interval = 60
	}
	resp.NextAnnounce = time.Now().Add(time.Second * time.Duration(interval))
	return
}
<|endoftext|>


<commit_before>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file implements instantiation of generic types
// through substitution of type parameters by type arguments.

package types2

import (
	"cmd/compile/internal/syntax"
	"errors"
	"fmt"
)

// Instantiate instantiates the type typ with the given type arguments targs.
// typ must be a *Named or a *Signature type, and its number of type parameters
// must match the number of provided type arguments. The result is a new,
// instantiated (not parameterized) type of the same kind (either a *Named or a
// *Signature). Any methods attached to a *Named are simply copied; they are
// not instantiated.
//
// If ctxt is non-nil, it may be used to de-dupe the instance against previous
// instances with the same identity.
//
// If verify is set and constraint satisfaction fails, the returned error may
// wrap an *ArgumentError indicating which type argument did not satisfy its
// corresponding type parameter constraint, and why.
//
// TODO(rfindley): change this function to also return an error if lengths of
// tparams and targs do not match.
func Instantiate(ctxt *Context, typ Type, targs []Type, validate bool) (Type, error) {
	inst := (*Checker)(nil).instance(nopos, typ, targs, ctxt)

	var err error
	if validate {
		var tparams []*TypeParam
		switch t := typ.(type) {
		case *Named:
			tparams = t.TypeParams().list()
		case *Signature:
			tparams = t.TypeParams().list()
		}
		if i, err := (*Checker)(nil).verify(nopos, tparams, targs); err != nil {
			return inst, &ArgumentError{i, err}
		}
	}

	return inst, err
}

// instance creates a type or function instance using the given original type
// typ and arguments targs. For Named types the resulting instance will be
// unexpanded.
func (check *Checker) instance(pos syntax.Pos, orig Type, targs []Type, ctxt *Context) (res Type) {
	var h string
	if ctxt != nil {
		h = ctxt.instanceHash(orig, targs)
		// typ may already have been instantiated with identical type arguments. In
		// that case, re-use the existing instance.
		if inst := ctxt.lookup(h, orig, targs); inst != nil {
			return inst
		}
	}

	switch orig := orig.(type) {
	case *Named:
		tname := NewTypeName(pos, orig.obj.pkg, orig.obj.name, nil)
		named := check.newNamed(tname, orig, nil, nil, nil) // underlying, tparams, and methods are set when named is resolved
		named.targs = NewTypeList(targs)
		named.resolver = func(ctxt *Context, n *Named) (*TypeParamList, Type, []*Func) {
			return expandNamed(ctxt, n, pos)
		}
		res = named

	case *Signature:
		tparams := orig.TypeParams()
		if !check.validateTArgLen(pos, tparams.Len(), len(targs)) {
			return Typ[Invalid]
		}
		if tparams.Len() == 0 {
			return orig // nothing to do (minor optimization)
		}
		sig := check.subst(pos, orig, makeSubstMap(tparams.list(), targs), ctxt).(*Signature)
		// If the signature doesn't use its type parameters, subst
		// will not make a copy. In that case, make a copy now (so
		// we can set tparams to nil w/o causing side-effects).
		if sig == orig {
			copy := *sig
			sig = &copy
		}
		// After instantiating a generic signature, it is not generic
		// anymore; we need to set tparams to nil.
		sig.tparams = nil
		res = sig
	default:
		// only types and functions can be generic
		panic(fmt.Sprintf("%v: cannot instantiate %v", pos, orig))
	}

	if ctxt != nil {
		// It's possible that we've lost a race to add named to the context.
		// In this case, use whichever instance is recorded in the context.
		res = ctxt.update(h, orig, targs, res)
	}

	return res
}

// validateTArgLen verifies that the length of targs and tparams matches,
// reporting an error if not. If validation fails and check is nil,
// validateTArgLen panics.
func (check *Checker) validateTArgLen(pos syntax.Pos, ntparams, ntargs int) bool {
	if ntargs != ntparams {
		// TODO(gri) provide better error message
		if check != nil {
			check.errorf(pos, "got %d arguments but %d type parameters", ntargs, ntparams)
			return false
		}
		panic(fmt.Sprintf("%v: got %d arguments but %d type parameters", pos, ntargs, ntparams))
	}
	return true
}

func (check *Checker) verify(pos syntax.Pos, tparams []*TypeParam, targs []Type) (int, error) {
	// TODO(rfindley): it would be great if users could pass in a qualifier here,
	// rather than falling back to verbose qualification. Maybe this can be part
	// of the shared context.
	var qf Qualifier
	if check != nil {
		qf = check.qualifier
	}

	smap := makeSubstMap(tparams, targs)
	for i, tpar := range tparams {
		// The type parameter bound is parameterized with the same type parameters
		// as the instantiated type; before we can use it for bounds checking we
		// need to instantiate it with the type arguments with which we instantiated
		// the parameterized type.
		bound := check.subst(pos, tpar.bound, smap, nil)
		if err := check.implements(targs[i], bound, qf); err != nil {
			return i, err
		}
	}
	return -1, nil
}

// implements checks if V implements T and reports an error if it doesn't.
// If a qualifier is provided, it is used in error formatting.
func (check *Checker) implements(V, T Type, qf Qualifier) error {
	Vu := under(V)
	Tu := under(T)
	if Vu == Typ[Invalid] || Tu == Typ[Invalid] {
		return nil
	}

	errorf := func(format string, args ...interface{}) error {
		return errors.New(sprintf(qf, false, format, args...))
	}

	Ti, _ := Tu.(*Interface)
	if Ti == nil {
		return errorf("%s is not an interface", T)
	}

	// Every type satisfies the empty interface.
	if Ti.Empty() {
		return nil
	}
	// T is not the empty interface (i.e., the type set of T is restricted)

	// An interface V with an empty type set satisfies any interface.
	// (The empty set is a subset of any set.)
	Vi, _ := Vu.(*Interface)
	if Vi != nil && Vi.typeSet().IsEmpty() {
		return nil
	}
	// type set of V is not empty

	// No type with non-empty type set satisfies the empty type set.
	if Ti.typeSet().IsEmpty() {
		return errorf("cannot implement %s (empty type set)", T)
	}

	// If T is comparable, V must be comparable.
	// TODO(gri) the error messages could be better, here
	if Ti.IsComparable() && !Comparable(V) {
		if Vi != nil && Vi.Empty() {
			return errorf("empty interface %s does not implement %s", V, T)
		}
		return errorf("%s does not implement comparable", V)
	}

	// V must implement T (methods)
	// - check only if we have methods
	if Ti.NumMethods() > 0 {
		if m, wrong := check.missingMethod(V, Ti, true); m != nil {
			// TODO(gri) needs to print updated name to avoid major confusion in error message!
			//           (print warning for now)
			// Old warning:
			// check.softErrorf(pos, "%s does not implement %s (warning: name not updated) = %s (missing method %s)", V, T, Ti, m)
			if wrong != nil {
				// TODO(gri) This can still report uninstantiated types which makes the error message
				//           more difficult to read then necessary.
				return errorf("%s does not implement %s: wrong method signature\n\tgot  %s\n\twant %s",
					V, T, wrong, m,
				)
			}
			return errorf("%s does not implement %s (missing method %s)", V, T, m.name)
		}
	}

	// V must also be in the set of types of T, if any.
	// Constraints with empty type sets were already excluded above.
	if !Ti.typeSet().hasTerms() {
		return nil // nothing to do
	}

	// If V is itself an interface, each of its possible types must be in the set
	// of T types (i.e., the V type set must be a subset of the T type set).
	// Interfaces V with empty type sets were already excluded above.
	if Vi != nil {
		if !Vi.typeSet().subsetOf(Ti.typeSet()) {
			// TODO(gri) report which type is missing
			return errorf("%s does not implement %s", V, T)
		}
		return nil
	}

	// Otherwise, V's type must be included in the iface type set.
	if !Ti.typeSet().includes(V) {
		// TODO(gri) report which type is missing
		return errorf("%s does not implement %s", V, T)
	}

	return nil
}
<commit_msg>cmd/compile/internal/types2: return an error from Instantiate on incorrect len(targs)<commit_after>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file implements instantiation of generic types
// through substitution of type parameters by type arguments.

package types2

import (
	"cmd/compile/internal/syntax"
	"errors"
	"fmt"
)

// Instantiate instantiates the type orig with the given type arguments targs.
// orig must be a *Named or a *Signature type. If there is no error, the
// resulting Type is a new, instantiated (not parameterized) type of the same
// kind (either a *Named or a *Signature). Methods attached to a *Named type
// are also instantiated, and associated with a new *Func that has the same
// position as the original method, but nil function scope.
//
// If ctxt is non-nil, it may be used to de-duplicate the instance against
// previous instances with the same identity. As a special case, generic
// *Signature origin types are only considered identical if they are pointer
// equivalent, so that instantiating distinct (but possibly identical)
// signatures will yield different instances.
//
// If validate is set, Instantiate verifies that the number of type arguments
// and parameters match, and that the type arguments satisfy their
// corresponding type constraints. If verification fails, the resulting error
// may wrap an *ArgumentError indicating which type argument did not satisfy
// its corresponding type parameter constraint, and why.
//
// If validate is not set, Instantiate does not verify the type argument count
// or whether the type arguments satisfy their constraints. Instantiate is
// guaranteed to not return an error, but may panic. Specifically, for
// *Signature types, Instantiate will panic immediately if the type argument
// count is incorrect; for *Named types, a panic may occur later inside the
// *Named API.
func Instantiate(ctxt *Context, orig Type, targs []Type, validate bool) (Type, error) {
	if validate {
		var tparams []*TypeParam
		switch t := orig.(type) {
		case *Named:
			tparams = t.TypeParams().list()
		case *Signature:
			tparams = t.TypeParams().list()
		}
		if len(targs) != len(tparams) {
			return nil, fmt.Errorf("got %d type arguments but %s has %d type parameters", len(targs), orig, len(tparams))
		}
		if i, err := (*Checker)(nil).verify(nopos, tparams, targs); err != nil {
			return nil, &ArgumentError{i, err}
		}
	}

	inst := (*Checker)(nil).instance(nopos, orig, targs, ctxt)
	return inst, nil
}

// instance creates a type or function instance using the given original type
// typ and arguments targs. For Named types the resulting instance will be
// unexpanded.
func (check *Checker) instance(pos syntax.Pos, orig Type, targs []Type, ctxt *Context) (res Type) {
	var h string
	if ctxt != nil {
		h = ctxt.instanceHash(orig, targs)
		// typ may already have been instantiated with identical type arguments. In
		// that case, re-use the existing instance.
		if inst := ctxt.lookup(h, orig, targs); inst != nil {
			return inst
		}
	}

	switch orig := orig.(type) {
	case *Named:
		tname := NewTypeName(pos, orig.obj.pkg, orig.obj.name, nil)
		named := check.newNamed(tname, orig, nil, nil, nil) // underlying, tparams, and methods are set when named is resolved
		named.targs = NewTypeList(targs)
		named.resolver = func(ctxt *Context, n *Named) (*TypeParamList, Type, []*Func) {
			return expandNamed(ctxt, n, pos)
		}
		res = named

	case *Signature:
		tparams := orig.TypeParams()
		if !check.validateTArgLen(pos, tparams.Len(), len(targs)) {
			return Typ[Invalid]
		}
		if tparams.Len() == 0 {
			return orig // nothing to do (minor optimization)
		}
		sig := check.subst(pos, orig, makeSubstMap(tparams.list(), targs), ctxt).(*Signature)
		// If the signature doesn't use its type parameters, subst
		// will not make a copy. In that case, make a copy now (so
		// we can set tparams to nil w/o causing side-effects).
		if sig == orig {
			copy := *sig
			sig = &copy
		}
		// After instantiating a generic signature, it is not generic
		// anymore; we need to set tparams to nil.
		sig.tparams = nil
		res = sig
	default:
		// only types and functions can be generic
		panic(fmt.Sprintf("%v: cannot instantiate %v", pos, orig))
	}

	if ctxt != nil {
		// It's possible that we've lost a race to add named to the context.
		// In this case, use whichever instance is recorded in the context.
		res = ctxt.update(h, orig, targs, res)
	}

	return res
}

// validateTArgLen verifies that the length of targs and tparams matches,
// reporting an error if not. If validation fails and check is nil,
// validateTArgLen panics.
func (check *Checker) validateTArgLen(pos syntax.Pos, ntparams, ntargs int) bool {
	if ntargs != ntparams {
		// TODO(gri) provide better error message
		if check != nil {
			check.errorf(pos, "got %d arguments but %d type parameters", ntargs, ntparams)
			return false
		}
		panic(fmt.Sprintf("%v: got %d arguments but %d type parameters", pos, ntargs, ntparams))
	}
	return true
}

func (check *Checker) verify(pos syntax.Pos, tparams []*TypeParam, targs []Type) (int, error) {
	// TODO(rfindley): it would be great if users could pass in a qualifier here,
	// rather than falling back to verbose qualification. Maybe this can be part
	// of the shared context.
	var qf Qualifier
	if check != nil {
		qf = check.qualifier
	}

	smap := makeSubstMap(tparams, targs)
	for i, tpar := range tparams {
		// The type parameter bound is parameterized with the same type parameters
		// as the instantiated type; before we can use it for bounds checking we
		// need to instantiate it with the type arguments with which we instantiated
		// the parameterized type.
		bound := check.subst(pos, tpar.bound, smap, nil)
		if err := check.implements(targs[i], bound, qf); err != nil {
			return i, err
		}
	}
	return -1, nil
}

// implements checks if V implements T and reports an error if it doesn't.
// If a qualifier is provided, it is used in error formatting.
func (check *Checker) implements(V, T Type, qf Qualifier) error {
	Vu := under(V)
	Tu := under(T)
	if Vu == Typ[Invalid] || Tu == Typ[Invalid] {
		return nil
	}

	errorf := func(format string, args ...interface{}) error {
		return errors.New(sprintf(qf, false, format, args...))
	}

	Ti, _ := Tu.(*Interface)
	if Ti == nil {
		return errorf("%s is not an interface", T)
	}

	// Every type satisfies the empty interface.
	if Ti.Empty() {
		return nil
	}
	// T is not the empty interface (i.e., the type set of T is restricted)

	// An interface V with an empty type set satisfies any interface.
	// (The empty set is a subset of any set.)
	Vi, _ := Vu.(*Interface)
	if Vi != nil && Vi.typeSet().IsEmpty() {
		return nil
	}
	// type set of V is not empty

	// No type with non-empty type set satisfies the empty type set.
	if Ti.typeSet().IsEmpty() {
		return errorf("cannot implement %s (empty type set)", T)
	}

	// If T is comparable, V must be comparable.
	// TODO(gri) the error messages could be better, here
	if Ti.IsComparable() && !Comparable(V) {
		if Vi != nil && Vi.Empty() {
			return errorf("empty interface %s does not implement %s", V, T)
		}
		return errorf("%s does not implement comparable", V)
	}

	// V must implement T (methods)
	// - check only if we have methods
	if Ti.NumMethods() > 0 {
		if m, wrong := check.missingMethod(V, Ti, true); m != nil {
			// TODO(gri) needs to print updated name to avoid major confusion in error message!
			//           (print warning for now)
			// Old warning:
			// check.softErrorf(pos, "%s does not implement %s (warning: name not updated) = %s (missing method %s)", V, T, Ti, m)
			if wrong != nil {
				// TODO(gri) This can still report uninstantiated types which makes the error message
				//           more difficult to read then necessary.
				return errorf("%s does not implement %s: wrong method signature\n\tgot  %s\n\twant %s",
					V, T, wrong, m,
				)
			}
			return errorf("%s does not implement %s (missing method %s)", V, T, m.name)
		}
	}

	// V must also be in the set of types of T, if any.
	// Constraints with empty type sets were already excluded above.
	if !Ti.typeSet().hasTerms() {
		return nil // nothing to do
	}

	// If V is itself an interface, each of its possible types must be in the set
	// of T types (i.e., the V type set must be a subset of the T type set).
	// Interfaces V with empty type sets were already excluded above.
	if Vi != nil {
		if !Vi.typeSet().subsetOf(Ti.typeSet()) {
			// TODO(gri) report which type is missing
			return errorf("%s does not implement %s", V, T)
		}
		return nil
	}

	// Otherwise, V's type must be included in the iface type set.
	if !Ti.typeSet().includes(V) {
		// TODO(gri) report which type is missing
		return errorf("%s does not implement %s", V, T)
	}

	return nil
}
<|endoftext|>


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

package e2e

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"strings"
	"sync"
	"time"

	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/gae/filter/featureBreaker"
	"go.chromium.org/luci/gae/filter/featureBreaker/flaky"
	"go.chromium.org/luci/gae/service/datastore"
	"go.chromium.org/luci/server/tq"
	"go.chromium.org/luci/server/tq/tqtesting"

	cfgpb "go.chromium.org/luci/cv/api/config/v2"
	diagnosticpb "go.chromium.org/luci/cv/api/diagnostic"
	migrationpb "go.chromium.org/luci/cv/api/migration"
	"go.chromium.org/luci/cv/internal/changelist"
	"go.chromium.org/luci/cv/internal/common"
	"go.chromium.org/luci/cv/internal/cvtesting"
	"go.chromium.org/luci/cv/internal/diagnostic"
	pollertask "go.chromium.org/luci/cv/internal/gerrit/poller/task"
	"go.chromium.org/luci/cv/internal/gerrit/updater"
	"go.chromium.org/luci/cv/internal/migration"
	"go.chromium.org/luci/cv/internal/prjmanager"
	pmimpl "go.chromium.org/luci/cv/internal/prjmanager/impl"
	"go.chromium.org/luci/cv/internal/prjmanager/prjpb"
	"go.chromium.org/luci/cv/internal/run"
	"go.chromium.org/luci/cv/internal/run/eventpb"
	runimpl "go.chromium.org/luci/cv/internal/run/impl"
)

const dsFlakinessFlagName = "cv.dsflakiness"

var dsFlakinessFlag = flag.Float64(dsFlakinessFlagName, 0, "DS flakiness probability between 0(default) and 1.0 (always fails)")

// Test encapsulates e2e setup for a CV test.
//
// Embeds cvtesting.Test, which sets CV's dependencies and some simple CV
// components (e.g. TreeClient), while this Test focuses on setup of CV's own
// components.
//
// Typical use:
//   ct := Test{CVDev: true}
//   ctx, cancel := ct.SetUp()
//   defer cancel()
type Test struct {
	*cvtesting.Test // auto-initialized if nil
	// CVDev if true sets e2e test to use `cv-dev` GAE app.
	// Defaults to `cv` GAE app.
	CVDev bool

	PMNotifier  *prjmanager.Notifier
	RunNotifier *run.Notifier

	DiagnosticServer diagnosticpb.DiagnosticServer
	MigrationServer  migrationpb.MigrationServer
	// TODO(tandrii): add CQD fake.

	// dsFlakiness enables ds flakiness for "RunUntil".
	dsFlakiness    float64
	dsFlakinesRand rand.Source
}

func (t *Test) SetUp() (ctx context.Context, deferme func()) {
	switch {
	case t.Test == nil:
		t.Test = &cvtesting.Test{}
	case t.Test.AppID != "":
		panic("overriding cvtesting.Test{AppID} in e2e not supported")
	}
	switch t.CVDev {
	case true:
		t.Test.AppID = "cv-dev"
	case false:
		t.Test.AppID = "cv"
	}

	// Delegate most setup to cvtesting.Test.
	ctx, cancel := t.Test.SetUp()

	if (*dsFlakinessFlag) != 0 {
		t.dsFlakiness = *dsFlakinessFlag
		if t.dsFlakiness < 0 || t.dsFlakiness > 1 {
			panic(fmt.Errorf("invalid %s %f: must be between 0.0 and 1.0", dsFlakinessFlagName, t.dsFlakiness))
		}
		logging.Warningf(ctx, "Using %.4f flaky Datastore", t.dsFlakiness)
		t.dsFlakinesRand = rand.NewSource(0)
	}

	t.PMNotifier = prjmanager.NewNotifier(t.TQDispatcher)
	t.RunNotifier = run.NewNotifier(t.TQDispatcher)
	clUpdater := updater.New(t.TQDispatcher, t.PMNotifier, t.RunNotifier)
	_ = pmimpl.New(t.PMNotifier, t.RunNotifier, clUpdater)
	_ = runimpl.New(t.RunNotifier, t.PMNotifier, clUpdater)

	t.MigrationServer = &migration.MigrationServer{}
	t.DiagnosticServer = &diagnostic.DiagnosticServer{}
	return ctx, cancel
}

// Now returns test clock time in UTC.
func (t *Test) Now() time.Time {
	return t.Clock.Now().UTC()
}

// RunAtLeastOncePM runs at least 1 PM task, possibly more or other tasks.
func (t *Test) RunAtLeastOncePM(ctx context.Context) {
	t.TQ.Run(ctx, tqtesting.StopAfterTask(prjpb.ManageProjectTaskClass))
}

// RunAtLeastOnceRun runs at least 1 Run task, possibly more or other tasks.
func (t *Test) RunAtLeastOnceRun(ctx context.Context) {
	t.TQ.Run(ctx, tqtesting.StopAfterTask(eventpb.ManageRunTaskClass))
}

// RunAtLeastOncePoller runs at least 1 Poller task, possibly more or other
// tasks.
func (t *Test) RunAtLeastOncePoller(ctx context.Context) {
	t.TQ.Run(ctx, tqtesting.StopAfterTask(pollertask.ClassID))
}

// RunUntil runs TQ tasks, while stopIf returns false.
//
// If `dsFlakinessFlag` is set, uses flaky datastore for running TQ tasks.
func (t *Test) RunUntil(ctx context.Context, stopIf func() bool) {
	maxTasks := 1000.0
	taskCtx := ctx
	if t.dsFlakiness > 0 {
		maxTasks *= math.Max(1.0, math.Round(100*t.dsFlakiness))
		taskCtx = t.flakifyDS(ctx)
	}
	i := 0
	tooLong := false

	t.sweepTTQ(ctx)
	var finished []string
	t.TQ.Run(
		taskCtx,
		// StopAfter must be first and also always return false s.t. we correctly
		// record all finished tasks.
		tqtesting.StopAfter(func(task *tqtesting.Task) bool {
			finished = append(finished, fmt.Sprintf("%30s (attempt# %d)", task.Class, task.Attempts))
			t.sweepTTQ(ctx)
			return false
		}),
		// StopBefore is actually used to for conditional stopping.
		// Note that it can `return true` (meaning stop) before any task was run at
		// all.
		tqtesting.StopBefore(func(t *tqtesting.Task) bool {
			switch {
			case stopIf():
				return true
			case float64(i) >= maxTasks:
				tooLong = true
				return true
			default:
				i++
				if i%1000 == 0 {
					logging.Debugf(ctx, "RunUntil is running %d task", i)
				}
				return false
			}
		}),
	)

	// Log only here after all tasks-in-progress are completed.
	outstanding := make([]string, len(t.TQ.Tasks().Pending()))
	for i, task := range t.TQ.Tasks().Pending() {
		outstanding[i] = task.Class
	}
	logging.Debugf(ctx, "RunUntil ran %d iterations, finished %d tasks, left %d tasks", i, len(finished), len(outstanding))
	for i, v := range finished {
		logging.Debugf(ctx, "    finished #%d task: %s", i, v)
	}
	if len(outstanding) > 0 {
		logging.Debugf(ctx, "  outstanding: %s", outstanding)
	}

	if tooLong {
		panic(errors.New("RunUntil ran for too long!"))
	}
	if err := ctx.Err(); err != nil {
		panic(err)
	}
}

// LoadProject returns Project entity or nil if not exists.
func (t *Test) LoadProject(ctx context.Context, lProject string) *prjmanager.Project {
	p, err := prjmanager.Load(ctx, lProject)
	if err != nil {
		panic(err)
	}
	return p
}

// LoadRun returns Run entity or nil if not exists.
func (t *Test) LoadRun(ctx context.Context, id common.RunID) *run.Run {
	r := &run.Run{ID: id}
	switch err := datastore.Get(ctx, r); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		panic(err)
	default:
		return r
	}
}

// LoadRunsOf loads all Runs of a project from Datastore.
func (t *Test) LoadRunsOf(ctx context.Context, lProject string) []*run.Run {
	var res []*run.Run
	err := datastore.GetAll(ctx, run.NewQueryWithLUCIProject(ctx, lProject), &res)
	if err != nil {
		panic(err)
	}
	return res
}

// EarliestCreatedRun returns the earliest created Run in a project.
//
// If there are several such runs, may return any one of them.
//
// Returns nil if there are no Runs.
func (t *Test) EarliestCreatedRunOf(ctx context.Context, lProject string) *run.Run {
	var earliest *run.Run
	for _, r := range t.LoadRunsOf(ctx, lProject) {
		if earliest == nil || earliest.CreateTime.After(r.CreateTime) {
			earliest = r
		}
	}
	return earliest
}

// LoadCL returns CL entity or nil if not exists.
func (t *Test) LoadCL(ctx context.Context, id common.CLID) *changelist.CL {
	cl := &changelist.CL{ID: id}
	switch err := datastore.Get(ctx, cl); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		panic(err)
	default:
		return cl
	}
}

// LoadGerritCL returns CL entity or nil if not exists.
func (t *Test) LoadGerritCL(ctx context.Context, gHost string, gChange int64) *changelist.CL {
	switch cl, err := changelist.MustGobID(gHost, gChange).Get(ctx); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		panic(err)
	default:
		return cl
	}
}

// LogPhase emits easy to recognize log like
// ===========================
// PHASE: ....
// ===========================
func (t *Test) LogPhase(ctx context.Context, format string, args ...interface{}) {
	line := strings.Repeat("=", 80)
	format = fmt.Sprintf("\n%s\nPHASE: %s\n%s", line, format, line)
	logging.Debugf(ctx, format, args...)
}

// MakeCfgSingular return project config with a single ConfigGroup.
func MakeCfgSingular(cgName, gHost, gRepo, gRef string) *cfgpb.Config {
	return &cfgpb.Config{
		ConfigGroups: []*cfgpb.ConfigGroup{
			{
				Name: cgName,
				Gerrit: []*cfgpb.ConfigGroup_Gerrit{
					{
						Url: "https://" + gHost + "/",
						Projects: []*cfgpb.ConfigGroup_Gerrit_Project{
							{
								Name:      gRepo,
								RefRegexp: []string{gRef},
							},
						},
					},
				},
			},
		},
	}
}

// implementation detail

// flakifyDS returns context with flaky Datastore.
func (t *Test) flakifyDS(ctx context.Context) context.Context {
	ctx, fb := featureBreaker.FilterRDS(ctx, nil)
	fb.BreakFeaturesWithCallback(
		flaky.Errors(flaky.Params{
			Rand:                             t.dsFlakinesRand,
			DeadlineProbability:              t.dsFlakiness,
			ConcurrentTransactionProbability: t.dsFlakiness,
		}),
		featureBreaker.DatastoreFeatures...,
	)
	return ctx
}

// sweepTTQ ensures all previously transactionally created tasks are actually
// added to TQ.
//
// This is critical when datastore is flaky, as this may result in transaction
// succeeding, but client receiving a transient error.
//
// Context passed must not have flaky datastore.
func (t *Test) sweepTTQ(ctx context.Context) {
	if t.dsFlakiness > 0 {
		// TODO(tandrii): find a way to instantiate && launch Sweeper per test with
		// a controlled lifetime.
		// Currently, tq.Default.Sweeper is global AND allows at most 1 concurrent
		// sweep, regardless of what `ctx` is passed to Sweep.
		// Furthermore, Sweep must be ran outside of
		// tqtesting.StopAfter/tqtesting.StopBefore callbacks to avoid deadlock.
		go func() {
			tqDefaultSweeperLock.Lock()
			defer tqDefaultSweeperLock.Unlock()
			if err := tq.Sweep(ctx); err != nil && err != ctx.Err() {
				logging.Errorf(ctx, "sweep failed: %s", err)
			}
		}()
	}
}

var tqDefaultSweeperLock sync.Mutex

func init() {
	tq.Default.Sweeper = tq.NewInProcSweeper(tq.InProcSweeperOptions{
		SweepShards:     1,
		SubmitBatchSize: 1,
	})
}
<commit_msg>[cv][e2e] efficient in-test TQ sweeping.<commit_after>// Copyright 2021 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package e2e

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"strings"
	"time"

	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry"
	"go.chromium.org/luci/common/sync/dispatcher"
	"go.chromium.org/luci/common/sync/dispatcher/buffer"
	"go.chromium.org/luci/gae/filter/featureBreaker"
	"go.chromium.org/luci/gae/filter/featureBreaker/flaky"
	"go.chromium.org/luci/gae/service/datastore"
	"go.chromium.org/luci/server/tq"
	"go.chromium.org/luci/server/tq/tqtesting"

	cfgpb "go.chromium.org/luci/cv/api/config/v2"
	diagnosticpb "go.chromium.org/luci/cv/api/diagnostic"
	migrationpb "go.chromium.org/luci/cv/api/migration"
	"go.chromium.org/luci/cv/internal/changelist"
	"go.chromium.org/luci/cv/internal/common"
	"go.chromium.org/luci/cv/internal/cvtesting"
	"go.chromium.org/luci/cv/internal/diagnostic"
	pollertask "go.chromium.org/luci/cv/internal/gerrit/poller/task"
	"go.chromium.org/luci/cv/internal/gerrit/updater"
	"go.chromium.org/luci/cv/internal/migration"
	"go.chromium.org/luci/cv/internal/prjmanager"
	pmimpl "go.chromium.org/luci/cv/internal/prjmanager/impl"
	"go.chromium.org/luci/cv/internal/prjmanager/prjpb"
	"go.chromium.org/luci/cv/internal/run"
	"go.chromium.org/luci/cv/internal/run/eventpb"
	runimpl "go.chromium.org/luci/cv/internal/run/impl"
)

const dsFlakinessFlagName = "cv.dsflakiness"

var dsFlakinessFlag = flag.Float64(dsFlakinessFlagName, 0, "DS flakiness probability between 0(default) and 1.0 (always fails)")

// Test encapsulates e2e setup for a CV test.
//
// Embeds cvtesting.Test, which sets CV's dependencies and some simple CV
// components (e.g. TreeClient), while this Test focuses on setup of CV's own
// components.
//
// Typical use:
//   ct := Test{CVDev: true}
//   ctx, cancel := ct.SetUp()
//   defer cancel()
type Test struct {
	*cvtesting.Test // auto-initialized if nil
	// CVDev if true sets e2e test to use `cv-dev` GAE app.
	// Defaults to `cv` GAE app.
	CVDev bool

	PMNotifier  *prjmanager.Notifier
	RunNotifier *run.Notifier

	DiagnosticServer diagnosticpb.DiagnosticServer
	MigrationServer  migrationpb.MigrationServer
	// TODO(tandrii): add CQD fake.

	// dsFlakiness enables ds flakiness for "RunUntil".
	dsFlakiness    float64
	dsFlakinesRand rand.Source
	tqSweepChannel dispatcher.Channel
}

func (t *Test) SetUp() (ctx context.Context, deferme func()) {
	switch {
	case t.Test == nil:
		t.Test = &cvtesting.Test{}
	case t.Test.AppID != "":
		panic("overriding cvtesting.Test{AppID} in e2e not supported")
	}
	switch t.CVDev {
	case true:
		t.Test.AppID = "cv-dev"
	case false:
		t.Test.AppID = "cv"
	}

	// Delegate most setup to cvtesting.Test.
	ctx, ctxCancel := t.Test.SetUp()
	deferme = ctxCancel

	if (*dsFlakinessFlag) != 0 {
		t.dsFlakiness = *dsFlakinessFlag
		if t.dsFlakiness < 0 || t.dsFlakiness > 1 {
			panic(fmt.Errorf("invalid %s %f: must be between 0.0 and 1.0", dsFlakinessFlagName, t.dsFlakiness))
		}
		logging.Warningf(ctx, "Using %.4f flaky Datastore", t.dsFlakiness)
		t.dsFlakinesRand = rand.NewSource(0)
		stopSweeping := t.startTQSweeping(ctx)
		deferme = func() {
			stopSweeping()
			ctxCancel()
		}
	}

	t.PMNotifier = prjmanager.NewNotifier(t.TQDispatcher)
	t.RunNotifier = run.NewNotifier(t.TQDispatcher)
	clUpdater := updater.New(t.TQDispatcher, t.PMNotifier, t.RunNotifier)
	_ = pmimpl.New(t.PMNotifier, t.RunNotifier, clUpdater)
	_ = runimpl.New(t.RunNotifier, t.PMNotifier, clUpdater)

	t.MigrationServer = &migration.MigrationServer{}
	t.DiagnosticServer = &diagnostic.DiagnosticServer{}
	return ctx, deferme
}

// Now returns test clock time in UTC.
func (t *Test) Now() time.Time {
	return t.Clock.Now().UTC()
}

// RunAtLeastOncePM runs at least 1 PM task, possibly more or other tasks.
func (t *Test) RunAtLeastOncePM(ctx context.Context) {
	t.TQ.Run(ctx, tqtesting.StopAfterTask(prjpb.ManageProjectTaskClass))
}

// RunAtLeastOnceRun runs at least 1 Run task, possibly more or other tasks.
func (t *Test) RunAtLeastOnceRun(ctx context.Context) {
	t.TQ.Run(ctx, tqtesting.StopAfterTask(eventpb.ManageRunTaskClass))
}

// RunAtLeastOncePoller runs at least 1 Poller task, possibly more or other
// tasks.
func (t *Test) RunAtLeastOncePoller(ctx context.Context) {
	t.TQ.Run(ctx, tqtesting.StopAfterTask(pollertask.ClassID))
}

// RunUntil runs TQ tasks, while stopIf returns false.
//
// If `dsFlakinessFlag` is set, uses flaky datastore for running TQ tasks.
func (t *Test) RunUntil(ctx context.Context, stopIf func() bool) {
	maxTasks := 1000.0
	taskCtx := ctx
	if t.dsFlakiness > 0 {
		maxTasks *= math.Max(1.0, math.Round(100*t.dsFlakiness))
		taskCtx = t.flakifyDS(ctx)
	}
	i := 0
	tooLong := false

	t.enqueueTQSweep(ctx)
	var finished []string
	t.TQ.Run(
		taskCtx,
		// StopAfter must be first and also always return false s.t. we correctly
		// record all finished tasks.
		tqtesting.StopAfter(func(task *tqtesting.Task) bool {
			finished = append(finished, fmt.Sprintf("%30s (attempt# %d)", task.Class, task.Attempts))
			t.enqueueTQSweep(ctx)
			return false
		}),
		// StopBefore is actually used to for conditional stopping.
		// Note that it can `return true` (meaning stop) before any task was run at
		// all.
		tqtesting.StopBefore(func(t *tqtesting.Task) bool {
			switch {
			case stopIf():
				return true
			case float64(i) >= maxTasks:
				tooLong = true
				return true
			default:
				i++
				if i%1000 == 0 {
					logging.Debugf(ctx, "RunUntil is running %d task", i)
				}
				return false
			}
		}),
	)

	// Log only here after all tasks-in-progress are completed.
	outstanding := make([]string, len(t.TQ.Tasks().Pending()))
	for i, task := range t.TQ.Tasks().Pending() {
		outstanding[i] = task.Class
	}
	logging.Debugf(ctx, "RunUntil ran %d iterations, finished %d tasks, left %d tasks", i, len(finished), len(outstanding))
	for i, v := range finished {
		logging.Debugf(ctx, "    finished #%d task: %s", i, v)
	}
	if len(outstanding) > 0 {
		logging.Debugf(ctx, "  outstanding: %s", outstanding)
	}

	if tooLong {
		panic(errors.New("RunUntil ran for too long!"))
	}
	if err := ctx.Err(); err != nil {
		panic(err)
	}
}

// LoadProject returns Project entity or nil if not exists.
func (t *Test) LoadProject(ctx context.Context, lProject string) *prjmanager.Project {
	p, err := prjmanager.Load(ctx, lProject)
	if err != nil {
		panic(err)
	}
	return p
}

// LoadRun returns Run entity or nil if not exists.
func (t *Test) LoadRun(ctx context.Context, id common.RunID) *run.Run {
	r := &run.Run{ID: id}
	switch err := datastore.Get(ctx, r); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		panic(err)
	default:
		return r
	}
}

// LoadRunsOf loads all Runs of a project from Datastore.
func (t *Test) LoadRunsOf(ctx context.Context, lProject string) []*run.Run {
	var res []*run.Run
	err := datastore.GetAll(ctx, run.NewQueryWithLUCIProject(ctx, lProject), &res)
	if err != nil {
		panic(err)
	}
	return res
}

// EarliestCreatedRun returns the earliest created Run in a project.
//
// If there are several such runs, may return any one of them.
//
// Returns nil if there are no Runs.
func (t *Test) EarliestCreatedRunOf(ctx context.Context, lProject string) *run.Run {
	var earliest *run.Run
	for _, r := range t.LoadRunsOf(ctx, lProject) {
		if earliest == nil || earliest.CreateTime.After(r.CreateTime) {
			earliest = r
		}
	}
	return earliest
}

// LoadCL returns CL entity or nil if not exists.
func (t *Test) LoadCL(ctx context.Context, id common.CLID) *changelist.CL {
	cl := &changelist.CL{ID: id}
	switch err := datastore.Get(ctx, cl); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		panic(err)
	default:
		return cl
	}
}

// LoadGerritCL returns CL entity or nil if not exists.
func (t *Test) LoadGerritCL(ctx context.Context, gHost string, gChange int64) *changelist.CL {
	switch cl, err := changelist.MustGobID(gHost, gChange).Get(ctx); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		panic(err)
	default:
		return cl
	}
}

// LogPhase emits easy to recognize log like
// ===========================
// PHASE: ....
// ===========================
func (t *Test) LogPhase(ctx context.Context, format string, args ...interface{}) {
	line := strings.Repeat("=", 80)
	format = fmt.Sprintf("\n%s\nPHASE: %s\n%s", line, format, line)
	logging.Debugf(ctx, format, args...)
}

// MakeCfgSingular return project config with a single ConfigGroup.
func MakeCfgSingular(cgName, gHost, gRepo, gRef string) *cfgpb.Config {
	return &cfgpb.Config{
		ConfigGroups: []*cfgpb.ConfigGroup{
			{
				Name: cgName,
				Gerrit: []*cfgpb.ConfigGroup_Gerrit{
					{
						Url: "https://" + gHost + "/",
						Projects: []*cfgpb.ConfigGroup_Gerrit_Project{
							{
								Name:      gRepo,
								RefRegexp: []string{gRef},
							},
						},
					},
				},
			},
		},
	}
}

// implementation detail

// flakifyDS returns context with flaky Datastore.
func (t *Test) flakifyDS(ctx context.Context) context.Context {
	ctx, fb := featureBreaker.FilterRDS(ctx, nil)
	fb.BreakFeaturesWithCallback(
		flaky.Errors(flaky.Params{
			Rand:                             t.dsFlakinesRand,
			DeadlineProbability:              t.dsFlakiness,
			ConcurrentTransactionProbability: t.dsFlakiness,
		}),
		featureBreaker.DatastoreFeatures...,
	)
	return ctx
}

// startTQSweeping starts asynchronous sweeping for the duration of the test.
//
// This is necessary if flaky DS is used.
func (t *Test) startTQSweeping(ctx context.Context) (deferme func()) {
	t.TQDispatcher.Sweeper = tq.NewInProcSweeper(tq.InProcSweeperOptions{
		SweepShards:     1,
		SubmitBatchSize: 1,
	})
	var err error
	t.tqSweepChannel, err = dispatcher.NewChannel(
		ctx,
		&dispatcher.Options{
			Buffer: buffer.Options{
				BatchSize: 1, // incoming event => sweep ASAP.
				MaxLeases: 1, // at most 1 sweep concurrently
				// 2+ outstanding requests to sweep should result in just 1 sweep.
				FullBehavior: &buffer.DropOldestBatch{MaxLiveItems: 1},
				// This is only useful if something is misconfigured to avoid pointless
				// retries, because the individual sweeps must not fail as we use
				// non-flaky Datastore for sweeping.
				Retry: retry.None,
			},
		},
		func(*buffer.Batch) error { return t.TQDispatcher.Sweep(ctx) },
	)
	if err != nil {
		panic(err)
	}
	return func() { t.tqSweepChannel.CloseAndDrain(ctx) }
}

// enqueueTQSweep ensures a TQ sweep will happen strictly afterwards.
//
// Noop if TQ sweeping is not required.
func (t *Test) enqueueTQSweep(ctx context.Context) {
	if t.TQDispatcher.Sweeper != nil {
		t.tqSweepChannel.C <- struct{}{}
	}
}
<|endoftext|>