TUN-5301: Separate datagram multiplex and session management logic from quic connection logic

connection/quic.go

@@ -16,6 +16,7 @@ import (
 	"github.com/rs/zerolog"
 	"golang.org/x/sync/errgroup"
 
+	"github.com/cloudflare/cloudflared/datagramsession"
 	quicpogs "github.com/cloudflare/cloudflared/quic"
 	tunnelpogs "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
 )
@@ -32,10 +33,11 @@ const (
 
 // QUICConnection represents the type that facilitates Proxying via QUIC streams.
 type QUICConnection struct {
-	session     quic.Session
-	logger      *zerolog.Logger
-	httpProxy   OriginProxy
-	udpSessions *udpSessions
+	session        quic.Session
+	logger         *zerolog.Logger
+	httpProxy      OriginProxy
+	sessionManager datagramsession.Manager
+	localIP        net.IP
 }
 
 // NewQUICConnection returns a new instance of QUICConnection.
@@ -49,12 +51,6 @@ func NewQUICConnection(
 	controlStreamHandler ControlStreamHandler,
 	observer *Observer,
 ) (*QUICConnection, error) {
-	localIP, err := GetLocalIP()
-	if err != nil {
-		return nil, err
-	}
-	observer.log.Info().Msgf("UDP proxy will use %s as packet source IP", localIP)
-	udpSessions := newUDPSessions(localIP)
 	session, err := quic.DialAddr(edgeAddr.String(), tlsConfig, quicConfig)
 	if err != nil {
 		return nil, fmt.Errorf("failed to dial to edge: %w", err)
@@ -71,24 +67,36 @@ func NewQUICConnection(
 		return nil, err
 	}
 
+	datagramMuxer, err := quicpogs.NewDatagramMuxer(session)
+	if err != nil {
+		return nil, err
+	}
+
+	sessionManager := datagramsession.NewManager(datagramMuxer, observer.log)
+
+	localIP, err := getLocalIP()
+	if err != nil {
+		return nil, err
+	}
+
 	return &QUICConnection{
-		session:     session,
-		httpProxy:   httpProxy,
-		logger:      observer.log,
-		udpSessions: udpSessions,
+		session:        session,
+		httpProxy:      httpProxy,
+		logger:         observer.log,
+		sessionManager: sessionManager,
+		localIP:        localIP,
 	}, nil
 }
 
 // Serve starts a QUIC session that begins accepting streams.
 func (q *QUICConnection) Serve(ctx context.Context) error {
 	errGroup, ctx := errgroup.WithContext(ctx)
-	errGroup.Go(func() error {
-		return q.listenEdgeDatagram()
-	})
-
 	errGroup.Go(func() error {
 		return q.acceptStream(ctx)
 	})
+	errGroup.Go(func() error {
+		return q.sessionManager.Serve(ctx)
+	})
 	return errGroup.Wait()
 }
 
@@ -111,26 +119,6 @@ func (q *QUICConnection) acceptStream(ctx context.Context) error {
 	}
 }
 
-// listenEdgeDatagram listens for datagram from edge, parse the session ID and find the UDPConn to send the payload
-func (q *QUICConnection) listenEdgeDatagram() error {
-	for {
-		msg, err := q.session.ReceiveMessage()
-		if err != nil {
-			return err
-		}
-		go func(msg []byte) {
-			sessionID, msgWithoutID, err := quicpogs.ExtractSessionID(msg)
-			if err != nil {
-				q.logger.Err(err).Msg("Failed to parse session ID from datagram")
-				return
-			}
-			if err := q.udpSessions.send(sessionID, msgWithoutID); err != nil {
-				q.logger.Err(err).Msg("Failed to send UDP to origin")
-			}
-		}(msg)
-	}
-}
-
 // Close closes the session with no errors specified.
 func (q *QUICConnection) Close() {
 	q.session.CloseWithError(0, "")
@@ -186,46 +174,29 @@ func (q *QUICConnection) handleRPCStream(rpcStream *quicpogs.RPCServerStream) er
 }
 
 func (q *QUICConnection) RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16) error {
-	udpConn, err := q.udpSessions.register(sessionID, dstIP, dstPort)
+	// Each session is a series of datagram from an eyeball to a dstIP:dstPort.
+	// (src port, dst IP, dst port) uniquely identifies a session, so it needs a dedicated connected socket.
+	originProxy, err := q.newUDPProxy(dstIP, dstPort)
 	if err != nil {
+		q.logger.Err(err).Msgf("Failed to create udp proxy to %s:%d", dstIP, dstPort)
 		return err
 	}
-	q.logger.Debug().Msgf("Register session %v, %v, %v", sessionID, dstIP, dstPort)
-	go q.listenOriginUDP(sessionID, udpConn)
+	session, err := q.sessionManager.RegisterSession(ctx, sessionID, originProxy)
+	if err != nil {
+		q.logger.Err(err).Msgf("Failed to register udp session %s", sessionID)
+		return err
+	}
+	go func() {
+		defer q.sessionManager.UnregisterSession(q.session.Context(), sessionID)
+		if err := session.Serve(q.session.Context()); err != nil {
+			q.logger.Debug().Err(err).Str("sessionID", sessionID.String()).Msg("session terminated")
+		}
+	}()
+	q.logger.Debug().Msgf("Registered session %v, %v, %v", sessionID, dstIP, dstPort)
 	return nil
 }
 
