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 openStreamTimeout = 30 * 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) { openStreamCtx, cancel := context.WithTimeout(ctx, openStreamTimeout) defer cancel() conn, err := h.newRPConn(openStreamCtx) 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(ctx context.Context) (*rpc.Conn, error) { stream, err := h.muxer.OpenStream(ctx, []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() } }