handshake/fragment: Carry oversized handshake messages as FEC coded fragments

handshake/fragment/assembly.go

@@ -0,0 +1,67 @@
+package fragment
+
+import (
+	"errors"
+
+	"github.com/slackhq/nebula/handshake/fec"
+	"github.com/slackhq/nebula/header"
+)
+
+// ErrMismatch is returned when a fragment disagrees with the message it claims to
+// belong to. The caller should discard the whole assembly.
+var ErrMismatch = errors.New("fragment does not match its message")
+
+// Assembly collects the fragments of one incoming message. Its memory is bounded
+// by the message length: at most K symbols are kept however many arrive.
+//
+// An Assembly is not safe for concurrent use.
+type Assembly struct {
+	fragID   uint32
+	msgLen   uint16
+	dec      *fec.Decoder
+	received int
+}
+
+// NewAssembly starts an assembly for the message described by h. The fragment
+// itself still has to be passed to Add.
+func NewAssembly(h Header) (*Assembly, error) {
+	dec, err := fec.NewDecoder(int(h.MsgLen), MaxSymbolSize)
+	if err != nil {
+		return nil, err
+	}
+	return &Assembly{fragID: h.FragID, msgLen: h.MsgLen, dec: dec}, nil
+}
+
+// Add records a fragment parsed by Parse. It returns true once enough symbols
+// are held to call Decode.
+func (a *Assembly) Add(h Header, sym []byte) (bool, error) {
+	if h.FragID != a.fragID || h.MsgLen != a.msgLen {
+		return a.dec.Ready(), ErrMismatch
+	}
+	a.received += header.Len + HeaderLen + len(sym)
+	return a.dec.Add(int(h.Index), sym)
+}
+
+// wouldStore returns how many bytes Add would retain for this fragment.
+func (a *Assembly) wouldStore(h Header, sym []byte) int {
+	if h.FragID != a.fragID || h.MsgLen != a.msgLen || a.dec.Ready() || a.dec.Has(int(h.Index)) {
+		return 0
+	}
+	return len(sym)
+}
+
+// Decode reconstructs the message. It requires Add to have reported ready.
+func (a *Assembly) Decode() ([]byte, error) {
+	return a.dec.Decode()
+}
+
+// Stored returns the number of symbol bytes held.
+func (a *Assembly) Stored() int {
+	return a.dec.Stored() * a.dec.Layout().SymbolSize
+}
+
+// Received returns the number of packet bytes received for this message,
+// duplicates included. It is the basis for the amplification budget.
+func (a *Assembly) Received() int {
+	return a.received
+}

handshake/fragment/fragment_test.go

