handshake/fragment: Remember answered messages, so repair fragments trailing a completed msg1 do not eat the source quota

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

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()
+}

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.

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
 	}

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)