TUN-6855: Add DatagramV2Type for IP packet with trace and tracing spans

Author: chungthuang

connection/quic.go

@@ -93,7 +93,8 @@ func NewQUICConnection(
 	sessionDemuxChan := make(chan *packet.Session, demuxChanCapacity)
 	datagramMuxer := quicpogs.NewDatagramMuxerV2(session, logger, sessionDemuxChan)
 	sessionManager := datagramsession.NewManager(logger, datagramMuxer.SendToSession, sessionDemuxChan)
-	packetRouter := packet.NewRouter(packetRouterConfig, datagramMuxer, &returnPipe{muxer: datagramMuxer}, logger, orchestrator.WarpRoutingEnabled)
+	muxer := muxerWrapper{muxer: datagramMuxer}
+	packetRouter := packet.NewRouter(packetRouterConfig, &muxer, &muxer, logger, orchestrator.WarpRoutingEnabled)
 
 	return &QUICConnection{
 		session:              session,
@@ -498,16 +499,28 @@ func (np *nopCloserReadWriter) Close() error {
 	return nil
 }
 
-// returnPipe wraps DatagramMuxerV2 to satisfy the packet.FunnelUniPipe interface
-type returnPipe struct {
+// muxerWrapper wraps DatagramMuxerV2 to satisfy the packet.FunnelUniPipe interface
+type muxerWrapper struct {
 	muxer *quicpogs.DatagramMuxerV2
 }
 
-func (rp *returnPipe) SendPacket(dst netip.Addr, pk packet.RawPacket) error {
-	return rp.muxer.SendPacket(pk)
+func (rp *muxerWrapper) SendPacket(dst netip.Addr, pk packet.RawPacket) error {
+	return rp.muxer.SendPacket(quicpogs.RawPacket(pk))
 }
 
-func (rp *returnPipe) Close() error {
+func (rp *muxerWrapper) ReceivePacket(ctx context.Context) (packet.RawPacket, error) {
+	pk, err := rp.muxer.ReceivePacket(ctx)
+	if err != nil {
+		return packet.RawPacket{}, err
+	}
+	rawPacket, ok := pk.(quicpogs.RawPacket)
+	if ok {
+		return packet.RawPacket(rawPacket), nil
+	}
+	return packet.RawPacket{}, fmt.Errorf("unexpected packet type %+v", pk)
+}
+
+func (rp *muxerWrapper) Close() error {
 	return nil
 }
 

quic/datagram.go

@@ -113,9 +113,12 @@ func extractSessionID(b []byte) (uuid.UUID, []byte, error) {
 // SuffixSessionID appends the session ID at the end of the payload. Suffix is more performant than prefix because
 // the payload slice might already have enough capacity to append the session ID at the end
 func SuffixSessionID(sessionID uuid.UUID, b []byte) ([]byte, error) {
-	if len(b)+len(sessionID) > MaxDatagramFrameSize {
+	return suffixMetadata(b, sessionID[:])
+}
+
+func suffixMetadata(payload, metadata []byte) ([]byte, error) {
+	if len(payload)+len(metadata) > MaxDatagramFrameSize {
 		return nil, fmt.Errorf("datagram size exceed %d", MaxDatagramFrameSize)
 	}
-	b = append(b, sessionID[:]...)
-	return b, nil
+	return append(payload, metadata...), nil
 }

quic/datagram_test.go

@@ -1,6 +1,7 @@
 package quic
 
 import (
+	"bytes"
 	"context"
 	"crypto/rand"
 	"crypto/rsa"
@@ -23,6 +24,7 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/cloudflare/cloudflared/packet"
+	"github.com/cloudflare/cloudflared/tracing"
 )
 
 var (
@@ -121,6 +123,15 @@ func testDatagram(t *testing.T, version uint8, sessionToPayloads []*packet.Sessi
 
 	logger := zerolog.Nop()
 
+	tracingIdentity, err := tracing.NewIdentity("ec31ad8a01fde11fdcabe2efdce36873:52726f6cabc144f5:0:1")
+	require.NoError(t, err)
+	serializedTracingID, err := tracingIdentity.MarshalBinary()
+	require.NoError(t, err)
+	tracingSpan := &TracingSpanPacket{
+		Spans:           []byte("tracing"),
+		TracingIdentity: serializedTracingID,
+	}
+
 	errGroup, ctx := errgroup.WithContext(context.Background())
 	// Run edge side of datagram muxer
 	errGroup.Go(func() error {
@@ -140,18 +151,17 @@ func testDatagram(t *testing.T, version uint8, sessionToPayloads []*packet.Sessi
 			muxer := NewDatagramMuxerV2(quicSession, &logger, sessionDemuxChan)
 			muxer.ServeReceive(ctx)
 
-			icmpDecoder := packet.NewICMPDecoder()
 			for _, pk := range packets {
 				received, err := muxer.ReceivePacket(ctx)
 				require.NoError(t, err)
-
-				receivedICMP, err := icmpDecoder.Decode(received)
+				validateIPPacket(t, received, &pk)
+				received, err = muxer.ReceivePacket(ctx)
 				require.NoError(t, err)
-				require.Equal(t, pk.IP, receivedICMP.IP)
-				require.Equal(t, pk.Type, receivedICMP.Type)
-				require.Equal(t, pk.Code, receivedICMP.Code)
-				require.Equal(t, pk.Body, receivedICMP.Body)
+				validateIPPacketWithTracing(t, received, &pk, serializedTracingID)
 			}
+			received, err := muxer.ReceivePacket(ctx)
+			require.NoError(t, err)
+			validateTracingSpans(t, received, tracingSpan)
 		default:
 			return fmt.Errorf("unknown datagram version %d", version)
 		}
@@ -188,10 +198,15 @@ func testDatagram(t *testing.T, version uint8, sessionToPayloads []*packet.Sessi
 			for _, pk := range packets {
 				encodedPacket, err := encoder.Encode(&pk)
 				require.NoError(t, err)
-				require.NoError(t, muxerV2.SendPacket(encodedPacket))
+				require.NoError(t, muxerV2.SendPacket(RawPacket(encodedPacket)))
+				require.NoError(t, muxerV2.SendPacket(&TracedPacket{
+					Packet:          encodedPacket,
+					TracingIdentity: serializedTracingID,
+				}))
 			}
+			require.NoError(t, muxerV2.SendPacket(tracingSpan))
 			// Payload larger than transport MTU, should not be sent
-			require.Error(t, muxerV2.SendPacket(packet.RawPacket{
+			require.Error(t, muxerV2.SendPacket(RawPacket{
 				Data: largePayload,
 			}))
 			muxer = muxerV2
@@ -217,6 +232,38 @@ func testDatagram(t *testing.T, version uint8, sessionToPayloads []*packet.Sessi
 	require.NoError(t, errGroup.Wait())
 }
 
+func validateIPPacket(t *testing.T, receivedPacket Packet, expectedICMP *packet.ICMP) {
+	require.Equal(t, DatagramTypeIP, receivedPacket.Type())
+	rawPacket := receivedPacket.(RawPacket)
+	decoder := packet.NewICMPDecoder()
+	receivedICMP, err := decoder.Decode(packet.RawPacket(rawPacket))
+	require.NoError(t, err)
+	validateICMP(t, expectedICMP, receivedICMP)
+}
+
+func validateIPPacketWithTracing(t *testing.T, receivedPacket Packet, expectedICMP *packet.ICMP, serializedTracingID []byte) {
+	require.Equal(t, DatagramTypeIPWithTrace, receivedPacket.Type())
+	tracedPacket := receivedPacket.(*TracedPacket)
+	decoder := packet.NewICMPDecoder()
+	receivedICMP, err := decoder.Decode(tracedPacket.Packet)
+	require.NoError(t, err)
+	validateICMP(t, expectedICMP, receivedICMP)
+	require.True(t, bytes.Equal(tracedPacket.TracingIdentity, serializedTracingID))
+}
+
+func validateICMP(t *testing.T, expected, actual *packet.ICMP) {
+	require.Equal(t, expected.IP, actual.IP)
+	require.Equal(t, expected.Type, actual.Type)
+	require.Equal(t, expected.Code, actual.Code)
+	require.Equal(t, expected.Body, actual.Body)
+}
+
+func validateTracingSpans(t *testing.T, receivedPacket Packet, expectedSpan *TracingSpanPacket) {
+	require.Equal(t, DatagramTypeTracingSpan, receivedPacket.Type())
+	tracingSpans := receivedPacket.(*TracingSpanPacket)
+	require.Equal(t, tracingSpans, expectedSpan)
+}
+
 func newQUICListener(t *testing.T, config *quic.Config) quic.Listener {
 	// Create a simple tls config.
 	tlsConfig := generateTLSConfig()

quic/datagramv2.go

@@ -9,15 +9,28 @@ import (
 	"github.com/rs/zerolog"
 
 	"github.com/cloudflare/cloudflared/packet"
+	"github.com/cloudflare/cloudflared/tracing"
 )
 
 type DatagramV2Type byte
 
 const (
+	// UDP payload
 	DatagramTypeUDP DatagramV2Type = iota
+	// Full IP packet
 	DatagramTypeIP
+	// DatagramTypeIP + tracing ID
+	DatagramTypeIPWithTrace
+	// Tracing spans in protobuf format
+	DatagramTypeTracingSpan
 )
 
+type Packet interface {
+	Type() DatagramV2Type
+	Payload() []byte
+	Metadata() []byte
+}
+
 const (
 	typeIDLen = 1
 	// Same as sessionDemuxChan capacity
@@ -41,7 +54,7 @@ type DatagramMuxerV2 struct {
 	session          quic.Connection
 	logger           *zerolog.Logger
 	sessionDemuxChan chan<- *packet.Session
-	packetDemuxChan  chan packet.RawPacket
+	packetDemuxChan  chan Packet
 }
 
 func NewDatagramMuxerV2(
@@ -54,7 +67,7 @@ func NewDatagramMuxerV2(
 		session:          quicSession,
 		logger:           &logger,
 		sessionDemuxChan: sessionDemuxChan,
-		packetDemuxChan:  make(chan packet.RawPacket, packetChanCapacity),
+		packetDemuxChan:  make(chan Packet, packetChanCapacity),
 	}
 }
 
@@ -79,14 +92,19 @@ func (dm *DatagramMuxerV2) SendToSession(session *packet.Session) error {
 	return nil
 }
 
-// SendPacket 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) SendPacket(pk packet.RawPacket) error {
-	payloadWithVersion, err := SuffixType(pk.Data, DatagramTypeIP)
+// SendPacket sends a packet with datagram version in the suffix. If ctx is a TracedContext, it adds the tracing
+// context between payload and datagram version.
+// 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) SendPacket(pk Packet) error {
+	payloadWithMetadata, err := suffixMetadata(pk.Payload(), pk.Metadata())
+	if err != nil {
+		return err
+	}
+	payloadWithMetadataAndType, err := SuffixType(payloadWithMetadata, pk.Type())
 	if err != nil {
 		return errors.Wrap(err, "Failed to suffix datagram type, it will be dropped")
 	}
-	if err := dm.session.SendMessage(payloadWithVersion); err != nil {
+	if err := dm.session.SendMessage(payloadWithMetadataAndType); err != nil {
 		return errors.Wrap(err, "Failed to send datagram back to edge")
 	}
 	return nil
@@ -108,10 +126,10 @@ func (dm *DatagramMuxerV2) ServeReceive(ctx context.Context) error {
 	}
 }
 
-func (dm *DatagramMuxerV2) ReceivePacket(ctx context.Context) (packet.RawPacket, error) {
+func (dm *DatagramMuxerV2) ReceivePacket(ctx context.Context) (pk Packet, err error) {
 	select {
 	case <-ctx.Done():
-		return packet.RawPacket{}, ctx.Err()
+		return nil, ctx.Err()
 	case pk := <-dm.packetDemuxChan:
 		return pk, nil
 	}
@@ -126,10 +144,8 @@ func (dm *DatagramMuxerV2) demux(ctx context.Context, msgWithType []byte) error
 	switch msgType {
 	case DatagramTypeUDP:
 		return dm.handleSession(ctx, msg)
-	case DatagramTypeIP:
-		return dm.handlePacket(ctx, msg)
 	default:
-		return fmt.Errorf("Unexpected datagram type %d", msgType)
+		return dm.handlePacket(ctx, msg, msgType)
 	}
 }
 
@@ -150,13 +166,93 @@ func (dm *DatagramMuxerV2) handleSession(ctx context.Context, session []byte) er
 	}
 }
 
-func (dm *DatagramMuxerV2) handlePacket(ctx context.Context, pk []byte) error {
+func (dm *DatagramMuxerV2) handlePacket(ctx context.Context, pk []byte, msgType DatagramV2Type) error {
+	var demuxedPacket Packet
+	switch msgType {
+	case DatagramTypeIP:
+		demuxedPacket = RawPacket(packet.RawPacket{Data: pk})
+	case DatagramTypeIPWithTrace:
+		tracingIdentity, payload, err := extractTracingIdentity(pk)
+		if err != nil {
+			return err
+		}
+		demuxedPacket = &TracedPacket{
+			Packet:          packet.RawPacket{Data: payload},
+			TracingIdentity: tracingIdentity,
+		}
+	case DatagramTypeTracingSpan:
+		tracingIdentity, spans, err := extractTracingIdentity(pk)
+		if err != nil {
+			return err
+		}
+		demuxedPacket = &TracingSpanPacket{
+			Spans:           spans,
+			TracingIdentity: tracingIdentity,
+		}
+	default:
+		return fmt.Errorf("Unexpected datagram type %d", msgType)
+	}
 	select {
 	case <-ctx.Done():
 		return ctx.Err()
-	case dm.packetDemuxChan <- packet.RawPacket{
-		Data: pk,
-	}:
+	case dm.packetDemuxChan <- demuxedPacket:
 		return nil
 	}
 }
+
+func extractTracingIdentity(pk []byte) (tracingIdentity []byte, payload []byte, err error) {
+	if len(pk) < tracing.IdentityLength {
+		return nil, nil, fmt.Errorf("packet with tracing context should have at least %d bytes, got %v", tracing.IdentityLength, pk)
+	}
+	tracingIdentity = pk[len(pk)-tracing.IdentityLength:]
+	payload = pk[:len(pk)-tracing.IdentityLength]
+	return tracingIdentity, payload, nil
+}
+
+type RawPacket packet.RawPacket
+
+func (rw RawPacket) Type() DatagramV2Type {
+	return DatagramTypeIP
+}
+
+func (rw RawPacket) Payload() []byte {
+	return rw.Data
+}
+
+func (rw RawPacket) Metadata() []byte {
+	return []byte{}
+}
+
+type TracedPacket struct {
+	Packet          packet.RawPacket
+	TracingIdentity []byte
+}
+
+func (tp *TracedPacket) Type() DatagramV2Type {
+	return DatagramTypeIPWithTrace
+}
+
+func (tp *TracedPacket) Payload() []byte {
+	return tp.Packet.Data
+}
+
+func (tp *TracedPacket) Metadata() []byte {
+	return tp.TracingIdentity
+}
+
+type TracingSpanPacket struct {
+	Spans           []byte
+	TracingIdentity []byte
+}
+
+func (tsp *TracingSpanPacket) Type() DatagramV2Type {
+	return DatagramTypeTracingSpan
+}
+
+func (tsp *TracingSpanPacket) Payload() []byte {
+	return tsp.Spans
+}
+
+func (tsp *TracingSpanPacket) Metadata() []byte {
+	return tsp.TracingIdentity
+}