diff --git a/client/internal/pqkem/convergence.go b/client/internal/pqkem/convergence.go index 3e1368ef7..63c0ba136 100644 --- a/client/internal/pqkem/convergence.go +++ b/client/internal/pqkem/convergence.go @@ -8,114 +8,108 @@ import ( // 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 { +func (m *Manager) handleOffer(remoteWgKey string, o *OfferMsg) error { + m.mu.Lock() + if ex := m.exchanges[remoteWgKey]; ex != nil && ex.id == o.ExchangeID { + state, last := ex.state, ex.lastSent + m.mu.Unlock() + if state == stateReserved { // 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). + // dropping avoids a second (randomized -> divergent) derivation. return nil } - return d.transport.Send(remoteWgKey, last) + return m.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(), - } - d.mu.Unlock() + // Reserve the slot so a concurrent duplicate offer bails. + m.exchanges[remoteWgKey] = &exchangeCtl{id: o.ExchangeID, state: stateReserved, startedAt: time.Now()} + m.mu.Unlock() - answer, err := d.mgr.HandleOffer(remoteWgKey, o) + answerBytes, psk, err := Respond(o.KEMOffer, m.binding(remoteWgKey)) if err != nil { return err } - raw, err := answer.Encode() + raw, err := (&AnswerMsg{ExchangeID: o.ExchangeID, KEMAnswer: answerBytes}).Encode() if err != nil { return err } - d.mu.Lock() - if ex := d.exchanges[remoteWgKey]; ex != nil && ex.id == o.ExchangeID { + m.mu.Lock() + if ex := m.exchanges[remoteWgKey]; ex != nil && ex.id == o.ExchangeID { + ex.state = stateAwaitingConfirm ex.lastSent = raw + ex.pendingPSK = psk } - d.mu.Unlock() + m.mu.Unlock() - return d.transport.Send(remoteWgKey, raw) + return m.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 { - // Ignore stale answers, answers to a responder-side exchange (cancel == nil), and - // duplicates once we are already in the confirm phase. A duplicate that still - // slips through is caught by the Manager, whose HandleAnswer bails under its lock. - d.mu.Lock() - ex := d.exchanges[remoteWgKey] - if ex == nil || ex.id != a.ExchangeID || ex.cancel == nil || isConfirmPhase(ex.lastSent) { - d.mu.Unlock() +// handleAnswer (initiator) derives the PSK, surfaces it, and sends the confirm, then +// switches the retransmit payload to the confirm. Only valid in stateAwaitingAnswer; +// advancing the state under the lock makes a concurrent/duplicate answer bail. +func (m *Manager) handleAnswer(remoteWgKey string, a *AnswerMsg) error { + m.mu.Lock() + ex := m.exchanges[remoteWgKey] + if ex == nil || ex.id != a.ExchangeID || ex.state != stateAwaitingAnswer { + m.mu.Unlock() return nil } - d.mu.Unlock() + ex.state = stateConfirming + init := ex.initiator + ex.initiator = nil + m.mu.Unlock() - psk, confirm, err := d.mgr.HandleAnswer(remoteWgKey, a) + psk, err := init.Finish(a.KEMAnswer, m.binding(remoteWgKey)) if err != nil { return err } - if err := d.wg.OnNewPSKReady(remoteWgKey, psk); err != nil { + if err := m.wg.OnNewPSKReady(remoteWgKey, psk); err != nil { return err } - raw, err := confirm.Encode() + raw, err := (&ConfirmMsg{ExchangeID: a.ExchangeID}).Encode() if err != nil { return err } - d.mu.Lock() - if ex := d.exchanges[remoteWgKey]; ex != nil && ex.id == a.ExchangeID { + m.mu.Lock() + if ex := m.exchanges[remoteWgKey]; ex != nil && ex.id == a.ExchangeID { ex.lastSent = raw // loop now retransmits the confirm - d.established[remoteWgKey] = true - d.failures[remoteWgKey] = 0 + m.established[remoteWgKey] = true + m.failures[remoteWgKey] = 0 _ = time.Since(ex.startedAt) // convergence latency (metrics hook, later step) } - d.mu.Unlock() + m.mu.Unlock() - return d.transport.Send(remoteWgKey, raw) + return m.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() +// handleConfirm (responder) commits the pending PSK. Only valid in +// stateAwaitingConfirm; duplicate confirms (the initiator best-effort resends it) +// find the exchange gone and are ignored. +func (m *Manager) handleConfirm(remoteWgKey string, c *ConfirmMsg) error { + m.mu.Lock() + ex := m.exchanges[remoteWgKey] + if ex == nil || ex.id != c.ExchangeID || ex.state != stateAwaitingConfirm { + m.mu.Unlock() return nil } - delete(d.exchanges, remoteWgKey) - d.established[remoteWgKey] = true - d.failures[remoteWgKey] = 0 + psk := ex.pendingPSK + delete(m.exchanges, remoteWgKey) + m.established[remoteWgKey] = true + m.failures[remoteWgKey] = 0 _ = time.Since(ex.startedAt) // convergence latency (metrics hook, later step) - d.mu.Unlock() + m.mu.Unlock() - psk, err := d.mgr.HandleConfirm(remoteWgKey, c) - if err != nil { - return err - } - return d.wg.OnNewPSKReady(remoteWgKey, psk) + return m.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) +// initiatorLoop retransmits the initiator's outstanding message, keyed off the +// exchange state: the offer while awaiting the answer (bounded by maxRetries -> +// failure), then the confirm a few best-effort times before stopping. The +// convergence deadline is thus derived from maxRetries * retryInterval. +func (m *Manager) initiatorLoop(ctx context.Context, remoteWgKey string, id ExchangeID) { + defer m.wait.Done() + t := time.NewTicker(m.retryInterval) defer t.Stop() offerAttempts, confirmsSent := 0, 0 @@ -124,37 +118,42 @@ func (d *Driver) initiatorLoop(ctx context.Context, remoteWgKey string, id Excha case <-ctx.Done(): return case <-t.C: - d.mu.Lock() - ex := d.exchanges[remoteWgKey] + m.mu.Lock() + ex := m.exchanges[remoteWgKey] if ex == nil || ex.id != id { - d.mu.Unlock() + m.mu.Unlock() return } - if !isConfirmPhase(ex.lastSent) { - if offerAttempts >= d.maxRetries { - delete(d.exchanges, remoteWgKey) - fail := d.registerFailureLocked(remoteWgKey) - d.mu.Unlock() - d.raiseFailure(remoteWgKey, fail) + switch ex.state { + case stateAwaitingAnswer: + if offerAttempts >= m.maxRetries { + delete(m.exchanges, remoteWgKey) + fail := m.registerFailureLocked(remoteWgKey) + m.mu.Unlock() + m.raiseFailure(remoteWgKey, fail) return } msg := ex.lastSent offerAttempts++ - d.mu.Unlock() - d.retransmit(remoteWgKey, msg) - continue - } + m.mu.Unlock() + m.retransmit(remoteWgKey, msg) - if confirmsSent >= confirmRetransmits { - delete(d.exchanges, remoteWgKey) - d.mu.Unlock() + case stateConfirming: + if confirmsSent >= confirmRetransmits { + delete(m.exchanges, remoteWgKey) + m.mu.Unlock() + return + } + msg := ex.lastSent + confirmsSent++ + m.mu.Unlock() + m.retransmit(remoteWgKey, msg) + + default: + m.mu.Unlock() return } - msg := ex.lastSent - confirmsSent++ - d.mu.Unlock() - d.retransmit(remoteWgKey, msg) } } } @@ -162,38 +161,31 @@ func (d *Driver) initiatorLoop(ctx context.Context, remoteWgKey string, id Excha // registerFailureLocked applies policy B and reports whether OnRekeyFailed is due: // an initial exchange (peer never established) fails immediately; a rekey tolerates // up to maxRekeyFailures consecutive misses (we stay on the still-valid previous -// PSK) before failing. Whether it is initial is read live from established. -// Assumes d.mu is held. -func (d *Driver) registerFailureLocked(remoteWgKey string) bool { - if !d.established[remoteWgKey] { +// PSK) before failing. Assumes m.mu is held. +func (m *Manager) registerFailureLocked(remoteWgKey string) bool { + if !m.established[remoteWgKey] { return true } - d.failures[remoteWgKey]++ - if d.failures[remoteWgKey] >= d.maxRekeyFailures { - d.failures[remoteWgKey] = 0 + m.failures[remoteWgKey]++ + if m.failures[remoteWgKey] >= m.maxRekeyFailures { + m.failures[remoteWgKey] = 0 return true } return false } -// isConfirmPhase reports whether the initiator's outstanding message is the confirm -// (as opposed to the offer), derived from the message's type byte. -func isConfirmPhase(lastSent []byte) bool { - return len(lastSent) > 0 && lastSent[0] == byte(MsgConfirm) -} - -func (d *Driver) raiseFailure(remoteWgKey string, fail bool) { +func (m *Manager) raiseFailure(remoteWgKey string, fail bool) { if !fail { - d.logger.Warn("pqkem rekey attempt timed out, will retry next cycle", "peer", remoteWgKey) + m.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) + if err := m.wg.OnRekeyFailed(remoteWgKey); err != nil { + m.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) +func (m *Manager) retransmit(remoteWgKey string, msg []byte) { + if err := m.transport.Send(remoteWgKey, msg); err != nil { + m.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 index e05331210..5bdd4ec20 100644 --- a/client/internal/pqkem/convergence_test.go +++ b/client/internal/pqkem/convergence_test.go @@ -15,7 +15,7 @@ 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 + peer *Manager drop atomic.Bool } @@ -27,9 +27,9 @@ func (g *gate) Send(remote string, msg []byte) error { return g.peer.HandleInbound(g.local, cp) } -func TestDriver_InitialTimeoutFailsImmediately(t *testing.T) { +func TestManager_InitialTimeoutFailsImmediately(t *testing.T) { wg := newFakeWG() - d := NewDriver("bbbb", dropTransport{}, wg, time.Hour, nil) // bbbb > aaaa -> initiator + d := NewManager("bbbb", dropTransport{}, wg, time.Hour, nil) // bbbb > aaaa -> initiator d.retryInterval = 5 * time.Millisecond d.maxRetries = 3 d.AddPeer("aaaa") @@ -45,14 +45,14 @@ func TestDriver_InitialTimeoutFailsImmediately(t *testing.T) { }, time.Second, 5*time.Millisecond) } -func TestDriver_RekeyToleratesKFailures(t *testing.T) { +func TestManager_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) + dA := NewManager("aaaa", gA, wgA, time.Hour, nil) + dB := NewManager("bbbb", gB, wgB, time.Hour, nil) gA.peer = dB gB.peer = dA dB.retryInterval = 5 * time.Millisecond diff --git a/client/internal/pqkem/driver.go b/client/internal/pqkem/driver.go deleted file mode 100644 index b4b6c4514..000000000 --- a/client/internal/pqkem/driver.go +++ /dev/null @@ -1,214 +0,0 @@ -package pqkem - -import ( - "context" - "fmt" - "log/slog" - "sync" - "time" -) - -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 the initiator retransmits its outstanding - // message (offer, then confirm) while an exchange is in flight. - DefaultRetryInterval = 2 * time.Second - // 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 times the initiator best-effort resends the - // confirm after converging, to cover its loss without a dedicated goroutine. - confirmRetransmits = 3 -) - -// Transport hands an already-encoded exchange message to the peer. The host routes -// it over the appropriate channel — Signal before the tunnel is up, the WireGuard -// tunnel for rekeys — so the driver never needs to know which is in use. -type Transport interface { - Send(remoteWgKey string, msg []byte) error -} - -// exchangeCtl tracks one in-flight exchange for a peer. A single lastSent holds the -// message currently being (re)transmitted: for the initiator the offer and then, -// once answered, the confirm; for the responder the answer, resent on a duplicate -// offer. Its type byte (lastSent[0]) also encodes the initiator's phase, so no -// separate "answered" flag is kept. Only the initiator runs a retransmit loop -// (cancel != nil); the responder is purely reactive. Whether an exchange is the -// peer's first is read live from d.established, not snapshotted here. -type exchangeCtl struct { - id ExchangeID - startedAt time.Time - lastSent []byte - cancel context.CancelFunc -} - -// Driver ties the pure Manager to the outside world: per-peer rekey timer, inbound -// dispatch, retransmit/convergence tracking, and surfacing events to the host via -// WGCallbackHandler. -type Driver struct { - mgr *Manager - transport Transport - wg WGCallbackHandler - logger *slog.Logger - - rekeyInterval time.Duration - retryInterval time.Duration - maxRetries int - maxRekeyFailures int - - rootCtx context.Context - rootCancel context.CancelFunc - - mu sync.Mutex - peers map[string]context.CancelFunc // per-peer rekey loop - exchanges map[string]*exchangeCtl // in-flight exchange per peer - established map[string]bool // peer has completed at least one exchange - failures map[string]int // consecutive rekey failures per peer - wait sync.WaitGroup -} - -// NewDriver builds a driver for the local peer. A zero interval falls back to -// 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 { - interval = DefaultRekeyInterval - } - if logger == nil { - logger = slog.Default() - } - ctx, cancel := context.WithCancel(context.Background()) - return &Driver{ - 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), - } -} - -// AddPeer registers a remote peer and starts its rekey timer. Re-adding is a no-op. -func (d *Driver) AddPeer(remoteWgKey string) { - d.mu.Lock() - defer d.mu.Unlock() - if _, ok := d.peers[remoteWgKey]; ok { - return - } - ctx, cancel := context.WithCancel(d.rootCtx) - d.peers[remoteWgKey] = cancel - d.wait.Add(1) - go d.rekeyLoop(ctx, remoteWgKey) -} - -// RemovePeer stops a peer's rekey timer and any in-flight exchange, and drops state. -func (d *Driver) RemovePeer(remoteWgKey string) { - d.mu.Lock() - if cancel, ok := d.peers[remoteWgKey]; ok { - cancel() - delete(d.peers, remoteWgKey) - } - if ex, ok := d.exchanges[remoteWgKey]; ok { - if ex.cancel != nil { - ex.cancel() - } - delete(d.exchanges, remoteWgKey) - } - delete(d.established, remoteWgKey) - delete(d.failures, remoteWgKey) - d.mu.Unlock() -} - -// Stop cancels all timers and in-flight exchanges and waits for goroutines to exit. -func (d *Driver) Stop() { - d.rootCancel() - d.wait.Wait() - d.mu.Lock() - d.peers = make(map[string]context.CancelFunc) - d.exchanges = make(map[string]*exchangeCtl) - d.mu.Unlock() -} - -// HandleInbound decodes an incoming message and drives the exchange, sending any -// response via the transport and surfacing derived PSKs / convergence to the host. -func (d *Driver) HandleInbound(remoteWgKey string, raw []byte) error { - typ, msg, err := Decode(raw) - if err != nil { - return fmt.Errorf("decode from %s: %w", remoteWgKey, err) - } - switch typ { - case MsgOffer: - return d.handleOffer(remoteWgKey, msg.(*OfferMsg)) - case MsgAnswer: - return d.handleAnswer(remoteWgKey, msg.(*AnswerMsg)) - case MsgConfirm: - return d.handleConfirm(remoteWgKey, msg.(*ConfirmMsg)) - default: - return fmt.Errorf("unhandled message type %d from %s", typ, remoteWgKey) - } -} - -// 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. -func (d *Driver) initiateRekey(remoteWgKey string) error { - if !d.mgr.IsInitiator(remoteWgKey) { - return nil - } - offer, err := d.mgr.StartExchange(remoteWgKey) - if err != nil { - return err - } - raw, err := offer.Encode() - if err != nil { - return err - } - - 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(), - lastSent: raw, - cancel: cancel, - } - d.mu.Unlock() - - d.wait.Add(1) - go d.initiatorLoop(ctx, remoteWgKey, offer.ExchangeID) - - return d.transport.Send(remoteWgKey, raw) -} - -func (d *Driver) rekeyLoop(ctx context.Context, remoteWgKey string) { - defer d.wait.Done() - t := time.NewTicker(d.rekeyInterval) - defer t.Stop() - for { - select { - case <-ctx.Done(): - return - case <-t.C: - if err := d.initiateRekey(remoteWgKey); err != nil { - d.logger.Error("pqkem rekey failed to start", "peer", remoteWgKey, "err", err) - } - } - } -} diff --git a/client/internal/pqkem/driver_test.go b/client/internal/pqkem/driver_test.go deleted file mode 100644 index 59fb381da..000000000 --- a/client/internal/pqkem/driver_test.go +++ /dev/null @@ -1,96 +0,0 @@ -package pqkem - -import ( - "sync" - "testing" - "time" - - "github.com/stretchr/testify/require" -) - -// loopback delivers a sent message synchronously to the peer driver's HandleInbound, -// attributing it to localKey (the sender). -type loopback struct { - localKey string - peer *Driver -} - -func (l *loopback) Send(remoteWgKey string, msg []byte) error { - cp := append([]byte(nil), msg...) - return l.peer.HandleInbound(l.localKey, cp) -} - -type fakeWG struct { - mu sync.Mutex - psks map[string]PSK - failed []string -} - -func newFakeWG() *fakeWG { return &fakeWG{psks: map[string]PSK{}} } - -func (f *fakeWG) OnNewPSKReady(remoteWgKey string, psk PSK) error { - f.mu.Lock() - defer f.mu.Unlock() - f.psks[remoteWgKey] = psk - return nil -} - -func (f *fakeWG) OnRekeyFailed(remoteWgKey string) error { - f.mu.Lock() - defer f.mu.Unlock() - f.failed = append(f.failed, remoteWgKey) - return nil -} - -func (f *fakeWG) psk(peer string) PSK { - f.mu.Lock() - defer f.mu.Unlock() - return f.psks[peer] -} - -func TestDriver_ExchangeConverges(t *testing.T) { - lbA := &loopback{localKey: "aaaa"} - lbB := &loopback{localKey: "bbbb"} - wgA := newFakeWG() - wgB := newFakeWG() - - // long interval so the ticker never fires during the test; we drive manually. - dA := NewDriver("aaaa", lbA, wgA, time.Hour, nil) - dB := NewDriver("bbbb", lbB, wgB, time.Hour, nil) - lbA.peer = dB // A sends -> B receives - lbB.peer = dA // B sends -> A receives - - dA.AddPeer("bbbb") - dB.AddPeer("aaaa") - defer dA.Stop() - defer dB.Stop() - - // B is the initiator ("bbbb" > "aaaa"). - require.NoError(t, dB.initiateRekey("aaaa")) - - pskB := wgB.psk("aaaa") // B committed on the answer - pskA := wgA.psk("bbbb") // A committed on the confirm - require.NotEqual(t, PSK{}, pskA, "responder A must have a PSK") - require.NotEqual(t, PSK{}, pskB, "initiator B must have a PSK") - require.Equal(t, pskB, pskA, "both sides converge on the same PSK") -} - -func TestDriver_NonInitiatorDoesNothing(t *testing.T) { - lbA := &loopback{localKey: "aaaa"} - wgA := newFakeWG() - dA := NewDriver("aaaa", lbA, wgA, time.Hour, nil) - // no peer driver wired; if A wrongly initiated, Send would nil-panic. - dA.AddPeer("bbbb") - defer dA.Stop() - - // A is NOT the initiator vs "bbbb" -> initiateRekey is a no-op, no Send. - require.NoError(t, dA.initiateRekey("bbbb")) - require.Equal(t, PSK{}, wgA.psk("bbbb")) -} - -func TestDriver_StopIsIdempotent(t *testing.T) { - dA := NewDriver("aaaa", &loopback{localKey: "aaaa"}, newFakeWG(), time.Hour, nil) - dA.AddPeer("bbbb") - dA.Stop() - dA.Stop() // must not panic or hang -} diff --git a/client/internal/pqkem/manager.go b/client/internal/pqkem/manager.go index cc746c423..cd48866f0 100644 --- a/client/internal/pqkem/manager.go +++ b/client/internal/pqkem/manager.go @@ -1,54 +1,123 @@ package pqkem import ( + "context" "crypto/rand" "fmt" + "log/slog" "sync" + "time" ) -// Manager drives the PQ-KEM exchange per remote peer. It is a pure state machine: -// it consumes and produces messages and returns the derived PSK at the point each -// side must commit it (initiator on the answer, responder on the confirm). It does -// no I/O — transport (Signal / tunnel) and the WireGuard SetPSK live in the caller -// (network module + wiring), so this stays extraction-ready as a standalone library. -// -// Commit ordering (see ConfirmMsg): the responder derives the PSK on the offer but -// must NOT commit it until the confirm arrives, otherwise a lost answer would leave -// it on a new PSK the initiator never got, breaking the tunnel at the next rekey. - -type sessionState uint8 - const ( - stateIdle sessionState = iota - stateAwaitingAnswer // initiator: sent offer, waiting for answer - stateAwaitingConfirm // responder: sent answer, holding pending PSK until confirm - stateEstablished + // 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 the initiator retransmits its outstanding + // message (offer, then confirm) while an exchange is in flight. + DefaultRetryInterval = 2 * time.Second + // 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 times the initiator best-effort resends the + // confirm after converging, to cover its loss without a dedicated goroutine. + confirmRetransmits = 3 ) -type peerSession struct { - state sessionState - exchangeID ExchangeID +// Transport hands an already-encoded exchange message to the peer. The host routes +// it over the appropriate channel — Signal before the tunnel is up, the WireGuard +// tunnel for rekeys — so the Manager never needs to know which is in use. It is the +// analogue of go-rosenpass's Conn seam. +type Transport interface { + Send(remoteWgKey string, msg []byte) error +} - // initiator side: in-flight ephemeral keypair between offer and answer. - initiator *Initiator - // responder side: PSK derived on the offer, committed only on the confirm. +// exchangeState is the single source of truth for an exchange's role and phase. +// Every handler and the retransmit loop key off it, so no role/phase is re-derived +// from other fields. +type exchangeState uint8 + +const ( + stateReserved exchangeState = iota // responder: deriving the answer + stateAwaitingAnswer // initiator: offer sent, awaiting the answer + stateAwaitingConfirm // responder: answer sent, awaiting the confirm + stateConfirming // initiator: answer in, PSK committed, flushing the confirm +) + +// exchangeCtl holds all state for one in-flight exchange with a peer, under the +// Manager's single lock. state drives every decision. lastSent is the message +// currently being (re)transmitted. The crypto payloads live here too: initiator is +// the ephemeral handle used at Finish (initiator side); pendingPSK is the derived +// key held until the confirm commits it (responder side). Only the initiator runs a +// retransmit loop, so only it sets cancel. +type exchangeCtl struct { + id ExchangeID + state exchangeState + startedAt time.Time + cancel context.CancelFunc + lastSent []byte + initiator *Initiator pendingPSK PSK } +// Manager is the stateful orchestrator — the analogue of go-rosenpass's Server. It +// runs the per-peer rekey timer, drives the X25519MLKEM768 exchange over a pluggable +// Transport, and surfaces the derived PSK to the host via WGCallbackHandler. The +// cryptography is the pure kem.go primitives; all per-exchange and per-peer state +// lives here under one lock. type Manager struct { localWgKey string + transport Transport + wg WGCallbackHandler + logger *slog.Logger - mu sync.Mutex - sessions map[string]*peerSession + rekeyInterval time.Duration + retryInterval time.Duration + maxRetries int + maxRekeyFailures int + + rootCtx context.Context + rootCancel context.CancelFunc + + mu sync.Mutex + peers map[string]context.CancelFunc // per-peer rekey loop + exchanges map[string]*exchangeCtl // in-flight exchange per peer + established map[string]bool // peer has completed at least one exchange + failures map[string]int // consecutive rekey failures per peer + wait sync.WaitGroup } -// NewManager creates a KEM manager for the local peer identified by its WireGuard -// public key (base64), used both for the deterministic initiator role and the -// identity binding of the derived PSK. -func NewManager(localWgKey string) *Manager { +// NewManager builds a manager for the local peer identified by its WireGuard public +// key (used for the deterministic initiator role and the identity binding). A zero +// interval falls back to DefaultRekeyInterval; a nil logger to slog.Default(). +// Retry/retries/K use their defaults and can be overridden before use. +func NewManager(localWgKey string, t Transport, h WGCallbackHandler, interval time.Duration, logger *slog.Logger) *Manager { + if interval <= 0 { + interval = DefaultRekeyInterval + } + if logger == nil { + logger = slog.Default() + } + ctx, cancel := context.WithCancel(context.Background()) return &Manager{ - localWgKey: localWgKey, - sessions: make(map[string]*peerSession), + localWgKey: 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), } } @@ -59,99 +128,121 @@ func (m *Manager) IsInitiator(remoteWgKey string) bool { return m.localWgKey > remoteWgKey } -// StartExchange begins (or restarts, on rotation) an exchange as the initiator and -// returns the offer to send. Any previous in-flight exchange for the peer is dropped. -func (m *Manager) StartExchange(remoteWgKey string) (*OfferMsg, error) { +// AddPeer registers a remote peer and starts its rekey timer. Re-adding is a no-op. +func (m *Manager) AddPeer(remoteWgKey string) { + m.mu.Lock() + defer m.mu.Unlock() + if _, ok := m.peers[remoteWgKey]; ok { + return + } + ctx, cancel := context.WithCancel(m.rootCtx) + m.peers[remoteWgKey] = cancel + m.wait.Add(1) + go m.rekeyLoop(ctx, remoteWgKey) +} + +// RemovePeer stops a peer's rekey timer and any in-flight exchange, and drops state. +func (m *Manager) RemovePeer(remoteWgKey string) { + m.mu.Lock() + if cancel, ok := m.peers[remoteWgKey]; ok { + cancel() + delete(m.peers, remoteWgKey) + } + if ex, ok := m.exchanges[remoteWgKey]; ok { + if ex.cancel != nil { + ex.cancel() + } + delete(m.exchanges, remoteWgKey) + } + delete(m.established, remoteWgKey) + delete(m.failures, remoteWgKey) + m.mu.Unlock() +} + +// Stop cancels all timers and in-flight exchanges and waits for goroutines to exit. +func (m *Manager) Stop() { + m.rootCancel() + m.wait.Wait() + m.mu.Lock() + m.peers = make(map[string]context.CancelFunc) + m.exchanges = make(map[string]*exchangeCtl) + m.mu.Unlock() +} + +// HandleInbound decodes an incoming message and drives the exchange, sending any +// response via the transport and surfacing derived PSKs / convergence to the host. +func (m *Manager) HandleInbound(remoteWgKey string, raw []byte) error { + typ, msg, err := Decode(raw) + if err != nil { + return fmt.Errorf("decode from %s: %w", remoteWgKey, err) + } + switch typ { + case MsgOffer: + return m.handleOffer(remoteWgKey, msg.(*OfferMsg)) + case MsgAnswer: + return m.handleAnswer(remoteWgKey, msg.(*AnswerMsg)) + case MsgConfirm: + return m.handleConfirm(remoteWgKey, msg.(*ConfirmMsg)) + default: + return fmt.Errorf("unhandled message type %d from %s", typ, remoteWgKey) + } +} + +// 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. +func (m *Manager) initiateRekey(remoteWgKey string) error { + if !m.IsInitiator(remoteWgKey) { + return nil + } init, err := NewInitiator() if err != nil { - return nil, err + return err } id, err := newExchangeID() if err != nil { - return nil, err + return err } - - m.mu.Lock() - m.sessions[remoteWgKey] = &peerSession{ - state: stateAwaitingAnswer, - exchangeID: id, - initiator: init, - } - m.mu.Unlock() - - return &OfferMsg{ExchangeID: id, KEMOffer: init.Offer()}, nil -} - -// 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. 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) { - answer, psk, err := Respond(o.KEMOffer, m.binding(remoteWgKey)) + raw, err := (&OfferMsg{ExchangeID: id, KEMOffer: init.Offer()}).Encode() if err != nil { - return nil, err + return err } + ctx, cancel := context.WithCancel(m.rootCtx) m.mu.Lock() - m.sessions[remoteWgKey] = &peerSession{ - state: stateAwaitingConfirm, - exchangeID: o.ExchangeID, - pendingPSK: psk, + if old := m.exchanges[remoteWgKey]; old != nil && old.cancel != nil { + old.cancel() + } + m.exchanges[remoteWgKey] = &exchangeCtl{ + id: id, + state: stateAwaitingAnswer, + startedAt: time.Now(), + cancel: cancel, + lastSent: raw, + initiator: init, } m.mu.Unlock() - return &AnswerMsg{ExchangeID: o.ExchangeID, KEMAnswer: answer}, nil + m.wait.Add(1) + go m.initiatorLoop(ctx, remoteWgKey, id) + + return m.transport.Send(remoteWgKey, raw) } -// HandleAnswer processes a received answer as the initiator. On success it returns -// the derived PSK (which the caller commits now) and the confirm to send. -func (m *Manager) HandleAnswer(remoteWgKey string, a *AnswerMsg) (PSK, *ConfirmMsg, error) { - m.mu.Lock() - s, ok := m.sessions[remoteWgKey] - if !ok || s.state != stateAwaitingAnswer { - m.mu.Unlock() - return PSK{}, nil, fmt.Errorf("no pending offer for peer %s", remoteWgKey) +func (m *Manager) rekeyLoop(ctx context.Context, remoteWgKey string) { + defer m.wait.Done() + t := time.NewTicker(m.rekeyInterval) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + if err := m.initiateRekey(remoteWgKey); err != nil { + m.logger.Error("pqkem rekey failed to start", "peer", remoteWgKey, "err", err) + } + } } - 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) - } - // 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. 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 { - return PSK{}, fmt.Errorf("no pending answer for peer %s", remoteWgKey) - } - if s.exchangeID != c.ExchangeID { - return PSK{}, fmt.Errorf("confirm exchangeID mismatch for peer %s", remoteWgKey) - } - - s.state = stateEstablished - psk := s.pendingPSK - s.pendingPSK = PSK{} - return psk, nil } func (m *Manager) binding(remoteWgKey string) Binding { diff --git a/client/internal/pqkem/manager_test.go b/client/internal/pqkem/manager_test.go index 3bdb0f413..e355b17fa 100644 --- a/client/internal/pqkem/manager_test.go +++ b/client/internal/pqkem/manager_test.go @@ -1,75 +1,96 @@ package pqkem import ( + "sync" "testing" + "time" "github.com/stretchr/testify/require" ) -func TestManager_FullExchange(t *testing.T) { - // deterministic roles: "bbbb" > "aaaa" -> B is initiator - a := NewManager("aaaa") - b := NewManager("bbbb") - require.True(t, b.IsInitiator("aaaa")) - require.False(t, a.IsInitiator("bbbb")) +// loopback delivers a sent message synchronously to the peer driver's HandleInbound, +// attributing it to localKey (the sender). +type loopback struct { + localKey string + peer *Manager +} - // B (initiator) -> A - offer, err := b.StartExchange("aaaa") - require.NoError(t, err) +func (l *loopback) Send(remoteWgKey string, msg []byte) error { + cp := append([]byte(nil), msg...) + return l.peer.HandleInbound(l.localKey, cp) +} - // A (responder) derives PSK but holds it pending - answer, err := a.HandleOffer("bbbb", offer) - require.NoError(t, err) +type fakeWG struct { + mu sync.Mutex + psks map[string]PSK + failed []string +} - // B derives PSK on the answer, commits now, produces confirm - pskB, confirm, err := b.HandleAnswer("aaaa", answer) - require.NoError(t, err) +func newFakeWG() *fakeWG { return &fakeWG{psks: map[string]PSK{}} } - // A commits only on the confirm - pskA, err := a.HandleConfirm("bbbb", confirm) - require.NoError(t, err) +func (f *fakeWG) OnNewPSKReady(remoteWgKey string, psk PSK) error { + f.mu.Lock() + defer f.mu.Unlock() + f.psks[remoteWgKey] = psk + return nil +} +func (f *fakeWG) OnRekeyFailed(remoteWgKey string) error { + f.mu.Lock() + defer f.mu.Unlock() + f.failed = append(f.failed, remoteWgKey) + return nil +} + +func (f *fakeWG) psk(peer string) PSK { + f.mu.Lock() + defer f.mu.Unlock() + return f.psks[peer] +} + +func TestManager_ExchangeConverges(t *testing.T) { + lbA := &loopback{localKey: "aaaa"} + lbB := &loopback{localKey: "bbbb"} + wgA := newFakeWG() + wgB := newFakeWG() + + // long interval so the ticker never fires during the test; we drive manually. + dA := NewManager("aaaa", lbA, wgA, time.Hour, nil) + dB := NewManager("bbbb", lbB, wgB, time.Hour, nil) + lbA.peer = dB // A sends -> B receives + lbB.peer = dA // B sends -> A receives + + dA.AddPeer("bbbb") + dB.AddPeer("aaaa") + defer dA.Stop() + defer dB.Stop() + + // B is the initiator ("bbbb" > "aaaa"). + require.NoError(t, dB.initiateRekey("aaaa")) + + pskB := wgB.psk("aaaa") // B committed on the answer + pskA := wgA.psk("bbbb") // A committed on the confirm + require.NotEqual(t, PSK{}, pskA, "responder A must have a PSK") + require.NotEqual(t, PSK{}, pskB, "initiator B must have a PSK") require.Equal(t, pskB, pskA, "both sides converge on the same PSK") - require.NotEqual(t, PSK{}, pskA) } -func TestManager_StaleAnswerDropped(t *testing.T) { - b := NewManager("bbbb") +func TestManager_NonInitiatorDoesNothing(t *testing.T) { + lbA := &loopback{localKey: "aaaa"} + wgA := newFakeWG() + dA := NewManager("aaaa", lbA, wgA, time.Hour, nil) + // no peer driver wired; if A wrongly initiated, Send would nil-panic. + dA.AddPeer("bbbb") + defer dA.Stop() - _, err := b.StartExchange("aaaa") - require.NoError(t, err) - - // craft an answer with a wrong exchangeID - stale := &AnswerMsg{ExchangeID: ExchangeID{0xFF}, KEMAnswer: make([]byte, AnswerSize)} - _, _, err = b.HandleAnswer("aaaa", stale) - require.Error(t, err) + // A is NOT the initiator vs "bbbb" -> initiateRekey is a no-op, no Send. + require.NoError(t, dA.initiateRekey("bbbb")) + require.Equal(t, PSK{}, wgA.psk("bbbb")) } -func TestManager_RestartResyncsViaNewExchange(t *testing.T) { - // A restarts (fresh manager) mid-flight; a new exchange from scratch converges. - a := NewManager("aaaa") - b := NewManager("bbbb") - - off1, err := b.StartExchange("aaaa") - require.NoError(t, err) - - // B "restarts": new manager, old in-flight state gone - b = NewManager("bbbb") - off2, err := b.StartExchange("aaaa") - require.NoError(t, err) - require.NotEqual(t, off1.ExchangeID, off2.ExchangeID) - - ans, err := a.HandleOffer("bbbb", off2) - require.NoError(t, err) - pskB, conf, err := b.HandleAnswer("aaaa", ans) - require.NoError(t, err) - pskA, err := a.HandleConfirm("bbbb", conf) - require.NoError(t, err) - require.Equal(t, pskB, pskA) -} - -func TestManager_UnknownPeerErrors(t *testing.T) { - a := NewManager("aaaa") - _, err := a.HandleConfirm("zzzz", &ConfirmMsg{}) - require.Error(t, err) +func TestManager_StopIsIdempotent(t *testing.T) { + dA := NewManager("aaaa", &loopback{localKey: "aaaa"}, newFakeWG(), time.Hour, nil) + dA.AddPeer("bbbb") + dA.Stop() + dA.Stop() // must not panic or hang }