Reassemble msg2 in a slot per responder FragID, so two hosts answering one msg1 do not wipe each other out

docs/pq-handshake.md

@@ -319,6 +319,12 @@ Implemented in `handshake/fragment` (`Sender`, `Assembly`, and the capped `Table
   FragID) with a total memory cap, per-source quotas (/32 for IPv4, /64 for IPv6), a
   sliding TTL of about 2 s refreshed by progress, and random eviction under pressure.
   msg2 and msg3 reassembly attaches to the existing pending handshake found by RemoteIndex.
+  The initiator keeps two msg2 slots keyed by the responder's FragID, each capped at 16 KiB.
+  When a stale lighthouse address for the peer belongs to another host now, both hosts answer
+  the same msg1, and their interleaved bursts would wipe each other out of a single slot that
+  emptied on every FragID change; with a slot each, whichever completes first is processed, and
+  a wrong host's answer takes the usual block-and-retry path. The responder's msg3 slot needs no
+  such care: it checks FragID against the initiator's index it learned from msg2.
 
 ### Retransmission and amplification
 

e2e/pq_hybrid_test.go

@@ -3,6 +3,7 @@
 package e2e
 
 import (
+	"fmt"
 	"net/netip"
 	"sync"
 	"testing"
@@ -501,3 +502,84 @@ func TestHybridRaceWithClassical(t *testing.T) {
 	myControl.Stop()
 	theirControl.Stop()
 }
+
+// TestHybridWrongResponder gives the initiator a stale address for its peer
+// that belongs to another v3 host now, so both answer its msg1. Their msg2
+// bursts interleave at the initiator and each must fill its own slot: the
+// tunnel to the right host comes up within two attempts, one if its answer
+// completes first and two if the other host's does, which the initiator
+// blocks and retries past as it would for a classical handshake. With require
+// there is no classical turn to fall back on, so a single slot that emptied
+// on every FragID change would run to the handshake timeout.
+func TestHybridWrongResponder(t *testing.T) {
+	t.Parallel()
+	for _, tc := range []struct {
+		name   string
+		cfg    m
+		pqOnly bool
+	}{
+		{"taking turns", pqConfig(true, true, false), false},
+		{"pq only", pqRequire(), true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			ca := newHybridCA(t)
+			myControl, myVpn, myUdp, _ := newSimpleServer(cert.Version3, ca.ca, ca.key, "me", "10.128.0.1/24", tc.cfg)
+			theirControl, theirVpn, theirUdp, _ := newSimpleServer(cert.Version3, ca.ca, ca.key, "them", "10.128.0.2/24", tc.cfg)
+			otherControl, otherVpn, otherUdp, _ := newSimpleServer(cert.Version3, ca.ca, ca.key, "other", "10.128.0.3/24", tc.cfg)
+			myControl.InjectLightHouseAddr(theirVpn[0].Addr(), otherUdp)
+			myControl.InjectLightHouseAddr(theirVpn[0].Addr(), theirUdp)
+
+			r := router.NewR(t, myControl, theirControl, otherControl)
+			defer r.RenderFlow()
+			myControl.Start()
+			theirControl.Start()
+			otherControl.Start()
+
+			var tally hybridTally
+			msg1ToThem := 0
+			msg2ByFragID := map[uint32]int{}
+			var classical []string
+			myControl.InjectTunPacket(BuildTunUDPPacket(theirVpn[0].Addr(), 80, myVpn[0].Addr(), 80, []byte("Hi from me")))
+			require.True(t, routeUntilData(t, r, theirControl, 5*time.Second, func(p *udp.Packet, receiver *nebula.Control) router.ExitType {
+				tally.observe(p)
+				if h, fh := handshakeOf(p); h != nil {
+					switch {
+					case fh == nil:
+						classical = append(classical, fmt.Sprintf("counter=%d to=%s", h.MessageCounter, p.To))
+					case h.MessageCounter == 1 && p.To == theirUdp:
+						msg1ToThem++
+					case h.MessageCounter == 2 && receiver == myControl:
+						msg2ByFragID[fh.FragID]++
+					}
+				}
+				return router.KeepRouting
+			}), "tunnel never came up")
+			t.Logf("msg1 fragments to them: %d, msg2 fragments by FragID: %v, classical: %v", msg1ToThem, msg2ByFragID, classical)
+
+			assertUdpPacket(t, []byte("Hi from me"), theirControl.GetFromTun(true), myVpn[0].Addr(), theirVpn[0].Addr(), 80, 80)
+			assert.GreaterOrEqual(t, len(msg2ByFragID), 2, "both hosts answered")
+			// A responder only resends msg2 when the initiator retries, so every
+			// FragID staying at one burst means each interleaved answer completed in
+			// its own slot. (The msg1 count is not a clean measure: a restart after
+			// the wrong host answers leaves the old handshake's timer entry in place,
+			// so the new one ticks twice, a few milliseconds apart.)
+			for fragID, n := range msg2ByFragID {
+				assert.LessOrEqual(t, n, 9+3, "responder %d had to resend msg2", fragID)
+			}
+			if tc.pqOnly {
+				assert.Empty(t, classical, "require never offers a classical attempt")
+			}
+
+			assertHostInfoPair(t, myUdp, theirUdp, myVpn, theirVpn, myControl, theirControl)
+			assertTunnel(t, myVpn[0].Addr(), theirVpn[0].Addr(), myControl, theirControl, r)
+			assert.Nil(t, myControl.GetHostInfoByVpnAddr(otherVpn[0].Addr(), false), "no tunnel to the other host")
+			assert.Empty(t, myControl.ListHostmapIndexes(true), "no pending index may outlive the handshake")
+			assert.Equal(t, cert.Version3, myControl.GetHostInfoByVpnAddr(theirVpn[0].Addr(), false).Cert.Version(), "the hybrid handshake won")
+
+			myControl.Stop()
+			theirControl.Stop()
+			otherControl.Stop()
+		})
+	}
+}

handshake/fec/fec.go

@@ -98,6 +98,21 @@ func (e *Encoder) Symbol(dst []byte, index int) ([]byte, error) {
 // 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.
+//
+// TODO(in-place decode): for receivers that already hold state for the message,
+// the initiator's msg2 slots and a pending responder's msg3, store into one
+// buffer of K*SymbolSize (plus room for the Nebula header) allocated from the
+// layout: data symbol i at offset i*SymbolSize, repair symbols parked in the
+// slots of data symbols not seen yet (moved to another empty slot if that data
+// symbol arrives; one always exists while fewer than K are held). Decode then
+// 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
+// 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.
 type Decoder struct {
 	layout Layout
 	have   [MaxSymbols]bool

handshake_manager_pq.go

@@ -28,13 +28,19 @@ const (
 	pqAmplification = 3
 )
 
+// pqMaxMsgLen bounds the fragmented handshake messages a node reassembles:
+// msg1 in the responder's table and msg2 in an initiator's slots. The largest
+// is msg2 at about 9.6KB, which leaves room for certificates carrying many
+// names and routes.
+const pqMaxMsgLen = 16384
+
 // pqTableConfig bounds the responder's reassembly of msg1 fragments from senders
 // that have not authenticated yet.
 var pqTableConfig = fragment.TableConfig{
 	MaxBytes:     4 << 20,
 	MaxEntries:   1024,
 	MaxPerSource: 4,
-	MaxMsgLen:    16384,
+	MaxMsgLen:    pqMaxMsgLen,
 	TTL:          3 * time.Second,
 }
 
@@ -84,6 +90,24 @@ func (hm *HandshakeManager) reloadPQ(c *config.C) {
 	)
 }
 
+// pqMsg2Slots is how many responders' msg2s one attempt reassembles at once.
+// Two hosts answer the same msg1 when a stale lighthouse address for the peer
+// belongs to another node now. Both answers name this attempt in RemoteIndex
+// and differ only in FragID, the responder's own index. A single assembly that
+// emptied on a FragID change let the two interleaved bursts wipe each other
+// out, and since hearing any msg2 fragment stops the classical turns, the
+// retries only repeated it (up to three extra attempts in e2e). One slot per
+// FragID lets whichever answer completes first be processed; a wrong host's
+// answer then takes the existing block-and-retry path.
+const pqMsg2Slots = 2
+
+// msg2Slot reassembles one responder's msg2.
+type msg2Slot struct {
+	fragID uint32
+	msgLen uint16
+	a      *fragment.Assembly
+}
+
 // pqAttempt is the initiator's post-quantum attempt for a pending handshake. It
 // lives alongside the classical attempt in HandshakeHostInfo and has its own
 // local index, so whichever answer arrives first wins.
@@ -92,9 +116,62 @@ type pqAttempt struct {
 	subtype        header.MessageSubType
 	index          uint32
 	sender         *fragment.Sender
-	msg2           *fragment.Assembly
-	cookieResendAt time.Time // last resend prompted by a cookie reply
-	heard          bool      // a msg2 fragment or cookie reply arrived, so the peer speaks PQ
+	msg2           [pqMsg2Slots]msg2Slot // an array, not a map: one slot is used nearly always
+	cookieResendAt time.Time             // last resend prompted by a cookie reply
+	heard          bool                  // a msg2 fragment or cookie reply arrived, so the peer speaks PQ
+}
+
+// slotFor returns the assembly collecting msg2 fragments with fh's FragID,
+// starting one if there is none. A responder never changes its FragID or
+// MsgLen, so a fragment whose MsgLen disagrees with its slot is a stray or a
+// forgery and starts that slot over. When every slot is taken the one holding
+// the least is replaced: a third responder is vanishingly rare, and a forged
+// FragID could already spoil an attempt. It returns nil if no assembly can be
+// made for fh.
+func (pq *pqAttempt) slotFor(fh fragment.Header) *fragment.Assembly {
+	victim := -1
+	for i := range pq.msg2 {
+		s := &pq.msg2[i]
+		if s.a != nil && s.fragID == fh.FragID {
+			if s.msgLen == fh.MsgLen {
+				return s.a
+			}
+			victim = i
+			break
+		}
+	}
+	if victim < 0 {
+		for i := range pq.msg2 {
+			if pq.msg2[i].a == nil {
+				victim = i
+				break
+			}
+		}
+	}
+	if victim < 0 {
+		victim = 0
+		for i := 1; i < len(pq.msg2); i++ {
+			if pq.msg2[i].a.Stored() < pq.msg2[victim].a.Stored() {
+				victim = i
+			}
+		}
+	}
+
+	a, err := fragment.NewAssembly(fh)
+	if err != nil {
+		return nil
+	}
+	pq.msg2[victim] = msg2Slot{fragID: fh.FragID, msgLen: fh.MsgLen, a: a}
+	return a
+}
+
+// dropSlot releases the slot holding a.
+func (pq *pqAttempt) dropSlot(a *fragment.Assembly) {
+	for i := range pq.msg2 {
+		if pq.msg2[i].a == a {
+			pq.msg2[i] = msg2Slot{}
+		}
+	}
 }
 
 // pqFragKey identifies one msg1 being reassembled.
@@ -392,8 +469,13 @@ func (hm *HandshakeManager) handleFragment(via ViaSender, packet []byte, h *head
 	}
 }
 
-// continuePQ collects msg2 fragments for the initiator's PQ attempt.
+// continuePQ collects msg2 fragments for the initiator's PQ attempt, one slot
+// per responder FragID, and processes the first msg2 to complete.
 func (hm *HandshakeManager) continuePQ(via ViaSender, hh *HandshakeHostInfo, h *header.H, fh fragment.Header, sym []byte) {
+	if int(fh.MsgLen) > pqMaxMsgLen {
+		return
+	}
+
 	hh.Lock()
 	defer hh.Unlock()
 
@@ -403,26 +485,17 @@ func (hm *HandshakeManager) continuePQ(via ViaSender, hh *HandshakeHostInfo, h *
 	}
 	pq.heard = true
 
-	if pq.msg2 == nil {
-		a, err := fragment.NewAssembly(fh)
-		if err != nil {
-			return
-		}
-		pq.msg2 = a
-	}
-	ready, err := pq.msg2.Add(fh, sym)
-	if errors.Is(err, fragment.ErrMismatch) {
-		// The responder's FragID and MsgLen never change, so this is a stray or a
-		// forgery. Start over from this fragment rather than wedge.
-		pq.msg2 = nil
+	a := pq.slotFor(fh)
+	if a == nil {
 		return
 	}
+	ready, err := a.Add(fh, sym)
 	if err != nil || !ready {
 		return
 	}
 
-	msg, err := pq.msg2.Decode()
-	pq.msg2 = nil
+	msg, err := a.Decode()
+	pq.dropSlot(a)
 	if err != nil {
 		hm.l.Debug("failed to decode handshake fragments", "from", via, "error", err)
 		return

handshake_manager_pq_test.go

@@ -397,3 +397,134 @@ func TestUsePQ(t *testing.T) {
 	hm.notePQPeer(header.HandshakeIXPSK2MLKEM1024, []netip.Addr{peer})
 	assert.Equal(t, []bool{true, true, true}, attempts(false, false, false), "then only PQ")
 }
+
+// msg2Fragments returns one attempt's packets carrying msg as a responder's msg2
+// under fragID, data symbols first and then repair.
+func msg2Fragments(t *testing.T, msg []byte, fragID uint32) [][]byte {
+	t.Helper()
+	s, err := fragment.NewSender(msg, header.HandshakeIXHybridMLKEM1024, 7, 2, fragID)
+	require.NoError(t, err)
+	now, later, err := s.Attempt(fragment.DefaultRepair(s.Layout().K))
+	require.NoError(t, err)
+	return append(now, later...)
+}
+
+// feedMsg2 puts one packet through an attempt's slots.
+func feedMsg2(t *testing.T, pq *pqAttempt, p []byte) (*fragment.Assembly, fragment.Header, bool) {
+	t.Helper()
+	fh, sym, err := fragment.Parse(p[header.Len:])
+	require.NoError(t, err)
+	a := pq.slotFor(fh)
+	require.NotNil(t, a)
+	ready, err := a.Add(fh, sym)
+	require.NoError(t, err)
+	return a, fh, ready
+}
+
+func slotsUsed(pq *pqAttempt) (n int, ids []uint32) {
+	for _, s := range pq.msg2 {
+		if s.a != nil {
+			n++
+			ids = append(ids, s.fragID)
+		}
+	}
+	return n, ids
+}
+
+// TestPQMsg2Slots interleaves two responders' msg2 bursts, the shape two hosts
+// answering the same msg1 produce on a LAN, and checks each fills its own slot.
+func TestPQMsg2Slots(t *testing.T) {
+	msgB := make([]byte, 9625) // the hybrid msg2: k = 9
+	msgC := make([]byte, 9625)
+	for i := range msgB {
+		msgB[i], msgC[i] = 'b', 'c'
+	}
+	pktsB := msg2Fragments(t, msgB, 1001)
+	pktsC := msg2Fragments(t, msgC, 2002)
+	require.Len(t, pktsB, 12)
+
+	// B0 B1 C0 B2 B3 B4 C1 C2 B5 B6 B7 B8 C3: no run of one FragID reaches nine,
+	// so a single slot emptied on every change would never complete.
+	order := []struct {
+		fromC bool
+		i     int
+	}{{false, 0}, {false, 1}, {true, 0}, {false, 2}, {false, 3}, {false, 4}, {true, 1}, {true, 2}, {false, 5}, {false, 6}, {false, 7}, {false, 8}, {true, 3}}
+
+	pq := &pqAttempt{}
+	var got []byte
+	for n, o := range order {
+		p := pktsB[o.i]
+		if o.fromC {
+			p = pktsC[o.i]
+		}
+		a, _, ready := feedMsg2(t, pq, p)
+		if !ready {
+			continue
+		}
+		require.Equal(t, 11, n, "B completes on its ninth symbol")
+		msg, err := a.Decode()
+		require.NoError(t, err)
+		got = msg
+		pq.dropSlot(a)
+	}
+	require.Equal(t, msgB, got)
+
+	n, ids := slotsUsed(pq)
+	assert.Equal(t, 1, n)
+	assert.Equal(t, []uint32{2002}, ids, "B's slot was released, C's kept")
+
+	// C's four symbols survived B's completion, so five more finish it.
+	var c *fragment.Assembly
+	var ready bool
+	for i := 4; i < 9; i++ {
+		c, _, ready = feedMsg2(t, pq, pktsC[i])
+		assert.Equal(t, i == 8, ready, "symbol %d", i)
+	}
+	msg, err := c.Decode()
+	require.NoError(t, err)
+	assert.Equal(t, msgC, msg)
+
+	t.Run("a third responder takes the slot holding the least", func(t *testing.T) {
+		pq := &pqAttempt{}
+		feedMsg2(t, pq, pktsB[0])
+		feedMsg2(t, pq, pktsB[1])
+		feedMsg2(t, pq, pktsC[0])
+		pktsD := msg2Fragments(t, msgC, 3003)
+		d, fh, _ := feedMsg2(t, pq, pktsD[0])
+		_, ids := slotsUsed(pq)
+		assert.ElementsMatch(t, []uint32{1001, 3003}, ids)
+
+		l, err := fh.Layout()
+		require.NoError(t, err)
+		b, _, _ := feedMsg2(t, pq, pktsB[2])
+		assert.Equal(t, 3*l.SymbolSize, b.Stored(), "B kept its symbols")
+		assert.Equal(t, l.SymbolSize, d.Stored())
+	})
+
+	t.Run("a different MsgLen under the same FragID starts the slot over", func(t *testing.T) {
+		pq := &pqAttempt{}
+		feedMsg2(t, pq, pktsB[0])
+		feedMsg2(t, pq, pktsB[1])
+		shorter := msg2Fragments(t, msgB[:9000], 1001)
+		a, fh, _ := feedMsg2(t, pq, shorter[0])
+		l, err := fh.Layout()
+		require.NoError(t, err)
+		assert.Equal(t, l.SymbolSize, a.Stored())
+		n, ids := slotsUsed(pq)
+		assert.Equal(t, 1, n)
+		assert.Equal(t, []uint32{1001}, ids)
+	})
+
+	t.Run("an oversize claim is dropped before it takes a slot", func(t *testing.T) {
+		hm, _ := newPQTestManager(t)
+		hh := &HandshakeHostInfo{hostinfo: &HostInfo{}, pq: &pqAttempt{index: 42}}
+		hm.indexes[42] = hh
+		fh := fragment.Header{FragID: 1001, MsgLen: pqMaxMsgLen + 1}
+		l, err := fh.Layout()
+		require.NoError(t, err)
+		hm.continuePQ(ViaSender{}, hh, &header.H{RemoteIndex: 42, MessageCounter: 2}, fh, make([]byte, l.SymbolSize))
+		n, _ := slotsUsed(hh.pq)
+		assert.Zero(t, n)
+		assert.False(t, hh.pq.heard)
+	})
+}