From de70516045cf57dfc3e12fcf99c39a4efc8518f7 Mon Sep 17 00:00:00 2001
From: riccardom <riccardomanfrin@gmail.com>
Date: Fri, 9 Oct 2026 09:54:28 +0200
Subject: [PATCH] [client] pqkem: drain in-flight signal handlers on Stop

Stop drained the data-path receive loop but not the host-driven signalling calls
(SignalOffer / SignalOnOffer / SignalOnAnswer). One in flight during shutdown
could still invoke OnNewPSKReady or OnRekeyFailed after Stop returned, touching a
WireGuard interface the engine was tearing down.

Register each signalling call in the wait group through beginCall, which refuses
once rootCtx is cancelled, mirroring startExchangeLocked's guard so no Add races
Stop's Wait. Stop already cancels, drains the transport, then Waits, so the Wait
now also covers these calls before state is cleared.

Found in cubic review on #7098 (client/internal/pqkem/manager.go:298).
---
 client/internal/pqkem/manager.go | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/client/internal/pqkem/manager.go b/client/internal/pqkem/manager.go
index db1a45c93..76a1a54c4 100644
--- a/client/internal/pqkem/manager.go
+++ b/client/internal/pqkem/manager.go
@@ -196,6 +196,21 @@ func (m *Manager) PSK(remoteID RemoteID) (PSK, bool) {
 	return psk, ok
 }
 
+// beginCall registers an in-flight host-driven signalling call in the wait group so Stop
+// drains it before clearing state, and refuses once the manager is stopping. It mirrors
+// startExchangeLocked's guard: rootCancel runs under m.mu before Stop's Wait, and this
+// Adds under the same lock after checking rootCtx, so no Add can race the Wait. The caller
+// must defer m.wait.Done() when this returns true.
+func (m *Manager) beginCall() bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.rootCtx.Err() != nil {
+		return false
+	}
+	m.wait.Add(1)
+	return true
+}
+
 // trace logs at LevelTrace, the verbose per-exchange lifecycle level gated by
 // NB_PQ_MLKEM_LOG_LEVEL=trace.
 func (m *Manager) trace(msg string, args ...any) {
@@ -313,6 +328,10 @@ func (m *Manager) Stop() {
 // PSK still survives idle in the manager (dropped only on account-level peer removal),
 // so a pure lazy wake with no re-negotiation reuses it via the conn's WG-config pull.
 func (m *Manager) SignalOffer(remoteID RemoteID) ([]byte, error) {
+	if !m.beginCall() {
+		return nil, fmt.Errorf("manager stopping")
+	}
+	defer m.wait.Done()
 	if !m.IsInitiator(remoteID) {
 		return nil, nil
 	}
@@ -370,6 +389,10 @@ func (m *Manager) ShouldSendBootstrapOffer(remoteID RemoteID) bool {
 // SignalOnOffer processes a KEM offer the host extracted from an incoming offer and
 // returns the KEM answer for the host to embed in its outgoing answer.
 func (m *Manager) SignalOnOffer(remoteID RemoteID, offer []byte) ([]byte, error) {
+	if !m.beginCall() {
+		return nil, fmt.Errorf("manager stopping")
+	}
+	defer m.wait.Done()
 	typ, msg, err := Decode(offer)
 	if err != nil {
 		return nil, fmt.Errorf("decode signal offer from %s: %w", remoteID, err)
@@ -397,6 +420,10 @@ func (m *Manager) SignalOnOffer(remoteID RemoteID, offer []byte) ([]byte, error)
 // SignalOnAnswer processes a KEM answer the host extracted from an incoming answer.
 // There is no reply: the next offer (over the data path) acknowledges this exchange.
 func (m *Manager) SignalOnAnswer(remoteID RemoteID, answer []byte) error {
+	if !m.beginCall() {
+		return fmt.Errorf("manager stopping")
+	}
+	defer m.wait.Done()
 	typ, msg, err := Decode(answer)
 	if err != nil {
 		return fmt.Errorf("decode signal answer from %s: %w", remoteID, err)