From 3b850c6ecf8fdebe4f8f706f1556d76159e99137 Mon Sep 17 00:00:00 2001 From: JackDoan Date: Wed, 14 Oct 2026 12:56:12 -0500 Subject: [PATCH] handshake/fragment: Remember answered messages, so repair fragments trailing a completed msg1 do not eat the source quota --- handshake/fec/fec.go | 5 ++ handshake/fragment/fragment_test.go | 84 +++++++++++++++++++++ handshake/fragment/table.go | 113 +++++++++++++++++++++++++++- handshake_manager_pq.go | 16 +++- handshake_manager_pq_test.go | 9 ++- 5 files changed, 217 insertions(+), 10 deletions(-) diff --git a/handshake/fec/fec.go b/handshake/fec/fec.go index 4007fc3f..2742f87c 100644 --- a/handshake/fec/fec.go +++ b/handshake/fec/fec.go @@ -128,6 +128,11 @@ func (d *Decoder) Has(index int) bool { return index >= 0 && index < MaxSymbols && d.have[index] } +// Seen returns which symbol indexes are held. +func (d *Decoder) Seen() [MaxSymbols]bool { + return d.have +} + // Ready reports whether enough symbols have arrived to decode. func (d *Decoder) Ready() bool { return len(d.syms) >= d.layout.K diff --git a/handshake/fragment/fragment_test.go b/handshake/fragment/fragment_test.go index 53e5e71f..470e2fda 100644 --- a/handshake/fragment/fragment_test.go +++ b/handshake/fragment/fragment_test.go @@ -2,6 +2,7 @@ package fragment import ( "bytes" + "crypto/sha256" "math/rand/v2" "sync" "testing" @@ -513,3 +514,86 @@ func FuzzTable(f *testing.F) { } }) } + +// TestTableAnswered sends many large messages from one source the way a responder +// sees back to back handshakes: each attempt's data fragments complete the +// message and its repair fragments trail. Remembering answered messages keeps the +// stragglers from piling up against the source quota, and a retry with fresh +// symbols still reports the message again. +func TestTableAnswered(t *testing.T) { + r := rand.New(rand.NewPCG(3, 3)) + tb := NewTable[testKey, string](testConfig()) + now := time.Unix(1000, 0) + const msgLen = 8006 // a hybrid msg1: K=7 data symbols, 3 repair + + sendAll := func(key testKey, pkts [][]byte) []*Result { + var out []*Result + for _, p := range pkts { + _, h, sym := parse(t, p) + res, err := tb.Add(key, key.src, now, h, sym) + require.NoError(t, err) + if res != nil { + out = append(out, res) + } + checkInvariants(t, tb) + } + return out + } + + for i := range 3 * testConfig().MaxPerSource { + msg := randMsg(r, msgLen) + key := testKey{"a", uint32(i + 1)} + s, err := NewSender(msg, testSubtype, 0, 1, key.fragID) + require.NoError(t, err) + data, repair, err := s.Attempt(DefaultRepair(s.Layout().K)) + require.NoError(t, err) + require.Len(t, data, 7) + require.Len(t, repair, 3) + + res := sendAll(key, data) + require.Len(t, res, 1, "message %d was refused", i) + require.Equal(t, msg, res[0].Msg) + tb.Answered(key, res[0], now) + + assert.Empty(t, sendAll(key, repair), "stragglers are not a new message") + assert.Zero(t, tb.Len(), "and hold no entry") + now = now.Add(100 * time.Millisecond) + } + + // A retry offers fresh symbols, and once K of them arrive it is a repeat. + msg := randMsg(r, msgLen) + key := testKey{"b", 99} + s, err := NewSender(msg, testSubtype, 0, 1, key.fragID) + require.NoError(t, err) + data, repair, err := s.Attempt(DefaultRepair(s.Layout().K)) + require.NoError(t, err) + res := sendAll(key, data) + require.Len(t, res, 1) + tb.Answered(key, res[0], now) + require.Empty(t, sendAll(key, repair)) + + data, repair, err = s.Attempt(DefaultRepair(s.Layout().K)) + require.NoError(t, err) + res = sendAll(key, append(data, repair...)) + require.Len(t, res, 1) + assert.True(t, res[0].Repeat) + assert.Nil(t, res[0].Msg) + assert.Equal(t, sha256.Sum256(msg), res[0].Hash) + assert.Positive(t, res[0].Received) + + // Forgotten, the message is reassembled from scratch again. + tb.Forget(key) + data, _, err = s.Attempt(DefaultRepair(s.Layout().K)) + require.NoError(t, err) + res = sendAll(key, data) + require.Len(t, res, 1) + assert.False(t, res[0].Repeat) + assert.Equal(t, msg, res[0].Msg) + + // Answered messages expire with the TTL like entries do. + tb.Answered(key, res[0], now) + tb.Expire(now.Add(time.Hour)) + tb.mu.Lock() + assert.Empty(t, tb.answered) + tb.mu.Unlock() +} diff --git a/handshake/fragment/table.go b/handshake/fragment/table.go index 0681c76c..5af147eb 100644 --- a/handshake/fragment/table.go +++ b/handshake/fragment/table.go @@ -1,10 +1,14 @@ package fragment import ( + "crypto/sha256" "errors" "math/rand/v2" "sync" "time" + + "github.com/slackhq/nebula/handshake/fec" + "github.com/slackhq/nebula/header" ) var ( @@ -30,6 +34,12 @@ type TableConfig struct { // K identifies one message, typically the source plus FragID. S identifies the // source for quotas, typically the address masked to /32 or /64. // +// A message the caller answered is remembered, see Answered, so the fragments +// still arriving for it do not start it over. Without that, the repair +// fragments trailing a message whose data fragments already completed it would +// open a new entry that, for K above the repair count, never completes: it holds +// symbols and a source's quota until it expires. +// // A Table is safe for concurrent use. type Table[K, S comparable] struct { mu sync.Mutex @@ -38,6 +48,20 @@ type Table[K, S comparable] struct { order []K perSource map[S]int bytes int + answered map[K]*answeredEntry +} + +// answeredEntry is what is kept of an answered message: no symbols, only which +// indexes have been seen, so a retry's fresh symbols can be told apart from +// stragglers of an attempt already decoded. +type answeredEntry struct { + hash [sha256.Size]byte + msgLen uint16 + k int + seen [fec.MaxSymbols]bool + fresh int + received int + lastSeen time.Time } type tableEntry[K, S comparable] struct { @@ -52,13 +76,24 @@ func NewTable[K, S comparable](cfg TableConfig) *Table[K, S] { cfg: cfg, entries: make(map[K]*tableEntry[K, S]), perSource: make(map[S]int), + answered: make(map[K]*answeredEntry), } } // Result describes a completed message. type Result struct { - Msg []byte // the decoded message - Received int // packet bytes received for it, for the amplification budget + Msg []byte // the decoded message, nil for a repeat + Hash [sha256.Size]byte // sha256 of the message + Received int // packet bytes received for it, for the amplification budget + + // Repeat means a message marked Answered was sent again: as many fresh + // symbols as it has data symbols arrived, which a retry offers and stragglers + // of an attempt already decoded do not. + Repeat bool + + msgLen uint16 + k int + seen [fec.MaxSymbols]bool } // Add records a fragment parsed by Parse. When the message is complete its entry @@ -70,6 +105,15 @@ func (t *Table[K, S]) Add(key K, source S, now time.Time, h Header, sym []byte) } t.mu.Lock() + if a, ok := t.answered[key]; ok { + if now.Sub(a.lastSeen) <= t.cfg.TTL { + res := a.add(now, h, sym) + t.mu.Unlock() + return res, nil + } + delete(t.answered, key) + } + e, ok := t.entries[key] if ok && now.Sub(e.lastSeen) > t.cfg.TTL { t.remove(key, e) @@ -111,7 +155,65 @@ func (t *Table[K, S]) Add(key K, source S, now time.Time, h Header, sym []byte) if err != nil { return nil, err } - return &Result{Msg: msg, Received: e.a.Received()}, nil + return &Result{ + Msg: msg, + Hash: sha256.Sum256(msg), + Received: e.a.Received(), + msgLen: e.a.msgLen, + k: e.a.dec.Layout().K, + seen: e.a.dec.Seen(), + }, nil +} + +func (a *answeredEntry) add(now time.Time, h Header, sym []byte) *Result { + if h.MsgLen != a.msgLen { + return nil + } + a.lastSeen = now + a.received += header.Len + HeaderLen + len(sym) + if a.seen[h.Index] { + return nil + } + a.seen[h.Index] = true + a.fresh++ + if a.fresh < a.k { + return nil + } + r := &Result{Hash: a.hash, Received: a.received, Repeat: true} + a.fresh = 0 + a.received = 0 + return r +} + +// Answered records that the caller answered res, the result of Add for key, so +// later fragments of it are counted rather than reassembled. A message is kept +// for the TTL after its last fragment, outside the entry, byte and source caps; +// at most MaxEntries are kept. +func (t *Table[K, S]) Answered(key K, res *Result, now time.Time) { + t.mu.Lock() + defer t.mu.Unlock() + // A straggler may have opened an entry while the caller worked. + if e, ok := t.entries[key]; ok { + t.remove(key, e) + } + if res.Repeat { + return + } + if _, ok := t.answered[key]; !ok && len(t.answered) >= t.cfg.MaxEntries { + for k := range t.answered { + delete(t.answered, k) + break + } + } + t.answered[key] = &answeredEntry{hash: res.Hash, msgLen: res.msgLen, k: res.k, seen: res.seen, lastSeen: now} +} + +// Forget drops what Answered recorded for key, so its fragments are reassembled +// again. +func (t *Table[K, S]) Forget(key K) { + t.mu.Lock() + defer t.mu.Unlock() + delete(t.answered, key) } // Expire drops every entry idle for longer than the TTL. @@ -128,6 +230,11 @@ func (t *Table[K, S]) Expire(now time.Time) { } i++ } + for key, a := range t.answered { + if now.Sub(a.lastSeen) > t.cfg.TTL { + delete(t.answered, key) + } + } } // Len returns the number of messages in progress. diff --git a/handshake_manager_pq.go b/handshake_manager_pq.go index 9c560994..ede6b93c 100644 --- a/handshake_manager_pq.go +++ b/handshake_manager_pq.go @@ -362,14 +362,23 @@ func (hm *HandshakeManager) handleFragment(via ViaSender, packet []byte, h *head if res == nil { return } - if existing := hm.pqAnswered.lookup(res.Msg, time.Now()); existing != nil { + if existing := hm.pqAnswered.lookup(res.Hash, time.Now()); existing != nil { // Already answered: resend, rate limited, with fresh symbols. if msg := existing.HandshakePacket[handshakePacketStage2]; msg != nil { hm.sendHandshakeResponse(via, msg, existing, true, res.Received) } + hm.pqTable.Answered(key, res, now) + return + } + if res.Repeat { + // Answered once, but long enough ago to be forgotten: start over. + hm.pqTable.Forget(key) return } hm.beginHandshake(via, reassembled(h, res.Msg), h, res.Received) + if hm.pqAnswered.lookup(res.Hash, now) != nil { + hm.pqTable.Answered(key, res, now) + } case 2: if hh := hm.queryIndex(h.RemoteIndex); hh != nil { @@ -553,10 +562,11 @@ func (a *pqAnswered) add(msg1 []byte, hostinfo *HostInfo, now time.Time) { a.entries[sha256.Sum256(msg1)] = pqAnsweredEntry{hostinfo: hostinfo, at: now} } -func (a *pqAnswered) lookup(msg1 []byte, now time.Time) *HostInfo { +// lookup finds the hostinfo that answered the msg1 with this sha256. +func (a *pqAnswered) lookup(hash [sha256.Size]byte, now time.Time) *HostInfo { a.Lock() defer a.Unlock() - e, ok := a.entries[sha256.Sum256(msg1)] + e, ok := a.entries[hash] if !ok || now.Sub(e.at) > a.ttl { return nil } diff --git a/handshake_manager_pq_test.go b/handshake_manager_pq_test.go index 9dcda5a2..2530b31b 100644 --- a/handshake_manager_pq_test.go +++ b/handshake_manager_pq_test.go @@ -1,6 +1,7 @@ package nebula import ( + "crypto/sha256" "math/rand/v2" "net/netip" "sync" @@ -126,11 +127,11 @@ func TestPQAnswered(t *testing.T) { a := newPQAnswered(5 * time.Second) hi := &HostInfo{} - assert.Nil(t, a.lookup([]byte("msg1"), now)) + assert.Nil(t, a.lookup(sha256.Sum256([]byte("msg1")), now)) a.add([]byte("msg1"), hi, now) - assert.Same(t, hi, a.lookup([]byte("msg1"), now.Add(time.Second))) - assert.Nil(t, a.lookup([]byte("other"), now)) - assert.Nil(t, a.lookup([]byte("msg1"), now.Add(6*time.Second)), "stale entries are not answered") + assert.Same(t, hi, a.lookup(sha256.Sum256([]byte("msg1")), now.Add(time.Second))) + assert.Nil(t, a.lookup(sha256.Sum256([]byte("other")), now)) + assert.Nil(t, a.lookup(sha256.Sum256([]byte("msg1")), now.Add(6*time.Second)), "stale entries are not answered") a.expire(now.Add(6 * time.Second)) assert.Empty(t, a.entries)