Bit of renaming

peer -> peerAddrs
have types for remoteID and localID
t.Close log error
Manager SetTransport -> Start

client/internal/engine.go

@@ -663,8 +663,8 @@ 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.pqkemManager = pqkem.NewManager(publicKey.String(), pqCallbackHandler{wg: e.wgInterface}, nil)
-			e.pqkemManager.SetTransport(tr)
+			e.pqkemManager = pqkem.NewManager(pqkem.LocalID(publicKey.String()), pqCallbackHandler{wg: e.wgInterface}, nil)
+			e.pqkemManager.Start(tr)
 			log.Infof("ML-KEM post-quantum exchange enabled (udp port %d on overlay %s)", e.pqkemManager.LocalPort(), e.config.WgAddr.IP)
 		}
 	}

client/internal/pqkem/callbacks.go

@@ -9,10 +9,10 @@ type CallbackHandler interface {
 	// and must be programmed into the consumer's secure channel. It is invoked at
 	// the commit point of each side: the initiator on receiving the answer, the
 	// responder on receiving the confirm.
-	OnNewPSKReady(remoteID string, psk PSK) error
+	OnNewPSKReady(remoteID RemoteID, psk PSK) error
 
 	// OnRekeyFailed fires when an exchange fails to converge within the allotted
 	// time. The host should tear the peer connection down so it re-establishes, and
 	// log a WARN. The library reports the event; it does not dictate the reaction.
-	OnRekeyFailed(remoteID string) error
+	OnRekeyFailed(remoteID RemoteID) error
 }

client/internal/pqkem/convergence.go

