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

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
+}