Demand cookies for post-quantum msg1s while under load, and reload handshakes.pq settings

docs/pq-handshake.md

@@ -66,12 +66,30 @@ signature, following PQNoise's pqIX, while keeping the classical DH tokens:
 - **msg3** (I to R): ciphertext to R's static KEM key, plus first data. R authenticates I
   when msg3 decrypts; I authenticates R implicitly when R's first reply decrypts.
 
+msg3 is about 1.6 KB, 2 data and 2 repair fragments, and adds no new message types:
+
+- **Responder.** After sending msg2 the responder keeps a pending Machine, found by its
+  local index like the initiator's today; this is the multi-message responder support the
+  manager lacks now. A repeat msg1 maps to it through the answered-msg1 cache and gets msg2
+  resent with fresh symbols, exactly as in phase 1. Until msg3 verifies, data for that
+  index is dropped silently: it must not draw a RecvError, which would tear down the
+  initiator's new tunnel. On completing it sends an empty `Test`/`TestReply` under the
+  new keys, so the initiator hears back even when there is no traffic.
+- **Initiator.** Writing msg3 yields the transport keys, so the hostinfo moves to the main
+  hostmap and data flows at once, readable only by the real responder. The handshake entry
+  stays behind in a confirming state: each retry tick resends msg3 with fresh symbols until
+  any packet from the responder decrypts on the new tunnel, and if that never happens
+  within the handshake timeout the tunnel is torn down. Losing msg3 costs the data sent
+  behind it, which the overlay's own transport retransmits.
+- msg3's fragments use the responder's index as RemoteIndex and land on the responder's
+  pending entry, so they never enter the pre-authentication table.
+
 The initiator never caches the responder's certificate, so every handshake has three
 messages and carries both full certificates. With ML-DSA-87:
 
 | Message | Size | Data fragments (k) |
 |---|---|---|
-| msg1 | ~8.2 KB | 7 |
+| msg1 | ~8.2 KB | 8 |
 | msg2 | ~9.8 KB | 9 |
 | msg3 | ~1.6 KB | 2 |
 
@@ -98,17 +116,29 @@ Every packet of a fragmented handshake subtype is:
 ```
 Nebula header (16)   Version | Type=Handshake | Subtype | Reserved=0 | RemoteIndex | Counter
 Fragment header (8)  FragID u32 | MsgLen u16 | Index u8 | Flags u8
+Cookie (16)          only when Flags has 0x01
 Symbol               SymbolSize bytes
 ```
 
 - `Counter` is the Noise message number (1, 2, 3). `RemoteIndex` is 0 for msg1, as today.
 - `FragID` is the sender's local index (random, non-zero).
 - `MsgLen` is the length of the whole message. Both ends derive the layout from it:
-  `K = ceil(MsgLen / 1176)`, `SymbolSize = ceil(MsgLen / K)`. Every symbol has the same
-  size, so a packet is at most 16 + 8 + 1176 = 1200 bytes.
+  `K = ceil(MsgLen / 1160)`, `SymbolSize = ceil(MsgLen / K)`. Every symbol has the same
+  size. Room for a cookie is always reserved, so a packet is at most
+  16 + 8 + 16 + 1160 = 1200 bytes with or without one.
 - `Index < K` is a data symbol (a slice of the message, the last one zero padded);
   `Index >= K` is a repair symbol.
-- `Flags` is reserved for a cookie extension and must be zero.
+- `Flags` bit 0x01 means a cookie follows the fragment header. Other bits must be zero.
+
+A cookie reply, sent by a responder under load, is a bare 32 byte packet:
+
+```
+Nebula header (16)   Version | Type=Handshake | Subtype | Reserved=0 | RemoteIndex=FragID | Counter=0
+Cookie (16)
+```
+
+It names the initiator's FragID, so the initiator finds its attempt the same way it finds
+msg2.
 
 ### Forward error correction
 
