TUN-4866: Add Control Stream for QUIC

This commit adds support to Register and Unregister Connections via RPC
on the QUIC transport protocol

Author: sudarshan-reddy

cmd/cloudflared/tunnel/configuration.go

@@ -248,10 +248,17 @@ func prepareTunnelConfig(
 
 	edgeTLSConfigs := make(map[connection.Protocol]*tls.Config, len(connection.ProtocolList))
 	for _, p := range connection.ProtocolList {
-		edgeTLSConfig, err := tlsconfig.CreateTunnelConfig(c, p.ServerName())
+		tlsSettings := p.TLSSettings()
+		if tlsSettings == nil {
+			return nil, ingress.Ingress{}, fmt.Errorf("%s has unknown TLS settings", p)
+		}
+		edgeTLSConfig, err := tlsconfig.CreateTunnelConfig(c, tlsSettings.ServerName)
 		if err != nil {
 			return nil, ingress.Ingress{}, errors.Wrap(err, "unable to create TLS config to connect with edge")
 		}
+		if len(edgeTLSConfig.NextProtos) > 0 {
+			edgeTLSConfig.NextProtos = tlsSettings.NextProtos
+		}
 		edgeTLSConfigs[p] = edgeTLSConfig
 	}
 

connection/control.go

@@ -0,0 +1,89 @@
+package connection
+
+import (
+	"context"
+	"io"
+	"time"
+
+	"github.com/rs/zerolog"
+
+	tunnelpogs "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
+)
+
+// RPCClientFunc derives a named tunnel rpc client that can then be used to register and unregister connections.
+type RPCClientFunc func(context.Context, io.ReadWriteCloser, *zerolog.Logger) NamedTunnelRPCClient
+
+type controlStream struct {
+	observer *Observer
+
+	connectedFuse     ConnectedFuse
+	namedTunnelConfig *NamedTunnelConfig
+	connIndex         uint8
+
+	newRPCClientFunc RPCClientFunc
+
+	gracefulShutdownC <-chan struct{}
+	gracePeriod       time.Duration
+	stoppedGracefully bool
+}
+
+// ControlStreamHandler registers connections with origintunneld and initiates graceful shutdown.
+type ControlStreamHandler interface {
+	ServeControlStream(ctx context.Context, rw io.ReadWriteCloser, connOptions *tunnelpogs.ConnectionOptions) error
+	IsStopped() bool
+}
+
+// NewControlStream returns a new instance of ControlStreamHandler
+func NewControlStream(
+	observer *Observer,
+	connectedFuse ConnectedFuse,
+	namedTunnelConfig *NamedTunnelConfig,
+	connIndex uint8,
+	newRPCClientFunc RPCClientFunc,
+	gracefulShutdownC <-chan struct{},
+	gracePeriod time.Duration,
+) ControlStreamHandler {
+	if newRPCClientFunc == nil {
+		newRPCClientFunc = newRegistrationRPCClient
+	}
+	return &controlStream{
+		observer:          observer,
+		connectedFuse:     connectedFuse,
+		namedTunnelConfig: namedTunnelConfig,
+		newRPCClientFunc:  newRPCClientFunc,
+		connIndex:         connIndex,
+		gracefulShutdownC: gracefulShutdownC,
+		gracePeriod:       gracePeriod,
+	}
+}
+
+func (c *controlStream) ServeControlStream(
+	ctx context.Context,
+	rw io.ReadWriteCloser,
+	connOptions *tunnelpogs.ConnectionOptions,
+) error {
+	rpcClient := c.newRPCClientFunc(ctx, rw, c.observer.log)
+	defer rpcClient.Close()
+
+	if err := rpcClient.RegisterConnection(ctx, c.namedTunnelConfig, connOptions, c.connIndex, c.observer); err != nil {
+		return err
+	}
+	c.connectedFuse.Connected()
+
+	// wait for connection termination or start of graceful shutdown
+	select {
+	case <-ctx.Done():
+		break
+	case <-c.gracefulShutdownC:
+		c.stoppedGracefully = true
+	}
+
+	c.observer.sendUnregisteringEvent(c.connIndex)
+	rpcClient.GracefulShutdown(ctx, c.gracePeriod)
+	c.observer.log.Info().Uint8(LogFieldConnIndex, c.connIndex).Msg("Unregistered tunnel connection")
+	return nil
+}
+
+func (c *controlStream) IsStopped() bool {
+	return c.stoppedGracefully
+}

connection/http2.go

@@ -30,52 +30,44 @@ var errEdgeConnectionClosed = fmt.Errorf("connection with edge closed")
 // HTTP2Connection represents a net.Conn that uses HTTP2 frames to proxy traffic from the edge to cloudflared on the
 // origin.
 type HTTP2Connection struct {
-	conn         net.Conn
-	server       *http2.Server
-	config       *Config
-	namedTunnel  *NamedTunnelConfig
-	connOptions  *tunnelpogs.ConnectionOptions
-	observer     *Observer
-	connIndexStr string
-	connIndex    uint8
+	conn        net.Conn
+	server      *http2.Server
+	config      *Config
+	connOptions *tunnelpogs.ConnectionOptions
+	observer    *Observer
+	connIndex   uint8
 	// newRPCClientFunc allows us to mock RPCs during testing
 	newRPCClientFunc func(context.Context, io.ReadWriteCloser, *zerolog.Logger) NamedTunnelRPCClient
 
-	log               *zerolog.Logger
-	activeRequestsWG  sync.WaitGroup
-	connectedFuse     ConnectedFuse
-	gracefulShutdownC <-chan struct{}
-	stoppedGracefully bool
-	controlStreamErr  error // result of running control stream handler
+	log                  *zerolog.Logger
+	activeRequestsWG     sync.WaitGroup
+	controlStreamHandler ControlStreamHandler
+	stoppedGracefully    bool
+	controlStreamErr     error // result of running control stream handler
 }
 
 // NewHTTP2Connection returns a new instance of HTTP2Connection.
 func NewHTTP2Connection(
 	conn net.Conn,
 	config *Config,
-	namedTunnelConfig *NamedTunnelConfig,
 	connOptions *tunnelpogs.ConnectionOptions,
 	observer *Observer,
 	connIndex uint8,
-	connectedFuse ConnectedFuse,
+	controlStreamHandler ControlStreamHandler,
 	log *zerolog.Logger,
-	gracefulShutdownC <-chan struct{},
 ) *HTTP2Connection {
 	return &HTTP2Connection{
 		conn: conn,
 		server: &http2.Server{
 			MaxConcurrentStreams: math.MaxUint32,
 		},
-		config:            config,
-		namedTunnel:       namedTunnelConfig,
-		connOptions:       connOptions,
-		observer:          observer,
-		connIndexStr:      uint8ToString(connIndex),
-		connIndex:         connIndex,
-		newRPCClientFunc:  newRegistrationRPCClient,
-		connectedFuse:     connectedFuse,
-		log:               log,
-		gracefulShutdownC: gracefulShutdownC,
+		config:               config,
+		connOptions:          connOptions,
+		observer:             observer,
+		connIndex:            connIndex,
+		newRPCClientFunc:     newRegistrationRPCClient,
+		controlStreamHandler: controlStreamHandler,
+		log:                  log,
 	}
 }
 
@@ -91,7 +83,7 @@ func (c *HTTP2Connection) Serve(ctx context.Context) error {
 	})
 
 	switch {
-	case c.stoppedGracefully:
+	case c.controlStreamHandler.IsStopped():
 		return nil
 	case c.controlStreamErr != nil:
 		return c.controlStreamErr
@@ -116,7 +108,7 @@ func (c *HTTP2Connection) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	switch connType {
 	case TypeControlStream:
-		if err := c.serveControlStream(r.Context(), respWriter); err != nil {
+		if err := c.controlStreamHandler.ServeControlStream(r.Context(), respWriter, c.connOptions); err != nil {
 			c.controlStreamErr = err
 			c.log.Error().Err(err)
 			respWriter.WriteErrorResponse()
@@ -154,29 +146,6 @@ func (c *HTTP2Connection) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func (c *HTTP2Connection) serveControlStream(ctx context.Context, respWriter *http2RespWriter) error {
-	rpcClient := c.newRPCClientFunc(ctx, respWriter, c.observer.log)
-	defer rpcClient.Close()
-
-	if err := rpcClient.RegisterConnection(ctx, c.namedTunnel, c.connOptions, c.connIndex, c.observer); err != nil {
-		return err
-	}
-	c.connectedFuse.Connected()
-
-	// wait for connection termination or start of graceful shutdown
-	select {
-	case <-ctx.Done():
-		break
-	case <-c.gracefulShutdownC:
-		c.stoppedGracefully = true
-	}
-
-	c.observer.sendUnregisteringEvent(c.connIndex)
-	rpcClient.GracefulShutdown(ctx, c.config.GracePeriod)
-	c.observer.log.Info().Uint8(LogFieldConnIndex, c.connIndex).Msg("Unregistered tunnel connection")
-	return nil
-}
-
 func (c *HTTP2Connection) close() {
 	// Wait for all serve HTTP handlers to return
 	c.activeRequestsWG.Wait()

connection/http2_test.go

@@ -30,16 +30,24 @@ func newTestHTTP2Connection() (*HTTP2Connection, net.Conn) {
 	edgeConn, originConn := net.Pipe()
 	var connIndex = uint8(0)
 	log := zerolog.Nop()
+	obs := NewObserver(&log, &log, false)
+	controlStream := NewControlStream(
+		obs,
+		mockConnectedFuse{},
+		&NamedTunnelConfig{},
+		connIndex,
+		nil,
+		nil,
+		1*time.Second,
+	)
 	return NewHTTP2Connection(
 		originConn,
 		testConfig,
-		&NamedTunnelConfig{},
 		&pogs.ConnectionOptions{},
-		NewObserver(&log, &log, false),
+		obs,
 		connIndex,
-		mockConnectedFuse{},
+		controlStream,
 		&log,
-		nil,
 	), edgeConn
 }
 
@@ -225,7 +233,18 @@ func TestServeControlStream(t *testing.T) {
 		registered:   make(chan struct{}),
 		unregistered: make(chan struct{}),
 	}
-	http2Conn.newRPCClientFunc = rpcClientFactory.newMockRPCClient
+
+	obs := NewObserver(&log, &log, false)
+	controlStream := NewControlStream(
+		obs,
+		mockConnectedFuse{},
+		&NamedTunnelConfig{},
+		1,
+		rpcClientFactory.newMockRPCClient,
+		nil,
+		1*time.Second,
+	)
+	http2Conn.controlStreamHandler = controlStream
 
 	ctx, cancel := context.WithCancel(context.Background())
 	var wg sync.WaitGroup
@@ -264,7 +283,18 @@ func TestFailRegistration(t *testing.T) {
 		registered:   make(chan struct{}),
 		unregistered: make(chan struct{}),
 	}
-	http2Conn.newRPCClientFunc = rpcClientFactory.newMockRPCClient
+
+	obs := NewObserver(&log, &log, false)
+	controlStream := NewControlStream(
+		obs,
+		mockConnectedFuse{},
+		&NamedTunnelConfig{},
+		http2Conn.connIndex,
+		rpcClientFactory.newMockRPCClient,
+		nil,
+		1*time.Second,
+	)
+	http2Conn.controlStreamHandler = controlStream
 
 	ctx, cancel := context.WithCancel(context.Background())
 	var wg sync.WaitGroup
@@ -297,10 +327,21 @@ func TestGracefulShutdownHTTP2(t *testing.T) {
 		unregistered: make(chan struct{}),
 	}
 	events := &eventCollectorSink{}
-	http2Conn.newRPCClientFunc = rpcClientFactory.newMockRPCClient
-	http2Conn.observer.RegisterSink(events)
+
 	shutdownC := make(chan struct{})
-	http2Conn.gracefulShutdownC = shutdownC
+	obs := NewObserver(&log, &log, false)
+	obs.RegisterSink(events)
+	controlStream := NewControlStream(
+		obs,
+		mockConnectedFuse{},
+		&NamedTunnelConfig{},
+		http2Conn.connIndex,
+		rpcClientFactory.newMockRPCClient,
+		shutdownC,
+		1*time.Second,
+	)
+
+	http2Conn.controlStreamHandler = controlStream
 
 	ctx, cancel := context.WithCancel(context.Background())
 	var wg sync.WaitGroup
@@ -339,7 +380,7 @@ func TestGracefulShutdownHTTP2(t *testing.T) {
 	case <-time.Tick(time.Second):
 		t.Fatal("timeout out waiting for unregistered signal")
 	}
-	assert.True(t, http2Conn.stoppedGracefully)
+	assert.True(t, controlStream.IsStopped())
 
 	cancel()
 	wg.Wait()