diff --git a/handshake/fragment/assembly.go b/handshake/fragment/assembly.go new file mode 100644 index 00000000..d6925cc5 --- /dev/null +++ b/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 +} diff --git a/handshake/fragment/fragment_test.go b/handshake/fragment/fragment_test.go new file mode 100644 index 00000000..e8e49d2e --- /dev/null +++ b/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) + } + }) +} diff --git a/handshake/fragment/sender.go b/handshake/fragment/sender.go new file mode 100644 index 00000000..16cafa37 --- /dev/null +++ b/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) +} diff --git a/handshake/fragment/table.go b/handshake/fragment/table.go new file mode 100644 index 00000000..94694dd9 --- /dev/null +++ b/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]-- + } +} diff --git a/handshake/fragment/wire.go b/handshake/fragment/wire.go new file mode 100644 index 00000000..a666e22c --- /dev/null +++ b/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 +}