@@ -151,7 +181,7 @@ Implemented in `handshake/fragment` (`Sender`, `Assembly`, and the capped `Table
   Nothing is allocated from a claimed size alone (FragmentSmack, game-engine "fragment
   bombs").
 - `MsgLen` must be identical for every fragment of a `FragID`; a mismatch drops the whole
-  entry. msg1 is capped at `MsgLen <= 16384` (K <= 14).
+  entry. msg1 is capped at `MsgLen <= 16384` (K <= 15).
 - The decoder keeps at most K distinct symbols, so an entry never exceeds about MsgLen bytes
   however many repair symbols arrive.
 - When K symbols are present the message is decoded and handed to the handshake `Machine`
@@ -183,12 +213,42 @@ Implemented in `handshake/fragment` (`Sender`, `Assembly`, and the capped `Table
 ### Denial of service
 
 The responder can no longer complete statelessly on msg1, since msg1 arrives in pieces.
-Defenses are the capped table above plus a **cookie under load**: when the table is under
-pressure, a fragment from an unvalidated address is answered with a stateless cookie
-(a MAC over source address, FragID and a time bucket), and the initiator resends with the
-cookie. This costs a round trip only while under attack, in the style of WireGuard's
-mac2, QUIC Retry and the SCTP state cookie. A CA-keyed pre-filter MAC was considered and
-rejected: networks commonly trust many CAs.
+Defenses are the capped table above plus a **cookie under load**, in the style of
+WireGuard's mac2, QUIC Retry and the SCTP state cookie:
+
+- Once the msg1 table is at least half full (by bytes or by entries), every msg1 fragment
+  must carry a valid cookie. One without is answered with a cookie reply and leaves no
+  state. Fragments that already belong to an entry need one too, so an attacker cannot
+  hold old entries open.
+- A cookie is `HMAC-SHA256(secret, source address and port, relay channel, FragID)`
+  truncated to 16 bytes. The secret rotates every 2 minutes and the previous one is still
+  accepted, so a cookie is good for 2 to 4 minutes.
+- Replies are never larger than the packet that prompted them (no amplification) and are
+  rate limited as a whole (a burst of 256, then 1000 per second), so spoofed floods cannot
+  turn the responder into a reflector.
+- The initiator attaches the cookie to its PQ attempt, so every later packet carries it,
+  and immediately resends fresh symbols to the address the reply came from, at most once
+  per retry interval. Under load a PQ handshake costs one extra round trip, which in
+  practice still beats the classical attempt on the next retry tick.
+- An attacker on the path can forge a cookie reply; the bogus cookie is simply rejected
+  and answered with a real one.
+
+A CA-keyed pre-filter MAC was considered and rejected: networks commonly trust many CAs.
+
+#### Injected symbols
+
+The decoder keeps the first K distinct symbols it sees, so one forged symbol in those K
+spoils that attempt's decode; Noise rejects the result and the next attempt brings fresh
+symbols. Forging one needs the attempt's FragID or index, which are random 32 bit values,
+and for msg1 the initiator's exact address and port too, so in practice only an attacker
+who sees the traffic can do it. An attacker who can also drop packets gains nothing new.
+One who can only inject, and races a forged symbol into every attempt, can stop a PQ
+handshake where it could not stop a classical one, since classical msg2s are tried one
+packet at a time. Keeping spare symbols to find a consistent subset only raises the number
+of forgeries needed per attempt, and authenticating symbols is impossible for msg1, which
+travels before any key exists. So the decoder deliberately does neither. A non-sticky
+initiator still gets a classical tunnel from its next attempt; a sticky one, by design,
+does not downgrade.
 
 ## Rollout
 
@@ -198,6 +258,8 @@ rejected: networks commonly trust many CAs.
 - Old nodes silently drop unknown subtypes, so discovery is speculative: an initiator with
   `initiate` enabled alternates PQ and classical attempts on successive retries, keeping
   both pending, and the first to complete wins.
+- All `handshakes.pq` settings reload without a restart. Handshakes in flight carry on;
+  the new settings apply to the next attempt and the next packet.
 - `handshakes.pq.sticky` (default off): once a peer has completed a PQ handshake, never fall
   back to classical for it again during this process lifetime. This resists active
   downgrade after first contact.
@@ -257,18 +319,18 @@ in every scenario above, and within three at least 99.99%.
     amplification budget.
   - **Relays** carry fragments unchanged; repair fragments trail via a timer.
 - The pre-authentication table is capped at 4 MiB, 1024 entries, 4 per source (IPv4 /32,
-  IPv6 /64, or relay channel), 16 KiB messages and a 3 s idle timeout.
+  IPv6 /64, or relay channel), 16 KiB messages and a 3 s idle timeout. Cookies are
+  demanded from half full.
 - Tests: Noise-level KEM tests, codec and fragment tests and fuzzers, manager unit tests
   for the response budget and rate limit, and e2e tests for a clean handshake, loss of
   data fragments in both directions, fallback to a peer that does not accept PQ, a lost
   msg2 recovered by fresh symbols, sticky mode under an attacker dropping PQ packets,
-  relays, and index cleanup on completion, fallback and timeout. ML-KEM and HKDF also
+  relays, index cleanup on completion, fallback and timeout, and a responder flooded from
+  spoofed sources demanding and accepting a cookie. Unit tests cover cookie binding and
+  rotation, the never-amplify rule, the initiator's rate limited resend, and reload. ML-KEM and HKDF also
   pass under `GOEXPERIMENT=boringcrypto` and `GODEBUG=fips140=on`.
 
 ## Open items
 
-1. Cookie wire format and the load threshold that triggers it.
-2. Whether the decoder keeps spare symbols to identify injected symbols.
-3. msg3 confirmation message format (phase 2).
-4. Configuration reload: the `handshakes.pq` settings are read at startup only, like the
-   rest of the `handshakes` section.
+None for phase 1. Phase 2 needs certificate v3, the KEM-aware Noise symmetric state, and
+the msg3 design above.

e2e/pq_handshakes_test.go

@@ -349,3 +349,55 @@ func TestPQTimeoutReleasesIndexes(t *testing.T) {
 	myControl.Stop()
 	theirControl.Stop()
 }
