[client] pqkem: report a peer as quantum resistant once its PSK lands

The per-peer "Quantum resistance" flag was computed only when the relay or ICE connection came up. When the KEM answer lands after the WireGuard endpoint is configured (the controller answering a crossing responder offer, the common case under lazy connections), the peer kept being reported as not converged while the PSK was programmed and traffic flowed. Reproduced in 4 of 7 fresh lab runs. Mark the peer in the status recorder once OnNewPSKReady has programmed the PSK, and OR that mark into every relay/ICE status update under the recorder's lock. The connection path reads the PQ state before configuring WireGuard and writes the status afterwards, so without the mark its stale "false" would clear the flag again. The mark is dropped with the peer, matching the PSK's lifetime, and the update is best-effort: a peer missing from the recorder never fails the PSK apply. Status-only change; PSK programming and strict gating are unaffected.
2026-10-22 13:29:08 +02:00 · 2026-10-19 15:08:34 +02:00
parent f669cba522
commit 2a2f7a9bbc
5 changed files with 173 additions and 8 deletions
@@ -608,6 +608,7 @@ func (e *Engine) startPQKEMManager(publicKey wgtypes.Key) error {
 	}
 	cbHandler := pqCallbackHandler{
 		wg:      e.wgInterface,
+		status:  e.statusRecorder,
 		applied: newAppliedGenerations(),
 		// 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
@@ -271,9 +271,12 @@ func (s *StatusChangeSubscription) deliver(snapshot *routerSnapshot, block bool)
 // every private-service request) don't contend against each other.
 // Pure read methods take RLock; anything that mutates state takes Lock.
 type Status struct {
-	mux              sync.RWMutex
-	muxRelays        sync.RWMutex
-	peers            map[string]State
+	mux       sync.RWMutex
+	muxRelays sync.RWMutex
+	peers     map[string]State
+	// pqMarked holds the peers with a programmed post-quantum PSK. It survives connection
+	// updates (see quantumResistantLocked) and is dropped with the peer.
+	pqMarked         map[string]struct{}
 	ipToKey          map[string]string
 	activeRoutePeers map[route.HAUniqueID]string
 	changeNotify     map[string]map[string]*StatusChangeSubscription // map[peerID]map[subscriptionID]*StatusChangeSubscription
@@ -338,6 +341,7 @@ type Status struct {
 func NewRecorder(mgmAddress string) *Status {
 	return &Status{
 		peers:                 make(map[string]State),
+		pqMarked:              make(map[string]struct{}),
 		ipToKey:               make(map[string]string),
 		activeRoutePeers:      make(map[route.HAUniqueID]string),
 		changeNotify:          make(map[string]map[string]*StatusChangeSubscription),
@@ -440,6 +444,7 @@ func (d *Status) RemovePeer(peerPubKey string) error {
 	}

 	delete(d.peers, peerPubKey)
+	delete(d.pqMarked, peerPubKey)
 	if mappedKey, exists := d.ipToKey[p.IP]; exists && mappedKey == peerPubKey {
 		delete(d.ipToKey, p.IP)
 	}
@@ -471,7 +476,7 @@ func (d *Status) UpdatePeerState(receivedState State) error {
 		peerState.LocalIceCandidateEndpoint = receivedState.LocalIceCandidateEndpoint
 		peerState.RemoteIceCandidateEndpoint = receivedState.RemoteIceCandidateEndpoint
 		peerState.RelayServerAddress = receivedState.RelayServerAddress
-		peerState.RosenpassEnabled = receivedState.RosenpassEnabled
+		peerState.RosenpassEnabled = d.quantumResistantLocked(receivedState.PubKey, receivedState.RosenpassEnabled)
 		peerState.clearLatency()
 	}

@@ -594,7 +599,7 @@ func (d *Status) UpdatePeerICEState(receivedState State) error {
 	peerState.RemoteIceCandidateType = receivedState.RemoteIceCandidateType
 	peerState.LocalIceCandidateEndpoint = receivedState.LocalIceCandidateEndpoint
 	peerState.RemoteIceCandidateEndpoint = receivedState.RemoteIceCandidateEndpoint
-	peerState.RosenpassEnabled = receivedState.RosenpassEnabled
+	peerState.RosenpassEnabled = d.quantumResistantLocked(receivedState.PubKey, receivedState.RosenpassEnabled)

 	d.peers[receivedState.PubKey] = peerState

@@ -631,7 +636,7 @@ func (d *Status) UpdatePeerRelayedState(receivedState State) error {
 	peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
 	peerState.Relayed = receivedState.Relayed
 	peerState.RelayServerAddress = receivedState.RelayServerAddress
-	peerState.RosenpassEnabled = receivedState.RosenpassEnabled
+	peerState.RosenpassEnabled = d.quantumResistantLocked(receivedState.PubKey, receivedState.RosenpassEnabled)

 	d.peers[receivedState.PubKey] = peerState

@@ -771,6 +776,34 @@ func (d *Status) UpdatePeerFQDN(peerPubKey, fqdn string) error {
 	return nil
 }

+// MarkPeerQuantumResistant flags the peer's tunnel as post-quantum protected. The flag is
+// otherwise computed only when the connection comes up, so a PSK that lands after that
+// (e.g. the KEM answer arriving after the WireGuard endpoint was configured) needs this
+// to show up in the status. The mark persists until the peer is removed.
+func (d *Status) MarkPeerQuantumResistant(peerPubKey string) error {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	peerState, ok := d.peers[peerPubKey]
+	if !ok {
+		return errors.New("peer doesn't exist")
+	}
+
+	d.pqMarked[peerPubKey] = struct{}{}
+	peerState.RosenpassEnabled = true
+	d.peers[peerPubKey] = peerState
+
+	return nil
+}
+
+// quantumResistantLocked merges a connection update's verdict with the PQ mark. The
+// connection path computes its verdict before configuring WireGuard, so it can write it
+// after the PSK was programmed; the mark keeps that stale false from clearing it.
+func (d *Status) quantumResistantLocked(peerPubKey string, fromConn bool) bool {
+	_, marked := d.pqMarked[peerPubKey]
+	return fromConn || marked
+}
+
 // UpdatePeerSSHHostKey updates peer's SSH host key
 func (d *Status) UpdatePeerSSHHostKey(peerPubKey string, sshHostKey []byte) error {
 	d.mux.Lock()
@@ -573,3 +573,46 @@ func TestActiveRoutePeers(t *testing.T) {
 	assert.False(t, ok)
 	assert.Equal(t, "peerB", active[netB])
 }
+
+// TestStatus_PQMarkSurvivesStaleConnectionUpdate reproduces a lab race: the relay/ICE
+// path builds its status update from a PQ state read before the KEM answer landed, and
+// writes it after the PSK was programmed and the peer marked quantum resistant. That
+// stale "false" must not clear the mark, since the PQ PSK lives until the peer is removed.
+func TestStatus_PQMarkSurvivesStaleConnectionUpdate(t *testing.T) {
+	const key = "eDQmZp9jhQGX7XdYpHGaFSpAouldr3Sbf7MZRcGTGGQ="
+	updates := map[string]func(*Status, State) error{
+		"relayed": (*Status).UpdatePeerRelayedState,
+		"ice":     (*Status).UpdatePeerICEState,
+		"state":   (*Status).UpdatePeerState,
+	}
+	for name, update := range updates {
+		t.Run(name, func(t *testing.T) {
+			status := NewRecorder("https://mgm")
+			require.NoError(t, status.AddPeer(key, "peer.netbird.cloud", "100.64.0.2", ""))
+			require.NoError(t, status.MarkPeerQuantumResistant(key))
+
+			require.NoError(t, update(status, State{PubKey: key, ConnStatus: StatusConnected, RosenpassEnabled: false}))
+
+			st, err := status.GetPeer(key)
+			require.NoError(t, err)
+			assert.True(t, st.RosenpassEnabled, "a stale connection update must not clear the PQ mark")
+		})
+	}
+}
+
+// TestStatus_PQMarkClearedOnRemove verifies the mark does not outlive the peer: once the
+// peer is removed and re-added, a connection update without PQ reports it as such.
+func TestStatus_PQMarkClearedOnRemove(t *testing.T) {
+	const key = "eDQmZp9jhQGX7XdYpHGaFSpAouldr3Sbf7MZRcGTGGQ="
+	status := NewRecorder("https://mgm")
+	require.NoError(t, status.AddPeer(key, "peer.netbird.cloud", "100.64.0.2", ""))
+	require.NoError(t, status.MarkPeerQuantumResistant(key))
+	require.NoError(t, status.RemovePeer(key))
+	require.NoError(t, status.AddPeer(key, "peer.netbird.cloud", "100.64.0.2", ""))
+
+	require.NoError(t, status.UpdatePeerRelayedState(State{PubKey: key, ConnStatus: StatusConnected, RosenpassEnabled: false}))
+
+	st, err := status.GetPeer(key)
+	require.NoError(t, err)
+	assert.False(t, st.RosenpassEnabled, "a re-added peer must not inherit the removed peer's PQ mark")
+}
@@ -42,10 +42,19 @@ type pqPresharedKeySetter interface {
 	SetPresharedKey(peerKey string, psk wgtypes.Key, updateOnly bool) error
 }

+// pqStatusMarker is the subset of the status recorder the ML-KEM callback needs: flagging
+// a peer as post-quantum protected once its PSK is programmed. *peer.Status satisfies it.
+type pqStatusMarker interface {
+	MarkPeerQuantumResistant(peerPubKey string) error
+}
+
 // pqCallbackHandler programs the derived PQ PSK onto the WireGuard peer. It is the
 // engine-side implementation of pqkem.CallbackHandler.
 type pqCallbackHandler struct {
 	wg pqPresharedKeySetter
+	// status reflects a newly programmed PSK in the peer's "Quantum resistance" flag.
+	// Nil disables it.
+	status pqStatusMarker
 	// applied drops an out-of-order PSK apply so a stale callback can't overwrite a newer
 	// key. Nil disables the guard.
 	applied *appliedGenerations
@@ -67,7 +76,17 @@ func (h pqCallbackHandler) OnNewPSKReady(remoteID pqkem.RemoteID, gen uint64, ps
 	// peer is not configured yet, so this is a no-op there and the PSK is instead
 	// pulled at peer-config time (pqHandshaker.PSK / conn.presharedKey).
 	log.Tracef("pqkem: programming PSK for peer %s (gen %d)", remoteID, gen)
-	return h.wg.SetPresharedKey(string(remoteID), wgtypes.Key(psk), true)
+	if err := h.wg.SetPresharedKey(string(remoteID), wgtypes.Key(psk), true); err != nil {
+		return err
+	}
+	// The status flag is computed when the connection comes up; a PSK landing afterwards
+	// would otherwise leave the peer reported as not quantum resistant.
+	if h.status != nil {
+		if err := h.status.MarkPeerQuantumResistant(string(remoteID)); err != nil {
+			log.Debugf("pqkem: mark peer %s quantum resistant in status: %v", remoteID, err)
+		}
+	}
+	return nil
 }

 // OnRekeyFailed reports a failed PQ (re)key convergence and re-bootstraps the KEM over
@@ -1,6 +1,75 @@
 package internal

-import "testing"
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+
+	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/internal/pqkem"
+)
+
+// fakePSKSetter records PSK programming and can be made to fail.
+type fakePSKSetter struct {
+	err   error
+	calls int
+}
+
+func (f *fakePSKSetter) SetPresharedKey(string, wgtypes.Key, bool) error {
+	f.calls++
+	return f.err
+}
+
+const pqTestPeerKey = "eDQmZp9jhQGX7XdYpHGaFSpAouldr3Sbf7MZRcGTGGQ="
+
+func newPQTestStatus(t *testing.T) *peer.Status {
+	t.Helper()
+	status := peer.NewRecorder("https://mgm")
+	require.NoError(t, status.AddPeer(pqTestPeerKey, "peer.netbird.cloud", "100.64.0.2", ""))
+	return status
+}
+
+// TestPQCallbackHandler_MarksPeerQuantumResistant reproduces the stale status flag: when
+// the KEM answer lands after the connection was set up, the per-peer "Quantum resistance"
+// flag was computed without a PSK and never refreshed. Programming the PSK must flag it.
+func TestPQCallbackHandler_MarksPeerQuantumResistant(t *testing.T) {
+	status := newPQTestStatus(t)
+	setter := &fakePSKSetter{}
+	h := pqCallbackHandler{wg: setter, applied: newAppliedGenerations(), status: status}
+
+	require.NoError(t, h.OnNewPSKReady(pqkem.RemoteID(pqTestPeerKey), 1, pqkem.PSK{1}))
+
+	st, err := status.GetPeer(pqTestPeerKey)
+	require.NoError(t, err)
+	assert.Equal(t, 1, setter.calls, "the PSK must be programmed once")
+	assert.True(t, st.RosenpassEnabled, "a programmed PQ PSK must mark the peer quantum resistant")
+}
+
+// TestPQCallbackHandler_FailedApplyKeepsFlag verifies the flag only follows a PSK that was
+// actually programmed: a failed SetPresharedKey must leave the peer unflagged.
+func TestPQCallbackHandler_FailedApplyKeepsFlag(t *testing.T) {
+	status := newPQTestStatus(t)
+	h := pqCallbackHandler{wg: &fakePSKSetter{err: errors.New("boom")}, applied: newAppliedGenerations(), status: status}
+
+	require.Error(t, h.OnNewPSKReady(pqkem.RemoteID(pqTestPeerKey), 1, pqkem.PSK{1}))
+
+	st, err := status.GetPeer(pqTestPeerKey)
+	require.NoError(t, err)
+	assert.False(t, st.RosenpassEnabled, "a PSK that failed to program must not mark the peer quantum resistant")
+}
+
+// TestPQCallbackHandler_UnknownStatusPeerDoesNotFailApply verifies the status update is
+// best-effort: a peer missing from the status recorder must not fail the PSK apply.
+func TestPQCallbackHandler_UnknownStatusPeerDoesNotFailApply(t *testing.T) {
+	setter := &fakePSKSetter{}
+	h := pqCallbackHandler{wg: setter, applied: newAppliedGenerations(), status: peer.NewRecorder("https://mgm")}
+
+	assert.NoError(t, h.OnNewPSKReady(pqkem.RemoteID(pqTestPeerKey), 1, pqkem.PSK{1}))
+	assert.Equal(t, 1, setter.calls, "the PSK must still be programmed")
+}

 // TestAppliedGenerations_DropsStale verifies the out-of-order guard: a generation is
 // accepted only when it is strictly newer than the last one applied for that peer, so a