Makes Transport just a UDP socket.

Manager owns maps for remoteID <-> remote UDP addr
Engine talks to manager only

client/internal/engine.go

@@ -199,9 +199,8 @@ type Engine struct {
 	rpManager *rosenpass.Manager
 
 	// pqkemManager runs the ML-KEM post-quantum PSK exchange (gated by NB_ENABLE_PQ_MLKEM).
+	// It owns the data-path transport and peer endpoint routing.
 	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
@@ -664,11 +663,9 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
 		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)
+			e.pqkemManager = pqkem.NewManager(publicKey.String(), pqCallbackHandler{wg: e.wgInterface}, nil)
+			e.pqkemManager.SetTransport(tr)
+			log.Infof("ML-KEM post-quantum exchange enabled (udp port %d on overlay %s)", e.pqkemManager.LocalPort(), e.config.WgAddr.IP)
 		}
 	}
 
@@ -2104,9 +2101,6 @@ func (e *Engine) close() {
 		_ = e.rpManager.Close()
 	}
 
-	if e.pqTransport != nil {
-		_ = e.pqTransport.Close()
-	}
 	if e.pqkemManager != nil {
 		e.pqkemManager.Stop()
 	}

client/internal/pqkem/convergence_test.go

@@ -1,36 +1,31 @@
 package pqkem
 
 import (
-	"sync/atomic"
+	"net/netip"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
 )
 
+// dropTransport is a pqkem.Transport that silently discards everything.
 type dropTransport struct{}
 
-func (dropTransport) SendDataPath(string, []byte) error { return nil }
+func (dropTransport) Send(netip.AddrPort, []byte) error       { return nil }
+func (dropTransport) LocalPort() int                          { return 0 }
+func (dropTransport) Run(func(netip.AddrPort, []byte))        {}
+func (dropTransport) Close() error                            { return nil }
 