@@ -9,7 +9,7 @@ import (
 // bootstrap) and returns the framed offer for the caller to send — pushed over the
 // data path for a chained rekey, or handed to the host for signalling when viaSignal
 // is set. Any previous in-flight exchange for the peer is cancelled.
-func (m *Manager) startExchange(remoteID string, viaSignal bool, ackID ExchangeID) ([]byte, error) {
+func (m *Manager) startExchange(remoteID RemoteID, viaSignal bool, ackID ExchangeID) ([]byte, error) {
 	init, err := NewInitiator()
 	if err != nil {
 		return nil, err
@@ -48,7 +48,7 @@ func (m *Manager) startExchange(remoteID string, viaSignal bool, ackID ExchangeI
 // (that offer riding the data path under the freshly adopted key proves it worked),
 // then derives the PSK for the new offer, commits it optimistically, and returns the
 // framed answer. A duplicate offer returns the cached answer without re-deriving.
-func (m *Manager) processOffer(remoteID string, o *OfferMsg) ([]byte, error) {
+func (m *Manager) processOffer(remoteID RemoteID, o *OfferMsg) ([]byte, error) {
 	if o.AckID != (ExchangeID{}) {
 		m.ackConverged(remoteID, o.AckID)
 	}
@@ -97,7 +97,7 @@ func (m *Manager) processOffer(remoteID string, o *OfferMsg) ([]byte, error) {
 // stateAwaitingRekey; the next offer (chained from OnDataPathRekeyed) will acknowledge
 // this exchange. Only valid in stateAwaitingAnswer; advancing the state under the
 // lock makes a concurrent/duplicate answer bail.
-func (m *Manager) processAnswer(remoteID string, a *AnswerMsg) error {
+func (m *Manager) processAnswer(remoteID RemoteID, a *AnswerMsg) error {
 	m.mu.Lock()
 	ex := m.exchanges[remoteID]
 	if ex == nil || ex.id != a.ExchangeID || ex.state != stateAwaitingAnswer {
@@ -126,7 +126,7 @@ func (m *Manager) processAnswer(remoteID string, a *AnswerMsg) error {
 // ackConverged (responder) records convergence of the exchange named by ackID: a
 // later offer acknowledging it proves both sides operate on that exchange's key. Only
 // acts on a matching stateAwaitingAck exchange; anything else is ignored.
-func (m *Manager) ackConverged(remoteID string, ackID ExchangeID) {
+func (m *Manager) ackConverged(remoteID RemoteID, ackID ExchangeID) {
 	m.mu.Lock()
 	ex := m.exchanges[remoteID]
 	if ex == nil || ex.id != ackID || ex.state != stateAwaitingAck {
@@ -146,7 +146,7 @@ func (m *Manager) ackConverged(remoteID string, ackID ExchangeID) {
 // here). Exhausting the deadline before the answer arrives is a failure. Once the
 // answer is in (state past awaitingAnswer) the loop exits: the next rotation is driven
 // by OnDataPathRekeyed, and the idle wait for it has no deadline.
-func (m *Manager) initiatorLoop(ctx context.Context, remoteID string, id ExchangeID) {
+func (m *Manager) initiatorLoop(ctx context.Context, remoteID RemoteID, id ExchangeID) {
 	defer m.wait.Done()
 	t := time.NewTicker(m.retryInterval)
 	defer t.Stop()
@@ -199,7 +199,7 @@ func (m *Manager) initiatorLoop(ctx context.Context, remoteID string, id Exchang
 // an initial exchange (peer never established) fails immediately; a rekey tolerates
 // up to maxRekeyFailures consecutive misses (we stay on the still-valid previous
 // PSK) before failing. Assumes m.mu is held.
-func (m *Manager) registerFailureLocked(remoteID string) bool {
+func (m *Manager) registerFailureLocked(remoteID RemoteID) bool {
 	if !m.established[remoteID] {
 		return true
 	}
@@ -211,7 +211,7 @@ func (m *Manager) registerFailureLocked(remoteID string) bool {
 	return false
 }
 
-func (m *Manager) raiseFailure(remoteID string, fail bool) {
+func (m *Manager) raiseFailure(remoteID RemoteID, fail bool) {
 	if !fail {
 		m.logger.Warn("pqkem rekey attempt timed out, will retry next cycle", "peer", remoteID)
 		return

client/internal/pqkem/convergence_test.go

@@ -25,7 +25,7 @@ func failedCount(f *fakeWG) int {
 func TestManager_InitialTimeoutFailsImmediately(t *testing.T) {
 	wg := newFakeWG()
 	d := NewManager("bbbb", wg, nil) // bbbb > aaaa -> initiator
-	d.SetTransport(dropTransport{})
+	d.Start(dropTransport{})
 	d.retryInterval = 5 * time.Millisecond
 	d.maxRetries = 3
 	defer d.Stop()

client/internal/pqkem/manager.go

@@ -22,11 +22,18 @@ const (
 	DefaultMaxRekeyFailures = 3
 )
 
+// LocalID and RemoteID are peer identity keys (e.g. WireGuard public keys). They are
+// distinct types so the local and a remote identity cannot be mixed up.
+type (
+	LocalID  string
+	RemoteID string
+)
+
 // 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).
+// belongs to the Manager (Run at Start, Close at Stop).
 type Transport interface {
 	// Send delivers msg to the given data-path endpoint.
 	Send(endpoint netip.AddrPort, msg []byte) error
@@ -75,7 +82,7 @@ type exchangeCtl struct {
 // (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
+	localID   LocalID
 	cbHandler CallbackHandler
 	logger    *slog.Logger
 
@@ -88,18 +95,18 @@ type Manager struct {
 
 	mu          sync.Mutex
 	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)
+	exchanges   map[RemoteID]*exchangeCtl   // in-flight exchange per peer
+	established map[RemoteID]bool           // peer has completed at least one exchange
+	failures    map[RemoteID]int            // consecutive rekey failures per peer
+	peerAddrs   map[RemoteID]netip.AddrPort // remoteID -> data-path endpoint (send routing)
+	peersByAddr map[netip.AddrPort]RemoteID // 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(). Set the data-path transport with SetTransport.
-func NewManager(localID string, h CallbackHandler, logger *slog.Logger) *Manager {
+// falls back to slog.Default(). Install the data-path transport with Start.
+func NewManager(localID LocalID, h CallbackHandler, logger *slog.Logger) *Manager {
 	if logger == nil {
 		logger = slog.Default()
 	}
@@ -113,17 +120,17 @@ func NewManager(localID string, h CallbackHandler, logger *slog.Logger) *Manager
 		maxRekeyFailures: DefaultMaxRekeyFailures,
 		rootCtx:          ctx,
 		rootCancel:       cancel,
-		exchanges:        make(map[string]*exchangeCtl),
-		established:      make(map[string]bool),
-		failures:         make(map[string]int),
-		peers:            make(map[string]netip.AddrPort),
-		peersByAddr:      make(map[netip.AddrPort]string),
+		exchanges:        make(map[RemoteID]*exchangeCtl),
+		established:      make(map[RemoteID]bool),
+		failures:         make(map[RemoteID]int),
+		peerAddrs:        make(map[RemoteID]netip.AddrPort),
+		peersByAddr:      make(map[netip.AddrPort]RemoteID),
 	}
 }
 
-// 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) {
+// Start installs the data-path transport and begins its inbound delivery. The Manager
+// owns it from here; Stop closes it. Start/Stop are the transport lifecycle pair.
+func (m *Manager) Start(t Transport) {
 	m.mu.Lock()
 	m.transport = t
 	m.mu.Unlock()
@@ -147,27 +154,27 @@ func (m *Manager) LocalPort() int {
 // 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.
-func (m *Manager) IsInitiator(remoteID string) bool {
-	return m.localID > remoteID
+func (m *Manager) IsInitiator(remoteID RemoteID) bool {
+	return string(m.localID) > string(remoteID)
 }
 
 // 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) {
+func (m *Manager) AddPeer(remoteID RemoteID, endpoint netip.AddrPort) {
 	if !endpoint.IsValid() {
 		return
 	}
 	m.mu.Lock()
-	if old, ok := m.peers[remoteID]; ok {
+	if old, ok := m.peerAddrs[remoteID]; ok {
 		delete(m.peersByAddr, old)
 	}
-	m.peers[remoteID] = endpoint
+	m.peerAddrs[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) {
+func (m *Manager) RemovePeer(remoteID RemoteID) {
 	m.mu.Lock()
 	if ex, ok := m.exchanges[remoteID]; ok {
 		if ex.cancel != nil {
@@ -177,9 +184,9 @@ func (m *Manager) RemovePeer(remoteID string) {
 	}
 	delete(m.established, remoteID)
 	delete(m.failures, remoteID)
-	if ep, ok := m.peers[remoteID]; ok {
+	if ep, ok := m.peerAddrs[remoteID]; ok {
 		delete(m.peersByAddr, ep)
-		delete(m.peers, remoteID)
+		delete(m.peerAddrs, remoteID)
 	}
 	m.mu.Unlock()
 }
@@ -192,10 +199,12 @@ func (m *Manager) Stop() {
 	m.mu.Lock()
 	t := m.transport
 	m.transport = nil
-	m.exchanges = make(map[string]*exchangeCtl)
+	m.exchanges = make(map[RemoteID]*exchangeCtl)
 	m.mu.Unlock()
 	if t != nil {
-		_ = t.Close()
+		if err := t.Close(); err != nil {
+			m.logger.Warn("pqkem: closing data-path transport", "err", err)
+		}
 	}
 }
 
@@ -205,7 +214,7 @@ func (m *Manager) Stop() {
 // remoteID (bootstrap). It returns (nil, nil) when the local peer is not the
 // initiator. It is idempotent for an in-flight bootstrap: a repeat call returns the
 // same offer rather than starting a new exchange.
-func (m *Manager) SignalOffer(remoteID string) ([]byte, error) {
+func (m *Manager) SignalOffer(remoteID RemoteID) ([]byte, error) {
 	if !m.IsInitiator(remoteID) {
 		return nil, nil
 	}
@@ -222,7 +231,7 @@ func (m *Manager) SignalOffer(remoteID string) ([]byte, error) {
 
 // SignalOnOffer processes a KEM offer the host extracted from an incoming offer and
 // returns the KEM answer for the host to embed in its outgoing answer.
-func (m *Manager) SignalOnOffer(remoteID string, offer []byte) ([]byte, error) {
+func (m *Manager) SignalOnOffer(remoteID RemoteID, offer []byte) ([]byte, error) {
 	typ, msg, err := Decode(offer)
 	if err != nil {
 		return nil, fmt.Errorf("decode signal offer from %s: %w", remoteID, err)
@@ -235,7 +244,7 @@ func (m *Manager) SignalOnOffer(remoteID string, offer []byte) ([]byte, error) {
 
 // SignalOnAnswer processes a KEM answer the host extracted from an incoming answer.
 // There is no reply: the next offer (over the data path) acknowledges this exchange.
-func (m *Manager) SignalOnAnswer(remoteID string, answer []byte) error {
+func (m *Manager) SignalOnAnswer(remoteID RemoteID, answer []byte) error {
 	typ, msg, err := Decode(answer)
 	if err != nil {
 		return fmt.Errorf("decode signal answer from %s: %w", remoteID, err)
@@ -252,9 +261,9 @@ func (m *Manager) SignalOnAnswer(remoteID string, answer []byte) error {
 // 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]
+	remoteID, ok := m.peersByAddr[src]
 	m.mu.Unlock()
-	if remoteID == "" {
+	if !ok {
 		return
 	}
 	if err := m.OnDataPathMessage(remoteID, msg); err != nil {
@@ -264,7 +273,7 @@ func (m *Manager) onDataPathInbound(src netip.AddrPort, msg []byte) {
 
 // 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 {
+func (m *Manager) OnDataPathMessage(remoteID RemoteID, raw []byte) error {
 	typ, msg, err := Decode(raw)
 	if err != nil {
 		return fmt.Errorf("decode data-path msg from %s: %w", remoteID, err)
@@ -291,7 +300,7 @@ func (m *Manager) OnDataPathMessage(remoteID string, raw []byte) error {
 // 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) {
+func (m *Manager) OnDataPathRekeyed(remoteID RemoteID) {
 	m.mu.Lock()
 	ex := m.exchanges[remoteID]
 	chain := ex != nil && ex.state == stateAwaitingRekey
@@ -317,15 +326,15 @@ func (m *Manager) OnDataPathRekeyed(remoteID string) {
 // 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) {}
+func (m *Manager) OnDataPathDown(remoteID RemoteID) {}
 
 // ---- internals ----
 
 // 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 {
+func (m *Manager) pushDataPath(remoteID RemoteID, msg []byte) error {
 	m.mu.Lock()
-	ep, ok := m.peers[remoteID]
+	ep, ok := m.peerAddrs[remoteID]
 	t := m.transport
 	m.mu.Unlock()
 	if !ok {
@@ -337,7 +346,7 @@ func (m *Manager) pushDataPath(remoteID string, msg []byte) error {
 	return t.Send(ep, msg)
 }
 
-func (m *Manager) binding(remoteID string) Binding {
+func (m *Manager) binding(remoteID RemoteID) Binding {
 	return Binding{LocalID: []byte(m.localID), RemoteID: []byte(remoteID)}
 }
 

client/internal/pqkem/manager_test.go

@@ -59,27 +59,27 @@ func (l *loopback) Close() error                               { return nil }
 
 type fakeWG struct {
 	mu     sync.Mutex
-	psks   map[string]PSK
-	failed []string
+	psks   map[RemoteID]PSK
+	failed []RemoteID
 }
 
-func newFakeWG() *fakeWG { return &fakeWG{psks: map[string]PSK{}} }
+func newFakeWG() *fakeWG { return &fakeWG{psks: map[RemoteID]PSK{}} }
 
-func (f *fakeWG) OnNewPSKReady(remoteID string, psk PSK) error {
+func (f *fakeWG) OnNewPSKReady(remoteID RemoteID, psk PSK) error {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 	f.psks[remoteID] = psk
 	return nil
 }
 
-func (f *fakeWG) OnRekeyFailed(remoteID string) error {
+func (f *fakeWG) OnRekeyFailed(remoteID RemoteID) error {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 	f.failed = append(f.failed, remoteID)
 	return nil
 }
 
-func (f *fakeWG) psk(peer string) PSK {
+func (f *fakeWG) psk(peer RemoteID) PSK {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 	return f.psks[peer]
@@ -100,9 +100,9 @@ func pair(t *testing.T) (dA, dB *Manager, wgA, wgB *fakeWG, lbB *loopback) {
 	wgB = newFakeWG()
 	dA = NewManager("aaaa", wgA, nil)
 	dB = NewManager("bbbb", wgB, nil)
-	dA.SetTransport(&loopback{ep: epA, sw: sw})
+	dA.Start(&loopback{ep: epA, sw: sw})
 	lbB = &loopback{ep: epB, sw: sw}
-	dB.SetTransport(lbB)
+	dB.Start(lbB)
 	dA.AddPeer("bbbb", epB)
 	dB.AddPeer("aaaa", epA)
 	return dA, dB, wgA, wgB, lbB
@@ -163,7 +163,7 @@ func TestManager_NonInitiatorReturnsNoOffer(t *testing.T) {
 
 func TestManager_StopIsIdempotent(t *testing.T) {
 	dA := NewManager("aaaa", newFakeWG(), nil)
-	dA.SetTransport(&loopback{ep: epA, sw: newSwitch()})
+	dA.Start(&loopback{ep: epA, sw: newSwitch()})
 	dA.Stop()
 	dA.Stop() // must not panic or hang
 }

client/internal/pqkem_adapter.go

@@ -21,13 +21,13 @@ type pqCallbackHandler struct {
 
 // OnNewPSKReady programs the freshly derived PSK for the peer (updateOnly: a no-op
 // if the peer is not present, mirroring Rosenpass). remoteID is the peer's WG pubkey.
-func (h pqCallbackHandler) OnNewPSKReady(remoteID string, psk pqkem.PSK) error {
-	return h.wg.SetPresharedKey(remoteID, wgtypes.Key(psk), true)
+func (h pqCallbackHandler) OnNewPSKReady(remoteID pqkem.RemoteID, psk pqkem.PSK) error {
+	return h.wg.SetPresharedKey(string(remoteID), wgtypes.Key(psk), true)
 }
 
 // OnRekeyFailed reports a failed PQ (re)key convergence.
 // TODO(NET-1406): tear the peer connection down / trigger ICE reconnect.
-func (h pqCallbackHandler) OnRekeyFailed(remoteID string) error {
+func (h pqCallbackHandler) OnRekeyFailed(remoteID pqkem.RemoteID) error {
 	log.Warnf("pqkem: post-quantum rekey failed for peer %s", remoteID)
 	return nil
 }