analyzer: make http3/quic handling more reliable

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
 	}
-
-	pl, err := quic.ReadCryptoPayload(data)
-	if err != nil || len(pl) < 4 { // FIXME: isn't length checked inside quic.ReadCryptoPayload? Also, what about error handling?
-		s.invalidCount++
-		return nil, s.invalidCount >= quicInvalidCountThreshold
+	for _, f := range frs {
+		s.mergeFrame(f.Offset, f.Data)
+	}
+	pl := s.contiguousPayloadFromZero()
+	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
+}