+
+// TestPQCookieUnderLoad floods the responder with msg1 fragments from spoofed
+// sources until it demands cookies, then checks a real initiator picks up the
+// cookie and still completes a PQ handshake.
+func TestPQCookieUnderLoad(t *testing.T) {
+	t.Parallel()
+	myControl, theirControl, myVpn, theirVpn, _, theirUdp, r := newPQPair(t, pqConfig(true, true, false), pqConfig(false, true, false))
+	defer r.RenderFlow()
+
+	// Half of the 1024 entry table is the cookie threshold. Stopping exactly there
+	// keeps them from answering the spoofed sources, which the router cannot reach.
+	t.Log("Flood them with the first fragment of msg1s nobody will finish")
+	for i := range 512 {
+		from := netip.AddrPortFrom(netip.AddrFrom4([4]byte{198, 18, byte(i >> 8), byte(i)}), 4242)
+		s, err := fragment.NewSender(make([]byte, 2000), header.HandshakeIXPSK2MLKEM1024, 0, 1, uint32(i+1))
+		require.NoError(t, err)
+		pkts, _, err := s.Attempt(0)
+		require.NoError(t, err)
+		theirControl.InjectUDPPacket(&udp.Packet{From: from, To: theirUdp, Data: pkts[0]})
+	}
+
+	var tally handshakeTally
+	var cookieReplies, msg2, withCookie int
+	myControl.InjectTunPacket(BuildTunUDPPacket(theirVpn[0].Addr(), 80, myVpn[0].Addr(), 80, []byte("Hi from me")))
+	require.True(t, r.RouteForAllExitFuncOrTimeout(5*time.Second, func(p *udp.Packet, receiver *nebula.Control) router.ExitType {
+		tally.observe(p)
+		if h, fh := handshakeOf(p); h != nil && h.Subtype == header.HandshakeIXPSK2MLKEM1024 {
+			switch {
+			case h.MessageCounter == 0:
+				cookieReplies++
+			case h.MessageCounter == 2:
+				msg2++
+			case fh != nil && fh.HasCookie():
+				withCookie++
+			}
+		}
+		if isTunnelData(p, receiver, theirControl) {
+			return router.RouteAndExit
+		}
+		return router.KeepRouting
+	}), "tunnel never came up")
+
+	assertUdpPacket(t, []byte("Hi from me"), theirControl.GetFromTun(true), myVpn[0].Addr(), theirVpn[0].Addr(), 80, 80)
+	assert.Positive(t, cookieReplies, "a loaded responder asks for a cookie")
+	assert.Positive(t, withCookie, "the initiator resends with it")
+	assert.Positive(t, msg2, "and the handshake completes post-quantum")
+	assert.Zero(t, tally.classical, "the cookie round trip beats the classical fallback")
+	assertTunnel(t, myVpn[0].Addr(), theirVpn[0].Addr(), myControl, theirControl, r)
+
+	myControl.Stop()
+	theirControl.Stop()
+}

handshake_manager.go

