Adds transport

client/internal/engine.go

@@ -200,6 +200,8 @@ type Engine struct {
 
 	// pqkemManager runs the ML-KEM post-quantum PSK exchange (gated by NB_ENABLE_PQ_MLKEM).
 	pqkemManager *pqkem.Manager
+	// pqTransport is the ML-KEM data-path UDP transport, bound on the WG overlay IP.
+	pqTransport *pqTransport
 
 	// syncMsgMux is used to guarantee sequential Management Service message processing
 	syncMsgMux *sync.Mutex
@@ -655,12 +657,19 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
 		e.rpManager.SetInterface(e.wgInterface)
 	}
 
-	// Start the ML-KEM PQ manager after the interface is up so its (dedicated UDP)
-	// transport can bind on the WG overlay IP. TODO(NET-1406): replace noopPQTransport
-	// with the real transport bound on e.config.WgAddr.IP.
+	// Start the ML-KEM PQ manager after the interface is up so its dedicated UDP
+	// transport can bind on the WG overlay IP.
 	if pqkem.Enabled() {
-		log.Infof("ML-KEM post-quantum exchange enabled")
-		e.pqkemManager = pqkem.NewManager(publicKey.String(), noopPQTransport{}, pqCallbackHandler{wg: e.wgInterface}, nil)
+		tr, pqErr := newPQTransport(e.config.WgAddr.IP)
+		if pqErr != nil {
+			log.Errorf("ML-KEM PQ transport bind failed, PQ exchange disabled: %v", pqErr)
+		} else {
+			e.pqTransport = tr
+			e.pqkemManager = pqkem.NewManager(publicKey.String(), tr, pqCallbackHandler{wg: e.wgInterface}, nil)
+			tr.setManager(e.pqkemManager)
+			go tr.run()
+			log.Infof("ML-KEM post-quantum exchange enabled (udp port %d on overlay %s)", tr.Port(), e.config.WgAddr.IP)
+		}
 	}
 
 	// if inbound conns are blocked there is no need to create the ACL manager
@@ -2095,6 +2104,9 @@ func (e *Engine) close() {
 		_ = e.rpManager.Close()
 	}
 
+	if e.pqTransport != nil {
+		_ = e.pqTransport.Close()
+	}
 	if e.pqkemManager != nil {
 		e.pqkemManager.Stop()
 	}

client/internal/pqkem_adapter.go

@@ -31,11 +31,3 @@ func (h pqCallbackHandler) OnRekeyFailed(remoteID string) error {
 	log.Warnf("pqkem: post-quantum rekey failed for peer %s", remoteID)
 	return nil
 }
-
-// noopPQTransport is a placeholder data-path transport for the ML-KEM manager.
-// TODO(NET-1406): replace with the real dedicated UDP transport bound on the WG
-// overlay IPv4 (send over the data path, feed inbound to Manager.OnDataPathMessage),
-// analogue of go-rosenpass's Conn.
-type noopPQTransport struct{}
-
-func (noopPQTransport) SendDataPath(remoteID string, msg []byte) error { return nil }

client/internal/pqkem_transport.go

