TUN-1626: Create new supervisor to establish connection with origintunneld

connection/connection.go

@@ -0,0 +1,159 @@
+package connection
+
+import (
+	"context"
+	"crypto/tls"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/cloudflare/cloudflared/h2mux"
+	"github.com/cloudflare/cloudflared/tunnelrpc"
+	tunnelpogs "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
+	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+
+	rpc "zombiezen.com/go/capnproto2/rpc"
+)
+
+const (
+	dialTimeout = 5 * time.Second
+)
+
+type dialError struct {
+	cause error
+}
+
+func (e dialError) Error() string {
+	return e.cause.Error()
+}
+
+type muxerShutdownError struct{}
+
+func (e muxerShutdownError) Error() string {
+	return "muxer shutdown"
+}
+
+type ConnectionConfig struct {
+	TLSConfig         *tls.Config
+	HeartbeatInterval time.Duration
+	MaxHeartbeats     uint64
+	Logger            *logrus.Entry
+}
+
+type connectionHandler interface {
+	serve(ctx context.Context) error
+	connect(ctx context.Context, parameters *tunnelpogs.ConnectParameters) (*tunnelpogs.ConnectResult, error)
+	shutdown()
+}
+
+type h2muxHandler struct {
+	muxer  *h2mux.Muxer
+	logger *logrus.Entry
+}
+
+type muxedStreamHandler struct {
+}
+
+// Implements MuxedStreamHandler interface
+func (h *muxedStreamHandler) ServeStream(stream *h2mux.MuxedStream) error {
+	return nil
+}
+
+func (h *h2muxHandler) serve(ctx context.Context) error {
+	// Serve doesn't return until h2mux is shutdown
+	if err := h.muxer.Serve(ctx); err != nil {
+		return err
+	}
+	return muxerShutdownError{}
+}
+
+// Connect is used to establish connections with cloudflare's edge network
+func (h *h2muxHandler) connect(ctx context.Context, parameters *tunnelpogs.ConnectParameters) (*tunnelpogs.ConnectResult, error) {
+	conn, err := h.newRPConn()
+	if err != nil {
+		return nil, errors.Wrap(err, "Failed to create new RPC connection")
+	}
+	defer conn.Close()
+	tsClient := tunnelpogs.TunnelServer_PogsClient{Client: conn.Bootstrap(ctx)}
+	return tsClient.Connect(ctx, parameters)
+}
+
+func (h *h2muxHandler) shutdown() {
+	h.muxer.Shutdown()
+}
+
+func (h *h2muxHandler) newRPConn() (*rpc.Conn, error) {
+	stream, err := h.muxer.OpenStream([]h2mux.Header{
+		{Name: ":method", Value: "RPC"},
+		{Name: ":scheme", Value: "capnp"},
+		{Name: ":path", Value: "*"},
+	}, nil)
+	if err != nil {
+		return nil, err
+	}
+	return rpc.NewConn(
+		tunnelrpc.NewTransportLogger(h.logger.WithField("subsystem", "rpc-register"), rpc.StreamTransport(stream)),
+		tunnelrpc.ConnLog(h.logger.WithField("subsystem", "rpc-transport")),
+	), nil
+}
+
+// NewConnectionHandler returns a connectionHandler, wrapping h2mux to make RPC calls
+func newH2MuxHandler(ctx context.Context,
+	config *ConnectionConfig,
+	edgeIP *net.TCPAddr,
+) (connectionHandler, error) {
+	// Inherit from parent context so we can cancel (Ctrl-C) while dialing
+	dialCtx, dialCancel := context.WithTimeout(ctx, dialTimeout)
+	defer dialCancel()
+	dialer := net.Dialer{DualStack: true}
+	plaintextEdgeConn, err := dialer.DialContext(dialCtx, "tcp", edgeIP.String())
+	if err != nil {
+		return nil, dialError{cause: errors.Wrap(err, "DialContext error")}
+	}
+	edgeConn := tls.Client(plaintextEdgeConn, config.TLSConfig)
+	edgeConn.SetDeadline(time.Now().Add(dialTimeout))
+	err = edgeConn.Handshake()
+	if err != nil {
+		return nil, dialError{cause: errors.Wrap(err, "Handshake with edge error")}
+	}
+	// clear the deadline on the conn; h2mux has its own timeouts
+	edgeConn.SetDeadline(time.Time{})
+	// Establish a muxed connection with the edge
+	// Client mux handshake with agent server
+	muxer, err := h2mux.Handshake(edgeConn, edgeConn, h2mux.MuxerConfig{
+		Timeout:           dialTimeout,
+		Handler:           &muxedStreamHandler{},
+		IsClient:          true,
+		HeartbeatInterval: config.HeartbeatInterval,
+		MaxHeartbeats:     config.MaxHeartbeats,
+		Logger:            config.Logger,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &h2muxHandler{
+		muxer:  muxer,
+		logger: config.Logger,
+	}, nil
+}
+
+// connectionPool is a pool of connection handlers
+type connectionPool struct {
+	sync.Mutex
+	connectionHandlers []connectionHandler
+}
+
+func (cp *connectionPool) put(h connectionHandler) {
+	cp.Lock()
+	defer cp.Unlock()
+	cp.connectionHandlers = append(cp.connectionHandlers, h)
+}
+
+func (cp *connectionPool) close() {
+	cp.Lock()
+	defer cp.Unlock()
+	for _, h := range cp.connectionHandlers {
+		h.shutdown()
+	}
+}