TUN-6584: Define QUIC datagram v2 format to support proxying IP packets

quic/datagram.go

@@ -1,6 +1,7 @@
 package quic
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/google/uuid"
@@ -13,53 +14,88 @@ const (
 	sessionIDLen = len(uuid.UUID{})
 )
 
+type SessionDatagram struct {
+	ID      uuid.UUID
+	Payload []byte
+}
+
+type BaseDatagramMuxer interface {
+	// MuxSession suffix the session ID to the payload so the other end of the QUIC connection can demultiplex the
+	// payload from multiple datagram sessions
+	MuxSession(sessionID uuid.UUID, payload []byte) error
+	// ServeReceive starts a loop to receive datagrams from the QUIC connection
+	ServeReceive(ctx context.Context) error
+}
+
 type DatagramMuxer struct {
-	session quic.Connection
-	logger  *zerolog.Logger
+	session   quic.Connection
+	logger    *zerolog.Logger
+	demuxChan chan<- *SessionDatagram
 }
 
-func NewDatagramMuxer(quicSession quic.Connection, logger *zerolog.Logger) (*DatagramMuxer, error) {
+func NewDatagramMuxer(quicSession quic.Connection, log *zerolog.Logger, demuxChan chan<- *SessionDatagram) *DatagramMuxer {
+	logger := log.With().Uint8("datagramVersion", 1).Logger()
 	return &DatagramMuxer{
-		session: quicSession,
-		logger:  logger,
-	}, nil
+		session:   quicSession,
+		logger:    &logger,
+		demuxChan: demuxChan,
+	}
 }
 
