TUN-8641: Expose methods to simplify V3 Datagram parsing on the edge

Author: GoncaloGarcia

quic/v3/datagram.go

@@ -27,7 +27,7 @@ const (
 	maxDatagramPayloadLen = 1280
 )
 
-func parseDatagramType(data []byte) (DatagramType, error) {
+func ParseDatagramType(data []byte) (DatagramType, error) {
 	if len(data) < datagramTypeLen {
 		return 0, ErrDatagramHeaderTooSmall
 	}
@@ -140,7 +140,7 @@ func (s *UDPSessionRegistrationDatagram) MarshalBinary() (data []byte, err error
 }
 
 func (s *UDPSessionRegistrationDatagram) UnmarshalBinary(data []byte) error {
-	datagramType, err := parseDatagramType(data)
+	datagramType, err := ParseDatagramType(data)
 	if err != nil {
 		return err
 	}
@@ -192,10 +192,10 @@ type UDPSessionPayloadDatagram struct {
 }
 
 const (
-	datagramPayloadHeaderLen = datagramTypeLen + datagramRequestIdLen
+	DatagramPayloadHeaderLen = datagramTypeLen + datagramRequestIdLen
 
 	// The maximum size that a proxied UDP payload can be in a [UDPSessionPayloadDatagram]
-	maxPayloadPlusHeaderLen = maxDatagramPayloadLen + datagramPayloadHeaderLen
+	maxPayloadPlusHeaderLen = maxDatagramPayloadLen + DatagramPayloadHeaderLen
 )
 
 // The datagram structure for UDPSessionPayloadDatagram is:
@@ -230,7 +230,7 @@ func MarshalPayloadHeaderTo(requestID RequestID, payload []byte) error {
 }
 
 func (s *UDPSessionPayloadDatagram) UnmarshalBinary(data []byte) error {
-	datagramType, err := parseDatagramType(data)
+	datagramType, err := ParseDatagramType(data)
 	if err != nil {
 		return err
 	}
@@ -332,7 +332,7 @@ func (s *UDPSessionRegistrationResponseDatagram) MarshalBinary() (data []byte, e
 }
 
 func (s *UDPSessionRegistrationResponseDatagram) UnmarshalBinary(data []byte) error {
-	datagramType, err := parseDatagramType(data)
+	datagramType, err := ParseDatagramType(data)
 	if err != nil {
 		return wrapUnmarshalErr(err)
 	}

quic/v3/muxer.go

@@ -122,7 +122,7 @@ func (c *datagramConn) Serve(ctx context.Context) error {
 
 		// Each incoming datagram will be processed in a new go routine to handle the demuxing and action associated.
 		go func() {
-			typ, err := parseDatagramType(datagram)
+			typ, err := ParseDatagramType(datagram)
 			if err != nil {
 				c.logger.Err(err).Msgf("unable to parse datagram type: %d", typ)
 				return

quic/v3/request.go

@@ -75,3 +75,15 @@ func (id RequestID) MarshalBinaryTo(data []byte) error {
 	binary.BigEndian.PutUint64(data[8:], id.lo)
 	return nil
 }
+
+func (id *RequestID) UnmarshalBinary(data []byte) error {
+	if len(data) != 16 {
+		return fmt.Errorf("invalid length slice provided to unmarshal: %d (expected 16)", len(data))
+	}
+
+	*id = RequestID{
+		binary.BigEndian.Uint64(data[:8]),
+		binary.BigEndian.Uint64(data[8:]),
+	}
+	return nil
+}

quic/v3/request_test.go

@@ -5,6 +5,8 @@ import (
 	"slices"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	v3 "github.com/cloudflare/cloudflared/quic/v3"
 )
 
@@ -48,3 +50,15 @@ func TestRequestIDParsing(t *testing.T) {
 		t.Fatalf("buf1 != buf2: %+v %+v", buf1, buf2)
 	}
 }
+
+func TestRequestID_MarshalBinary(t *testing.T) {
+	buf := make([]byte, 16)
+	err := testRequestID.MarshalBinaryTo(buf)
+	require.NoError(t, err)
+	require.Len(t, buf, 16)
+
+	parsed := v3.RequestID{}
+	err = parsed.UnmarshalBinary(buf)
+	require.NoError(t, err)
+	require.Equal(t, testRequestID, parsed)
+}

quic/v3/session.go

@@ -84,14 +84,14 @@ func (s *session) Serve(ctx context.Context) error {
 	go func() {
 		// QUIC implementation copies data to another buffer before returning https://github.com/quic-go/quic-go/blob/v0.24.0/session.go#L1967-L1975
 		// This makes it safe to share readBuffer between iterations
-		readBuffer := [maxOriginUDPPacketSize + datagramPayloadHeaderLen]byte{}
+		readBuffer := [maxOriginUDPPacketSize + DatagramPayloadHeaderLen]byte{}
 		// To perform a zero copy write when passing the datagram to the connection, we prepare the buffer with
 		// the required datagram header information. We can reuse this buffer for this session since the header is the
 		// same for the each read.
-		MarshalPayloadHeaderTo(s.id, readBuffer[:datagramPayloadHeaderLen])
+		MarshalPayloadHeaderTo(s.id, readBuffer[:DatagramPayloadHeaderLen])
 		for {
 			// Read from the origin UDP socket
-			n, err := s.origin.Read(readBuffer[datagramPayloadHeaderLen:])
+			n, err := s.origin.Read(readBuffer[DatagramPayloadHeaderLen:])
 			if errors.Is(err, net.ErrClosed) || errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
 				s.log.Debug().Msg("Session (origin) connection closed")
 			}
@@ -109,7 +109,7 @@ func (s *session) Serve(ctx context.Context) error {
 			}
 			// Sending a packet to the session does block on the [quic.Connection], however, this is okay because it
 			// will cause back-pressure to the kernel buffer if the writes are not fast enough to the edge.
-			err = s.eyeball.SendUDPSessionDatagram(readBuffer[:datagramPayloadHeaderLen+n])
+			err = s.eyeball.SendUDPSessionDatagram(readBuffer[:DatagramPayloadHeaderLen+n])
 			if err != nil {
 				s.closeChan <- err
 				return