From 9074f36761bb6c74bd4a9124d3f554549079d720 Mon Sep 17 00:00:00 2001 From: riccardom Date: Sat, 25 Jul 2026 11:30:39 +0200 Subject: [PATCH] Adds transport --- client/internal/engine.go | 22 +++-- client/internal/pqkem_adapter.go | 8 -- client/internal/pqkem_transport.go | 132 +++++++++++++++++++++++++++++ 3 files changed, 149 insertions(+), 13 deletions(-) create mode 100644 client/internal/pqkem_transport.go diff --git a/client/internal/engine.go b/client/internal/engine.go index f1715830f..a03ff40f9 100644 --- a/client/internal/engine.go +++ b/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() } diff --git a/client/internal/pqkem_adapter.go b/client/internal/pqkem_adapter.go index 7410915ea..1a94c65a4 100644 --- a/client/internal/pqkem_adapter.go +++ b/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 } diff --git a/client/internal/pqkem_transport.go b/client/internal/pqkem_transport.go new file mode 100644 index 000000000..da2e24f0d --- /dev/null +++ b/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() }