From cac03bd80cf3a1ac920dc87d6cd26d619331b11a 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 cdaa5057c..1ab629466 100644
--- a/client/internal/engine.go
+++ b/client/internal/engine.go
@@ -932,6 +932,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 c454c014b..7c6ba1215 100644
--- a/client/internal/peer/conn.go
+++ b/client/internal/peer/conn.go
@@ -90,6 +90,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 23f744909..767688375 100644
--- a/client/internal/peer/handshaker.go
+++ b/client/internal/peer/handshaker.go
@@ -130,6 +130,8 @@ func (h *Handshaker) Listen(ctx context.Context) {
 
 			h.updateRemoteICEState(&remoteOfferAnswer)
 
+			h.pqRegisterEndpoint(remoteOfferAnswer.MlkemPort)
+
 			if h.relayListener != nil {
 				h.relayListener.Notify(&remoteOfferAnswer)
 			}
@@ -152,6 +154,8 @@ func (h *Handshaker) Listen(ctx context.Context) {
 
 			h.updateRemoteICEState(&remoteOfferAnswer)
 
+			h.pqRegisterEndpoint(remoteOfferAnswer.MlkemPort)
+
 			if h.relayListener != nil {
 				h.relayListener.Notify(&remoteOfferAnswer)
 			}
@@ -170,6 +174,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)))
+}