diff --git a/client/internal/pqkem/manager.go b/client/internal/pqkem/manager.go new file mode 100644 index 000000000..beefe9bee --- /dev/null +++ b/client/internal/pqkem/manager.go @@ -0,0 +1,163 @@ +package pqkem + +import ( + "crypto/rand" + "fmt" + "sync" +) + +// 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 +) + +type peerSession struct { + state sessionState + exchangeID ExchangeID + + // initiator side: in-flight ephemeral keypair between offer and answer. + initiator *Initiator + // responder side: PSK derived on the offer, committed only on the confirm. + pendingPSK PSK +} + +type Manager struct { + localWgKey string + + mu sync.Mutex + sessions map[string]*peerSession +} + +// 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 { + return &Manager{ + localWgKey: localWgKey, + sessions: make(map[string]*peerSession), + } +} + +// IsInitiator reports whether the local peer drives the exchange for this remote +// peer. Roles are deterministic (lexicographic WG-key compare) so exactly one side +// initiates, mirroring how Rosenpass picks its handshake initiator. +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) { + init, err := NewInitiator() + if err != nil { + return nil, err + } + id, err := newExchangeID() + if err != nil { + return nil, 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. +func (m *Manager) HandleOffer(remoteWgKey string, o *OfferMsg) (*AnswerMsg, error) { + answer, psk, err := Respond(o.KEMOffer, m.binding(remoteWgKey)) + if err != nil { + return nil, err + } + + m.mu.Lock() + m.sessions[remoteWgKey] = &peerSession{ + state: stateAwaitingConfirm, + exchangeID: o.ExchangeID, + pendingPSK: psk, + } + m.mu.Unlock() + + return &AnswerMsg{ExchangeID: o.ExchangeID, KEMAnswer: answer}, nil +} + +// 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] + m.mu.Unlock() + + if !ok || s.state != stateAwaitingAnswer { + 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. + 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() + s.state = stateEstablished + s.initiator = nil + m.mu.Unlock() + + return psk, &ConfirmMsg{ExchangeID: a.ExchangeID}, nil +} + +// HandleConfirm processes a received confirm as the responder and returns the PSK +// to commit now. 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 { + return Binding{LocalWgPub: []byte(m.localWgKey), RemoteWgPub: []byte(remoteWgKey)} +} + +func newExchangeID() (ExchangeID, error) { + var id ExchangeID + if _, err := rand.Read(id[:]); err != nil { + return ExchangeID{}, fmt.Errorf("generate exchange id: %w", err) + } + return id, nil +} diff --git a/client/internal/pqkem/manager_test.go b/client/internal/pqkem/manager_test.go new file mode 100644 index 000000000..3bdb0f413 --- /dev/null +++ b/client/internal/pqkem/manager_test.go @@ -0,0 +1,75 @@ +package pqkem + +import ( + "testing" + + "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")) + + // B (initiator) -> A + offer, err := b.StartExchange("aaaa") + require.NoError(t, err) + + // A (responder) derives PSK but holds it pending + answer, err := a.HandleOffer("bbbb", offer) + require.NoError(t, err) + + // B derives PSK on the answer, commits now, produces confirm + pskB, confirm, err := b.HandleAnswer("aaaa", answer) + require.NoError(t, err) + + // A commits only on the confirm + pskA, err := a.HandleConfirm("bbbb", confirm) + require.NoError(t, err) + + 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") + + _, 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) +} + +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) +}