analyzer: make http3/quic handling more reliable

analyzer/udp/internal/quic/payload.go

@@ -13,6 +13,19 @@ import (
 )
 
 func ReadCryptoPayload(packet []byte) ([]byte, error) {
+	frs, err := ReadCryptoFrames(packet)
+	if err != nil {
+		return nil, err
+	}
+	data := assembleCryptoFrames(frs)
+	if data == nil {
+		return nil, errors.New("unable to assemble crypto frames")
+	}
+	return data, nil
+}
+
+// ReadCryptoFrames decrypts a QUIC Initial client packet and returns CRYPTO frames.
+func ReadCryptoFrames(packet []byte) ([]CryptoFrame, error) {
 	hdr, offset, err := ParseInitialHeader(packet)
 	if err != nil {
 		return nil, err
@@ -46,11 +59,7 @@ func ReadCryptoPayload(packet []byte) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	data := assembleCryptoFrames(frs)
+	return frs, nil
-	if data == nil {
-		return nil, errors.New("unable to assemble crypto frames")
-	}
-	return data, nil
 }
 
 const (
@@ -59,25 +68,78 @@ const (
 	cryptoFrameType  = 0x06
 )
 
-type cryptoFrame struct {
+type CryptoFrame struct {
 	Offset int64
 	Data   []byte
 }
 
-func extractCryptoFrames(r *bytes.Reader) ([]cryptoFrame, error) {
+func extractCryptoFrames(r *bytes.Reader) ([]CryptoFrame, error) {
-	var frames []cryptoFrame
+	var frames []CryptoFrame
 	for r.Len() > 0 {
 		typ, err := quicvarint.Read(r)
 		if err != nil {
 			return nil, err
 		}
-		if typ == paddingFrameType || typ == pingFrameType {
+		switch typ {
+		case paddingFrameType, pingFrameType, 0x1e:
+			// PADDING, PING, HANDSHAKE_DONE: no payload.
 			continue
+		case 0x02, 0x03:
+			// ACK, ACK_ECN
+			if _, err := quicvarint.Read(r); err != nil { // Largest Acknowledged
+				return nil, err
 			}
-		if typ != cryptoFrameType {
+			if _, err := quicvarint.Read(r); err != nil { // ACK Delay
-			return nil, fmt.Errorf("encountered unexpected frame type: %d", typ)
+				return nil, err
 			}
-		var frame cryptoFrame
+			ackRangeCount, err := quicvarint.Read(r)
+			if err != nil {
+				return nil, err
+			}
+			if _, err := quicvarint.Read(r); err != nil { // First ACK Range
+				return nil, err
+			}
+			for i := uint64(0); i < ackRangeCount; i++ {
+				if _, err := quicvarint.Read(r); err != nil { // Gap
+					return nil, err
+				}
+				if _, err := quicvarint.Read(r); err != nil { // ACK Range Length
+					return nil, err
+				}
+			}
+			if typ == 0x03 {
+				if _, err := quicvarint.Read(r); err != nil { // ECT0 Count
+					return nil, err
+				}
+				if _, err := quicvarint.Read(r); err != nil { // ECT1 Count
+					return nil, err
+				}
+				if _, err := quicvarint.Read(r); err != nil { // ECN-CE Count
+					return nil, err
+				}
+			}
+		case 0x04:
+			// RESET_STREAM
+			if _, err := quicvarint.Read(r); err != nil { // Stream ID
+				return nil, err
+			}
+			if _, err := quicvarint.Read(r); err != nil { // Application Error Code
+				return nil, err
+			}
+			if _, err := quicvarint.Read(r); err != nil { // Final Size
+				return nil, err
+			}
+		case 0x05:
+			// STOP_SENDING
+			if _, err := quicvarint.Read(r); err != nil { // Stream ID
+				return nil, err
+			}
+			if _, err := quicvarint.Read(r); err != nil { // Application Error Code
+				return nil, err
+			}
+		case cryptoFrameType:
+			// CRYPTO
+			var frame CryptoFrame
 			offset, err := quicvarint.Read(r)
 			if err != nil {
 				return nil, err
@@ -92,13 +154,139 @@ func extractCryptoFrames(r *bytes.Reader) ([]cryptoFrame, error) {
 				return nil, err
 			}
 			frames = append(frames, frame)
+		case 0x07:
+			// NEW_TOKEN
+			tokenLen, err := quicvarint.Read(r)
+			if err != nil {
+				return nil, err
+			}
+			if err := skipN(r, tokenLen); err != nil {
+				return nil, err
+			}
+		case 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f:
+			// STREAM
+			if _, err := quicvarint.Read(r); err != nil { // Stream ID
+				return nil, err
+			}
+			hasOffset := typ&0x04 != 0
+			hasLength := typ&0x02 != 0
+			if hasOffset {
+				if _, err := quicvarint.Read(r); err != nil { // Offset
+					return nil, err
+				}
+			}
+			var dataLen uint64
+			if hasLength {
+				n, err := quicvarint.Read(r)
+				if err != nil {
+					return nil, err
+				}
+				dataLen = n
+			} else {
+				dataLen = uint64(r.Len())
+			}
+			if err := skipN(r, dataLen); err != nil {
+				return nil, err
+			}
+		case 0x10, 0x12, 0x13, 0x14, 0x16, 0x17, 0x19:
+			// MAX_DATA, MAX_STREAMS_*, DATA_BLOCKED, STREAMS_BLOCKED_*, RETIRE_CONNECTION_ID
+			if _, err := quicvarint.Read(r); err != nil {
+				return nil, err
+			}
+		case 0x11, 0x15:
+			// MAX_STREAM_DATA, STREAM_DATA_BLOCKED
+			if _, err := quicvarint.Read(r); err != nil {
+				return nil, err
+			}
+			if _, err := quicvarint.Read(r); err != nil {
+				return nil, err
+			}
+		case 0x18:
+			// NEW_CONNECTION_ID
+			if _, err := quicvarint.Read(r); err != nil { // Sequence Number
+				return nil, err
+			}
+			if _, err := quicvarint.Read(r); err != nil { // Retire Prior To
+				return nil, err
+			}
+			cidLen, err := r.ReadByte()
+			if err != nil {
+				return nil, err
+			}
+			if cidLen > 20 {
+				return nil, fmt.Errorf("invalid connection ID length: %d", cidLen)
+			}
+			if err := skipN(r, uint64(cidLen)); err != nil { // Connection ID
+				return nil, err
+			}
+			if err := skipN(r, 16); err != nil { // Stateless Reset Token
+				return nil, err
+			}
+		case 0x1a, 0x1b:
+			// PATH_CHALLENGE, PATH_RESPONSE
+			if err := skipN(r, 8); err != nil {
+				return nil, err
+			}
+		case 0x1c:
+			// CONNECTION_CLOSE (transport)
+			if _, err := quicvarint.Read(r); err != nil { // Error Code
+				return nil, err
+			}
+			if _, err := quicvarint.Read(r); err != nil { // Frame Type
+				return nil, err
+			}
+			reasonLen, err := quicvarint.Read(r)
+			if err != nil {
+				return nil, err
+			}
+			if err := skipN(r, reasonLen); err != nil {
+				return nil, err
+			}
+		case 0x1d:
+			// CONNECTION_CLOSE (application)
+			if _, err := quicvarint.Read(r); err != nil { // Error Code
+				return nil, err
+			}
+			reasonLen, err := quicvarint.Read(r)
+			if err != nil {
+				return nil, err
+			}
+			if err := skipN(r, reasonLen); err != nil {
+				return nil, err
+			}
+		case 0x30, 0x31:
+			// DATAGRAM
+			var dataLen uint64
+			if typ&0x01 != 0 {
+				n, err := quicvarint.Read(r)
+				if err != nil {
+					return nil, err
+				}
+				dataLen = n
+			} else {
+				dataLen = uint64(r.Len())
+			}
+			if err := skipN(r, dataLen); err != nil {
+				return nil, err
+			}
+		default:
+			return nil, fmt.Errorf("unsupported frame type: %d", typ)
+		}
 	}
 	return frames, nil
 }
 
+func skipN(r *bytes.Reader, n uint64) error {
+	if n > uint64(r.Len()) {
+		return io.EOF
+	}
+	_, err := r.Seek(int64(n), io.SeekCurrent)
+	return err
+}
+
 // assembleCryptoFrames assembles multiple crypto frames into a single slice (if possible).
 // It returns an error if the frames cannot be assembled. This can happen if the frames are not contiguous.
-func assembleCryptoFrames(frames []cryptoFrame) []byte {
+func assembleCryptoFrames(frames []CryptoFrame) []byte {
 	if len(frames) == 0 {
 		return nil
 	}

analyzer/udp/internal/quic/payload_test.go

@@ -0,0 +1,39 @@
+package quic
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestExtractCryptoFrames_Interleaved(t *testing.T) {
+	// PADDING, CRYPTO(offset=0,"abc"), PING, STREAM(len=2), CRYPTO(offset=3,"de")
+	payload := []byte{
+		0x00,
+		0x06, 0x00, 0x03, 'a', 'b', 'c',
+		0x01,
+		0x0a, 0x01, 0x02, 'z', 'z',
+		0x06, 0x03, 0x02, 'd', 'e',
+	}
+
+	frames, err := extractCryptoFrames(bytes.NewReader(payload))
+	if err != nil {
+		t.Fatalf("extractCryptoFrames() error = %v", err)
+	}
+	if len(frames) != 2 {
+		t.Fatalf("extractCryptoFrames() len = %d, want 2", len(frames))
+	}
+	if frames[0].Offset != 0 || string(frames[0].Data) != "abc" {
+		t.Fatalf("frame0 = %+v, want offset=0 data=abc", frames[0])
+	}
+	if frames[1].Offset != 3 || string(frames[1].Data) != "de" {
+		t.Fatalf("frame1 = %+v, want offset=3 data=de", frames[1])
+	}
+}
+
+func TestExtractCryptoFrames_UnsupportedFrame(t *testing.T) {
+	// 0x20 is currently unsupported in this parser.
+	_, err := extractCryptoFrames(bytes.NewReader([]byte{0x20}))
+	if err == nil {
+		t.Fatal("extractCryptoFrames() error = nil, want non-nil")
+	}
+}

analyzer/udp/quic.go

@@ -1,6 +1,8 @@
 package udp
 
 import (
+	"sort"
+
 	"git.difuse.io/Difuse/Mellaris/analyzer"
 	"git.difuse.io/Difuse/Mellaris/analyzer/internal"
 	"git.difuse.io/Difuse/Mellaris/analyzer/udp/internal/quic"
@@ -8,7 +10,8 @@ import (
 )
 
 const (
-	quicInvalidCountThreshold = 4
+	quicInvalidCountThreshold = 16
+	quicMaxCryptoDataLen      = 256 * 1024
 )
 
 var (
@@ -27,12 +30,17 @@ func (a *QUICAnalyzer) Limit() int {
 }
 
 func (a *QUICAnalyzer) NewUDP(info analyzer.UDPInfo, logger analyzer.Logger) analyzer.UDPStream {
-	return &quicStream{logger: logger}
+	return &quicStream{
+		logger: logger,
+		frames: make(map[int64][]byte),
+	}
 }
 
 type quicStream struct {
 	logger       analyzer.Logger
 	invalidCount int
+	frames       map[int64][]byte
+	maxEnd       int64
 }
 
 func (s *quicStream) Feed(rev bool, data []byte) (u *analyzer.PropUpdate, done bool) {
@@ -42,15 +50,21 @@ func (s *quicStream) Feed(rev bool, data []byte) (u *analyzer.PropUpdate, done b
 	const minDataSize = 41
 
 	if rev {
-		// We don't support server direction for now
+		// Ignore server direction; we only parse client initial data.
+		return nil, false
+	}
+
+	frs, err := quic.ReadCryptoFrames(data)
+	if err != nil || len(frs) == 0 {
 		s.invalidCount++
 		return nil, s.invalidCount >= quicInvalidCountThreshold
 	}
-
+	for _, f := range frs {
-	pl, err := quic.ReadCryptoPayload(data)
+		s.mergeFrame(f.Offset, f.Data)
-	if err != nil || len(pl) < 4 { // FIXME: isn't length checked inside quic.ReadCryptoPayload? Also, what about error handling?
+	}
-		s.invalidCount++
+	pl := s.contiguousPayloadFromZero()
-		return nil, s.invalidCount >= quicInvalidCountThreshold
+	if len(pl) < 4 {
+		return nil, false
 	}
 
 	if pl[0] != internal.TypeClientHello {
@@ -63,8 +77,12 @@ func (s *quicStream) Feed(rev bool, data []byte) (u *analyzer.PropUpdate, done b
 		s.invalidCount++
 		return nil, s.invalidCount >= quicInvalidCountThreshold
 	}
+	if len(pl) < 4+chLen {
+		// Wait for more CRYPTO data from subsequent packets.
+		return nil, false
+	}
 
-	m := internal.ParseTLSClientHelloMsgData(&utils.ByteBuffer{Buf: pl[4:]})
+	m := internal.ParseTLSClientHelloMsgData(&utils.ByteBuffer{Buf: pl[4 : 4+chLen]})
 	if m == nil {
 		s.invalidCount++
 		return nil, s.invalidCount >= quicInvalidCountThreshold
@@ -79,3 +97,53 @@ func (s *quicStream) Feed(rev bool, data []byte) (u *analyzer.PropUpdate, done b
 func (s *quicStream) Close(limited bool) *analyzer.PropUpdate {
 	return nil
 }
+
+func (s *quicStream) mergeFrame(offset int64, data []byte) {
+	if len(data) == 0 || offset < 0 {
+		return
+	}
+	if s.frames == nil {
+		s.frames = make(map[int64][]byte)
+	}
+	if _, exists := s.frames[offset]; exists {
+		return
+	}
+	s.frames[offset] = append([]byte(nil), data...)
+	end := offset + int64(len(data))
+	if end > s.maxEnd {
+		s.maxEnd = end
+	}
+}
+
+func (s *quicStream) contiguousPayloadFromZero() []byte {
+	if len(s.frames) == 0 || s.maxEnd <= 0 || s.maxEnd > quicMaxCryptoDataLen {
+		return nil
+	}
+	keys := make([]int64, 0, len(s.frames))
+	for k := range s.frames {
+		keys = append(keys, k)
+	}
+	sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
+	if keys[0] != 0 {
+		return nil
+	}
+	out := make([]byte, 0, s.maxEnd)
+	next := int64(0)
+	for _, k := range keys {
+		if k > next {
+			break
+		}
+		frame := s.frames[k]
+		frameEnd := k + int64(len(frame))
+		if frameEnd <= next {
+			continue
+		}
+		start := next - k
+		out = append(out, frame[start:]...)
+		next = frameEnd
+		if next >= quicMaxCryptoDataLen {
+			break
+		}
+	}
+	return out
+}