Anticipates PSK before WG does handshake so it finds it to set it

This commit is contained in:
riccardom
2026-08-05 16:39:55 +02:00
parent 2df8e69f59
commit 726ea030ab
3 changed files with 25 additions and 12 deletions

View File

@@ -138,6 +138,15 @@ func (h *Handshaker) Listen(ctx context.Context) {
h.pqRegisterEndpoint(remoteOfferAnswer.MlkemPort)
// Derive+store the KEM PSK (inside sendAnswer's AnswerPayload) BEFORE bringing
// up the connection: the relay/ICE workers configure the WG endpoint, which
// pulls the PSK for the first handshake. Notifying them first would race the
// KEM exchange and hand the first handshake a not-yet-derived key.
if err := h.sendAnswer(&remoteOfferAnswer); err != nil {
h.log.Errorf("failed to send remote offer confirmation: %s", err)
continue
}
if h.relayListener != nil {
h.relayListener.Notify(&remoteOfferAnswer)
}
@@ -145,11 +154,6 @@ func (h *Handshaker) Listen(ctx context.Context) {
if h.iceListener != nil && h.RemoteICESupported() {
h.iceListener(&remoteOfferAnswer)
}
if err := h.sendAnswer(&remoteOfferAnswer); err != nil {
h.log.Errorf("failed to send remote offer confirmation: %s", err)
continue
}
case remoteOfferAnswer := <-h.remoteAnswerCh:
h.log.Infof("received answer, running version %s, remote WireGuard listen port %d, session id: %s, remote ICE supported: %t", remoteOfferAnswer.Version, remoteOfferAnswer.WgListenPort, remoteOfferAnswer.SessionIDString(), remoteOfferAnswer.hasICECredentials())
@@ -162,6 +166,13 @@ func (h *Handshaker) Listen(ctx context.Context) {
h.pqRegisterEndpoint(remoteOfferAnswer.MlkemPort)
// Feed the KEM answer (derive+store PSK) BEFORE bringing up the connection so
// the WG endpoint config pulls the real PSK for the first handshake instead of
// racing ahead of the KEM exchange.
if h.config.PQ != nil {
h.config.PQ.OnAnswer(h.config.Key, remoteOfferAnswer.MlkemPayload)
}
if h.relayListener != nil {
h.relayListener.Notify(&remoteOfferAnswer)
}
@@ -169,10 +180,6 @@ func (h *Handshaker) Listen(ctx context.Context) {
if h.iceListener != nil && h.RemoteICESupported() {
h.iceListener(&remoteOfferAnswer)
}
if h.config.PQ != nil {
h.config.PQ.OnAnswer(h.config.Key, remoteOfferAnswer.MlkemPayload)
}
case <-ctx.Done():
h.log.Infof("stop listening for remote offers and answers")
return

View File

@@ -58,7 +58,8 @@ func TestManager_EstablishedPeerNotDowngraded(t *testing.T) {
dB.MarkNonCapable("aaaa") // stray zero after establishment
offer, err := dB.SignalOffer("aaaa")
require.NoError(t, err)
require.NotNil(t, offer, "an established peer must keep running the KEM despite a stray zero")
// The peer keeps its derived PSK (MarkNonCapable is a no-op once established).
psk, ok := dB.PSK("aaaa")
require.True(t, ok, "an established peer must keep its PSK despite a stray zero")
require.NotEqual(t, PSK{}, psk)
}

View File

@@ -271,6 +271,11 @@ func (m *Manager) Stop() {
// 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 returns the
// same offer rather than starting a new exchange.
//
// A signal re-negotiation always re-bootstraps (fresh exchange): the remote may have
// restarted and lost its PSK, so reusing a locally frozen one would desync. The derived
// PSK still survives idle in the manager (dropped only on account-level peer removal),
// so a pure lazy wake with no re-negotiation reuses it via the conn's WG-config pull.
func (m *Manager) SignalOffer(remoteID RemoteID) ([]byte, error) {
if !m.IsInitiator(remoteID) {
return nil, nil