nebula/handshake/fragment/fragment_test.go

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: 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) {
	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())

	// 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)
		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("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)
		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)
		}
	})
}