[client] pqkem: apply PSK generations to WireGuard in order

Two derivations for the same peer can overlap, for example a rotation and
a Signal re-bootstrap. The generation guard checked and recorded the
generation under a lock but called SetPresharedKey after releasing it, so an
older apply already inside SetPresharedKey could finish after a newer one
and leave WireGuard on the stale key, splitting the peers until the next
exchange (reported in the security review of #7098).

Hold the lock across the generation check and the SetPresharedKey call. The
lock only serialises PSK applies, which are quick local calls, so this does
not reach into the manager's own locking.

client/internal/pqkem_adapter.go

@@ -23,17 +23,20 @@ func newAppliedGenerations() *appliedGenerations {
 	return &appliedGenerations{last: make(map[string]uint64)}
 }
 
-// claim reports whether gen is newer than the last applied for the peer; when it is, it
-// records gen as the newest and returns true. The check and record are atomic, so the
-// slow SetPresharedKey call runs outside this lock.
-func (a *appliedGenerations) claim(peer string, gen uint64) bool {
+// apply runs set when gen is newer than the last generation applied for the peer, and
+// reports whether it ran. The check and set run under one lock: with only the check
+// atomic, an older apply already inside set could finish after a newer one and restore
+// its stale key. The generation is recorded before set runs, so a failed newer apply
+// still keeps an older one from landing afterwards. The lock is shared by all peers,
+// which is cheap because set is a quick local interface call.
+func (a *appliedGenerations) apply(peer string, gen uint64, set func() error) (bool, error) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 	if gen <= a.last[peer] {
-		return false
+		return false, nil
 	}
 	a.last[peer] = gen
-	return true
+	return true, set()
 }
 
 // pqPresharedKeySetter is the subset of the WireGuard interface the ML-KEM callback
@@ -68,16 +71,26 @@ type pqCallbackHandler struct {
 // newer than the last applied for the peer is dropped, so a reordered apply can't
 // restore an older PSK over a newer one.
 func (h pqCallbackHandler) OnNewPSKReady(remoteID pqkem.RemoteID, gen uint64, psk pqkem.PSK) error {
-	if h.applied != nil && !h.applied.claim(string(remoteID), gen) {
-		log.Tracef("pqkem: dropping stale PSK apply for peer %s (gen %d)", remoteID, gen)
-		return nil
-	}
 	// 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.Tracef("pqkem: programming PSK for peer %s (gen %d)", remoteID, gen)
-	if err := h.wg.SetPresharedKey(string(remoteID), wgtypes.Key(psk), true); err != nil {
-		return err
+	set := func() error {
+		log.Tracef("pqkem: programming PSK for peer %s (gen %d)", remoteID, gen)
+		return h.wg.SetPresharedKey(string(remoteID), wgtypes.Key(psk), true)
+	}
+	if h.applied == nil {
+		if err := set(); err != nil {
+			return err
+		}
+	} else {
+		applied, err := h.applied.apply(string(remoteID), gen, set)
+		if err != nil {
+			return err
+		}
+		if !applied {
+			log.Tracef("pqkem: dropping stale PSK apply for peer %s (gen %d)", remoteID, gen)
+			return nil
+		}
 	}
 	// The status flag is computed when the connection comes up; a PSK landing afterwards
 	// would otherwise leave the peer reported as not quantum resistant.

client/internal/pqkem_adapter_test.go

@@ -2,7 +2,9 @@ package internal
 
 import (
 	"errors"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -76,25 +78,86 @@ func TestPQCallbackHandler_UnknownStatusPeerDoesNotFailApply(t *testing.T) {
 // reordered PSK callback cannot restore an older key over a newer one.
 func TestAppliedGenerations_DropsStale(t *testing.T) {
 	a := newAppliedGenerations()
+	accept := func(peer string, gen uint64) bool {
+		ran, err := a.apply(peer, gen, func() error { return nil })
+		require.NoError(t, err)
+		return ran
+	}
 
-	if !a.claim("peerA", 1) {
+	if !accept("peerA", 1) {
 		t.Fatal("first generation must be accepted")
 	}
-	if !a.claim("peerA", 2) {
+	if !accept("peerA", 2) {
 		t.Fatal("a newer generation must be accepted")
 	}
-	if a.claim("peerA", 2) {
+	if accept("peerA", 2) {
 		t.Fatal("re-applying the same generation must be dropped")
 	}
-	if a.claim("peerA", 1) {
+	if accept("peerA", 1) {
 		t.Fatal("an older generation arriving late must be dropped")
 	}
-	if !a.claim("peerA", 3) {
+	if !accept("peerA", 3) {
 		t.Fatal("a newer generation after a dropped stale one must still be accepted")
 	}
 
 	// Generations are tracked independently per peer.
-	if !a.claim("peerB", 1) {
+	if !accept("peerB", 1) {
 		t.Fatal("a different peer's first generation must be accepted")
 	}
 }
+
+// blockingPSKSetter holds the first SetPresharedKey call until released, recording the
+// key WireGuard ends up with.
+type blockingPSKSetter struct {
+	mu      sync.Mutex
+	calls   int
+	last    wgtypes.Key
+	entered chan struct{}
+	release chan struct{}
+}
+
+func (s *blockingPSKSetter) SetPresharedKey(_ string, psk wgtypes.Key, _ bool) error {
+	s.mu.Lock()
+	s.calls++
+	first := s.calls == 1
+	s.mu.Unlock()
+	if first {
+		close(s.entered)
+		<-s.release
+	}
+	s.mu.Lock()
+	s.last = psk
+	s.mu.Unlock()
+	return nil
+}
+
+// TestPQCallbackHandler_AppliesGenerationsInOrder covers two PSK callbacks for one peer
+// overlapping. The newer generation must be the one WireGuard keeps: with the generation
+// checked under a lock but the key set outside it, an older apply that was already
+// writing could finish last and restore its stale key.
+func TestPQCallbackHandler_AppliesGenerationsInOrder(t *testing.T) {
+	setter := &blockingPSKSetter{entered: make(chan struct{}), release: make(chan struct{})}
+	h := pqCallbackHandler{wg: setter, applied: newAppliedGenerations()}
+	older, newer := pqkem.PSK{3}, pqkem.PSK{4}
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		assert.NoError(t, h.OnNewPSKReady(pqkem.RemoteID(pqTestPeerKey), 3, older))
+	}()
+	<-setter.entered // the older apply is now inside SetPresharedKey
+
+	go func() {
+		defer wg.Done()
+		assert.NoError(t, h.OnNewPSKReady(pqkem.RemoteID(pqTestPeerKey), 4, newer))
+	}()
+	// Give the newer apply the chance to overtake the blocked older one.
+	time.Sleep(50 * time.Millisecond)
+	close(setter.release)
+	wg.Wait()
+
+	setter.mu.Lock()
+	defer setter.mu.Unlock()
+	assert.Equal(t, wgtypes.Key(newer), setter.last, "WireGuard must keep the newest generation's PSK")
+}