diff --git a/client/internal/pqkem_adapter.go b/client/internal/pqkem_adapter.go index b15370d90..ad65a895c 100644 --- a/client/internal/pqkem_adapter.go +++ b/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. diff --git a/client/internal/pqkem_adapter_test.go b/client/internal/pqkem_adapter_test.go index b2e23684e..2deda1125 100644 --- a/client/internal/pqkem_adapter_test.go +++ b/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") +}