mirror of
https://github.com/slackhq/nebula.git
synced 2026-11-15 12:56:52 +01:00
handshake/fragment: Carry oversized handshake messages as FEC coded fragments
This commit is contained in:
JackDoan committed
1 parent
8f9a631f6e
commit
31f3d6435e
5 files changed
+944
No files matched your search
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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]--
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in new issue
Block a user