From efd63d87462e8a91be8bb2f6efaf02b68b8f1dab Mon Sep 17 00:00:00 2001 From: wwqgtxx Date: Tue, 28 Jul 2026 18:26:08 +0800 Subject: [PATCH] chore: extract bitmap handling from QUIC crypto reassembly --- component/sniffer/bitmap.go | 84 +++++++++++++++++++++++++++++++ component/sniffer/bitmap_test.go | 53 +++++++++++++++++++ component/sniffer/quic_sniffer.go | 48 +++--------------- 3 files changed, 143 insertions(+), 42 deletions(-) create mode 100644 component/sniffer/bitmap.go create mode 100644 component/sniffer/bitmap_test.go diff --git a/component/sniffer/bitmap.go b/component/sniffer/bitmap.go new file mode 100644 index 00000000..707617fe --- /dev/null +++ b/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] + } +} diff --git a/component/sniffer/bitmap_test.go b/component/sniffer/bitmap_test.go new file mode 100644 index 00000000..d99ec0b2 --- /dev/null +++ b/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) }) + }) +} diff --git a/component/sniffer/quic_sniffer.go b/component/sniffer/quic_sniffer.go index 3cf01e6c..a9f5ac26 100644 --- a/component/sniffer/quic_sniffer.go +++ b/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 }