chore: extract bitmap handling from QUIC crypto reassembly

component/sniffer/bitmap.go

@@ -0,0 +1,84 @@
+package sniffer
+
+import "math/bits"
+
+// bitmap is a growable bitmap. Its zero value is ready for use.
+type bitmap struct {
+	words []uint64
+}
+
+// setRange sets every bit in the half-open interval [start, end). It panics if
+// start is negative or end is less than start.
+func (b *bitmap) setRange(start, end int) {
+	if start < 0 || end < start {
+		panic("invalid bitmap range")
+	}
+	if start == end {
+		return
+	}
+
+	wordCount := (end-1)/64 + 1
+	if wordCount > len(b.words) {
+		newWordCount := uint(1) << bits.Len(uint(wordCount-1))
+		if newWordCount > ^uint(0)>>1 {
+			newWordCount = uint(wordCount)
+		}
+		words := make([]uint64, int(newWordCount))
+		copy(words, b.words)
+		b.words = words
+	}
+
+	firstWord := start / 64
+	lastWord := (end - 1) / 64
+	firstMask := ^uint64(0) << (start % 64)
+	lastMask := ^uint64(0)
+	if endBit := end % 64; endBit != 0 {
+		lastMask = (uint64(1) << endBit) - 1
+	}
+	if firstWord == lastWord {
+		b.words[firstWord] |= firstMask & lastMask
+		return
+	}
+
+	b.words[firstWord] |= firstMask
+	for i := firstWord + 1; i < lastWord; i++ {
+		b.words[i] = ^uint64(0)
+	}
+	b.words[lastWord] |= lastMask
+}
+
+// firstUnset returns the first unset bit in [start, end), or end if all bits
+// in the interval are set. It panics if start is negative or end is less than
+// start.
+func (b bitmap) firstUnset(start, end int) int {
+	if start < 0 || end < start {
+		panic("invalid bitmap range")
+	}
+	if start == end {
+		return end
+	}
+
+	wordIndex := start / 64
+	if wordIndex >= len(b.words) {
+		return start
+	}
+	word := b.words[wordIndex] | ((uint64(1) << (start % 64)) - 1)
+	for {
+		if word != ^uint64(0) {
+			unset := wordIndex*64 + bits.TrailingZeros64(^word)
+			if unset < end {
+				return unset
+			}
+			return end
+		}
+
+		wordIndex++
+		if wordIndex*64 >= end {
+			return end
+		}
+		if wordIndex >= len(b.words) {
+			return wordIndex * 64
+		}
+		word = b.words[wordIndex]
+	}
+}

component/sniffer/bitmap_test.go

@@ -0,0 +1,53 @@
+package sniffer
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBitmap(t *testing.T) {
+	t.Run("zero value", func(t *testing.T) {
+		var coverage bitmap
+		assert.Equal(t, 0, coverage.firstUnset(0, 128))
+		assert.Equal(t, 0, coverage.firstUnset(0, 0))
+	})
+
+	t.Run("single word", func(t *testing.T) {
+		var coverage bitmap
+		coverage.setRange(1, 5)
+		assert.Equal(t, 0, coverage.firstUnset(0, 5))
+		assert.Equal(t, 5, coverage.firstUnset(1, 5))
+		assert.Equal(t, 5, coverage.firstUnset(1, 6))
+	})
+
+	t.Run("cross word", func(t *testing.T) {
+		var coverage bitmap
+		coverage.setRange(1, 130)
+		assert.Equal(t, 130, coverage.firstUnset(1, 130))
+		assert.Equal(t, 130, coverage.firstUnset(1, 131))
+	})
+
+	t.Run("one bit gap", func(t *testing.T) {
+		var coverage bitmap
+		coverage.setRange(0, 64)
+		coverage.setRange(65, 130)
+		assert.Equal(t, 64, coverage.firstUnset(0, 130))
+	})
+
+	t.Run("overlap survives growth", func(t *testing.T) {
+		var coverage bitmap
+		coverage.setRange(0, 10)
+		coverage.setRange(64, 70)
+		coverage.setRange(8, 65)
+		assert.Equal(t, 70, coverage.firstUnset(0, 71))
+	})
+
+	t.Run("invalid range", func(t *testing.T) {
+		var coverage bitmap
+		assert.Panics(t, func() { coverage.setRange(-1, 0) })
+		assert.Panics(t, func() { coverage.setRange(1, 0) })
+		assert.Panics(t, func() { coverage.firstUnset(-1, 0) })
+		assert.Panics(t, func() { coverage.firstUnset(1, 0) })
+	})
+}