@@ -10,6 +10,7 @@ import (
 	"net/netip"
 	"slices"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/rcrowley/go-metrics"
@@ -17,6 +18,7 @@ import (
 	"github.com/slackhq/nebula/handshake"
 	"github.com/slackhq/nebula/handshake/fragment"
 	"github.com/slackhq/nebula/header"
+	"github.com/slackhq/nebula/ratelimit"
 	"github.com/slackhq/nebula/udp"
 )
 
@@ -77,12 +79,17 @@ type HandshakeManager struct {
 	// can be used to trigger outbound handshake for the given vpnIp
 	trigger chan netip.Addr
 
+	// pq holds the post-quantum settings, which change on reload.
+	pq atomic.Pointer[pqHandshakeConfig]
 	// pqTable reassembles post-quantum msg1s from senders that have not authenticated yet.
 	pqTable *fragment.Table[pqFragKey, pqFragSource]
 	// pqPeers holds vpn addresses that completed a post-quantum handshake, for sticky mode.
 	pqPeers sync.Map
 	// pqAnswered recognises repeats of post-quantum msg1s already answered.
 	pqAnswered *pqAnswered
+	// pqCookies and pqCookieLimiter serve cookies while pqTable is under load.
+	pqCookies       *pqCookies
+	pqCookieLimiter *ratelimit.Limiter
 }
 
 type HandshakeHostInfo struct {
@@ -127,7 +134,7 @@ func (hh *HandshakeHostInfo) cachePacket(l *slog.Logger, t header.MessageType, s
 }
 
 func NewHandshakeManager(l *slog.Logger, mainHostMap *HostMap, lightHouse *LightHouse, outside udp.Conn, config HandshakeConfig) *HandshakeManager {
-	return &HandshakeManager{
+	hm := &HandshakeManager{
 		vpnIps:                 map[netip.Addr]*HandshakeHostInfo{},
 		indexes:                map[uint32]*HandshakeHostInfo{},
 		mainHostMap:            mainHostMap,
@@ -143,6 +150,13 @@ func NewHandshakeManager(l *slog.Logger, mainHostMap *HostMap, lightHouse *Light
 		pqTable:                fragment.NewTable[pqFragKey, pqFragSource](pqTableConfig),
 		pqAnswered:             newPQAnswered(hsTimeout(config.retries, config.tryInterval)),
 	}
+	pq := config.pq
+	hm.pq.Store(&pq)
+	hm.pqCookies = newPQCookies(time.Now())
+	// Cookie replies are tiny, but a flood of spoofed fragments must not turn
+	// into a flood of replies.
+	hm.pqCookieLimiter, _ = ratelimit.New(256, time.Millisecond)
+	return hm
 }
 
 func (hm *HandshakeManager) Run(ctx context.Context) {

handshake_manager_pq.go

@@ -2,13 +2,17 @@ package nebula
 
 import (
 	"context"
+	"crypto/hmac"
+	"crypto/rand"
 	"crypto/sha256"
+	"encoding/binary"
 	"errors"
 	"log/slog"
 	"net/netip"
 	"sync"
 	"time"
 
+	"github.com/slackhq/nebula/config"
 	"github.com/slackhq/nebula/handshake"
 	"github.com/slackhq/nebula/handshake/fragment"
 	"github.com/slackhq/nebula/header"
@@ -43,14 +47,46 @@ type pqHandshakeConfig struct {
 	repairDelay time.Duration // delay between data and repair fragments
 }
 
+func pqHandshakeConfigFrom(c *config.C) pqHandshakeConfig {
+	return pqHandshakeConfig{
+		initiate:    c.GetBool("handshakes.pq.initiate", false),
+		accept:      c.GetBool("handshakes.pq.accept", false),
+		sticky:      c.GetBool("handshakes.pq.sticky", false),
+		repairDelay: c.GetDuration("handshakes.pq.repair_delay", DefaultPQRepairDelay),
+	}
+}
+
+// pqConfig returns the post-quantum settings in effect.
+func (hm *HandshakeManager) pqConfig() *pqHandshakeConfig {
+	return hm.pq.Load()
+}
+
+// reloadPQ applies changed post-quantum settings. Handshakes already in flight
+// keep going; the new settings apply to the next attempt and packet.
+func (hm *HandshakeManager) reloadPQ(c *config.C) {
+	if !c.HasChanged("handshakes.pq") {
+		return
+	}
+	old := hm.pqConfig()
+	cfg := pqHandshakeConfigFrom(c)
+	hm.pq.Store(&cfg)
+	hm.l.Info("Post-quantum handshake settings changed",
+		"initiate", cfg.initiate, "oldInitiate", old.initiate,
+		"accept", cfg.accept, "oldAccept", old.accept,
+		"sticky", cfg.sticky, "oldSticky", old.sticky,
+		"repairDelay", cfg.repairDelay, "oldRepairDelay", old.repairDelay,
+	)
+}
+
 // 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.
 type pqAttempt struct {
-	machine *handshake.Machine
-	index   uint32
-	sender  *fragment.Sender
-	msg2    *fragment.Assembly
+	machine        *handshake.Machine
+	index          uint32
+	sender         *fragment.Sender
+	msg2           *fragment.Assembly
+	cookieResendAt time.Time // last resend prompted by a cookie reply
 }
 
 // pqFragKey identifies one msg1 being reassembled.
@@ -94,10 +130,11 @@ func reassembled(h *header.H, msg []byte) []byte {
 // drops it. A peer that already completed a PQ handshake never gets a classical
 // attempt when sticky is on.
 func (hm *HandshakeManager) usePQ(hh *HandshakeHostInfo, vpnIp netip.Addr) bool {
-	if !hm.config.pq.initiate {
+	cfg := hm.pqConfig()
+	if !cfg.initiate {
 		return false
 	}
-	if hm.config.pq.sticky {
+	if cfg.sticky {
 		if _, ok := hm.pqPeers.Load(vpnIp); ok {
 			return true
 		}
@@ -207,6 +244,11 @@ func (hh *HandshakeHostInfo) pqIndex() uint32 {
 
 // handleFragment routes one fragment of a PQ handshake message.
 func (hm *HandshakeManager) handleFragment(via ViaSender, packet []byte, h *header.H) {
+	if h.MessageCounter == pqCookieCounter {
+		hm.handleCookie(via, packet, h)
+		return
+	}
+
 	fh, sym, err := fragment.Parse(packet[header.Len:])
 	if err != nil {
 		hm.l.Debug("dropping malformed handshake fragment", "from", via, "error", err)
@@ -215,7 +257,7 @@ func (hm *HandshakeManager) handleFragment(via ViaSender, packet []byte, h *head
 
 	switch h.MessageCounter {
 	case 1:
-		if !hm.config.pq.accept {
+		if !hm.pqConfig().accept {
 			hm.l.Debug("dropping post-quantum handshake, not accepting them", "from", via)
 			return
 		}
@@ -224,8 +266,13 @@ func (hm *HandshakeManager) handleFragment(via ViaSender, packet []byte, h *head
 				"from", via, "remoteIndex", h.RemoteIndex)
 			return
 		}
+		now := time.Now()
 		key, src := pqFragKeyFor(via, fh.FragID)
-		res, err := hm.pqTable.Add(key, src, time.Now(), fh, sym)
+		if hm.pqTable.Load() >= pqCookieLoad && !(fh.HasCookie() && hm.pqCookies.valid(key, fh.Cookie, now)) {
+			hm.sendCookie(via, fh, key, len(packet), now)
+			return
+		}
+		res, err := hm.pqTable.Add(key, src, now, fh, sym)
 		if err != nil {
 			hm.l.Debug("dropping handshake fragment", "from", via, "error", err)
 			return
@@ -303,12 +350,13 @@ func (hm *HandshakeManager) sendFragmented(now, later [][]byte, send func([]byte
 			}
 		}
 	}
+	delay := hm.pqConfig().repairDelay
 	switch {
 	case len(later) == 0:
-	case hm.config.pq.repairDelay <= 0:
+	case delay <= 0:
 		sendLater()
 	default:
-		time.AfterFunc(hm.config.pq.repairDelay, sendLater)
+		time.AfterFunc(delay, sendLater)
 	}
 	return first
 }
@@ -356,30 +404,22 @@ func (hm *HandshakeManager) sendFragmentedResponse(via ViaSender, msg []byte, ho
 		return err
 	}
 
-	f := hm.f
-	var send func([]byte) error
-	if !via.IsRelayed {
-		send = func(p []byte) error {
-			return f.outside.WriteTo(p, via.UdpAddr)
-		}
-	} else {
+	if via.IsRelayed {
 		if via.relay == nil {
 			return errors.New("both addr and via.relay are nil")
 		}
+		// We received a valid handshake on this relay, so make sure the relay
+		// state reflects that, in case it had been marked Disestablished.
 		hostinfo.relayState.InsertRelayTo(via.relayHI.vpnAddrs[0])
 		via.relayHI.relayState.UpdateRelayForByIdxState(via.relay.LocalIndex, Established)
-		send = func(p []byte) error {
-			f.SendVia(via.relayHI, via.relay, p, make([]byte, 12), make([]byte, mtu), false, 0)
-			return nil
-		}
 	}
-	f.messageMetrics.Tx(header.Handshake, h.Subtype, int64(len(pkts)+len(later)))
-	return hm.sendFragmented(pkts, later, send)
+	hm.f.messageMetrics.Tx(header.Handshake, h.Subtype, int64(len(pkts)+len(later)))
+	return hm.sendFragmented(pkts, later, hm.sendTo(via))
 }
 
 // notePQPeer remembers peers that completed a PQ handshake, for sticky mode.
 func (hm *HandshakeManager) notePQPeer(subtype header.MessageSubType, vpnAddrs []netip.Addr) {
-	if !hm.config.pq.sticky || subtype != pqSubtype {
+	if !hm.pqConfig().sticky || subtype != pqSubtype {
 		return
 	}
 	for _, a := range vpnAddrs {
@@ -444,3 +484,157 @@ func (a *pqAnswered) expire(now time.Time) {
 		}
 	}
 }
+
+// A responder whose msg1 table is under load stops keeping state for senders
+// that have not shown they receive at their address. It answers a fragment
+// without a valid cookie with a cookie reply and keeps nothing; the initiator
+// resends with the cookie attached. This is WireGuard's under-load cookie, QUIC's
+// Retry and SCTP's state cookie: one extra round trip, only while under attack.
+const (
+	// pqCookieLoad is the table load above which msg1 fragments need a cookie.
+	pqCookieLoad = 0.5
+
+	// pqCookieCounter marks a cookie reply in the header's message counter.
+	// Handshake messages count from 1.
+	pqCookieCounter = 0
+
+	// pqCookieRotate is how often the cookie secret changes. A cookie stays
+	// valid for between one and two rotations.
+	pqCookieRotate = 2 * time.Minute
+
+	pqCookieReplyLen = header.Len + fragment.CookieLen
+)
+
+// pqCookies mints and checks cookies: an HMAC over the sender's address (or relay
+// channel) and FragID under a rotating secret, with the previous secret still
+// accepted.
+type pqCookies struct {
+	sync.Mutex
+	secrets   [2][32]byte
+	rotatedAt time.Time
+}
+
+func newPQCookies(now time.Time) *pqCookies {
+	c := &pqCookies{rotatedAt: now}
+	_, _ = rand.Read(c.secrets[0][:])
+	_, _ = rand.Read(c.secrets[1][:])
+	return c
+}
+
+func (c *pqCookies) rotate(now time.Time) {
+	if now.Sub(c.rotatedAt) < pqCookieRotate {
+		return
+	}
+	c.secrets[1] = c.secrets[0]
+	_, _ = rand.Read(c.secrets[0][:])
+	c.rotatedAt = now
+}
+
+func cookieMAC(secret []byte, key pqFragKey) fragment.Cookie {
+	mac := hmac.New(sha256.New, secret)
+	addr := key.addr.Addr().As16()
+	mac.Write(addr[:])
+	var b [10]byte
+	binary.BigEndian.PutUint16(b[0:], key.addr.Port())
+	binary.BigEndian.PutUint32(b[2:], key.relay)
+	binary.BigEndian.PutUint32(b[6:], key.fragID)
+	mac.Write(b[:])
+	var c fragment.Cookie
+	copy(c[:], mac.Sum(nil))
+	return c
+}
+
+func (c *pqCookies) make(key pqFragKey, now time.Time) fragment.Cookie {
+	c.Lock()
+	defer c.Unlock()
+	c.rotate(now)
+	return cookieMAC(c.secrets[0][:], key)
+}
+
+func (c *pqCookies) valid(key pqFragKey, cookie fragment.Cookie, now time.Time) bool {
+	c.Lock()
+	defer c.Unlock()
+	c.rotate(now)
+	for _, s := range c.secrets {
+		want := cookieMAC(s[:], key)
+		if hmac.Equal(want[:], cookie[:]) {
+			return true
+		}
+	}
+	return false
+}
+
+// sendCookie answers a msg1 fragment that arrived without a valid cookie while
+// under load. The reply is never larger than what triggered it, and replies are
+// rate limited as a whole.
+func (hm *HandshakeManager) sendCookie(via ViaSender, fh fragment.Header, key pqFragKey, received int, now time.Time) {
+	if received < pqCookieReplyLen {
+		return
+	}
+	if _, ok := hm.pqCookieLimiter.Allow(); !ok {
+		return
+	}
+	c := hm.pqCookies.make(key, now)
+	reply := make([]byte, header.Len, pqCookieReplyLen)
+	header.Encode(reply, header.Version, header.Handshake, pqSubtype, fh.FragID, pqCookieCounter)
+	reply = append(reply, c[:]...)
+	hm.f.messageMetrics.Tx(header.Handshake, pqSubtype, 1)
+	if err := hm.sendTo(via)(reply); err != nil {
+		hm.l.Debug("Failed to send handshake cookie", "to", via, "error", err)
+	}
+}
+
+// handleCookie attaches a responder's cookie to the PQ attempt it names and
+// resends right away, at most once per retry interval, rather than waiting for
+// the next attempt.
+func (hm *HandshakeManager) handleCookie(via ViaSender, packet []byte, h *header.H) {
+	if len(packet) != pqCookieReplyLen {
+		return
+	}
+	hh := hm.queryIndex(h.RemoteIndex)
+	if hh == nil {
+		return
+	}
+	hh.Lock()
+	defer hh.Unlock()
+
+	pq := hh.pq
+	if pq == nil || pq.index != h.RemoteIndex || !hm.stillPending(hh, pq.index) {
+		return
+	}
+	var c fragment.Cookie
+	copy(c[:], packet[header.Len:])
+	pq.sender.SetCookie(c)
+
+	now := time.Now()
+	if now.Sub(pq.cookieResendAt) < hm.config.tryInterval {
+		return
+	}
+	pq.cookieResendAt = now
+	pkts, later, err := pq.sender.Attempt(fragment.DefaultRepair(pq.sender.Layout().K))
+	if err != nil {
+		return
+	}
+	hm.messageMetrics.Tx(header.Handshake, pqSubtype, int64(len(pkts)+len(later)))
+	if err := hm.sendFragmented(pkts, later, hm.sendTo(via)); err != nil {
+		hm.l.Debug("Failed to resend handshake with cookie", "to", via, "error", err)
+	}
+}
+
+// sendTo returns a function sending a packet back the way via came: directly, or
+// through the relay it arrived on.
+func (hm *HandshakeManager) sendTo(via ViaSender) func([]byte) error {
+	f := hm.f
+	if !via.IsRelayed {
+		return func(p []byte) error {
+			return f.outside.WriteTo(p, via.UdpAddr)
+		}
+	}
+	return func(p []byte) error {
+		if via.relay == nil {
+			return errors.New("both addr and via.relay are nil")
+		}
+		f.SendVia(via.relayHI, via.relay, p, make([]byte, 12), make([]byte, mtu), false, 0)
+		return nil
+	}
+}

handshake_manager_pq_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/slackhq/nebula/config"
 	"github.com/slackhq/nebula/handshake/fragment"
 	"github.com/slackhq/nebula/header"
 	"github.com/slackhq/nebula/test"
@@ -153,3 +154,197 @@ func TestPQFragKeyFor(t *testing.T) {
 	_, s5 := pqFragKeyFor(ViaSender{UdpAddr: netip.MustParseAddrPort("[::ffff:192.0.2.1]:1")}, 9)
 	assert.Equal(t, s1, s5, "v4-mapped addresses count as IPv4")
 }
+
+func TestPQCookies(t *testing.T) {
+	now := time.Unix(1000, 0)
+	c := newPQCookies(now)
+	key := pqFragKey{addr: netip.MustParseAddrPort("192.0.2.1:4242"), fragID: 9}
+
+	cookie := c.make(key, now)
+	assert.True(t, c.valid(key, cookie, now))
+
+	for _, other := range []pqFragKey{
+		{addr: netip.MustParseAddrPort("192.0.2.2:4242"), fragID: 9},
+		{addr: netip.MustParseAddrPort("192.0.2.1:4243"), fragID: 9},
+		{addr: netip.MustParseAddrPort("192.0.2.1:4242"), fragID: 10},
+		{addr: netip.MustParseAddrPort("192.0.2.1:4242"), relay: 1, fragID: 9},
+	} {
+		assert.False(t, c.valid(other, cookie, now), "%+v", other)
+	}
+
+	now = now.Add(pqCookieRotate)
+	assert.True(t, c.valid(key, cookie, now), "the previous secret is still accepted")
+	assert.NotEqual(t, cookie, c.make(key, now), "new cookies come from the new secret")
+
+	now = now.Add(pqCookieRotate)
+	assert.False(t, c.valid(key, cookie, now), "two rotations retire a cookie")
+}
+
+// pqMsg1Fragment builds the first fragment of a msg1, optionally with a cookie.
+func pqMsg1Fragment(t *testing.T, fragID uint32, cookie *fragment.Cookie) ([]byte, *header.H) {
+	t.Helper()
+	s, err := fragment.NewSender(make([]byte, 2000), pqSubtype, 0, 1, fragID)
+	require.NoError(t, err)
+	if cookie != nil {
+		s.SetCookie(*cookie)
+	}
+	pkts, _, err := s.Attempt(0)
+	require.NoError(t, err)
+	var h header.H
+	require.NoError(t, h.Parse(pkts[0]))
+	return pkts[0], &h
+}
+
+// fillPQTable puts the responder's msg1 table at the cookie threshold.
+func fillPQTable(t *testing.T, hm *HandshakeManager) {
+	t.Helper()
+	now := time.Now()
+	for i := 0; hm.pqTable.Load() < pqCookieLoad; i++ {
+		p, _ := pqMsg1Fragment(t, uint32(i+1), nil)
+		fh, sym, err := fragment.Parse(p[header.Len:])
+		require.NoError(t, err)
+		addr := netip.AddrPortFrom(netip.AddrFrom4([4]byte{198, 18, byte(i >> 8), byte(i)}), 4242)
+		key, src := pqFragKeyFor(ViaSender{UdpAddr: addr}, fh.FragID)
+		_, err = hm.pqTable.Add(key, src, now, fh, sym)
+		require.NoError(t, err)
+	}
+}
+
+func TestPQCookieUnderLoad(t *testing.T) {
+	hm, conn := newPQTestManager(t)
+	via := ViaSender{UdpAddr: netip.MustParseAddrPort("192.0.2.1:4242")}
+
+	t.Run("no cookie needed below the threshold", func(t *testing.T) {
+		p, h := pqMsg1Fragment(t, 5000, nil)
+		before := hm.pqTable.Len()
+		hm.handleFragment(via, p, h)
+		assert.Empty(t, conn.take())
+		assert.Equal(t, before+1, hm.pqTable.Len())
+	})
+
+	fillPQTable(t, hm)
+
+	var cookie fragment.Cookie
+	t.Run("a fragment without a cookie gets one and leaves no state", func(t *testing.T) {
+		p, h := pqMsg1Fragment(t, 7, nil)
+		before := hm.pqTable.Len()
+		hm.handleFragment(via, p, h)
+		assert.Equal(t, before, hm.pqTable.Len())
+
+		sent := conn.take()
+		require.Len(t, sent, 1)
+		reply := sent[0]
+		require.Len(t, reply, pqCookieReplyLen)
+		var rh header.H
+		require.NoError(t, rh.Parse(reply))
+		assert.Equal(t, header.Handshake, rh.Type)
+		assert.Equal(t, pqSubtype, rh.Subtype)
+		assert.Equal(t, uint32(7), rh.RemoteIndex, "the reply names the initiator's FragID")
+		assert.Equal(t, uint64(pqCookieCounter), rh.MessageCounter)
+		copy(cookie[:], reply[header.Len:])
+	})
+
+	t.Run("a forged cookie gets a fresh one", func(t *testing.T) {
+		bad := cookie
+		bad[0] ^= 1
+		p, h := pqMsg1Fragment(t, 7, &bad)
+		before := hm.pqTable.Len()
+		hm.handleFragment(via, p, h)
+		assert.Equal(t, before, hm.pqTable.Len())
+		assert.Len(t, conn.take(), 1)
+	})
+
+	t.Run("the cookie gets the fragment in", func(t *testing.T) {
+		p, h := pqMsg1Fragment(t, 7, &cookie)
+		before := hm.pqTable.Len()
+		hm.handleFragment(via, p, h)
+		assert.Equal(t, before+1, hm.pqTable.Len())
+		assert.Empty(t, conn.take())
+	})
+
+	t.Run("a cookie is bound to its address", func(t *testing.T) {
+		p, h := pqMsg1Fragment(t, 7, &cookie)
+		hm.handleFragment(ViaSender{UdpAddr: netip.MustParseAddrPort("192.0.2.99:4242")}, p, h)
+		assert.Len(t, conn.take(), 1, "another address gets its own cookie instead")
+	})
+
+	t.Run("a reply never outgrows what prompted it", func(t *testing.T) {
+		p, h := pqMsg1Fragment(t, 8, nil)
+		hm.sendCookie(via, fragment.Header{FragID: 8}, pqFragKey{}, pqCookieReplyLen-1, time.Now())
+		assert.Empty(t, conn.take())
+		hm.handleFragment(via, p, h)
+		assert.Len(t, conn.take(), 1)
+	})
+}
+
+func TestPQHandleCookie(t *testing.T) {
+	hm, conn := newPQTestManager(t)
+	via := ViaSender{UdpAddr: netip.MustParseAddrPort("192.0.2.1:4242")}
+
+	s, err := fragment.NewSender(make([]byte, 2000), pqSubtype, 0, 1, 99)
+	require.NoError(t, err)
+	_, _, err = s.Attempt(2)
+	require.NoError(t, err)
+	hh := &HandshakeHostInfo{hostinfo: &HostInfo{}, pq: &pqAttempt{index: 99, sender: s}}
+	hm.indexes[99] = hh
+
+	reply := make([]byte, header.Len, pqCookieReplyLen)
+	header.Encode(reply, header.Version, header.Handshake, pqSubtype, 99, pqCookieCounter)
+	reply = append(reply, bytes16(0xab)...)
+	var h header.H
+	require.NoError(t, h.Parse(reply))
+
+	hm.handleFragment(via, reply, &h)
+	sent := conn.take()
+	require.Len(t, sent, 4, "an immediate resend of fresh symbols")
+	for i, p := range sent {
+		fh, _, err := fragment.Parse(p[header.Len:])
+		require.NoError(t, err)
+		assert.True(t, fh.HasCookie())
+		assert.Equal(t, fragment.Cookie(bytes16(0xab)), fh.Cookie)
+		assert.Equal(t, uint8(4+i), fh.Index)
+	}
+
+	hm.handleFragment(via, reply, &h)
+	assert.Empty(t, conn.take(), "immediate resends are rate limited")
+
+	t.Run("replies for other indexes are ignored", func(t *testing.T) {
+		other := append([]byte(nil), reply...)
+		header.Encode(other, header.Version, header.Handshake, pqSubtype, 98, pqCookieCounter)
+		var oh header.H
+		require.NoError(t, oh.Parse(other))
+		hh.pq.cookieResendAt = time.Time{}
+		hm.handleFragment(via, other, &oh)
+		assert.Empty(t, conn.take())
+	})
+}
+
+func bytes16(b byte) []byte {
+	out := make([]byte, 16)
+	for i := range out {
+		out[i] = b
+	}
+	return out
+}
+
+func TestReloadPQ(t *testing.T) {
+	l := test.NewLogger()
+	c := config.NewC(l)
+	require.NoError(t, c.LoadString(`handshakes: {pq: {initiate: false}}`))
+	cfg := defaultHandshakeConfig
+	cfg.pq = pqHandshakeConfigFrom(c)
+	hm := NewHandshakeManager(l, newHostMap(l), newTestLighthouse(), &udp.NoopConn{}, cfg)
+	c.RegisterReloadCallback(hm.reloadPQ)
+
+	assert.Equal(t, pqHandshakeConfig{repairDelay: DefaultPQRepairDelay}, *hm.pqConfig())
+
+	require.NoError(t, c.ReloadConfigString(`handshakes: {pq: {initiate: true, accept: true, sticky: true, repair_delay: 25ms}}`))
+	assert.Equal(t, pqHandshakeConfig{initiate: true, accept: true, sticky: true, repairDelay: 25 * time.Millisecond}, *hm.pqConfig())
+
+	before := hm.pqConfig()
+	require.NoError(t, c.ReloadConfigString(`handshakes: {try_interval: 1s, pq: {initiate: true, accept: true, sticky: true, repair_delay: 25ms}}`))
+	assert.Same(t, before, hm.pqConfig(), "an unrelated change leaves the settings alone")
+
+	require.NoError(t, c.ReloadConfigString(`handshakes: {}`))
+	assert.Equal(t, pqHandshakeConfig{repairDelay: DefaultPQRepairDelay}, *hm.pqConfig(), "removing the section restores the defaults")
+}

main.go

@@ -232,15 +232,11 @@ func Main(c *config.C, configTest bool, buildVersion string, l *slog.Logger, dev
 		retries:        int64(c.GetInt("handshakes.retries", DefaultHandshakeRetries)),
 		triggerBuffer:  c.GetInt("handshakes.trigger_buffer", DefaultHandshakeTriggerBuffer),
 		messageMetrics: messageMetrics,
-		pq: pqHandshakeConfig{
-			initiate:    c.GetBool("handshakes.pq.initiate", false),
-			accept:      c.GetBool("handshakes.pq.accept", false),
-			sticky:      c.GetBool("handshakes.pq.sticky", false),
-			repairDelay: c.GetDuration("handshakes.pq.repair_delay", DefaultPQRepairDelay),
-		},
+		pq:             pqHandshakeConfigFrom(c),
 	}
 
 	handshakeManager := NewHandshakeManager(l, hostMap, lightHouse, udpConns[0], handshakeConfig)
+	c.RegisterReloadCallback(handshakeManager.reloadPQ)
 	lightHouse.handshakeTrigger = handshakeManager.trigger
 
 	ds, err := newDnsServerFromConfig(ctx, l, pki, hostMap, c)