From cbcf441d7c65a5e2a27e92ed998386d468de9c61 Mon Sep 17 00:00:00 2001 From: JackDoan Date: Wed, 14 Oct 2026 15:30:15 -0500 Subject: [PATCH] Decode a fragmented handshake message straight into its packet, behind room for the Nebula header, instead of copying it there --- handshake/fec/fec.go | 18 ++++++++++++------ handshake/fec/fec_test.go | 22 +++++++++++++++------- handshake/fragment/assembly.go | 7 ++++--- handshake/fragment/fragment_test.go | 2 ++ handshake/fragment/table.go | 5 ++++- handshake_manager_hybrid.go | 4 ++-- handshake_manager_pq.go | 17 ++++++++--------- handshake_manager_pq_test.go | 4 ++-- 8 files changed, 49 insertions(+), 30 deletions(-) diff --git a/handshake/fec/fec.go b/handshake/fec/fec.go index 1e807bff..15122a83 100644 --- a/handshake/fec/fec.go +++ b/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 } diff --git a/handshake/fec/fec_test.go b/handshake/fec/fec_test.go index 91dde6d6..faf4a9f0 100644 --- a/handshake/fec/fec_test.go +++ b/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) } } diff --git a/handshake/fragment/assembly.go b/handshake/fragment/assembly.go index d6925cc5..e170f8f0 100644 --- a/handshake/fragment/assembly.go +++ b/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. diff --git a/handshake/fragment/fragment_test.go b/handshake/fragment/fragment_test.go index a71a7445..7d6b7ece 100644 --- a/handshake/fragment/fragment_test.go +++ b/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) diff --git a/handshake/fragment/table.go b/handshake/fragment/table.go index 565861a4..30452d0a 100644 --- a/handshake/fragment/table.go +++ b/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, diff --git a/handshake_manager_hybrid.go b/handshake_manager_hybrid.go index af24af28..4fee8a93 100644 --- a/handshake_manager_hybrid.go +++ b/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", diff --git a/handshake_manager_pq.go b/handshake_manager_pq.go index f6beee03..08373b97 100644 --- a/handshake_manager_pq.go +++ b/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 diff --git a/handshake_manager_pq_test.go b/handshake_manager_pq_test.go index f50262aa..a802d1e3 100644 --- a/handshake_manager_pq_test.go +++ b/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)