From e641c2d8a5bb436fbae9c212f1792b99c3c71212 Mon Sep 17 00:00:00 2001 From: riccardom Date: Wed, 22 Jul 2026 09:35:05 +0200 Subject: [PATCH] Manages convergence --- client/internal/pqkem/convergence.go | 193 ++++++++++++++++++++++ client/internal/pqkem/convergence_test.go | 88 ++++++++++ client/internal/pqkem/driver.go | 143 +++++++--------- client/internal/pqkem/manager.go | 52 +++--- 4 files changed, 360 insertions(+), 116 deletions(-) create mode 100644 client/internal/pqkem/convergence.go create mode 100644 client/internal/pqkem/convergence_test.go diff --git a/client/internal/pqkem/convergence.go b/client/internal/pqkem/convergence.go new file mode 100644 index 000000000..8b2ac93bb --- /dev/null +++ b/client/internal/pqkem/convergence.go @@ -0,0 +1,193 @@ +package pqkem + +import ( + "context" + "time" +) + +// handleOffer (responder) derives and sends the answer for a new exchange, or +// resends the cached answer for a duplicate offer (same exchangeID) without +// re-deriving. The responder is purely reactive: no retransmit loop, no deadline. +func (d *Driver) handleOffer(remoteWgKey string, o *OfferMsg) error { + d.mu.Lock() + if ex := d.exchanges[remoteWgKey]; ex != nil && ex.id == o.ExchangeID { + last := ex.lastSent + d.mu.Unlock() + if last == nil { + // another goroutine reserved this exchange and is deriving the answer; + // it will send it. Dropping this duplicate avoids a second derivation + // (Respond is randomized -> a different PSK). + return nil + } + return d.transport.Send(remoteWgKey, last) + } + // Reserve the slot (lastSent nil = deriving) so a concurrent duplicate offer bails. + d.exchanges[remoteWgKey] = &exchangeCtl{ + id: o.ExchangeID, + startedAt: time.Now(), + isInitial: !d.established[remoteWgKey], + } + d.mu.Unlock() + + answer, err := d.mgr.HandleOffer(remoteWgKey, o) + if err != nil { + return err + } + raw, err := answer.Encode() + if err != nil { + return err + } + + d.mu.Lock() + if ex := d.exchanges[remoteWgKey]; ex != nil && ex.id == o.ExchangeID { + ex.lastSent = raw + } + d.mu.Unlock() + + return d.transport.Send(remoteWgKey, raw) +} + +// handleAnswer (initiator) derives the PSK, surfaces it, and sends the confirm. The +// exchange then switches its retransmit payload to the confirm. Stale or duplicate +// answers are ignored. +func (d *Driver) handleAnswer(remoteWgKey string, a *AnswerMsg) error { + // Claim the answer under the lock (set answered) so a concurrent/duplicate answer + // bails before calling the Manager. + d.mu.Lock() + ex := d.exchanges[remoteWgKey] + if ex == nil || ex.id != a.ExchangeID || ex.cancel == nil || ex.answered { + d.mu.Unlock() + return nil + } + ex.answered = true + d.mu.Unlock() + + psk, confirm, err := d.mgr.HandleAnswer(remoteWgKey, a) + if err != nil { + return err + } + if err := d.wg.OnNewPSKReady(remoteWgKey, psk); err != nil { + return err + } + raw, err := confirm.Encode() + if err != nil { + return err + } + + d.mu.Lock() + if ex := d.exchanges[remoteWgKey]; ex != nil && ex.id == a.ExchangeID { + ex.lastSent = raw // loop now retransmits the confirm + d.established[remoteWgKey] = true + d.failures[remoteWgKey] = 0 + _ = time.Since(ex.startedAt) // convergence latency (metrics hook, later step) + } + d.mu.Unlock() + + return d.transport.Send(remoteWgKey, raw) +} + +// handleConfirm (responder) commits the pending PSK. Stale or duplicate confirms +// (the initiator best-effort resends it) are ignored. +func (d *Driver) handleConfirm(remoteWgKey string, c *ConfirmMsg) error { + // Claim-and-remove the exchange under the lock so a duplicate confirm (the + // initiator best-effort resends it) bails before calling the Manager. + d.mu.Lock() + ex := d.exchanges[remoteWgKey] + if ex == nil || ex.id != c.ExchangeID { + d.mu.Unlock() + return nil + } + delete(d.exchanges, remoteWgKey) + d.established[remoteWgKey] = true + d.failures[remoteWgKey] = 0 + _ = time.Since(ex.startedAt) // convergence latency (metrics hook, later step) + d.mu.Unlock() + + psk, err := d.mgr.HandleConfirm(remoteWgKey, c) + if err != nil { + return err + } + return d.wg.OnNewPSKReady(remoteWgKey, psk) +} + +// initiatorLoop retransmits the initiator's outstanding message: the offer until +// the answer arrives (bounded by maxRetries -> failure), then the confirm a few +// best-effort times before stopping. The convergence deadline is thus derived from +// maxRetries * retryInterval; there is no separate deadline timer. +func (d *Driver) initiatorLoop(ctx context.Context, remoteWgKey string, id ExchangeID) { + defer d.wait.Done() + t := time.NewTicker(d.retryInterval) + defer t.Stop() + + offerAttempts, confirmsSent := 0, 0 + for { + select { + case <-ctx.Done(): + return + case <-t.C: + d.mu.Lock() + ex := d.exchanges[remoteWgKey] + if ex == nil || ex.id != id { + d.mu.Unlock() + return + } + + if !ex.answered { + if offerAttempts >= d.maxRetries { + delete(d.exchanges, remoteWgKey) + fail := d.registerFailureLocked(remoteWgKey, ex.isInitial) + d.mu.Unlock() + d.raiseFailure(remoteWgKey, fail) + return + } + msg := ex.lastSent + offerAttempts++ + d.mu.Unlock() + d.retransmit(remoteWgKey, msg) + continue + } + + if confirmsSent >= confirmRetransmits { + delete(d.exchanges, remoteWgKey) + d.mu.Unlock() + return + } + msg := ex.lastSent + confirmsSent++ + d.mu.Unlock() + d.retransmit(remoteWgKey, msg) + } + } +} + +// registerFailureLocked applies policy B and reports whether OnRekeyFailed is due: +// the initial exchange fails immediately; a rekey tolerates up to maxRekeyFailures +// consecutive misses (we stay on the still-valid previous PSK) before failing. +// Assumes d.mu is held. +func (d *Driver) registerFailureLocked(remoteWgKey string, isInitial bool) bool { + if isInitial { + return true + } + d.failures[remoteWgKey]++ + if d.failures[remoteWgKey] >= d.maxRekeyFailures { + d.failures[remoteWgKey] = 0 + return true + } + return false +} + +func (d *Driver) raiseFailure(remoteWgKey string, fail bool) { + if !fail { + d.logger.Warn("pqkem rekey attempt timed out, will retry next cycle", "peer", remoteWgKey) + return + } + if err := d.wg.OnRekeyFailed(remoteWgKey); err != nil { + d.logger.Error("pqkem OnRekeyFailed handler error", "peer", remoteWgKey, "err", err) + } +} + +func (d *Driver) retransmit(remoteWgKey string, msg []byte) { + if err := d.transport.Send(remoteWgKey, msg); err != nil { + d.logger.Warn("pqkem retransmit failed", "peer", remoteWgKey, "err", err) + } +} diff --git a/client/internal/pqkem/convergence_test.go b/client/internal/pqkem/convergence_test.go new file mode 100644 index 000000000..e05331210 --- /dev/null +++ b/client/internal/pqkem/convergence_test.go @@ -0,0 +1,88 @@ +package pqkem + +import ( + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type dropTransport struct{} + +func (dropTransport) Send(string, []byte) error { return nil } + +// gate is a loopback transport with a switchable drop flag. +type gate struct { + local string + peer *Driver + drop atomic.Bool +} + +func (g *gate) Send(remote string, msg []byte) error { + if g.drop.Load() { + return nil + } + cp := append([]byte(nil), msg...) + return g.peer.HandleInbound(g.local, cp) +} + +func TestDriver_InitialTimeoutFailsImmediately(t *testing.T) { + wg := newFakeWG() + d := NewDriver("bbbb", dropTransport{}, wg, time.Hour, nil) // bbbb > aaaa -> initiator + d.retryInterval = 5 * time.Millisecond + d.maxRetries = 3 + d.AddPeer("aaaa") + defer d.Stop() + + require.NoError(t, d.initiateRekey("aaaa")) + + // no answer will ever come -> the initial exchange fails fast. + require.Eventually(t, func() bool { + wg.mu.Lock() + defer wg.mu.Unlock() + return len(wg.failed) == 1 + }, time.Second, 5*time.Millisecond) +} + +func TestDriver_RekeyToleratesKFailures(t *testing.T) { + gA := &gate{local: "aaaa"} + gB := &gate{local: "bbbb"} + wgA := newFakeWG() + wgB := newFakeWG() + + dA := NewDriver("aaaa", gA, wgA, time.Hour, nil) + dB := NewDriver("bbbb", gB, wgB, time.Hour, nil) + gA.peer = dB + gB.peer = dA + dB.retryInterval = 5 * time.Millisecond + dB.maxRetries = 2 + dA.AddPeer("bbbb") + dB.AddPeer("aaaa") + defer dA.Stop() + defer dB.Stop() + + // first exchange succeeds -> peer becomes established (subsequent failures are rekeys). + require.NoError(t, dB.initiateRekey("aaaa")) + require.NotEqual(t, PSK{}, wgB.psk("aaaa")) + + // now drop B's outbound: rekeys can no longer converge. + gB.drop.Store(true) + + // K-1 failures must NOT raise OnRekeyFailed. + for i := 0; i < DefaultMaxRekeyFailures-1; i++ { + require.NoError(t, dB.initiateRekey("aaaa")) + time.Sleep(50 * time.Millisecond) + } + require.Equal(t, 0, failedCount(wgB), "no failure before K attempts") + + // the K-th failure raises it once. + require.NoError(t, dB.initiateRekey("aaaa")) + require.Eventually(t, func() bool { return failedCount(wgB) == 1 }, time.Second, 5*time.Millisecond) +} + +func failedCount(f *fakeWG) int { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.failed) +} diff --git a/client/internal/pqkem/driver.go b/client/internal/pqkem/driver.go index c26d4ec94..8e0980dbc 100644 --- a/client/internal/pqkem/driver.go +++ b/client/internal/pqkem/driver.go @@ -12,15 +12,18 @@ const ( // DefaultRekeyInterval matches WireGuard's own REKEY_AFTER_TIME so the freshly // rotated PSK is naturally adopted by WG's next handshake without forcing one. DefaultRekeyInterval = 2 * time.Minute - // DefaultRetryInterval is how often an in-flight exchange retransmits. + // DefaultRetryInterval is how often the initiator retransmits its outstanding + // message (offer, then confirm) while an exchange is in flight. DefaultRetryInterval = 2 * time.Second - // DefaultConvergenceTimeout bounds a single exchange before it is declared failed. - DefaultConvergenceTimeout = 20 * time.Second - // DefaultMaxRekeyFailures is how many consecutive rekey (non-initial) failures are - // tolerated before OnRekeyFailed is raised. The initial exchange fails immediately. + // DefaultMaxRetries bounds how many times the offer is retransmitted before the + // exchange is declared failed. The convergence deadline is thus derived as + // MaxRetries * RetryInterval — there is no separate deadline timer. + DefaultMaxRetries = 10 + // DefaultMaxRekeyFailures is how many consecutive rekey (non-initial) failures + // are tolerated before OnRekeyFailed. The initial exchange fails immediately. DefaultMaxRekeyFailures = 3 - // confirmRetransmits is how many extra times the initiator best-effort resends the - // confirm (the exchange's unacked last message) to reduce a stuck responder. + // confirmRetransmits is how many times the initiator best-effort resends the + // confirm after converging, to cover its loss without a dedicated goroutine. confirmRetransmits = 3 ) @@ -31,18 +34,18 @@ type Transport interface { Send(remoteWgKey string, msg []byte) error } -type encodableMsg interface { - Encode() ([]byte, error) -} - -// exchangeCtl tracks one in-flight exchange for a peer: its id, the cancel for its -// retransmit/deadline goroutine, when it started (for convergence latency) and -// whether it is the peer's first (initial) exchange (which fails hard on timeout). +// exchangeCtl tracks one in-flight exchange for a peer. A single lastSent holds the +// message currently being (re)transmitted: for the initiator it is the offer and +// then, once answered, the confirm; for the responder it is the answer, resent on a +// duplicate offer. Only the initiator runs a retransmit loop (cancel != nil); the +// responder is purely reactive. type exchangeCtl struct { id ExchangeID - cancel context.CancelFunc startedAt time.Time isInitial bool + lastSent []byte + cancel context.CancelFunc + answered bool // initiator: the answer arrived; retransmit phase is now the confirm } // Driver ties the pure Manager to the outside world: per-peer rekey timer, inbound @@ -54,10 +57,10 @@ type Driver struct { wg WGCallbackHandler logger *slog.Logger - rekeyInterval time.Duration - retryInterval time.Duration - convergenceTimeout time.Duration - maxRekeyFailures int + rekeyInterval time.Duration + retryInterval time.Duration + maxRetries int + maxRekeyFailures int rootCtx context.Context rootCancel context.CancelFunc @@ -71,7 +74,7 @@ type Driver struct { } // NewDriver builds a driver for the local peer. A zero interval falls back to -// DefaultRekeyInterval; a nil logger falls back to slog.Default(). Retry/timeout/K +// DefaultRekeyInterval; a nil logger falls back to slog.Default(). Retry/retries/K // use their defaults and can be overridden on the returned struct before use. func NewDriver(localWgKey string, t Transport, h WGCallbackHandler, interval time.Duration, logger *slog.Logger) *Driver { if interval <= 0 { @@ -82,20 +85,20 @@ func NewDriver(localWgKey string, t Transport, h WGCallbackHandler, interval tim } ctx, cancel := context.WithCancel(context.Background()) return &Driver{ - mgr: NewManager(localWgKey), - transport: t, - wg: h, - logger: logger, - rekeyInterval: interval, - retryInterval: DefaultRetryInterval, - convergenceTimeout: DefaultConvergenceTimeout, - maxRekeyFailures: DefaultMaxRekeyFailures, - rootCtx: ctx, - rootCancel: cancel, - peers: make(map[string]context.CancelFunc), - exchanges: make(map[string]*exchangeCtl), - established: make(map[string]bool), - failures: make(map[string]int), + mgr: NewManager(localWgKey), + transport: t, + wg: h, + logger: logger, + rekeyInterval: interval, + retryInterval: DefaultRetryInterval, + maxRetries: DefaultMaxRetries, + maxRekeyFailures: DefaultMaxRekeyFailures, + rootCtx: ctx, + rootCancel: cancel, + peers: make(map[string]context.CancelFunc), + exchanges: make(map[string]*exchangeCtl), + established: make(map[string]bool), + failures: make(map[string]int), } } @@ -120,7 +123,9 @@ func (d *Driver) RemovePeer(remoteWgKey string) { delete(d.peers, remoteWgKey) } if ex, ok := d.exchanges[remoteWgKey]; ok { - ex.cancel() + if ex.cancel != nil { + ex.cancel() + } delete(d.exchanges, remoteWgKey) } delete(d.established, remoteWgKey) @@ -145,7 +150,6 @@ func (d *Driver) HandleInbound(remoteWgKey string, raw []byte) error { if err != nil { return fmt.Errorf("decode from %s: %w", remoteWgKey, err) } - switch typ { case MsgOffer: return d.handleOffer(remoteWgKey, msg.(*OfferMsg)) @@ -158,46 +162,6 @@ func (d *Driver) HandleInbound(remoteWgKey string, raw []byte) error { } } -func (d *Driver) handleOffer(remoteWgKey string, o *OfferMsg) error { - answer, err := d.mgr.HandleOffer(remoteWgKey, o) - if err != nil { - return err - } - // Arm a responder exchange (deadline waiting for the confirm) unless one for this - // id is already armed — a retransmitted offer just re-sends the same answer. - d.armExchange(remoteWgKey, o.ExchangeID, nil) - return d.send(remoteWgKey, answer) -} - -func (d *Driver) handleAnswer(remoteWgKey string, a *AnswerMsg) error { - psk, confirm, err := d.mgr.HandleAnswer(remoteWgKey, a) - if err != nil { - return err - } - if err := d.wg.OnNewPSKReady(remoteWgKey, psk); err != nil { - return err - } - // Initiator has the PSK now -> this side has converged. - d.onConverged(remoteWgKey, a.ExchangeID) - if err := d.send(remoteWgKey, confirm); err != nil { - return err - } - d.retransmitConfirm(remoteWgKey, confirm) - return nil -} - -func (d *Driver) handleConfirm(remoteWgKey string, c *ConfirmMsg) error { - psk, err := d.mgr.HandleConfirm(remoteWgKey, c) - if err != nil { - return err - } - if err := d.wg.OnNewPSKReady(remoteWgKey, psk); err != nil { - return err - } - d.onConverged(remoteWgKey, c.ExchangeID) - return nil -} - // initiateRekey starts a fresh exchange when the local peer is the initiator for // this remote peer; the responder waits for the offer instead. Exposed (unexported // but directly callable) so tests can drive a rekey without waiting on the ticker. @@ -213,7 +177,24 @@ func (d *Driver) initiateRekey(remoteWgKey string) error { if err != nil { return err } - d.armExchange(remoteWgKey, offer.ExchangeID, raw) + + ctx, cancel := context.WithCancel(d.rootCtx) + d.mu.Lock() + if old := d.exchanges[remoteWgKey]; old != nil && old.cancel != nil { + old.cancel() + } + d.exchanges[remoteWgKey] = &exchangeCtl{ + id: offer.ExchangeID, + startedAt: time.Now(), + isInitial: !d.established[remoteWgKey], + lastSent: raw, + cancel: cancel, + } + d.mu.Unlock() + + d.wait.Add(1) + go d.initiatorLoop(ctx, remoteWgKey, offer.ExchangeID) + return d.transport.Send(remoteWgKey, raw) } @@ -232,11 +213,3 @@ func (d *Driver) rekeyLoop(ctx context.Context, remoteWgKey string) { } } } - -func (d *Driver) send(remoteWgKey string, m encodableMsg) error { - raw, err := m.Encode() - if err != nil { - return err - } - return d.transport.Send(remoteWgKey, raw) -} diff --git a/client/internal/pqkem/manager.go b/client/internal/pqkem/manager.go index bfdd1a1e2..cc746c423 100644 --- a/client/internal/pqkem/manager.go +++ b/client/internal/pqkem/manager.go @@ -31,11 +31,8 @@ type peerSession struct { // initiator side: in-flight ephemeral keypair between offer and answer. initiator *Initiator - // responder side: the derived PSK and the answer for this exchange. Both are - // retained after establishment so retransmitted offers/confirms (same exchangeID) - // are handled idempotently — never re-derive a different key. + // responder side: PSK derived on the offer, committed only on the confirm. pendingPSK PSK - answer *AnswerMsg } type Manager struct { @@ -87,33 +84,24 @@ func (m *Manager) StartExchange(remoteWgKey string) (*OfferMsg, error) { // HandleOffer processes a received offer as the responder, derives the PSK, and // returns the answer to send. The PSK is held pending and NOT returned: the caller -// must not program it until HandleConfirm. A retransmitted offer (same exchangeID) -// returns the cached answer without re-deriving, so both sides keep the same key. +// must not program it until HandleConfirm. Retransmitted offers are deduplicated by +// the driver (via exchangeID) so this is called once per exchange and never +// re-derives a different key for the same round. func (m *Manager) HandleOffer(remoteWgKey string, o *OfferMsg) (*AnswerMsg, error) { - m.mu.Lock() - if s, ok := m.sessions[remoteWgKey]; ok && s.exchangeID == o.ExchangeID && s.answer != nil { - ans := s.answer - m.mu.Unlock() - return ans, nil - } - m.mu.Unlock() - answer, psk, err := Respond(o.KEMOffer, m.binding(remoteWgKey)) if err != nil { return nil, err } - ansMsg := &AnswerMsg{ExchangeID: o.ExchangeID, KEMAnswer: answer} m.mu.Lock() m.sessions[remoteWgKey] = &peerSession{ state: stateAwaitingConfirm, exchangeID: o.ExchangeID, pendingPSK: psk, - answer: ansMsg, } m.mu.Unlock() - return ansMsg, nil + return &AnswerMsg{ExchangeID: o.ExchangeID, KEMAnswer: answer}, nil } // HandleAnswer processes a received answer as the initiator. On success it returns @@ -121,39 +109,39 @@ func (m *Manager) HandleOffer(remoteWgKey string, o *OfferMsg) (*AnswerMsg, erro func (m *Manager) HandleAnswer(remoteWgKey string, a *AnswerMsg) (PSK, *ConfirmMsg, error) { m.mu.Lock() s, ok := m.sessions[remoteWgKey] - m.mu.Unlock() - if !ok || s.state != stateAwaitingAnswer { + m.mu.Unlock() return PSK{}, nil, fmt.Errorf("no pending offer for peer %s", remoteWgKey) } if s.exchangeID != a.ExchangeID { // stale answer (e.g. to a pre-restart offer) — drop, keep waiting. + m.mu.Unlock() return PSK{}, nil, fmt.Errorf("answer exchangeID mismatch for peer %s", remoteWgKey) } - - psk, err := s.initiator.Finish(a.KEMAnswer, m.binding(remoteWgKey)) - if err != nil { - return PSK{}, nil, err - } - - m.mu.Lock() + // Claim the exchange under the lock so a concurrent answer bails, and copy out + // the initiator so Finish (expensive crypto) runs without holding the lock. + init := s.initiator s.state = stateEstablished s.initiator = nil m.mu.Unlock() + psk, err := init.Finish(a.KEMAnswer, m.binding(remoteWgKey)) + if err != nil { + return PSK{}, nil, err + } + return psk, &ConfirmMsg{ExchangeID: a.ExchangeID}, nil } // HandleConfirm processes a received confirm as the responder and returns the PSK -// to commit now. It is idempotent for the current exchangeID: a retransmitted -// confirm returns the same PSK again (programming the same PSK is harmless). It -// errors on a stale/unknown confirm so the caller ignores it. +// to commit now. Retransmitted confirms are deduplicated by the driver, so this is +// called once. It errors on a stale/unknown confirm so the caller ignores it. func (m *Manager) HandleConfirm(remoteWgKey string, c *ConfirmMsg) (PSK, error) { m.mu.Lock() defer m.mu.Unlock() s, ok := m.sessions[remoteWgKey] - if !ok || (s.state != stateAwaitingConfirm && s.state != stateEstablished) { + if !ok || s.state != stateAwaitingConfirm { return PSK{}, fmt.Errorf("no pending answer for peer %s", remoteWgKey) } if s.exchangeID != c.ExchangeID { @@ -161,7 +149,9 @@ func (m *Manager) HandleConfirm(remoteWgKey string, c *ConfirmMsg) (PSK, error) } s.state = stateEstablished - return s.pendingPSK, nil + psk := s.pendingPSK + s.pendingPSK = PSK{} + return psk, nil } func (m *Manager) binding(remoteWgKey string) Binding {