Decode a fragmented handshake message straight into its packet, behind room for the Nebula header, instead of copying it there

handshake/fec/fec.go

@@ -11,6 +11,7 @@ package fec
 import (
 	"errors"
 	"fmt"
+	"slices"
 )
 
 // MaxSymbols is the number of distinct symbols, data and repair together, one
@@ -108,8 +109,8 @@ func (e *Encoder) Symbol(dst []byte, index int) ([]byte, error) {
 // folds the known columns into each parked repair symbol in place and runs
 // Gauss-Jordan across those rows in place; a Cauchy submatrix needs no
 // pivoting, so each solution lands in the slot of its own column and the
-// buffer is the packet. That removes the per-symbol copies, the output buffer
-// and reassembled()'s copy, about 20KB and a dozen allocations per msg2, at the
+// buffer is the packet. That removes the per-symbol copies, about 10KB and a
+// dozen allocations per msg2, at the
 // cost of holding K*SymbolSize from the first fragment and destroying the
 // input on a failed decode. It is not for the responder's msg1 table, which
 // must never allocate from a claimed size.
@@ -170,19 +171,24 @@ func (d *Decoder) Add(index int, sym []byte) (bool, error) {
 	return d.Ready(), nil
 }
 
-// Decode reconstructs the message from the stored symbols.
-func (d *Decoder) Decode() ([]byte, error) {
+// Decode reconstructs the message from the stored symbols and appends it to
+// dst. Decoding needs the padding of the last symbol too, so dst has to have
+// room for K*SymbolSize bytes or it is grown to.
+func (d *Decoder) Decode(dst []byte) ([]byte, error) {
 	if !d.Ready() {
 		return nil, ErrNotEnough
 	}
 	l := d.layout
 
-	out := make([]byte, l.K*l.SymbolSize)
+	start := len(dst)
+	dst = slices.Grow(dst, l.K*l.SymbolSize)[:start+l.K*l.SymbolSize]
+	out := dst[start:]
 	var missing, repair []int
 	for i := range l.K {
 		if s, ok := d.syms[i]; ok {
 			copy(out[i*l.SymbolSize:], s)
 		} else {
+			clear(out[i*l.SymbolSize : (i+1)*l.SymbolSize])
 			missing = append(missing, i)
 		}
 	}
@@ -232,5 +238,5 @@ func (d *Decoder) Decode() ([]byte, error) {
 			return nil, ErrCorrupt
 		}
 	}
-	return out[:l.MsgLen], nil
+	return dst[:start+l.MsgLen], nil
 }

handshake/fec/fec_test.go

@@ -93,9 +93,17 @@ func decodeWith(t *testing.T, enc *Encoder, msgLen, maxSym int, indexes []int) [
 		require.NoError(t, err)
 	}
 	require.True(t, dec.Ready())
-	out, err := dec.Decode()
+	// Decode after a prefix, into capacity holding garbage, which must neither
+	// disturb the prefix nor leak into the message.
+	dst := make([]byte, 3, 3+dec.Layout().K*dec.Layout().SymbolSize)
+	for i := range dst[:cap(dst)] {
+		dst[:cap(dst)][i] = 0xa5
+	}
+	out, err := dec.Decode(dst)
 	require.NoError(t, err)
-	return out
+	require.Equal(t, []byte{0xa5, 0xa5, 0xa5}, out[:3])
+	require.Equal(t, &dst[:1][0], &out[:1][0], "decoded in place")
+	return out[3:]
 }
 
 // TestExhaustiveSubsets decodes every k-subset of the first k+m symbols for
@@ -156,7 +164,7 @@ func TestDecoderBounds(t *testing.T) {
 	dec, err := NewDecoder(3000, 1176)
 	require.NoError(t, err)
 
-	_, err = dec.Decode()
+	_, err = dec.Decode(nil)
 	require.ErrorIs(t, err, ErrNotEnough)
 
 	_, err = dec.Add(-1, make([]byte, 1000))
@@ -190,7 +198,7 @@ func TestDecoderBounds(t *testing.T) {
 	require.True(t, ready)
 	require.Len(t, dec.syms, 3)
 
-	out, err := dec.Decode()
+	out, err := dec.Decode(nil)
 	require.NoError(t, err)
 	require.Equal(t, msg, out)
 }
@@ -207,7 +215,7 @@ func TestCorruptPadding(t *testing.T) {
 		_, err = dec.Add(i, s)
 		require.NoError(t, err)
 	}
-	_, err = dec.Decode()
+	_, err = dec.Decode(nil)
 	require.ErrorIs(t, err, ErrCorrupt)
 }
 
@@ -274,7 +282,7 @@ func FuzzDecoder(f *testing.F) {
 				t.Fatal(err)
 			}
 		}
-		out, err := dec.Decode()
+		out, err := dec.Decode(nil)
 		if err == nil {
 			require.Len(t, out, int(msgLen))
 		}
@@ -311,7 +319,7 @@ func BenchmarkDecode(b *testing.B) {
 		for i, x := range idx {
 			_, _ = dec.Add(x, syms[i])
 		}
-		if _, err := dec.Decode(); err != nil {
+		if _, err := dec.Decode(nil); err != nil {
 			b.Fatal(err)
 		}
 	}

handshake/fragment/assembly.go

@@ -50,9 +50,10 @@ func (a *Assembly) wouldStore(h Header, sym []byte) int {
 	return len(sym)
 }
 
-// Decode reconstructs the message. It requires Add to have reported ready.
-func (a *Assembly) Decode() ([]byte, error) {
-	return a.dec.Decode()
+// Decode reconstructs the message and appends it to dst, see fec.Decoder.Decode.
+// It requires Add to have reported ready.
+func (a *Assembly) Decode(dst []byte) ([]byte, error) {
+	return a.dec.Decode(dst)
 }
 
 // Stored returns the number of symbol bytes held.

handshake/fragment/fragment_test.go

@@ -241,6 +241,8 @@ func TestLossy(t *testing.T) {
 					// duplicate can complete one again; the handshake layer
 					// recognises the repeat.
 					require.True(t, bytes.Equal(msg, res.Msg))
+					require.Len(t, res.Packet, header.Len+len(msg))
+					require.Equal(t, &res.Packet[header.Len], &res.Msg[0], "Msg follows room for a header")
 					got = res
 				}
 				checkInvariants(t, tb)

handshake/fragment/table.go

@@ -91,6 +91,7 @@ func NewTable[K, S comparable](cfg TableConfig) *Table[K, S] {
 // Result describes a completed message.
 type Result struct {
 	Msg      []byte            // the decoded message, nil for a repeat
+	Packet   []byte            // header.Len bytes for the caller's Nebula header, then Msg
 	Hash     [sha256.Size]byte // sha256 of the message
 	Received int               // packet bytes received for it, for the amplification budget
 
@@ -159,12 +160,14 @@ func (t *Table[K, S]) Add(key K, source S, now time.Time, h Header, sym []byte)
 	t.remove(key, e)
 	t.mu.Unlock()
 
-	msg, err := e.a.Decode()
+	pkt, err := e.a.Decode(make([]byte, header.Len))
 	if err != nil {
 		return nil, err
 	}
+	msg := pkt[header.Len:]
 	return &Result{
 		Msg:      msg,
+		Packet:   pkt,
 		Hash:     sha256.Sum256(msg),
 		Received: e.a.Received(),
 		msgLen:   e.a.msgLen,