handshake/fragment: Carry an optional cookie, with room for it reserved in every fragment

handshake/fragment/fragment_test.go

@@ -97,9 +97,28 @@ func TestParse(t *testing.T) {
 	_, _, err = Parse(bad)
 	require.ErrorIs(t, err, ErrMsgLen)
 
-	bad = Header{FragID: 1, MsgLen: 3000, Flags: 1}.Encode(nil)
+	bad = Header{FragID: 1, MsgLen: 3000, Flags: 0x02}.Encode(nil)
 	_, _, err = Parse(append(bad, make([]byte, 1000)...))
 	require.ErrorIs(t, err, ErrFlags)
+
+	t.Run("cookie", func(t *testing.T) {
+		h := Header{FragID: 7, MsgLen: 3000, Index: 4, Flags: FlagCookie}
+		for i := range h.Cookie {
+			h.Cookie[i] = byte(i + 1)
+		}
+		b := h.Encode(nil)
+		require.Len(t, b, HeaderLen+CookieLen)
+		b = append(b, make([]byte, 1000)...)
+
+		got, sym, err := Parse(b)
+		require.NoError(t, err)
+		assert.Equal(t, h, got)
+		assert.True(t, got.HasCookie())
+		assert.Len(t, sym, 1000)
+
+		_, _, err = Parse(b[:HeaderLen+CookieLen-1])
+		require.ErrorIs(t, err, ErrNoCookie)
+	})
 }
 
 func TestDefaultRepair(t *testing.T) {
@@ -158,6 +177,20 @@ func TestSender(t *testing.T) {
 	s2.SetNext(-1)
 	assert.Equal(t, 255, s2.Next())
 
+	// A cookie rides on every later packet, which still fits MaxPacketLen even
+	// with a full size symbol.
+	full, err := NewSender(make([]byte, 2*MaxSymbolSize), testSubtype, 0, 1, 9)
+	require.NoError(t, err)
+	full.SetCookie(Cookie{1, 2, 3})
+	pkts, _, err := full.Attempt(1)
+	require.NoError(t, err)
+	for _, p := range pkts {
+		assert.Len(t, p, MaxPacketLen)
+		_, h, _ := parse(t, p)
+		assert.True(t, h.HasCookie())
+		assert.Equal(t, Cookie{1, 2, 3}, h.Cookie)
+	}
+
 	// Indexes wrap once the space is used up.
 	for range 19 {
 		_, _, err = s.Attempt(3)
@@ -380,6 +413,22 @@ func TestTableLimits(t *testing.T) {
 		checkInvariants(t, tb)
 	})
 
+	t.Run("load", func(t *testing.T) {
+		cfg := testConfig()
+		cfg.MaxEntries = 4
+		cfg.MaxBytes = 10000
+		tb := NewTable[testKey, string](cfg)
+		assert.Zero(t, tb.Load())
+		h, sym := firstFragment(t, 5000, 1)
+		_, err := tb.Add(testKey{"a", 1}, "a", now, h, sym)
+		require.NoError(t, err)
+		assert.InDelta(t, 0.25, tb.Load(), 0.001, "one of four entries")
+		h, sym = firstFragment(t, 5000, 2)
+		_, err = tb.Add(testKey{"b", 2}, "b", now, h, sym)
+		require.NoError(t, err)
+		assert.InDelta(t, 0.5, tb.Load(), 0.001)
+	})
+
 	t.Run("ttl", func(t *testing.T) {
 		tb := NewTable[testKey, string](testConfig())
 		s, err := NewSender(make([]byte, 5000), testSubtype, 0, 1, 1)

handshake/fragment/sender.go

@@ -30,6 +30,7 @@ type Sender struct {
 	counter     uint64
 	fragID      uint32
 	next        int
+	cookie      *Cookie
 }
 
 // NewSender prepares msg for sending. subtype, remoteIndex and counter fill the
@@ -64,6 +65,11 @@ func (s *Sender) SetNext(index int) {
 	s.next = ((index % fec.MaxSymbols) + fec.MaxSymbols) % fec.MaxSymbols
 }
 
+// SetCookie attaches a cookie from a responder under load to every later packet.
+func (s *Sender) SetCookie(c Cookie) {
+	s.cookie = &c
+}
+
 // Next returns the index the next attempt starts from.
 func (s *Sender) Next() int {
 	return s.next
@@ -96,12 +102,17 @@ func (s *Sender) Attempt(repair int) (now, later [][]byte, err error) {
 
 func (s *Sender) packet(index int) ([]byte, error) {
 	l := s.enc.Layout()
-	p := make([]byte, header.Len, header.Len+HeaderLen+l.SymbolSize)
+	p := make([]byte, header.Len, header.Len+HeaderLen+CookieLen+l.SymbolSize)
 	header.Encode(p, header.Version, header.Handshake, s.subtype, s.remoteIndex, s.counter)
-	p = Header{
+	h := Header{
 		FragID: s.fragID,
 		MsgLen: uint16(l.MsgLen),
 		Index:  uint8(index),
-	}.Encode(p)
+	}
+	if s.cookie != nil {
+		h.Flags |= FlagCookie
+		h.Cookie = *s.cookie
+	}
+	p = h.Encode(p)
 	return s.enc.Symbol(p, index)
 }

handshake/fragment/table.go

@@ -137,6 +137,14 @@ func (t *Table[K, S]) Len() int {
 	return len(t.entries)
 }
 
+// Load returns how full the table is, from 0 to 1, by whichever of its byte and
+// entry limits is closer.
+func (t *Table[K, S]) Load() float64 {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return max(float64(t.bytes)/float64(t.cfg.MaxBytes), float64(len(t.entries))/float64(t.cfg.MaxEntries))
+}
+
 // Bytes returns the number of symbol bytes held.
 func (t *Table[K, S]) Bytes() int {
 	t.mu.Lock()

handshake/fragment/wire.go

@@ -2,11 +2,13 @@
 // 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
+//	Nebula header (16) | Fragment header (8) | [cookie (16)] | 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.
+// value never has to be sent. FlagCookie marks a cookie a responder under load
+// asked for; room for it is always reserved, so a fragment never outgrows
+// MaxPacketLen.
 package fragment
 
 import (
@@ -21,13 +23,19 @@ const (
 	// HeaderLen is the size of the fragment header that follows the Nebula header.
 	HeaderLen = 8
 
+	// CookieLen is the size of the cookie a fragment carries when FlagCookie is set.
+	CookieLen = 16
+
+	// FlagCookie marks a fragment that carries a cookie after its header.
+	FlagCookie uint8 = 0x01
+
 	// 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
+	// MaxSymbolSize is the largest symbol that fits in MaxPacketLen alongside a cookie.
+	MaxSymbolSize = MaxPacketLen - header.Len - HeaderLen - CookieLen
 )
 
 var (
@@ -35,22 +43,35 @@ var (
 	ErrFragID        = errors.New("fragment has a zero FragID")
 	ErrMsgLen        = errors.New("fragment has a zero MsgLen")
 	ErrFlags         = errors.New("fragment has reserved flags set")
+	ErrNoCookie      = errors.New("fragment is too short for its cookie")
 	ErrSymbolSize    = errors.New("fragment symbol size does not match its MsgLen")
 )
 
+// Cookie proves a sender received a responder's reply at its address.
+type Cookie [CookieLen]byte
+
 // 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
+	Flags  uint8  // FlagCookie, other bits reserved and zero
+	Cookie Cookie // valid when Flags has FlagCookie
 }
 
-// Encode appends the header to b.
+func (h Header) HasCookie() bool {
+	return h.Flags&FlagCookie != 0
+}
+
+// Encode appends the header, and its cookie if it has one, 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)
+	b = append(b, h.Index, h.Flags)
+	if h.HasCookie() {
+		b = append(b, h.Cookie[:]...)
+	}
+	return b
 }
 
 // Layout returns the symbol layout the header's MsgLen implies.
@@ -76,15 +97,23 @@ func Parse(b []byte) (Header, []byte, error) {
 	if h.MsgLen == 0 {
 		return Header{}, nil, ErrMsgLen
 	}
-	if h.Flags != 0 {
+	if h.Flags&^FlagCookie != 0 {
 		return Header{}, nil, ErrFlags
 	}
+	b = b[HeaderLen:]
+	if h.HasCookie() {
+		if len(b) < CookieLen {
+			return Header{}, nil, ErrNoCookie
+		}
+		copy(h.Cookie[:], b)
+		b = b[CookieLen:]
+	}
 
 	l, err := h.Layout()
 	if err != nil {
 		return Header{}, nil, err
 	}
-	sym := b[HeaderLen:]
+	sym := b
 	if len(sym) != l.SymbolSize {
 		return Header{}, nil, ErrSymbolSize
 	}