handshake/fec: Systematic Cauchy Reed-Solomon erasure code over GF(2^8)

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
+}

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<<n; mask++ {
+				var idx []int
+				for i := range n {
+					if mask&(1<<i) != 0 {
+						idx = append(idx, i)
+					}
+				}
+				if len(idx) != k {
+					continue
+				}
+				require.Equal(t, msg, decodeWith(t, enc, msgLen, maxSym, idx), "k=%d m=%d idx=%v", k, m, idx)
+			}
+		}
+	}
+}
+
+// TestRandomHighIndexes draws symbols from the whole index space, the shape
+// retransmits produce once fresh repair symbols are in play.
+func TestRandomHighIndexes(t *testing.T) {
+	r := rand.New(rand.NewPCG(3, 4))
+	for range 2000 {
+		msgLen := 1 + r.IntN(65535)
+		msg := randMsg(r, msgLen)
+		enc, err := NewEncoder(msg, 1176)
+		require.NoError(t, err)
+		k := enc.Layout().K
+
+		idx := r.Perm(MaxSymbols)[:k]
+		require.Equal(t, msg, decodeWith(t, enc, msgLen, 1176, idx), "k=%d idx=%v", k, idx)
+	}
+}
+
+func TestDecoderBounds(t *testing.T) {
+	r := rand.New(rand.NewPCG(5, 6))
+	msg := randMsg(r, 3000)
+	enc, err := NewEncoder(msg, 1176)
+	require.NoError(t, err)
+	l := enc.Layout()
+	require.Equal(t, 3, l.K)
+	require.Equal(t, 1000, l.SymbolSize)
+
+	dec, err := NewDecoder(3000, 1176)
+	require.NoError(t, err)
+
+	_, err = dec.Decode()
+	require.ErrorIs(t, err, ErrNotEnough)
+
+	_, err = dec.Add(-1, make([]byte, 1000))
+	require.ErrorIs(t, err, ErrIndexRange)
+	_, err = dec.Add(MaxSymbols, make([]byte, 1000))
+	require.ErrorIs(t, err, ErrIndexRange)
+	_, err = dec.Add(0, make([]byte, 999))
+	require.ErrorIs(t, err, ErrSymbolSize)
+
+	s5, _ := enc.Symbol(nil, 5)
+	ready, err := dec.Add(5, s5)
+	require.NoError(t, err)
+	require.False(t, ready)
+
+	// A duplicate index is ignored, even with different contents.
+	ready, err = dec.Add(5, make([]byte, 1000))
+	require.NoError(t, err)
+	require.False(t, ready)
+
+	s0, _ := enc.Symbol(nil, 0)
+	s9, _ := enc.Symbol(nil, 9)
+	_, err = dec.Add(0, s0)
+	require.NoError(t, err)
+	ready, err = dec.Add(9, s9)
+	require.NoError(t, err)
+	require.True(t, ready)
+
+	// Symbols beyond K are not stored.
+	ready, err = dec.Add(1, make([]byte, 1000))
+	require.NoError(t, err)
+	require.True(t, ready)
+	require.Len(t, dec.syms, 3)
+
+	out, err := dec.Decode()
+	require.NoError(t, err)
+	require.Equal(t, msg, out)
+}
+
+func TestCorruptPadding(t *testing.T) {
+	// 3001 bytes in 3 symbols of 1001 leaves 2 bytes of padding in the last one.
+	dec, err := NewDecoder(3001, 1176)
+	require.NoError(t, err)
+	require.Equal(t, 1001, dec.Layout().SymbolSize)
+
+	last := make([]byte, 1001)
+	last[1000] = 1
+	for i, s := range [][]byte{make([]byte, 1001), make([]byte, 1001), last} {
+		_, err = dec.Add(i, s)
+		require.NoError(t, err)
+	}
+	_, err = dec.Decode()
+	require.ErrorIs(t, err, ErrCorrupt)
+}
+
+// TestKnownAnswer pins the encoding. Both ends of a handshake must produce
+// identical repair symbols, so any change here is a wire format change.
+func TestKnownAnswer(t *testing.T) {
+	msg := make([]byte, 3000)
+	for i := range msg {
+		msg[i] = byte(i*7 + 3)
+	}
+	enc, err := NewEncoder(msg, 1176)
+	require.NoError(t, err)
+
+	// Cross-checked against an independent bit-serial GF(2^8) implementation.
+	want := map[int]string{
+		3:   "52124f58276ec6338d3b3f4eed273bc7cabfe4f0f467928962b046ad1b348f50",
+		4:   "a50928b04738981985a7d4a17ffaf8d7e30014ca64fc53c99800ffb61aec4aa7",
+		5:   "ada4c95a7facbc542aaef584e6f0d081b7b0d0eed69ded03b57d0753a119cfc9",
+		255: "839066cdf94f963b223fb1c1c1360f698c14f6349b25653ff5b521791a4021b1",
+	}
+	for idx, h := range want {
+		sym, err := enc.Symbol(nil, idx)
+		require.NoError(t, err)
+		sum := sha256.Sum256(sym)
+		assert.Equal(t, h, hex.EncodeToString(sum[:]), "symbol %d", idx)
+	}
+}
+
+func FuzzRoundTrip(f *testing.F) {
+	f.Add([]byte("hello"), uint16(2), uint64(0))
+	f.Add(make([]byte, 9000), uint16(1176), uint64(7))
+	f.Fuzz(func(t *testing.T, msg []byte, maxSym uint16, seed uint64) {
+		if len(msg) == 0 || maxSym == 0 {
+			return
+		}
+		enc, err := NewEncoder(msg, int(maxSym))
+		if err != nil {
+			return
+		}
+		r := rand.New(rand.NewPCG(seed, seed))
+		idx := r.Perm(MaxSymbols)[:enc.Layout().K]
+		require.Equal(t, msg, decodeWith(t, enc, len(msg), int(maxSym), idx))
+	})
+}
+
+// FuzzDecoder feeds arbitrary symbols to a decoder; it must never panic and
+// must either fail or return a message of the declared length.
+func FuzzDecoder(f *testing.F) {
+	f.Add(uint16(3000), []byte{0, 5, 9, 200}, []byte("some symbol bytes"))
+	f.Fuzz(func(t *testing.T, msgLen uint16, indexes []byte, fill []byte) {
+		dec, err := NewDecoder(int(msgLen), 1176)
+		if err != nil {
+			return
+		}
+		size := dec.Layout().SymbolSize
+		for i, idx := range indexes {
+			sym := make([]byte, size)
+			for j := range sym {
+				if len(fill) > 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)
+		}
+	}
+}

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
+}