-// listenOriginUDP reads UDP from origin in a loop, and returns when it cannot write to edge or cannot read from origin
-func (q *QUICConnection) listenOriginUDP(sessionID uuid.UUID, conn *net.UDPConn) {
-	defer func() {
-		q.udpSessions.unregister(sessionID)
-		conn.Close()
-	}()
-	readBuffer := make([]byte, MaxDatagramFrameSize)
-	for {
-		n, err := conn.Read(readBuffer)
-		if n > 0 {
-			if n > MaxDatagramFrameSize-sessionIDLen {
-				// TODO: TUN-5302 return ICMP packet too big message
-				q.logger.Error().Msgf("Origin UDP payload has %d bytes, which exceeds transport MTU %d", n, MaxDatagramFrameSize-sessionIDLen)
-				continue
-			}
-			msgWithID, err := quicpogs.SuffixSessionID(sessionID, readBuffer[:n])
-			if err != nil {
-				q.logger.Err(err).Msg("Failed to suffix session ID to datagram, it will be dropped")
-				continue
-			}
-			if err := q.session.SendMessage(msgWithID); err != nil {
-				q.logger.Err(err).Msg("Failed to send datagram back to edge")
-				return
-			}
-		}
-		if err != nil {
-			q.logger.Err(err).Msg("Failed to read UDP from origin")
-			return
-		}
-	}
-}
+// TODO: TUN-5422 Implement UnregisterUdpSession RPC
 
 // streamReadWriteAcker is a light wrapper over QUIC streams with a callback to send response back to
 // the client.
@@ -320,3 +291,35 @@ func isTransferEncodingChunked(req *http.Request) bool {
 	// separated value as well.
 	return strings.Contains(strings.ToLower(transferEncodingVal), "chunked")
 }
+
+// TODO: TUN-5303: Define an UDPProxy in ingress package
+func (q *QUICConnection) newUDPProxy(dstIP net.IP, dstPort uint16) (*net.UDPConn, error) {
+	dstAddr := &net.UDPAddr{
+		IP:   dstIP,
+		Port: int(dstPort),
+	}
+	return net.DialUDP("udp", nil, dstAddr)
+}
+
+// TODO: TUN-5303: Find the local IP once in ingress package
+// TODO: TUN-5421 allow user to specify which IP to bind to
+func getLocalIP() (net.IP, error) {
+	addrs, err := net.InterfaceAddrs()
+	if err != nil {
+		return nil, err
+	}
+	for _, addr := range addrs {
+		// Find the IP that is not loop back
+		var ip net.IP
+		switch v := addr.(type) {
+		case *net.IPNet:
+			ip = v.IP
+		case *net.IPAddr:
+			ip = v.IP
+		}
+		if !ip.IsLoopback() {
+			return ip, nil
+		}
+	}
+	return nil, fmt.Errorf("cannot determine IP to bind to")
+}

connection/udp_session.go

@@ -1,90 +0,0 @@
-package connection
-
-import (
-	"fmt"
-	"net"
-	"sync"
-
-	"github.com/google/uuid"
-)
-
-// TODO: TUN-5422 Unregister session
-const (
-	sessionIDLen = len(uuid.UUID{})
-)
-
-type udpSessions struct {
-	lock     sync.RWMutex
-	sessions map[uuid.UUID]*net.UDPConn
-	localIP  net.IP
-}
-
-func newUDPSessions(localIP net.IP) *udpSessions {
-	return &udpSessions{
-		sessions: make(map[uuid.UUID]*net.UDPConn),
-		localIP:  localIP,
-	}
-}
-
-func (us *udpSessions) register(id uuid.UUID, dstIP net.IP, dstPort uint16) (*net.UDPConn, error) {
-	us.lock.Lock()
-	defer us.lock.Unlock()
-	dstAddr := &net.UDPAddr{
-		IP:   dstIP,
-		Port: int(dstPort),
-	}
-	conn, err := net.DialUDP("udp", us.localAddr(), dstAddr)
-	if err != nil {
-		return nil, err
-	}
-	us.sessions[id] = conn
-	return conn, nil
-}
-
-func (us *udpSessions) unregister(id uuid.UUID) {
-	us.lock.Lock()
-	defer us.lock.Unlock()
-	delete(us.sessions, id)
-}
-
-func (us *udpSessions) send(id uuid.UUID, payload []byte) error {
-	us.lock.RLock()
-	defer us.lock.RUnlock()
-	conn, ok := us.sessions[id]
-	if !ok {
-		return fmt.Errorf("session %s not found", id)
-	}
-	_, err := conn.Write(payload)
-	return err
-}
-
-func (ud *udpSessions) localAddr() *net.UDPAddr {
-	// TODO: Determine the IP to bind to
-
-	return &net.UDPAddr{
-		IP:   ud.localIP,
-		Port: 0,
-	}
-}
-
-// TODO: TUN-5421 allow user to specify which IP to bind to
-func GetLocalIP() (net.IP, error) {
-	addrs, err := net.InterfaceAddrs()
-	if err != nil {
-		return nil, err
-	}
-	for _, addr := range addrs {
-		// Find the IP that is not loop back
-		var ip net.IP
-		switch v := addr.(type) {
-		case *net.IPNet:
-			ip = v.IP
-		case *net.IPAddr:
-			ip = v.IP
-		}
-		if !ip.IsLoopback() {
-			return ip, nil
-		}
-	}
-	return nil, fmt.Errorf("cannot determine IP to bind to")
-}