Collapse Driver and Manager in one.

- Have just one manager => one lock
 - Session state is needed in driver to => we have it available now.
 - Isomorphically align to rosenpass components and functionality

File	Role	                                  rosenpass equivalent
kem.go	primitive pure X25519MLKEM768	          crypto.go/handshake
message.go	Offer/Answer/Confirm + Encode/Decode  messages.go
manager.go	Manager stateful, single lock	      server logic
callbacks.go	WGCallbackHandler (seam output)	  Handler
Transport (interfaccia)	seam trasporto pluggable  Conn

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

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

client/internal/pqkem/driver.go

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

client/internal/pqkem/driver_test.go

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

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 {

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
 }