From c49666c859ec45b65529d73716d9e87e038b38d2 Mon Sep 17 00:00:00 2001
From: riccardom <riccardomanfrin@gmail.com>
Date: Mon, 27 Jul 2026 16:18:35 +0200
Subject: [PATCH] pqkem: register data-path endpoint from signalling

Learn the peer's data-path endpoint from the signalling offer/answer: its WG
overlay IP combined with the advertised pq UDP port (SetRemotePort -> AddPeer).
Registering here is safe before the tunnel is up because sends only ever fire
once it is (clocked by OnDataPathRekeyed). RemovePeer is wired at peer teardown
(engine.removePeer), not on transient disconnect.
---
 client/internal/engine.go          |  4 ++++
 client/internal/peer/conn.go       |  3 +++
 client/internal/peer/handshaker.go | 13 +++++++++++++
 client/internal/pqkem_adapter.go   | 12 ++++++++++++
 4 files changed, 32 insertions(+)

diff --git a/client/internal/engine.go b/client/internal/engine.go
index 216ffa05e..728943119 100644
--- a/client/internal/engine.go
+++ b/client/internal/engine.go
@@ -934,6 +934,10 @@ func (e *Engine) removePeer(peerKey string) error {
 
 	e.connMgr.RemovePeerConn(peerKey)
 
+	if e.pqkemManager != nil {
+		e.pqkemManager.RemovePeer(pqkem.RemoteID(peerKey))
+	}
+
 	err := e.statusRecorder.RemovePeer(peerKey)
 	if err != nil {
 		log.Warnf("received error when removing peer %s from status recorder: %v", peerKey, err)
diff --git a/client/internal/peer/conn.go b/client/internal/peer/conn.go
index cbb6746a0..9c3e26d8e 100644
--- a/client/internal/peer/conn.go
+++ b/client/internal/peer/conn.go
@@ -91,6 +91,9 @@ type PQHandshaker interface {
 	// PSK returns the peer's latest derived post-quantum PSK to program at WG
 	// peer-config time (the pull path). ok is false until one has been derived.
 	PSK(remoteKey string) (wgtypes.Key, bool)
+	// SetRemotePort registers the peer's data-path endpoint from signalling: the peer's
+	// WG overlay IP combined with its advertised pq UDP port.
+	SetRemotePort(remoteKey string, overlayIP netip.Addr, port int)
 }
 
 // ConnConfig is a peer Connection configuration
diff --git a/client/internal/peer/handshaker.go b/client/internal/peer/handshaker.go
index 82b35dd79..7fa366932 100644
--- a/client/internal/peer/handshaker.go
+++ b/client/internal/peer/handshaker.go
@@ -135,6 +135,8 @@ func (h *Handshaker) Listen(ctx context.Context) {
 
 			h.updateRemoteICEState(&remoteOfferAnswer)
 
+			h.pqRegisterEndpoint(remoteOfferAnswer.MlkemPort)
+
 			if h.relayListener != nil {
 				h.relayListener.Notify(&remoteOfferAnswer)
 			}
@@ -157,6 +159,8 @@ func (h *Handshaker) Listen(ctx context.Context) {
 
 			h.updateRemoteICEState(&remoteOfferAnswer)
 
+			h.pqRegisterEndpoint(remoteOfferAnswer.MlkemPort)
+
 			if h.relayListener != nil {
 				h.relayListener.Notify(&remoteOfferAnswer)
 			}
@@ -175,6 +179,15 @@ func (h *Handshaker) Listen(ctx context.Context) {
 	}
 }
 
+// pqRegisterEndpoint feeds the post-quantum handshaker the peer's data-path endpoint
+// (its WG overlay IP plus the advertised pq UDP port) learned from a remote offer/answer.
+func (h *Handshaker) pqRegisterEndpoint(remotePort int) {
+	if h.config.PQ == nil || remotePort <= 0 || len(h.config.WgConfig.AllowedIps) == 0 {
+		return
+	}
+	h.config.PQ.SetRemotePort(h.config.Key, h.config.WgConfig.AllowedIps[0].Addr(), remotePort)
+}
+
 func (h *Handshaker) SendOffer() error {
 	h.mu.Lock()
 	defer h.mu.Unlock()
diff --git a/client/internal/pqkem_adapter.go b/client/internal/pqkem_adapter.go
index 611e72498..4de9a0ab1 100644
--- a/client/internal/pqkem_adapter.go
+++ b/client/internal/pqkem_adapter.go
@@ -1,6 +1,8 @@
 package internal
 
 import (
+	"net/netip"
+
 	log "github.com/sirupsen/logrus"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
@@ -78,3 +80,13 @@ func (p pqHandshaker) PSK(remoteKey string) (wgtypes.Key, bool) {
 	}
 	return wgtypes.Key(psk), true
 }
+
+// SetRemotePort registers the peer's data-path endpoint (overlay IP + pq UDP port)
+// learned from signalling. Sends only ever fire once the tunnel is up (clocked by
+// OnDataPathRekeyed), so registering here is safe even before connection-up.
+func (p pqHandshaker) SetRemotePort(remoteKey string, overlayIP netip.Addr, port int) {
+	if port <= 0 || port > 65535 || !overlayIP.IsValid() {
+		return
+	}
+	p.mgr.AddPeer(pqkem.RemoteID(remoteKey), netip.AddrPortFrom(overlayIP, uint16(port)))
+}