@@ -0,0 +1,466 @@
+package fragment
+
+import (
+	"bytes"
+	"math/rand/v2"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/slackhq/nebula/header"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const testSubtype = header.MessageSubType(9)
+
+type testKey struct {
+	src    string
+	fragID uint32
+}
+
+func randMsg(r *rand.Rand, n int) []byte {
+	b := make([]byte, n)
+	for i := range b {
+		b[i] = byte(r.Uint32())
+	}
+	return b
+}
+
+func testConfig() TableConfig {
+	return TableConfig{
+		MaxBytes:     1 << 20,
+		MaxEntries:   64,
+		MaxPerSource: 4,
+		MaxMsgLen:    16384,
+		TTL:          2 * time.Second,
+	}
+}
+
+// checkInvariants verifies the table's bookkeeping against its contents.
+func checkInvariants[K, S comparable](t *testing.T, tb *Table[K, S]) {
+	t.Helper()
+	tb.mu.Lock()
+	defer tb.mu.Unlock()
+
+	require.Len(t, tb.order, len(tb.entries))
+	require.LessOrEqual(t, len(tb.entries), tb.cfg.MaxEntries)
+	bytes := 0
+	perSource := map[S]int{}
+	for i, key := range tb.order {
+		e, ok := tb.entries[key]
+		require.True(t, ok)
+		require.Equal(t, i, e.pos)
+		bytes += e.a.Stored()
+		perSource[e.source]++
+	}
+	require.Equal(t, bytes, tb.bytes)
+	require.LessOrEqual(t, tb.bytes, tb.cfg.MaxBytes)
+	require.Equal(t, perSource, tb.perSource)
+	for _, n := range perSource {
+		require.LessOrEqual(t, n, tb.cfg.MaxPerSource)
+	}
+}
+
+func parse(t *testing.T, p []byte) (header.H, Header, []byte) {
+	t.Helper()
+	var nh header.H
+	require.NoError(t, nh.Parse(p))
+	h, sym, err := Parse(p[header.Len:])
+	require.NoError(t, err)
+	return nh, h, sym
+}
+
+func TestParse(t *testing.T) {
+	h := Header{FragID: 0xdeadbeef, MsgLen: 3000, Index: 7}
+	b := h.Encode(nil)
+	require.Len(t, b, HeaderLen)
+	b = append(b, make([]byte, 1000)...)
+
+	got, sym, err := Parse(b)
+	require.NoError(t, err)
+	assert.Equal(t, h, got)
+	assert.Len(t, sym, 1000)
+
+	_, _, err = Parse(b[:HeaderLen-1])
+	require.ErrorIs(t, err, ErrShortFragment)
+	_, _, err = Parse(b[:len(b)-1])
+	require.ErrorIs(t, err, ErrSymbolSize)
+	_, _, err = Parse(append(b, 0))
+	require.ErrorIs(t, err, ErrSymbolSize)
+
+	bad := Header{FragID: 0, MsgLen: 3000}.Encode(nil)
+	_, _, err = Parse(append(bad, make([]byte, 1000)...))
+	require.ErrorIs(t, err, ErrFragID)
+
+	bad = Header{FragID: 1, MsgLen: 0}.Encode(nil)
+	_, _, err = Parse(bad)
+	require.ErrorIs(t, err, ErrMsgLen)
+
+	bad = Header{FragID: 1, MsgLen: 3000, Flags: 1}.Encode(nil)
+	_, _, err = Parse(append(bad, make([]byte, 1000)...))
+	require.ErrorIs(t, err, ErrFlags)
+}
+
+func TestDefaultRepair(t *testing.T) {
+	for k, want := range map[int]int{1: 2, 2: 2, 6: 2, 7: 3, 9: 3, 10: 4, 14: 5} {
+		assert.Equal(t, want, DefaultRepair(k), "k=%d", k)
+	}
+}
+
+func TestSender(t *testing.T) {
+	_, err := NewSender(nil, testSubtype, 0, 1, 1)
+	require.ErrorIs(t, err, ErrMsgLen)
+	_, err = NewSender(make([]byte, 65536), testSubtype, 0, 1, 1)
+	require.ErrorIs(t, err, ErrMsgLen)
+	_, err = NewSender(make([]byte, 10), testSubtype, 0, 1, 0)
+	require.ErrorIs(t, err, ErrFragID)
+
+	msg := randMsg(rand.New(rand.NewPCG(1, 1)), 9800)
+	s, err := NewSender(msg, testSubtype, 77, 2, 1234)
+	require.NoError(t, err)
+	require.Equal(t, 9, s.Layout().K)
+
+	indexes := func(ps [][]byte) []int {
+		var out []int
+		for _, p := range ps {
+			require.LessOrEqual(t, len(p), MaxPacketLen)
+			nh, h, _ := parse(t, p)
+			assert.Equal(t, header.Handshake, nh.Type)
+			assert.Equal(t, testSubtype, nh.Subtype)
+			assert.Equal(t, uint32(77), nh.RemoteIndex)
+			assert.Equal(t, uint64(2), nh.MessageCounter)
+			assert.Equal(t, uint32(1234), h.FragID)
+			assert.Equal(t, uint16(9800), h.MsgLen)
+			out = append(out, int(h.Index))
+		}
+		return out
+	}
+
+	now, later, err := s.Attempt(3)
+	require.NoError(t, err)
+	assert.Equal(t, []int{0, 1, 2, 3, 4, 5, 6, 7, 8}, indexes(now))
+	assert.Equal(t, []int{9, 10, 11}, indexes(later))
+
+	// A retry offers only indexes the receiver has not seen.
+	now, later, err = s.Attempt(3)
+	require.NoError(t, err)
+	assert.Equal(t, []int{12, 13, 14, 15, 16, 17, 18, 19, 20}, indexes(now))
+	assert.Equal(t, []int{21, 22, 23}, indexes(later))
+
+	// A sender rebuilt from the same message resumes where it was told to.
+	s2, err := NewSender(msg, testSubtype, 77, 2, 1234)
+	require.NoError(t, err)
+	s2.SetNext(s.Next())
+	now2, _, err := s2.Attempt(0)
+	require.NoError(t, err)
+	assert.Equal(t, []int{24, 25, 26, 27, 28, 29, 30, 31, 32}, indexes(now2))
+	s2.SetNext(-1)
+	assert.Equal(t, 255, s2.Next())
+
+	// Indexes wrap once the space is used up.
+	for range 19 {
+		_, _, err = s.Attempt(3)
+		require.NoError(t, err)
+	}
+	now, _, err = s.Attempt(3)
+	require.NoError(t, err)
+	assert.Equal(t, []int{252, 253, 254, 255, 0, 1, 2, 3, 4}, indexes(now))
+}
+
+// TestLossy runs whole messages through heavy loss, reordering and duplication,
+// merging across attempts, and checks the table ends empty.
+func TestLossy(t *testing.T) {
+	r := rand.New(rand.NewPCG(2, 2))
+	tb := NewTable[testKey, string](testConfig())
+	now := time.Unix(1000, 0)
+
+	for i := range 300 {
+		msg := randMsg(r, 1+r.IntN(16384))
+		fragID := uint32(i + 1)
+		s, err := NewSender(msg, testSubtype, 0, 1, fragID)
+		require.NoError(t, err)
+
+		var got *Result
+		for attempt := 0; got == nil; attempt++ {
+			require.Less(t, attempt, 20, "message never completed")
+			a, b, err := s.Attempt(DefaultRepair(s.Layout().K))
+			require.NoError(t, err)
+			var pkts [][]byte
+			for _, p := range append(a, b...) {
+				if r.Float64() < 0.3 {
+					continue
+				}
+				pkts = append(pkts, p)
+				if r.Float64() < 0.2 {
+					pkts = append(pkts, p)
+				}
+			}
+			r.Shuffle(len(pkts), func(i, j int) { pkts[i], pkts[j] = pkts[j], pkts[i] })
+
+			for _, p := range pkts {
+				_, h, sym := parse(t, p)
+				res, err := tb.Add(testKey{"a", h.FragID}, "a", now, h, sym)
+				require.NoError(t, err)
+				if res != nil {
+					// The table keeps no memory of completed messages, so a late
+					// duplicate can complete one again; the handshake layer
+					// recognises the repeat.
+					require.True(t, bytes.Equal(msg, res.Msg))
+					got = res
+				}
+				checkInvariants(t, tb)
+			}
+			now = now.Add(100 * time.Millisecond)
+		}
+		require.Positive(t, got.Received)
+		// Late fragments of a completed message start a fresh entry; drop it.
+		tb.Expire(now.Add(time.Hour))
+	}
+	assert.Zero(t, tb.Len())
+	assert.Zero(t, tb.Bytes())
+}
+
+// TestConcurrent spreads each message's fragments across goroutines, as
+// multiple UDP reader routines would.
+func TestConcurrent(t *testing.T) {
+	tb := NewTable[testKey, string](testConfig())
+	now := time.Unix(1000, 0)
+	r := rand.New(rand.NewPCG(3, 3))
+
+	const msgs = 50
+	want := make(map[uint32][]byte, msgs)
+	var pkts [][]byte
+	for id := uint32(1); id <= msgs; id++ {
+		msg := randMsg(r, 1+r.IntN(16384))
+		want[id] = msg
+		s, err := NewSender(msg, testSubtype, 0, 1, id)
+		require.NoError(t, err)
+		a, b, err := s.Attempt(DefaultRepair(s.Layout().K))
+		require.NoError(t, err)
+		pkts = append(pkts, a...)
+		pkts = append(pkts, b...)
+	}
+	r.Shuffle(len(pkts), func(i, j int) { pkts[i], pkts[j] = pkts[j], pkts[i] })
+
+	var mu sync.Mutex
+	done := map[uint32]bool{}
+	var wg sync.WaitGroup
+	for w := range 8 {
+		wg.Go(func() {
+			for i := w; i < len(pkts); i += 8 {
+				h, sym, err := Parse(pkts[i][16:])
+				if err != nil {
+					t.Error(err)
+					return
+				}
+				src := string(rune('a' + h.FragID))
+				res, err := tb.Add(testKey{src, h.FragID}, src, now, h, sym)
+				if err != nil {
+					t.Error(err)
+					return
+				}
+				if res != nil {
+					if !bytes.Equal(want[h.FragID], res.Msg) {
+						t.Errorf("message %d decoded wrong", h.FragID)
+					}
+					mu.Lock()
+					done[h.FragID] = true
+					mu.Unlock()
+				}
+			}
+		})
+	}
+	wg.Wait()
+	assert.Len(t, done, msgs)
+	checkInvariants(t, tb)
+}
+
+func firstFragment(t *testing.T, msgLen int, fragID uint32) (Header, []byte) {
+	t.Helper()
+	s, err := NewSender(make([]byte, msgLen), testSubtype, 0, 1, fragID)
+	require.NoError(t, err)
+	now, _, err := s.Attempt(0)
+	require.NoError(t, err)
+	_, h, sym := parse(t, now[0])
+	return h, sym
+}
+
+func TestTableLimits(t *testing.T) {
+	now := time.Unix(1000, 0)
+
+	t.Run("max message length", func(t *testing.T) {
+		tb := NewTable[testKey, string](testConfig())
+		h, sym := firstFragment(t, 16385, 1)
+		_, err := tb.Add(testKey{"a", 1}, "a", now, h, sym)
+		require.ErrorIs(t, err, ErrTooLarge)
+		assert.Zero(t, tb.Len())
+	})
+
+	t.Run("per source quota", func(t *testing.T) {
+		cfg := testConfig()
+		cfg.MaxPerSource = 2
+		tb := NewTable[testKey, string](cfg)
+		for id := uint32(1); id <= 3; id++ {
+			h, sym := firstFragment(t, 5000, id)
+			_, err := tb.Add(testKey{"a", id}, "a", now, h, sym)
+			if id <= 2 {
+				require.NoError(t, err)
+			} else {
+				require.ErrorIs(t, err, ErrSourceQuota)
+			}
+		}
+		h, sym := firstFragment(t, 5000, 4)
+		_, err := tb.Add(testKey{"b", 4}, "b", now, h, sym)
+		require.NoError(t, err)
+		assert.Equal(t, 3, tb.Len())
+		checkInvariants(t, tb)
+	})
+
+	t.Run("max entries evicts", func(t *testing.T) {
+		cfg := testConfig()
+		cfg.MaxEntries = 4
+		tb := NewTable[testKey, string](cfg)
+		for id := uint32(1); id <= 10; id++ {
+			h, sym := firstFragment(t, 5000, id)
+			src := string(rune('a' + id))
+			_, err := tb.Add(testKey{src, id}, src, now, h, sym)
+			require.NoError(t, err)
+			checkInvariants(t, tb)
+		}
+		assert.Equal(t, 4, tb.Len())
+	})
+
+	t.Run("max bytes evicts others", func(t *testing.T) {
+		cfg := testConfig()
+		cfg.MaxBytes = 3 * 1000
+		tb := NewTable[testKey, string](cfg)
+		for id := uint32(1); id <= 10; id++ {
+			h, sym := firstFragment(t, 5000, id) // 5 symbols of 1000
+			src := string(rune('a' + id))
+			_, err := tb.Add(testKey{src, id}, src, now, h, sym)
+			require.NoError(t, err)
+			checkInvariants(t, tb)
+		}
+		assert.Equal(t, 3000, tb.Bytes())
+	})
+
+	t.Run("no room removes new entry", func(t *testing.T) {
+		cfg := testConfig()
+		cfg.MaxBytes = 500
+		tb := NewTable[testKey, string](cfg)
+		h, sym := firstFragment(t, 5000, 1)
+		_, err := tb.Add(testKey{"a", 1}, "a", now, h, sym)
+		require.ErrorIs(t, err, ErrNoRoom)
+		assert.Zero(t, tb.Len())
+		checkInvariants(t, tb)
+	})
+
+	t.Run("duplicates are not charged", func(t *testing.T) {
+		tb := NewTable[testKey, string](testConfig())
+		h, sym := firstFragment(t, 5000, 1)
+		for range 5 {
+			_, err := tb.Add(testKey{"a", 1}, "a", now, h, sym)
+			require.NoError(t, err)
+		}
+		assert.Equal(t, 1000, tb.Bytes())
+	})
+
+	t.Run("mismatch drops the entry", func(t *testing.T) {
+		tb := NewTable[testKey, string](testConfig())
+		h, sym := firstFragment(t, 5000, 1)
+		_, err := tb.Add(testKey{"a", 1}, "a", now, h, sym)
+		require.NoError(t, err)
+
+		h2, sym2 := firstFragment(t, 6000, 1)
+		_, err = tb.Add(testKey{"a", 1}, "a", now, h2, sym2)
+		require.ErrorIs(t, err, ErrMismatch)
+		assert.Zero(t, tb.Len())
+		assert.Zero(t, tb.Bytes())
+		checkInvariants(t, tb)
+	})
+
+	t.Run("ttl", func(t *testing.T) {
+		tb := NewTable[testKey, string](testConfig())
+		s, err := NewSender(make([]byte, 5000), testSubtype, 0, 1, 1)
+		require.NoError(t, err)
+		pkts, _, err := s.Attempt(0)
+		require.NoError(t, err)
+
+		_, h, sym := parse(t, pkts[0])
+		_, err = tb.Add(testKey{"a", 1}, "a", now, h, sym)
+		require.NoError(t, err)
+
+		// Activity within the TTL keeps the entry and slides its deadline.
+		_, h, sym = parse(t, pkts[1])
+		_, err = tb.Add(testKey{"a", 1}, "a", now.Add(1500*time.Millisecond), h, sym)
+		require.NoError(t, err)
+		assert.Equal(t, 2000, tb.Bytes())
+
+		tb.Expire(now.Add(3 * time.Second))
+		assert.Equal(t, 1, tb.Len())
+
+		// A stale entry is restarted rather than merged.
+		_, h, sym = parse(t, pkts[2])
+		_, err = tb.Add(testKey{"a", 1}, "a", now.Add(10*time.Second), h, sym)
+		require.NoError(t, err)
+		assert.Equal(t, 1000, tb.Bytes())
+
+		tb.Expire(now.Add(time.Minute))
+		assert.Zero(t, tb.Len())
+		checkInvariants(t, tb)
+	})
+}
+
+func FuzzParse(f *testing.F) {
+	f.Add(append(Header{FragID: 1, MsgLen: 3000}.Encode(nil), make([]byte, 1000)...))
+	f.Add([]byte{0, 0, 0, 1, 0, 1, 0, 0, 7})
+	f.Fuzz(func(t *testing.T, b []byte) {
+		h, sym, err := Parse(b)
+		if err != nil {
+			return
+		}
+		require.Equal(t, b, append(h.Encode(nil), sym...))
+		l, err := h.Layout()
+		require.NoError(t, err)
+		require.LessOrEqual(t, l.SymbolSize, MaxSymbolSize)
+	})
+}
+
+// FuzzTable drives a small table with arbitrary fragments and checks its
+// bookkeeping after every step.
+func FuzzTable(f *testing.F) {
+	f.Add([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})
+	f.Fuzz(func(t *testing.T, ops []byte) {
+		tb := NewTable[testKey, string](TableConfig{
+			MaxBytes:     4000,
+			MaxEntries:   3,
+			MaxPerSource: 2,
+			MaxMsgLen:    6000,
+			TTL:          time.Second,
+		})
+		now := time.Unix(0, 0)
+		for len(ops) >= 4 {
+			src := string(rune('a' + ops[0]%4))
+			fragID := uint32(ops[1]%5) + 1
+			msgLen := uint16(ops[2])*32 + 1
+			index := ops[3]
+			ops = ops[4:]
+
+			h := Header{FragID: fragID, MsgLen: msgLen, Index: index}
+			l, err := h.Layout()
+			require.NoError(t, err)
+			sym := make([]byte, l.SymbolSize)
+			if index%7 == 0 {
+				now = now.Add(600 * time.Millisecond)
+				tb.Expire(now)
+			}
+			res, err := tb.Add(testKey{src, fragID}, src, now, h, sym)
+			if res != nil {
+				require.NoError(t, err)
+				require.Len(t, res.Msg, int(msgLen))
+			}
+			checkInvariants(t, tb)
+		}
+	})
+}

handshake/fragment/sender.go

@@ -0,0 +1,107 @@
+package fragment
+
+import (
+	"errors"
+	"math"
+
+	"github.com/slackhq/nebula/handshake/fec"
+	"github.com/slackhq/nebula/header"
+)
+
+// DefaultRepair is the starting number of repair symbols per attempt for a
+// message of k data symbols. Small messages get proportionally more repair.
+func DefaultRepair(k int) int {
+	return max(2, (k+2)/3)
+}
+
+// Sender produces the fragment packets of one outgoing handshake message.
+//
+// The first attempt sends the K data symbols followed by repair symbols. Every
+// later attempt sends symbols with indexes the receiver has not been offered yet,
+// so whatever it already holds still counts toward K. Symbols are a deterministic
+// function of the message, which keeps retransmits idempotent. Once all
+// fec.MaxSymbols indexes are used the sequence wraps and repeats.
+//
+// A Sender is not safe for concurrent use.
+type Sender struct {
+	enc         *fec.Encoder
+	subtype     header.MessageSubType
+	remoteIndex uint32
+	counter     uint64
+	fragID      uint32
+	next        int
+}
+
+// NewSender prepares msg for sending. subtype, remoteIndex and counter fill the
+// Nebula header of every fragment; fragID is the sender's local index.
+func NewSender(msg []byte, subtype header.MessageSubType, remoteIndex uint32, counter uint64, fragID uint32) (*Sender, error) {
+	if len(msg) == 0 || len(msg) > math.MaxUint16 {
+		return nil, ErrMsgLen
+	}
+	if fragID == 0 {
+		return nil, ErrFragID
+	}
+	enc, err := fec.NewEncoder(msg, MaxSymbolSize)
+	if err != nil {
+		return nil, err
+	}
+	return &Sender{
+		enc:         enc,
+		subtype:     subtype,
+		remoteIndex: remoteIndex,
+		counter:     counter,
+		fragID:      fragID,
+	}, nil
+}
+
+func (s *Sender) Layout() fec.Layout {
+	return s.enc.Layout()
+}
+
+// SetNext sets the index the next attempt starts from, so a sender rebuilt from
+// a cached message can carry on offering fresh symbols.
+func (s *Sender) SetNext(index int) {
+	s.next = ((index % fec.MaxSymbols) + fec.MaxSymbols) % fec.MaxSymbols
+}
+
+// Next returns the index the next attempt starts from.
+func (s *Sender) Next() int {
+	return s.next
+}
+
+// Attempt returns the packets for one transmission attempt carrying K+repair
+// symbols. now should be sent immediately and later after the repair delay, so
+// a tail drop at the end of the first burst does not also take the repair.
+func (s *Sender) Attempt(repair int) (now, later [][]byte, err error) {
+	if repair < 0 {
+		return nil, nil, errors.New("negative repair count")
+	}
+	k := s.enc.Layout().K
+	now = make([][]byte, 0, k)
+	later = make([][]byte, 0, repair)
+	for i := range k + repair {
+		p, err := s.packet(s.next)
+		if err != nil {
+			return nil, nil, err
+		}
+		s.next = (s.next + 1) % fec.MaxSymbols
+		if i < k {
+			now = append(now, p)
+		} else {
+			later = append(later, p)
+		}
+	}
+	return now, later, nil
+}
+
+func (s *Sender) packet(index int) ([]byte, error) {
+	l := s.enc.Layout()
+	p := make([]byte, header.Len, header.Len+HeaderLen+l.SymbolSize)
+	header.Encode(p, header.Version, header.Handshake, s.subtype, s.remoteIndex, s.counter)
+	p = Header{
+		FragID: s.fragID,
+		MsgLen: uint16(l.MsgLen),
+		Index:  uint8(index),
+	}.Encode(p)
+	return s.enc.Symbol(p, index)
+}

handshake/fragment/table.go

@@ -0,0 +1,212 @@
+package fragment
+
+import (
+	"errors"
+	"math/rand/v2"
+	"sync"
+	"time"
+)
+
+var (
+	ErrTooLarge    = errors.New("fragmented message exceeds the size limit")
+	ErrSourceQuota = errors.New("source has too many messages in progress")
+	ErrNoRoom      = errors.New("no room for fragment")
+)
+
+// TableConfig bounds a Table.
+type TableConfig struct {
+	MaxBytes     int           // symbol bytes held across all entries
+	MaxEntries   int           // messages in progress across all sources
+	MaxPerSource int           // messages in progress per source
+	MaxMsgLen    int           // largest message accepted
+	TTL          time.Duration // an entry idle this long is dropped
+}
+
+// Table reassembles messages from senders that have not authenticated yet, the
+// responder's view of msg1. Every resource is capped. When a cap is hit a
+// random entry is evicted; random eviction is harder to aim than LRU, which a
+// flood can use to flush specific entries.
+//
+// K identifies one message, typically the source plus FragID. S identifies the
+// source for quotas, typically the address masked to /32 or /64.
+//
+// A Table is safe for concurrent use.
+type Table[K, S comparable] struct {
+	mu        sync.Mutex
+	cfg       TableConfig
+	entries   map[K]*tableEntry[K, S]
+	order     []K
+	perSource map[S]int
+	bytes     int
+}
+
+type tableEntry[K, S comparable] struct {
+	a        *Assembly
+	source   S
+	lastSeen time.Time
+	pos      int
+}
+
+func NewTable[K, S comparable](cfg TableConfig) *Table[K, S] {
+	return &Table[K, S]{
+		cfg:       cfg,
+		entries:   make(map[K]*tableEntry[K, S]),
+		perSource: make(map[S]int),
+	}
+}
+
+// Result describes a completed message.
+type Result struct {
+	Msg      []byte // the decoded message
+	Received int    // packet bytes received for it, for the amplification budget
+}
+
+// Add records a fragment parsed by Parse. When the message is complete its entry
+// is removed and the result returned; otherwise the result is nil. A decode
+// error means the message is lost and must be retransmitted.
+func (t *Table[K, S]) Add(key K, source S, now time.Time, h Header, sym []byte) (*Result, error) {
+	if int(h.MsgLen) > t.cfg.MaxMsgLen {
+		return nil, ErrTooLarge
+	}
+
+	t.mu.Lock()
+	e, ok := t.entries[key]
+	if ok && now.Sub(e.lastSeen) > t.cfg.TTL {
+		t.remove(key, e)
+		ok = false
+	}
+	if !ok {
+		var err error
+		e, err = t.insert(key, source, h)
+		if err != nil {
+			t.mu.Unlock()
+			return nil, err
+		}
+	}
+	e.lastSeen = now
+
+	if !t.makeRoom(key, e.a.wouldStore(h, sym)) {
+		if e.a.Stored() == 0 {
+			t.remove(key, e)
+		}
+		t.mu.Unlock()
+		return nil, ErrNoRoom
+	}
+
+	before := e.a.Stored()
+	ready, err := e.a.Add(h, sym)
+	t.bytes += e.a.Stored() - before
+	if err != nil || !ready {
+		if errors.Is(err, ErrMismatch) || (err != nil && e.a.Stored() == 0) {
+			t.remove(key, e)
+		}
+		t.mu.Unlock()
+		return nil, err
+	}
+
+	t.remove(key, e)
+	t.mu.Unlock()
+
+	msg, err := e.a.Decode()
+	if err != nil {
+		return nil, err
+	}
+	return &Result{Msg: msg, Received: e.a.Received()}, nil
+}
+
+// Expire drops every entry idle for longer than the TTL.
+func (t *Table[K, S]) Expire(now time.Time) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	for i := 0; i < len(t.order); {
+		key := t.order[i]
+		e := t.entries[key]
+		if now.Sub(e.lastSeen) > t.cfg.TTL {
+			// remove swaps the last entry into position i, so look at i again.
+			t.remove(key, e)
+			continue
+		}
+		i++
+	}
+}
+
+// Len returns the number of messages in progress.
+func (t *Table[K, S]) Len() int {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return len(t.entries)
+}
+
+// Bytes returns the number of symbol bytes held.
+func (t *Table[K, S]) Bytes() int {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.bytes
+}
+
+func (t *Table[K, S]) insert(key K, source S, h Header) (*tableEntry[K, S], error) {
+	if t.perSource[source] >= t.cfg.MaxPerSource {
+		return nil, ErrSourceQuota
+	}
+	a, err := NewAssembly(h)
+	if err != nil {
+		return nil, err
+	}
+	for len(t.entries) >= t.cfg.MaxEntries {
+		if !t.evictRandom(nil) {
+			return nil, ErrNoRoom
+		}
+	}
+
+	e := &tableEntry[K, S]{a: a, source: source, pos: len(t.order)}
+	t.entries[key] = e
+	t.order = append(t.order, key)
+	t.perSource[source]++
+	return e, nil
+}
+
+// makeRoom evicts random entries other than keep until n more bytes fit.
+func (t *Table[K, S]) makeRoom(keep K, n int) bool {
+	for t.bytes+n > t.cfg.MaxBytes {
+		if !t.evictRandom(&keep) {
+			return false
+		}
+	}
+	return true
+}
+
+func (t *Table[K, S]) evictRandom(keep *K) bool {
+	candidates := len(t.order)
+	if keep != nil {
+		if _, ok := t.entries[*keep]; ok {
+			candidates--
+		}
+	}
+	if candidates <= 0 {
+		return false
+	}
+	for {
+		key := t.order[rand.IntN(len(t.order))]
+		if keep != nil && key == *keep {
+			continue
+		}
+		t.remove(key, t.entries[key])
+		return true
+	}
+}
+
+func (t *Table[K, S]) remove(key K, e *tableEntry[K, S]) {
+	last := len(t.order) - 1
+	moved := t.order[last]
+	t.order[e.pos] = moved
+	t.entries[moved].pos = e.pos
+	t.order = t.order[:last]
+	delete(t.entries, key)
+
+	t.bytes -= e.a.Stored()
+	if t.perSource[e.source] <= 1 {
+		delete(t.perSource, e.source)
+	} else {
+		t.perSource[e.source]--
+	}
+}

