Removes confirm. Uses next offer to deliver confirmation/ack of previous round

We clock the next Offer initiation to the OnDataPathRekeyed, so we have 2 minutes
ahead of us to do our attempts and stuff before to give up.
On failure, we will know because we will not receive a new answer.. but more importantly
the wg handshake will fail :D
This commit is contained in:
riccardom
2026-08-03 12:48:46 +02:00
parent 2f76c9a013
commit a5be4c61c8
6 changed files with 179 additions and 258 deletions
+39 -42
View File
@@ -5,11 +5,11 @@ import (
"time" "time"
) )
// startExchange creates a fresh initiator exchange and returns the framed offer for // startExchange creates a fresh initiator exchange (acknowledging ackID, zero for a
// the caller to send (pushed over the data path for a rekey, or handed to the host // bootstrap) and returns the framed offer for the caller to send — pushed over the
// for the signalling channel when viaSignal is set). Any previous in-flight exchange // data path for a chained rekey, or handed to the host for signalling when viaSignal
// for the peer is cancelled. // is set. Any previous in-flight exchange for the peer is cancelled.
func (m *Manager) startExchange(remoteID string, viaSignal bool) ([]byte, error) { func (m *Manager) startExchange(remoteID string, viaSignal bool, ackID ExchangeID) ([]byte, error) {
init, err := NewInitiator() init, err := NewInitiator()
if err != nil { if err != nil {
return nil, err return nil, err
@@ -18,7 +18,7 @@ func (m *Manager) startExchange(remoteID string, viaSignal bool) ([]byte, error)
if err != nil { if err != nil {
return nil, err return nil, err
} }
raw, err := (&OfferMsg{ExchangeID: id, KEMOffer: init.Offer()}).Encode() raw, err := (&OfferMsg{ExchangeID: id, AckID: ackID, KEMOffer: init.Offer()}).Encode()
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -44,10 +44,15 @@ func (m *Manager) startExchange(remoteID string, viaSignal bool) ([]byte, error)
return raw, nil return raw, nil
} }
// processOffer (responder) derives the PSK, commits it optimistically, and returns // processOffer (responder) first acknowledges the previous exchange the offer names
// the framed answer for the caller to send. A duplicate offer (same exchangeID) // (that offer riding the data path under the freshly adopted key proves it worked),
// returns the cached answer without re-deriving; a still-reserved slot returns nil. // then derives the PSK for the new offer, commits it optimistically, and returns the
// framed answer. A duplicate offer returns the cached answer without re-deriving.
func (m *Manager) processOffer(remoteID string, o *OfferMsg) ([]byte, error) { func (m *Manager) processOffer(remoteID string, o *OfferMsg) ([]byte, error) {
if o.AckID != (ExchangeID{}) {
m.ackConverged(remoteID, o.AckID)
}
m.mu.Lock() m.mu.Lock()
if ex := m.exchanges[remoteID]; ex != nil && ex.id == o.ExchangeID { if ex := m.exchanges[remoteID]; ex != nil && ex.id == o.ExchangeID {
state, last := ex.state, ex.lastSent state, last := ex.state, ex.lastSent
@@ -76,7 +81,7 @@ func (m *Manager) processOffer(remoteID string, o *OfferMsg) ([]byte, error) {
m.mu.Unlock() m.mu.Unlock()
return nil, nil return nil, nil
} }
ex.state = stateAwaitingConfirm ex.state = stateAwaitingAck
ex.lastSent = raw ex.lastSent = raw
ex.pendingPSK = psk ex.pendingPSK = psk
m.mu.Unlock() m.mu.Unlock()
@@ -88,9 +93,9 @@ func (m *Manager) processOffer(remoteID string, o *OfferMsg) ([]byte, error) {
return raw, nil return raw, nil
} }
// processAnswer (initiator) derives and commits the PSK, then parks in // processAnswer (initiator) derives and commits the PSK and parks in
// stateAwaitingRekey; the confirm is sent later, over the data path, from // stateAwaitingRekey; the next offer (chained from OnDataPathRekeyed) will acknowledge
// OnDataPathRekeyed. Only valid in stateAwaitingAnswer; advancing the state under the // this exchange. Only valid in stateAwaitingAnswer; advancing the state under the
// lock makes a concurrent/duplicate answer bail. // lock makes a concurrent/duplicate answer bail.
func (m *Manager) processAnswer(remoteID string, a *AnswerMsg) error { func (m *Manager) processAnswer(remoteID string, a *AnswerMsg) error {
m.mu.Lock() m.mu.Lock()
@@ -108,40 +113,45 @@ func (m *Manager) processAnswer(remoteID string, a *AnswerMsg) error {
if err != nil { if err != nil {
return err return err
} }
// The initiator has converged: the responder must have derived the key to answer.
m.mu.Lock()
m.established[remoteID] = true
m.failures[remoteID] = 0
m.mu.Unlock()
return m.cbHandler.OnNewPSKReady(remoteID, psk) return m.cbHandler.OnNewPSKReady(remoteID, psk)
} }
// processConfirm (responder) records convergence. The PSK was already committed in // ackConverged (responder) records convergence of the exchange named by ackID: a
// processOffer; a confirm arriving over the data path with a matching exchangeID // later offer acknowledging it proves both sides operate on that exchange's key. Only
// proves we operate on the new key from this exchange. Stale/duplicate confirms find // acts on a matching stateAwaitingAck exchange; anything else is ignored.
// the exchange gone and are ignored. func (m *Manager) ackConverged(remoteID string, ackID ExchangeID) {
func (m *Manager) processConfirm(remoteID string, c *ConfirmMsg) error {
m.mu.Lock() m.mu.Lock()
ex := m.exchanges[remoteID] ex := m.exchanges[remoteID]
if ex == nil || ex.id != c.ExchangeID || ex.state != stateAwaitingConfirm { if ex == nil || ex.id != ackID || ex.state != stateAwaitingAck {
m.mu.Unlock() m.mu.Unlock()
return nil return
} }
delete(m.exchanges, remoteID) delete(m.exchanges, remoteID)
m.established[remoteID] = true m.established[remoteID] = true
m.failures[remoteID] = 0 m.failures[remoteID] = 0
_ = time.Since(ex.startedAt) // convergence latency (metrics hook, later step) _ = time.Since(ex.startedAt) // convergence latency (metrics hook, later step)
m.mu.Unlock() m.mu.Unlock()
return nil
} }
// initiatorLoop enforces the convergence deadline and retransmits the initiator's // initiatorLoop enforces the convergence deadline and retransmits the initiator's
// outstanding DATA-PATH message: it resends the offer while awaiting the answer only // outstanding data-path offer while awaiting the answer (a signalling-bootstrapped
// when the offer went over the data path (a signalling-bootstrapped offer is // offer is retransmitted by the host, so it is not resent here). It then waits,
// retransmitted by the host with its own negotiation); it waits, counting toward the // counting toward the deadline, in stateAwaitingRekey until OnDataPathRekeyed chains
// deadline, while awaiting the data-path rekey; it resends the confirm a few times // the next exchange (which supersedes and cancels this loop). Exhausting the deadline
// once sent. Exhausting the deadline before the confirm phase is a failure. // is a failure.
func (m *Manager) initiatorLoop(ctx context.Context, remoteID string, id ExchangeID) { func (m *Manager) initiatorLoop(ctx context.Context, remoteID string, id ExchangeID) {
defer m.wait.Done() defer m.wait.Done()
t := time.NewTicker(m.retryInterval) t := time.NewTicker(m.retryInterval)
defer t.Stop() defer t.Stop()
attempts, confirmsSent := 0, 0 attempts := 0
for { for {
select { select {
case <-ctx.Done(): case <-ctx.Done():
@@ -174,8 +184,8 @@ func (m *Manager) initiatorLoop(ctx context.Context, remoteID string, id Exchang
} }
case stateAwaitingRekey: case stateAwaitingRekey:
// Waiting for OnDataPathRekeyed; nothing to send, but the deadline // Waiting for OnDataPathRekeyed to chain the next exchange; the
// still applies (the data path may never come up with the new key). // deadline still applies (the data path may never adopt the new key).
if attempts >= m.maxRetries { if attempts >= m.maxRetries {
delete(m.exchanges, remoteID) delete(m.exchanges, remoteID)
fail := m.registerFailureLocked(remoteID) fail := m.registerFailureLocked(remoteID)
@@ -186,19 +196,6 @@ func (m *Manager) initiatorLoop(ctx context.Context, remoteID string, id Exchang
attempts++ attempts++
m.mu.Unlock() m.mu.Unlock()
case stateConfirming:
if confirmsSent >= confirmRetransmits {
delete(m.exchanges, remoteID)
m.mu.Unlock()
return
}
msg := ex.lastSent
confirmsSent++
m.mu.Unlock()
if err := m.pushDataPath(remoteID, msg); err != nil {
m.logger.Warn("pqkem confirm retransmit failed", "peer", remoteID, "err", err)
}
default: default:
m.mu.Unlock() m.mu.Unlock()
return return
+9 -13
View File
@@ -30,10 +30,9 @@ func (g *gate) SendDataPath(remoteID string, msg []byte) error {
func TestManager_InitialTimeoutFailsImmediately(t *testing.T) { func TestManager_InitialTimeoutFailsImmediately(t *testing.T) {
wg := newFakeWG() wg := newFakeWG()
d := NewManager("bbbb", dropTransport{}, wg, time.Hour, nil) // bbbb > aaaa -> initiator d := NewManager("bbbb", dropTransport{}, wg, nil) // bbbb > aaaa -> initiator
d.retryInterval = 5 * time.Millisecond d.retryInterval = 5 * time.Millisecond
d.maxRetries = 3 d.maxRetries = 3
d.AddPeer("aaaa")
defer d.Stop() defer d.Stop()
// Bootstrap offer is produced for signalling; no answer ever comes back -> the // Bootstrap offer is produced for signalling; no answer ever comes back -> the
@@ -55,41 +54,38 @@ func TestManager_RekeyToleratesKFailures(t *testing.T) {
wgA := newFakeWG() wgA := newFakeWG()
wgB := newFakeWG() wgB := newFakeWG()
dA := NewManager("aaaa", gA, wgA, time.Hour, nil) dA := NewManager("aaaa", gA, wgA, nil)
dB := NewManager("bbbb", gB, wgB, time.Hour, nil) dB := NewManager("bbbb", gB, wgB, nil)
gA.peer = dB gA.peer = dB
gB.peer = dA gB.peer = dA
dB.retryInterval = 5 * time.Millisecond dB.retryInterval = 5 * time.Millisecond
dB.maxRetries = 2 dB.maxRetries = 2
dA.AddPeer("bbbb")
dB.AddPeer("aaaa")
defer dA.Stop() defer dA.Stop()
defer dB.Stop() defer dB.Stop()
// Establish via a signalling bootstrap + data-path-rekeyed, so B becomes // Bootstrap over signalling -> B becomes established.
// established and its data path is up.
offer, err := dB.SignalOffer("aaaa") offer, err := dB.SignalOffer("aaaa")
require.NoError(t, err) require.NoError(t, err)
answer, err := dA.SignalOnOffer("bbbb", offer) answer, err := dA.SignalOnOffer("bbbb", offer)
require.NoError(t, err) require.NoError(t, err)
require.NoError(t, dB.SignalOnAnswer("aaaa", answer)) require.NoError(t, dB.SignalOnAnswer("aaaa", answer))
dA.OnDataPathRekeyed("bbbb")
dB.OnDataPathRekeyed("aaaa")
require.NotEqual(t, PSK{}, wgB.psk("aaaa")) require.NotEqual(t, PSK{}, wgB.psk("aaaa"))
// Now drop B's data-path delivery: rekeys can no longer converge. // Bring the data path up on both, then drop B's delivery so rekeys can't converge.
dA.OnDataPathRekeyed("bbbb")
dB.OnDataPathRekeyed("aaaa")
gB.drop.Store(true) gB.drop.Store(true)
// K-1 data-path rekeys must NOT raise OnRekeyFailed. // K-1 data-path rekeys must NOT raise OnRekeyFailed.
for i := 0; i < DefaultMaxRekeyFailures-1; i++ { for i := 0; i < DefaultMaxRekeyFailures-1; i++ {
_, err := dB.startExchange("aaaa", false) _, err := dB.startExchange("aaaa", false, ExchangeID{})
require.NoError(t, err) require.NoError(t, err)
time.Sleep(50 * time.Millisecond) time.Sleep(50 * time.Millisecond)
} }
require.Equal(t, 0, failedCount(wgB), "no failure before K attempts") require.Equal(t, 0, failedCount(wgB), "no failure before K attempts")
// The K-th failure raises it once. // The K-th failure raises it once.
_, err = dB.startExchange("aaaa", false) _, err = dB.startExchange("aaaa", false, ExchangeID{})
require.NoError(t, err) require.NoError(t, err)
require.Eventually(t, func() bool { return failedCount(wgB) == 1 }, time.Second, 5*time.Millisecond) require.Eventually(t, func() bool { return failedCount(wgB) == 1 }, time.Second, 5*time.Millisecond)
} }
+53 -119
View File
@@ -10,12 +10,8 @@ import (
) )
const ( const (
// DefaultRekeyInterval is the default PSK rotation cadence (~2 min), chosen so a
// rotated PSK is adopted by the consumer's next transport handshake without
// forcing one.
DefaultRekeyInterval = 2 * time.Minute
// DefaultRetryInterval is how often the initiator retransmits its outstanding // DefaultRetryInterval is how often the initiator retransmits its outstanding
// data-path message while an exchange is in flight. // data-path offer while awaiting the answer.
DefaultRetryInterval = 2 * time.Second DefaultRetryInterval = 2 * time.Second
// DefaultMaxRetries bounds how many ticks an exchange may run before it is // DefaultMaxRetries bounds how many ticks an exchange may run before it is
// declared failed. The convergence deadline is thus MaxRetries * RetryInterval. // declared failed. The convergence deadline is thus MaxRetries * RetryInterval.
@@ -23,16 +19,12 @@ const (
// DefaultMaxRekeyFailures is how many consecutive rekey (non-initial) failures // DefaultMaxRekeyFailures is how many consecutive rekey (non-initial) failures
// are tolerated before OnRekeyFailed. The initial exchange fails immediately. // are tolerated before OnRekeyFailed. The initial exchange fails immediately.
DefaultMaxRekeyFailures = 3 DefaultMaxRekeyFailures = 3
// confirmRetransmits is how many times the initiator best-effort resends the
// confirm over the data path, to cover its loss.
confirmRetransmits = 3
) )
// Transport pushes a message over the peer's data path (e.g. a WireGuard tunnel). // Transport pushes a message over the peer's data path (e.g. a WireGuard tunnel).
// It is the library's only outbound send: the control-plane (signalling) channel is // It is the library's only outbound send: the control-plane (signalling) channel is
// host-driven — the library hands the host offer/answer payloads to piggyback on the // host-driven — the library hands the host offer/answer payloads to piggyback on the
// host's own offer/answer, it never pushes there itself (that would drive the host's // host's own negotiation, it never pushes there itself.
// connection negotiation, which is not the library's to control).
type Transport interface { type Transport interface {
SendDataPath(remoteID string, msg []byte) error SendDataPath(remoteID string, msg []byte) error
} }
@@ -41,20 +33,19 @@ type Transport interface {
type exchangeState uint8 type exchangeState uint8
const ( const (
stateReserved exchangeState = iota // responder: deriving the answer stateReserved exchangeState = iota // responder: deriving the answer
stateAwaitingAnswer // initiator: offer sent, awaiting the answer stateAwaitingAnswer // initiator: offer sent, awaiting the answer
stateAwaitingRekey // initiator: PSK derived+set, awaiting OnDataPathRekeyed to send the confirm stateAwaitingRekey // initiator: PSK derived+set, awaiting OnDataPathRekeyed to chain the next offer
stateConfirming // initiator: confirm sent over the data path, best-effort retransmit stateAwaitingAck // responder: answer sent, awaiting the next offer that acks this exchange
stateAwaitingConfirm // responder: answer sent, awaiting the confirm over the data path
) )
// exchangeCtl holds all state for one in-flight exchange with a peer, under the // 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 current // Manager's single lock. state drives every decision. lastSent is the current
// data-path retransmit payload. initiator is the ephemeral handle used at Finish; // data-path retransmit payload (the offer, for the initiator). initiator is the
// pendingPSK is the responder's derived key. viaSignal records that the offer was // ephemeral handle used at Finish; pendingPSK is the responder's derived key.
// handed to the host for the signalling channel (bootstrap), so the loop does not // viaSignal records that the offer went to the host for the signalling channel, so
// retransmit it on the data path (the host retransmits it with its own negotiation). // the loop does not retransmit it on the data path. Only the initiator runs a
// Only the initiator runs a retransmit loop, so only it sets cancel. // retransmit loop, so only it sets cancel.
type exchangeCtl struct { type exchangeCtl struct {
id ExchangeID id ExchangeID
state exchangeState state exchangeState
@@ -67,16 +58,17 @@ type exchangeCtl struct {
} }
// Manager is the stateful orchestrator — the analogue of go-rosenpass's Server. It // Manager is the stateful orchestrator — the analogue of go-rosenpass's Server. It
// runs the per-peer rekey timer, drives the X25519MLKEM768 exchange, and surfaces the // drives the X25519MLKEM768 exchange and surfaces the derived PSK and convergence to
// derived PSK and convergence to the host via CallbackHandler. The cryptography is // the host via CallbackHandler. It is fully event-driven: the bootstrap is triggered
// the pure kem.go primitives; all state lives here under one lock. // by the host (SignalOffer) and each rotation is clocked by OnDataPathRekeyed (the
// consumer's transport rekey). The cryptography is the pure kem.go primitives; all
// state lives here under one lock.
type Manager struct { type Manager struct {
localID string localID string
transport Transport transport Transport
cbHandler CallbackHandler cbHandler CallbackHandler
logger *slog.Logger logger *slog.Logger
rekeyInterval time.Duration
retryInterval time.Duration retryInterval time.Duration
maxRetries int maxRetries int
maxRekeyFailures int maxRekeyFailures int
@@ -85,24 +77,20 @@ type Manager struct {
rootCancel context.CancelFunc rootCancel context.CancelFunc
mu sync.Mutex mu sync.Mutex
peers map[string]context.CancelFunc // per-peer rekey loop exchanges map[string]*exchangeCtl // in-flight exchange per peer
exchanges map[string]*exchangeCtl // in-flight exchange per peer established map[string]bool // peer has completed at least one exchange
established map[string]bool // peer has completed at least one exchange failures map[string]int // consecutive rekey failures per peer
failures map[string]int // consecutive rekey failures per peer // dataSend holds the peer's data-path sender when the data path is up; nil
// dataSend holds the peer's data-path sender when the data path is up; nil (absent) // (absent) means it is down. Toggled by OnDataPathRekeyed / OnDataPathDown.
// means it is down. Toggled by OnDataPathRekeyed / OnDataPathDown. Its presence is
// the "a data path exists" signal.
dataSend map[string]func(string, []byte) error dataSend map[string]func(string, []byte) error
wait sync.WaitGroup wait sync.WaitGroup
} }
// NewManager builds a manager for the local peer identified by its peer identity key // NewManager builds a manager for the local peer identified by its peer identity key
// (used for the deterministic initiator role and the identity binding). A zero // (used for the deterministic initiator role and the identity binding). A nil logger
// interval falls back to DefaultRekeyInterval; a nil logger to slog.Default(). // falls back to slog.Default(). Retry/retries/K use their defaults and can be
func NewManager(localID string, t Transport, h CallbackHandler, interval time.Duration, logger *slog.Logger) *Manager { // overridden before use.
if interval <= 0 { func NewManager(localID string, t Transport, h CallbackHandler, logger *slog.Logger) *Manager {
interval = DefaultRekeyInterval
}
if logger == nil { if logger == nil {
logger = slog.Default() logger = slog.Default()
} }
@@ -112,13 +100,11 @@ func NewManager(localID string, t Transport, h CallbackHandler, interval time.Du
transport: t, transport: t,
cbHandler: h, cbHandler: h,
logger: logger, logger: logger,
rekeyInterval: interval,
retryInterval: DefaultRetryInterval, retryInterval: DefaultRetryInterval,
maxRetries: DefaultMaxRetries, maxRetries: DefaultMaxRetries,
maxRekeyFailures: DefaultMaxRekeyFailures, maxRekeyFailures: DefaultMaxRekeyFailures,
rootCtx: ctx, rootCtx: ctx,
rootCancel: cancel, rootCancel: cancel,
peers: make(map[string]context.CancelFunc),
exchanges: make(map[string]*exchangeCtl), exchanges: make(map[string]*exchangeCtl),
established: make(map[string]bool), established: make(map[string]bool),
failures: make(map[string]int), failures: make(map[string]int),
@@ -133,26 +119,9 @@ func (m *Manager) IsInitiator(remoteID string) bool {
return m.localID > remoteID return m.localID > remoteID
} }
// AddPeer registers a remote peer and starts its rekey timer. Re-adding is a no-op. // RemovePeer stops any in-flight exchange for a peer and drops its state.
func (m *Manager) AddPeer(remoteID string) {
m.mu.Lock()
defer m.mu.Unlock()
if _, ok := m.peers[remoteID]; ok {
return
}
ctx, cancel := context.WithCancel(m.rootCtx)
m.peers[remoteID] = cancel
m.wait.Add(1)
go m.rekeyLoop(ctx, remoteID)
}
// RemovePeer stops a peer's rekey timer and any in-flight exchange, and drops state.
func (m *Manager) RemovePeer(remoteID string) { func (m *Manager) RemovePeer(remoteID string) {
m.mu.Lock() m.mu.Lock()
if cancel, ok := m.peers[remoteID]; ok {
cancel()
delete(m.peers, remoteID)
}
if ex, ok := m.exchanges[remoteID]; ok { if ex, ok := m.exchanges[remoteID]; ok {
if ex.cancel != nil { if ex.cancel != nil {
ex.cancel() ex.cancel()
@@ -165,23 +134,21 @@ func (m *Manager) RemovePeer(remoteID string) {
m.mu.Unlock() m.mu.Unlock()
} }
// Stop cancels all timers and in-flight exchanges and waits for goroutines to exit. // Stop cancels all in-flight exchanges and waits for their goroutines to exit.
func (m *Manager) Stop() { func (m *Manager) Stop() {
m.rootCancel() m.rootCancel()
m.wait.Wait() m.wait.Wait()
m.mu.Lock() m.mu.Lock()
m.peers = make(map[string]context.CancelFunc)
m.exchanges = make(map[string]*exchangeCtl) m.exchanges = make(map[string]*exchangeCtl)
m.mu.Unlock() m.mu.Unlock()
} }
// ---- Signalling channel (host-driven; rides the host's offer/answer) ---- // ---- Signalling channel (host-driven; rides the host's negotiation) ----
// SignalOffer returns the KEM offer for the host to embed in its outgoing offer to // SignalOffer returns the KEM offer for the host to embed in its outgoing offer to
// remoteID (bootstrap). It returns (nil, nil) when the local peer is not the // remoteID (bootstrap). It returns (nil, nil) when the local peer is not the
// initiator. It is idempotent for an in-flight bootstrap: a repeat call (e.g. the // initiator. It is idempotent for an in-flight bootstrap: a repeat call returns the
// host retransmitting its offer) returns the same offer rather than starting a new // same offer rather than starting a new exchange.
// exchange.
func (m *Manager) SignalOffer(remoteID string) ([]byte, error) { func (m *Manager) SignalOffer(remoteID string) ([]byte, error) {
if !m.IsInitiator(remoteID) { if !m.IsInitiator(remoteID) {
return nil, nil return nil, nil
@@ -193,7 +160,8 @@ func (m *Manager) SignalOffer(remoteID string) ([]byte, error) {
return last, nil return last, nil
} }
m.mu.Unlock() m.mu.Unlock()
return m.startExchange(remoteID, true) // bootstrap offer acknowledges nothing (zero AckID).
return m.startExchange(remoteID, true, ExchangeID{})
} }
// SignalOnOffer processes a KEM offer the host extracted from an incoming offer and // SignalOnOffer processes a KEM offer the host extracted from an incoming offer and
@@ -210,7 +178,7 @@ func (m *Manager) SignalOnOffer(remoteID string, offer []byte) ([]byte, error) {
} }
// SignalOnAnswer processes a KEM answer the host extracted from an incoming answer. // SignalOnAnswer processes a KEM answer the host extracted from an incoming answer.
// There is no reply: the confirm rides the data path after OnDataPathRekeyed. // There is no reply: the next offer (over the data path) acknowledges this exchange.
func (m *Manager) SignalOnAnswer(remoteID string, answer []byte) error { func (m *Manager) SignalOnAnswer(remoteID string, answer []byte) error {
typ, msg, err := Decode(answer) typ, msg, err := Decode(answer)
if err != nil { if err != nil {
@@ -243,46 +211,42 @@ func (m *Manager) OnDataPathMessage(remoteID string, raw []byte) error {
return m.pushDataPath(remoteID, answer) return m.pushDataPath(remoteID, answer)
case MsgAnswer: case MsgAnswer:
return m.processAnswer(remoteID, msg.(*AnswerMsg)) return m.processAnswer(remoteID, msg.(*AnswerMsg))
case MsgConfirm:
return m.processConfirm(remoteID, msg.(*ConfirmMsg))
default: default:
return fmt.Errorf("unhandled data-path message type %d from %s", typ, remoteID) return fmt.Errorf("unhandled data-path message type %d from %s", typ, remoteID)
} }
} }
// OnDataPathRekeyed notifies that the peer's data path is up and freshly keyed with // OnDataPathRekeyed notifies that the peer's data path is up and freshly keyed with
// the latest PSK (fired on first establishment AND every rekey — the same event). It // the latest PSK (fired on first establishment AND every rekey). It marks the data
// marks the data path usable and, if we are the initiator waiting to confirm, sends // path usable and, if we are the initiator that just derived a PSK, chains the next
// the confirm over the data path (its arrival proves to the responder that we operate // exchange: a fresh offer over the data path that acknowledges the just-completed one
// on the new key, correlated by exchangeID). // (its arrival under the new key proves to the responder that the key works).
func (m *Manager) OnDataPathRekeyed(remoteID string) { func (m *Manager) OnDataPathRekeyed(remoteID string) {
m.mu.Lock() m.mu.Lock()
m.dataSend[remoteID] = m.transport.SendDataPath m.dataSend[remoteID] = m.transport.SendDataPath
ex := m.exchanges[remoteID] ex := m.exchanges[remoteID]
if ex == nil || ex.state != stateAwaitingRekey { chain := ex != nil && ex.state == stateAwaitingRekey
m.mu.Unlock() var ackID ExchangeID
return if chain {
ackID = ex.id
} }
confirm, err := (&ConfirmMsg{ExchangeID: ex.id}).Encode()
if err != nil {
m.mu.Unlock()
m.logger.Error("pqkem encode confirm", "peer", remoteID, "err", err)
return
}
ex.state = stateConfirming
ex.lastSent = confirm
m.established[remoteID] = true
m.failures[remoteID] = 0
_ = time.Since(ex.startedAt) // convergence latency (metrics hook, later step)
m.mu.Unlock() m.mu.Unlock()
if err := m.pushDataPath(remoteID, confirm); err != nil { if !chain {
m.logger.Warn("pqkem send confirm failed", "peer", remoteID, "err", err) return
}
offer, err := m.startExchange(remoteID, false, ackID)
if err != nil {
m.logger.Error("pqkem chain offer failed to start", "peer", remoteID, "err", err)
return
}
if err := m.pushDataPath(remoteID, offer); err != nil {
m.logger.Warn("pqkem send chain offer failed", "peer", remoteID, "err", err)
} }
} }
// OnDataPathDown notifies that the peer's data path went down; further rekeys wait // OnDataPathDown notifies that the peer's data path went down; rotations pause until
// until it is up again (the host re-bootstraps over signalling on reconnect). // it is up again (the host re-bootstraps over signalling on reconnect).
func (m *Manager) OnDataPathDown(remoteID string) { func (m *Manager) OnDataPathDown(remoteID string) {
m.mu.Lock() m.mu.Lock()
delete(m.dataSend, remoteID) delete(m.dataSend, remoteID)
@@ -291,36 +255,6 @@ func (m *Manager) OnDataPathDown(remoteID string) {
// ---- internals ---- // ---- internals ----
func (m *Manager) rekeyLoop(ctx context.Context, remoteID string) {
defer m.wait.Done()
t := time.NewTicker(m.rekeyInterval)
defer t.Stop()
for {
select {
case <-ctx.Done():
return
case <-t.C:
m.mu.Lock()
_, dpUp := m.dataSend[remoteID]
_, inFlight := m.exchanges[remoteID]
m.mu.Unlock()
// Rekeys ride the data path only; skip when it is down (the host will
// re-bootstrap over signalling on reconnect) or an exchange is in flight.
if !dpUp || inFlight || !m.IsInitiator(remoteID) {
continue
}
offer, err := m.startExchange(remoteID, false)
if err != nil {
m.logger.Error("pqkem rekey failed to start", "peer", remoteID, "err", err)
continue
}
if err := m.pushDataPath(remoteID, offer); err != nil {
m.logger.Warn("pqkem send rekey offer failed", "peer", remoteID, "err", err)
}
}
}
}
// pushDataPath sends over the peer's data path, erroring if it is down. // pushDataPath sends over the peer's data path, erroring if it is down.
func (m *Manager) pushDataPath(remoteID string, msg []byte) error { func (m *Manager) pushDataPath(remoteID string, msg []byte) error {
m.mu.Lock() m.mu.Lock()
+48 -34
View File
@@ -3,14 +3,13 @@ package pqkem
import ( import (
"sync" "sync"
"testing" "testing"
"time"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
// loopback is a data-path transport: a SendDataPath delivers synchronously to the // loopback is a data-path transport: SendDataPath delivers synchronously to the peer
// peer manager's OnDataPathMessage, attributing it to localID (the sender). The // manager's OnDataPathMessage, attributing it to localID (the sender). The signalling
// signalling channel is driven by the test directly via the SignalX methods. // channel is driven by the test directly via the SignalX methods.
type loopback struct { type loopback struct {
localID string localID string
peer *Manager peer *Manager
@@ -49,60 +48,75 @@ func (f *fakeWG) psk(peer string) PSK {
return f.psks[peer] return f.psks[peer]
} }
func TestManager_ExchangeConverges(t *testing.T) { // pair builds two wired managers (B is the initiator, "bbbb" > "aaaa").
func pair(t *testing.T) (dA, dB *Manager, wgA, wgB *fakeWG) {
t.Helper()
lbA := &loopback{localID: "aaaa"} lbA := &loopback{localID: "aaaa"}
lbB := &loopback{localID: "bbbb"} lbB := &loopback{localID: "bbbb"}
wgA := newFakeWG() wgA = newFakeWG()
wgB := newFakeWG() wgB = newFakeWG()
dA = NewManager("aaaa", lbA, wgA, nil)
dB = NewManager("bbbb", lbB, wgB, nil)
lbA.peer = dB
lbB.peer = dA
return dA, dB, wgA, wgB
}
dA := NewManager("aaaa", lbA, wgA, time.Hour, nil) // bootstrap runs the signalling offer/answer (the test plays the host carrying bytes).
dB := NewManager("bbbb", lbB, wgB, time.Hour, nil) func bootstrap(t *testing.T, dA, dB *Manager) {
lbA.peer = dB // A's data-path sends -> B receives t.Helper()
lbB.peer = dA // B's data-path sends -> A receives
dA.AddPeer("bbbb")
dB.AddPeer("aaaa")
defer dA.Stop()
defer dB.Stop()
// Bootstrap over the signalling channel (the test plays the host carrying bytes).
// B is the initiator ("bbbb" > "aaaa").
offer, err := dB.SignalOffer("aaaa") offer, err := dB.SignalOffer("aaaa")
require.NoError(t, err) require.NoError(t, err)
require.NotNil(t, offer) require.NotNil(t, offer)
answer, err := dA.SignalOnOffer("bbbb", offer) answer, err := dA.SignalOnOffer("bbbb", offer)
require.NoError(t, err) require.NoError(t, err)
require.NotNil(t, answer) require.NotNil(t, answer)
require.NoError(t, dB.SignalOnAnswer("aaaa", answer)) require.NoError(t, dB.SignalOnAnswer("aaaa", answer))
}
// Data path comes up on both sides; this makes B send the confirm over the data func TestManager_BootstrapDerivesSamePSK(t *testing.T) {
// path, converging A. dA, dB, wgA, wgB := pair(t)
defer dA.Stop()
defer dB.Stop()
bootstrap(t, dA, dB)
pskA := wgA.psk("bbbb")
pskB := wgB.psk("aaaa")
require.NotEqual(t, PSK{}, pskA)
require.Equal(t, pskB, pskA, "both sides derive the same PSK from the bootstrap exchange")
}
func TestManager_ChainRotatesAndAcks(t *testing.T) {
dA, dB, wgA, wgB := pair(t)
defer dA.Stop()
defer dB.Stop()
bootstrap(t, dA, dB)
psk1 := wgB.psk("aaaa")
// Data path up on both sides; B chains the next offer (acking exchange 1) over the
// data path, which rotates both to a fresh PSK and acknowledges A.
dA.OnDataPathRekeyed("bbbb") dA.OnDataPathRekeyed("bbbb")
dB.OnDataPathRekeyed("aaaa") dB.OnDataPathRekeyed("aaaa")
pskB := wgB.psk("aaaa") psk2A := wgA.psk("bbbb")
pskA := wgA.psk("bbbb") psk2B := wgB.psk("aaaa")
require.NotEqual(t, PSK{}, pskA, "responder A must have a PSK") require.Equal(t, psk2B, psk2A, "both sides converge on the rotated PSK")
require.NotEqual(t, PSK{}, pskB, "initiator B must have a PSK") require.NotEqual(t, psk1, psk2B, "the chain rotated to a new PSK")
require.Equal(t, pskB, pskA, "both sides converge on the same PSK")
} }
func TestManager_NonInitiatorReturnsNoOffer(t *testing.T) { func TestManager_NonInitiatorReturnsNoOffer(t *testing.T) {
dA := NewManager("aaaa", &loopback{localID: "aaaa"}, newFakeWG(), time.Hour, nil) dA := NewManager("aaaa", &loopback{localID: "aaaa"}, newFakeWG(), nil)
dA.AddPeer("bbbb")
defer dA.Stop() defer dA.Stop()
// A is NOT the initiator vs "bbbb" -> no offer to send. offer, err := dA.SignalOffer("bbbb") // not the initiator vs "bbbb"
offer, err := dA.SignalOffer("bbbb")
require.NoError(t, err) require.NoError(t, err)
require.Nil(t, offer) require.Nil(t, offer)
} }
func TestManager_StopIsIdempotent(t *testing.T) { func TestManager_StopIsIdempotent(t *testing.T) {
dA := NewManager("aaaa", &loopback{localID: "aaaa"}, newFakeWG(), time.Hour, nil) dA := NewManager("aaaa", &loopback{localID: "aaaa"}, newFakeWG(), nil)
dA.AddPeer("bbbb")
dA.Stop() dA.Stop()
dA.Stop() // must not panic or hang dA.Stop() // must not panic or hang
} }
+26 -41
View File
@@ -6,39 +6,45 @@ import (
) )
// Wire framing for the PQ-KEM exchange. Messages are self-contained, versioned, // Wire framing for the PQ-KEM exchange. Messages are self-contained, versioned,
// transport-agnostic byte blobs: the same bytes ride the Signal offer/answer // transport-agnostic byte blobs: the same bytes ride the signalling channel
// (initial, pre-tunnel) or a data-tunnel packet (rekey). They are NOT a gRPC // (initial bootstrap) or a data-tunnel packet (rekey). The library only ever sees
// service — the network layer only sees opaque []byte. // opaque []byte at the transport seam.
// //
// Layout (all messages): [type:1][version:1][exchangeID:16][payload...] // Layout (all messages): [type:1][version:1][exchangeID:16][payload...]
//
// There is no confirm message: an exchange is acknowledged by the NEXT offer, which
// carries the acked exchange's id (see OfferMsg.AckID) and — riding the data path
// under the freshly adopted key — proves that key works.
const ( const (
// ProtocolVersion is bumped on any wire-incompatible change; a peer rejects // ProtocolVersion is bumped on any wire-incompatible change; a peer rejects
// messages it does not understand rather than misparsing them. // messages it does not understand rather than misparsing them.
ProtocolVersion uint8 = 1 ProtocolVersion uint8 = 1
// ExchangeIDSize identifies one offer/answer/confirm round so stale answers // ExchangeIDSize identifies one exchange so answers/acks correlate and stale
// (e.g. an answer to a pre-restart offer) are dropped instead of applied. // messages are dropped.
ExchangeIDSize = 16 ExchangeIDSize = 16
headerSize = 1 + 1 + ExchangeIDSize headerSize = 1 + 1 + ExchangeIDSize
) )
// MsgType tags the three message kinds of the exchange. // MsgType tags the two message kinds of the exchange.
type MsgType uint8 type MsgType uint8
const ( const (
MsgOffer MsgType = iota + 1 MsgOffer MsgType = iota + 1
MsgAnswer MsgAnswer
MsgConfirm
) )
// ExchangeID is the per-round correlator echoed by the answer and the confirm. // ExchangeID is the per-exchange correlator. The zero value means "none" (an offer
// that acknowledges nothing, i.e. the first exchange of a connection).
type ExchangeID [ExchangeIDSize]byte type ExchangeID [ExchangeIDSize]byte
// OfferMsg carries the initiator's public material (X25519 pub ‖ ML-KEM encap key). // OfferMsg carries the initiator's public material (X25519 pub ‖ ML-KEM encap key)
// and AckID, the id of the previous exchange this offer acknowledges (zero if none).
type OfferMsg struct { type OfferMsg struct {
ExchangeID ExchangeID ExchangeID ExchangeID
AckID ExchangeID
// KEMOffer is the raw Initiator.Offer() blob (OfferSize bytes). // KEMOffer is the raw Initiator.Offer() blob (OfferSize bytes).
KEMOffer []byte KEMOffer []byte
} }
@@ -51,27 +57,15 @@ type AnswerMsg struct {
KEMAnswer []byte KEMAnswer []byte
} }
// ConfirmMsg is sent ONE WAY, initiator -> responder, over the tunnel under the // Encode serialises the offer with its framed header (payload = AckID ‖ KEMOffer).
// NEW PSK, for the round identified by ExchangeID. It carries no key material.
//
// It is one-way because of the knowledge asymmetry of a KEM: the initiator learns
// the responder holds the key simply by receiving the answer (the responder had to
// derive it to encapsulate), so the initiator needs no confirmation. Only the
// responder is left unsure whether the initiator received the answer and committed
// the key — this message resolves that. As a bonus, being sent under the new PSK it
// exercises the consumer's channel handshake with the new key: the responder
// converges on receiving it, and the initiator converges by observing that handshake
// succeed (which fails on a PSK mismatch, so success proves the responder committed too).
type ConfirmMsg struct {
ExchangeID ExchangeID
}
// Encode serialises the offer with its framed header.
func (m *OfferMsg) Encode() ([]byte, error) { func (m *OfferMsg) Encode() ([]byte, error) {
if len(m.KEMOffer) != OfferSize { if len(m.KEMOffer) != OfferSize {
return nil, fmt.Errorf("offer payload: got %d, want %d", len(m.KEMOffer), OfferSize) return nil, fmt.Errorf("offer payload: got %d, want %d", len(m.KEMOffer), OfferSize)
} }
return frame(MsgOffer, m.ExchangeID, m.KEMOffer), nil payload := make([]byte, 0, ExchangeIDSize+OfferSize)
payload = append(payload, m.AckID[:]...)
payload = append(payload, m.KEMOffer...)
return frame(MsgOffer, m.ExchangeID, payload), nil
} }
// Encode serialises the answer with its framed header. // Encode serialises the answer with its framed header.
@@ -82,13 +76,7 @@ func (m *AnswerMsg) Encode() ([]byte, error) {
return frame(MsgAnswer, m.ExchangeID, m.KEMAnswer), nil return frame(MsgAnswer, m.ExchangeID, m.KEMAnswer), nil
} }
// Encode serialises the confirm with its framed header. It returns an error only // Decode parses a framed message into one of *OfferMsg / *AnswerMsg.
// for signature uniformity with the other messages; it never actually fails.
func (m *ConfirmMsg) Encode() ([]byte, error) {
return frame(MsgConfirm, m.ExchangeID, nil), nil
}
// Decode parses a framed message into one of *OfferMsg / *AnswerMsg / *ConfirmMsg.
func Decode(buf []byte) (MsgType, any, error) { func Decode(buf []byte) (MsgType, any, error) {
if len(buf) < headerSize { if len(buf) < headerSize {
return 0, nil, fmt.Errorf("message too short: %d bytes", len(buf)) return 0, nil, fmt.Errorf("message too short: %d bytes", len(buf))
@@ -104,20 +92,17 @@ func Decode(buf []byte) (MsgType, any, error) {
switch typ { switch typ {
case MsgOffer: case MsgOffer:
if len(payload) != OfferSize { if len(payload) != ExchangeIDSize+OfferSize {
return typ, nil, fmt.Errorf("offer payload: got %d, want %d", len(payload), OfferSize) return typ, nil, fmt.Errorf("offer payload: got %d, want %d", len(payload), ExchangeIDSize+OfferSize)
} }
return typ, &OfferMsg{ExchangeID: id, KEMOffer: payload}, nil var ack ExchangeID
copy(ack[:], payload[:ExchangeIDSize])
return typ, &OfferMsg{ExchangeID: id, AckID: ack, KEMOffer: payload[ExchangeIDSize:]}, nil
case MsgAnswer: case MsgAnswer:
if len(payload) != AnswerSize { if len(payload) != AnswerSize {
return typ, nil, fmt.Errorf("answer payload: got %d, want %d", len(payload), AnswerSize) return typ, nil, fmt.Errorf("answer payload: got %d, want %d", len(payload), AnswerSize)
} }
return typ, &AnswerMsg{ExchangeID: id, KEMAnswer: payload}, nil return typ, &AnswerMsg{ExchangeID: id, KEMAnswer: payload}, nil
case MsgConfirm:
if len(payload) != 0 {
return typ, nil, fmt.Errorf("confirm payload must be empty, got %d bytes", len(payload))
}
return typ, &ConfirmMsg{ExchangeID: id}, nil
default: default:
return typ, nil, fmt.Errorf("unknown message type %d", typ) return typ, nil, fmt.Errorf("unknown message type %d", typ)
} }
+4 -9
View File
@@ -13,13 +13,15 @@ func TestMessageRoundTrip(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
id := ExchangeID{1, 2, 3, 4} id := ExchangeID{1, 2, 3, 4}
ack := ExchangeID{9, 9, 9}
offBytes, err := (&OfferMsg{ExchangeID: id, KEMOffer: init.Offer()}).Encode() offBytes, err := (&OfferMsg{ExchangeID: id, AckID: ack, KEMOffer: init.Offer()}).Encode()
require.NoError(t, err) require.NoError(t, err)
typ, decoded, err := Decode(offBytes) typ, decoded, err := Decode(offBytes)
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, MsgOffer, typ) require.Equal(t, MsgOffer, typ)
require.Equal(t, id, decoded.(*OfferMsg).ExchangeID) require.Equal(t, id, decoded.(*OfferMsg).ExchangeID)
require.Equal(t, ack, decoded.(*OfferMsg).AckID)
require.Equal(t, init.Offer(), decoded.(*OfferMsg).KEMOffer) require.Equal(t, init.Offer(), decoded.(*OfferMsg).KEMOffer)
ansBytes, err := (&AnswerMsg{ExchangeID: id, KEMAnswer: answer}).Encode() ansBytes, err := (&AnswerMsg{ExchangeID: id, KEMAnswer: answer}).Encode()
@@ -28,13 +30,6 @@ func TestMessageRoundTrip(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, MsgAnswer, typ) require.Equal(t, MsgAnswer, typ)
require.Equal(t, answer, decoded.(*AnswerMsg).KEMAnswer) require.Equal(t, answer, decoded.(*AnswerMsg).KEMAnswer)
confBytes, err := (&ConfirmMsg{ExchangeID: id}).Encode()
require.NoError(t, err)
typ, decoded, err = Decode(confBytes)
require.NoError(t, err)
require.Equal(t, MsgConfirm, typ)
require.Equal(t, id, decoded.(*ConfirmMsg).ExchangeID)
} }
func TestDecodeRejects(t *testing.T) { func TestDecodeRejects(t *testing.T) {
@@ -43,7 +38,7 @@ func TestDecodeRejects(t *testing.T) {
require.Error(t, err) require.Error(t, err)
// wrong version // wrong version
bad := make([]byte, headerSize+OfferSize) bad := make([]byte, headerSize+ExchangeIDSize+OfferSize)
bad[0] = byte(MsgOffer) bad[0] = byte(MsgOffer)
bad[1] = ProtocolVersion + 1 bad[1] = ProtocolVersion + 1
_, _, err = Decode(bad) _, _, err = Decode(bad)