diff --git a/handshake/fec/fec.go b/handshake/fec/fec.go new file mode 100644 index 00000000..4007fc3f --- /dev/null +++ b/handshake/fec/fec.go @@ -0,0 +1,216 @@ +// Package fec implements the systematic Reed-Solomon erasure code used to carry +// handshake messages that do not fit in a single UDP packet. +// +// A message is split into K equally sized data symbols (the last one zero +// padded). Symbol indexes 0..K-1 are the data symbols themselves; indexes +// K..MaxSymbols-1 are repair symbols. Any K distinct symbols reconstruct the +// message. Repair symbols are generated on demand, so a sender can emit fresh +// ones on every retransmit and a receiver can merge everything it has seen. +package fec + +import ( + "errors" + "fmt" +) + +// MaxSymbols is the number of distinct symbols, data and repair together, one +// message can be encoded into. It is bounded by the size of GF(2^8). +const MaxSymbols = 256 + +var ( + ErrInvalidLayout = errors.New("invalid fec layout") + ErrIndexRange = errors.New("symbol index out of range") + ErrSymbolSize = errors.New("symbol has the wrong size") + ErrNotEnough = errors.New("not enough symbols to decode") + ErrCorrupt = errors.New("decoded message has non-zero padding") +) + +// Layout describes how a message is split into symbols. Both ends derive it +// from the message length alone, so it never has to be sent. +type Layout struct { + MsgLen int // length of the message in bytes + K int // number of data symbols + SymbolSize int // size of every symbol, data and repair +} + +// NewLayout returns the layout for a message of msgLen bytes carried in symbols +// of at most maxSymbolSize bytes. Symbols are balanced, so padding is always +// less than K bytes. +func NewLayout(msgLen, maxSymbolSize int) (Layout, error) { + if msgLen <= 0 || maxSymbolSize <= 0 { + return Layout{}, ErrInvalidLayout + } + k := (msgLen + maxSymbolSize - 1) / maxSymbolSize + if k >= MaxSymbols { + return Layout{}, fmt.Errorf("%w: %d data symbols", ErrInvalidLayout, k) + } + return Layout{ + MsgLen: msgLen, + K: k, + SymbolSize: (msgLen + k - 1) / k, + }, nil +} + +// Encoder produces symbols for one message. +type Encoder struct { + layout Layout + data [][]byte +} + +// NewEncoder copies msg into padded data symbols. +func NewEncoder(msg []byte, maxSymbolSize int) (*Encoder, error) { + l, err := NewLayout(len(msg), maxSymbolSize) + if err != nil { + return nil, err + } + + buf := make([]byte, l.K*l.SymbolSize) + copy(buf, msg) + data := make([][]byte, l.K) + for i := range data { + data[i] = buf[i*l.SymbolSize : (i+1)*l.SymbolSize] + } + return &Encoder{layout: l, data: data}, nil +} + +func (e *Encoder) Layout() Layout { + return e.layout +} + +// Symbol appends the symbol with the given index to dst and returns the result. +func (e *Encoder) Symbol(dst []byte, index int) ([]byte, error) { + if index < 0 || index >= MaxSymbols { + return nil, ErrIndexRange + } + if index < e.layout.K { + return append(dst, e.data[index]...), nil + } + + start := len(dst) + dst = append(dst, make([]byte, e.layout.SymbolSize)...) + out := dst[start:] + for col, d := range e.data { + mulAdd(out, d, coefficient(index, col)) + } + return dst, nil +} + +// Decoder collects symbols for one message until it can reconstruct it. +// It stores at most K symbols, so its memory is bounded by the layout no matter +// how many symbols arrive. +type Decoder struct { + layout Layout + have [MaxSymbols]bool + syms map[int][]byte +} + +// NewDecoder returns a decoder for a message of msgLen bytes that was encoded +// with the same maxSymbolSize. +func NewDecoder(msgLen, maxSymbolSize int) (*Decoder, error) { + l, err := NewLayout(msgLen, maxSymbolSize) + if err != nil { + return nil, err + } + return &Decoder{layout: l, syms: make(map[int][]byte, l.K)}, nil +} + +func (d *Decoder) Layout() Layout { + return d.layout +} + +// Stored returns the number of symbols held. +func (d *Decoder) Stored() int { + return len(d.syms) +} + +// Has reports whether the symbol with the given index is held. +func (d *Decoder) Has(index int) bool { + return index >= 0 && index < MaxSymbols && d.have[index] +} + +// Ready reports whether enough symbols have arrived to decode. +func (d *Decoder) Ready() bool { + return len(d.syms) >= d.layout.K +} + +// Add stores a copy of the symbol. Duplicates, and anything beyond the first K +// distinct symbols, are ignored. It returns whether the decoder is ready. +func (d *Decoder) Add(index int, sym []byte) (bool, error) { + if index < 0 || index >= MaxSymbols { + return d.Ready(), ErrIndexRange + } + if len(sym) != d.layout.SymbolSize { + return d.Ready(), ErrSymbolSize + } + if d.have[index] || d.Ready() { + return d.Ready(), nil + } + d.have[index] = true + d.syms[index] = append([]byte(nil), sym...) + return d.Ready(), nil +} + +// Decode reconstructs the message from the stored symbols. +func (d *Decoder) Decode() ([]byte, error) { + if !d.Ready() { + return nil, ErrNotEnough + } + l := d.layout + + out := make([]byte, l.K*l.SymbolSize) + var missing, repair []int + for i := range l.K { + if s, ok := d.syms[i]; ok { + copy(out[i*l.SymbolSize:], s) + } else { + missing = append(missing, i) + } + } + + if len(missing) > 0 { + for i := l.K; i < MaxSymbols && len(repair) < len(missing); i++ { + if d.have[i] { + repair = append(repair, i) + } + } + if len(repair) < len(missing) { + return nil, ErrNotEnough + } + + // Each repair symbol r satisfies r = sum_col coefficient(r, col) * data[col]. + // Move the known data columns to the left hand side, leaving a square + // Cauchy system in the missing columns. + rhs := make([][]byte, len(repair)) + m := make([][]byte, len(repair)) + for ri, row := range repair { + rhs[ri] = append([]byte(nil), d.syms[row]...) + m[ri] = make([]byte, len(missing)) + mi := 0 + for col := range l.K { + if mi < len(missing) && missing[mi] == col { + m[ri][mi] = coefficient(row, col) + mi++ + continue + } + mulAdd(rhs[ri], out[col*l.SymbolSize:(col+1)*l.SymbolSize], coefficient(row, col)) + } + } + if !invert(m) { + // Unreachable for a Cauchy matrix; kept so a broken invariant fails loudly. + return nil, errors.New("fec: singular decode matrix") + } + for mi, col := range missing { + dst := out[col*l.SymbolSize : (col+1)*l.SymbolSize] + for ri := range repair { + mulAdd(dst, rhs[ri], m[mi][ri]) + } + } + } + + for _, b := range out[l.MsgLen:] { + if b != 0 { + return nil, ErrCorrupt + } + } + return out[:l.MsgLen], nil +} diff --git a/handshake/fec/fec_test.go b/handshake/fec/fec_test.go new file mode 100644 index 00000000..91dde6d6 --- /dev/null +++ b/handshake/fec/fec_test.go @@ -0,0 +1,318 @@ +package fec + +import ( + "crypto/sha256" + "encoding/hex" + "math/rand/v2" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// slowMul is carry-less multiplication reduced by gfPoly, used to check the tables. +func slowMul(a, b byte) byte { + var p int + x, y := int(a), int(b) + for y > 0 { + if y&1 != 0 { + p ^= x + } + x <<= 1 + if x&0x100 != 0 { + x ^= gfPoly + } + y >>= 1 + } + return byte(p) +} + +func TestGF(t *testing.T) { + for a := range 256 { + for b := range 256 { + require.Equal(t, slowMul(byte(a), byte(b)), gfMul(byte(a), byte(b)), "%d*%d", a, b) + } + if a != 0 { + require.Equal(t, byte(1), gfMul(byte(a), gfInv(byte(a))), "inv(%d)", a) + } + } + + src := make([]byte, 256) + for i := range src { + src[i] = byte(i) + } + for c := range 256 { + dst := make([]byte, 256) + mulAdd(dst, src, byte(c)) + for i := range src { + require.Equal(t, gfMul(byte(c), src[i]), dst[i]) + } + } +} + +func TestLayout(t *testing.T) { + for msgLen := 1; msgLen <= 20000; msgLen++ { + l, err := NewLayout(msgLen, 1176) + require.NoError(t, err) + require.Equal(t, (msgLen+1175)/1176, l.K) + require.LessOrEqual(t, l.SymbolSize, 1176) + require.GreaterOrEqual(t, l.K*l.SymbolSize, msgLen) + require.Less(t, l.K*l.SymbolSize-msgLen, l.K, "padding must be less than one byte per symbol") + } + + l, err := NewLayout(65535, 1176) + require.NoError(t, err) + assert.Equal(t, 56, l.K) + + _, err = NewLayout(0, 1176) + require.ErrorIs(t, err, ErrInvalidLayout) + _, err = NewLayout(100, 0) + require.ErrorIs(t, err, ErrInvalidLayout) + _, err = NewLayout(256, 1) + require.ErrorIs(t, err, ErrInvalidLayout) + _, err = NewLayout(255, 1) + require.NoError(t, err) +} + +func randMsg(r *rand.Rand, n int) []byte { + b := make([]byte, n) + for i := range b { + b[i] = byte(r.Uint32()) + } + return b +} + +func decodeWith(t *testing.T, enc *Encoder, msgLen, maxSym int, indexes []int) []byte { + t.Helper() + dec, err := NewDecoder(msgLen, maxSym) + require.NoError(t, err) + for _, idx := range indexes { + sym, err := enc.Symbol(nil, idx) + require.NoError(t, err) + _, err = dec.Add(idx, sym) + require.NoError(t, err) + } + require.True(t, dec.Ready()) + out, err := dec.Decode() + require.NoError(t, err) + return out +} + +// TestExhaustiveSubsets decodes every k-subset of the first k+m symbols for +// small codes, which checks the MDS property directly. +func TestExhaustiveSubsets(t *testing.T) { + r := rand.New(rand.NewPCG(1, 2)) + const maxSym = 16 + for k := 1; k <= 10; k++ { + for m := 1; m <= 5; m++ { + msgLen := (k-1)*maxSym + 1 + r.IntN(maxSym) + msg := randMsg(r, msgLen) + enc, err := NewEncoder(msg, maxSym) + require.NoError(t, err) + require.Equal(t, k, enc.Layout().K) + + n := k + m + for mask := 0; mask < 1< 0 { + sym[j] = fill[(i+j)%len(fill)] + } + } + if _, err := dec.Add(int(idx), sym); err != nil { + t.Fatal(err) + } + } + out, err := dec.Decode() + if err == nil { + require.Len(t, out, int(msgLen)) + } + }) +} + +// BenchmarkRepair and BenchmarkDecode use a v3-sized msg2 (~9.8KB, k=9) with +// three repair symbols, decoding with three data symbols missing. +func BenchmarkRepair(b *testing.B) { + msg := randMsg(rand.New(rand.NewPCG(7, 8)), 9800) + enc, err := NewEncoder(msg, 1176) + require.NoError(b, err) + buf := make([]byte, 0, 1176) + b.SetBytes(int64(len(msg))) + for b.Loop() { + for idx := 9; idx < 12; idx++ { + buf, _ = enc.Symbol(buf[:0], idx) + } + } +} + +func BenchmarkDecode(b *testing.B) { + msg := randMsg(rand.New(rand.NewPCG(7, 8)), 9800) + enc, err := NewEncoder(msg, 1176) + require.NoError(b, err) + idx := []int{0, 1, 2, 3, 5, 6, 9, 10, 11} + syms := make([][]byte, len(idx)) + for i, x := range idx { + syms[i], _ = enc.Symbol(nil, x) + } + b.SetBytes(int64(len(msg))) + for b.Loop() { + dec, _ := NewDecoder(len(msg), 1176) + for i, x := range idx { + _, _ = dec.Add(x, syms[i]) + } + if _, err := dec.Decode(); err != nil { + b.Fatal(err) + } + } +} diff --git a/handshake/fec/gf.go b/handshake/fec/gf.go new file mode 100644 index 00000000..027c73e3 --- /dev/null +++ b/handshake/fec/gf.go @@ -0,0 +1,115 @@ +package fec + +// Arithmetic over GF(2^8) with the reducing polynomial x^8+x^4+x^3+x^2+1 (0x11d). +// Addition is XOR. Multiplication uses log/exp tables; the exp table is doubled so +// the sum of two logs never needs reducing mod 255. + +const gfPoly = 0x11d + +var ( + gfExp [510]byte + gfLog [256]byte +) + +func init() { + x := 1 + for i := range 255 { + gfExp[i] = byte(x) + gfExp[i+255] = byte(x) + gfLog[x] = byte(i) + x <<= 1 + if x&0x100 != 0 { + x ^= gfPoly + } + } +} + +func gfMul(a, b byte) byte { + if a == 0 || b == 0 { + return 0 + } + return gfExp[int(gfLog[a])+int(gfLog[b])] +} + +// gfInv returns the multiplicative inverse of a. a must be non-zero. +func gfInv(a byte) byte { + return gfExp[255-int(gfLog[a])] +} + +// mulAdd computes dst ^= c*src element-wise. dst must be at least as long as src. +func mulAdd(dst, src []byte, c byte) { + switch c { + case 0: + return + case 1: + for i, s := range src { + dst[i] ^= s + } + return + } + + var t [256]byte + lc := int(gfLog[c]) + for s := 1; s < 256; s++ { + t[s] = gfExp[int(gfLog[s])+lc] + } + for i, s := range src { + dst[i] ^= t[s] + } +} + +// coefficient returns the generator matrix entry for repair symbol index row and +// data symbol index col. Repair rows are a Cauchy matrix 1/(x_row + y_col) with +// x_row = row and y_col = col: rows are always >= k and columns < k so the two +// sets are disjoint, which makes every square submatrix invertible and the +// systematic code [I; C] MDS. A row depends only on its index, never on how many +// repair symbols exist, so new repair symbols can be generated at any time. +func coefficient(row, col int) byte { + return gfInv(byte(row) ^ byte(col)) +} + +// invert inverts the n x n matrix m in place with Gauss-Jordan elimination. +// It returns false if m is singular. +func invert(m [][]byte) bool { + n := len(m) + inv := make([][]byte, n) + for i := range inv { + inv[i] = make([]byte, n) + inv[i][i] = 1 + } + + for col := range n { + pivot := -1 + for r := col; r < n; r++ { + if m[r][col] != 0 { + pivot = r + break + } + } + if pivot < 0 { + return false + } + m[col], m[pivot] = m[pivot], m[col] + inv[col], inv[pivot] = inv[pivot], inv[col] + + scale := gfInv(m[col][col]) + for c := range n { + m[col][c] = gfMul(m[col][c], scale) + inv[col][c] = gfMul(inv[col][c], scale) + } + + for r := 0; r < n; r++ { + if r == col || m[r][col] == 0 { + continue + } + f := m[r][col] + for c := range n { + m[r][c] ^= gfMul(f, m[col][c]) + inv[r][c] ^= gfMul(f, inv[col][c]) + } + } + } + + copy(m, inv) + return true +}