component/sniffer/quic_sniffer.go

@@ -8,7 +8,6 @@ import (
 	"encoding/binary"
 	"errors"
 	"io"
-	"math/bits"
 	"sync"
 	"time"
 
@@ -141,7 +140,7 @@ type quicPacketSender struct {
 	structure           *quicStructure
 	lock                sync.RWMutex
 	buffer              []byte
-	receivedCryptoData  []uint64
+	receivedCryptoData  bitmap
 	contiguousCryptoEnd uint64
 	result              string
 
@@ -208,7 +207,7 @@ func (q *quicPacketSender) close() {
 			_ = pool.Put(q.buffer)
 			q.buffer = nil
 		}
-		q.receivedCryptoData = nil
+		q.receivedCryptoData = bitmap{}
 		q.contiguousCryptoEnd = 0
 	}
 }
@@ -511,47 +510,12 @@ func (q *quicPacketSender) addCryptoData(offset uint64, data []byte) error {
 	}
 	copy(q.buffer[offset:end], data)
 
-	wordCount := (cap(q.buffer) + 63) / 64
-	if wordCount > len(q.receivedCryptoData) {
-		receivedCryptoData := make([]uint64, wordCount)
-		copy(receivedCryptoData, q.receivedCryptoData)
-		q.receivedCryptoData = receivedCryptoData
-	}
-	firstWord := int(offset / 64)
-	lastWord := int((end - 1) / 64)
-	firstMask := ^uint64(0) << (offset % 64)
-	lastMask := ^uint64(0)
-	if endBit := end % 64; endBit != 0 {
-		lastMask = (uint64(1) << endBit) - 1
-	}
-	if firstWord == lastWord {
-		q.receivedCryptoData[firstWord] |= firstMask & lastMask
-	} else {
-		q.receivedCryptoData[firstWord] |= firstMask
-		for i := firstWord + 1; i < lastWord; i++ {
-			q.receivedCryptoData[i] = ^uint64(0)
-		}
-		q.receivedCryptoData[lastWord] |= lastMask
-	}
+	q.receivedCryptoData.setRange(int(offset), int(end))
 
-	// The contiguous prefix only moves forward, so each retained byte is
-	// checked at most once regardless of fragment order or retransmission.
+	// The contiguous prefix only moves forward, allowing the bitmap to skip
+	// complete words without rescanning the assembled prefix.
 	if offset <= q.contiguousCryptoEnd && end > q.contiguousCryptoEnd {
-		for q.contiguousCryptoEnd < uint64(len(q.buffer)) {
-			word := q.receivedCryptoData[q.contiguousCryptoEnd/64] >> (q.contiguousCryptoEnd % 64)
-			covered := uint64(bits.TrailingZeros64(^word))
-			remaining := uint64(len(q.buffer)) - q.contiguousCryptoEnd
-			if wordRemaining := 64 - q.contiguousCryptoEnd%64; remaining > wordRemaining {
-				remaining = wordRemaining
-			}
-			if covered > remaining {
-				covered = remaining
-			}
-			q.contiguousCryptoEnd += covered
-			if covered < remaining {
-				break
-			}
-		}
+		q.contiguousCryptoEnd = uint64(q.receivedCryptoData.firstUnset(int(q.contiguousCryptoEnd), len(q.buffer)))
 	}
 	return nil
 }