From 8b7c105b5eb71e9640fec15c97a526b3e779dd4b Mon Sep 17 00:00:00 2001 From: riccardom Date: Sun, 26 Jul 2026 14:04:51 +0200 Subject: [PATCH] pqkem: apply derived PSK at WG peer-config time (pull) + keep push for rekey --- client/internal/peer/conn.go | 12 ++++++++++++ client/internal/pqkem/convergence.go | 2 ++ client/internal/pqkem/manager.go | 13 +++++++++++++ client/internal/pqkem_adapter.go | 13 +++++++++++++ 4 files changed, 40 insertions(+) diff --git a/client/internal/peer/conn.go b/client/internal/peer/conn.go index 34d128ca1..7b9f22108 100644 --- a/client/internal/peer/conn.go +++ b/client/internal/peer/conn.go @@ -87,6 +87,9 @@ type PQHandshaker interface { AnswerPayload(remoteKey string, recvOffer []byte) (payload []byte, port int) // OnAnswer feeds a received KEM answer (nil if absent). OnAnswer(remoteKey string, recvAnswer []byte) + // 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) } // ConnConfig is a peer Connection configuration @@ -1007,6 +1010,15 @@ func (conn *Conn) AgentVersionString() string { } func (conn *Conn) presharedKey(remoteRosenpassKey []byte) *wgtypes.Key { + // Post-quantum: once the ML-KEM exchange has derived a PSK for this peer, program + // it here so the peer's next WireGuard handshake adopts it. Applied at peer-config + // time (bootstrap / reconnect); steady-state rotation is pushed separately. + if conn.config.PQ != nil { + if psk, ok := conn.config.PQ.PSK(conn.config.Key); ok { + return &psk + } + } + if conn.config.RosenpassConfig.PubKey == nil { return conn.config.WgConfig.PreSharedKey } diff --git a/client/internal/pqkem/convergence.go b/client/internal/pqkem/convergence.go index dc6364a8b..63e1d61e9 100644 --- a/client/internal/pqkem/convergence.go +++ b/client/internal/pqkem/convergence.go @@ -84,6 +84,7 @@ func (m *Manager) processOffer(remoteID RemoteID, o *OfferMsg) ([]byte, error) { ex.state = stateAwaitingAck ex.lastSent = raw ex.pendingPSK = psk + m.psks[remoteID] = psk m.mu.Unlock() // Commit optimistically so our data path can rekey to the new PSK. @@ -118,6 +119,7 @@ func (m *Manager) processAnswer(remoteID RemoteID, a *AnswerMsg) error { m.mu.Lock() m.established[remoteID] = true m.failures[remoteID] = 0 + m.psks[remoteID] = psk m.mu.Unlock() return m.cbHandler.OnNewPSKReady(remoteID, psk) diff --git a/client/internal/pqkem/manager.go b/client/internal/pqkem/manager.go index de0615bde..9cc06a00a 100644 --- a/client/internal/pqkem/manager.go +++ b/client/internal/pqkem/manager.go @@ -98,6 +98,7 @@ type Manager struct { exchanges map[RemoteID]*exchangeCtl // in-flight exchange per peer established map[RemoteID]bool // peer has completed at least one exchange failures map[RemoteID]int // consecutive rekey failures per peer + psks map[RemoteID]PSK // latest derived PSK per peer (pulled at WG peer-config time) peerAddrs map[RemoteID]netip.AddrPort // remoteID -> data-path endpoint (send routing) peersByAddr map[netip.AddrPort]RemoteID // reverse: source endpoint -> remoteID (inbound) wait sync.WaitGroup @@ -123,6 +124,7 @@ func NewManager(localID LocalID, h CallbackHandler, logger *slog.Logger) *Manage exchanges: make(map[RemoteID]*exchangeCtl), established: make(map[RemoteID]bool), failures: make(map[RemoteID]int), + psks: make(map[RemoteID]PSK), peerAddrs: make(map[RemoteID]netip.AddrPort), peersByAddr: make(map[netip.AddrPort]RemoteID), } @@ -158,6 +160,15 @@ func (m *Manager) IsInitiator(remoteID RemoteID) bool { return string(m.localID) > string(remoteID) } +// PSK returns the latest PSK derived for the peer, for the host to program at WG +// peer-config time (the pull path). ok is false until an exchange has derived one. +func (m *Manager) PSK(remoteID RemoteID) (PSK, bool) { + m.mu.Lock() + defer m.mu.Unlock() + psk, ok := m.psks[remoteID] + return psk, ok +} + // AddPeer registers where a peer's data-path messages are sent and received: its // overlay endpoint (IP:port). Re-adding updates the endpoint. func (m *Manager) AddPeer(remoteID RemoteID, endpoint netip.AddrPort) { @@ -184,6 +195,7 @@ func (m *Manager) RemovePeer(remoteID RemoteID) { } delete(m.established, remoteID) delete(m.failures, remoteID) + delete(m.psks, remoteID) if ep, ok := m.peerAddrs[remoteID]; ok { delete(m.peersByAddr, ep) delete(m.peerAddrs, remoteID) @@ -200,6 +212,7 @@ func (m *Manager) Stop() { t := m.transport m.transport = nil m.exchanges = make(map[RemoteID]*exchangeCtl) + m.psks = make(map[RemoteID]PSK) m.mu.Unlock() if t != nil { if err := t.Close(); err != nil { diff --git a/client/internal/pqkem_adapter.go b/client/internal/pqkem_adapter.go index 4373520d7..611e72498 100644 --- a/client/internal/pqkem_adapter.go +++ b/client/internal/pqkem_adapter.go @@ -22,6 +22,10 @@ type pqCallbackHandler struct { // OnNewPSKReady programs the freshly derived PSK for the peer (updateOnly: a no-op // if the peer is not present, mirroring Rosenpass). remoteID is the peer's WG pubkey. func (h pqCallbackHandler) OnNewPSKReady(remoteID pqkem.RemoteID, psk pqkem.PSK) error { + // updateOnly: applies to an already-configured peer (rotation). At bootstrap the + // 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.Debugf("pqkem: programming PSK for peer %s", remoteID) return h.wg.SetPresharedKey(string(remoteID), wgtypes.Key(psk), true) } @@ -65,3 +69,12 @@ func (p pqHandshaker) OnAnswer(remoteKey string, recvAnswer []byte) { log.Warnf("pqkem: process answer from %s: %v", remoteKey, err) } } + +// PSK exposes the peer's derived PSK for the conn to program at WG peer-config time. +func (p pqHandshaker) PSK(remoteKey string) (wgtypes.Key, bool) { + psk, ok := p.mgr.PSK(pqkem.RemoteID(remoteKey)) + if !ok { + return wgtypes.Key{}, false + } + return wgtypes.Key(psk), true +}