open-consul/agent/grpc/client.go

package grpc

import (
	"context"
	"fmt"
	"net"
	"sync"

	"google.golang.org/grpc"

	"github.com/hashicorp/consul/agent/metadata"
	"github.com/hashicorp/consul/agent/pool"
)

// ClientConnPool creates and stores a connection for each datacenter.
type ClientConnPool struct {
	dialer    dialer
	servers   ServerLocator
	conns     map[string]*grpc.ClientConn
	connsLock sync.Mutex
}

type ServerLocator interface {
	// ServerForAddr is used to look up server metadata from an address.
	ServerForAddr(addr string) (*metadata.Server, error)
	// Scheme returns the url scheme to use to dial the server. This is primarily
	// needed for testing multiple agents in parallel, because gRPC requires the
	// resolver to be registered globally.
	Scheme() string
}

// TLSWrapper wraps a non-TLS connection and returns a connection with TLS
// enabled.
type TLSWrapper func(dc string, conn net.Conn) (net.Conn, error)

type dialer func(context.Context, string) (net.Conn, error)

func NewClientConnPool(servers ServerLocator, tls TLSWrapper) *ClientConnPool {
	return &ClientConnPool{
		dialer:  newDialer(servers, tls),
		servers: servers,
		conns:   make(map[string]*grpc.ClientConn),
	}
}

// ClientConn returns a grpc.ClientConn for the datacenter. If there are no
// existing connections in the pool, a new one will be created, stored in the pool,
// then returned.
func (c *ClientConnPool) ClientConn(datacenter string) (*grpc.ClientConn, error) {
	c.connsLock.Lock()
	defer c.connsLock.Unlock()

	if conn, ok := c.conns[datacenter]; ok {
		return conn, nil
	}

	conn, err := grpc.Dial(
		fmt.Sprintf("%s:///server.%s", c.servers.Scheme(), datacenter),
		// use WithInsecure mode here because we handle the TLS wrapping in the
		// custom dialer based on logic around whether the server has TLS enabled.
		grpc.WithInsecure(),
		grpc.WithContextDialer(c.dialer),
		grpc.WithDisableRetry(),
		grpc.WithStatsHandler(newStatsHandler(defaultMetrics)),
		// nolint:staticcheck // there is no other supported alternative to WithBalancerName
		grpc.WithBalancerName("pick_first"))
	if err != nil {
		return nil, err
	}

	c.conns[datacenter] = conn
	return conn, nil
}

// newDialer returns a gRPC dialer function that conditionally wraps the connection
// with TLS based on the Server.useTLS value.
func newDialer(servers ServerLocator, wrapper TLSWrapper) func(context.Context, string) (net.Conn, error) {
	return func(ctx context.Context, addr string) (net.Conn, error) {
		d := net.Dialer{}
		conn, err := d.DialContext(ctx, "tcp", addr)
		if err != nil {
			return nil, err
		}

		server, err := servers.ServerForAddr(addr)
		if err != nil {
			conn.Close()
			return nil, err
		}

		if server.UseTLS {
			if wrapper == nil {
				conn.Close()
				return nil, fmt.Errorf("TLS enabled but got nil TLS wrapper")
			}

			// Switch the connection into TLS mode
			if _, err := conn.Write([]byte{byte(pool.RPCTLS)}); err != nil {
				conn.Close()
				return nil, err
			}

			// Wrap the connection in a TLS client
			tlsConn, err := wrapper(server.Datacenter, conn)
			if err != nil {
				conn.Close()
				return nil, err
			}
			conn = tlsConn
		}

		_, err = conn.Write([]byte{pool.RPCGRPC})
		if err != nil {
			conn.Close()
			return nil, err
		}

		return conn, nil
	}
}