-// gate is a data-path loopback with a switchable drop flag. When dropping it reports
-// success but does not deliver (mimics a lossy/broken tunnel).
-type gate struct {
-	localID string
-	peer    *Manager
-	drop    atomic.Bool
-}
-
-func (g *gate) SendDataPath(remoteID string, msg []byte) error {
-	if g.drop.Load() {
-		return nil
-	}
-	cp := append([]byte(nil), msg...)
-	return g.peer.OnDataPathMessage(g.localID, cp)
+func failedCount(f *fakeWG) int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return len(f.failed)
 }
 
 func TestManager_InitialTimeoutFailsImmediately(t *testing.T) {
 	wg := newFakeWG()
-	d := NewManager("bbbb", dropTransport{}, wg, nil) // bbbb > aaaa -> initiator
+	d := NewManager("bbbb", wg, nil) // bbbb > aaaa -> initiator
+	d.SetTransport(dropTransport{})
 	d.retryInterval = 5 * time.Millisecond
 	d.maxRetries = 3
 	defer d.Stop()
@@ -41,40 +36,25 @@ func TestManager_InitialTimeoutFailsImmediately(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, offer)
 
-	require.Eventually(t, func() bool {
-		wg.mu.Lock()
-		defer wg.mu.Unlock()
-		return len(wg.failed) == 1
-	}, time.Second, 5*time.Millisecond)
+	require.Eventually(t, func() bool { return failedCount(wg) == 1 }, time.Second, 5*time.Millisecond)
 }
 
 func TestManager_RekeyToleratesKFailures(t *testing.T) {
-	gA := &gate{localID: "aaaa"}
-	gB := &gate{localID: "bbbb"}
-	wgA := newFakeWG()
-	wgB := newFakeWG()
-
-	dA := NewManager("aaaa", gA, wgA, nil)
-	dB := NewManager("bbbb", gB, wgB, nil)
-	gA.peer = dB
-	gB.peer = dA
-	dB.retryInterval = 5 * time.Millisecond
-	dB.maxRetries = 2
+	dA, dB, _, wgB, lbB := pair(t)
 	defer dA.Stop()
 	defer dB.Stop()
 
-	// Bootstrap over signalling -> B becomes established.
-	offer, err := dB.SignalOffer("aaaa")
-	require.NoError(t, err)
-	answer, err := dA.SignalOnOffer("bbbb", offer)
-	require.NoError(t, err)
-	require.NoError(t, dB.SignalOnAnswer("aaaa", answer))
-	require.NotEqual(t, PSK{}, wgB.psk("aaaa"))
-
-	// Bring the data path up on both, then drop B's delivery so rekeys can't converge.
+	// Establish: bootstrap + data-path-rekeyed so B becomes established and its data
+	// path is usable.
+	bootstrap(t, dA, dB)
 	dA.OnDataPathRekeyed("bbbb")
 	dB.OnDataPathRekeyed("aaaa")
-	gB.drop.Store(true)
+	require.NotEqual(t, PSK{}, wgB.psk("aaaa"))
+
+	// Tighten timings and drop B's outbound so rekeys can no longer converge.
+	dB.retryInterval = 5 * time.Millisecond
+	dB.maxRetries = 2
+	lbB.drop.Store(true)
 
 	// K-1 data-path rekeys must NOT raise OnRekeyFailed.
 	for i := 0; i < DefaultMaxRekeyFailures-1; i++ {
@@ -85,13 +65,7 @@ func TestManager_RekeyToleratesKFailures(t *testing.T) {
 	require.Equal(t, 0, failedCount(wgB), "no failure before K attempts")
 
 	// The K-th failure raises it once.
-	_, err = dB.startExchange("aaaa", false, ExchangeID{})
+	_, err := dB.startExchange("aaaa", false, ExchangeID{})
 	require.NoError(t, err)
 	require.Eventually(t, func() bool { return failedCount(wgB) == 1 }, time.Second, 5*time.Millisecond)
 }
-
-func failedCount(f *fakeWG) int {
-	f.mu.Lock()
-	defer f.mu.Unlock()
-	return len(f.failed)
-}

client/internal/pqkem/manager.go

@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"fmt"
 	"log/slog"
+	"net/netip"
 	"sync"
 	"time"
 )
@@ -21,12 +22,22 @@ const (
 	DefaultMaxRekeyFailures = 3
 )
 
-// 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
-// host-driven — the library hands the host offer/answer payloads to piggyback on the
-// host's own negotiation, it never pushes there itself.
+// Transport is the data-path socket the Manager drives (the analogue of
+// go-rosenpass's Conn). It is a dumb mover of bytes to/from endpoints: the Manager
+// owns the remoteID<->endpoint routing and hands the transport a resolved endpoint
+// to Send, and reverse-resolves the source of each inbound datagram. Its lifecycle
+// belongs to the Manager (Run at SetTransport, Close at Stop).
 type Transport interface {
-	SendDataPath(remoteID string, msg []byte) error
+	// Send delivers msg to the given data-path endpoint.
+	Send(endpoint netip.AddrPort, msg []byte) error
+	// LocalPort is the bound local UDP port, announced to peers so they know where
+	// to send data-path messages.
+	LocalPort() int
+	// Run starts delivering inbound datagrams as (source endpoint, msg) to onInbound
+	// and returns immediately; it runs until Close.
+	Run(onInbound func(src netip.AddrPort, msg []byte))
+	// Close stops delivery and releases the socket.
+	Close() error
 }
 
 // exchangeState is the single source of truth for an exchange's role and phase.
@@ -58,14 +69,13 @@ type exchangeCtl struct {
 }
 
 // Manager is the stateful orchestrator — the analogue of go-rosenpass's Server. It
-// drives the X25519MLKEM768 exchange and surfaces the derived PSK and convergence to
-// the host via CallbackHandler. It is fully event-driven: the bootstrap is triggered
-// 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.
+// drives the X25519MLKEM768 exchange, owns the peer endpoint routing and the data-path
+// transport, and surfaces the derived PSK and convergence to the host via
+// CallbackHandler. It is event-driven: the bootstrap is triggered by the host
+// (SignalOffer) and each rotation is clocked by OnDataPathRekeyed. The cryptography is
+// the pure kem.go primitives; all state lives here under one lock.
 type Manager struct {
 	localID   string
-	transport Transport
 	cbHandler CallbackHandler
 	logger    *slog.Logger
 
@@ -77,27 +87,25 @@ type Manager struct {
 	rootCancel context.CancelFunc
 
 	mu          sync.Mutex
-	exchanges   map[string]*exchangeCtl // in-flight exchange per peer
-	established map[string]bool         // peer has completed at least one exchange
-	failures    map[string]int          // consecutive rekey failures per peer
-	// dataSend holds the peer's data-path sender when the data path is up; nil
-	// (absent) means it is down. Toggled by OnDataPathRekeyed / OnDataPathDown.
-	dataSend map[string]func(string, []byte) error
-	wait     sync.WaitGroup
+	transport   Transport
+	exchanges   map[string]*exchangeCtl   // in-flight exchange per peer
+	established  map[string]bool          // peer has completed at least one exchange
+	failures    map[string]int            // consecutive rekey failures per peer
+	peers       map[string]netip.AddrPort // remoteID -> data-path endpoint (send routing)
+	peersByAddr map[netip.AddrPort]string // reverse: source endpoint -> remoteID (inbound)
+	wait        sync.WaitGroup
 }
 
 // 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 nil logger
-// falls back to slog.Default(). Retry/retries/K use their defaults and can be
-// overridden before use.
-func NewManager(localID string, t Transport, h CallbackHandler, logger *slog.Logger) *Manager {
+// falls back to slog.Default(). Set the data-path transport with SetTransport.
+func NewManager(localID string, h CallbackHandler, logger *slog.Logger) *Manager {
 	if logger == nil {
 		logger = slog.Default()
 	}
 	ctx, cancel := context.WithCancel(context.Background())
 	return &Manager{
 		localID:          localID,
-		transport:        t,
 		cbHandler:        h,
 		logger:           logger,
 		retryInterval:    DefaultRetryInterval,
@@ -108,10 +116,34 @@ func NewManager(localID string, t Transport, h CallbackHandler, logger *slog.Log
 		exchanges:        make(map[string]*exchangeCtl),
 		established:      make(map[string]bool),
 		failures:         make(map[string]int),
-		dataSend:         make(map[string]func(string, []byte) error),
+		peers:            make(map[string]netip.AddrPort),
+		peersByAddr:      make(map[netip.AddrPort]string),
 	}
 }
 
+// SetTransport installs the data-path transport and starts its inbound delivery. The
+// Manager owns it from here: Stop closes it.
+func (m *Manager) SetTransport(t Transport) {
+	m.mu.Lock()
+	m.transport = t
+	m.mu.Unlock()
+	if t != nil {
+		t.Run(m.onDataPathInbound)
+	}
+}
+
+// LocalPort is the data-path transport's bound UDP port (0 if no transport), to be
+// announced to peers.
+func (m *Manager) LocalPort() int {
+	m.mu.Lock()
+	t := m.transport
+	m.mu.Unlock()
+	if t == nil {
+		return 0
+	}
+	return t.LocalPort()
+}
+
 // IsInitiator reports whether the local peer drives the exchange for this remote
 // peer. Roles are deterministic (lexicographic identity-key compare) so exactly one
 // side initiates, mirroring how Rosenpass picks its handshake initiator.
@@ -119,7 +151,22 @@ func (m *Manager) IsInitiator(remoteID string) bool {
 	return m.localID > remoteID
 }
 
-// RemovePeer stops any in-flight exchange for a peer and drops its state.
+// AddPeer registers where a peer's data-path messages are sent and received: its
+// overlay endpoint (IP:port). Re-adding updates the endpoint.
+func (m *Manager) AddPeer(remoteID string, endpoint netip.AddrPort) {
+	if !endpoint.IsValid() {
+		return
+	}
+	m.mu.Lock()
+	if old, ok := m.peers[remoteID]; ok {
+		delete(m.peersByAddr, old)
+	}
+	m.peers[remoteID] = endpoint
+	m.peersByAddr[endpoint] = remoteID
+	m.mu.Unlock()
+}
+
+// RemovePeer stops any in-flight exchange for a peer and drops its state and routing.
 func (m *Manager) RemovePeer(remoteID string) {
 	m.mu.Lock()
 	if ex, ok := m.exchanges[remoteID]; ok {
@@ -130,17 +177,26 @@ func (m *Manager) RemovePeer(remoteID string) {
 	}
 	delete(m.established, remoteID)
 	delete(m.failures, remoteID)
-	delete(m.dataSend, remoteID)
+	if ep, ok := m.peers[remoteID]; ok {
+		delete(m.peersByAddr, ep)
+		delete(m.peers, remoteID)
+	}
 	m.mu.Unlock()
 }
 
-// Stop cancels all in-flight exchanges and waits for their goroutines to exit.
+// Stop cancels all in-flight exchanges, closes the transport, and waits for the
+// exchange goroutines to exit.
 func (m *Manager) Stop() {
 	m.rootCancel()
 	m.wait.Wait()
 	m.mu.Lock()
+	t := m.transport
+	m.transport = nil
 	m.exchanges = make(map[string]*exchangeCtl)
 	m.mu.Unlock()
+	if t != nil {
+		_ = t.Close()
+	}
 }
 
 // ---- Signalling channel (host-driven; rides the host's negotiation) ----
@@ -190,10 +246,24 @@ func (m *Manager) SignalOnAnswer(remoteID string, answer []byte) error {
 	return m.processAnswer(remoteID, msg.(*AnswerMsg))
 }
 
-// ---- Data path (library-driven push) ----
+// ---- Data path ----
 
-// OnDataPathMessage feeds a KEM message received over the data path (tunnel) and
-// pushes any reply back over the data path.
+// onDataPathInbound is the transport's inbound handler: it reverse-resolves the
+// source endpoint to a peer and dispatches. Unknown sources are dropped.
+func (m *Manager) onDataPathInbound(src netip.AddrPort, msg []byte) {
+	m.mu.Lock()
+	remoteID := m.peersByAddr[src]
+	m.mu.Unlock()
+	if remoteID == "" {
+		return
+	}
+	if err := m.OnDataPathMessage(remoteID, msg); err != nil {
+		m.logger.Debug("pqkem inbound", "peer", remoteID, "err", err)
+	}
+}
+
+// OnDataPathMessage handles a KEM message received over the data path from remoteID
+// and pushes any reply back over the data path.
 func (m *Manager) OnDataPathMessage(remoteID string, raw []byte) error {
 	typ, msg, err := Decode(raw)
 	if err != nil {
@@ -217,13 +287,12 @@ func (m *Manager) OnDataPathMessage(remoteID string, raw []byte) error {
 }
 
 // OnDataPathRekeyed notifies that the peer's data path is up and freshly keyed with
-// the latest PSK (fired on first establishment AND every rekey). It marks the data
-// path usable and, if we are the initiator that just derived a PSK, chains the next
-// exchange: a fresh offer over the data path that acknowledges the just-completed one
-// (its arrival under the new key proves to the responder that the key works).
+// the latest PSK (fired on first establishment AND every rekey). If we are the
+// initiator that just derived a PSK, it chains the next exchange: a fresh offer over
+// the data path that acknowledges the just-completed one (its arrival under the new
+// key proves to the responder that the key works).
 func (m *Manager) OnDataPathRekeyed(remoteID string) {
 	m.mu.Lock()
-	m.dataSend[remoteID] = m.transport.SendDataPath
 	ex := m.exchanges[remoteID]
 	chain := ex != nil && ex.state == stateAwaitingRekey
 	var ackID ExchangeID
@@ -245,25 +314,27 @@ func (m *Manager) OnDataPathRekeyed(remoteID string) {
 	}
 }
 
-// OnDataPathDown notifies that the peer's data path went down; rotations pause until
-// it is up again (the host re-bootstraps over signalling on reconnect).
-func (m *Manager) OnDataPathDown(remoteID string) {
-	m.mu.Lock()
-	delete(m.dataSend, remoteID)
-	m.mu.Unlock()
-}
+// OnDataPathDown notifies that the peer's data path went down. Rotations resume once
+// the host re-bootstraps over signalling on reconnect; in-flight data-path sends will
+// simply fail until then. Reserved as an explicit hook.
+func (m *Manager) OnDataPathDown(remoteID string) {}
 
 // ---- internals ----
 
-// pushDataPath sends over the peer's data path, erroring if it is down.
+// pushDataPath resolves the peer's endpoint and sends over the data-path transport,
+// erroring if the peer is unknown or no transport is set.
 func (m *Manager) pushDataPath(remoteID string, msg []byte) error {
 	m.mu.Lock()
-	send := m.dataSend[remoteID]
+	ep, ok := m.peers[remoteID]
+	t := m.transport
 	m.mu.Unlock()
-	if send == nil {
-		return fmt.Errorf("no data path for peer %s", remoteID)
+	if !ok {
+		return fmt.Errorf("no data-path endpoint for peer %s", remoteID)
 	}
-	return send(remoteID, msg)
+	if t == nil {
+		return fmt.Errorf("no data-path transport")
+	}
+	return t.Send(ep, msg)
 }
 
 func (m *Manager) binding(remoteID string) Binding {

client/internal/pqkem/manager_test.go

@@ -1,25 +1,62 @@
 package pqkem
 
 import (
+	"fmt"
+	"net/netip"
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
-// loopback is a data-path transport: SendDataPath delivers synchronously to the peer
-// manager's OnDataPathMessage, attributing it to localID (the sender). The signalling
-// channel is driven by the test directly via the SignalX methods.
-type loopback struct {
-	localID string
-	peer    *Manager
+// netSwitch is an in-memory UDP fabric: transports register their endpoint and get
+// datagrams delivered to their inbound handler.
+type netSwitch struct {
+	mu sync.Mutex
+	h  map[netip.AddrPort]func(netip.AddrPort, []byte)
 }
 
-func (l *loopback) SendDataPath(remoteID string, msg []byte) error {
-	cp := append([]byte(nil), msg...)
-	return l.peer.OnDataPathMessage(l.localID, cp)
+func newSwitch() *netSwitch {
+	return &netSwitch{h: map[netip.AddrPort]func(netip.AddrPort, []byte){}}
 }
 
+func (s *netSwitch) register(ep netip.AddrPort, fn func(netip.AddrPort, []byte)) {
+	s.mu.Lock()
+	s.h[ep] = fn
+	s.mu.Unlock()
+}
+
+func (s *netSwitch) deliver(dst, src netip.AddrPort, msg []byte) error {
+	s.mu.Lock()
+	fn := s.h[dst]
+	s.mu.Unlock()
+	if fn == nil {
+		return fmt.Errorf("no route to %s", dst)
+	}
+	fn(src, msg)
+	return nil
+}
+
+// loopback is an endpoint-based pqkem.Transport over a netSwitch, with a switchable
+// drop flag.
+type loopback struct {
+	ep   netip.AddrPort
+	sw   *netSwitch
+	drop atomic.Bool
+}
+
+func (l *loopback) Send(dst netip.AddrPort, msg []byte) error {
+	if l.drop.Load() {
+		return nil
+	}
+	return l.sw.deliver(dst, l.ep, append([]byte(nil), msg...))
+}
+
+func (l *loopback) LocalPort() int                             { return int(l.ep.Port()) }
+func (l *loopback) Run(onInbound func(netip.AddrPort, []byte)) { l.sw.register(l.ep, onInbound) }
+func (l *loopback) Close() error                               { return nil }
+
 type fakeWG struct {
 	mu     sync.Mutex
 	psks   map[string]PSK
@@ -48,18 +85,27 @@ func (f *fakeWG) psk(peer string) PSK {
 	return f.psks[peer]
 }
 
-// pair builds two wired managers (B is the initiator, "bbbb" > "aaaa").
-func pair(t *testing.T) (dA, dB *Manager, wgA, wgB *fakeWG) {
+var (
+	epA = netip.MustParseAddrPort("100.64.0.1:51833")
+	epB = netip.MustParseAddrPort("100.64.0.2:51833")
+)
+
+// pair builds two wired managers (B is the initiator, "bbbb" > "aaaa") sharing a
+// netSwitch, with each peer's data-path endpoint registered. lbB is B's loopback
+// (for toggling drop).
+func pair(t *testing.T) (dA, dB *Manager, wgA, wgB *fakeWG, lbB *loopback) {
 	t.Helper()
-	lbA := &loopback{localID: "aaaa"}
-	lbB := &loopback{localID: "bbbb"}
+	sw := newSwitch()
 	wgA = 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", wgA, nil)
+	dB = NewManager("bbbb", wgB, nil)
+	dA.SetTransport(&loopback{ep: epA, sw: sw})
+	lbB = &loopback{ep: epB, sw: sw}
+	dB.SetTransport(lbB)
+	dA.AddPeer("bbbb", epB)
+	dB.AddPeer("aaaa", epA)
+	return dA, dB, wgA, wgB, lbB
 }
 
 // bootstrap runs the signalling offer/answer (the test plays the host carrying bytes).
@@ -75,7 +121,7 @@ func bootstrap(t *testing.T, dA, dB *Manager) {
 }
 
 func TestManager_BootstrapDerivesSamePSK(t *testing.T) {
-	dA, dB, wgA, wgB := pair(t)
+	dA, dB, wgA, wgB, _ := pair(t)
 	defer dA.Stop()
 	defer dB.Stop()
 
@@ -88,15 +134,15 @@ func TestManager_BootstrapDerivesSamePSK(t *testing.T) {
 }
 
 func TestManager_ChainRotatesAndAcks(t *testing.T) {
-	dA, dB, wgA, wgB := pair(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.
+	// Data path up: B (initiator) chains the next offer over the data path, which
+	// rotates both to a fresh PSK and acknowledges A.
 	dA.OnDataPathRekeyed("bbbb")
 	dB.OnDataPathRekeyed("aaaa")
 
@@ -107,7 +153,7 @@ func TestManager_ChainRotatesAndAcks(t *testing.T) {
 }
 
 func TestManager_NonInitiatorReturnsNoOffer(t *testing.T) {
-	dA := NewManager("aaaa", &loopback{localID: "aaaa"}, newFakeWG(), nil)
+	dA := NewManager("aaaa", newFakeWG(), nil)
 	defer dA.Stop()
 
 	offer, err := dA.SignalOffer("bbbb") // not the initiator vs "bbbb"
@@ -116,7 +162,8 @@ func TestManager_NonInitiatorReturnsNoOffer(t *testing.T) {
 }
 
 func TestManager_StopIsIdempotent(t *testing.T) {
-	dA := NewManager("aaaa", &loopback{localID: "aaaa"}, newFakeWG(), nil)
+	dA := NewManager("aaaa", newFakeWG(), nil)
+	dA.SetTransport(&loopback{ep: epA, sw: newSwitch()})
 	dA.Stop()
 	dA.Stop() // must not panic or hang
 }

client/internal/pqkem_transport.go

@@ -4,11 +4,8 @@ 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
@@ -17,25 +14,18 @@ import (
 // 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).
+// pqTransport is the ML-KEM data-path transport: a dumb UDP socket bound on the WG
+// overlay IP. It implements pqkem.Transport — the manager owns the remoteID<->endpoint
+// routing and drives this socket's lifecycle (Run / Close).
 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.
+// falling back to an OS-assigned ephemeral port if it is in use. Call it after the WG
+// interface is up so the overlay IP is assigned; when the bound port is not
+// DefaultPort it must be 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")
@@ -43,90 +33,40 @@ func newPQTransport(overlayIP netip.Addr) (*pqTransport, error) {
 	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
+	return &pqTransport{conn: conn, port: conn.LocalAddr().(*net.UDPAddr).Port}, 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)
+// Send implements pqkem.Transport.
+func (t *pqTransport) Send(endpoint netip.AddrPort, msg []byte) error {
+	_, err := t.conn.WriteToUDPAddrPort(msg, endpoint)
 	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
+// LocalPort implements pqkem.Transport.
+func (t *pqTransport) LocalPort() int { return t.port }
+
+// Run implements pqkem.Transport: the receive loop, delivering each datagram as
+// (source endpoint, msg). Exits when the socket is closed.
+func (t *pqTransport) Run(onInbound func(src netip.AddrPort, msg []byte)) {
+	go func() {
+		buf := make([]byte, 2048)
+		for {
+			n, src, err := t.conn.ReadFromUDPAddrPort(buf)
+			if err != nil {
+				return
+			}
+			msg := make([]byte, n)
+			copy(msg, buf[:n])
+			onInbound(src, msg)
 		}
-		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)
-		}
-	}
+	}()
 }
 
+// Close implements pqkem.Transport.
 func (t *pqTransport) Close() error { return t.conn.Close() }