Manages convergence

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)
+	}
+}

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)
+}

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)
-}

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 {