-// SendTo suffix the session ID to the payload so the other end of the QUIC session can demultiplex
-// the payload from multiple datagram sessions
-func (dm *DatagramMuxer) SendTo(sessionID uuid.UUID, payload []byte) error {
-	if len(payload) > maxDatagramPayloadSize {
+// Maximum application payload to send to / receive from QUIC datagram frame
+func (dm *DatagramMuxer) mtu() int {
+	return maxDatagramPayloadSize
+}
+
+func (dm *DatagramMuxer) MuxSession(sessionID uuid.UUID, payload []byte) error {
+	if len(payload) > dm.mtu() {
 		// TODO: TUN-5302 return ICMP packet too big message
-		return fmt.Errorf("origin UDP payload has %d bytes, which exceeds transport MTU %d", len(payload), dm.MTU())
+		// drop packet for now, eventually reply with ICMP for PMTUD
+		return fmt.Errorf("origin UDP payload has %d bytes, which exceeds transport MTU %d", len(payload), dm.mtu())
 	}
-	msgWithID, err := suffixSessionID(sessionID, payload)
+	payloadWithMetadata, err := suffixSessionID(sessionID, payload)
 	if err != nil {
 		return errors.Wrap(err, "Failed to suffix session ID to datagram, it will be dropped")
 	}
-	if err := dm.session.SendMessage(msgWithID); err != nil {
+	if err := dm.session.SendMessage(payloadWithMetadata); err != nil {
 		return errors.Wrap(err, "Failed to send datagram back to edge")
 	}
 	return nil
 }
 
-// ReceiveFrom extracts datagram session ID, then sends the session ID and payload to session manager
-// which determines how to proxy to the origin. It assumes the datagram session has already been
-// registered with session manager through other side channel
-func (dm *DatagramMuxer) ReceiveFrom() (uuid.UUID, []byte, error) {
-	msg, err := dm.session.ReceiveMessage()
-	if err != nil {
-		return uuid.Nil, nil, err
+func (dm *DatagramMuxer) ServeReceive(ctx context.Context) error {
+	for {
+		// Extracts datagram session ID, then sends the session ID and payload to receiver
+		// which determines how to proxy to the origin. It assumes the datagram session has already been
+		// registered with receiver through other side channel
+		msg, err := dm.session.ReceiveMessage()
+		if err != nil {
+			return err
+		}
+		if err := dm.demux(ctx, msg); err != nil {
+			dm.logger.Error().Err(err).Msg("Failed to demux datagram")
+			if err == context.Canceled {
+				return err
+			}
+		}
 	}
-	sessionID, payload, err := extractSessionID(msg)
-	if err != nil {
-		return uuid.Nil, nil, err
-	}
-	return sessionID, payload, nil
 }
 
-// Maximum application payload to send to / receive from QUIC datagram frame
-func (dm *DatagramMuxer) MTU() int {
-	return maxDatagramPayloadSize
+func (dm *DatagramMuxer) demux(ctx context.Context, msg []byte) error {
+	sessionID, payload, err := extractSessionID(msg)
+	if err != nil {
+		return err
+	}
+	sessionDatagram := SessionDatagram{
+		ID:      sessionID,
+		Payload: payload,
+	}
+	select {
+	case dm.demuxChan <- &sessionDatagram:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
 // Each QUIC datagram should be suffixed with session ID.

quic/datagram_test.go

@@ -8,6 +8,7 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/pem"
+	"fmt"
 	"math/big"
 	"testing"
 	"time"
@@ -52,9 +53,29 @@ func TestSuffixSessionIDError(t *testing.T) {
 	require.Error(t, err)
 }
 
-func TestMaxDatagramPayload(t *testing.T) {
-	payload := make([]byte, maxDatagramPayloadSize)
+func TestDatagram(t *testing.T) {
+	maxPayload := make([]byte, maxDatagramPayloadSize)
+	noPayloadSession := uuid.New()
+	maxPayloadSession := uuid.New()
+	sessionToPayload := []*SessionDatagram{
+		{
+			ID:      noPayloadSession,
+			Payload: make([]byte, 0),
+		},
+		{
+			ID:      maxPayloadSession,
+			Payload: maxPayload,
+		},
+	}
+	flowPayloads := [][]byte{
+		maxPayload,
+	}
 
+	testDatagram(t, 1, sessionToPayload, nil)
+	testDatagram(t, 2, sessionToPayload, flowPayloads)
+}
+
+func testDatagram(t *testing.T, version uint8, sessionToPayloads []*SessionDatagram, packetPayloads [][]byte) {
 	quicConfig := &quic.Config{
 		KeepAlivePeriod:      5 * time.Millisecond,
 		EnableDatagrams:      true,
@@ -63,6 +84,8 @@ func TestMaxDatagramPayload(t *testing.T) {
 	quicListener := newQUICListener(t, quicConfig)
 	defer quicListener.Close()
 
+	logger := zerolog.Nop()
+
 	errGroup, ctx := errgroup.WithContext(context.Background())
 	// Run edge side of datagram muxer
 	errGroup.Go(func() error {
@@ -72,22 +95,32 @@ func TestMaxDatagramPayload(t *testing.T) {
 			return err
 		}
 
-		logger := zerolog.Nop()
-		muxer, err := NewDatagramMuxer(quicSession, &logger)
-		if err != nil {
-			return err
+		sessionDemuxChan := make(chan *SessionDatagram, 16)
+
+		switch version {
+		case 1:
+			muxer := NewDatagramMuxer(quicSession, &logger, sessionDemuxChan)
+			muxer.ServeReceive(ctx)
+		case 2:
+			packetDemuxChan := make(chan []byte, len(packetPayloads))
+			muxer := NewDatagramMuxerV2(quicSession, &logger, sessionDemuxChan, packetDemuxChan)
+			muxer.ServeReceive(ctx)
+
+			for _, expectedPayload := range packetPayloads {
+				require.Equal(t, expectedPayload, <-packetDemuxChan)
+			}
+		default:
+			return fmt.Errorf("unknown datagram version %d", version)
 		}
 
-		sessionID, receivedPayload, err := muxer.ReceiveFrom()
-		if err != nil {
-			return err
+		for _, expectedPayload := range sessionToPayloads {
+			actualPayload := <-sessionDemuxChan
+			require.Equal(t, expectedPayload, actualPayload)
 		}
-		require.Equal(t, testSessionID, sessionID)
-		require.True(t, bytes.Equal(payload, receivedPayload))
-
 		return nil
 	})
 
+	largePayload := make([]byte, MaxDatagramFrameSize)
 	// Run cloudflared side of datagram muxer
 	errGroup.Go(func() error {
 		tlsClientConfig := &tls.Config{
@@ -97,24 +130,35 @@ func TestMaxDatagramPayload(t *testing.T) {
 		// Establish quic connection
 		quicSession, err := quic.DialAddrEarly(quicListener.Addr().String(), tlsClientConfig, quicConfig)
 		require.NoError(t, err)
-
-		logger := zerolog.Nop()
-		muxer, err := NewDatagramMuxer(quicSession, &logger)
-		if err != nil {
-			return err
-		}
+		defer quicSession.CloseWithError(0, "")
 
 		// Wait a few milliseconds for MTU discovery to take place
 		time.Sleep(time.Millisecond * 100)
-		err = muxer.SendTo(testSessionID, payload)
-		if err != nil {
-			return err
+
+		var muxer BaseDatagramMuxer
+		switch version {
+		case 1:
+			muxer = NewDatagramMuxer(quicSession, &logger, nil)
+		case 2:
+			muxerV2 := NewDatagramMuxerV2(quicSession, &logger, nil, nil)
+			for _, payload := range packetPayloads {
+				require.NoError(t, muxerV2.MuxPacket(payload))
+			}
+			// Payload larger than transport MTU, should not be sent
+			require.Error(t, muxerV2.MuxPacket(largePayload))
+			muxer = muxerV2
+		default:
+			return fmt.Errorf("unknown datagram version %d", version)
 		}
 
-		// Payload larger than transport MTU, should return an error
-		largePayload := make([]byte, MaxDatagramFrameSize)
-		err = muxer.SendTo(testSessionID, largePayload)
-		require.Error(t, err)
+		for _, sessionDatagram := range sessionToPayloads {
+			require.NoError(t, muxer.MuxSession(sessionDatagram.ID, sessionDatagram.Payload))
+		}
+		// Payload larger than transport MTU, should not be sent
+		require.Error(t, muxer.MuxSession(testSessionID, largePayload))
+
+		// Wait for edge to finish receiving the messages
+		time.Sleep(time.Millisecond * 100)
 
 		return nil
 	})
@@ -154,3 +198,35 @@ func generateTLSConfig() *tls.Config {
 		NextProtos:   []string{"argotunnel"},
 	}
 }
+
+type sessionMuxer interface {
+	SendToSession(sessionID uuid.UUID, payload []byte) error
+}
+
+type mockSessionReceiver struct {
+	expectedSessionToPayload map[uuid.UUID][]byte
+	receivedCount            int
+}
+
+func (msr *mockSessionReceiver) ReceiveDatagram(sessionID uuid.UUID, payload []byte) error {
+	expectedPayload := msr.expectedSessionToPayload[sessionID]
+	if !bytes.Equal(expectedPayload, payload) {
+		return fmt.Errorf("expect %v to have payload %s, got %s", sessionID, string(expectedPayload), string(payload))
+	}
+	msr.receivedCount++
+	return nil
+}
+
+type mockFlowReceiver struct {
+	expectedPayloads [][]byte
+	receivedCount    int
+}
+
+func (mfr *mockFlowReceiver) ReceiveFlow(payload []byte) error {
+	expectedPayload := mfr.expectedPayloads[mfr.receivedCount]
+	if !bytes.Equal(expectedPayload, payload) {
+		return fmt.Errorf("expect flow %d to have payload %s, got %s", mfr.receivedCount, string(expectedPayload), string(payload))
+	}
+	mfr.receivedCount++
+	return nil
+}

quic/datagramv2.go

@@ -0,0 +1,136 @@
+package quic
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/google/uuid"
+	"github.com/lucas-clemente/quic-go"
+	"github.com/pkg/errors"
+	"github.com/rs/zerolog"
+)
+
+type datagramV2Type byte
+
+const (
+	udp datagramV2Type = iota
+	ip
+)
+
+func suffixType(b []byte, datagramType datagramV2Type) ([]byte, error) {
+	if len(b)+1 > MaxDatagramFrameSize {
+		return nil, fmt.Errorf("datagram size %d exceeds max frame size %d", len(b), MaxDatagramFrameSize)
+	}
+	b = append(b, byte(datagramType))
+	return b, nil
+}
+
+// Maximum application payload to send to / receive from QUIC datagram frame
+func (dm *DatagramMuxerV2) mtu() int {
+	return maxDatagramPayloadSize
+}
+
+type DatagramMuxerV2 struct {
+	session          quic.Connection
+	logger           *zerolog.Logger
+	sessionDemuxChan chan<- *SessionDatagram
+	packetDemuxChan  chan<- []byte
+}
+
+func NewDatagramMuxerV2(
+	quicSession quic.Connection,
+	log *zerolog.Logger,
+	sessionDemuxChan chan<- *SessionDatagram,
+	packetDemuxChan chan<- []byte) *DatagramMuxerV2 {
+	logger := log.With().Uint8("datagramVersion", 2).Logger()
+	return &DatagramMuxerV2{
+		session:          quicSession,
+		logger:           &logger,
+		sessionDemuxChan: sessionDemuxChan,
+		packetDemuxChan:  packetDemuxChan,
+	}
+}
+
+// MuxSession suffix the session ID and datagram version to the payload so the other end of the QUIC connection can
+// demultiplex the payload from multiple datagram sessions
+func (dm *DatagramMuxerV2) MuxSession(sessionID uuid.UUID, payload []byte) error {
+	if len(payload) > dm.mtu() {
+		// TODO: TUN-5302 return ICMP packet too big message
+		return fmt.Errorf("origin UDP payload has %d bytes, which exceeds transport MTU %d", len(payload), dm.mtu())
+	}
+	msgWithID, err := suffixSessionID(sessionID, payload)
+	if err != nil {
+		return errors.Wrap(err, "Failed to suffix session ID to datagram, it will be dropped")
+	}
+	msgWithIDAndType, err := suffixType(msgWithID, udp)
+	if err != nil {
+		return errors.Wrap(err, "Failed to suffix datagram type, it will be dropped")
+	}
+	if err := dm.session.SendMessage(msgWithIDAndType); err != nil {
+		return errors.Wrap(err, "Failed to send datagram back to edge")
+	}
+	return nil
+}
+
+// MuxPacket suffix the datagram type to the packet. The other end of the QUIC connection can demultiplex by parsing
+// the payload as IP and look at the source and destination.
+func (dm *DatagramMuxerV2) MuxPacket(packet []byte) error {
+	payloadWithVersion, err := suffixType(packet, ip)
+	if err != nil {
+		return errors.Wrap(err, "Failed to suffix datagram type, it will be dropped")
+	}
+	if err := dm.session.SendMessage(payloadWithVersion); err != nil {
+		return errors.Wrap(err, "Failed to send datagram back to edge")
+	}
+	return nil
+}
+
+// Demux reads datagrams from the QUIC connection and demuxes depending on whether it's a session or packet
+func (dm *DatagramMuxerV2) ServeReceive(ctx context.Context) error {
+	for {
+		msg, err := dm.session.ReceiveMessage()
+		if err != nil {
+			return err
+		}
+		if err := dm.demux(ctx, msg); err != nil {
+			dm.logger.Error().Err(err).Msg("Failed to demux datagram")
+			if err == context.Canceled {
+				return err
+			}
+		}
+	}
+}
+
+func (dm *DatagramMuxerV2) demux(ctx context.Context, msgWithType []byte) error {
+	if len(msgWithType) < 1 {
+		return fmt.Errorf("QUIC datagram should have at least 1 byte")
+	}
+	msgType := datagramV2Type(msgWithType[len(msgWithType)-1])
+	msg := msgWithType[0 : len(msgWithType)-1]
+	switch msgType {
+	case udp:
+		sessionID, payload, err := extractSessionID(msg)
+		if err != nil {
+			return err
+		}
+		sessionDatagram := SessionDatagram{
+			ID:      sessionID,
+			Payload: payload,
+		}
+		select {
+		case dm.sessionDemuxChan <- &sessionDatagram:
+			return nil
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	case ip:
+		select {
+		case dm.packetDemuxChan <- msg:
+			return nil
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	default:
+		return fmt.Errorf("Unexpected datagram type %d", msgType)
+	}
+}