diff --git a/client/internal/engine.go b/client/internal/engine.go index 89768da54..e6e760cd0 100644 --- a/client/internal/engine.go +++ b/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() } diff --git a/client/internal/peer/conn.go b/client/internal/peer/conn.go index e0e303d8a..9de4c8000 100644 --- a/client/internal/peer/conn.go +++ b/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; diff --git a/client/internal/peer/conn_pq_test.go b/client/internal/peer/conn_pq_test.go index ae11f70e8..6c14fe6cd 100644 --- a/client/internal/peer/conn_pq_test.go +++ b/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) {} diff --git a/client/internal/pqkem_adapter.go b/client/internal/pqkem_adapter.go index 8ec4a326c..8d7a66950 100644 --- a/client/internal/pqkem_adapter.go +++ b/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 diff --git a/client/internal/pqkem_adapter_test.go b/client/internal/pqkem_adapter_test.go index 2deda1125..0492f687f 100644 --- a/client/internal/pqkem_adapter_test.go +++ b/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") +}