@@ -0,0 +1,132 @@
+package internal
+
+import (
+	"fmt"
+	"net"
+	"net/netip"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/internal/pqkem"
+)
+
+// DefaultPort is the preferred UDP port for the ML-KEM data-path service, bound on
+// the WG overlay IP. Since each client owns a distinct overlay IP, this port is
+// almost always free, so it need not be announced (peers assume it). A peer only
+// announces Body.mlkemPort when a collision forced it onto a different port.
+const DefaultPort = 51833
+
+// pqTransport is the ML-KEM data-path transport: a dedicated UDP socket bound on the
+// WireGuard overlay IP. Rekey messages travel through the tunnel to each peer's
+// overlay IP and announced pqkem port. It implements pqkem.Transport and feeds
+// inbound datagrams to the manager (set via setManager after construction, since the
+// manager is built with this transport).
+type pqTransport struct {
+	conn *net.UDPConn
+	port int
+
+	mu     sync.RWMutex
+	mgr    *pqkem.Manager
+	peers  map[string]*net.UDPAddr // remoteID (WG pubkey) -> overlay UDP addr to send to
+	byAddr map[string]string       // source addr string -> remoteID (inbound dispatch)
+}
+
+// newPQTransport binds a UDP socket on the WG overlay IP, preferring DefaultPort and
+// falling back to an OS-assigned ephemeral port if it is in use. It must be called
+// after the WG interface is up so the overlay IP is assigned; when the bound port is
+// not DefaultPort it is announced to peers via Body.mlkemPort.
+func newPQTransport(overlayIP netip.Addr) (*pqTransport, error) {
+	if !overlayIP.IsValid() {
+		return nil, fmt.Errorf("invalid overlay IP for pqkem transport")
+	}
+	ip := net.IP(overlayIP.AsSlice())
+	conn, err := net.ListenUDP("udp4", &net.UDPAddr{IP: ip, Port: DefaultPort})
+	if err != nil {
+		// Default port unavailable (rare on a dedicated overlay IP): fall back to an
+		// ephemeral port, which will be announced to peers.
+		log.Debugf("pqkem: default port %d unavailable on %s (%v), using an ephemeral port", DefaultPort, overlayIP, err)
+		conn, err = net.ListenUDP("udp4", &net.UDPAddr{IP: ip, Port: 0})
+		if err != nil {
+			return nil, fmt.Errorf("bind pqkem udp on overlay %s: %w", overlayIP, err)
+		}
+	}
+	return &pqTransport{
+		conn:   conn,
+		port:   conn.LocalAddr().(*net.UDPAddr).Port,
+		peers:  make(map[string]*net.UDPAddr),
+		byAddr: make(map[string]string),
+	}, nil
+}
+
+// Port is the bound UDP port, announced to peers as Body.mlkemPort.
+func (t *pqTransport) Port() int { return t.port }
+
+func (t *pqTransport) setManager(m *pqkem.Manager) {
+	t.mu.Lock()
+	t.mgr = m
+	t.mu.Unlock()
+}
+
+// AddPeer records where a peer's rekey messages are sent: its overlay IP and the
+// pqkem port it announced. A zero port or invalid IP is ignored.
+func (t *pqTransport) AddPeer(remoteID string, overlayIP netip.Addr, port int) {
+	if !overlayIP.IsValid() || port <= 0 {
+		return
+	}
+	addr := &net.UDPAddr{IP: net.IP(overlayIP.AsSlice()), Port: port}
+	t.mu.Lock()
+	if old, ok := t.peers[remoteID]; ok {
+		delete(t.byAddr, old.String())
+	}
+	t.peers[remoteID] = addr
+	t.byAddr[addr.String()] = remoteID
+	t.mu.Unlock()
+}
+
+func (t *pqTransport) RemovePeer(remoteID string) {
+	t.mu.Lock()
+	if a, ok := t.peers[remoteID]; ok {
+		delete(t.byAddr, a.String())
+		delete(t.peers, remoteID)
+	}
+	t.mu.Unlock()
+}
+
+// SendDataPath implements pqkem.Transport.
+func (t *pqTransport) SendDataPath(remoteID string, msg []byte) error {
+	t.mu.RLock()
+	addr := t.peers[remoteID]
+	t.mu.RUnlock()
+	if addr == nil {
+		return fmt.Errorf("no data-path address for peer %s", remoteID)
+	}
+	_, err := t.conn.WriteToUDP(msg, addr)
+	return err
+}
+
+// run is the receive loop: it maps each datagram's source overlay address to a peer
+// and feeds it to the manager. Exits when the socket is closed.
+func (t *pqTransport) run() {
+	buf := make([]byte, 2048)
+	for {
+		n, src, err := t.conn.ReadFromUDP(buf)
+		if err != nil {
+			return
+		}
+		t.mu.RLock()
+		remoteID := t.byAddr[src.String()]
+		mgr := t.mgr
+		t.mu.RUnlock()
+		if remoteID == "" || mgr == nil {
+			continue
+		}
+		msg := make([]byte, n)
+		copy(msg, buf[:n])
+		if err := mgr.OnDataPathMessage(remoteID, msg); err != nil {
+			log.Debugf("pqkem: inbound from %s: %v", remoteID, err)
+		}
+	}
+}
+
+func (t *pqTransport) Close() error { return t.conn.Close() }