handshake/fragment/wire.go

@@ -0,0 +1,92 @@
+// Package fragment carries handshake messages that do not fit in one UDP
+// packet. A message is Reed-Solomon coded into equally sized symbols (see
+// handshake/fec) and each symbol travels in its own packet:
+//
+//	Nebula header (16) | Fragment header (8) | symbol
+//
+// The fragment header is FragID u32 | MsgLen u16 | Index u8 | Flags u8, all big
+// endian. Both ends derive the symbol layout from MsgLen, so any MsgLen-derived
+// value never has to be sent.
+package fragment
+
+import (
+	"encoding/binary"
+	"errors"
+
+	"github.com/slackhq/nebula/handshake/fec"
+	"github.com/slackhq/nebula/header"
+)
+
+const (
+	// HeaderLen is the size of the fragment header that follows the Nebula header.
+	HeaderLen = 8
+
+	// MaxPacketLen bounds every fragment packet, Nebula header included. A relayed
+	// fragment gains a 16 byte header and a 16 byte tag, which lands exactly on
+	// the 1232 byte UDP payload available on an IPv6 minimum MTU path.
+	MaxPacketLen = 1200
+
+	// MaxSymbolSize is the largest symbol that fits in MaxPacketLen.
+	MaxSymbolSize = MaxPacketLen - header.Len - HeaderLen
+)
+
+var (
+	ErrShortFragment = errors.New("fragment is too short")
+	ErrFragID        = errors.New("fragment has a zero FragID")
+	ErrMsgLen        = errors.New("fragment has a zero MsgLen")
+	ErrFlags         = errors.New("fragment has reserved flags set")
+	ErrSymbolSize    = errors.New("fragment symbol size does not match its MsgLen")
+)
+
+// Header is the fragment header.
+type Header struct {
+	FragID uint32 // the sender's local index for this handshake, never zero
+	MsgLen uint16 // length of the whole message
+	Index  uint8  // symbol index, data below K and repair from K
+	Flags  uint8  // reserved, must be zero
+}
+
+// Encode appends the header to b.
+func (h Header) Encode(b []byte) []byte {
+	b = binary.BigEndian.AppendUint32(b, h.FragID)
+	b = binary.BigEndian.AppendUint16(b, h.MsgLen)
+	return append(b, h.Index, h.Flags)
+}
+
+// Layout returns the symbol layout the header's MsgLen implies.
+func (h Header) Layout() (fec.Layout, error) {
+	return fec.NewLayout(int(h.MsgLen), MaxSymbolSize)
+}
+
+// Parse validates the bytes following the Nebula header and splits them into
+// the fragment header and the symbol. The symbol aliases b.
+func Parse(b []byte) (Header, []byte, error) {
+	if len(b) < HeaderLen {
+		return Header{}, nil, ErrShortFragment
+	}
+	h := Header{
+		FragID: binary.BigEndian.Uint32(b[0:4]),
+		MsgLen: binary.BigEndian.Uint16(b[4:6]),
+		Index:  b[6],
+		Flags:  b[7],
+	}
+	if h.FragID == 0 {
+		return Header{}, nil, ErrFragID
+	}
+	if h.MsgLen == 0 {
+		return Header{}, nil, ErrMsgLen
+	}
+	if h.Flags != 0 {
+		return Header{}, nil, ErrFlags
+	}
+
+	l, err := h.Layout()
+	if err != nil {
+		return Header{}, nil, err
+	}
+	sym := b[HeaderLen:]
+	if len(sym) != l.SymbolSize {
+		return Header{}, nil, ErrSymbolSize
+	}
+	return h, sym, nil
+}