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,

handshake_manager_hybrid.go

@@ -214,14 +214,14 @@ func (hm *HandshakeManager) continueHybridResponder(via ViaSender, h *header.H,
 	if err != nil || !ready {
 		return
 	}
-	msg, err := r.msg3.Decode()
+	pkt, err := r.msg3.Decode(make([]byte, header.Len))
 	r.msg3 = nil
 	if err != nil {
 		hm.l.Debug("failed to decode handshake fragments", "from", via, "error", err)
 		return
 	}
 
-	_, result, err := hh.machine.ProcessPacket(nil, reassembled(h, msg))
+	_, result, err := hh.machine.ProcessPacket(nil, reassembled(h, pkt))
 	if err != nil {
 		if hh.machine.Failed() {
 			hm.l.Warn("Failed to process handshake packet, abandoning",

handshake_manager_pq.go

@@ -219,12 +219,11 @@ func pqFragKeyFor(via ViaSender, fragID uint32) (pqFragKey, pqFragSource) {
 	return pqFragKey{addr: via.UdpAddr, fragID: fragID}, pqFragSource{prefix: prefix}
 }
 
-// reassembled rebuilds the packet the handshake Machine expects from the header
-// of a fragment and the decoded message.
-func reassembled(h *header.H, msg []byte) []byte {
-	p := make([]byte, header.Len, header.Len+len(msg))
-	header.Encode(p, header.Version, header.Handshake, h.Subtype, h.RemoteIndex, h.MessageCounter)
-	return append(p, msg...)
+// reassembled turns a decoded message into the packet the handshake Machine
+// expects, by writing the header of one of its fragments into the header.Len
+// bytes left for it in front of the message.
+func reassembled(h *header.H, pkt []byte) []byte {
+	return header.Encode(pkt, header.Version, header.Handshake, h.Subtype, h.RemoteIndex, h.MessageCounter)[:len(pkt)]
 }
 
 // usePQ decides whether this outbound attempt is post-quantum. A peer that does
@@ -474,7 +473,7 @@ func (hm *HandshakeManager) handleFragment(via ViaSender, packet []byte, h *head
 			hm.pqTable.Forget(key)
 			return
 		}
-		hm.beginHandshake(via, reassembled(h, res.Msg), h, res.Received)
+		hm.beginHandshake(via, reassembled(h, res.Packet), h, res.Received)
 		if hm.pqAnswered.lookup(res.Hash, now) != nil {
 			hm.pqTable.Answered(key, res, now)
 		}
@@ -516,13 +515,13 @@ func (hm *HandshakeManager) continuePQ(via ViaSender, hh *HandshakeHostInfo, h *
 		return
 	}
 
-	msg, err := a.Decode()
+	pkt, err := a.Decode(make([]byte, header.Len))
 	pq.dropSlot(a)
 	if err != nil {
 		hm.l.Debug("failed to decode handshake fragments", "from", via, "error", err)
 		return
 	}
-	hm.processResponse(via, hh, pq.machine, reassembled(h, msg))
+	hm.processResponse(via, hh, pq.machine, reassembled(h, pkt))
 }
 
 // sendFragmented sends packets to a peer: now immediately and later after the

handshake_manager_pq_test.go

@@ -470,7 +470,7 @@ func TestPQMsg2Slots(t *testing.T) {
 			continue
 		}
 		require.Equal(t, 11, n, "B completes on its ninth symbol")
-		msg, err := a.Decode()
+		msg, err := a.Decode(nil)
 		require.NoError(t, err)
 		got = msg
 		pq.dropSlot(a)
@@ -488,7 +488,7 @@ func TestPQMsg2Slots(t *testing.T) {
 		c, _, ready = feedMsg2(t, pq, pktsC[i])
 		assert.Equal(t, i == 8, ready, "symbol %d", i)
 	}
-	msg, err := c.Decode()
+	msg, err := c.Decode(nil)
 	require.NoError(t, err)
 	assert.Equal(t, msgC, msg)