[client] pqkem: reprogram the latest PSK after a stale WireGuard configure

The conn reads the PSK before configuring WireGuard. A rotation committed
in between pushes its key first, and the configure then overwrites it with
the old one, splitting the keys until the next reconnect.

client/internal/engine.go

@@ -220,6 +220,9 @@ type Engine struct {
 	// pqkemManager runs the ML-KEM post-quantum PSK exchange (gated by NB_ENABLE_PQ_MLKEM).
 	// It owns the data-path transport and peer endpoint routing.
 	pqkemManager *pqkem.Manager
+	// pqPSKReapplier restores the latest PQ PSK after a conn configures WireGuard with a
+	// stale one. Set together with pqkemManager.
+	pqPSKReapplier *pqPSKReapplier
 
 	// syncMsgMux is used to guarantee sequential Management Service message processing
 	syncMsgMux *sync.Mutex
@@ -606,10 +609,11 @@ func (e *Engine) startPQKEMManager(publicKey wgtypes.Key) error {
 		log.Errorf("pqkem: transport bind failed, exchange disabled: %v", err)
 		return nil
 	}
+	applied := newAppliedGenerations()
 	cbHandler := pqCallbackHandler{
 		wg:      e.wgInterface,
 		status:  e.statusRecorder,
-		applied: newAppliedGenerations(),
+		applied: applied,
 		// On a persistent rekey failure, re-bootstrap the KEM over Signal: a fresh
 		// signalling offer starts a new exchange that overwrites the stalled PSK on both
 		// sides, recovering from a data-path desync.
@@ -626,6 +630,8 @@ func (e *Engine) startPQKEMManager(publicKey wgtypes.Key) error {
 		opts = append(opts, pqkem.WithPresharedSecret(e.config.PreSharedKey[:]))
 	}
 	e.pqkemManager = pqkem.NewManager(pqkem.LocalID(publicKey.String()), cbHandler, pqkem.NewLogger(), opts...)
+	// Shares the apply lock with the rotation push, so a reapply and a push are ordered.
+	e.pqPSKReapplier = &pqPSKReapplier{psks: e.pqkemManager, wg: e.wgInterface, applied: applied}
 	e.pqkemManager.Start(tr)
 	log.Infof("pqkem: enabled (udp port %d on overlay %s)", e.pqkemManager.LocalPort(), e.config.WgAddr.IP)
 	return nil
@@ -2095,7 +2101,7 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentV
 		NetMgr:    e.netMgr,
 	}
 	if e.pqkemManager != nil {
-		config.PQ = pqHandshaker{mgr: e.pqkemManager}
+		config.PQ = pqHandshaker{mgr: e.pqkemManager, reapplier: e.pqPSKReapplier}
 		config.PQStrict = pqkem.Strict()
 	}
 

client/internal/peer/conn.go

@@ -97,6 +97,10 @@ type PQHandshaker interface {
 	// PSK returns the peer's latest derived post-quantum PSK to program at WG
 	// peer-config time (the pull path). ok is false until one has been derived.
 	PSK(remoteKey string) (wgtypes.Key, bool)
+	// ReapplyPSK runs after the conn configured WireGuard with programmed, a PSK it read
+	// beforehand. A rotation committed in between was overwritten by that stale value, so
+	// the latest PSK is programmed again when it differs.
+	ReapplyPSK(remoteKey string, programmed *wgtypes.Key)
 	// SetRemoteAddr registers the peer's data-path endpoint learned from signalling:
 	// its WG overlay IP with the announced pq UDP port (port 0 means the peer omitted
 	// it and is on the default port).
@@ -538,6 +542,7 @@ func (conn *Conn) onICEConnectionIsReady(priority conntype.ConnPriority, iceConn
 		return
 	}
 	wgConfigWorkaround()
+	conn.reapplyPQPSK(presharedKey)
 
 	if conn.wgProxyRelay != nil {
 		conn.Log.Debugf("redirect packets from relayed conn to WireGuard")
@@ -585,6 +590,8 @@ func (conn *Conn) onICEStateDisconnected(sessionChanged bool) {
 		presharedKey := conn.presharedKey(conn.rosenpassRemoteKey, pqPSK)
 		if err := conn.endpointUpdater.SwitchWGEndpoint(conn.wgProxyRelay.EndpointAddr(), presharedKey); err != nil {
 			conn.Log.Errorf("failed to switch to relay conn: %v", err)
+		} else {
+			conn.reapplyPQPSK(presharedKey)
 		}
 
 		conn.currentConnPriority = conntype.Relay
@@ -664,7 +671,8 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 	}
 	updateTime := time.Now()
 	conn.enableWgWatcherIfNeeded(updateTime)
-	if err := conn.endpointUpdater.ConfigureWGEndpoint(wgProxy.EndpointAddr(), conn.presharedKey(rci.rosenpassPubKey, pqPSK)); err != nil {
+	presharedKey := conn.presharedKey(rci.rosenpassPubKey, pqPSK)
+	if err := conn.endpointUpdater.ConfigureWGEndpoint(wgProxy.EndpointAddr(), presharedKey); err != nil {
 		if err := wgProxy.CloseConn(); err != nil {
 			conn.Log.Warnf("Failed to close relay connection: %v", err)
 		}
@@ -676,6 +684,7 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 	}
 
 	wgConfigWorkaround()
+	conn.reapplyPQPSK(presharedKey)
 
 	conn.injectPendingFirstPacket(wgProxy, nil)
 
@@ -1148,6 +1157,16 @@ func (conn *Conn) pqPSK() (*wgtypes.Key, bool) {
 	return &psk, true
 }
 
+// reapplyPQPSK restores the latest PQ PSK after a WireGuard configure that used
+// programmed. It runs under conn.mu; the lock order conn.mu -> PSK apply -> manager holds,
+// since neither of the other two calls back into Conn while locked.
+func (conn *Conn) reapplyPQPSK(programmed *wgtypes.Key) {
+	if conn.config.PQ == nil {
+		return
+	}
+	conn.config.PQ.ReapplyPSK(conn.config.Key, programmed)
+}
+
 // IdlePresharedKey is the preshared key an idle (lazy) WireGuard peer is installed with.
 // An idle peer accepts a handshake the remote initiates directly, so in ML-KEM strict
 // mode it must already carry the PQ PSK, or the strict sentinel until one exists;

client/internal/peer/conn_pq_test.go

@@ -28,6 +28,7 @@ func (f fakePQ) ShouldSendBootstrapOffer(string) bool          { return false }
 func (f fakePQ) AnswerPayload(string, []byte) ([]byte, uint16) { return nil, 0 }
 func (f fakePQ) OnAnswer(string, []byte)                       {}
 func (f fakePQ) PSK(string) (wgtypes.Key, bool)                { return f.psk, f.ok }
+func (f fakePQ) ReapplyPSK(string, *wgtypes.Key)               {}
 func (f fakePQ) SetRemoteAddr(string, netip.AddrPort)          {}
 func (f fakePQ) OnDataPathRekeyed(string, time.Duration)       {}
 func (f fakePQ) OnDataPathDown(string)                         {}

client/internal/pqkem_adapter.go

@@ -39,6 +39,43 @@ func (a *appliedGenerations) apply(peer string, gen uint64, set func() error) (b
 	return true, set()
 }
 
+// locked runs fn under the apply lock, so it is ordered with every PSK push.
+func (a *appliedGenerations) locked(fn func()) {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	fn()
+}
+
+// pqPSKSource is the subset of the ML-KEM manager the reapply needs. *pqkem.Manager
+// satisfies it.
+type pqPSKSource interface {
+	PSK(remoteID pqkem.RemoteID) (pqkem.PSK, bool)
+}
+
+// pqPSKReapplier programs the latest PQ PSK again after a WireGuard configure that may
+// have used a stale one. The conn reads the PSK before configuring, and a rotation that
+// commits in between pushes its key first, which the configure then overwrites.
+type pqPSKReapplier struct {
+	psks    pqPSKSource
+	wg      pqPresharedKeySetter
+	applied *appliedGenerations
+}
+
+// reapply reads the latest PSK under the apply lock: a push already done is visible in
+// it, and a later push of a newer key runs after this one and wins.
+func (r *pqPSKReapplier) reapply(peer string, programmed *wgtypes.Key) {
+	r.applied.locked(func() {
+		psk, ok := r.psks.PSK(pqkem.RemoteID(peer))
+		if !ok || (programmed != nil && wgtypes.Key(psk) == *programmed) {
+			return
+		}
+		log.Debugf("pqkem: reprogramming the latest PSK for peer %s after a stale configure", peer)
+		if err := r.wg.SetPresharedKey(peer, wgtypes.Key(psk), true); err != nil {
+			log.Warnf("pqkem: reprogram PSK for peer %s: %v", peer, err)
+		}
+	})
+}
+
 // pqPresharedKeySetter is the subset of the WireGuard interface the ML-KEM callback
 // needs: programming a peer's preshared key. *iface.WGIface satisfies it.
 type pqPresharedKeySetter interface {
@@ -119,6 +156,8 @@ func (h pqCallbackHandler) OnRekeyFailed(remoteID pqkem.RemoteID) error {
 // wiring the host's signalling offers/answers to the KEM exchange.
 type pqHandshaker struct {
 	mgr *pqkem.Manager
+	// reapplier restores the latest PSK after a stale WireGuard configure. Nil disables it.
+	reapplier *pqPSKReapplier
 }
 
 // announcedPort is the PQ data-path port to advertise to peers. It is omitted (0) when
@@ -198,6 +237,14 @@ func (p pqHandshaker) PSK(remoteKey string) (wgtypes.Key, bool) {
 	return wgtypes.Key(psk), true
 }
 
+// ReapplyPSK programs the peer's latest PSK again when the conn configured WireGuard with
+// a different one.
+func (p pqHandshaker) ReapplyPSK(remoteKey string, programmed *wgtypes.Key) {
+	if p.reapplier != nil {
+		p.reapplier.reapply(remoteKey, programmed)
+	}
+}
+
 // SetRemoteAddr registers the peer's data-path endpoint learned from signalling. A
 // zero port means the peer omitted it (it is on DefaultPort), so we resolve it here —
 // DefaultPort lives in this package, not in peer. Sends only ever fire once the tunnel

client/internal/pqkem_adapter_test.go

@@ -161,3 +161,94 @@ func TestPQCallbackHandler_AppliesGenerationsInOrder(t *testing.T) {
 	defer setter.mu.Unlock()
 	assert.Equal(t, wgtypes.Key(newer), setter.last, "WireGuard must keep the newest generation's PSK")
 }
+
+// fakePSKSource returns a fixed PSK for every peer.
+type fakePSKSource struct {
+	psk pqkem.PSK
+	ok  bool
+}
+
+func (f fakePSKSource) PSK(pqkem.RemoteID) (pqkem.PSK, bool) { return f.psk, f.ok }
+
+// recordingPSKSetter records the last key programmed.
+type recordingPSKSetter struct {
+	calls int
+	last  wgtypes.Key
+}
+
+func (r *recordingPSKSetter) SetPresharedKey(_ string, psk wgtypes.Key, _ bool) error {
+	r.calls++
+	r.last = psk
+	return nil
+}
+
+// TestPQPSKReapplier_RestoresRotatedKey reproduces a configure that read the PSK before a
+// rotation committed: WireGuard was configured with the old key while the manager already
+// holds the new one, and the push of the new key was overwritten. The reapply must
+// program the new key again.
+func TestPQPSKReapplier_RestoresRotatedKey(t *testing.T) {
+	stale, latest := wgtypes.Key{1}, pqkem.PSK{2}
+	setter := &recordingPSKSetter{}
+	r := &pqPSKReapplier{psks: fakePSKSource{psk: latest, ok: true}, wg: setter, applied: newAppliedGenerations()}
+
+	r.reapply(pqTestPeerKey, &stale)
+
+	assert.Equal(t, 1, setter.calls, "a stale configure must be followed by one reprogram")
+	assert.Equal(t, wgtypes.Key(latest), setter.last, "WireGuard must end on the latest PSK")
+}
+
+// TestPQPSKReapplier_StrictSentinelReplaced covers strict mode: the configure programmed
+// the sentinel because no PSK existed yet, and the first PSK landed right after.
+func TestPQPSKReapplier_StrictSentinelReplaced(t *testing.T) {
+	sentinel, first := wgtypes.Key{9}, pqkem.PSK{3}
+	setter := &recordingPSKSetter{}
+	r := &pqPSKReapplier{psks: fakePSKSource{psk: first, ok: true}, wg: setter, applied: newAppliedGenerations()}
+
+	r.reapply(pqTestPeerKey, &sentinel)
+
+	assert.Equal(t, wgtypes.Key(first), setter.last, "the sentinel must be replaced by the first PQ PSK")
+}
+
+// TestPQPSKReapplier_NoWriteWhenCurrent: the common case, where the configure already used
+// the latest PSK or none exists yet, must not touch WireGuard.
+func TestPQPSKReapplier_NoWriteWhenCurrent(t *testing.T) {
+	current := pqkem.PSK{4}
+	programmed := wgtypes.Key(current)
+
+	setter := &recordingPSKSetter{}
+	r := &pqPSKReapplier{psks: fakePSKSource{psk: current, ok: true}, wg: setter, applied: newAppliedGenerations()}
+	r.reapply(pqTestPeerKey, &programmed)
+	assert.Zero(t, setter.calls, "a configure with the latest PSK needs no reprogram")
+
+	r = &pqPSKReapplier{psks: fakePSKSource{}, wg: setter, applied: newAppliedGenerations()}
+	r.reapply(pqTestPeerKey, nil)
+	assert.Zero(t, setter.calls, "with no PSK derived yet there is nothing to reprogram")
+}
+
+// lockCheckingSource records whether the apply lock was held when the PSK was read.
+type lockCheckingSource struct {
+	applied    *appliedGenerations
+	lockedRead bool
+}
+
+func (l *lockCheckingSource) PSK(pqkem.RemoteID) (pqkem.PSK, bool) {
+	if l.applied.mu.TryLock() {
+		l.applied.mu.Unlock()
+	} else {
+		l.lockedRead = true
+	}
+	return pqkem.PSK{5}, true
+}
+
+// TestPQPSKReapplier_ReadsUnderApplyLock checks the reapply reads the PSK while holding the
+// lock the rotation push takes. Reading it outside would let a push land between the read
+// and the write, and the reapply would restore the key it was meant to replace.
+func TestPQPSKReapplier_ReadsUnderApplyLock(t *testing.T) {
+	applied := newAppliedGenerations()
+	src := &lockCheckingSource{applied: applied}
+	r := &pqPSKReapplier{psks: src, wg: &recordingPSKSetter{}, applied: applied}
+
+	r.reapply(pqTestPeerKey, &wgtypes.Key{1})
+
+	assert.True(t, src.lockedRead, "the PSK must be read under the apply lock")
+}