Convey MDM overlay config to Debug Bundle output

Aligns to other clients OSes behavior

client/android/client.go

@@ -23,6 +23,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
 	"github.com/netbirdio/netbird/client/internal/routemanager"
 	"github.com/netbirdio/netbird/client/internal/stdnet"
+	"github.com/netbirdio/netbird/client/mdm"
 	"github.com/netbirdio/netbird/client/net"
 	"github.com/netbirdio/netbird/client/system"
 	"github.com/netbirdio/netbird/formatter"
@@ -75,6 +76,13 @@ type Client struct {
 	connectClient *internal.ConnectClient
 	config        *profilemanager.Config
 	cacheDir      string
+
+	// mdmLoader holds the per-Client MDM policy source. Set by
+	// SetMDMPolicyFetcher (called from the Kotlin side). Each Run
+	// passes this loader to the resolved Config so applyMDMPolicy
+	// picks up the active overlay. Nil means "MDM enforcement off
+	// for this Client".
+	mdmLoader *mdm.Loader
 }
 
 func (c *Client) setState(cfg *profilemanager.Config, cacheDir string, cc *internal.ConnectClient) {
@@ -129,6 +137,7 @@ func (c *Client) Run(platformFiles PlatformFiles, urlOpener URLOpener, isAndroid
 	if err != nil {
 		return err
 	}
+	c.applyMDMOverlay(cfg)
 	c.recorder.UpdateManagementAddress(cfg.ManagementURL.String())
 	c.recorder.UpdateRosenpass(cfg.RosenpassEnabled, cfg.RosenpassPermissive)
 
@@ -173,6 +182,7 @@ func (c *Client) RunWithoutLogin(platformFiles PlatformFiles, dns *DNSList, dnsR
 	if err != nil {
 		return err
 	}
+	c.applyMDMOverlay(cfg)
 	c.recorder.UpdateManagementAddress(cfg.ManagementURL.String())
 	c.recorder.UpdateRosenpass(cfg.RosenpassEnabled, cfg.RosenpassPermissive)
 
@@ -230,6 +240,7 @@ func (c *Client) DebugBundle(platformFiles PlatformFiles, anonymize bool) (strin
 		if err != nil {
 			return "", fmt.Errorf("load config: %w", err)
 		}
+		c.applyMDMOverlay(cfg)
 		cacheDir = platformFiles.CacheDir()
 	}
 

client/android/mdm.go

@@ -0,0 +1,80 @@
+//go:build android
+
+package android
+
+import (
+	"encoding/json"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/internal/profilemanager"
+	"github.com/netbirdio/netbird/client/mdm"
+)
+
+// PolicyFetcher is the mobile-side bridge for the MDM managed-config
+// snapshot. The native layer (Kotlin) implements this and registers
+// the instance per Client via Client.SetMDMPolicyFetcher. Every
+// invocation of fetchJSON must read the current RestrictionsManager
+// state and return the result as a JSON-encoded map[string]any string.
+//
+// JSON is used because gomobile does not support map[string]any
+// crossing the JNI boundary — the adapter on the Go side parses the
+// string back into the map[string]any expected by mdm.Loader.
+//
+// Return value contract:
+//   - "" (empty)         : interpreted as "no MDM source / no managed keys"
+//   - "{}"               : managed config explicitly empty
+//   - "{...}"            : JSON object with key/value pairs
+//   - malformed JSON     : logged and treated as empty
+type PolicyFetcher interface {
+	FetchJSON() string
+}
+
+// jsonFetcherAdapter wraps a gomobile-exposed PolicyFetcher into the
+// internal mdm.PolicyFetcher interface, taking care of JSON decoding
+// on every Fetch.
+type jsonFetcherAdapter struct {
+	inner PolicyFetcher
+}
+
+func (a *jsonFetcherAdapter) Fetch() map[string]any {
+	raw := a.inner.FetchJSON()
+	if raw == "" {
+		return nil
+	}
+	var out map[string]any
+	if err := json.Unmarshal([]byte(raw), &out); err != nil {
+		log.Warnf("MDM mobile fetcher: invalid JSON payload from native: %v", err)
+		return nil
+	}
+	return out
+}
+
+// SetMDMPolicyFetcher registers the native-provided MDM policy fetcher
+// on this Client. Call once from the gomobile-init code (Kotlin
+// Application.onCreate or Service onCreate) before invoking Run /
+// RunWithoutLogin. Passing nil disables MDM enforcement on this
+// Client.
+//
+// The fetcher is held as a *mdm.Loader instance on the Client (no
+// package-level state) — multiple Clients in the same process get
+// independent Loaders, and tests can inject fakes per Client.
+func (c *Client) SetMDMPolicyFetcher(p PolicyFetcher) {
+	if p == nil {
+		c.mdmLoader = mdm.NewLoader(nil)
+		return
+	}
+	c.mdmLoader = mdm.NewLoader(&jsonFetcherAdapter{inner: p})
+}
+
+// applyMDMOverlay applies the Client-held MDM Loader's current policy
+// on top of the just-read Config. Called immediately after every
+// UpdateOrCreateConfig — profilemanager's apply() initialises the
+// policy to empty and leaves overlay responsibility to the lifecycle
+// owner. No-op when no fetcher was registered.
+func (c *Client) applyMDMOverlay(cfg *profilemanager.Config) {
+	if cfg == nil || c.mdmLoader == nil {
+		return
+	}
+	cfg.ApplyMDMPolicy(c.mdmLoader.Load())
+}

client/cmd/login.go

@@ -17,6 +17,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal"
 	"github.com/netbirdio/netbird/client/internal/auth"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
+	"github.com/netbirdio/netbird/client/mdm"
 	nbnet "github.com/netbirdio/netbird/client/net"
 	"github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/client/server"
@@ -332,6 +333,11 @@ func doForegroundLogin(ctx context.Context, cmd *cobra.Command, setupKey string,
 	if err != nil {
 		return fmt.Errorf("read config file %s: %v", configFilePath, err)
 	}
+	// CLI standalone login: profilemanager no longer auto-applies MDM,
+	// so layer in the OS-native policy here. Desktop builds construct
+	// a Loader with no fetcher — the build-tagged loadPlatform reads
+	// the registry/plist directly.
+	config.ApplyMDMPolicy(mdm.NewLoader(nil).Load())
 
 	// Mirror runInForegroundMode: recover residual state (DNS, firewall,
 	// ssh config, legacy routing) from a previous unclean shutdown and

client/cmd/up.go

@@ -21,6 +21,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
+	"github.com/netbirdio/netbird/client/mdm"
 	"github.com/netbirdio/netbird/client/proto"
 	nbnet "github.com/netbirdio/netbird/client/net"
 	"github.com/netbirdio/netbird/client/server"
@@ -228,6 +229,10 @@ func runInForegroundMode(ctx context.Context, cmd *cobra.Command, activeProf *pr
 	if err != nil {
 		return fmt.Errorf("get config file: %v", err)
 	}
+	// CLI foreground path runs without the daemon Server: layer in the
+	// active MDM policy explicitly so a forced ManagementURL / PSK /
+	// other managed key actually takes effect on this run.
+	config.ApplyMDMPolicy(mdm.NewLoader(nil).Load())
 
 	_, _ = profilemanager.UpdateOldManagementURL(ctx, config, configFilePath)
 

client/embed/embed.go

@@ -21,6 +21,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/auth"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
+	"github.com/netbirdio/netbird/client/mdm"
 	sshcommon "github.com/netbirdio/netbird/client/ssh"
 	"github.com/netbirdio/netbird/client/system"
 	"github.com/netbirdio/netbird/shared/management/domain"
@@ -215,6 +216,10 @@ func New(opts Options) (*Client, error) {
 	if err != nil {
 		return nil, fmt.Errorf("create config: %w", err)
 	}
+	// Embedded path runs without the daemon Server: apply the active
+	// MDM policy explicitly so a forced ManagementURL / PSK / other
+	// managed key takes effect on this embedded engine instance.
+	config.ApplyMDMPolicy(mdm.NewLoader(nil).Load())
 
 	if opts.PrivateKey != "" {
 		config.PrivateKey = opts.PrivateKey

client/internal/engine.go

@@ -50,7 +50,6 @@ import (
 	icemaker "github.com/netbirdio/netbird/client/internal/peer/ice"
 	"github.com/netbirdio/netbird/client/internal/peerstore"
 	"github.com/netbirdio/netbird/client/internal/portforward"
-	"github.com/netbirdio/netbird/client/internal/pqkem"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
 	"github.com/netbirdio/netbird/client/internal/relay"
 	"github.com/netbirdio/netbird/client/internal/rosenpass"
@@ -198,10 +197,6 @@ type Engine struct {
 	// rpManager is a Rosenpass manager
 	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
-
 	// syncMsgMux is used to guarantee sequential Management Service message processing
 	syncMsgMux *sync.Mutex
 
@@ -656,19 +651,6 @@ 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.
-	if pqkem.Enabled() {
-		tr, pqErr := newPQTransport(e.config.WgAddr.IP)
-		if pqErr != nil {
-			log.Errorf("pqkem: transport bind failed, exchange disabled: %v", pqErr)
-		} else {
-			e.pqkemManager = pqkem.NewManager(pqkem.LocalID(publicKey.String()), pqCallbackHandler{wg: e.wgInterface}, pqkem.NewLogger())
-			e.pqkemManager.Start(tr)
-			log.Infof("pqkem: enabled (udp port %d on overlay %s)", e.pqkemManager.LocalPort(), e.config.WgAddr.IP)
-		}
-	}
-
 	// if inbound conns are blocked there is no need to create the ACL manager
 	if e.firewall != nil && !e.config.BlockInbound {
 		e.acl = acl.NewDefaultManager(e.firewall)
@@ -932,10 +914,6 @@ 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)
@@ -1922,9 +1900,6 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentV
 		},
 		ICEConfig: e.createICEConfig(),
 	}
-	if e.pqkemManager != nil {
-		config.PQ = pqHandshaker{mgr: e.pqkemManager}
-	}
 
 	serviceDependencies := peer.ServiceDependencies{
 		StatusRecorder:     e.statusRecorder,
@@ -2108,10 +2083,6 @@ func (e *Engine) close() {
 		_ = e.rpManager.Close()
 	}
 
-	if e.pqkemManager != nil {
-		e.pqkemManager.Stop()
-	}
-
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 	if err := e.portForwardManager.GracefullyStop(ctx); err != nil {
@@ -2924,8 +2895,6 @@ func convertToOfferAnswer(msg *sProto.Message) (*peer.OfferAnswer, error) {
 		Version:         msg.GetBody().GetNetBirdVersion(),
 		RosenpassPubKey: rosenpassPubKey,
 		RosenpassAddr:   rosenpassAddr,
-		MlkemPayload:    msg.GetBody().GetMlkemPayload(),
-		MlkemPort:       int(msg.GetBody().GetMlkemPort()),
 		RelaySrvAddress: msg.GetBody().GetRelayServerAddress(),
 		RelaySrvIP:      relayIP,
 		SessionID:       sessionID,

client/internal/peer/conn.go

@@ -74,32 +74,6 @@ type RosenpassConfig struct {
 	PermissiveMode bool
 }
 
-// PQHandshaker attaches post-quantum ML-KEM material to signalling offers/answers and
-// feeds received material back. It is implemented by the engine over the pqkem
-// manager and is nil when the PQ exchange is disabled. remoteKey is the peer's
-// WireGuard public key.
-type PQHandshaker interface {
-	// OfferPayload returns the KEM offer to embed in an outgoing offer (nil if this
-	// peer is not the KEM initiator) and the local PQ data-path port to announce.
-	OfferPayload(remoteKey string) (payload []byte, port int)
-	// AnswerPayload processes a received KEM offer (nil if absent) and returns the KEM
-	// answer to embed in the outgoing answer (nil if none) and the local PQ port.
-	AnswerPayload(remoteKey string, recvOffer []byte) (payload []byte, port int)
-	// OnAnswer feeds a received KEM answer (nil if absent).
-	OnAnswer(remoteKey string, recvAnswer []byte)
-	// 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)
-	// SetRemoteAddr registers the peer's data-path endpoint learned from signalling:
-	// its WG overlay IP with the advertised pq UDP port.
-	SetRemoteAddr(remoteKey string, addr netip.AddrPort)
-	// OnDataPathRekeyed signals a fresh WireGuard handshake for the peer; it clocks the
-	// next chained PSK rotation pushed over the data path.
-	OnDataPathRekeyed(remoteKey string)
-	// OnDataPathDown signals the peer's tunnel went down.
-	OnDataPathDown(remoteKey string)
-}
-
 // ConnConfig is a peer Connection configuration
 type ConnConfig struct {
 	// Key is a public key of a remote peer
@@ -117,9 +91,6 @@ type ConnConfig struct {
 
 	RosenpassConfig RosenpassConfig
 
-	// PQ carries post-quantum ML-KEM material on offers/answers; nil when disabled.
-	PQ PQHandshaker
-
 	// ICEConfig ICE protocol configuration
 	ICEConfig icemaker.Config
 }
@@ -710,10 +681,6 @@ func (conn *Conn) onWGDisconnected(watcherCtx context.Context) {
 
 	conn.Log.Warnf("WireGuard handshake timeout detected, closing current connection")
 
-	if conn.config.PQ != nil {
-		conn.config.PQ.OnDataPathDown(conn.config.Key)
-	}
-
 	// Close the active connection based on current priority
 	switch conn.currentConnPriority {
 	case conntype.Relay:
@@ -979,11 +946,6 @@ func (conn *Conn) onWGCheckSuccess() {
 	conn.mu.Lock()
 	conn.wgTimeouts = 0
 	conn.mu.Unlock()
-
-	// A fresh WireGuard handshake is the clock for the post-quantum PSK rotation.
-	if conn.config.PQ != nil {
-		conn.config.PQ.OnDataPathRekeyed(conn.config.Key)
-	}
 }
 
 // recordConnectionMetrics records connection stage timestamps as metrics
@@ -1025,15 +987,6 @@ func (conn *Conn) AgentVersionString() string {
 }
 
 func (conn *Conn) presharedKey(remoteRosenpassKey []byte) *wgtypes.Key {
-	// Post-quantum: once the ML-KEM exchange has derived a PSK for this peer, program
-	// it here so the peer's next WireGuard handshake adopts it. Applied at peer-config
-	// time (bootstrap / reconnect); steady-state rotation is pushed separately.
-	if conn.config.PQ != nil {
-		if psk, ok := conn.config.PQ.PSK(conn.config.Key); ok {
-			return &psk
-		}
-	}
-
 	if conn.config.RosenpassConfig.PubKey == nil {
 		return conn.config.WgConfig.PreSharedKey
 	}

client/internal/peer/handshaker.go

@@ -39,16 +39,6 @@ type OfferAnswer struct {
 	// This value is the local Rosenpass server address when sending the message
 	RosenpassAddr string
 
-	// MlkemPayload carries the post-quantum X25519MLKEM768 handshake message
-	// (pqkem-framed offer on an OFFER, answer on an ANSWER) that seeds the
-	// WireGuard PSK. Opaque here — the pqkem library frames and parses it. Nil
-	// when the peer does not run the ML-KEM PQ exchange.
-	MlkemPayload []byte
-
-	// MlkemPort is the peer's ML-KEM PQ service UDP port (bound on its WG overlay
-	// IP) where data-path rekey messages are sent. Zero when not running the exchange.
-	MlkemPort int
-
 	// relay server address
 	RelaySrvAddress string
 	// RelaySrvIP is the IP the remote peer is connected to on its
@@ -130,8 +120,6 @@ func (h *Handshaker) Listen(ctx context.Context) {
 
 			h.updateRemoteICEState(&remoteOfferAnswer)
 
-			h.pqRegisterEndpoint(remoteOfferAnswer.MlkemPort)
-
 			if h.relayListener != nil {
 				h.relayListener.Notify(&remoteOfferAnswer)
 			}
@@ -140,7 +128,7 @@ func (h *Handshaker) Listen(ctx context.Context) {
 				h.iceListener(&remoteOfferAnswer)
 			}
 
-			if err := h.sendAnswer(&remoteOfferAnswer); err != nil {
+			if err := h.sendAnswer(); err != nil {
 				h.log.Errorf("failed to send remote offer confirmation: %s", err)
 				continue
 			}
@@ -154,8 +142,6 @@ func (h *Handshaker) Listen(ctx context.Context) {
 
 			h.updateRemoteICEState(&remoteOfferAnswer)
 
-			h.pqRegisterEndpoint(remoteOfferAnswer.MlkemPort)
-
 			if h.relayListener != nil {
 				h.relayListener.Notify(&remoteOfferAnswer)
 			}
@@ -163,10 +149,6 @@ func (h *Handshaker) Listen(ctx context.Context) {
 			if h.iceListener != nil && h.RemoteICESupported() {
 				h.iceListener(&remoteOfferAnswer)
 			}
-
-			if h.config.PQ != nil {
-				h.config.PQ.OnAnswer(h.config.Key, remoteOfferAnswer.MlkemPayload)
-			}
 		case <-ctx.Done():
 			h.log.Infof("stop listening for remote offers and answers")
 			return
@@ -174,16 +156,6 @@ 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 || remotePort > 65535 || len(h.config.WgConfig.AllowedIps) == 0 {
-		return
-	}
-	addr := netip.AddrPortFrom(h.config.WgConfig.AllowedIps[0].Addr(), uint16(remotePort))
-	h.config.PQ.SetRemoteAddr(h.config.Key, addr)
-}
-
 func (h *Handshaker) SendOffer() error {
 	h.mu.Lock()
 	defer h.mu.Unlock()
@@ -223,23 +195,13 @@ func (h *Handshaker) sendOffer() error {
 	}
 
 	offer := h.buildOfferAnswer()
-	if h.config.PQ != nil {
-		offer.MlkemPayload, offer.MlkemPort = h.config.PQ.OfferPayload(h.config.Key)
-	}
 	h.log.Debugf("sending offer with serial: %s", offer.SessionIDString())
 
 	return h.signaler.SignalOffer(offer, h.config.Key)
 }
 
-func (h *Handshaker) sendAnswer(remoteOffer *OfferAnswer) error {
+func (h *Handshaker) sendAnswer() error {
 	answer := h.buildOfferAnswer()
-	if h.config.PQ != nil {
-		var recvOffer []byte
-		if remoteOffer != nil {
-			recvOffer = remoteOffer.MlkemPayload
-		}
-		answer.MlkemPayload, answer.MlkemPort = h.config.PQ.AnswerPayload(h.config.Key, recvOffer)
-	}
 	h.log.Debugf("sending answer with serial: %s", answer.SessionIDString())
 
 	return h.signaler.SignalAnswer(answer, h.config.Key)

client/internal/peer/signaler.go

@@ -63,8 +63,6 @@ func (s *Signaler) signalOfferAnswer(offerAnswer OfferAnswer, remoteKey string,
 		},
 		RosenpassPubKey: offerAnswer.RosenpassPubKey,
 		RosenpassAddr:   offerAnswer.RosenpassAddr,
-		MlkemPayload:    offerAnswer.MlkemPayload,
-		MlkemPort:       offerAnswer.MlkemPort,
 		RelaySrvAddress: offerAnswer.RelaySrvAddress,
 		RelaySrvIP:      offerAnswer.RelaySrvIP,
 		SessionID:       sessionIDBytes,

client/internal/pqkem/bench_test.go

@@ -1,59 +0,0 @@
-package pqkem
-
-import (
-	"crypto/ecdh"
-	"crypto/mlkem"
-	"crypto/rand"
-	"testing"
-)
-
-func BenchmarkX25519Keygen(b *testing.B) {
-	c := ecdh.X25519()
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		if _, err := c.GenerateKey(rand.Reader); err != nil {
-			b.Fatal(err)
-		}
-	}
-}
-
-func BenchmarkX25519ECDH(b *testing.B) {
-	c := ecdh.X25519()
-	a, _ := c.GenerateKey(rand.Reader)
-	p, _ := c.GenerateKey(rand.Reader)
-	pub := p.PublicKey()
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		if _, err := a.ECDH(pub); err != nil {
-			b.Fatal(err)
-		}
-	}
-}
-
-func BenchmarkMLKEMKeygen(b *testing.B) {
-	for i := 0; i < b.N; i++ {
-		if _, err := mlkem.GenerateKey768(); err != nil {
-			b.Fatal(err)
-		}
-	}
-}
-
-func BenchmarkMLKEMEncaps(b *testing.B) {
-	dk, _ := mlkem.GenerateKey768()
-	ek := dk.EncapsulationKey()
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		_, _ = ek.Encapsulate()
-	}
-}
-
-func BenchmarkMLKEMDecaps(b *testing.B) {
-	dk, _ := mlkem.GenerateKey768()
-	_, ct := dk.EncapsulationKey().Encapsulate()
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		if _, err := dk.Decapsulate(ct); err != nil {
-			b.Fatal(err)
-		}
-	}
-}

client/internal/pqkem/callbacks.go

@@ -1,18 +0,0 @@
-package pqkem
-
-// CallbackHandler is implemented by the host and invoked by the library. The
-// library only reports events; the host owns the reaction. Keeping this an
-// interface — rather than touching the transport or keying directly — is what lets
-// the KEM code be extracted as a standalone library.
-type CallbackHandler interface {
-	// OnNewPSKReady fires when a fresh post-quantum PSK has been derived for a peer
-	// 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 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 RemoteID) error
-}

client/internal/pqkem/convergence.go

@@ -1,224 +0,0 @@
-package pqkem
-
-import (
-	"context"
-	"time"
-)
-
-// startExchange creates a fresh initiator exchange (acknowledging ackID, zero for a
-// 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 RemoteID, viaSignal bool, ackID ExchangeID) ([]byte, error) {
-	init, err := NewInitiator()
-	if err != nil {
-		return nil, err
-	}
-	id, err := newExchangeID()
-	if err != nil {
-		return nil, err
-	}
-	raw, err := (&OfferMsg{ExchangeID: id, AckID: ackID, KEMOffer: init.Offer()}).Encode()
-	if err != nil {
-		return nil, err
-	}
-
-	ctx, cancel := context.WithCancel(m.rootCtx)
-	m.mu.Lock()
-	if old := m.exchanges[remoteID]; old != nil && old.cancel != nil {
-		old.cancel()
-	}
-	m.exchanges[remoteID] = &exchangeCtl{
-		id:        id,
-		state:     stateAwaitingAnswer,
-		startedAt: time.Now(),
-		cancel:    cancel,
-		lastSent:  raw,
-		initiator: init,
-		viaSignal: viaSignal,
-	}
-	m.mu.Unlock()
-
-	m.wait.Add(1)
-	go m.initiatorLoop(ctx, remoteID, id)
-	return raw, nil
-}
-
-// processOffer (responder) first acknowledges the previous exchange the offer names
-// (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 RemoteID, o *OfferMsg) ([]byte, error) {
-	if o.AckID != (ExchangeID{}) {
-		m.ackConverged(remoteID, o.AckID)
-	}
-
-	m.mu.Lock()
-	if ex := m.exchanges[remoteID]; ex != nil && ex.id == o.ExchangeID {
-		state, last := ex.state, ex.lastSent
-		m.mu.Unlock()
-		if state == stateReserved {
-			return nil, nil
-		}
-		return last, nil
-	}
-	// Reserve the slot so a concurrent duplicate offer bails.
-	m.exchanges[remoteID] = &exchangeCtl{id: o.ExchangeID, state: stateReserved, startedAt: time.Now()}
-	m.mu.Unlock()
-
-	answerBytes, psk, err := Respond(o.KEMOffer, m.binding(remoteID))
-	if err != nil {
-		return nil, err
-	}
-	raw, err := (&AnswerMsg{ExchangeID: o.ExchangeID, KEMAnswer: answerBytes}).Encode()
-	if err != nil {
-		return nil, err
-	}
-
-	m.mu.Lock()
-	ex := m.exchanges[remoteID]
-	if ex == nil || ex.id != o.ExchangeID {
-		m.mu.Unlock()
-		return nil, nil
-	}
-	ex.state = stateAwaitingAck
-	ex.lastSent = raw
-	ex.pendingPSK = psk
-	m.psks[remoteID] = psk
-	m.mu.Unlock()
-
-	// Commit optimistically so our data path can rekey to the new PSK.
-	if err := m.cbHandler.OnNewPSKReady(remoteID, psk); err != nil {
-		return nil, err
-	}
-	return raw, nil
-}
-
-// processAnswer (initiator) derives and commits the PSK and parks in
-// 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 RemoteID, a *AnswerMsg) error {
-	m.mu.Lock()
-	ex := m.exchanges[remoteID]
-	if ex == nil || ex.id != a.ExchangeID || ex.state != stateAwaitingAnswer {
-		m.mu.Unlock()
-		return nil
-	}
-	ex.state = stateAwaitingRekey
-	init := ex.initiator
-	ex.initiator = nil
-	m.mu.Unlock()
-
-	psk, err := init.Finish(a.KEMAnswer, m.binding(remoteID))
-	if err != nil {
-		return err
-	}
-
-	// The initiator has converged: the responder must have derived the key to answer.
-	m.mu.Lock()
-	m.established[remoteID] = true
-	m.failures[remoteID] = 0
-	m.psks[remoteID] = psk
-	m.mu.Unlock()
-
-	return m.cbHandler.OnNewPSKReady(remoteID, psk)
-}
-
-// 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 RemoteID, ackID ExchangeID) {
-	m.mu.Lock()
-	ex := m.exchanges[remoteID]
-	if ex == nil || ex.id != ackID || ex.state != stateAwaitingAck {
-		m.mu.Unlock()
-		return
-	}
-	delete(m.exchanges, remoteID)
-	m.established[remoteID] = true
-	m.failures[remoteID] = 0
-	_ = time.Since(ex.startedAt) // convergence latency (metrics hook, later step)
-	m.mu.Unlock()
-}
-
-// initiatorLoop enforces the offer->answer convergence deadline and retransmits the
-// initiator's outstanding data-path offer while awaiting the answer (a
-// signalling-bootstrapped offer is retransmitted by the host, so it is not resent
-// 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 RemoteID, id ExchangeID) {
-	defer m.wait.Done()
-	t := time.NewTicker(m.retryInterval)
-	defer t.Stop()
-
-	attempts := 0
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case <-t.C:
-			m.mu.Lock()
-			ex := m.exchanges[remoteID]
-			if ex == nil || ex.id != id {
-				m.mu.Unlock()
-				return
-			}
-
-			switch ex.state {
-			case stateAwaitingAnswer:
-				if attempts >= m.maxRetries {
-					delete(m.exchanges, remoteID)
-					fail := m.registerFailureLocked(remoteID)
-					m.mu.Unlock()
-					m.raiseFailure(remoteID, fail)
-					return
-				}
-				viaSignal := ex.viaSignal
-				msg := ex.lastSent
-				attempts++
-				m.mu.Unlock()
-				if !viaSignal {
-					if err := m.pushDataPath(remoteID, msg); err != nil {
-						m.logger.Warn("pqkem: offer retransmit failed", "peer", remoteID, "err", err)
-					}
-				}
-
-			default:
-				// Past awaiting the answer (converged) or superseded: the loop's job
-				// is done. The next rotation is driven externally by OnDataPathRekeyed,
-				// so there is no deadline while idle-waiting for it (that wait can be
-				// as long as the transport's natural rekey interval).
-				m.mu.Unlock()
-				return
-			}
-		}
-	}
-}
-
-// registerFailureLocked applies policy B and reports whether OnRekeyFailed is due:
-// 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 RemoteID) bool {
-	if !m.established[remoteID] {
-		return true
-	}
-	m.failures[remoteID]++
-	if m.failures[remoteID] >= m.maxRekeyFailures {
-		m.failures[remoteID] = 0
-		return true
-	}
-	return false
-}
-
-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
-	}
-	if err := m.cbHandler.OnRekeyFailed(remoteID); err != nil {
-		m.logger.Error("pqkem: OnRekeyFailed handler error", "peer", remoteID, "err", err)
-	}
-}

client/internal/pqkem/convergence_test.go

@@ -1,74 +0,0 @@
-package pqkem
-
-import (
-	"net/netip"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-)
-
-// dropTransport is a pqkem.Transport that silently discards everything.
-type dropTransport struct{}
-
-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 }
-
-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", wg, nil) // bbbb > aaaa -> initiator
-	d.Start(dropTransport{})
-	d.retryInterval = 5 * time.Millisecond
-	d.maxRetries = 3
-	defer d.Stop()
-
-	// Bootstrap offer is produced for signalling; no answer ever comes back -> the
-	// initial exchange fails fast.
-	offer, err := d.SignalOffer("aaaa")
-	require.NoError(t, err)
-	require.NotNil(t, offer)
-
-	require.Eventually(t, func() bool { return failedCount(wg) == 1 }, time.Second, 5*time.Millisecond)
-}
-
-func TestManager_RekeyToleratesKFailures(t *testing.T) {
-	dA, dB, _, wgB, lbB := pair(t)
-	defer dA.Stop()
-	defer dB.Stop()
-
-	// Tighten B's timings before any exchange loop spawns (the loop reads these
-	// fields, so writing them after a loop is running would race).
-	dB.retryInterval = 5 * time.Millisecond
-	dB.maxRetries = 2
-
-	// 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")
-	require.NotEqual(t, PSK{}, wgB.psk("aaaa"))
-
-	// Drop B's outbound so rekeys can no longer converge.
-	lbB.drop.Store(true)
-
-	// K-1 data-path rekeys must NOT raise OnRekeyFailed.
-	for i := 0; i < DefaultMaxRekeyFailures-1; i++ {
-		_, err := dB.startExchange("aaaa", false, ExchangeID{})
-		require.NoError(t, err)
-		time.Sleep(50 * time.Millisecond)
-	}
-	require.Equal(t, 0, failedCount(wgB), "no failure before K attempts")
-
-	// The K-th failure raises it once.
-	_, err := dB.startExchange("aaaa", false, ExchangeID{})
-	require.NoError(t, err)
-	require.Eventually(t, func() bool { return failedCount(wgB) == 1 }, time.Second, 5*time.Millisecond)
-}

client/internal/pqkem/env.go

@@ -1,59 +0,0 @@
-package pqkem
-
-import (
-	"log/slog"
-	"os"
-	"strconv"
-	"strings"
-
-	log "github.com/sirupsen/logrus"
-)
-
-// EnvEnabled is the environment variable that turns the ML-KEM post-quantum
-// exchange on for this client. Accepts on/off aliases plus anything
-// strconv.ParseBool understands (true/false/1/0).
-const EnvEnabled = "NB_ENABLE_PQ_MLKEM"
-
-// Enabled reports whether the ML-KEM PQ exchange is enabled via the environment.
-// An empty or unrecognized value is treated as disabled.
-func Enabled() bool {
-	raw := strings.ToLower(strings.TrimSpace(os.Getenv(EnvEnabled)))
-	switch raw {
-	case "":
-		return false
-	case "on":
-		return true
-	case "off":
-		return false
-	}
-	enabled, err := strconv.ParseBool(raw)
-	if err != nil {
-		log.Warnf("failed to parse %s value %q: %v", EnvEnabled, raw, err)
-		return false
-	}
-	return enabled
-}
-
-// EnvLogLevel overrides the ML-KEM manager's slog level (debug/info/warn/error).
-// Defaults to info.
-const EnvLogLevel = "NB_PQ_MLKEM_LOG_LEVEL"
-
-// NewLogger builds the slog logger for the ML-KEM manager: a text handler to stdout
-// at the level from EnvLogLevel. Mirrors the Rosenpass manager's logger setup so PQ
-// components log consistently.
-func NewLogger() *slog.Logger {
-	return slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: logLevel()}))
-}
-
-func logLevel() slog.Level {
-	switch strings.ToLower(strings.TrimSpace(os.Getenv(EnvLogLevel))) {
-	case "debug":
-		return slog.LevelDebug
-	case "warn":
-		return slog.LevelWarn
-	case "error":
-		return slog.LevelError
-	default:
-		return slog.LevelInfo
-	}
-}

client/internal/pqkem/kem.go

@@ -1,167 +0,0 @@
-// Package pqkem is a spike (NET-1406) for a post-quantum pre-shared-key exchange
-// that could replace Rosenpass. It performs an X25519MLKEM768 hybrid key
-// encapsulation and derives a 32-byte pre-shared key (PSK).
-//
-// The exchange is a single round trip designed to ride the (already
-// authenticated) Signal offer/answer channel:
-//
-//	initiator --Offer(1216B)--> responder
-//	initiator <--Answer(1120B)-- responder
-//
-// Both sides then hold the same PSK, which is bound to the two peers' identities
-// (their peer identity keys) so the derived key cannot be transplanted
-// to a different peer pair even if the transport authentication were bypassed.
-//
-// Combiner note: this follows draft-ietf-tls-ecdhe-mlkem for X25519MLKEM768 — on
-// the wire ML-KEM ‖ X25519 (the draft deliberately reversed the share order for
-// this group), and ML-KEM_ss ‖ X25519_ss fed into the KDF. The spike uses SHA-256
-// (also binding the transcript and peer identities); a production version should
-// use HKDF — see TODO below.
-package pqkem
-
-import (
-	"crypto/ecdh"
-	"crypto/mlkem"
-	"crypto/rand"
-	"crypto/sha256"
-	"fmt"
-)
-
-const (
-	// OfferSize is the initiator message: ML-KEM-768 encapsulation key ‖ X25519 public key
-	// (share order per draft-ietf-tls-ecdhe-mlkem for X25519MLKEM768).
-	OfferSize = mlkem.EncapsulationKeySize768 + 32 // 1216
-	// AnswerSize is the responder message: ML-KEM-768 ciphertext ‖ X25519 public key.
-	AnswerSize = mlkem.CiphertextSize768 + 32 // 1120
-
-	pskLabel = "netbird-pq-psk-v1"
-)
-
-// PSK is the 32-byte derived pre-shared key handed to the consumer to key its channel.
-type PSK [32]byte
-
-// Binding identifies the peer pair the PSK is derived for. Callers set both
-// peer identity keys; the order does not matter (it is canonicalised).
-type Binding struct {
-	LocalID  []byte
-	RemoteID []byte
-}
-
-// Initiator holds the ephemeral secrets between Offer and Finish.
-type Initiator struct {
-	x25519  *ecdh.PrivateKey
-	mlkemDK *mlkem.DecapsulationKey768
-	offer   []byte
-}
-
-// NewInitiator generates the ephemeral X25519 + ML-KEM-768 keypairs.
-func NewInitiator() (*Initiator, error) {
-	x, err := ecdh.X25519().GenerateKey(rand.Reader)
-	if err != nil {
-		return nil, fmt.Errorf("x25519 keygen: %w", err)
-	}
-	dk, err := mlkem.GenerateKey768()
-	if err != nil {
-		return nil, fmt.Errorf("ml-kem keygen: %w", err)
-	}
-
-	offer := make([]byte, 0, OfferSize)
-	offer = append(offer, dk.EncapsulationKey().Bytes()...)
-	offer = append(offer, x.PublicKey().Bytes()...)
-
-	return &Initiator{x25519: x, mlkemDK: dk, offer: offer}, nil
-}
-
-// Offer returns the initiator message to send over Signal.
-func (i *Initiator) Offer() []byte {
-	return i.offer
-}
-
-// Finish consumes the responder's answer and derives the PSK.
-func (i *Initiator) Finish(answer []byte, b Binding) (PSK, error) {
-	if len(answer) != AnswerSize {
-		return PSK{}, fmt.Errorf("answer: got %d bytes, want %d", len(answer), AnswerSize)
-	}
-	ct := answer[:mlkem.CiphertextSize768]
-	peerX := answer[mlkem.CiphertextSize768:]
-
-	ssMLKEM, err := i.mlkemDK.Decapsulate(ct)
-	if err != nil {
-		return PSK{}, fmt.Errorf("ml-kem decapsulate: %w", err)
-	}
-	pub, err := ecdh.X25519().NewPublicKey(peerX)
-	if err != nil {
-		return PSK{}, fmt.Errorf("parse peer x25519: %w", err)
-	}
-	ssX, err := i.x25519.ECDH(pub)
-	if err != nil {
-		return PSK{}, fmt.Errorf("x25519 ecdh: %w", err)
-	}
-
-	return derivePSK(ssMLKEM, ssX, i.offer, answer, b), nil
-}
-
-// Respond consumes an initiator offer, produces the answer, and derives the PSK.
-func Respond(offer []byte, b Binding) (answer []byte, psk PSK, err error) {
-	if len(offer) != OfferSize {
-		return nil, PSK{}, fmt.Errorf("offer: got %d bytes, want %d", len(offer), OfferSize)
-	}
-	peerEK := offer[:mlkem.EncapsulationKeySize768]
-	peerX := offer[mlkem.EncapsulationKeySize768:]
-
-	ek, err := mlkem.NewEncapsulationKey768(peerEK)
-	if err != nil {
-		return nil, PSK{}, fmt.Errorf("parse peer ml-kem key: %w", err)
-	}
-	ssMLKEM, ct := ek.Encapsulate()
-
-	x, err := ecdh.X25519().GenerateKey(rand.Reader)
-	if err != nil {
-		return nil, PSK{}, fmt.Errorf("x25519 keygen: %w", err)
-	}
-	pub, err := ecdh.X25519().NewPublicKey(peerX)
-	if err != nil {
-		return nil, PSK{}, fmt.Errorf("parse peer x25519: %w", err)
-	}
-	ssX, err := x.ECDH(pub)
-	if err != nil {
-		return nil, PSK{}, fmt.Errorf("x25519 ecdh: %w", err)
-	}
-
-	answer = make([]byte, 0, AnswerSize)
-	answer = append(answer, ct...)
-	answer = append(answer, x.PublicKey().Bytes()...)
-
-	// derivePSK uses the same argument order on both sides; the responder's local
-	// binding is the mirror of the initiator's, canonicalised inside derivePSK.
-	return answer, derivePSK(ssMLKEM, ssX, offer, answer, b), nil
-}
-
-// derivePSK combines the two shared secrets and binds the result to the full
-// transcript (offer ‖ answer) and the canonicalised peer identities.
-//
-// TODO(NET-1406): replace the SHA-256 concat with the RFC HKDF combiner
-// (crypto/hkdf, Go 1.24+) and proper labels before this leaves spike status.
-func derivePSK(ssMLKEM, ssX, offer, answer []byte, b Binding) PSK {
-	lo, hi := canonicalPair(b.LocalID, b.RemoteID)
-
-	h := sha256.New()
-	h.Write([]byte(pskLabel))
-	h.Write(ssMLKEM)
-	h.Write(ssX)
-	h.Write(offer)
-	h.Write(answer)
-	h.Write(lo)
-	h.Write(hi)
-
-	var psk PSK
-	copy(psk[:], h.Sum(nil))
-	return psk
-}
-
-func canonicalPair(a, b []byte) (lo, hi []byte) {
-	if string(a) <= string(b) {
-		return a, b
-	}
-	return b, a
-}

client/internal/pqkem/kem_test.go

@@ -1,89 +0,0 @@
-package pqkem
-
-import (
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-)
-
-var (
-	wgA = []byte("peer-A-wireguard-pubkey-32bytes!")
-	wgB = []byte("peer-B-wireguard-pubkey-32bytes!")
-)
-
-func TestExchange_DerivesMatchingPSK(t *testing.T) {
-	init, err := NewInitiator()
-	require.NoError(t, err)
-
-	require.Len(t, init.Offer(), OfferSize)
-
-	answer, pskB, err := Respond(init.Offer(), Binding{LocalID: wgB, RemoteID: wgA})
-	require.NoError(t, err)
-	require.Len(t, answer, AnswerSize)
-
-	pskA, err := init.Finish(answer, Binding{LocalID: wgA, RemoteID: wgB})
-	require.NoError(t, err)
-
-	require.Equal(t, pskB, pskA, "both sides must derive the same PSK")
-	require.NotEqual(t, PSK{}, pskA, "PSK must not be zero")
-}
-
-func TestExchange_PSKBoundToPeerIdentities(t *testing.T) {
-	init, err := NewInitiator()
-	require.NoError(t, err)
-
-	// responder computes with the honest pair...
-	_, pskHonest, err := Respond(init.Offer(), Binding{LocalID: wgB, RemoteID: wgA})
-	require.NoError(t, err)
-
-	// ...a second responder run with a different peer identity yields a different PSK,
-	// even though the KEM material would otherwise combine identically.
-	wgC := []byte("peer-C-wireguard-pubkey-32bytes!")
-	_, pskWrong, err := Respond(init.Offer(), Binding{LocalID: wgC, RemoteID: wgA})
-	require.NoError(t, err)
-
-	require.NotEqual(t, pskHonest, pskWrong, "PSK must be bound to the peer pair")
-}
-
-func TestExchange_RejectsMalformedMessages(t *testing.T) {
-	init, err := NewInitiator()
-	require.NoError(t, err)
-
-	_, _, err = Respond(init.Offer()[:10], Binding{})
-	require.Error(t, err)
-
-	_, err = init.Finish([]byte("too short"), Binding{})
-	require.Error(t, err)
-}
-
-// TestExchange_ReportSizesAndTiming is a spike measurement, not a pass/fail gate.
-// Run with: go test -run TestExchange_ReportSizesAndTiming -v ./client/internal/pqkem/
-func TestExchange_ReportSizesAndTiming(t *testing.T) {
-	const iters = 200
-
-	var tInit, tResp, tFinish time.Duration
-	for i := 0; i < iters; i++ {
-		s0 := time.Now()
-		init, err := NewInitiator()
-		require.NoError(t, err)
-		tInit += time.Since(s0)
-
-		s1 := time.Now()
-		answer, _, err := Respond(init.Offer(), Binding{LocalID: wgB, RemoteID: wgA})
-		require.NoError(t, err)
-		tResp += time.Since(s1)
-
-		s2 := time.Now()
-		_, err = init.Finish(answer, Binding{LocalID: wgA, RemoteID: wgB})
-		require.NoError(t, err)
-		tFinish += time.Since(s2)
-	}
-
-	t.Logf("wire sizes:   offer=%d B   answer=%d B   (Rosenpass static pubkey ~524160 B)", OfferSize, AnswerSize)
-	t.Logf("total on-wire per handshake: %d B  (~%.0fx smaller than RP static key)", OfferSize+AnswerSize, 524160.0/float64(OfferSize+AnswerSize))
-	t.Logf("avg NewInitiator (keygen): %s", tInit/iters)
-	t.Logf("avg Respond (encaps+dh):   %s", tResp/iters)
-	t.Logf("avg Finish (decaps+dh):    %s", tFinish/iters)
-	t.Logf("avg full handshake CPU:    %s", (tInit+tResp+tFinish)/iters)
-}

client/internal/pqkem/manager.go

@@ -1,372 +0,0 @@
-package pqkem
-
-import (
-	"context"
-	"crypto/rand"
-	"fmt"
-	"log/slog"
-	"net/netip"
-	"sync"
-	"time"
-)
-
-const (
-	// DefaultRetryInterval is how often the initiator retransmits its outstanding
-	// data-path offer while awaiting the answer.
-	DefaultRetryInterval = 2 * time.Second
-	// DefaultMaxRetries bounds how many ticks an exchange may run before it is
-	// declared failed. The convergence deadline is thus MaxRetries * RetryInterval.
-	DefaultMaxRetries = 10
-	// DefaultMaxRekeyFailures is how many consecutive rekey (non-initial) failures
-	// are tolerated before OnRekeyFailed. The initial exchange fails immediately.
-	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 Start, Close at Stop).
-type Transport interface {
-	// 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.
-type exchangeState uint8
-
-const (
-	stateReserved       exchangeState = iota // responder: deriving the answer
-	stateAwaitingAnswer                      // initiator: offer sent, awaiting the answer
-	stateAwaitingRekey                       // initiator: PSK derived+set, awaiting OnDataPathRekeyed to chain the next offer
-	stateAwaitingAck                         // responder: answer sent, awaiting the next offer that acks this exchange
-)
-
-// exchangeCtl holds all state for one in-flight exchange with a peer, under the
-// Manager's single lock. state drives every decision. lastSent is the current
-// data-path retransmit payload (the offer, for the initiator). initiator is the
-// ephemeral handle used at Finish; pendingPSK is the responder's derived key.
-// viaSignal records that the offer went to the host for the signalling channel, so
-// the loop does not retransmit it on the data path. Only the initiator runs a
-// retransmit loop, so only it sets cancel.
-type exchangeCtl struct {
-	id         ExchangeID
-	state      exchangeState
-	startedAt  time.Time
-	cancel     context.CancelFunc
-	lastSent   []byte
-	initiator  *Initiator
-	pendingPSK PSK
-	viaSignal  bool
-}
-
-// Manager is the stateful orchestrator — the analogue of go-rosenpass's Server. It
-// 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   LocalID
-	cbHandler CallbackHandler
-	logger    *slog.Logger
-
-	retryInterval    time.Duration
-	maxRetries       int
-	maxRekeyFailures int
-
-	rootCtx    context.Context
-	rootCancel context.CancelFunc
-
-	mu          sync.Mutex
-	transport   Transport
-	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
-	psks        map[RemoteID]PSK            // latest derived PSK per peer (pulled at WG peer-config time)
-	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(). Install the data-path transport with Start.
-func NewManager(localID LocalID, h CallbackHandler, logger *slog.Logger) *Manager {
-	if logger == nil {
-		logger = slog.Default()
-	}
-	ctx, cancel := context.WithCancel(context.Background())
-	return &Manager{
-		localID:          localID,
-		cbHandler:        h,
-		logger:           logger,
-		retryInterval:    DefaultRetryInterval,
-		maxRetries:       DefaultMaxRetries,
-		maxRekeyFailures: DefaultMaxRekeyFailures,
-		rootCtx:          ctx,
-		rootCancel:       cancel,
-		exchanges:        make(map[RemoteID]*exchangeCtl),
-		established:      make(map[RemoteID]bool),
-		failures:         make(map[RemoteID]int),
-		psks:             make(map[RemoteID]PSK),
-		peerAddrs:        make(map[RemoteID]netip.AddrPort),
-		peersByAddr:      make(map[netip.AddrPort]RemoteID),
-	}
-}
-
-// 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()
-	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.
-func (m *Manager) IsInitiator(remoteID RemoteID) bool {
-	return string(m.localID) > string(remoteID)
-}
-
-// PSK returns the latest PSK derived for the peer, for the host to program at WG
-// peer-config time (the pull path). ok is false until an exchange has derived one.
-func (m *Manager) PSK(remoteID RemoteID) (PSK, bool) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-	psk, ok := m.psks[remoteID]
-	return psk, ok
-}
-
-// 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 RemoteID, endpoint netip.AddrPort) {
-	if !endpoint.IsValid() {
-		return
-	}
-	m.mu.Lock()
-	if old, ok := m.peerAddrs[remoteID]; ok {
-		delete(m.peersByAddr, old)
-	}
-	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 RemoteID) {
-	m.mu.Lock()
-	if ex, ok := m.exchanges[remoteID]; ok {
-		if ex.cancel != nil {
-			ex.cancel()
-		}
-		delete(m.exchanges, remoteID)
-	}
-	delete(m.established, remoteID)
-	delete(m.failures, remoteID)
-	delete(m.psks, remoteID)
-	if ep, ok := m.peerAddrs[remoteID]; ok {
-		delete(m.peersByAddr, ep)
-		delete(m.peerAddrs, remoteID)
-	}
-	m.mu.Unlock()
-}
-
-// 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[RemoteID]*exchangeCtl)
-	m.psks = make(map[RemoteID]PSK)
-	m.mu.Unlock()
-	if t != nil {
-		if err := t.Close(); err != nil {
-			m.logger.Warn("pqkem: closing data-path transport", "err", err)
-		}
-	}
-}
-
-// ---- Signalling channel (host-driven; rides the host's negotiation) ----
-
-// SignalOffer returns the KEM offer for the host to embed in its outgoing offer to
-// 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 RemoteID) ([]byte, error) {
-	if !m.IsInitiator(remoteID) {
-		return nil, nil
-	}
-	m.mu.Lock()
-	if ex := m.exchanges[remoteID]; ex != nil && ex.viaSignal && ex.state == stateAwaitingAnswer {
-		last := ex.lastSent
-		m.mu.Unlock()
-		return last, nil
-	}
-	m.mu.Unlock()
-	// bootstrap offer acknowledges nothing (zero AckID).
-	return m.startExchange(remoteID, true, ExchangeID{})
-}
-
-// 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 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)
-	}
-	if typ != MsgOffer {
-		return nil, fmt.Errorf("expected offer from %s, got type %d", remoteID, typ)
-	}
-	return m.processOffer(remoteID, msg.(*OfferMsg))
-}
-
-// 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 RemoteID, answer []byte) error {
-	typ, msg, err := Decode(answer)
-	if err != nil {
-		return fmt.Errorf("decode signal answer from %s: %w", remoteID, err)
-	}
-	if typ != MsgAnswer {
-		return fmt.Errorf("expected answer from %s, got type %d", remoteID, typ)
-	}
-	return m.processAnswer(remoteID, msg.(*AnswerMsg))
-}
-
-// ---- 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, ok := m.peersByAddr[src]
-	m.mu.Unlock()
-	if !ok {
-		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 RemoteID, raw []byte) error {
-	typ, msg, err := Decode(raw)
-	if err != nil {
-		return fmt.Errorf("decode data-path msg from %s: %w", remoteID, err)
-	}
-	switch typ {
-	case MsgOffer:
-		answer, err := m.processOffer(remoteID, msg.(*OfferMsg))
-		if err != nil {
-			return err
-		}
-		if answer == nil {
-			return nil
-		}
-		return m.pushDataPath(remoteID, answer)
-	case MsgAnswer:
-		return m.processAnswer(remoteID, msg.(*AnswerMsg))
-	default:
-		return fmt.Errorf("unhandled data-path message type %d from %s", typ, remoteID)
-	}
-}
-
-// OnDataPathRekeyed notifies that the peer's data path is up and freshly keyed with
-// 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 RemoteID) {
-	m.mu.Lock()
-	ex := m.exchanges[remoteID]
-	chain := ex != nil && ex.state == stateAwaitingRekey
-	var ackID ExchangeID
-	if chain {
-		ackID = ex.id
-	}
-	m.mu.Unlock()
-
-	if !chain {
-		return
-	}
-	offer, err := m.startExchange(remoteID, false, ackID)
-	if err != nil {
-		m.logger.Error("pqkem: chain offer failed to start", "peer", remoteID, "err", err)
-		return
-	}
-	if err := m.pushDataPath(remoteID, offer); err != nil {
-		m.logger.Warn("pqkem: send chain offer failed", "peer", remoteID, "err", err)
-	}
-}
-
-// 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 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 RemoteID, msg []byte) error {
-	m.mu.Lock()
-	ep, ok := m.peerAddrs[remoteID]
-	t := m.transport
-	m.mu.Unlock()
-	if !ok {
-		return fmt.Errorf("no data-path endpoint for peer %s", remoteID)
-	}
-	if t == nil {
-		return fmt.Errorf("no data-path transport")
-	}
-	return t.Send(ep, msg)
-}
-
-func (m *Manager) binding(remoteID RemoteID) Binding {
-	return Binding{LocalID: []byte(m.localID), RemoteID: []byte(remoteID)}
-}
-
-func newExchangeID() (ExchangeID, error) {
-	var id ExchangeID
-	if _, err := rand.Read(id[:]); err != nil {
-		return ExchangeID{}, fmt.Errorf("generate exchange id: %w", err)
-	}
-	return id, nil
-}

client/internal/pqkem/manager_test.go

@@ -1,169 +0,0 @@
-package pqkem
-
-import (
-	"fmt"
-	"net/netip"
-	"sync"
-	"sync/atomic"
-	"testing"
-
-	"github.com/stretchr/testify/require"
-)
-
-// 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 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[RemoteID]PSK
-	failed []RemoteID
-}
-
-func newFakeWG() *fakeWG { return &fakeWG{psks: map[RemoteID]PSK{}} }
-
-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 RemoteID) error {
-	f.mu.Lock()
-	defer f.mu.Unlock()
-	f.failed = append(f.failed, remoteID)
-	return nil
-}
-
-func (f *fakeWG) psk(peer RemoteID) PSK {
-	f.mu.Lock()
-	defer f.mu.Unlock()
-	return f.psks[peer]
-}
-
-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()
-	sw := newSwitch()
-	wgA = newFakeWG()
-	wgB = newFakeWG()
-	dA = NewManager("aaaa", wgA, nil)
-	dB = NewManager("bbbb", wgB, nil)
-	dA.Start(&loopback{ep: epA, sw: sw})
-	lbB = &loopback{ep: epB, sw: sw}
-	dB.Start(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).
-func bootstrap(t *testing.T, dA, dB *Manager) {
-	t.Helper()
-	offer, err := dB.SignalOffer("aaaa")
-	require.NoError(t, err)
-	require.NotNil(t, offer)
-	answer, err := dA.SignalOnOffer("bbbb", offer)
-	require.NoError(t, err)
-	require.NotNil(t, answer)
-	require.NoError(t, dB.SignalOnAnswer("aaaa", answer))
-}
-
-func TestManager_BootstrapDerivesSamePSK(t *testing.T) {
-	dA, dB, wgA, wgB, _ := pair(t)
-	defer dA.Stop()
-	defer dB.Stop()
-
-	bootstrap(t, dA, dB)
-
-	pskA := wgA.psk("bbbb")
-	pskB := wgB.psk("aaaa")
-	require.NotEqual(t, PSK{}, pskA)
-	require.Equal(t, pskB, pskA, "both sides derive the same PSK from the bootstrap exchange")
-}
-
-func TestManager_ChainRotatesAndAcks(t *testing.T) {
-	dA, dB, wgA, wgB, _ := pair(t)
-	defer dA.Stop()
-	defer dB.Stop()
-
-	bootstrap(t, dA, dB)
-	psk1 := wgB.psk("aaaa")
-
-	// 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")
-
-	psk2A := wgA.psk("bbbb")
-	psk2B := wgB.psk("aaaa")
-	require.Equal(t, psk2B, psk2A, "both sides converge on the rotated PSK")
-	require.NotEqual(t, psk1, psk2B, "the chain rotated to a new PSK")
-}
-
-func TestManager_NonInitiatorReturnsNoOffer(t *testing.T) {
-	dA := NewManager("aaaa", newFakeWG(), nil)
-	defer dA.Stop()
-
-	offer, err := dA.SignalOffer("bbbb") // not the initiator vs "bbbb"
-	require.NoError(t, err)
-	require.Nil(t, offer)
-}
-
-func TestManager_StopIsIdempotent(t *testing.T) {
-	dA := NewManager("aaaa", newFakeWG(), nil)
-	dA.Start(&loopback{ep: epA, sw: newSwitch()})
-	dA.Stop()
-	dA.Stop() // must not panic or hang
-}

client/internal/pqkem/message.go

@@ -1,121 +0,0 @@
-package pqkem
-
-import (
-	"crypto/mlkem"
-	"fmt"
-)
-
-// Wire framing for the PQ-KEM exchange. Messages are self-contained, versioned,
-// transport-agnostic byte blobs: the same bytes ride the signalling channel
-// (initial bootstrap) or a data-tunnel packet (rekey). The library only ever sees
-// opaque []byte at the transport seam.
-//
-// Layout (all messages): [type:1][version:1][exchangeID:16][payload...]
-//
-// There is no confirm message: an exchange is acknowledged by the NEXT offer, which
-// carries the acked exchange's id (see OfferMsg.AckID) and — riding the data path
-// under the freshly adopted key — proves that key works.
-
-const (
-	// ProtocolVersion is bumped on any wire-incompatible change; a peer rejects
-	// messages it does not understand rather than misparsing them.
-	ProtocolVersion uint8 = 1
-
-	// ExchangeIDSize identifies one exchange so answers/acks correlate and stale
-	// messages are dropped.
-	ExchangeIDSize = 16
-
-	headerSize = 1 + 1 + ExchangeIDSize
-)
-
-// MsgType tags the two message kinds of the exchange.
-type MsgType uint8
-
-const (
-	MsgOffer MsgType = iota + 1
-	MsgAnswer
-)
-
-// ExchangeID is the per-exchange correlator. The zero value means "none" (an offer
-// that acknowledges nothing, i.e. the first exchange of a connection).
-type ExchangeID [ExchangeIDSize]byte
-
-// OfferMsg carries the initiator's public material (X25519 pub ‖ ML-KEM encap key)
-// and AckID, the id of the previous exchange this offer acknowledges (zero if none).
-type OfferMsg struct {
-	ExchangeID ExchangeID
-	AckID      ExchangeID
-	// KEMOffer is the raw Initiator.Offer() blob (OfferSize bytes).
-	KEMOffer []byte
-}
-
-// AnswerMsg carries the responder's reply (ML-KEM ciphertext ‖ X25519 pub) for the
-// round identified by ExchangeID.
-type AnswerMsg struct {
-	ExchangeID ExchangeID
-	// KEMAnswer is the raw Respond() answer blob (AnswerSize bytes).
-	KEMAnswer []byte
-}
-
-// Encode serialises the offer with its framed header (payload = AckID ‖ KEMOffer).
-func (m *OfferMsg) Encode() ([]byte, error) {
-	if len(m.KEMOffer) != OfferSize {
-		return nil, fmt.Errorf("offer payload: got %d, want %d", len(m.KEMOffer), OfferSize)
-	}
-	payload := make([]byte, 0, ExchangeIDSize+OfferSize)
-	payload = append(payload, m.AckID[:]...)
-	payload = append(payload, m.KEMOffer...)
-	return frame(MsgOffer, m.ExchangeID, payload), nil
-}
-
-// Encode serialises the answer with its framed header.
-func (m *AnswerMsg) Encode() ([]byte, error) {
-	if len(m.KEMAnswer) != AnswerSize {
-		return nil, fmt.Errorf("answer payload: got %d, want %d", len(m.KEMAnswer), AnswerSize)
-	}
-	return frame(MsgAnswer, m.ExchangeID, m.KEMAnswer), nil
-}
-
-// Decode parses a framed message into one of *OfferMsg / *AnswerMsg.
-func Decode(buf []byte) (MsgType, any, error) {
-	if len(buf) < headerSize {
-		return 0, nil, fmt.Errorf("message too short: %d bytes", len(buf))
-	}
-	typ := MsgType(buf[0])
-	if ver := buf[1]; ver != ProtocolVersion {
-		return typ, nil, fmt.Errorf("unsupported protocol version %d (want %d)", ver, ProtocolVersion)
-	}
-
-	var id ExchangeID
-	copy(id[:], buf[2:headerSize])
-	payload := buf[headerSize:]
-
-	switch typ {
-	case MsgOffer:
-		if len(payload) != ExchangeIDSize+OfferSize {
-			return typ, nil, fmt.Errorf("offer payload: got %d, want %d", len(payload), ExchangeIDSize+OfferSize)
-		}
-		var ack ExchangeID
-		copy(ack[:], payload[:ExchangeIDSize])
-		return typ, &OfferMsg{ExchangeID: id, AckID: ack, KEMOffer: payload[ExchangeIDSize:]}, nil
-	case MsgAnswer:
-		if len(payload) != AnswerSize {
-			return typ, nil, fmt.Errorf("answer payload: got %d, want %d", len(payload), AnswerSize)
-		}
-		return typ, &AnswerMsg{ExchangeID: id, KEMAnswer: payload}, nil
-	default:
-		return typ, nil, fmt.Errorf("unknown message type %d", typ)
-	}
-}
-
-func frame(typ MsgType, id ExchangeID, payload []byte) []byte {
-	buf := make([]byte, headerSize+len(payload))
-	buf[0] = byte(typ)
-	buf[1] = ProtocolVersion
-	copy(buf[2:], id[:])
-	copy(buf[headerSize:], payload)
-	return buf
-}
-
-// compile-time assurance the KEM blob sizes referenced here stay in sync with kem.go.
-var _ = [1]struct{}{}[OfferSize-(32+mlkem.EncapsulationKeySize768)]

client/internal/pqkem/message_test.go

@@ -1,57 +0,0 @@
-package pqkem
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/require"
-)
-
-func TestMessageRoundTrip(t *testing.T) {
-	init, err := NewInitiator()
-	require.NoError(t, err)
-	answer, _, err := Respond(init.Offer(), Binding{LocalID: wgB, RemoteID: wgA})
-	require.NoError(t, err)
-
-	id := ExchangeID{1, 2, 3, 4}
-	ack := ExchangeID{9, 9, 9}
-
-	offBytes, err := (&OfferMsg{ExchangeID: id, AckID: ack, KEMOffer: init.Offer()}).Encode()
-	require.NoError(t, err)
-	typ, decoded, err := Decode(offBytes)
-	require.NoError(t, err)
-	require.Equal(t, MsgOffer, typ)
-	require.Equal(t, id, decoded.(*OfferMsg).ExchangeID)
-	require.Equal(t, ack, decoded.(*OfferMsg).AckID)
-	require.Equal(t, init.Offer(), decoded.(*OfferMsg).KEMOffer)
-
-	ansBytes, err := (&AnswerMsg{ExchangeID: id, KEMAnswer: answer}).Encode()
-	require.NoError(t, err)
-	typ, decoded, err = Decode(ansBytes)
-	require.NoError(t, err)
-	require.Equal(t, MsgAnswer, typ)
-	require.Equal(t, answer, decoded.(*AnswerMsg).KEMAnswer)
-}
-
-func TestDecodeRejects(t *testing.T) {
-	// too short
-	_, _, err := Decode([]byte{1, 1})
-	require.Error(t, err)
-
-	// wrong version
-	bad := make([]byte, headerSize+ExchangeIDSize+OfferSize)
-	bad[0] = byte(MsgOffer)
-	bad[1] = ProtocolVersion + 1
-	_, _, err = Decode(bad)
-	require.Error(t, err)
-
-	// unknown type
-	bad2 := make([]byte, headerSize)
-	bad2[0] = 99
-	bad2[1] = ProtocolVersion
-	_, _, err = Decode(bad2)
-	require.Error(t, err)
-
-	// offer with wrong payload size
-	_, err = (&OfferMsg{KEMOffer: []byte{1, 2, 3}}).Encode()
-	require.Error(t, err)
-}

client/internal/pqkem_adapter.go

@@ -1,102 +0,0 @@
-package internal
-
-import (
-	"net/netip"
-
-	log "github.com/sirupsen/logrus"
-	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-
-	"github.com/netbirdio/netbird/client/internal/pqkem"
-)
-
-// pqPresharedKeySetter is the subset of the WireGuard interface the ML-KEM callback
-// needs: programming a peer's preshared key. *iface.WGIface satisfies it.
-type pqPresharedKeySetter interface {
-	SetPresharedKey(peerKey string, psk wgtypes.Key, updateOnly bool) error
-}
-
-// pqCallbackHandler programs the derived PQ PSK onto the WireGuard peer. It is the
-// engine-side implementation of pqkem.CallbackHandler.
-type pqCallbackHandler struct {
-	wg pqPresharedKeySetter
-}
-
-// 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 pqkem.RemoteID, psk pqkem.PSK) error {
-	// updateOnly: applies to an already-configured peer (rotation). At bootstrap the
-	// peer is not configured yet, so this is a no-op there and the PSK is instead
-	// pulled at peer-config time (pqHandshaker.PSK / conn.presharedKey).
-	log.Debugf("pqkem: programming PSK for peer %s", remoteID)
-	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 pqkem.RemoteID) error {
-	log.Warnf("pqkem: post-quantum rekey failed for peer %s", remoteID)
-	return nil
-}
-
-// pqHandshaker adapts the pqkem manager to peer.PQHandshaker (string peer keys),
-// wiring the host's signalling offers/answers to the KEM exchange.
-type pqHandshaker struct {
-	mgr *pqkem.Manager
-}
-
-func (p pqHandshaker) OfferPayload(remoteKey string) ([]byte, int) {
-	payload, err := p.mgr.SignalOffer(pqkem.RemoteID(remoteKey))
-	if err != nil {
-		log.Warnf("pqkem: build offer for %s: %v", remoteKey, err)
-	}
-	return payload, p.mgr.LocalPort()
-}
-
-func (p pqHandshaker) AnswerPayload(remoteKey string, recvOffer []byte) ([]byte, int) {
-	if len(recvOffer) == 0 {
-		return nil, p.mgr.LocalPort()
-	}
-	payload, err := p.mgr.SignalOnOffer(pqkem.RemoteID(remoteKey), recvOffer)
-	if err != nil {
-		log.Warnf("pqkem: build answer for %s: %v", remoteKey, err)
-	}
-	return payload, p.mgr.LocalPort()
-}
-
-func (p pqHandshaker) OnAnswer(remoteKey string, recvAnswer []byte) {
-	if len(recvAnswer) == 0 {
-		return
-	}
-	if err := p.mgr.SignalOnAnswer(pqkem.RemoteID(remoteKey), recvAnswer); err != nil {
-		log.Warnf("pqkem: process answer from %s: %v", remoteKey, err)
-	}
-}
-
-// PSK exposes the peer's derived PSK for the conn to program at WG peer-config time.
-func (p pqHandshaker) PSK(remoteKey string) (wgtypes.Key, bool) {
-	psk, ok := p.mgr.PSK(pqkem.RemoteID(remoteKey))
-	if !ok {
-		return wgtypes.Key{}, false
-	}
-	return wgtypes.Key(psk), true
-}
-
-// SetRemoteAddr 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) SetRemoteAddr(remoteKey string, addr netip.AddrPort) {
-	if !addr.IsValid() || addr.Port() == 0 {
-		return
-	}
-	p.mgr.AddPeer(pqkem.RemoteID(remoteKey), addr)
-}
-
-// OnDataPathRekeyed clocks the next chained PSK rotation on a fresh WG handshake.
-func (p pqHandshaker) OnDataPathRekeyed(remoteKey string) {
-	p.mgr.OnDataPathRekeyed(pqkem.RemoteID(remoteKey))
-}
-
-// OnDataPathDown signals the peer's tunnel went down.
-func (p pqHandshaker) OnDataPathDown(remoteKey string) {
-	p.mgr.OnDataPathDown(pqkem.RemoteID(remoteKey))
-}

client/internal/pqkem_transport.go

@@ -1,72 +0,0 @@
-package internal
-
-import (
-	"fmt"
-	"net"
-	"net/netip"
-
-	log "github.com/sirupsen/logrus"
-)
-
-// 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 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
-}
-
-// 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. 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")
-	}
-	ip := net.IP(overlayIP.AsSlice())
-	conn, err := net.ListenUDP("udp4", &net.UDPAddr{IP: ip, Port: DefaultPort})
-	if err != nil {
-		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}, nil
-}
-
-// Send implements pqkem.Transport.
-func (t *pqTransport) Send(endpoint netip.AddrPort, msg []byte) error {
-	_, err := t.conn.WriteToUDPAddrPort(msg, endpoint)
-	return err
-}
-
-// 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)
-		}
-	}()
-}
-
-// Close implements pqkem.Transport.
-func (t *pqTransport) Close() error { return t.conn.Close() }

client/internal/profilemanager/config.go

@@ -58,10 +58,6 @@ var DefaultInterfaceBlacklist = []string{
 	"Tailscale", "tailscale", "docker", "veth", "br-", "lo",
 }
 
-// loadMDMPolicy is the package-level indirection used by apply() to read the
-// active MDM policy. Tests override this to inject a fake policy.
-var loadMDMPolicy = mdm.LoadPolicy
-
 // ConfigInput carries configuration changes to the client
 type ConfigInput struct {
 	ManagementURL                 string
@@ -186,14 +182,27 @@ type Config struct {
 
 	MTU uint16
 
-	// policy is the MDM policy that produced the currently-set values for
-	// any MDM-enforced fields. Set by applyMDMPolicy at the tail of apply()
-	// and reset on every apply() invocation. Never persisted to disk.
-	// Callers query enforcement state via Policy() and the mdm.Policy API
-	// (HasKey, ManagedKeys, IsEmpty).
+	// policy is the MDM policy that produced the currently-set values
+	// for any MDM-enforced fields. Set by ApplyMDMPolicy on every
+	// invocation. Never persisted to disk. Callers query enforcement
+	// state via Policy() and the mdm.Policy API (HasKey, ManagedKeys,
+	// IsEmpty).
 	policy *mdm.Policy `json:"-"`
 }
 
+// ApplyMDMPolicy overlays the supplied MDM Policy on top of the
+// currently resolved Config values. Idempotent — pass an empty Policy
+// to clear any prior overlay. The lifecycle owner (Server.getConfig
+// on desktop, the Client.Run path on mobile) calls this with
+// loader.Load() once the per-process Loader is known; the Config
+// itself holds no reference to the Loader.
+func (config *Config) ApplyMDMPolicy(policy *mdm.Policy) {
+	if config == nil {
+		return
+	}
+	config.applyMDMPolicy(policy)
+}
+
 // Policy returns the MDM policy applied to this Config. Returns a non-nil
 // empty Policy when MDM enforcement is inactive; callers can always invoke
 // HasKey / ManagedKeys / IsEmpty without a nil check.
@@ -650,9 +659,11 @@ func (config *Config) apply(input ConfigInput) (updated bool, err error) {
 		updated = true
 	}
 
-	// MDM is the last override layer: any key present in the policy
-	// supersedes defaults, on-disk config, env vars and CLI input.
-	config.applyMDMPolicy(loadMDMPolicy())
+	// Initialise the MDM overlay to "no enforcement" so Config.Policy()
+	// never returns a stale or nil policy on a freshly applied Config.
+	// Lifecycle owners that want to enforce a real MDM policy invoke
+	// Config.ApplyMDMPolicy(loader.Load()) after this returns.
+	config.applyMDMPolicy(mdm.NewPolicy(nil))
 
 	return updated, nil
 }

client/internal/profilemanager/config_mdm_test.go

@@ -10,24 +10,58 @@ import (
 	"github.com/netbirdio/netbird/client/mdm"
 )
 
-// withMDMPolicy temporarily overrides the package-level loadMDMPolicy hook so
-// apply() observes the supplied Policy. The original loader is restored at
-// test cleanup.
-func withMDMPolicy(t *testing.T, policy *mdm.Policy) {
+// fakeFetcher implements mdm.PolicyFetcher returning a pre-set policy
+// map. Test helper used to construct a Loader without touching the OS
+// or any package-level state.
+type fakeFetcher struct{ values map[string]any }
+
+func (f *fakeFetcher) Fetch() map[string]any { return f.values }
+
+// loaderFor builds an mdm.Loader whose loadPlatform returns the
+// supplied Policy's underlying values.
+func loaderFor(policy *mdm.Policy) *mdm.Loader {
+	if policy == nil || policy.IsEmpty() {
+		return mdm.NewLoader(&fakeFetcher{values: nil})
+	}
+	values := make(map[string]any)
+	for _, k := range policy.ManagedKeys() {
+		if v, ok := policy.GetString(k); ok {
+			values[k] = v
+			continue
+		}
+		if v, ok := policy.GetBool(k); ok {
+			values[k] = v
+			continue
+		}
+		if v, ok := policy.GetInt(k); ok {
+			values[k] = v
+			continue
+		}
+		if v, ok := policy.GetStringSlice(k); ok {
+			values[k] = v
+		}
+	}
+	return mdm.NewLoader(&fakeFetcher{values: values})
+}
+
+// configWithMDM is the test convenience that builds a Config via
+// UpdateOrCreateConfig and overlays the supplied MDM policy on top —
+// mirrors the production pattern (Server.getConfig / Client.applyMDMOverlay)
+// where the Loader lives outside Config and the apply step is driven
+// by the lifecycle owner.
+func configWithMDM(t *testing.T, input ConfigInput, policy *mdm.Policy) *Config {
 	t.Helper()
-	prev := loadMDMPolicy
-	loadMDMPolicy = func() *mdm.Policy { return policy }
-	t.Cleanup(func() { loadMDMPolicy = prev })
+	cfg, err := UpdateOrCreateConfig(input)
+	require.NoError(t, err)
+	require.NotNil(t, cfg)
+	cfg.ApplyMDMPolicy(loaderFor(policy).Load())
+	return cfg
 }
 
 func TestApply_MDMEmpty_NoEnforcement(t *testing.T) {
-	withMDMPolicy(t, mdm.NewPolicy(nil))
-
-	cfg, err := UpdateOrCreateConfig(ConfigInput{
+	cfg := configWithMDM(t, ConfigInput{
 		ConfigPath: filepath.Join(t.TempDir(), "config.json"),
-	})
-	require.NoError(t, err)
-	require.NotNil(t, cfg)
+	}, mdm.NewPolicy(nil))
 
 	assert.True(t, cfg.Policy().IsEmpty(), "no MDM source ⇒ empty Policy")
 	assert.False(t, cfg.Policy().HasKey(mdm.KeyManagementURL))
@@ -39,18 +73,15 @@ func TestApply_MDMEmpty_NoEnforcement(t *testing.T) {
 
 func TestApply_MDMOnly_OverridesDefaults(t *testing.T) {
 	const mdmURL = "https://corp.mdm.example.com:443"
-	withMDMPolicy(t, mdm.NewPolicy(map[string]any{
+
+	cfg := configWithMDM(t, ConfigInput{
+		ConfigPath: filepath.Join(t.TempDir(), "config.json"),
+	}, mdm.NewPolicy(map[string]any{
 		mdm.KeyManagementURL:       mdmURL,
 		mdm.KeyDisableClientRoutes: true,
 		mdm.KeyBlockInbound:        true,
 	}))
 
-	cfg, err := UpdateOrCreateConfig(ConfigInput{
-		ConfigPath: filepath.Join(t.TempDir(), "config.json"),
-	})
-	require.NoError(t, err)
-	require.NotNil(t, cfg)
-
 	assert.Equal(t, mdmURL, cfg.ManagementURL.String())
 	assert.True(t, cfg.DisableClientRoutes)
 	assert.True(t, cfg.BlockInbound)
@@ -65,16 +96,12 @@ func TestApply_MDMBeatsCLIInput(t *testing.T) {
 	const mdmURL = "https://mdm.example.com:443"
 	const cliURL = "https://cli.example.com:443"
 
-	withMDMPolicy(t, mdm.NewPolicy(map[string]any{
-		mdm.KeyManagementURL: mdmURL,
-	}))
-
-	cfg, err := UpdateOrCreateConfig(ConfigInput{
+	cfg := configWithMDM(t, ConfigInput{
 		ConfigPath:    filepath.Join(t.TempDir(), "config.json"),
 		ManagementURL: cliURL,
-	})
-	require.NoError(t, err)
-	require.NotNil(t, cfg)
+	}, mdm.NewPolicy(map[string]any{
+		mdm.KeyManagementURL: mdmURL,
+	}))
 
 	// MDM wins over CLI-supplied management URL.
 	assert.Equal(t, mdmURL, cfg.ManagementURL.String())
@@ -82,16 +109,12 @@ func TestApply_MDMBeatsCLIInput(t *testing.T) {
 }
 
 func TestApply_MDMInvalidURL_KeepsPreviousValue(t *testing.T) {
-	withMDMPolicy(t, mdm.NewPolicy(map[string]any{
+	cfg := configWithMDM(t, ConfigInput{
+		ConfigPath: filepath.Join(t.TempDir(), "config.json"),
+	}, mdm.NewPolicy(map[string]any{
 		mdm.KeyManagementURL: "not-a-url",
 	}))
 
-	cfg, err := UpdateOrCreateConfig(ConfigInput{
-		ConfigPath: filepath.Join(t.TempDir(), "config.json"),
-	})
-	require.NoError(t, err)
-	require.NotNil(t, cfg)
-
 	// Invalid MDM URL is logged and skipped: default URL stays in place
 	// to keep the client functional.
 	assert.Equal(t, DefaultManagementURL, cfg.ManagementURL.String())
@@ -106,24 +129,20 @@ func TestApply_MDMBoolKeysOverrideOnDiskValue(t *testing.T) {
 	tmp := filepath.Join(t.TempDir(), "config.json")
 
 	// Seed without MDM.
-	withMDMPolicy(t, mdm.NewPolicy(nil))
-	_, err := UpdateOrCreateConfig(ConfigInput{
+	configWithMDM(t, ConfigInput{
 		ConfigPath:          tmp,
 		DisableClientRoutes: boolPtr(false),
 		RosenpassEnabled:    boolPtr(false),
-	})
-	require.NoError(t, err)
+	}, mdm.NewPolicy(nil))
 
 	// Now enable MDM enforcement for these keys.
-	withMDMPolicy(t, mdm.NewPolicy(map[string]any{
+	cfg := configWithMDM(t, ConfigInput{
+		ConfigPath: tmp,
+	}, mdm.NewPolicy(map[string]any{
 		mdm.KeyDisableClientRoutes: true,
 		mdm.KeyRosenpassEnabled:    true,
 	}))
 
-	cfg, err := UpdateOrCreateConfig(ConfigInput{ConfigPath: tmp})
-	require.NoError(t, err)
-	require.NotNil(t, cfg)
-
 	assert.True(t, cfg.DisableClientRoutes, "MDM override should flip on-disk false to true")
 	assert.True(t, cfg.RosenpassEnabled)
 	assert.True(t, cfg.Policy().HasKey(mdm.KeyDisableClientRoutes))
@@ -164,16 +183,12 @@ func TestApply_MDMLazyConnection(t *testing.T) {
 func TestApply_MDMPreSharedKeyRedactionSentinelRejected(t *testing.T) {
 	const maskSentinel = "**********"
 
-	withMDMPolicy(t, mdm.NewPolicy(map[string]any{
+	cfg := configWithMDM(t, ConfigInput{
+		ConfigPath: filepath.Join(t.TempDir(), "config.json"),
+	}, mdm.NewPolicy(map[string]any{
 		mdm.KeyPreSharedKey: maskSentinel,
 	}))
 
-	cfg, err := UpdateOrCreateConfig(ConfigInput{
-		ConfigPath: filepath.Join(t.TempDir(), "config.json"),
-	})
-	require.NoError(t, err)
-	require.NotNil(t, cfg)
-
 	// Mask sentinel must not be persisted as the actual PSK.
 	assert.NotEqual(t, maskSentinel, cfg.PreSharedKey)
 	// Key still marked managed so user writes are still rejected.

client/internal/routemanager/dnsinterceptor/handler.go

@@ -95,7 +95,7 @@ func (d *DnsInterceptor) RemoveRoute() error {
 
 			// AllowedIPs should use real IPs
 			if d.currentPeerKey != "" {
-				if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
+				if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
 					merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
 				}
 			}
@@ -172,7 +172,7 @@ func (d *DnsInterceptor) removeAllowedIP(realPrefix netip.Prefix) error {
 	}
 
 	// AllowedIPs use real IPs
-	if _, err := d.allowedIPsRefcounter.Decrement(realPrefix); err != nil {
+	if _, err := d.allowedIPsRefcounter.Decrement(realPrefix, d.currentPeerKey); err != nil {
 		return fmt.Errorf("remove allowed IP %s: %v", realPrefix, err)
 	}
 
@@ -205,7 +205,7 @@ func (d *DnsInterceptor) RemoveAllowedIPs() error {
 	for _, prefixes := range d.interceptedDomains {
 		for _, prefix := range prefixes {
 			// AllowedIPs use real IPs
-			if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
+			if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
 				merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
 			}
 		}

client/internal/routemanager/dynamic/route.go

@@ -135,7 +135,7 @@ func (r *Route) RemoveAllowedIPs() error {
 	var merr *multierror.Error
 	for _, domainPrefixes := range r.dynamicDomains {
 		for _, prefix := range domainPrefixes {
-			if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
+			if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil {
 				merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
 			}
 		}
@@ -320,7 +320,7 @@ func (r *Route) removeRoutes(prefixes []netip.Prefix) ([]netip.Prefix, error) {
 			merr = multierror.Append(merr, fmt.Errorf("remove dynamic route for IP %s: %w", prefix, err))
 		}
 		if r.currentPeerKey != "" {
-			if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
+			if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil {
 				merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
 			}
 		}

client/internal/routemanager/manager.go

@@ -216,7 +216,7 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) {
 		)
 	}
 
-	m.allowedIPsRefCounter = refcounter.New(
+	m.allowedIPsRefCounter = refcounter.NewAllowedIPs(
 		func(prefix netip.Prefix, peerKey string) (string, error) {
 			// save peerKey to use it in the remove function
 			return peerKey, m.wgInterface.AddAllowedIP(peerKey, prefix)

client/internal/routemanager/refcounter/allowedips.go

@@ -0,0 +1,185 @@
+package refcounter
+
+import (
+	"errors"
+	"fmt"
+	"net/netip"
+	"sort"
+	"sync"
+
+	"github.com/hashicorp/go-multierror"
+
+	nberrors "github.com/netbirdio/netbird/client/errors"
+)
+
+// allowedIPsEntry holds the per-peer reference counts for a single prefix and which peer is
+// currently installed in WireGuard. WireGuard allows a prefix on exactly one peer, so at most
+// one peer is active at a time even when several peers reference the prefix.
+type allowedIPsEntry struct {
+	// peers maps a peerKey to the number of references holding the prefix for that peer.
+	peers map[string]int
+	// active is the peerKey currently installed in WireGuard for this prefix ("" if none).
+	active string
+	// total is the sum of all per-peer reference counts (kept in sync with peers).
+	total int
+}
+
+// AllowedIPsRefCounter is a peer-aware reference counter for WireGuard AllowedIPs.
+//
+// The generic Counter keys only by prefix and remembers a single Out value set by the first
+// caller, which it never changes. That is wrong for AllowedIPs: two independent watchers (or
+// multiple resolved domains) can reference the same prefix through different peers, and when the
+// peer currently installed in WireGuard releases its last reference the prefix must be handed over
+// to a surviving peer instead of being left pointing at the released one.
+//
+// It calls add/remove (which program WireGuard) only on the transitions that matter:
+//   - add on the first reference for a prefix, or when swapping the active peer;
+//   - remove on the last reference for a prefix, or on the old peer during a swap.
+type AllowedIPsRefCounter struct {
+	mu      sync.Mutex
+	entries map[netip.Prefix]*allowedIPsEntry
+	add     AddFunc[netip.Prefix, string, string]
+	remove  RemoveFunc[netip.Prefix, string]
+}
+
+// NewAllowedIPs creates a new peer-aware AllowedIPs reference counter.
+// add programs a prefix on a peer in WireGuard and returns the peerKey to store as the active peer.
+// remove unprograms the prefix from the given peer.
+func NewAllowedIPs(add AddFunc[netip.Prefix, string, string], remove RemoveFunc[netip.Prefix, string]) *AllowedIPsRefCounter {
+	return &AllowedIPsRefCounter{
+		entries: map[netip.Prefix]*allowedIPsEntry{},
+		add:     add,
+		remove:  remove,
+	}
+}
+
+// Increment adds a reference to prefix for peerKey. WireGuard is programmed only for the first
+// reference to a prefix; while a different peer is already installed the prefix is left with it
+// (first peer wins, HA at the WireGuard layer is not possible) and only the reference count is kept.
+func (rm *AllowedIPsRefCounter) Increment(prefix netip.Prefix, peerKey string) (Ref[string], error) {
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+
+	e, ok := rm.entries[prefix]
+	if !ok {
+		e = &allowedIPsEntry{peers: map[string]int{}}
+		rm.entries[prefix] = e
+	}
+
+	logCallerF("Increasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]",
+		prefix, peerKey, e.peers[peerKey], e.peers[peerKey]+1, e.total, e.total+1, e.active)
+
+	// Program WireGuard only when nothing is installed yet for this prefix.
+	if e.active == "" {
+		out, err := rm.add(prefix, peerKey)
+		if errors.Is(err, ErrIgnore) {
+			if e.total == 0 {
+				delete(rm.entries, prefix)
+			}
+			return Ref[string]{Count: e.total, Out: e.active}, nil
+		}
+		if err != nil {
+			if e.total == 0 {
+				delete(rm.entries, prefix)
+			}
+			return Ref[string]{}, fmt.Errorf("failed to add allowed IP %v for peer %s: %w", prefix, peerKey, err)
+		}
+		e.active = out
+	}
+
+	e.peers[peerKey]++
+	e.total++
+
+	return Ref[string]{Count: e.total, Out: e.active}, nil
+}
+
+// Decrement removes a reference to prefix for peerKey. When the peer currently installed in
+// WireGuard releases its last reference, the prefix is swapped to a surviving peer if one exists,
+// otherwise it is removed from WireGuard.
+func (rm *AllowedIPsRefCounter) Decrement(prefix netip.Prefix, peerKey string) (Ref[string], error) {
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+
+	e, ok := rm.entries[prefix]
+	if !ok {
+		logCallerF("No allowed IP reference found for prefix %v", prefix)
+		return Ref[string]{}, nil
+	}
+
+	if e.peers[peerKey] > 0 {
+		logCallerF("Decreasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]",
+			prefix, peerKey, e.peers[peerKey], e.peers[peerKey]-1, e.total, e.total-1, e.active)
+		e.peers[peerKey]--
+		e.total--
+		if e.peers[peerKey] == 0 {
+			delete(e.peers, peerKey)
+		}
+	} else {
+		logCallerF("No allowed IP reference found for prefix %v peer %s", prefix, peerKey)
+	}
+
+	// If the peer currently installed in WireGuard still holds references, nothing to reprogram.
+	// Keying the check on the active peer (not the one just released) makes this self-healing:
+	// a prior swap whose remove/add failed leaves e.active pointing at a peer with no references,
+	// and this retries the hand-off on the next Decrement instead of getting stuck.
+	if e.active != "" && e.peers[e.active] > 0 {
+		return Ref[string]{Count: e.total, Out: e.active}, nil
+	}
+
+	// Detach the stale/gone active peer from WireGuard before reprogramming.
+	if e.active != "" {
+		if err := rm.remove(prefix, e.active); err != nil {
+			return Ref[string]{Count: e.total, Out: e.active}, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err)
+		}
+		e.active = ""
+	}
+
+	// Hand the prefix over to a surviving peer, or drop the entry when none remain.
+	if survivor, ok := pickSurvivor(e.peers); ok {
+		out, err := rm.add(prefix, survivor)
+		if err != nil {
+			return Ref[string]{Count: e.total, Out: ""}, fmt.Errorf("swap allowed IP %v to peer %s: %w", prefix, survivor, err)
+		}
+		e.active = out
+		return Ref[string]{Count: e.total, Out: e.active}, nil
+	}
+
+	delete(rm.entries, prefix)
+	return Ref[string]{Count: 0, Out: ""}, nil
+}
+
+// Flush removes all prefixes from WireGuard and clears the counter.
+func (rm *AllowedIPsRefCounter) Flush() error {
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+
+	var merr *multierror.Error
+	for prefix, e := range rm.entries {
+		if e.active == "" {
+			continue
+		}
+		logCallerF("Flushing allowed IP for prefix %v peer %s", prefix, e.active)
+		if err := rm.remove(prefix, e.active); err != nil {
+			merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err))
+		}
+	}
+
+	clear(rm.entries)
+
+	return nberrors.FormatErrorOrNil(merr)
+}
+
+// pickSurvivor deterministically selects a peer still referencing the prefix. WireGuard cannot do
+// multipath for a single prefix, so any surviving peer is a valid winner; the choice is made stable
+// (lowest peerKey) for predictable behavior and testability.
+func pickSurvivor(peers map[string]int) (string, bool) {
+	if len(peers) == 0 {
+		return "", false
+	}
+	keys := make([]string, 0, len(peers))
+	for k := range peers {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	return keys[0], true
+}

client/internal/routemanager/refcounter/allowedips_test.go

@@ -0,0 +1,241 @@
+package refcounter
+
+import (
+	"errors"
+	"net/netip"
+	"testing"
+)
+
+// fakeWG models WireGuard's cryptokey routing: a prefix can be installed on exactly one peer.
+// failAdd/failRemove make the next add/remove fail once, to exercise the self-healing error paths.
+type fakeWG struct {
+	installed  map[netip.Prefix]string
+	adds       int
+	removes    int
+	failAdd    bool
+	failRemove bool
+}
+
+func newFakeWG() *fakeWG {
+	return &fakeWG{installed: map[netip.Prefix]string{}}
+}
+
+func (f *fakeWG) counter() *AllowedIPsRefCounter {
+	return NewAllowedIPs(
+		func(prefix netip.Prefix, peerKey string) (string, error) {
+			if f.failAdd {
+				f.failAdd = false
+				return "", errors.New("add failed")
+			}
+			f.adds++
+			f.installed[prefix] = peerKey
+			return peerKey, nil
+		},
+		func(prefix netip.Prefix, peerKey string) error {
+			if f.failRemove {
+				f.failRemove = false
+				return errors.New("remove failed")
+			}
+			f.removes++
+			// only clear if this peer is the one installed, mirroring wg semantics
+			if f.installed[prefix] == peerKey {
+				delete(f.installed, prefix)
+			}
+			return nil
+		},
+	)
+}
+
+func mustPrefix(t *testing.T, s string) netip.Prefix {
+	t.Helper()
+	p, err := netip.ParsePrefix(s)
+	if err != nil {
+		t.Fatalf("parse prefix %q: %v", s, err)
+	}
+	return p
+}
+
+func mustIncrement(t *testing.T, c *AllowedIPsRefCounter, p netip.Prefix, peer string) Ref[string] {
+	t.Helper()
+	ref, err := c.Increment(p, peer)
+	if err != nil {
+		t.Fatalf("Increment(%v, %s): %v", p, peer, err)
+	}
+	return ref
+}
+
+func mustDecrement(t *testing.T, c *AllowedIPsRefCounter, p netip.Prefix, peer string) Ref[string] {
+	t.Helper()
+	ref, err := c.Decrement(p, peer)
+	if err != nil {
+		t.Fatalf("Decrement(%v, %s): %v", p, peer, err)
+	}
+	return ref
+}
+
+// TestAllowedIPs_SwapOnActivePeerRemoval reproduces the reported bug: two networks with the same
+// prefix routed by different peers. Removing the network whose peer is installed must hand the
+// prefix over to the surviving peer instead of leaving it on the removed one.
+func TestAllowedIPs_SwapOnActivePeerRemoval(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	mustIncrement(t, c, p, "peerA")
+	mustIncrement(t, c, p, "peerB")
+	// First peer wins while both are present.
+	if got := f.installed[p]; got != "peerA" {
+		t.Fatalf("expected peerA installed, got %q", got)
+	}
+
+	// Remove the active peer's network -> must swap to peerB.
+	mustDecrement(t, c, p, "peerA")
+	if got := f.installed[p]; got != "peerB" {
+		t.Fatalf("BUG: prefix stuck on removed peer, want peerB got %q", got)
+	}
+
+	// Remove the last one -> prefix gone.
+	mustDecrement(t, c, p, "peerB")
+	if _, ok := f.installed[p]; ok {
+		t.Fatalf("expected prefix removed, still installed on %q", f.installed[p])
+	}
+}
+
+// TestAllowedIPs_RemoveNonActivePeer removing a non-installed peer must not touch WireGuard.
+func TestAllowedIPs_RemoveNonActivePeer(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	mustIncrement(t, c, p, "peerA")
+	mustIncrement(t, c, p, "peerB")
+	removesBefore := f.removes
+
+	mustDecrement(t, c, p, "peerB")
+	if f.installed[p] != "peerA" {
+		t.Fatalf("active peer must stay peerA, got %q", f.installed[p])
+	}
+	if f.removes != removesBefore {
+		t.Fatalf("removing a non-active peer must not call wg remove")
+	}
+}
+
+// TestAllowedIPs_SamePeerMultipleRefs two references via the same peer must keep the prefix until
+// the last reference is released (the reason the per-peer count must be an int, not a set).
+func TestAllowedIPs_SamePeerMultipleRefs(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	mustIncrement(t, c, p, "peerA")
+	mustIncrement(t, c, p, "peerA")
+	if f.adds != 1 {
+		t.Fatalf("expected a single wg add for the same peer, got %d", f.adds)
+	}
+
+	mustDecrement(t, c, p, "peerA")
+	if f.installed[p] != "peerA" {
+		t.Fatalf("prefix must stay while a reference remains, got %q", f.installed[p])
+	}
+	if f.removes != 0 {
+		t.Fatalf("no wg remove expected while a reference remains, got %d", f.removes)
+	}
+
+	mustDecrement(t, c, p, "peerA")
+	if _, ok := f.installed[p]; ok {
+		t.Fatalf("prefix must be removed after last reference")
+	}
+}
+
+// TestAllowedIPs_RefCountAndActive checks the Ref returned to callers (used for the HA-disabled log).
+func TestAllowedIPs_RefCountAndActive(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	ref := mustIncrement(t, c, p, "peerA")
+	if ref.Count != 1 || ref.Out != "peerA" {
+		t.Fatalf("want {1, peerA}, got {%d, %q}", ref.Count, ref.Out)
+	}
+	ref = mustIncrement(t, c, p, "peerB")
+	if ref.Count != 2 || ref.Out != "peerA" {
+		t.Fatalf("want {2, peerA}, got {%d, %q}", ref.Count, ref.Out)
+	}
+}
+
+// TestAllowedIPs_Flush removes everything installed and clears the counter.
+func TestAllowedIPs_Flush(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p1 := mustPrefix(t, "10.44.8.0/24")
+	p2 := mustPrefix(t, "10.44.9.0/24")
+
+	mustIncrement(t, c, p1, "peerA")
+	mustIncrement(t, c, p2, "peerB")
+
+	if err := c.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if len(f.installed) != 0 {
+		t.Fatalf("expected all prefixes removed, got %v", f.installed)
+	}
+	// After flush, a fresh increment must add again.
+	mustIncrement(t, c, p1, "peerC")
+	if f.installed[p1] != "peerC" {
+		t.Fatalf("counter not reset after flush")
+	}
+}
+
+// TestAllowedIPs_SelfHealAfterSwapAddError ensures a failed add during a swap does not permanently
+// strand the prefix: the next Decrement (or Increment) must retry and install a surviving peer.
+func TestAllowedIPs_SelfHealAfterSwapAddError(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	mustIncrement(t, c, p, "peerA")
+	mustIncrement(t, c, p, "peerB")
+	mustIncrement(t, c, p, "peerC")
+
+	// Removing the active peerA triggers a swap to a survivor; make the add fail once.
+	f.failAdd = true
+	if _, err := c.Decrement(p, "peerA"); err == nil {
+		t.Fatalf("expected error from failed swap add")
+	}
+	if _, ok := f.installed[p]; ok {
+		t.Fatalf("nothing should be installed after a failed swap add, got %q", f.installed[p])
+	}
+
+	// A later Decrement of a non-active survivor must retry the hand-off (self-heal), not stay stuck.
+	ref := mustDecrement(t, c, p, "peerC")
+	if got := f.installed[p]; got == "" {
+		t.Fatalf("self-heal failed: prefix left unrouted after add recovered")
+	}
+	if ref.Out == "" {
+		t.Fatalf("expected an active peer after self-heal, got empty")
+	}
+}
+
+// TestAllowedIPs_SelfHealAfterRemoveError ensures a failed remove during a swap is retried instead
+// of leaving e.active stuck on a peer that no longer holds references.
+func TestAllowedIPs_SelfHealAfterRemoveError(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	mustIncrement(t, c, p, "peerA")
+	mustIncrement(t, c, p, "peerB")
+
+	// Releasing active peerA must detach it (remove) then add peerB; fail the remove once.
+	f.failRemove = true
+	if _, err := c.Decrement(p, "peerA"); err == nil {
+		t.Fatalf("expected error from failed remove")
+	}
+
+	// Next Decrement of the non-active survivor retries: removes stale peerA, installs peerB.
+	mustDecrement(t, c, p, "peerB")
+	// peerB had only one ref, so after retry the prefix is fully released.
+	if _, ok := f.installed[p]; ok {
+		t.Fatalf("expected prefix released after self-heal, still on %q", f.installed[p])
+	}
+}

client/internal/routemanager/refcounter/types.go

@@ -5,5 +5,7 @@ import "net/netip"
 // RouteRefCounter is a Counter for Route, it doesn't take any input on Increment and doesn't use any output on Decrement
 type RouteRefCounter = Counter[netip.Prefix, struct{}, struct{}]
 
-// AllowedIPsRefCounter is a Counter for AllowedIPs, it takes a peer key on Increment and passes it back to Decrement
-type AllowedIPsRefCounter = Counter[netip.Prefix, string, string]
+// AllowedIPsRefCounter tracks WireGuard AllowedIPs per prefix. Unlike the generic Counter it is peer-aware:
+// a prefix can be claimed by several peers at once and WireGuard allows a given prefix on exactly one peer,
+// so the counter records the per-peer reference count and swaps the installed peer when the active one is released.
+// See allowedips.go.

client/internal/routemanager/static/route.go

@@ -15,6 +15,11 @@ type Route struct {
 	route                *route.Route
 	routeRefCounter      *refcounter.RouteRefCounter
 	allowedIPsRefcounter *refcounter.AllowedIPsRefCounter
+	// currentPeerKey is the routing peer this watcher currently has the prefix installed on
+	// (the HA winner elected by the watcher). It can differ from route.Peer and change on
+	// failover, so it is recorded on AddAllowedIPs and used on RemoveAllowedIPs to decrement
+	// the exact peer that was incremented.
+	currentPeerKey string
 }
 
 func NewRoute(params common.HandlerParams) *Route {
@@ -52,12 +57,15 @@ func (r *Route) AddAllowedIPs(peerKey string) error {
 			ref.Out,
 		)
 	}
+	r.currentPeerKey = peerKey
 	return nil
 }
 
 func (r *Route) RemoveAllowedIPs() error {
-	if _, err := r.allowedIPsRefcounter.Decrement(r.route.Network); err != nil {
-		return err
+	var err error
+	if _, decErr := r.allowedIPsRefcounter.Decrement(r.route.Network, r.currentPeerKey); decErr != nil {
+		err = fmt.Errorf("remove allowed IP %s: %w", r.route.Network, decErr)
 	}
-	return nil
+	r.currentPeerKey = ""
+	return err
 }

client/ios/NetBirdSDK/client.go

@@ -22,6 +22,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/listener"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
+	"github.com/netbirdio/netbird/client/mdm"
 	"github.com/netbirdio/netbird/client/system"
 	"github.com/netbirdio/netbird/formatter"
 	"github.com/netbirdio/netbird/route"
@@ -83,6 +84,13 @@ type Client struct {
 	// preloadedConfig holds config loaded from JSON (used on tvOS where file writes are blocked)
 	preloadedConfig *profilemanager.Config
 
+	// mdmLoader holds the per-Client MDM policy source. Set by
+	// SetMDMPolicyFetcher (called from the Swift side at extension
+	// init). Each Run passes this loader to the resolved Config so
+	// applyMDMPolicy picks up the active overlay. Nil means "MDM
+	// enforcement off for this Client".
+	mdmLoader *mdm.Loader
+
 	stateMu       sync.RWMutex
 	connectClient *internal.ConnectClient
 	config        *profilemanager.Config
@@ -143,6 +151,7 @@ func (c *Client) Run(fd int32, interfaceName string, envList *EnvList) error {
 		if err != nil {
 			return err
 		}
+		c.applyMDMOverlay(cfg)
 	}
 	c.recorder.UpdateManagementAddress(cfg.ManagementURL.String())
 	c.recorder.UpdateRosenpass(cfg.RosenpassEnabled, cfg.RosenpassPermissive)
@@ -215,6 +224,7 @@ func (c *Client) DebugBundle(anonymize bool) (string, error) {
 				return "", fmt.Errorf("load config: %w", err)
 			}
 		}
+		c.applyMDMOverlay(cfg)
 	}
 
 	deps := debug.GeneratorDependencies{
@@ -364,6 +374,7 @@ func (c *Client) IsLoginRequired() bool {
 			// If we can't load config, assume login is required
 			return true
 		}
+		c.applyMDMOverlay(cfg)
 	}
 
 	if cfg == nil {
@@ -412,6 +423,7 @@ func (c *Client) LoginForMobile() string {
 		log.Errorf("LoginForMobile: failed to load config: %v", err)
 		return fmt.Sprintf("failed to load config: %v", err)
 	}
+	c.applyMDMOverlay(cfg)
 
 	oAuthFlow, err := auth.NewOAuthFlow(ctx, cfg, false, false, "")
 	if err != nil {

client/ios/NetBirdSDK/mdm.go

@@ -0,0 +1,82 @@
+//go:build ios
+
+package NetBirdSDK
+
+import (
+	"encoding/json"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/internal/profilemanager"
+	"github.com/netbirdio/netbird/client/mdm"
+)
+
+// PolicyFetcher is the mobile-side bridge for the MDM managed-config
+// snapshot. The native layer (Swift) implements this and registers
+// the instance per Client via Client.SetMDMPolicyFetcher. Every
+// invocation of fetchJSON must read the current
+// UserDefaults.standard.dictionary(forKey: "com.apple.configuration.managed")
+// and return the result as a JSON-encoded map[string]any string.
+//
+// JSON is used because gomobile does not support map[string]any
+// crossing the Objective-C boundary — the adapter on the Go side
+// parses the string back into the map[string]any expected by
+// mdm.Loader.
+//
+// Return value contract:
+//   - "" (empty)         : interpreted as "no MDM source / no managed keys"
+//   - "{}"               : managed config explicitly empty
+//   - "{...}"            : JSON object with key/value pairs
+//   - malformed JSON     : logged and treated as empty
+type PolicyFetcher interface {
+	FetchJSON() string
+}
+
+// jsonFetcherAdapter wraps a gomobile-exposed PolicyFetcher into the
+// internal mdm.PolicyFetcher interface, taking care of JSON decoding
+// on every Fetch.
+type jsonFetcherAdapter struct {
+	inner PolicyFetcher
+}
+
+func (a *jsonFetcherAdapter) Fetch() map[string]any {
+	raw := a.inner.FetchJSON()
+	if raw == "" {
+		return nil
+	}
+	var out map[string]any
+	if err := json.Unmarshal([]byte(raw), &out); err != nil {
+		log.Warnf("MDM mobile fetcher: invalid JSON payload from native: %v", err)
+		return nil
+	}
+	return out
+}
+
+// SetMDMPolicyFetcher registers the native-provided MDM policy fetcher
+// on this Client. Call once from the gomobile-init code (Swift
+// AppDelegate / PacketTunnelProvider.startTunnel) before invoking Run.
+// Passing nil disables MDM enforcement on this Client.
+//
+// The fetcher is held as a *mdm.Loader instance on the Client (no
+// package-level state) — multiple Clients in the same process get
+// independent Loaders, and tests can inject fakes per Client.
+func (c *Client) SetMDMPolicyFetcher(p PolicyFetcher) {
+	if p == nil {
+		c.mdmLoader = mdm.NewLoader(nil)
+		return
+	}
+	c.mdmLoader = mdm.NewLoader(&jsonFetcherAdapter{inner: p})
+}
+
+// applyMDMOverlay applies the Client-held MDM Loader's current policy
+// on top of the just-read Config. Called immediately after every
+// UpdateOrCreateConfig / DirectUpdateOrCreateConfig — profilemanager's
+// apply() initialises the policy to empty and leaves overlay
+// responsibility to the lifecycle owner. No-op when no fetcher was
+// registered.
+func (c *Client) applyMDMOverlay(cfg *profilemanager.Config) {
+	if cfg == nil || c.mdmLoader == nil {
+		return
+	}
+	cfg.ApplyMDMPolicy(c.mdmLoader.Load())
+}

client/mdm/policy.go

@@ -104,16 +104,48 @@ func NewPolicy(values map[string]any) *Policy {
 	return &Policy{values: values}
 }
 
-// LoadPolicy reads the platform-native MDM configuration. Returns an
-// empty (but non-nil) Policy when no source is present, the source is
-// empty, or the platform is unsupported.
+// PolicyFetcher is implemented by mobile platforms (Android / iOS) that
+// push the OS-managed configuration into the Go runtime instead of
+// having Go read an on-disk source directly. Desktop platforms ignore
+// this interface — Loader.loadPlatform on windows/darwin reads the
+// registry / plist on its own. A Loader constructed with a non-nil
+// fetcher delegates to it on mobile; passing nil disables MDM
+// enforcement (loadPlatform returns nil values).
+type PolicyFetcher interface {
+	Fetch() map[string]any
+}
+
+// Loader is the DI-friendly entry point for reading the active MDM
+// policy. Construct one at the daemon's lifecycle owner (Server on
+// desktop, gomobile-exposed bridge on mobile) and pass it to anything
+// that needs to read MDM state (the reload ticker, profilemanager's
+// Config). Each callsite has the Loader handed in instead of looking
+// up package-level state.
+type Loader struct {
+	fetcher PolicyFetcher
+}
+
+// NewLoader constructs a Loader. The fetcher is consulted only on
+// mobile builds (ios || android); on desktop it is unused but accepted
+// to keep a single constructor signature across platforms — pass nil
+// on desktop.
+func NewLoader(f PolicyFetcher) *Loader {
+	return &Loader{fetcher: f}
+}
+
+// Load reads the platform-native MDM configuration and returns a
+// Policy. Returns an empty (but non-nil) Policy when no source is
+// present, the source is empty, or the platform is unsupported.
 //
 // Diagnostic logging differentiates the three states:
 //   - source absent / unsupported platform: trace log only
 //   - source present, zero keys:             info "MDM enrolled (no managed keys)"
 //   - source present, N keys:                info "MDM enrolled with N managed keys: [...]"
-func LoadPolicy() *Policy {
-	values, err := loadPlatformPolicy()
+func (l *Loader) Load() *Policy {
+	if l == nil {
+		return &Policy{values: map[string]any{}}
+	}
+	values, err := l.loadPlatform()
 	if err != nil {
 		log.Tracef("MDM policy load: %v", err)
 		return &Policy{values: map[string]any{}}

client/mdm/policy_darwin.go

@@ -25,8 +25,10 @@ import (
 // writable plist, as a defense against tampered installs.
 const policyPlistPath = "/Library/Managed Preferences/io.netbird.client.plist"
 
-// loadPlatformPolicy reads the MDM-managed configuration from the macOS
-// managed-preferences plist at policyPlistPath. Returns:
+// loadPlatform reads the MDM-managed configuration from the macOS
+// managed-preferences plist at policyPlistPath. The Loader's fetcher
+// field is unused on this platform — the plist is the authoritative
+// source. Returns:
 //   - (nil, nil)  when the plist is absent (device not MDM-enrolled for
 //     NetBird, or admin has not yet pushed a payload)
 //   - (map, nil)  with N entries when N managed values are present
@@ -39,7 +41,13 @@ const policyPlistPath = "/Library/Managed Preferences/io.netbird.client.plist"
 // skipped so a stray entry in the payload does not block startup.
 // Native plist value types map naturally onto the Policy accessor
 // expectations (GetString / GetBool / GetInt / GetStringSlice).
-func loadPlatformPolicy() (map[string]any, error) {
+func (l *Loader) loadPlatform() (map[string]any, error) {
+	// Honour the injected fetcher when present so tests (and any
+	// future non-macOS MDM channel) can short-circuit the plist read
+	// with a scripted policy.
+	if l != nil && l.fetcher != nil {
+		return l.fetcher.Fetch(), nil
+	}
 	f, err := os.Open(policyPlistPath)
 	if err != nil {
 		if errors.Is(err, fs.ErrNotExist) {

client/mdm/policy_mobile.go

@@ -2,13 +2,14 @@
 
 package mdm
 
-// loadPlatformPolicy is unused on mobile: the native layer (Swift on iOS,
-// Kotlin/Java on Android) reads the OS managed-config store and pushes the
-// resulting dictionary in-process via a gomobile entry point that lands in
-// Phase 5 / Phase 6. The stub keeps the package compilable for mobile
-// builds and returns (nil, nil) — the platform-absent sentinel that
-// LoadPolicy in policy.go treats as "no MDM source present".
-func loadPlatformPolicy() (map[string]any, error) {
-	//nolint:nilnil // (nil, nil) is the documented platform-absent sentinel; see LoadPolicy.
-	return nil, nil
+// loadPlatform reads the OS-managed configuration via the native
+// PolicyFetcher injected at Loader construction. Returns
+// (nil, nil) — the platform-absent sentinel that Loader.Load treats as
+// "no MDM source present" — when no fetcher was provided.
+func (l *Loader) loadPlatform() (map[string]any, error) {
+	if l == nil || l.fetcher == nil {
+		//nolint:nilnil // (nil, nil) is the documented platform-absent sentinel; see Loader.Load.
+		return nil, nil
+	}
+	return l.fetcher.Fetch(), nil
 }

client/mdm/policy_other.go

@@ -2,13 +2,17 @@
 
 package mdm
 
-// loadPlatformPolicy returns no policy on platforms without an MDM channel
-// (Linux, FreeBSD). MDM enforcement is off and the client behaves as if
-// the feature did not exist. Returns (nil, nil) — the platform-absent
-// sentinel the caller (LoadPolicy in policy.go) treats as "no MDM
-// source present"; an error here would just translate to the same
-// outcome with an extra log line.
-func loadPlatformPolicy() (map[string]any, error) {
-	//nolint:nilnil // (nil, nil) is the documented platform-absent sentinel; see LoadPolicy.
+// loadPlatform reads the MDM policy on platforms without a native MDM
+// channel (Linux, FreeBSD). When no fetcher was injected the policy is
+// (nil, nil) — the platform-absent sentinel that Loader.Load treats as
+// "MDM enforcement disabled". A non-nil fetcher takes precedence: it
+// is the test-seam used by unit tests to inject a scripted policy
+// without touching the OS, and the same hook supports any future
+// non-mobile OS that grows an out-of-band MDM channel.
+func (l *Loader) loadPlatform() (map[string]any, error) {
+	if l != nil && l.fetcher != nil {
+		return l.fetcher.Fetch(), nil
+	}
+	//nolint:nilnil // (nil, nil) is the documented platform-absent sentinel; see Loader.Load.
 	return nil, nil
 }

client/mdm/policy_test.go

@@ -155,10 +155,12 @@ func TestPolicy_GetStringSlice(t *testing.T) {
 	})
 }
 
-func TestLoadPolicy_PlatformStubReturnsEmpty(t *testing.T) {
-	// loadPlatformPolicy is a stub on every OS for Phase 1. LoadPolicy must
-	// degrade gracefully and never return nil.
-	p := LoadPolicy()
+func TestLoader_NilFetcherReturnsEmpty(t *testing.T) {
+	// Loader.Load with no fetcher (desktop construction) must degrade
+	// gracefully and never return nil; on linux loadPlatform is a stub
+	// returning (nil, nil), and Load is expected to translate that
+	// into a non-nil empty Policy.
+	p := NewLoader(nil).Load()
 	require.NotNil(t, p)
 	assert.True(t, p.IsEmpty())
 	assert.Empty(t, p.ManagedKeys())

client/mdm/policy_windows.go

@@ -61,8 +61,10 @@ func readRegistryValue(k registry.Key, name, canonical string, out map[string]an
 	}
 }
 
-// loadPlatformPolicy reads the MDM-managed configuration from the
-// Windows registry under HKLM\Software\Policies\NetBird. Returns:
+// loadPlatform reads the MDM-managed configuration from the Windows
+// registry under HKLM\Software\Policies\NetBird. The Loader's fetcher
+// field is unused on this platform — the registry is the
+// authoritative source. Returns:
 //   - (nil, nil)  when the key is absent (device not MDM-enrolled for NetBird)
 //   - (map, nil)  with N entries when N managed values are set (N may be 0)
 //   - (nil, err)  on open / enumerate registry errors
@@ -70,7 +72,13 @@ func readRegistryValue(k registry.Key, name, canonical string, out map[string]an
 // Per-value type coercion + skip-on-error is delegated to
 // readRegistryValue. Unknown value names are logged and skipped so a
 // malformed deployment does not block startup.
-func loadPlatformPolicy() (map[string]any, error) {
+func (l *Loader) loadPlatform() (map[string]any, error) {
+	// Honour the injected fetcher when present so tests (and any
+	// future non-Windows MDM channel) can short-circuit the registry
+	// read with a scripted policy.
+	if l != nil && l.fetcher != nil {
+		return l.fetcher.Fetch(), nil
+	}
 	k, err := registry.OpenKey(registry.LOCAL_MACHINE, policyRegistryPath, registry.QUERY_VALUE)
 	if err != nil {
 		if errors.Is(err, registry.ErrNotExist) {

client/mdm/ticker.go

@@ -15,33 +15,33 @@ import (
 // instead, hence anticipating the ticker mechanism entirely.
 const DefaultReloadInterval = 1 * time.Minute
 
-// policyLoader is the indirection through which the ticker reads the
-// OS-native policy, both for the initial observation and on every tick.
-// Production points it at LoadPolicy; tests in this package override it to
-// feed a scripted sequence of policies without touching the real OS store.
-var policyLoader = LoadPolicy
-
-// Ticker periodically re-reads the OS-native MDM policy via LoadPolicy and
-// invokes the onChange callback (supplied to Run) whenever the observed
-// Policy diverges from the last observation (added / removed / changed
-// keys). Launch with Run from a goroutine; cancel the supplied context
-// to stop.
+// Ticker periodically re-reads the OS-native MDM policy via the
+// injected Loader and invokes the onChange callback (supplied to Run)
+// whenever the observed Policy diverges from the last observation
+// (added / removed / changed keys). Launch with Run from a goroutine;
+// cancel the supplied context to stop.
 type Ticker struct {
 	interval time.Duration
+	loader   *Loader
 	prev     *Policy
 }
 
 // NewTicker constructs a Ticker that will re-read the OS-native policy
-// every reloadInterval once Run is called.
-// The initial snapshot is populated by calling policyLoader at
+// every reloadInterval once Run is called. The Loader is injected so
+// the ticker doesn't depend on any package-level state — production
+// passes the daemon-owned Loader, tests pass a fake Loader (built with
+// a fake PolicyFetcher).
+//
+// The initial snapshot is populated by calling loader.Load() at
 // construction time so the first tick only fires
 // onChange when the policy actually changed since boot — without
 // this baseline the first tick would report every currently-managed
 // key as "added" and trigger a spurious engine restart.
-func NewTicker(reloadInterval time.Duration) *Ticker {
+func NewTicker(reloadInterval time.Duration, loader *Loader) *Ticker {
 	return &Ticker{
 		interval: reloadInterval,
-		prev:     policyLoader(),
+		loader:   loader,
+		prev:     loader.Load(),
 	}
 }
 
@@ -58,7 +58,7 @@ func (t *Ticker) Run(ctx context.Context, onChange func(prev, curr *Policy) erro
 			log.Info("MDM policy reload ticker stopped")
 			return
 		case <-tk.C:
-			curr := policyLoader()
+			curr := t.loader.Load()
 			if policiesEqual(t.prev, curr) {
 				continue
 			}

client/mdm/ticker_test.go

@@ -13,28 +13,40 @@ import (
 // testReloadInterval for speeding up the ticker cadence under `go test`
 const testReloadInterval = 1 * time.Second
 
-// withPolicyLoader overrides the package-level policyLoader for the duration
-// of the test so the ticker observes a scripted policy instead of the real
-// OS-native store. The original loader is restored on cleanup.
-func withPolicyLoader(t *testing.T, fn func() *Policy) {
-	t.Helper()
-	prev := policyLoader
-	policyLoader = fn
-	t.Cleanup(func() { policyLoader = prev })
+// fakePolicyFetcher implements PolicyFetcher returning a scripted
+// policy map. Goroutine-safe so the test can mutate the script while
+// the ticker is observing it.
+type fakePolicyFetcher struct {
+	mu     sync.Mutex
+	values map[string]any
+}
+
+func (f *fakePolicyFetcher) Fetch() map[string]any {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if f.values == nil {
+		return nil
+	}
+	out := make(map[string]any, len(f.values))
+	for k, v := range f.values {
+		out[k] = v
+	}
+	return out
+}
+
+func (f *fakePolicyFetcher) set(values map[string]any) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.values = values
 }
 
 func TestTicker_FiresOnChangeWithDelta(t *testing.T) {
-	var mu sync.Mutex
-	current := NewPolicy(nil) // initial observation: empty (no enforcement)
-	withPolicyLoader(t, func() *Policy {
-		mu.Lock()
-		defer mu.Unlock()
-		return current
-	})
+	fetcher := &fakePolicyFetcher{} // initial observation: empty (no enforcement)
+	loader := NewLoader(fetcher)
 
 	type change struct{ prev, curr *Policy }
 	changes := make(chan change, 1)
-	tk := NewTicker(testReloadInterval)
+	tk := NewTicker(testReloadInterval, loader)
 	require.Equal(t, testReloadInterval, tk.interval)
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -49,15 +61,13 @@ func TestTicker_FiresOnChangeWithDelta(t *testing.T) {
 		})
 		close(done)
 	}()
-	// Stop Run and wait for it to exit before returning, so the policyLoader
-	// restore in t.Cleanup can't race the ticker goroutine still reading it.
+	// Stop Run and wait for it to exit before returning, so the test
+	// goroutine doesn't race the still-running ticker.
 	defer func() { cancel(); <-done }()
 
-	// Flip the OS-observed policy from empty to one managed key. The next
-	// tick must detect the diff and invoke onChange.
-	mu.Lock()
-	current = NewPolicy(map[string]any{KeyManagementURL: "https://mdm.example.com:443"})
-	mu.Unlock()
+	// Flip the OS-observed policy from empty to one managed key. The
+	// next tick must detect the diff and invoke onChange.
+	fetcher.set(map[string]any{KeyManagementURL: "https://mdm.example.com:443"})
 
 	select {
 	case c := <-changes:
@@ -69,12 +79,11 @@ func TestTicker_FiresOnChangeWithDelta(t *testing.T) {
 }
 
 func TestTicker_NoCallbackWhenPolicyUnchanged(t *testing.T) {
-	withPolicyLoader(t, func() *Policy {
-		return NewPolicy(map[string]any{KeyBlockInbound: true})
-	})
+	fetcher := &fakePolicyFetcher{values: map[string]any{KeyBlockInbound: true}}
+	loader := NewLoader(fetcher)
 
 	fired := make(chan struct{}, 1)
-	tk := NewTicker(testReloadInterval)
+	tk := NewTicker(testReloadInterval, loader)
 
 	ctx, cancel := context.WithCancel(context.Background())
 	done := make(chan struct{})
@@ -90,8 +99,8 @@ func TestTicker_NoCallbackWhenPolicyUnchanged(t *testing.T) {
 	}()
 	defer func() { cancel(); <-done }()
 
-	// Over ~2 ticks at the 1s test cadence the policy never changes, so the
-	// diff guard must suppress the callback entirely.
+	// Over ~2 ticks at the 1s test cadence the policy never changes,
+	// so the diff guard must suppress the callback entirely.
 	select {
 	case <-fired:
 		t.Fatal("onChange fired despite an unchanged policy")

client/server/mdm.go

@@ -21,10 +21,6 @@ import (
 // a no-op echo, never as a conflict with the policy.
 const preSharedKeyRedactedSentinel = "**********"
 
-// loadMDMPolicy is the indirection used by server handlers to read the
-// active MDM policy. Tests override this to inject a fake policy.
-var loadMDMPolicy = mdm.LoadPolicy
-
 // conflictCheck is a value-aware comparison between a single field in
 // the incoming request and the corresponding MDM-enforced value. It
 // runs only when the field was actually set in the request (presence

client/server/server.go

@@ -123,6 +123,15 @@ type Server struct {
 	// stopped by the rootCtx cancellation.
 	mdmTicker *mdm.Ticker
 
+	// mdmLoader is the daemon-owned source of the active MDM policy.
+	// Constructed once during Server.Start (with a nil PolicyFetcher on
+	// desktop — the build-tagged Loader.loadPlatform reads the OS
+	// registry / plist directly) and injected into every consumer:
+	// mdmTicker for its periodic reload, the SetConfig / Login MDM
+	// gates for conflict detection, and every Config produced via
+	// getConfig() so its apply() picks up the same overlay.
+	mdmLoader *mdm.Loader
+
 	updateManager *updater.Manager
 
 	jwtCache *jwtCache
@@ -197,8 +206,14 @@ func (s *Server) Start() error {
 	// Runs re-resolves Config (re-running profilemanager.Config.apply which
 	// applies the freshly-read MDM policy as the last layer) and brings
 	// the engine back with the new values.
+	if s.mdmLoader == nil {
+		// Desktop builds pass a nil PolicyFetcher: the Loader's
+		// build-tagged loadPlatform reads the OS source directly
+		// (registry on Windows, plist on macOS, no-op elsewhere).
+		s.mdmLoader = mdm.NewLoader(nil)
+	}
 	if s.mdmTicker == nil {
-		s.mdmTicker = mdm.NewTicker(mdm.DefaultReloadInterval)
+		s.mdmTicker = mdm.NewTicker(mdm.DefaultReloadInterval, s.mdmLoader)
 		go s.mdmTicker.Run(s.rootCtx, s.onMDMPolicyChange)
 	}
 
@@ -406,7 +421,7 @@ func (s *Server) SetConfig(callerCtx context.Context, msg *proto.SetConfigReques
 	// by the active MDM policy. The error carries an MDMManagedFields-
 	// Violation detail listing the offending key names. Non-conflicting
 	// fields in the same request are not applied either.
-	policy := loadMDMPolicy()
+	policy := s.mdmLoader.Load()
 	if err := rejectMDMManagedFieldConflicts(mdmManagedFieldConflicts(msg, policy)); err != nil {
 		return nil, err
 	}
@@ -531,7 +546,7 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro
 		if s.checkUpdateSettingsDisabled() {
 			return nil, gstatus.Errorf(codes.Unavailable, errUpdateSettingsDisabled)
 		}
-		policy := loadMDMPolicy()
+		policy := s.mdmLoader.Load()
 		if err := rejectMDMManagedFieldConflicts(loginRequestMDMConflicts(msg, policy)); err != nil {
 			return nil, err
 		}
@@ -1260,6 +1275,12 @@ func (s *Server) getConfig(activeProf *profilemanager.ActiveProfileState) (*prof
 		return nil, false, fmt.Errorf("failed to get config: %w", err)
 	}
 
+	// Apply the daemon-owned MDM policy on top of the just-resolved
+	// Config. profilemanager's apply() initialises the policy to
+	// empty — the Loader lives outside Config, so this overlay step
+	// is driven externally here.
+	config.ApplyMDMPolicy(s.mdmLoader.Load())
+
 	return config, configExisted, nil
 }
 
@@ -1309,6 +1330,9 @@ func (s *Server) logoutFromProfile(ctx context.Context, profile *profilemanager.
 	if err != nil {
 		return fmt.Errorf("profile '%s' not found", profile.ID)
 	}
+	// Honour any MDM-enforced ManagementURL when issuing the logout
+	// RPC: the user-stored value may have been overridden by policy.
+	config.ApplyMDMPolicy(s.mdmLoader.Load())
 
 	return s.sendLogoutRequestWithConfig(ctx, config)
 }
@@ -1909,6 +1933,11 @@ func (s *Server) GetConfig(ctx context.Context, req *proto.GetConfigRequest) (*p
 		log.Errorf("failed to get active profile config: %v", err)
 		return nil, fmt.Errorf("failed to get active profile config: %w", err)
 	}
+	// Overlay the active MDM policy so the response's MDMManagedFields
+	// list reflects what the GUI / CLI must render as read-only.
+	// profilemanager.GetConfig itself returns a Config without the
+	// overlay (Loader lives outside profilemanager).
+	cfg.ApplyMDMPolicy(s.mdmLoader.Load())
 	managementURL := cfg.ManagementURL
 	adminURL := cfg.AdminURL
 

client/server/setconfig_mdm_test.go

@@ -16,14 +16,40 @@ import (
 	"github.com/netbirdio/netbird/client/proto"
 )
 
-// withMDMPolicy temporarily overrides the server-package loadMDMPolicy hook
-// so SetConfig observes the supplied Policy. Restores the original loader
-// at test cleanup.
-func withMDMPolicy(t *testing.T, policy *mdm.Policy) {
+// fakeMDMFetcher implements mdm.PolicyFetcher returning a pre-set
+// policy map. Tests build one per Server instance to inject a
+// scripted MDM overlay via a Loader rather than via package-level state.
+type fakeMDMFetcher struct{ values map[string]any }
+
+func (f *fakeMDMFetcher) Fetch() map[string]any { return f.values }
+
+// withMDMPolicy installs an mdm.Loader on the given Server whose
+// loadPlatform returns the supplied Policy's underlying values. Use
+// after setupServerWithProfile to inject the scripted policy the
+// SetConfig / Login MDM gates will observe.
+func withMDMPolicy(t *testing.T, s *Server, policy *mdm.Policy) {
 	t.Helper()
-	prev := loadMDMPolicy
-	loadMDMPolicy = func() *mdm.Policy { return policy }
-	t.Cleanup(func() { loadMDMPolicy = prev })
+	values := map[string]any{}
+	if policy != nil {
+		for _, k := range policy.ManagedKeys() {
+			if v, ok := policy.GetString(k); ok {
+				values[k] = v
+				continue
+			}
+			if v, ok := policy.GetBool(k); ok {
+				values[k] = v
+				continue
+			}
+			if v, ok := policy.GetInt(k); ok {
+				values[k] = v
+				continue
+			}
+			if v, ok := policy.GetStringSlice(k); ok {
+				values[k] = v
+			}
+		}
+	}
+	s.mdmLoader = mdm.NewLoader(&fakeMDMFetcher{values: values})
 }
 
 // setupServerWithProfile mirrors the boilerplate of TestSetConfig_AllFieldsSaved:
@@ -89,12 +115,11 @@ func extractViolation(t *testing.T, err error) *proto.MDMManagedFieldsViolation
 }
 
 func TestSetConfig_MDMReject_SingleField(t *testing.T) {
-	withMDMPolicy(t, mdm.NewPolicy(map[string]any{
+	s, ctx, profName, username, _ := setupServerWithProfile(t)
+	withMDMPolicy(t, s, mdm.NewPolicy(map[string]any{
 		mdm.KeyManagementURL: "https://mdm.example.com:443",
 	}))
 
-	s, ctx, profName, username, _ := setupServerWithProfile(t)
-
 	_, err := s.SetConfig(ctx, &proto.SetConfigRequest{
 		ProfileName:   profName,
 		Username:      username,
@@ -106,14 +131,13 @@ func TestSetConfig_MDMReject_SingleField(t *testing.T) {
 }
 
 func TestSetConfig_MDMReject_MultipleFields(t *testing.T) {
-	withMDMPolicy(t, mdm.NewPolicy(map[string]any{
+	s, ctx, profName, username, _ := setupServerWithProfile(t)
+	withMDMPolicy(t, s, mdm.NewPolicy(map[string]any{
 		mdm.KeyManagementURL:    "https://mdm.example.com:443",
 		mdm.KeyBlockInbound:     true,
 		mdm.KeyRosenpassEnabled: true,
 	}))
 
-	s, ctx, profName, username, _ := setupServerWithProfile(t)
-
 	blockInbound := false
 	rosenpassEnabled := false
 	_, err := s.SetConfig(ctx, &proto.SetConfigRequest{
@@ -137,12 +161,11 @@ func TestSetConfig_MDMReject_AllOrNothing(t *testing.T) {
 	// enforced field AND a non-enforced field (RosenpassEnabled).
 	// The whole request must be rejected — non-conflicting fields are not
 	// applied either.
-	withMDMPolicy(t, mdm.NewPolicy(map[string]any{
+	s, ctx, profName, username, cfgPath := setupServerWithProfile(t)
+	withMDMPolicy(t, s, mdm.NewPolicy(map[string]any{
 		mdm.KeyManagementURL: "https://mdm.example.com:443",
 	}))
 
-	s, ctx, profName, username, cfgPath := setupServerWithProfile(t)
-
 	rosenpassEnabled := true
 	_, err := s.SetConfig(ctx, &proto.SetConfigRequest{
 		ProfileName:      profName,
@@ -164,12 +187,11 @@ func TestSetConfig_MDMReject_AllOrNothing(t *testing.T) {
 func TestSetConfig_MDMAllow_NonManagedFields(t *testing.T) {
 	// MDM enforces ManagementURL but the user only writes RosenpassEnabled.
 	// Request must succeed.
-	withMDMPolicy(t, mdm.NewPolicy(map[string]any{
+	s, ctx, profName, username, _ := setupServerWithProfile(t)
+	withMDMPolicy(t, s, mdm.NewPolicy(map[string]any{
 		mdm.KeyManagementURL: "https://mdm.example.com:443",
 	}))
 
-	s, ctx, profName, username, _ := setupServerWithProfile(t)
-
 	rosenpassEnabled := true
 	resp, err := s.SetConfig(ctx, &proto.SetConfigRequest{
 		ProfileName:      profName,
@@ -198,12 +220,11 @@ func TestSetConfig_MDMAllow_ManagementURLPortNormalized(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			withMDMPolicy(t, mdm.NewPolicy(map[string]any{
+			s, ctx, profName, username, _ := setupServerWithProfile(t)
+			withMDMPolicy(t, s, mdm.NewPolicy(map[string]any{
 				mdm.KeyManagementURL: tc.mdmURL,
 			}))
 
-			s, ctx, profName, username, _ := setupServerWithProfile(t)
-
 			rosenpassEnabled := true
 			resp, err := s.SetConfig(ctx, &proto.SetConfigRequest{
 				ProfileName:      profName,
@@ -220,9 +241,8 @@ func TestSetConfig_MDMAllow_ManagementURLPortNormalized(t *testing.T) {
 
 func TestSetConfig_MDMEmpty_NoEnforcement(t *testing.T) {
 	// No MDM policy active: any field can be written.
-	withMDMPolicy(t, mdm.NewPolicy(nil))
-
 	s, ctx, profName, username, _ := setupServerWithProfile(t)
+	withMDMPolicy(t, s, mdm.NewPolicy(nil))
 
 	resp, err := s.SetConfig(ctx, &proto.SetConfigRequest{
 		ProfileName:   profName,

client/ui/frontend/src/modules/auto-update/UpdateVersionCard.tsx

@@ -2,6 +2,7 @@ import { type ReactNode } from "react";
 import { useTranslation } from "react-i18next";
 import { Browser } from "@wailsio/runtime";
 import { DownloadIcon, NotepadText } from "lucide-react";
+import { Update as UpdateSvc } from "@bindings/services";
 import { Button } from "@/components/buttons/Button";
 import { useClientVersion } from "@/contexts/ClientVersionContext";
 import { cn } from "@/lib/cn";
@@ -14,6 +15,12 @@ function openUrl(url: string) {
     });
 }
 
+function openInstallerDownload() {
+    UpdateSvc.DownloadURL()
+        .then(openUrl)
+        .catch(() => openUrl(GITHUB_RELEASES));
+}
+
 export function UpdateVersionCard() {
     const { t } = useTranslation();
     const { updateVersion, enforced, triggerUpdate } = useClientVersion();
@@ -37,11 +44,7 @@ export function UpdateVersionCard() {
                         {t("update.card.installNow")}
                     </Button>
                 ) : (
-                    <Button
-                        variant={"primary"}
-                        size={"xs"}
-                        onClick={() => openUrl(GITHUB_RELEASES)}
-                    >
+                    <Button variant={"primary"} size={"xs"} onClick={openInstallerDownload}>
                         <DownloadIcon size={14} />
                         {t("update.card.getInstaller")}
                     </Button>

client/ui/main.go

@@ -281,6 +281,9 @@ func newApplication(onSecondInstance func()) *application.App {
 		Linux: application.LinuxOptions{
 			ProgramName: "netbird",
 		},
+		Windows: application.WindowsOptions{
+			WndProcInterceptor: endSessionInterceptor(),
+		},
 		SingleInstance: &application.SingleInstanceOptions{
 			UniqueID: "io.netbird.ui",
 			OnSecondInstanceLaunch: func(_ application.SecondInstanceData) {
@@ -367,6 +370,9 @@ func newMainWindow(app *application.App, prefStore *preferences.Store) *applicat
 
 	// Hide instead of quit on close; "really quit" is reached via tray -> Quit.
 	window.RegisterHook(events.Common.WindowClosing, func(e *application.WindowEvent) {
+		if services.ShuttingDown() {
+			return
+		}
 		e.Cancel()
 		window.Hide()
 	})

client/ui/services/shutdown.go

@@ -0,0 +1,24 @@
+package services
+
+import "sync/atomic"
+
+var (
+	sessionEnding atomic.Bool
+	quitting      atomic.Bool
+)
+
+func BeginSessionEnd() {
+	sessionEnding.Store(true)
+}
+
+func AbortSessionEnd() {
+	sessionEnding.Store(false)
+}
+
+func BeginShutdown() {
+	quitting.Store(true)
+}
+
+func ShuttingDown() bool {
+	return sessionEnding.Load() || quitting.Load()
+}

client/ui/services/update.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/client/ui/updater"
+	"github.com/netbirdio/netbird/version"
 )
 
 // UpdateResult mirrors TriggerUpdateResponse.
@@ -33,6 +34,12 @@ func (s *Update) GetState() updater.State {
 	return s.holder.Get()
 }
 
+// DownloadURL returns the platform-appropriate installer download link for
+// manual (non-enforced) updates.
+func (s *Update) DownloadURL() string {
+	return version.DownloadUrl()
+}
+
 // Quit exits the app. Scheduled off the calling goroutine so the JS caller's
 // response returns before the runtime tears down.
 func (s *Update) Quit() {

client/ui/services/windowmanager.go

@@ -154,6 +154,9 @@ func NewWindowManager(app *application.App, mainWindow *application.WebviewWindo
 	})
 	// Hide (not destroy) on close to keep React state; reset to General for a flash-free reopen.
 	s.settings.RegisterHook(events.Common.WindowClosing, func(e *application.WindowEvent) {
+		if ShuttingDown() {
+			return
+		}
 		e.Cancel()
 		s.app.Event.Emit(EventSettingsOpen, "general")
 		s.settings.Hide()
@@ -393,6 +396,9 @@ func (s *WindowManager) CloseWelcome() {
 // OpenError shows the custom error dialog; title/message are pre-localised and ride in the
 // start URL. A second error replaces the open one via SetURL. Singleton, destroyed on close.
 func (s *WindowManager) OpenError(title, message string) {
+	if ShuttingDown() {
+		return
+	}
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	startURL := errorDialogURL(title, message)

client/ui/shutdown_other.go

@@ -0,0 +1,7 @@
+//go:build !windows && !android && !ios && !freebsd && !js
+
+package main
+
+func endSessionInterceptor() func(hwnd uintptr, msg uint32, wParam, lParam uintptr) (uintptr, bool) {
+	return nil
+}

client/ui/shutdown_windows.go

@@ -0,0 +1,36 @@
+//go:build windows
+
+package main
+
+import (
+	"os"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/ui/services"
+)
+
+const (
+	wmQueryEndSession = 0x0011
+	wmEndSession      = 0x0016
+)
+
+func endSessionInterceptor() func(hwnd uintptr, msg uint32, wParam, lParam uintptr) (uintptr, bool) {
+	return func(_ uintptr, msg uint32, wParam, _ uintptr) (uintptr, bool) {
+		switch msg {
+		case wmQueryEndSession:
+			services.BeginSessionEnd()
+			return 1, true
+		case wmEndSession:
+			if wParam == 0 {
+				services.AbortSessionEnd()
+				return 0, true
+			}
+			log.Info("windows session is ending; exiting immediately")
+			os.Exit(0)
+			return 0, true
+		default:
+			return 0, false
+		}
+	}
+}

client/ui/tray.go

@@ -32,9 +32,8 @@ const (
 
 	quitDownTimeout = 5 * time.Second
 
-	urlGitHubRepo     = "https://github.com/netbirdio/netbird"
-	urlGitHubReleases = "https://github.com/netbirdio/netbird/releases/latest"
-	urlDocs           = "https://docs.netbird.io"
+	urlGitHubRepo = "https://github.com/netbirdio/netbird"
+	urlDocs       = "https://docs.netbird.io"
 )
 
 // TrayServices bundles the services the tray menu needs, grouped so NewTray
@@ -453,6 +452,7 @@ func (t *Tray) buildMenu() *application.Menu {
 }
 
 func (t *Tray) handleQuit() {
+	services.BeginShutdown()
 	t.profileMu.Lock()
 	if t.switchCancel != nil {
 		t.switchCancel()

client/ui/tray_notify.go

@@ -25,6 +25,9 @@ type sendFn func(notifications.NotificationOptions) error
 // event-dispatch goroutine that panic is fatal process-wide; recover() turns
 // it into a logged no-op.
 func safeSendNotification(send sendFn, what string, opts notifications.NotificationOptions) (err error) {
+	if services.ShuttingDown() {
+		return nil
+	}
 	defer func() {
 		if r := recover(); r != nil {
 			log.Errorf("notify %s: recovered from panic (notification bus unavailable): %v", what, r)

client/ui/tray_update.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/netbirdio/netbird/client/ui/services"
 	"github.com/netbirdio/netbird/client/ui/updater"
+	"github.com/netbirdio/netbird/version"
 )
 
 // trayUpdater owns the tray UI that reacts to auto-update. Composed inside Tray.
@@ -76,15 +77,15 @@ func (u *trayUpdater) applyLanguage() {
 	u.refreshMenuItem(state)
 }
 
-// handleClick opens the GitHub releases page when not Enforced, otherwise shows
-// the progress page and asks the daemon to start the installer.
+// handleClick opens the installer download link when not Enforced, otherwise
+// shows the progress page and asks the daemon to start the installer.
 func (u *trayUpdater) handleClick() {
 	u.mu.Lock()
 	state := u.state
 	u.mu.Unlock()
 
 	if !state.Enforced {
-		_ = u.app.Browser.OpenURL(urlGitHubReleases)
+		_ = u.app.Browser.OpenURL(version.DownloadUrl())
 		return
 	}
 

docs/agent-networks/01-end-to-end-flows.md

@@ -109,7 +109,7 @@ sequenceDiagram
         Chk->>Inj: continue
         Inj->>Inj: inject NetBird identity headers per provider config
         Inj->>Grd: continue
-        Grd->>Grd: enforce model allowlist
+        Grd->>Grd: enforce per-provider allowlist (fail-closed backstop)
         Grd->>Up: forward (over WireGuard)
         Up-->>Resp: response (JSON or SSE stream)
         Resp->>Resp: parse usage tokens, completion
@@ -135,6 +135,21 @@ sequenceDiagram
   (`redact_pii = settings.RedactPii`). Phones, emails, credit cards,
   PII names — see `redact.go` for the full set. See
   [`modules/31-proxy-middleware-builtin.md`](modules/31-proxy-middleware-builtin.md).
+- The model allowlist is enforced in TWO places. `CheckLLMPolicyLimits`
+  is authoritative: it resolves the policy that governs this
+  (provider, caller-groups) and denies (`deny_code = llm_policy.model_blocked`)
+  when no applicable policy permits the model — so an allowlist scoped to
+  one group/provider never leaks to another, and an un-guardrailed policy
+  is genuinely unrestricted. `llm_guardrail` is a per-provider fail-closed
+  backstop: it only carries an allowlist for a provider every authorising
+  policy restricts, and blocks unknown/undetermined models even when
+  management is unreachable. Because that backstop allowlist is the UNION
+  of every restricting policy's models, per-group narrowing lives only in
+  the authoritative check: during a `CheckLLMPolicyLimits` outage
+  `llm_limit_check` fails open, so a caller can reach any model in the
+  provider's union — a group scoped to model A could reach model B if
+  another group restricts the same provider to B. This is the documented
+  fail-open trade-off; a future flag may switch it to fail-closed.
 - SSE streaming requires special handling on the response side; the
   parser must handle partial chunks without buffering the whole
   stream. See [`modules/32-proxy-llm-parsers.md`](modules/32-proxy-llm-parsers.md).

docs/agent-networks/modules/21-management-agentnetwork.md

@@ -122,7 +122,7 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
   | on_request | 1 | `llm_router` | `{"providers":[{id, models[], upstream_*, auth_header_*, allowed_group_ids[]}]}` | **true** |
   | on_request | 2 | `llm_limit_check` | `{}` | – |
   | on_request | 3 | `llm_identity_inject` | `{"providers":[{provider_id, header_pair?, json_metadata?, extra_headers?}]}` | **true** |
-  | on_request | 4 | `llm_guardrail` | `{"model_allowlist"?, "prompt_capture":{enabled,redact_pii}}` | – |
+  | on_request | 4 | `llm_guardrail` | `{"provider_allowlists"?: {providerID: []model}, "prompt_capture":{enabled,redact_pii}}` | – |
   | on_response | 5 | `llm_limit_record` | `{}` (runs LAST at runtime) | – |
   | on_response | 6 | `cost_meter` | `{}` | – |
   | on_response | 7 | `llm_response_parser` | `{"capture_completion": <bool>, "redact_pii"?: true}` | – |

docs/agent-networks/modules/31-proxy-middleware-builtin.md

@@ -244,7 +244,7 @@ no mocks. Tests: `TestChain_AllowPath_StampsAttributionAndRecordsCounter`
 | `llm_router` | `{providers: [{id, models, upstream_scheme, upstream_host, upstream_path?, auth_header_name, auth_header_value, allowed_group_ids}]}` |
 | `llm_limit_check` | `{}` — pulls `MgmtClient` from `FactoryContext` |
 | `llm_identity_inject` | `{providers: [{provider_id, header_pair?|json_metadata?, extra_headers?}]}` |
-| `llm_guardrail` | `{model_allowlist: []string, prompt_capture: {enabled, redact_pii}}` |
+| `llm_guardrail` | `{provider_allowlists: {providerID: []string}, prompt_capture: {enabled, redact_pii}}` — allowlist keyed by resolved provider id; a provider absent from the map is unrestricted (fail-closed backstop; authoritative per-policy/group check is management's `CheckLLMPolicyLimits`) |
 | `llm_response_parser` | `{redact_pii?, capture_completion?: *bool}` |
 | `cost_meter` | `{pricing_path?}` (basename inside data-dir; defaults `pricing.yaml`) |
 | `llm_limit_record` | `{}` — same pattern as `llm_limit_check` |

e2e/agentnetwork/guardrail_block_test.go

@@ -0,0 +1,209 @@
+//go:build e2e
+
+package agentnetwork
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/e2e/harness"
+	"github.com/netbirdio/netbird/shared/management/http/api"
+)
+
+// pathRoutedGuardrailCase is one provider's self-contained scenario: its own
+// provider, its own guardrail whose allowlist holds ONLY that provider's
+// allowed model, and its own policy. Each case runs in isolation (its own
+// proxy + client), so the guardrail the proxy enforces contains exactly this
+// provider's model — never a mixed cross-provider list.
+type pathRoutedGuardrailCase struct {
+	name       string
+	catalogID  string // agent-network catalog provider id
+	wire       string // harness.WireVertex | harness.WireBedrock
+	allowEntry string // the single model id put on the guardrail allowlist
+	allowModel string // model id sent that MUST be served (200)
+	blockModel string // model id sent that MUST be denied (403 model_blocked)
+}
+
+// TestGuardrailBlocksUnselectedModel_PathRouted is the end-to-end regression
+// guard for the customer report that a model-allowlist guardrail attached to a
+// policy has no effect for PATH-ROUTED providers — where the model travels in
+// the URL, not the JSON body: Google Vertex (…/models/{model}:rawPredict) and
+// AWS Bedrock (/model/{id}/invoke).
+//
+// Each provider is tested in isolation with a guardrail allowlisting a single
+// model of its own: the allowed model (in the URL path) is served (200) and an
+// unselected model (in the URL path) is denied 403 by the guardrail
+// (llm_policy.model_blocked) before the upstream. The Vertex case mirrors the
+// customer verbatim — allow Sonnet, and the unselected model is the exact
+// claude-opus-4-6 they reported reaching the model unblocked. The Bedrock case
+// sends a region-prefixed, versioned inference-profile id so URL-path model
+// normalization is exercised too.
+//
+// The provider is catch-all (no models), so the router forwards any model and a
+// 403 can only come from the guardrail, never model_not_routable. Only the
+// upstream LLM is mocked (the vLLM nginx answers any path with 200); management
+// synth/reconcile, the proxy middleware chain (URL-path model extraction,
+// router, guardrail) and the tunnel are all real, and the guardrail denies
+// before the upstream is dialed so the mock cannot influence the block. A
+// static bearer api key is used so the router injects a static Authorization
+// header instead of minting a GCP token — the only reason path-routed providers
+// normally need live credentials — so the test runs with none and is always on.
+func TestGuardrailBlocksUnselectedModel_PathRouted(t *testing.T) {
+	cases := []pathRoutedGuardrailCase{
+		{
+			name:       "vertex",
+			catalogID:  "vertex_ai_api",
+			wire:       harness.WireVertex,
+			allowEntry: "claude-sonnet-4-5",
+			allowModel: "claude-sonnet-4-5",
+			blockModel: "claude-opus-4-6", // the customer-reported model
+		},
+		{
+			name:       "bedrock",
+			catalogID:  "bedrock_api",
+			wire:       harness.WireBedrock,
+			allowEntry: "anthropic.claude-sonnet-4-5", // normalized catalog id
+			allowModel: "us.anthropic.claude-sonnet-4-5-v1:0",
+			blockModel: "us.anthropic.claude-opus-4-8-v1:0",
+		},
+	}
+
+	for _, tc := range cases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			runPathRoutedGuardrailCase(t, tc)
+		})
+	}
+}
+
+func runPathRoutedGuardrailCase(t *testing.T, tc pathRoutedGuardrailCase) {
+	t.Helper()
+
+	const (
+		vertexProject = "e2e-project"
+		vertexRegion  = "global"
+	)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+	defer cancel()
+
+	vllm, err := harness.StartVLLM(ctx, srv)
+	require.NoError(t, err, "start mock upstream")
+	t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
+
+	grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-" + tc.name})
+	require.NoError(t, err, "create group")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
+
+	ephemeral := false
+	sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
+		Name:       "e2e-guardrail-" + tc.name + "-client",
+		Type:       "reusable",
+		ExpiresIn:  86400,
+		UsageLimit: 0,
+		AutoGroups: []string{grp.Id},
+		Ephemeral:  &ephemeral,
+	})
+	require.NoError(t, err, "mint setup key")
+	require.NotEmpty(t, sk.Key, "setup key plaintext")
+
+	// Catch-all provider (no models) so the router forwards any model; a static
+	// bearer key means the router injects a static auth header instead of minting
+	// a GCP token. Bootstraps the cluster if it isn't already.
+	staticKey := "static-e2e-token"
+	prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:             tc.name,
+		ProviderId:       tc.catalogID,
+		UpstreamUrl:      vllm.URL,
+		ApiKey:           &staticKey,
+		Enabled:          ptr(true),
+		BootstrapCluster: ptr(harness.AgentNetworkCluster),
+	})
+	require.NoError(t, err, "create %s provider", tc.name)
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
+
+	// Guardrail allowlisting ONLY this provider's allowed model.
+	var gr api.AgentNetworkGuardrailRequest
+	gr.Name = "e2e-guardrail-" + tc.name
+	gr.Checks.ModelAllowlist.Enabled = true
+	gr.Checks.ModelAllowlist.Models = []string{tc.allowEntry}
+	guard, err := srv.CreateGuardrail(ctx, gr)
+	require.NoError(t, err, "create guardrail")
+	t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
+
+	enabled := true
+	pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-guardrail-" + tc.name,
+		Enabled:                &enabled,
+		SourceGroups:           []string{grp.Id},
+		DestinationProviderIds: []string{prov.Id},
+		GuardrailIds:           &[]string{guard.Id},
+	})
+	require.NoError(t, err, "create policy")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
+
+	settings, err := srv.GetSettings(ctx)
+	require.NoError(t, err, "read settings")
+	require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
+
+	proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-"+tc.name+"-proxy")
+	require.NoError(t, err, "mint proxy token")
+	px, err := harness.StartProxy(ctx, srv, proxyToken)
+	require.NoError(t, err, "start proxy")
+	t.Cleanup(func() { _ = px.Terminate(context.Background()) })
+
+	cl, err := harness.StartClient(ctx, srv, sk.Key)
+	require.NoError(t, err, "start client")
+	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
+
+	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	// Probe first: the GET resolves the endpoint and its first packet wakes the
+	// lazy proxy peer, so WaitProxyPeer then observes it connected.
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
+	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
+		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
+	}
+
+	send := func(model string) (int, string) {
+		var code int
+		var body string
+		var cerr error
+		switch tc.wire {
+		case harness.WireVertex:
+			code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
+		case harness.WireBedrock:
+			code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
+		default:
+			t.Fatalf("unsupported wire %q", tc.wire)
+		}
+		require.NoError(t, cerr, "request must reach the proxy for %s", tc.name)
+		return code, body
+	}
+
+	// Allowed model (in the URL path) is served. Retry to absorb tunnel/DNS
+	// jitter on the first call over the freshly warmed tunnel.
+	var code int
+	var body string
+	deadline := time.Now().Add(90 * time.Second)
+	for time.Now().Before(deadline) {
+		code, body = send(tc.allowModel)
+		if code == 200 {
+			break
+		}
+		time.Sleep(5 * time.Second)
+	}
+	assert.Equal(t, 200, code,
+		"allowed %s model (URL path) must be served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
+
+	// Unselected model (in the URL path) must be blocked by the guardrail.
+	code, body = send(tc.blockModel)
+	assert.Equal(t, 403, code,
+		"unselected %s model (URL path) must be denied, not served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
+	assert.Contains(t, body, "llm_policy.model_blocked",
+		"%s denial must come from the guardrail allowlist, not routing; body: %s", tc.name, body)
+}

e2e/agentnetwork/guardrail_groupswitch_test.go

@@ -0,0 +1,205 @@
+//go:build e2e
+
+package agentnetwork
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/e2e/harness"
+	"github.com/netbirdio/netbird/shared/management/http/api"
+)
+
+// TestGuardrailGroupSwitchTakesEffectAfterTTL proves that moving a peer between
+// groups flips its model-allowlist decision once the proxy's tunnel-peer cache
+// expires. The peer's groups reach the guardrail via ValidateTunnelPeer, which
+// the proxy caches; the switch is invisible until that cache expires. The proxy
+// runs with a short NB_PROXY_TUNNEL_CACHE_TTL so the flip happens in seconds
+// instead of the 5-minute default.
+//
+// Setup: one catch-all provider declaring modelA + modelB; polA (grpA -> allow
+// modelA) and polB (grpB -> allow modelB). The client starts in grpA. modelA is
+// served and modelB denied; after switching the client grpA -> grpB, modelB is
+// served and modelA denied. The cross-group deny comes from management's
+// per-policy/group CheckLLMPolicyLimits (the proxy backstop carries the union).
+func TestGuardrailGroupSwitchTakesEffectAfterTTL(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+	defer cancel()
+
+	const (
+		modelA = "e2e-model-a"
+		modelB = "e2e-model-b"
+		// Short tunnel-cache TTL so a group switch propagates in seconds.
+		// Exercises the NB_PROXY_TUNNEL_CACHE_TTL override.
+		cacheTTL = 3 * time.Second
+	)
+
+	vllm, err := harness.StartVLLM(ctx, srv)
+	require.NoError(t, err, "start mock upstream")
+	t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
+
+	grpA, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-a"})
+	require.NoError(t, err, "create group A")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpA.Id) })
+
+	grpB, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-b"})
+	require.NoError(t, err, "create group B")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpB.Id) })
+
+	ephemeral := false
+	sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
+		Name:       "e2e-gswitch-client",
+		Type:       "reusable",
+		ExpiresIn:  86400,
+		UsageLimit: 0,
+		AutoGroups: []string{grpA.Id}, // client starts in group A
+		Ephemeral:  &ephemeral,
+	})
+	require.NoError(t, err, "mint setup key")
+	require.NotEmpty(t, sk.Key, "setup key plaintext")
+
+	staticKey := "static-e2e-token"
+	prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:        "gswitch",
+		ProviderId:  "openai_api",
+		UpstreamUrl: vllm.URL,
+		ApiKey:      &staticKey,
+		Enabled:     ptr(true),
+		Models: &[]api.AgentNetworkProviderModel{
+			{Id: modelA, InputPer1k: 0.001, OutputPer1k: 0.001},
+			{Id: modelB, InputPer1k: 0.001, OutputPer1k: 0.001},
+		},
+		BootstrapCluster: ptr(harness.AgentNetworkCluster),
+	})
+	require.NoError(t, err, "create provider")
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
+
+	mkGuard := func(name, model string) api.AgentNetworkGuardrail {
+		var gr api.AgentNetworkGuardrailRequest
+		gr.Name = name
+		gr.Checks.ModelAllowlist.Enabled = true
+		gr.Checks.ModelAllowlist.Models = []string{model}
+		g, gerr := srv.CreateGuardrail(ctx, gr)
+		require.NoError(t, gerr, "create guardrail %s", name)
+		t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
+		return g
+	}
+	gA := mkGuard("e2e-gswitch-a", modelA)
+	gB := mkGuard("e2e-gswitch-b", modelB)
+
+	enabled := true
+	polA, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-gswitch-a",
+		Enabled:                &enabled,
+		SourceGroups:           []string{grpA.Id},
+		DestinationProviderIds: []string{prov.Id},
+		GuardrailIds:           &[]string{gA.Id},
+	})
+	require.NoError(t, err, "create policy A")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polA.Id) })
+
+	polB, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-gswitch-b",
+		Enabled:                &enabled,
+		SourceGroups:           []string{grpB.Id},
+		DestinationProviderIds: []string{prov.Id},
+		GuardrailIds:           &[]string{gB.Id},
+	})
+	require.NoError(t, err, "create policy B")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polB.Id) })
+
+	settings, err := srv.GetSettings(ctx)
+	require.NoError(t, err, "read settings")
+	require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
+
+	proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-gswitch-proxy")
+	require.NoError(t, err, "mint proxy token")
+	px, err := harness.StartProxy(ctx, srv, proxyToken, map[string]string{
+		"NB_PROXY_TUNNEL_CACHE_TTL": cacheTTL.String(),
+	})
+	require.NoError(t, err, "start proxy")
+	t.Cleanup(func() { _ = px.Terminate(context.Background()) })
+
+	cl, err := harness.StartClient(ctx, srv, sk.Key)
+	require.NoError(t, err, "start client")
+	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
+
+	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
+	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
+		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
+	}
+
+	send := func(model string) (int, string) {
+		code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
+		require.NoError(t, cerr, "request must reach the proxy")
+		return code, body
+	}
+	sendUntil := func(model string, want int, timeout time.Duration) (int, string) {
+		var code int
+		var body string
+		deadline := time.Now().Add(timeout)
+		for time.Now().Before(deadline) {
+			code, body = send(model)
+			if code == want {
+				return code, body
+			}
+			time.Sleep(2 * time.Second)
+		}
+		return code, body
+	}
+
+	// Phase 1 — client is in group A: modelA served, modelB denied.
+	code, body := sendUntil(modelA, 200, 90*time.Second)
+	assert.Equal(t, 200, code,
+		"group-A model must be served while the client is in group A; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+	code, body = send(modelB)
+	assert.Equal(t, 403, code,
+		"group-B model must be denied while the client is in group A; body: %s", body)
+	assert.Contains(t, body, "llm_policy.model_blocked")
+
+	// Switch the client peer from group A to group B.
+	peerID := clientPeerInGroup(t, ctx, grpA.Id)
+	_, err = srv.API().Groups.Update(ctx, grpB.Id, api.PutApiGroupsGroupIdJSONRequestBody{
+		Name:  grpB.Name,
+		Peers: &[]string{peerID},
+	})
+	require.NoError(t, err, "add peer to group B")
+	_, err = srv.API().Groups.Update(ctx, grpA.Id, api.PutApiGroupsGroupIdJSONRequestBody{
+		Name:  grpA.Name,
+		Peers: &[]string{},
+	})
+	require.NoError(t, err, "remove peer from group A")
+
+	// Phase 2 — after the short TTL expires the proxy re-validates the peer,
+	// sees group B, and the decision flips. Poll to absorb TTL + re-validation.
+	code, body = sendUntil(modelB, 200, 60*time.Second)
+	assert.Equal(t, 200, code,
+		"after the group switch + TTL, the group-B model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+	code, body = send(modelA)
+	assert.Equal(t, 403, code,
+		"after the switch, the old group-A model must be denied; body: %s", body)
+	assert.Contains(t, body, "llm_policy.model_blocked")
+}
+
+// clientPeerInGroup returns the id of the single peer that is a member of the
+// given group — the test client. The proxy peer is never added to test groups.
+func clientPeerInGroup(t *testing.T, ctx context.Context, groupID string) string {
+	t.Helper()
+	peers, err := srv.API().Peers.List(ctx)
+	require.NoError(t, err, "list peers")
+	for _, p := range peers {
+		for _, g := range p.Groups {
+			if g.Id == groupID {
+				return p.Id
+			}
+		}
+	}
+	t.Fatalf("no peer found in group %s", groupID)
+	return ""
+}

e2e/agentnetwork/guardrail_multipolicy_test.go

@@ -0,0 +1,201 @@
+//go:build e2e
+
+package agentnetwork
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/e2e/harness"
+	"github.com/netbirdio/netbird/shared/management/http/api"
+)
+
+// TestGuardrailMultiPolicyModelAllowlist: modelSelected served (200), grpOther's
+// modelOther denied for the grpMain client (403 model_blocked, no cross-group
+// leak), and openModel on the un-guardrailed policy's provider served (200).
+func TestGuardrailMultiPolicyModelAllowlist(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+	defer cancel()
+
+	const (
+		modelSelected = "e2e-selected"
+		modelOther    = "e2e-other"
+		openModel     = "e2e-open"
+	)
+
+	vllm, err := harness.StartVLLM(ctx, srv)
+	require.NoError(t, err, "start mock upstream")
+	t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
+
+	grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-main"})
+	require.NoError(t, err, "create main group")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
+
+	grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-other"})
+	require.NoError(t, err, "create other group")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
+
+	ephemeral := false
+	sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
+		Name:       "e2e-guardrail-mp-client",
+		Type:       "reusable",
+		ExpiresIn:  86400,
+		UsageLimit: 0,
+		AutoGroups: []string{grpMain.Id}, // client joins grpMain only
+		Ephemeral:  &ephemeral,
+	})
+	require.NoError(t, err, "mint setup key")
+	require.NotEmpty(t, sk.Key, "setup key plaintext")
+
+	staticKey := "static-e2e-token"
+	models := func(ids ...string) *[]api.AgentNetworkProviderModel {
+		out := make([]api.AgentNetworkProviderModel, 0, len(ids))
+		for _, id := range ids {
+			out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
+		}
+		return &out
+	}
+
+	// pRestricted declares the two guardrailed models so routing is deterministic
+	// (model -> provider). Created first, so it carries the bootstrap cluster.
+	pRestricted, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:             "restricted",
+		ProviderId:       "openai_api",
+		UpstreamUrl:      vllm.URL,
+		ApiKey:           &staticKey,
+		Enabled:          ptr(true),
+		Models:           models(modelSelected, modelOther),
+		BootstrapCluster: ptr(harness.AgentNetworkCluster),
+	})
+	require.NoError(t, err, "create restricted provider")
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pRestricted.Id) })
+
+	pOpen, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:        "open",
+		ProviderId:  "openai_api",
+		UpstreamUrl: vllm.URL,
+		ApiKey:      &staticKey,
+		Enabled:     ptr(true),
+		Models:      models(openModel),
+	})
+	require.NoError(t, err, "create open provider")
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pOpen.Id) })
+
+	mkGuardrail := func(name, model string) api.AgentNetworkGuardrail {
+		var gr api.AgentNetworkGuardrailRequest
+		gr.Name = name
+		gr.Checks.ModelAllowlist.Enabled = true
+		gr.Checks.ModelAllowlist.Models = []string{model}
+		g, gerr := srv.CreateGuardrail(ctx, gr)
+		require.NoError(t, gerr, "create guardrail %s", name)
+		t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
+		return g
+	}
+	gMain := mkGuardrail("e2e-guardrail-mp-main", modelSelected)
+	gOther := mkGuardrail("e2e-guardrail-mp-other", modelOther)
+
+	enabled := true
+	// polMain: grpMain restricted to modelSelected on pRestricted.
+	polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-guardrail-mp-main",
+		Enabled:                &enabled,
+		SourceGroups:           []string{grpMain.Id},
+		DestinationProviderIds: []string{pRestricted.Id},
+		GuardrailIds:           &[]string{gMain.Id},
+	})
+	require.NoError(t, err, "create main policy")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
+
+	// polOther: grpOther restricted to modelOther on the SAME provider. The
+	// client is not in grpOther, so modelOther must never be usable by it.
+	polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-guardrail-mp-other",
+		Enabled:                &enabled,
+		SourceGroups:           []string{grpOther.Id},
+		DestinationProviderIds: []string{pRestricted.Id},
+		GuardrailIds:           &[]string{gOther.Id},
+	})
+	require.NoError(t, err, "create other policy")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
+
+	// polOpen: grpMain on pOpen with NO guardrail — unrestricted.
+	polOpen, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-guardrail-mp-open",
+		Enabled:                &enabled,
+		SourceGroups:           []string{grpMain.Id},
+		DestinationProviderIds: []string{pOpen.Id},
+	})
+	require.NoError(t, err, "create open policy")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOpen.Id) })
+
+	settings, err := srv.GetSettings(ctx)
+	require.NoError(t, err, "read settings")
+	require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
+
+	proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-mp-proxy")
+	require.NoError(t, err, "mint proxy token")
+	px, err := harness.StartProxy(ctx, srv, proxyToken)
+	require.NoError(t, err, "start proxy")
+	t.Cleanup(func() { _ = px.Terminate(context.Background()) })
+
+	cl, err := harness.StartClient(ctx, srv, sk.Key)
+	require.NoError(t, err, "start client")
+	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
+
+	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
+	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
+		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
+	}
+
+	send := func(model string) (int, string) {
+		code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
+		require.NoError(t, cerr, "request must reach the proxy")
+		return code, body
+	}
+	// sendUntil200 absorbs first-call tunnel/DNS jitter on the freshly warmed tunnel.
+	sendUntil200 := func(model string) (int, string) {
+		var code int
+		var body string
+		deadline := time.Now().Add(90 * time.Second)
+		for time.Now().Before(deadline) {
+			code, body = send(model)
+			if code == 200 {
+				break
+			}
+			time.Sleep(5 * time.Second)
+		}
+		return code, body
+	}
+
+	t.Run("selected model allowed for its group", func(t *testing.T) {
+		code, body := sendUntil200(modelSelected)
+		assert.Equal(t, 200, code,
+			"grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+	})
+
+	t.Run("other group's model does not leak", func(t *testing.T) {
+		// modelOther is allowlisted only for grpOther. The grpMain client must be
+		// denied by management's per-policy/group check — not waved through by an
+		// account-wide union. This is the security-critical wrong-ALLOW guard.
+		code, body := send(modelOther)
+		assert.Equal(t, 403, code,
+			"another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+		assert.Contains(t, body, "llm_policy.model_blocked",
+			"denial must be a model-allowlist decision; body: %s", body)
+	})
+
+	t.Run("unguarded policy leaves its provider unrestricted", func(t *testing.T) {
+		// polOpen carries no guardrail, so pOpen is unrestricted for grpMain. The
+		// old account-wide union would have blocked openModel (it is on no
+		// allowlist); it must now be served — the false-DENY guard.
+		code, body := sendUntil200(openModel)
+		assert.Equal(t, 200, code,
+			"an un-guardrailed policy's provider must not be blocked by another policy's allowlist; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+	})
+}

e2e/agentnetwork/guardrail_pergroup_providers_test.go

@@ -0,0 +1,422 @@
+//go:build e2e
+
+package agentnetwork
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/e2e/harness"
+	"github.com/netbirdio/netbird/shared/management/http/api"
+)
+
+// pergroupCase describes one provider surface for the per-group allowlist matrix.
+// selectedReq/otherReq are the model identifiers as they travel in the request
+// (URL path for Bedrock/Vertex, body "model" for chat/messages). selectedAllow/
+// otherAllow are the (normalized) forms the guardrail allowlist holds — for
+// Bedrock these differ from the request form so path normalization is exercised.
+type pergroupCase struct {
+	name          string
+	catalogID     string
+	wire          string // "chat", "messages", "vertex", "bedrock"
+	models        *[]api.AgentNetworkProviderModel
+	selectedReq   string
+	selectedAllow string
+	otherReq      string
+	otherAllow    string
+
+	providerID string // filled during setup
+}
+
+// TestGuardrailPerGroupAllowlist_AllProviders proves the per-policy/group model
+// allowlist end to end across every always-on provider surface, including the
+// path-routed ones (Vertex, Bedrock) where the model travels in the URL.
+//
+// For each provider two policies target it: grpMain (the client) is allowed only
+// selectedReq; grpOther (which the client is NOT in) is allowed only otherReq.
+// The client must get selectedReq served (200) and otherReq denied (403,
+// llm_policy.model_blocked) — the cross-group no-leak property. The deny is the
+// authoritative per-policy/group decision from management (the proxy per-provider
+// backstop carries the union of both models), so this also confirms management
+// receives the correct normalized model for path-routed providers.
+func TestGuardrailPerGroupAllowlist_AllProviders(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
+	defer cancel()
+
+	const (
+		vertexProject = "e2e-project"
+		vertexRegion  = "global"
+	)
+
+	priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
+		out := make([]api.AgentNetworkProviderModel, 0, len(ids))
+		for _, id := range ids {
+			out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
+		}
+		return &out
+	}
+
+	cases := []*pergroupCase{
+		{
+			name: "openai", catalogID: "openai_api", wire: harness.WireChat,
+			models:      priced("oai-model-a", "oai-model-b"),
+			selectedReq: "oai-model-a", selectedAllow: "oai-model-a",
+			otherReq: "oai-model-b", otherAllow: "oai-model-b",
+		},
+		{
+			name: "anthropic", catalogID: "anthropic_api", wire: harness.WireMessages,
+			models:      priced("ant-model-a", "ant-model-b"),
+			selectedReq: "ant-model-a", selectedAllow: "ant-model-a",
+			otherReq: "ant-model-b", otherAllow: "ant-model-b",
+		},
+		{
+			// Vertex catalog ids travel bare in the rawPredict path.
+			name: "vertex", catalogID: "vertex_ai_api", wire: "vertex",
+			selectedReq: "claude-sonnet-4-5", selectedAllow: "claude-sonnet-4-5",
+			otherReq: "claude-opus-4-6", otherAllow: "claude-opus-4-6",
+		},
+		{
+			// Bedrock request ids are region-prefixed/versioned; the parser
+			// normalizes them to the catalog key the allowlist holds.
+			name: "bedrock", catalogID: "bedrock_api", wire: "bedrock",
+			selectedReq: "us.anthropic.claude-sonnet-4-5-v1:0", selectedAllow: "anthropic.claude-sonnet-4-5",
+			otherReq: "us.anthropic.claude-opus-4-8-v1:0", otherAllow: "anthropic.claude-opus-4-8",
+		},
+	}
+
+	vllm, err := harness.StartVLLM(ctx, srv)
+	require.NoError(t, err, "start mock upstream")
+	t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
+
+	grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-main"})
+	require.NoError(t, err, "create main group")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
+
+	grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-other"})
+	require.NoError(t, err, "create other group")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
+
+	ephemeral := false
+	sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
+		Name:       "e2e-pergroup-client",
+		Type:       "reusable",
+		ExpiresIn:  86400,
+		UsageLimit: 0,
+		AutoGroups: []string{grpMain.Id},
+		Ephemeral:  &ephemeral,
+	})
+	require.NoError(t, err, "mint setup key")
+	require.NotEmpty(t, sk.Key, "setup key plaintext")
+
+	staticKey := "static-e2e-token"
+	enabled := true
+
+	for i, c := range cases {
+		req := api.AgentNetworkProviderRequest{
+			Name:        "e2e-pergroup-" + c.name,
+			ProviderId:  c.catalogID,
+			UpstreamUrl: vllm.URL,
+			ApiKey:      &staticKey,
+			Enabled:     ptr(true),
+			Models:      c.models,
+		}
+		if i == 0 {
+			req.BootstrapCluster = ptr(harness.AgentNetworkCluster)
+		}
+		prov, perr := srv.CreateProvider(ctx, req)
+		require.NoError(t, perr, "create provider %s", c.name)
+		c.providerID = prov.Id
+		t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
+
+		gSel := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-sel", c.selectedAllow)
+		gOth := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-oth", c.otherAllow)
+
+		polMain, merr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+			Name:                   "e2e-pergroup-" + c.name + "-main",
+			Enabled:                &enabled,
+			SourceGroups:           []string{grpMain.Id},
+			DestinationProviderIds: []string{prov.Id},
+			GuardrailIds:           &[]string{gSel.Id},
+		})
+		require.NoError(t, merr, "create main policy %s", c.name)
+		t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
+
+		polOther, oerr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+			Name:                   "e2e-pergroup-" + c.name + "-other",
+			Enabled:                &enabled,
+			SourceGroups:           []string{grpOther.Id},
+			DestinationProviderIds: []string{prov.Id},
+			GuardrailIds:           &[]string{gOth.Id},
+		})
+		require.NoError(t, oerr, "create other policy %s", c.name)
+		t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
+	}
+
+	settings, err := srv.GetSettings(ctx)
+	require.NoError(t, err, "read settings")
+	require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
+
+	proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-pergroup-proxy")
+	require.NoError(t, err, "mint proxy token")
+	px, err := harness.StartProxy(ctx, srv, proxyToken)
+	require.NoError(t, err, "start proxy")
+	t.Cleanup(func() { _ = px.Terminate(context.Background()) })
+
+	cl, err := harness.StartClient(ctx, srv, sk.Key)
+	require.NoError(t, err, "start client")
+	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
+
+	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
+	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
+		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
+	}
+
+	send := func(c *pergroupCase, model string) (int, string) {
+		var code int
+		var body string
+		var cerr error
+		switch c.wire {
+		case "vertex":
+			code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
+		case "bedrock":
+			code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
+		default:
+			code, body, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, c.wire, model, "Reply with exactly: pong", "")
+		}
+		require.NoError(t, cerr, "request must reach the proxy for %s", c.name)
+		return code, body
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			// grpMain's own model is served. Retry to absorb tunnel/DNS jitter on
+			// the first call over the freshly warmed tunnel.
+			var code int
+			var body string
+			deadline := time.Now().Add(90 * time.Second)
+			for time.Now().Before(deadline) {
+				code, body = send(c, c.selectedReq)
+				if code == 200 {
+					break
+				}
+				time.Sleep(5 * time.Second)
+			}
+			assert.Equal(t, 200, code,
+				"%s: grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
+
+			// grpOther's model must NOT leak to the grpMain client.
+			code, body = send(c, c.otherReq)
+			assert.Equal(t, 403, code,
+				"%s: another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
+			assert.Contains(t, body, "llm_policy.model_blocked",
+				"%s: denial must be a model-allowlist decision, not routing; body: %s", c.name, body)
+		})
+	}
+}
+
+// TestGuardrailMultiGroupUser proves the per-policy/group decision for a caller
+// that belongs to MULTIPLE groups at once. Two scenarios, one shared stack:
+//
+//   - union across the user's groups: the client is in gUX and gUY, each with
+//     its own policy+guardrail on provider P1 (gUX->union-a, gUY->union-b). The
+//     client may use BOTH models (the union of its groups' allowlists) while a
+//     third, un-allowlisted model is denied.
+//   - an un-guardrailed group lifts the restriction: the client is in gMP and
+//     gMQ on provider P2, where gMP restricts to mix-a but gMQ's policy carries
+//     NO guardrail. Because one applicable policy is unrestricted, the client may
+//     use a model on no allowlist (mix-z) as well as mix-a.
+func TestGuardrailMultiGroupUser(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+	defer cancel()
+
+	const (
+		unionA = "mg-union-a"
+		unionB = "mg-union-b"
+		unionC = "mg-union-c" // allowlisted by neither group
+		mixA   = "mg-mix-a"
+		mixZ   = "mg-mix-z" // on no allowlist; reachable only via the un-guardrailed policy
+	)
+
+	priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
+		out := make([]api.AgentNetworkProviderModel, 0, len(ids))
+		for _, id := range ids {
+			out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
+		}
+		return &out
+	}
+
+	vllm, err := harness.StartVLLM(ctx, srv)
+	require.NoError(t, err, "start mock upstream")
+	t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
+
+	mkGroup := func(name string) *api.Group {
+		g, gerr := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: name})
+		require.NoError(t, gerr, "create group %s", name)
+		t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), g.Id) })
+		return g
+	}
+	gUX := mkGroup("e2e-mg-union-x")
+	gUY := mkGroup("e2e-mg-union-y")
+	gMP := mkGroup("e2e-mg-mix-p")
+	gMQ := mkGroup("e2e-mg-mix-q")
+
+	ephemeral := false
+	sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
+		Name:       "e2e-mg-client",
+		Type:       "reusable",
+		ExpiresIn:  86400,
+		UsageLimit: 0,
+		AutoGroups: []string{gUX.Id, gUY.Id, gMP.Id, gMQ.Id}, // client in all four groups
+		Ephemeral:  &ephemeral,
+	})
+	require.NoError(t, err, "mint setup key")
+	require.NotEmpty(t, sk.Key, "setup key plaintext")
+
+	staticKey := "static-e2e-token"
+	enabled := true
+
+	// P1 — union scenario: two restricting policies, one per group.
+	p1, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:             "e2e-mg-union",
+		ProviderId:       "openai_api",
+		UpstreamUrl:      vllm.URL,
+		ApiKey:           &staticKey,
+		Enabled:          ptr(true),
+		Models:           priced(unionA, unionB, unionC),
+		BootstrapCluster: ptr(harness.AgentNetworkCluster),
+	})
+	require.NoError(t, err, "create union provider")
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p1.Id) })
+
+	polUX, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-mg-union-x",
+		Enabled:                &enabled,
+		SourceGroups:           []string{gUX.Id},
+		DestinationProviderIds: []string{p1.Id},
+		GuardrailIds:           &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-x", unionA).Id},
+	})
+	require.NoError(t, err, "create union policy X")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUX.Id) })
+
+	polUY, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-mg-union-y",
+		Enabled:                &enabled,
+		SourceGroups:           []string{gUY.Id},
+		DestinationProviderIds: []string{p1.Id},
+		GuardrailIds:           &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-y", unionB).Id},
+	})
+	require.NoError(t, err, "create union policy Y")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUY.Id) })
+
+	// P2 — mixed scenario: one restricting policy + one un-guardrailed policy.
+	p2, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:        "e2e-mg-mix",
+		ProviderId:  "openai_api",
+		UpstreamUrl: vllm.URL,
+		ApiKey:      &staticKey,
+		Enabled:     ptr(true),
+		Models:      priced(mixA, mixZ),
+	})
+	require.NoError(t, err, "create mix provider")
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p2.Id) })
+
+	polMP, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-mg-mix-p",
+		Enabled:                &enabled,
+		SourceGroups:           []string{gMP.Id},
+		DestinationProviderIds: []string{p2.Id},
+		GuardrailIds:           &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-mix-p", mixA).Id},
+	})
+	require.NoError(t, err, "create mix policy P")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMP.Id) })
+
+	polMQ, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-mg-mix-q",
+		Enabled:                &enabled,
+		SourceGroups:           []string{gMQ.Id},
+		DestinationProviderIds: []string{p2.Id}, // NO guardrail -> unrestricted
+	})
+	require.NoError(t, err, "create mix policy Q")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMQ.Id) })
+
+	settings, err := srv.GetSettings(ctx)
+	require.NoError(t, err, "read settings")
+	require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
+
+	proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-mg-proxy")
+	require.NoError(t, err, "mint proxy token")
+	px, err := harness.StartProxy(ctx, srv, proxyToken)
+	require.NoError(t, err, "start proxy")
+	t.Cleanup(func() { _ = px.Terminate(context.Background()) })
+
+	cl, err := harness.StartClient(ctx, srv, sk.Key)
+	require.NoError(t, err, "start client")
+	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
+
+	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
+	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
+		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
+	}
+
+	send := func(model string) (int, string) {
+		code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
+		require.NoError(t, cerr, "request must reach the proxy")
+		return code, body
+	}
+	sendUntil200 := func(model string) (int, string) {
+		var code int
+		var body string
+		deadline := time.Now().Add(90 * time.Second)
+		for time.Now().Before(deadline) {
+			code, body = send(model)
+			if code == 200 {
+				break
+			}
+			time.Sleep(5 * time.Second)
+		}
+		return code, body
+	}
+
+	t.Run("union across the user's groups", func(t *testing.T) {
+		code, body := sendUntil200(unionA)
+		assert.Equal(t, 200, code, "model allowed by group X must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+
+		code, body = sendUntil200(unionB)
+		assert.Equal(t, 200, code, "model allowed by group Y must also be served (union across the user's groups); body: %s", body)
+
+		code, body = send(unionC)
+		assert.Equal(t, 403, code, "a model on neither group's allowlist must be denied; body: %s", body)
+		assert.Contains(t, body, "llm_policy.model_blocked")
+	})
+
+	t.Run("an un-guardrailed group lifts the restriction", func(t *testing.T) {
+		code, body := sendUntil200(mixA)
+		assert.Equal(t, 200, code, "the restricted group's model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+
+		code, body = sendUntil200(mixZ)
+		assert.Equal(t, 200, code,
+			"a non-allowlisted model must be served because the user is also in a group whose policy has no guardrail; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+	})
+}
+
+// mkAllowGuardrail creates a guardrail whose model allowlist is enabled and holds
+// exactly the given model, registering cleanup.
+func mkAllowGuardrail(t *testing.T, ctx context.Context, name, model string) api.AgentNetworkGuardrail {
+	t.Helper()
+	var gr api.AgentNetworkGuardrailRequest
+	gr.Name = name
+	gr.Checks.ModelAllowlist.Enabled = true
+	gr.Checks.ModelAllowlist.Models = []string{model}
+	g, err := srv.CreateGuardrail(ctx, gr)
+	require.NoError(t, err, "create guardrail %s", name)
+	t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
+	return g
+}

e2e/harness/proxy.go

@@ -38,7 +38,11 @@ type Proxy struct {
 // network, registered via the given account proxy token and serving the
 // AgentNetworkCluster over a self-signed wildcard cert. It does not wait for
 // peer connectivity — callers poll management for the proxy peer.
-func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, error) {
+// StartProxy launches the reverse-proxy container. Optional envOverrides are
+// merged into the container environment after the defaults, so callers can set
+// or override any NB_PROXY_* var (e.g. NB_PROXY_TUNNEL_CACHE_TTL for tests that
+// need a short authorization-cache window).
+func StartProxy(ctx context.Context, c *Combined, proxyToken string, envOverrides ...map[string]string) (*Proxy, error) {
 	root, err := repoRoot()
 	if err != nil {
 		return nil, err
@@ -93,6 +97,12 @@ func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, er
 		WaitingFor: wait.ForLog("Initial mapping sync complete").WithStartupTimeout(90 * time.Second),
 	}
 
+	for _, ov := range envOverrides {
+		for k, v := range ov {
+			req.Env[k] = v
+		}
+	}
+
 	ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
 		ContainerRequest: req,
 		Started:          true,

management/internals/modules/agentnetwork/manager.go

@@ -86,12 +86,17 @@ type Manager interface {
 
 // PolicySelectionInput is the per-request selection envelope. The
 // proxy populates it from CapturedData (account, user, groups) plus
-// the provider llm_router resolved.
+// the provider llm_router resolved and the model it extracted.
 type PolicySelectionInput struct {
 	AccountID  string
 	UserID     string
 	GroupIDs   []string
 	ProviderID string
+	// Model is the already-normalised upstream model id the proxy extracted
+	// (parser strips Bedrock region/version, Vertex @version), so a
+	// case-insensitive compare suffices. Empty = undetermined → not permitted
+	// (fail closed).
+	Model string
 }
 
 // PolicySelectionResult names the policy that "pays" for this request

management/internals/modules/agentnetwork/policyselect.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"math"
 	"sort"
+	"strings"
 	"time"
 
 	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
@@ -35,6 +36,10 @@ const (
 	denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded"
 	//nolint:gosec // account deny code label, not a credential
 	denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded"
+	// denyCodeModelBlocked is returned when policies govern the request's
+	// (provider, caller-groups) but none permits the model. Matches the proxy
+	// guardrail's code so both layers surface the same label.
+	denyCodeModelBlocked = "llm_policy.model_blocked"
 )
 
 // consumptionCache holds the consumption counters prefetched for one
@@ -159,6 +164,25 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec
 	}
 	candidates := filterApplicablePolicies(policies, in)
 
+	// Model-allowlist gate scoped to the matched policies: keep candidates whose
+	// guardrails permit the model (none enabled = unrestricted), deny when
+	// policies apply but none permits it. Skip the load when none has a guardrail.
+	if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) {
+		guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID)
+		if gErr != nil {
+			return nil, gErr
+		}
+		permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model)
+		if len(permitted) == 0 {
+			return &PolicySelectionResult{
+				Allow:      false,
+				DenyCode:   denyCodeModelBlocked,
+				DenyReason: modelBlockedReason(in.Model),
+			}, nil
+		}
+		candidates = permitted
+	}
+
 	// Prefetch every consumption counter the ceiling + candidate policies will
 	// read, in a single store round-trip, then score against the cache.
 	cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now)
@@ -250,6 +274,90 @@ func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput)
 	return out
 }
 
+// anyPolicyHasGuardrails reports whether any policy references at least one
+// guardrail, so the selector can skip loading guardrails when none do.
+func anyPolicyHasGuardrails(policies []*types.Policy) bool {
+	for _, p := range policies {
+		if p != nil && len(p.GuardrailIDs) > 0 {
+			return true
+		}
+	}
+	return false
+}
+
+// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
+// model-allowlist gate to resolve each candidate policy's attached guardrails.
+func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
+	guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID)
+	if err != nil {
+		return nil, fmt.Errorf("list account guardrails: %w", err)
+	}
+	byID := make(map[string]*types.Guardrail, len(guardrails))
+	for _, g := range guardrails {
+		if g != nil {
+			byID[g.ID] = g
+		}
+	}
+	return byID, nil
+}
+
+// filterModelPermittedPolicies returns the subset of policies whose guardrails
+// permit the model. Order is preserved so downstream scoring is unaffected.
+func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy {
+	out := make([]*types.Policy, 0, len(policies))
+	for _, p := range policies {
+		if policyPermitsModel(p, byID, model) {
+			out = append(out, p)
+		}
+	}
+	return out
+}
+
+// policyPermitsModel reports whether a policy permits the model. No
+// allowlist-enabled guardrail = unrestricted (permits any, incl. empty);
+// otherwise the model must be in the union of its allowlists, so an
+// empty/undetermined model fails closed.
+func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool {
+	if p == nil {
+		return false
+	}
+	wanted := normaliseModelID(model)
+	restricted := false
+	for _, gID := range p.GuardrailIDs {
+		g, ok := byID[gID]
+		if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
+			continue
+		}
+		restricted = true
+		if wanted == "" {
+			continue
+		}
+		for _, allowed := range g.Checks.ModelAllowlist.Models {
+			if normaliseModelID(allowed) == wanted {
+				return true
+			}
+		}
+	}
+	return !restricted
+}
+
+// normaliseModelID lowercases and trims a model identifier so the allowlist
+// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's
+// normaliseModel so both layers agree on what "same model" means.
+func normaliseModelID(model string) string {
+	return strings.ToLower(strings.TrimSpace(model))
+}
+
+// modelBlockedReason builds the human-readable deny reason for a model-allowlist
+// rejection. The model is quoted when known; an undetermined model is reported
+// as such so the access log distinguishes "wrong model" from "no model".
+func modelBlockedReason(model string) string {
+	if normaliseModelID(model) == "" {
+		return "request model could not be determined for the policy allowlist"
+	}
+	return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model)
+}
+
 // candidate is the per-policy intermediate the selector ranks. A
 // policy that's been exhausted on any enabled cap never makes it
 // into this slice; the selector's deny envelope carries the latest

management/internals/modules/agentnetwork/policyselect_model_test.go

@@ -0,0 +1,329 @@
+package agentnetwork
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	"github.com/netbirdio/netbird/management/server/store"
+)
+
+// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups
+// to reach providerID under the given guardrails. Uncapped keeps the selector's
+// headroom scoring trivial so these tests isolate the model-allowlist gate.
+func guardedPolicy(id, account string, sourceGroups []string, providerID string, guardrailIDs ...string) *types.Policy {
+	return &types.Policy{
+		ID:                     id,
+		AccountID:              account,
+		Enabled:                true,
+		SourceGroups:           sourceGroups,
+		DestinationProviderIDs: []string{providerID},
+		GuardrailIDs:           guardrailIDs,
+		CreatedAt:              time.Now().UTC(),
+	}
+}
+
+// allowlistGuardrail builds a guardrail whose model allowlist is enabled and
+// carries the given models.
+func allowlistGuardrail(id, account string, models ...string) *types.Guardrail {
+	return &types.Guardrail{
+		ID:        id,
+		AccountID: account,
+		Checks: types.GuardrailChecks{
+			ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: models},
+		},
+	}
+}
+
+func expectPolicies(mockStore *store.MockStore, account string, policies ...*types.Policy) {
+	mockStore.EXPECT().
+		GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), account).
+		Return(policies, nil)
+}
+
+func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...*types.Guardrail) {
+	mockStore.EXPECT().
+		GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), account).
+		Return(guardrails, nil)
+}
+
+// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist
+// decision: a policy authorises the (provider, group) but restricts the model,
+// and the requested model isn't on the list, so the request is denied.
+func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		UserID:     "user-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4",
+	})
+	require.NoError(t, err)
+	assert.False(t, res.Allow, "a model outside the only applicable policy's allowlist must be denied")
+	assert.Equal(t, denyCodeModelBlocked, res.DenyCode, "deny code must be model_blocked")
+	assert.NotEmpty(t, res.DenyReason, "deny reason must be populated")
+}
+
+// TestSelectPolicy_ModelAllowedByAllowlist is the allow counterpart: the model
+// is on the applicable policy's allowlist, so selection proceeds normally.
+func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4"))
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		UserID:     "user-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4",
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "a model on the applicable policy's allowlist must be allowed")
+	assert.Equal(t, "pol-A", res.SelectedPolicyID)
+}
+
+// TestSelectPolicy_CaseInsensitiveModelMatch proves the compare tolerates case
+// and surrounding whitespace, matching the proxy guardrail's normalisation.
+func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "  GPT-4o  "))
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "gpt-4o",
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "case/whitespace variants must match the allowlist entry")
+}
+
+// TestSelectPolicy_UnguardedPolicyIsUnrestricted is the false-deny fix: when two
+// policies authorise the same (provider, group) and one has no guardrail, that
+// policy makes the request unrestricted — not caught by the other's allowlist.
+func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	restricted := guardedPolicy("pol-restricted", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail
+	expectPolicies(mockStore, "acc-1", restricted, open)
+	expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4",
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "an un-guardrailed policy for the same (provider, group) must leave the request unrestricted")
+	assert.Equal(t, "pol-open", res.SelectedPolicyID, "the unrestricted policy must be the one that pays")
+}
+
+// TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups is the false-allow fix: a
+// model allowlisted only for grp-b must not be usable by a grp-a caller. The
+// selector considers only policies applicable to the caller's groups.
+func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	polA := guardedPolicy("pol-a", "acc-1", []string{"grp-a"}, "prov-1", "g-a")
+	polB := guardedPolicy("pol-b", "acc-1", []string{"grp-b"}, "prov-1", "g-b")
+	expectPolicies(mockStore, "acc-1", polA, polB)
+	expectGuardrails(mockStore, "acc-1",
+		allowlistGuardrail("g-a", "acc-1", "gpt-4o"),
+		allowlistGuardrail("g-b", "acc-1", "claude-opus-4"),
+	)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-a"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4", // only allowed for grp-b
+	})
+	require.NoError(t, err)
+	assert.False(t, res.Allow, "grp-b's allowlisted model must not leak to a grp-a caller")
+	assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
+}
+
+// TestSelectPolicy_UndeterminedModelFailsClosed proves the fail-closed contract
+// mirrors the proxy: with a restricted applicable policy and an empty model
+// (e.g. a path-routed shape the parser couldn't map), the request is denied.
+func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "", // undetermined
+	})
+	require.NoError(t, err)
+	assert.False(t, res.Allow, "an undetermined model must fail closed against a restricted policy")
+	assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
+}
+
+// TestSelectPolicy_DisabledAllowlistDoesNotRestrict proves a guardrail whose
+// model allowlist is disabled imposes no model restriction, even though the
+// policy references it.
+func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	disabled := &types.Guardrail{
+		ID:        "g-1",
+		AccountID: "acc-1",
+		Checks: types.GuardrailChecks{
+			ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}},
+		},
+	}
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1", disabled)
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "anything-goes",
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "a disabled allowlist must not restrict the model")
+	assert.Equal(t, "pol-A", res.SelectedPolicyID)
+}
+
+// TestSelectPolicy_UnionAcrossPolicyGuardrails proves a policy with multiple
+// allowlist guardrails permits the union of their models (not just the first).
+func TestSelectPolicy_UnionAcrossPolicyGuardrails(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1", "g-2")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1",
+		allowlistGuardrail("g-1", "acc-1", "gpt-4o"),
+		allowlistGuardrail("g-2", "acc-1", "claude-opus-4"),
+	)
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4", // only in the second guardrail's list
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "a model in any of the policy's allowlist guardrails must be permitted")
+	assert.Equal(t, "pol-A", res.SelectedPolicyID)
+}
+
+// TestSelectPolicy_GuardrailLookupErrorPropagates proves a store failure while
+// resolving the candidate policies' guardrails surfaces as an error, not a
+// silent allow/deny.
+func TestSelectPolicy_GuardrailLookupErrorPropagates(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	expectPolicies(mockStore, "acc-1", policy)
+	mockStore.EXPECT().
+		GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), "acc-1").
+		Return(nil, errors.New("store unavailable"))
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "gpt-4o",
+	})
+	require.Error(t, err, "a guardrail-lookup failure must surface as an error")
+	assert.Nil(t, res)
+}
+
+// TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted proves a
+// policy referencing a guardrail ID absent from the account's set (a stale
+// reference) imposes no model restriction — same as no guardrail.
+func TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-missing")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1")
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "anything-goes",
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "an orphaned guardrail reference must not restrict the model")
+	assert.Equal(t, "pol-A", res.SelectedPolicyID)
+}
+
+// TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter proves the model
+// gate narrows candidates before cap scoring: the permitting policy is selected
+// even though the blocked one has a larger, more attractive cap.
+func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	polBig := guardedPolicy("pol-big", "acc-1", []string{"grp-eng"}, "prov-1", "g-restrict")
+	polBig.Limits = types.PolicyLimits{
+		TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 1_000_000, WindowSeconds: 3600},
+	}
+	polSmall := guardedPolicy("pol-small", "acc-1", []string{"grp-eng"}, "prov-1", "g-permit")
+	polSmall.Limits = types.PolicyLimits{
+		TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 100, WindowSeconds: 3600},
+	}
+	expectPolicies(mockStore, "acc-1", polBig, polSmall)
+	expectGuardrails(mockStore, "acc-1",
+		allowlistGuardrail("g-restrict", "acc-1", "gpt-4o"),
+		allowlistGuardrail("g-permit", "acc-1", "claude-opus-4"),
+	)
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4", // only pol-small's guardrail permits this
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow)
+	assert.Equal(t, "pol-small", res.SelectedPolicyID,
+		"the model filter must exclude pol-big before cap scoring")
+}

management/internals/modules/agentnetwork/synthesizer.go

@@ -235,7 +235,12 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
 
 	mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
 	applyAccountCollectionControls(&mergedGuardrails, settings)
-	guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails)
+	// The proxy guardrail is a per-provider fail-closed backstop; the
+	// authoritative per-policy/group decision is management's
+	// SelectPolicyForRequest. A provider lands in this map only when every
+	// authorising policy restricts models.
+	providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
+	guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture)
 	if err != nil {
 		return nil, err
 	}
@@ -780,10 +785,12 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt
 
 // guardrailConfig is the JSON shape the proxy-side llm_guardrail
 // middleware expects. Mirrors the proxy registration documented in
-// the management→proxy contract.
+// the management→proxy contract. provider_allowlists is keyed by the
+// resolved provider id llm_router stamps; a provider absent from the map is
+// unrestricted at the proxy layer.
 type guardrailConfig struct {
-	ModelAllowlist []string               `json:"model_allowlist,omitempty"`
-	PromptCapture  guardrailPromptCapture `json:"prompt_capture"`
+	ProviderAllowlists map[string][]string    `json:"provider_allowlists,omitempty"`
+	PromptCapture      guardrailPromptCapture `json:"prompt_capture"`
 }
 
 type guardrailPromptCapture struct {
@@ -828,13 +835,10 @@ func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Se
 	merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii
 }
 
-func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
+func marshalGuardrailConfig(providerAllowlists map[string][]string, capture MergedPromptCapture) ([]byte, error) {
 	cfg := guardrailConfig{
-		ModelAllowlist: merged.ModelAllowlist,
-		PromptCapture: guardrailPromptCapture{
-			Enabled:   merged.PromptCapture.Enabled,
-			RedactPii: merged.PromptCapture.RedactPii,
-		},
+		ProviderAllowlists: providerAllowlists,
+		PromptCapture:      guardrailPromptCapture(capture),
 	}
 	out, err := json.Marshal(cfg)
 	if err != nil {
@@ -843,6 +847,74 @@ func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
 	return out, nil
 }
 
+// buildProviderAllowlists returns the proxy's per-provider backstop: a provider
+// is included only when every authorising policy restricts models (their union);
+// if any leaves it unrestricted it is omitted, so management decides per group.
+func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string {
+	type providerAcc struct {
+		models          map[string]struct{}
+		anyUnrestricted bool
+	}
+	accs := make(map[string]*providerAcc)
+	for _, p := range policies {
+		if p == nil {
+			continue
+		}
+		restricted, models := policyModelAllowlist(p, byID)
+		for _, providerID := range p.DestinationProviderIDs {
+			if providerID == "" {
+				continue
+			}
+			acc, ok := accs[providerID]
+			if !ok {
+				acc = &providerAcc{models: make(map[string]struct{})}
+				accs[providerID] = acc
+			}
+			if !restricted {
+				acc.anyUnrestricted = true
+				continue
+			}
+			for _, m := range models {
+				acc.models[m] = struct{}{}
+			}
+		}
+	}
+	out := make(map[string][]string, len(accs))
+	for providerID, acc := range accs {
+		if acc.anyUnrestricted {
+			continue
+		}
+		models := make([]string, 0, len(acc.models))
+		for m := range acc.models {
+			models = append(models, m)
+		}
+		sort.Strings(models)
+		out[providerID] = models
+	}
+	return out
+}
+
+// policyModelAllowlist reports whether a policy restricts models (has an
+// allowlist-enabled guardrail) and the union of allowed models. Models are
+// verbatim; the proxy factory lowercases/trims them at decode time.
+func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) {
+	restricted := false
+	var models []string
+	for _, gID := range p.GuardrailIDs {
+		g, ok := byID[gID]
+		if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
+			continue
+		}
+		restricted = true
+		for _, m := range g.Checks.ModelAllowlist.Models {
+			if m != "" {
+				models = append(models, m)
+			}
+		}
+	}
+	return restricted, models
+}
+
 // buildAccountService composes the per-account gateway Service. The
 // target carries the noop placeholder URL — the router middleware
 // rewrites every request to the matched provider's upstream before the
@@ -986,38 +1058,11 @@ func unionSourceGroups(policies []*types.Policy) []string {
 	return out
 }
 
-// MergedGuardrails is the JSON shape passed to the proxy via the
-// guardrail middleware's config_json. Mirrors the proxy-side
-// expectations and is intentionally distinct from
-// types.GuardrailChecks so we can evolve either side independently.
+// MergedGuardrails is the synthesiser's fold target. Only prompt capture is
+// merged here — the model allowlist is emitted per-provider, and
+// token/budget/retention moved onto Policy.Limits and account Settings.
 type MergedGuardrails struct {
-	ModelAllowlist []string            `json:"model_allowlist,omitempty"`
-	TokenLimits    MergedTokenLimits   `json:"token_limits"`
-	Budget         MergedBudget        `json:"budget"`
-	PromptCapture  MergedPromptCapture `json:"prompt_capture"`
-	Retention      MergedRetention     `json:"retention"`
-}
-
-type MergedTokenLimits struct {
-	Hourly  *MergedTokenWindow `json:"hourly,omitempty"`
-	Daily   *MergedTokenWindow `json:"daily,omitempty"`
-	Monthly *MergedTokenWindow `json:"monthly,omitempty"`
-}
-
-type MergedTokenWindow struct {
-	MaxInputTokens  int `json:"max_input_tokens,omitempty"`
-	MaxOutputTokens int `json:"max_output_tokens,omitempty"`
-}
-
-type MergedBudget struct {
-	Hourly  *MergedBudgetWindow `json:"hourly,omitempty"`
-	Daily   *MergedBudgetWindow `json:"daily,omitempty"`
-	Monthly *MergedBudgetWindow `json:"monthly,omitempty"`
-}
-
-type MergedBudgetWindow struct {
-	SoftCapUSD float64 `json:"soft_cap_usd,omitempty"`
-	HardCapUSD float64 `json:"hard_cap_usd,omitempty"`
+	PromptCapture MergedPromptCapture
 }
 
 type MergedPromptCapture struct {
@@ -1025,64 +1070,31 @@ type MergedPromptCapture struct {
 	RedactPii bool `json:"redact_pii"`
 }
 
-type MergedRetention struct {
-	Enabled bool `json:"enabled"`
-	Days    int  `json:"days"`
-}
-
-// mergeGuardrails computes the effective guardrail spec applied at the
-// proxy, given the referencing policies and the account's guardrail
-// catalogue. Policy enabled-ness is the caller's responsibility — only
-// enabled policies should be passed in.
+// mergeGuardrails folds the referencing policies' guardrails into the
+// prompt-capture decision only. The model allowlist is enforced per-policy/group
+// in management and shipped per-provider; token/budget/retention live off
+// guardrails now.
 //
-// Merge rules:
-//   - Model allowlist:   union of allowlists across policies that enable it.
-//   - Token / Budget:    most-restrictive (min of non-zero caps) per window.
-//   - Prompt capture:    enabled if any policy enables it; redact_pii sticks
-//     if any enabling policy turns it on.
-//   - Retention:         enabled if any enables it; smallest non-zero days wins.
+// Merge rule — prompt capture: enabled if any policy enables it; redact_pii
+// sticks if any enabling policy turns it on.
 func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails {
 	merged := MergedGuardrails{}
-	allowlist := make(map[string]struct{})
-	allowlistEnabled := false
-
 	for _, policy := range policies {
 		for _, gID := range policy.GuardrailIDs {
 			g, ok := byID[gID]
 			if !ok || g == nil {
 				continue
 			}
-			mergeGuardrail(g, &merged, allowlist, &allowlistEnabled)
+			mergeGuardrail(g, &merged)
 		}
 	}
-
-	if allowlistEnabled {
-		merged.ModelAllowlist = make([]string, 0, len(allowlist))
-		for m := range allowlist {
-			merged.ModelAllowlist = append(merged.ModelAllowlist, m)
-		}
-		sort.Strings(merged.ModelAllowlist)
-	}
 	return merged
 }
 
-// mergeGuardrail folds a single guardrail's enabled checks into the
-// running merge: model-allowlist models join the shared set (and flip
-// allowlistEnabled), and prompt-capture / redact-pii stick once any
-// enabling guardrail turns them on.
-//
-// TokenLimits, Budget, and Retention have moved off guardrails — token
-// and budget caps now live on the Policy itself (Policy.Limits) and
-// retention moves to account-level Settings — so they are not merged here.
-func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) {
-	if g.Checks.ModelAllowlist.Enabled {
-		*allowlistEnabled = true
-		for _, m := range g.Checks.ModelAllowlist.Models {
-			if m != "" {
-				allowlist[m] = struct{}{}
-			}
-		}
-	}
+// mergeGuardrail folds a single guardrail's prompt-capture settings into the
+// running merge: enabled / redact-pii stick once any enabling guardrail turns
+// them on.
+func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
 	if g.Checks.PromptCapture.Enabled {
 		merged.PromptCapture.Enabled = true
 		if g.Checks.PromptCapture.RedactPii {

management/internals/modules/agentnetwork/synthesizer_guardrail_realstore_test.go

@@ -81,8 +81,8 @@ func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testi
 	require.Len(t, services, 1)
 
 	cfg := decodeServiceGuardrailConfig(t, services[0])
-	assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist,
-		"model allowlist is a pure policy guardrail and must always reach the config")
+	assert.Equal(t, map[string][]string{"prov-1": {"gpt-5.4"}}, cfg.ProviderAllowlists,
+		"model allowlist is a pure policy guardrail and must reach the per-provider config")
 	assert.False(t, cfg.PromptCapture.Enabled,
 		"prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail")
 }
@@ -172,7 +172,7 @@ func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) {
 	require.Len(t, services, 1, "exactly one synth service expected")
 
 	cfg := decodeServiceGuardrailConfig(t, services[0])
-	assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist")
+	assert.Empty(t, cfg.ProviderAllowlists, "no guardrail → provider unrestricted (absent from map)")
 	assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default")
 	assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default")
 }

management/internals/modules/agentnetwork/synthesizer_provider_allowlist_test.go

@@ -0,0 +1,95 @@
+package agentnetwork
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+)
+
+// policyForProviders builds an enabled policy authorising the given providers
+// under the given guardrails (both optional). Groups are irrelevant to
+// buildProviderAllowlists, which keys purely on destination provider.
+func policyForProviders(id string, guardrailIDs []string, providerIDs ...string) *types.Policy {
+	return &types.Policy{
+		ID:                     id,
+		Enabled:                true,
+		DestinationProviderIDs: providerIDs,
+		GuardrailIDs:           guardrailIDs,
+	}
+}
+
+func TestBuildProviderAllowlists(t *testing.T) {
+	byID := map[string]*types.Guardrail{
+		"g-4o":       allowlistGuardrail("g-4o", "acc-1", "gpt-4o"),
+		"g-opus":     allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"),
+		"g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}},
+	}
+
+	t.Run("all authorising policies restrict yields per-provider union", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-4o"}, "prov-x"),
+			policyForProviders("p2", []string{"g-opus"}, "prov-x"),
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got,
+			"a provider every policy restricts carries the sorted union of their models")
+	})
+
+	t.Run("any un-guardrailed policy leaves the provider unrestricted (omitted)", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-4o"}, "prov-x"),
+			policyForProviders("p2", nil, "prov-x"), // no guardrail
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.NotContains(t, got, "prov-x",
+			"a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted")
+	})
+
+	t.Run("a disabled allowlist counts as unrestricted", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-disabled"}, "prov-x"),
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.NotContains(t, got, "prov-x",
+			"a policy whose only guardrail has a disabled allowlist is unrestricted")
+	})
+
+	t.Run("providers are isolated from one another", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-4o"}, "prov-x"),
+			policyForProviders("p2", []string{"g-opus"}, "prov-y"),
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model")
+		assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model")
+	})
+
+	t.Run("one policy authorising two providers restricts both", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"),
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.Equal(t, []string{"gpt-4o"}, got["prov-x"])
+		assert.Equal(t, []string{"gpt-4o"}, got["prov-y"])
+	})
+
+	t.Run("union across a single policy's guardrails", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-4o", "g-opus"}, "prov-x"),
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.ElementsMatch(t, []string{"claude-opus-4", "gpt-4o"}, got["prov-x"],
+			"a policy's own multiple allowlist guardrails union together")
+	})
+
+	t.Run("an enabled allowlist with no models denies everything", func(t *testing.T) {
+		empty := map[string]*types.Guardrail{"g-empty": allowlistGuardrail("g-empty", "acc-1")}
+		got := buildProviderAllowlists([]*types.Policy{
+			policyForProviders("p1", []string{"g-empty"}, "prov-x"),
+		}, empty)
+		assert.Equal(t, map[string][]string{"prov-x": {}}, got,
+			"an enabled-but-empty allowlist is restricted with an empty set, not unrestricted")
+	})
+}

management/internals/modules/agentnetwork/synthesizer_test.go

@@ -1031,8 +1031,12 @@ func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *t
 
 	var cfg guardrailConfig
 	require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly")
-	assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist,
-		"model allowlist union must keep both models")
+	// Both policies restrict the same provider, so the per-provider backstop
+	// carries the union of their models — a coarse gate that management's
+	// per-policy/group check narrows; it only blocks models outside the union
+	// when management is down.
+	assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ProviderAllowlists["prov-1"],
+		"per-provider allowlist union must keep both models")
 }
 
 func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) {

management/internals/shared/grpc/proxy.go

@@ -285,6 +285,7 @@ func (s *ProxyServiceServer) CheckLLMPolicyLimits(ctx context.Context, req *prot
 		UserID:     req.GetUserId(),
 		GroupIDs:   req.GetGroupIds(),
 		ProviderID: req.GetProviderId(),
+		Model:      req.GetModel(),
 	})
 	if err != nil {
 		log.WithContext(ctx).Errorf("select policy for request: %v", err)

management/internals/shared/grpc/proxy_llm_policy_limits_test.go

@@ -0,0 +1,138 @@
+package grpc
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	grpcstatus "google.golang.org/grpc/status"
+
+	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// fakeAgentNetworkLimits records the PolicySelectionInput it was invoked with
+// and returns a pre-programmed result, so tests can assert what the handler
+// forwards to the selector.
+type fakeAgentNetworkLimits struct {
+	gotInput agentnetwork.PolicySelectionInput
+	result   *agentnetwork.PolicySelectionResult
+	err      error
+}
+
+func (f *fakeAgentNetworkLimits) SelectPolicyForRequest(_ context.Context, in agentnetwork.PolicySelectionInput) (*agentnetwork.PolicySelectionResult, error) {
+	f.gotInput = in
+	if f.err != nil {
+		return nil, f.err
+	}
+	return f.result, nil
+}
+
+func (f *fakeAgentNetworkLimits) RecordUsage(_ context.Context, _ agentnetwork.RecordUsageInput) error {
+	return nil
+}
+
+// TestCheckLLMPolicyLimits_ForwardsModelToSelector proves the wiring added here:
+// the model the proxy extracted must reach the selector's Model unchanged,
+// alongside the account/user/group/provider fields.
+func TestCheckLLMPolicyLimits_ForwardsModelToSelector(t *testing.T) {
+	fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true, SelectedPolicyID: "pol-1"}}
+	s := &ProxyServiceServer{}
+	s.SetAgentNetworkLimitsService(fake)
+
+	req := &proto.CheckLLMPolicyLimitsRequest{
+		AccountId:  "acc-1",
+		UserId:     "user-1",
+		GroupIds:   []string{"grp-a", "grp-b"},
+		ProviderId: "prov-1",
+		Model:      "claude-opus-4",
+	}
+
+	resp, err := s.CheckLLMPolicyLimits(context.Background(), req)
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+
+	assert.Equal(t, "acc-1", fake.gotInput.AccountID)
+	assert.Equal(t, "user-1", fake.gotInput.UserID)
+	assert.Equal(t, []string{"grp-a", "grp-b"}, fake.gotInput.GroupIDs)
+	assert.Equal(t, "prov-1", fake.gotInput.ProviderID)
+	assert.Equal(t, "claude-opus-4", fake.gotInput.Model,
+		"the request's model must be forwarded to the selector")
+}
+
+// TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty proves an undetermined
+// model (empty string) is forwarded as-is; the selector decides how to treat it.
+func TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty(t *testing.T) {
+	fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true}}
+	s := &ProxyServiceServer{}
+	s.SetAgentNetworkLimitsService(fake)
+
+	req := &proto.CheckLLMPolicyLimitsRequest{
+		AccountId:  "acc-1",
+		ProviderId: "prov-1",
+	}
+
+	_, err := s.CheckLLMPolicyLimits(context.Background(), req)
+	require.NoError(t, err)
+	assert.Equal(t, "", fake.gotInput.Model, "an absent model must be forwarded as empty, not fabricated")
+}
+
+// TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode proves the deny
+// envelope surfaces the model-allowlist deny code + reason through the response.
+func TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode(t *testing.T) {
+	fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{
+		Allow:      false,
+		DenyCode:   "llm_policy.model_blocked",
+		DenyReason: `model "claude-opus-4" is not permitted by any applicable policy allowlist`,
+	}}
+	s := &ProxyServiceServer{}
+	s.SetAgentNetworkLimitsService(fake)
+
+	resp, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
+		AccountId:  "acc-1",
+		ProviderId: "prov-1",
+		Model:      "claude-opus-4",
+	})
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+	assert.Equal(t, "deny", resp.Decision)
+	assert.Equal(t, "llm_policy.model_blocked", resp.DenyCode)
+	assert.NotEmpty(t, resp.DenyReason)
+	assert.Empty(t, resp.SelectedPolicyId, "a denied request must carry no selected policy")
+}
+
+// TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal proves a selector
+// failure surfaces as an Internal gRPC error rather than a silent allow.
+func TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal(t *testing.T) {
+	fake := &fakeAgentNetworkLimits{err: errors.New("boom")}
+	s := &ProxyServiceServer{}
+	s.SetAgentNetworkLimitsService(fake)
+
+	_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
+		AccountId:  "acc-1",
+		ProviderId: "prov-1",
+		Model:      "gpt-4o",
+	})
+	require.Error(t, err)
+	st, ok := grpcstatus.FromError(err)
+	require.True(t, ok)
+	assert.Equal(t, codes.Internal, st.Code(), "selector errors must never fail open on the hot path")
+}
+
+// TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented locks the
+// fallback: with no limits service wired the RPC returns Unimplemented.
+func TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented(t *testing.T) {
+	s := &ProxyServiceServer{}
+
+	_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
+		AccountId:  "acc-1",
+		ProviderId: "prov-1",
+	})
+	require.Error(t, err)
+	st, ok := grpcstatus.FromError(err)
+	require.True(t, ok)
+	assert.Equal(t, codes.Unimplemented, st.Code())
+}

proxy/internal/auth/tunnel_cache.go

@@ -3,20 +3,30 @@ package auth
 import (
 	"context"
 	"net/netip"
+	"os"
+	"strings"
 	"sync"
 	"time"
 
+	log "github.com/sirupsen/logrus"
 	"golang.org/x/sync/singleflight"
 
 	"github.com/netbirdio/netbird/proxy/internal/types"
 	"github.com/netbirdio/netbird/shared/management/proto"
 )
 
-// tunnelCacheTTL caps how long a positive ValidateTunnelPeer result is
-// reused before re-fetching from management. 5 minutes balances freshness
-// against management load on busy mesh networks.
+// tunnelCacheTTL is the default cap on how long a positive ValidateTunnelPeer
+// result is reused before re-fetching from management. 5 minutes balances
+// freshness against management load on busy mesh networks. Override it with
+// envTunnelCacheTTL when an account needs authorization changes to take effect
+// sooner (at the cost of more ValidateTunnelPeer RPCs).
 const tunnelCacheTTL = 300 * time.Second
 
+// envTunnelCacheTTL overrides tunnelCacheTTL. The value is a Go duration string
+// (e.g. "30s", "2m"); an unset, unparseable, or non-positive value keeps the
+// default.
+const envTunnelCacheTTL = "NB_PROXY_TUNNEL_CACHE_TTL"
+
 // tunnelCachePerAccount caps the number of cached identities per account.
 // Bounded eviction avoids memory growth in pathological cases (huge peer
 // roster, brief request bursts) while staying generous for normal use.
@@ -60,16 +70,35 @@ type accountBucket struct {
 	order []tunnelCacheKey
 }
 
-// newTunnelValidationCache constructs a cache with default TTL and bounds.
+// newTunnelValidationCache constructs a cache with the configured TTL
+// (envTunnelCacheTTL override or default) and default bounds.
 func newTunnelValidationCache() *tunnelValidationCache {
 	return &tunnelValidationCache{
 		entries: make(map[types.AccountID]*accountBucket),
-		ttl:     tunnelCacheTTL,
+		ttl:     tunnelCacheTTLFromEnv(),
 		maxSize: tunnelCachePerAccount,
 		now:     time.Now,
 	}
 }
 
+// tunnelCacheTTLFromEnv returns the tunnel-cache TTL, honoring the
+// envTunnelCacheTTL override. The override must be a positive Go duration
+// string (e.g. "30s", "2m"); anything unset, unparseable, or non-positive
+// falls back to tunnelCacheTTL.
+func tunnelCacheTTLFromEnv() time.Duration {
+	raw := strings.TrimSpace(os.Getenv(envTunnelCacheTTL))
+	if raw == "" {
+		return tunnelCacheTTL
+	}
+	d, err := time.ParseDuration(raw)
+	if err != nil || d <= 0 {
+		log.Warnf("ignoring invalid %s=%q (want a positive Go duration like 30s or 2m); using default %s",
+			envTunnelCacheTTL, raw, tunnelCacheTTL)
+		return tunnelCacheTTL
+	}
+	return d
+}
+
 // get returns a cached response for the key, or nil when missing or
 // expired. Expired entries are evicted lazily on read.
 func (c *tunnelValidationCache) get(key tunnelCacheKey) *proto.ValidateTunnelPeerResponse {

proxy/internal/auth/tunnel_cache_test.go

@@ -169,3 +169,32 @@ func TestTunnelCache_BoundedSizeEvictsOldest(t *testing.T) {
 	assert.NotNil(t, cache.get(keys[1]), "second-newest must remain cached")
 	assert.NotNil(t, cache.get(keys[2]), "newest must remain cached")
 }
+
+func TestTunnelCacheTTLFromEnv(t *testing.T) {
+	t.Run("unset uses default", func(t *testing.T) {
+		t.Setenv(envTunnelCacheTTL, "")
+		assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
+	})
+	t.Run("valid duration overrides", func(t *testing.T) {
+		t.Setenv(envTunnelCacheTTL, "45s")
+		assert.Equal(t, 45*time.Second, tunnelCacheTTLFromEnv())
+	})
+	t.Run("whitespace trimmed", func(t *testing.T) {
+		t.Setenv(envTunnelCacheTTL, "  2m ")
+		assert.Equal(t, 2*time.Minute, tunnelCacheTTLFromEnv())
+	})
+	t.Run("unparseable uses default", func(t *testing.T) {
+		t.Setenv(envTunnelCacheTTL, "nonsense")
+		assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
+	})
+	t.Run("non-positive uses default", func(t *testing.T) {
+		t.Setenv(envTunnelCacheTTL, "0s")
+		assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
+		t.Setenv(envTunnelCacheTTL, "-30s")
+		assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
+	})
+	t.Run("constructor honors override", func(t *testing.T) {
+		t.Setenv(envTunnelCacheTTL, "90s")
+		assert.Equal(t, 90*time.Second, newTunnelValidationCache().ttl)
+	})
+}

proxy/internal/middleware/builtin/llm_guardrail/factory.go

@@ -10,11 +10,15 @@ import (
 )
 
 // Config is the JSON-decoded shape accepted by the factory. The
-// runtime path consumes the normalised allowlist; raw config is not
+// runtime path consumes the normalised allowlists; raw config is not
 // retained beyond construction.
 type Config struct {
-	ModelAllowlist []string      `json:"model_allowlist"`
-	PromptCapture  PromptCapture `json:"prompt_capture"`
+	// ProviderAllowlists maps a resolved provider id (KeyLLMResolvedProviderID) to
+	// its model allowlist. A provider present is restricted to those models; one
+	// absent is unrestricted. Kept per-provider so one provider's list can't leak
+	// onto another.
+	ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
+	PromptCapture      PromptCapture       `json:"prompt_capture"`
 }
 
 // PromptCapture toggles the optional prompt capture + redaction step
@@ -54,21 +58,28 @@ func isEmptyJSON(raw []byte) bool {
 	return false
 }
 
-// normaliseConfig lowercases and trims allowlist entries so the runtime
-// match is case-insensitive. Empty entries are dropped.
+// normaliseConfig lowercases and trims allowlist entries for case-insensitive
+// matching; empty entries drop. A provider whose entries all drop keeps an empty
+// (non-nil) list — "deny every model" — distinct from an absent provider
+// (unrestricted).
 func normaliseConfig(cfg Config) Config {
-	if len(cfg.ModelAllowlist) == 0 {
+	if len(cfg.ProviderAllowlists) == 0 {
+		cfg.ProviderAllowlists = nil
 		return cfg
 	}
-	cleaned := make([]string, 0, len(cfg.ModelAllowlist))
-	for _, entry := range cfg.ModelAllowlist {
-		n := normaliseModel(entry)
-		if n == "" {
-			continue
+	cleaned := make(map[string][]string, len(cfg.ProviderAllowlists))
+	for provider, models := range cfg.ProviderAllowlists {
+		list := make([]string, 0, len(models))
+		for _, entry := range models {
+			n := normaliseModel(entry)
+			if n == "" {
+				continue
+			}
+			list = append(list, n)
 		}
-		cleaned = append(cleaned, n)
+		cleaned[provider] = list
 	}
-	cfg.ModelAllowlist = cleaned
+	cfg.ProviderAllowlists = cleaned
 	return cfg
 }
 

proxy/internal/middleware/builtin/llm_guardrail/middleware.go

@@ -83,8 +83,9 @@ func (m *Middleware) MutationsSupported() bool { return false }
 // prompt capture only affects the metadata emitted alongside an allow.
 func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
 	model, modelPresent := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
+	providerID, _ := lookupMetadata(in.Metadata, middleware.KeyLLMResolvedProviderID)
 
-	if denial := m.evaluateAllowlist(model, modelPresent); denial != nil {
+	if denial := m.evaluateAllowlist(providerID, model, modelPresent); denial != nil {
 		return denial, nil
 	}
 
@@ -110,20 +111,32 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
 // is a no-op.
 func (m *Middleware) Close() error { return nil }
 
-// evaluateAllowlist returns a deny Output when the configured allowlist
-// rejects the model. A nil return means the request should proceed.
-func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middleware.Output {
-	if len(m.cfg.ModelAllowlist) == 0 {
+// evaluateAllowlist denies when the resolved provider's allowlist rejects the
+// model; nil means proceed. Scoped to the provider llm_router resolved, so an
+// unrestricted provider (absent from config) is never caught by another's list.
+func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bool) *middleware.Output {
+	if len(m.cfg.ProviderAllowlists) == 0 {
 		return nil
 	}
-	// Fail closed: with an allowlist configured, a request whose model the
-	// upstream parser could not extract (absent or empty) must be denied rather
-	// than allowed. This is what enforces the allowlist for URL/path-routed
-	// providers (Bedrock, Vertex, ...) whose model lives outside the JSON body.
+	// Restrictions exist but the resolved provider is unknown, so we can't tell
+	// if this request targets a restricted provider — fail closed. llm_router
+	// normally stamps the provider first, so this is a defensive guard.
+	if providerID == "" {
+		return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
+	}
+	allowlist, restricted := m.cfg.ProviderAllowlists[providerID]
+	if !restricted {
+		// This provider has no allowlist (some authorising policy left it
+		// unrestricted); management owns any per-policy/group decision.
+		return nil
+	}
+	// Fail closed: with an allowlist in effect for this provider, a request whose
+	// model the parser couldn't extract (absent/empty) is denied. This enforces
+	// the allowlist for path-routed providers (Bedrock, Vertex) with no body model.
 	if !modelPresent || normaliseModel(model) == "" {
 		return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
 	}
-	if m.modelInAllowlist(model) {
+	if modelInAllowlist(allowlist, model) {
 		return nil
 	}
 	return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
@@ -151,14 +164,15 @@ func denyModel(model, code, message, reason string) *middleware.Output {
 	}
 }
 
-// modelInAllowlist reports whether the model matches any allowlist
-// entry under the case-insensitive, trim-tolerant comparison rule.
-func (m *Middleware) modelInAllowlist(model string) bool {
+// modelInAllowlist reports whether the model matches any entry in the supplied
+// (already-normalised) allowlist under the case-insensitive, trim-tolerant
+// comparison rule.
+func modelInAllowlist(allowlist []string, model string) bool {
 	normalised := normaliseModel(model)
 	if normalised == "" {
 		return false
 	}
-	for _, allowed := range m.cfg.ModelAllowlist {
+	for _, allowed := range allowlist {
 		if allowed == normalised {
 			return true
 		}

proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go

@@ -26,6 +26,25 @@ func newInput(meta ...middleware.KV) *middleware.Input {
 	return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: meta}
 }
 
+const (
+	testProvider  = "prov-1"
+	otherProvider = "prov-2"
+)
+
+// providerCfg builds a Config restricting testProvider to the given models.
+func providerCfg(models ...string) Config {
+	return Config{ProviderAllowlists: map[string][]string{testProvider: models}}
+}
+
+// newInputProvider builds an input that carries a resolved provider id (as
+// llm_router would stamp) plus any extra metadata.
+func newInputProvider(provider string, meta ...middleware.KV) *middleware.Input {
+	all := make([]middleware.KV, 0, len(meta)+1)
+	all = append(all, middleware.KV{Key: middleware.KeyLLMResolvedProviderID, Value: provider})
+	all = append(all, meta...)
+	return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: all}
+}
+
 func TestMiddlewareIdentity(t *testing.T) {
 	mw := New(Config{})
 	assert.Equal(t, ID, mw.ID(), "middleware ID must be llm_guardrail")
@@ -47,12 +66,12 @@ func TestMiddlewareIdentity(t *testing.T) {
 
 func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
 	mw := New(Config{})
-	out, err := mw.Invoke(context.Background(), newInput(
+	out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
 		middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
 	))
 	require.NoError(t, err)
 	require.NotNil(t, out)
-	assert.Equal(t, middleware.DecisionAllow, out.Decision, "empty allowlist must allow any model")
+	assert.Equal(t, middleware.DecisionAllow, out.Decision, "no provider allowlists must allow any model")
 	v, ok := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision)
 	require.True(t, ok, "decision metadata must be emitted")
 	assert.Equal(t, "allow", v, "decision must be allow")
@@ -62,8 +81,8 @@ func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
 }
 
 func TestAllowlistMatchAllows(t *testing.T) {
-	mw := New(Config{ModelAllowlist: []string{"gpt-4o", "claude-opus-4"}})
-	out, err := mw.Invoke(context.Background(), newInput(
+	mw := New(providerCfg("gpt-4o", "claude-opus-4"))
+	out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
 		middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
 	))
 	require.NoError(t, err)
@@ -71,8 +90,8 @@ func TestAllowlistMatchAllows(t *testing.T) {
 }
 
 func TestAllowlistMissDenies(t *testing.T) {
-	mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
-	out, err := mw.Invoke(context.Background(), newInput(
+	mw := New(providerCfg("gpt-4o"))
+	out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
 		middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
 	))
 	require.NoError(t, err)
@@ -91,10 +110,10 @@ func TestAllowlistMissDenies(t *testing.T) {
 }
 
 func TestAllowlistCaseInsensitive(t *testing.T) {
-	mw := New(Config{ModelAllowlist: []string{"  GPT-4o  ", "Claude-OPUS-4"}})
+	mw := New(providerCfg("  GPT-4o  ", "Claude-OPUS-4"))
 	cases := []string{"gpt-4o", "GPT-4O", "  claude-opus-4 "}
 	for _, model := range cases {
-		out, err := mw.Invoke(context.Background(), newInput(
+		out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
 			middleware.KV{Key: middleware.KeyLLMModel, Value: model},
 		))
 		require.NoError(t, err)
@@ -103,14 +122,15 @@ func TestAllowlistCaseInsensitive(t *testing.T) {
 }
 
 func TestAllowlistMissingModelKeyDenies(t *testing.T) {
-	// Fail closed: with an allowlist configured, a request whose model the
-	// parser could not extract (URL/path-routed providers such as Bedrock or
-	// Vertex whose shape wasn't recognised) must be denied, not allowed.
-	mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
-	out, err := mw.Invoke(context.Background(), newInput())
+	// Fail closed: with an allowlist in effect for the resolved provider, a
+	// request whose model the parser could not extract (URL/path-routed
+	// providers such as Bedrock or Vertex whose shape wasn't recognised) must be
+	// denied, not allowed.
+	mw := New(providerCfg("gpt-4o"))
+	out, err := mw.Invoke(context.Background(), newInputProvider(testProvider))
 	require.NoError(t, err)
 	require.NotNil(t, out)
-	assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when an allowlist is set")
+	assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when the provider is restricted")
 	assert.Equal(t, 403, out.DenyStatus, "deny status must be 403")
 	require.NotNil(t, out.DenyReason, "deny reason must be populated")
 	assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
@@ -122,26 +142,101 @@ func TestAllowlistMissingModelKeyDenies(t *testing.T) {
 
 func TestAllowlistEmptyModelValueDenies(t *testing.T) {
 	// A present-but-empty model is as undeterminable as an absent one.
-	mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
-	out, err := mw.Invoke(context.Background(), newInput(
+	mw := New(providerCfg("gpt-4o"))
+	out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
 		middleware.KV{Key: middleware.KeyLLMModel, Value: "   "},
 	))
 	require.NoError(t, err)
 	require.NotNil(t, out)
-	assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when an allowlist is set")
+	assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when the provider is restricted")
 	require.NotNil(t, out.DenyReason, "deny reason must be populated")
 	assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
 }
 
 func TestAllowlistEmptyListAllowsMissingModel(t *testing.T) {
-	// Without an allowlist there is nothing to enforce, so a missing model is
-	// still allowed — the fail-closed rule only applies when a list is set.
+	// Without any provider allowlists there is nothing to enforce, so a missing
+	// model is still allowed — the fail-closed rule only applies when a
+	// restriction is in effect.
 	mw := New(Config{})
 	out, err := mw.Invoke(context.Background(), newInput())
 	require.NoError(t, err)
 	assert.Equal(t, middleware.DecisionAllow, out.Decision, "no allowlist must allow even without a model")
 }
 
+func TestUnrestrictedProviderAllowsAnyModel(t *testing.T) {
+	// The request resolved to otherProvider, which has no allowlist, so its
+	// traffic must not be caught by testProvider's restriction — the
+	// cross-provider-leak / false-deny guard.
+	mw := New(providerCfg("gpt-4o"))
+	out, err := mw.Invoke(context.Background(), newInputProvider(otherProvider,
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
+	))
+	require.NoError(t, err)
+	assert.Equal(t, middleware.DecisionAllow, out.Decision, "an unrestricted provider must not inherit another provider's allowlist")
+}
+
+func TestPerProviderAllowlistsAreIsolated(t *testing.T) {
+	// gpt-4o is allowed only on testProvider; claude-opus-4 only on
+	// otherProvider. A model allowlisted for one provider must not be usable on
+	// the other — the fail-closed layer never unions allowlists across providers.
+	mw := New(Config{ProviderAllowlists: map[string][]string{
+		testProvider:  {"gpt-4o"},
+		otherProvider: {"claude-opus-4"},
+	}})
+	out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
+	))
+	require.NoError(t, err)
+	assert.Equal(t, middleware.DecisionDeny, out.Decision, "claude-opus-4 is allowed only on otherProvider, not testProvider")
+	require.NotNil(t, out.DenyReason)
+	assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "cross-provider model must be blocked, not model_unknown")
+}
+
+func TestRestrictionsButNoResolvedProviderFailsClosed(t *testing.T) {
+	// Restrictions exist for the account but the resolved provider id is absent,
+	// so the request cannot be scoped to a provider. Fail closed rather than
+	// wave it through.
+	mw := New(providerCfg("gpt-4o"))
+	out, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+	))
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	assert.Equal(t, middleware.DecisionDeny, out.Decision, "missing resolved provider must fail closed when restrictions exist")
+	require.NotNil(t, out.DenyReason)
+	assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
+}
+
+func TestEnabledButEmptyAllowlistDeniesEveryModel(t *testing.T) {
+	// An allowlist-enabled provider with zero models is distinct from an
+	// unrestricted (absent) provider: it must deny every model.
+	mw := New(providerCfg())
+	out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+	))
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	assert.Equal(t, middleware.DecisionDeny, out.Decision, "an enabled-but-empty allowlist must deny every model")
+	require.NotNil(t, out.DenyReason)
+	assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked, not model_unknown")
+}
+
+func TestFactoryAllEmptyEntriesDenyEveryModel(t *testing.T) {
+	// All the provider's entries are blank; they collapse to a non-nil empty
+	// list (deny everything for that provider), not "no restriction".
+	raw := []byte(`{"provider_allowlists":{"prov-1":["","   "]}}`)
+	mw, err := Factory{}.New(raw)
+	require.NoError(t, err)
+	out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+	))
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	assert.Equal(t, middleware.DecisionDeny, out.Decision, "all-blank allowlist entries must still restrict the provider")
+	require.NotNil(t, out.DenyReason)
+	assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked")
+}
+
 func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) {
 	mw := New(Config{})
 	out, err := mw.Invoke(context.Background(), newInput(
@@ -217,8 +312,8 @@ func TestFactoryAcceptsZeroConfigs(t *testing.T) {
 
 func TestFactoryDecodesValidConfig(t *testing.T) {
 	cfg := Config{
-		ModelAllowlist: []string{"gpt-4o"},
-		PromptCapture:  PromptCapture{Enabled: true, RedactPii: true},
+		ProviderAllowlists: map[string][]string{testProvider: {"gpt-4o"}},
+		PromptCapture:      PromptCapture{Enabled: true, RedactPii: true},
 	}
 	raw, err := json.Marshal(cfg)
 	require.NoError(t, err, "marshalling test config must succeed")
@@ -234,15 +329,15 @@ func TestFactoryRejectsMalformedJSON(t *testing.T) {
 }
 
 func TestFactoryNormalisesAllowlist(t *testing.T) {
-	raw := []byte(`{"model_allowlist":["  GPT-4o  ","",""," Claude-3 "]}`)
+	raw := []byte(`{"provider_allowlists":{"prov-1":["  GPT-4o  ","",""," Claude-3 "]}}`)
 	mw, err := Factory{}.New(raw)
 	require.NoError(t, err)
-	out, err := mw.Invoke(context.Background(), newInput(
+	out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
 		middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
 	))
 	require.NoError(t, err)
 	assert.Equal(t, middleware.DecisionAllow, out.Decision, "factory must lowercase + trim allowlist entries")
-	out2, err := mw.Invoke(context.Background(), newInput(
+	out2, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
 		middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-3"},
 	))
 	require.NoError(t, err)

proxy/internal/middleware/builtin/llm_limit_check/middleware.go

@@ -175,7 +175,7 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou
 		DenyStatus: 403,
 		DenyReason: &middleware.DenyReason{
 			Code:    code,
-			Message: "LLM policy limit exceeded",
+			Message: denyMessageForCode(code),
 		},
 		Metadata: []middleware.KV{
 			{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
@@ -184,6 +184,21 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou
 	}
 }
 
+// denyMessageForCode maps a management deny code to a public message.
+// Model-allowlist rejections get a model-specific message matching the
+// local guardrail; everything else keeps the generic quota wording. The
+// message stays generic so it never leaks internal quota detail.
+func denyMessageForCode(code string) string {
+	switch code {
+	case "llm_policy.model_blocked":
+		return "model is not in the policy allowlist"
+	case "llm_policy.model_unknown":
+		return "request model could not be determined for the policy allowlist"
+	default:
+		return "LLM policy limit exceeded"
+	}
+}
+
 // lookupKV returns the value associated with key, or the empty
 // string when absent.
 func lookupKV(kvs []middleware.KV, key string) string {

proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go

@@ -115,6 +115,46 @@ func TestInvoke_DenyConvertsToProxyDeny(t *testing.T) {
 	assert.NotContains(t, out.DenyReason.Message, "1000", "internal cap numbers must not reach the caller")
 }
 
+// TestInvoke_ModelDenyMessages proves a model-allowlist rejection gets a
+// model-specific public message rather than the generic quota wording, so a
+// blocked or undetermined model reads consistently with the local guardrail.
+func TestInvoke_ModelDenyMessages(t *testing.T) {
+	cases := []struct {
+		name    string
+		code    string
+		message string
+	}{
+		{"blocked", "llm_policy.model_blocked", "model is not in the policy allowlist"},
+		{"unknown", "llm_policy.model_unknown", "request model could not be determined for the policy allowlist"},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			mgmt := &fakeMgmt{
+				checkResp: &proto.CheckLLMPolicyLimitsResponse{
+					Decision: "deny",
+					DenyCode: tc.code,
+				},
+			}
+			m := New(mgmt, nil)
+
+			out := runInvoke(t, m, &middleware.Input{
+				AccountID:  "acc-1",
+				UserGroups: []string{"grp-engineers"},
+				Metadata: []middleware.KV{
+					{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"},
+					{Key: middleware.KeyLLMModel, Value: "some-model"},
+				},
+			})
+
+			assert.Equal(t, middleware.DecisionDeny, out.Decision)
+			require.NotNil(t, out.DenyReason, "deny envelope must carry a reason payload")
+			assert.Equal(t, tc.code, out.DenyReason.Code, "canonical deny code surfaces to the caller")
+			assert.Equal(t, tc.message, out.DenyReason.Message,
+				"model denials must use a model-specific message, matching the local guardrail")
+		})
+	}
+}
+
 // TestInvoke_NoMgmtClientPassesThrough proves the partial-wiring
 // safety: a middleware constructed without a management client
 // allows every request without attribution. This makes a half-set-up

proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go

@@ -25,10 +25,17 @@ func runParserGuardrail(t *testing.T, url string, body []byte, allowlist []strin
 	})
 	require.NoError(t, err, "parser must not error")
 
-	guard := llm_guardrail.New(llm_guardrail.Config{ModelAllowlist: allowlist})
+	const providerID = "prov-under-test"
+	guard := llm_guardrail.New(llm_guardrail.Config{
+		ProviderAllowlists: map[string][]string{providerID: allowlist},
+	})
+	// The real chain has llm_router stamp the resolved provider id before the
+	// guardrail runs; the parser doesn't, so add it here so the guardrail can
+	// scope the allowlist to this provider.
+	meta := append([]middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: providerID}}, parsed.Metadata...)
 	out, err := guard.Invoke(context.Background(), &middleware.Input{
 		Slot:     middleware.SlotOnRequest,
-		Metadata: parsed.Metadata,
+		Metadata: meta,
 	})
 	require.NoError(t, err, "guardrail must not error")
 	require.NotNil(t, out, "guardrail must return an output")

shared/signal/client/client.go

@@ -52,11 +52,6 @@ type CredentialPayload struct {
 	Credential      *Credential
 	RosenpassPubKey []byte
 	RosenpassAddr   string
-	// MlkemPayload is the opaque post-quantum KEM handshake message riding this
-	// OFFER/ANSWER (see Body.mlkemPayload). Nil when not running the PQ exchange.
-	MlkemPayload []byte
-	// MlkemPort is the sender's ML-KEM PQ service UDP port (0 when not running).
-	MlkemPort       int
 	RelaySrvAddress string
 	RelaySrvIP      netip.Addr
 	SessionID       []byte
@@ -94,13 +89,6 @@ func MarshalCredential(myKey wgtypes.Key, remoteKey string, p CredentialPayload)
 	if p.RelaySrvIP.IsValid() {
 		body.RelayServerIP = p.RelaySrvIP.Unmap().AsSlice()
 	}
-	if len(p.MlkemPayload) > 0 {
-		body.MlkemPayload = p.MlkemPayload
-	}
-	if p.MlkemPort > 0 {
-		port := uint32(p.MlkemPort)
-		body.MlkemPort = &port
-	}
 	return &proto.Message{
 		Key:       myKey.PublicKey().String(),
 		RemoteKey: remoteKey,

shared/signal/proto/signalexchange.pb.go

@@ -239,16 +239,6 @@ type Body struct {
 	// fallback dial target when DNS resolution of relayServerAddress fails.
 	// SNI/TLS verification still uses relayServerAddress.
 	RelayServerIP []byte `protobuf:"bytes,11,opt,name=relayServerIP,proto3,oneof" json:"relayServerIP,omitempty"`
-	// mlkemPayload carries a post-quantum X25519MLKEM768 handshake message that
-	// seeds the WireGuard PSK, riding this Body's OFFER/ANSWER: on an OFFER it is
-	// the KEM offer, on an ANSWER the KEM answer. It is opaque to signal — the
-	// pqkem library frames and parses it. Absent when the sender does not run the
-	// ML-KEM PQ exchange; unknown to older clients, which ignore it.
-	MlkemPayload []byte `protobuf:"bytes,12,opt,name=mlkemPayload,proto3,oneof" json:"mlkemPayload,omitempty"`
-	// mlkemPort is the UDP port of the sender's ML-KEM PQ service, bound on its
-	// WireGuard overlay IP. Peers send subsequent rekey messages there over the
-	// data path. Zero/absent when the ML-KEM PQ exchange is not running.
-	MlkemPort *uint32 `protobuf:"varint,13,opt,name=mlkemPort,proto3,oneof" json:"mlkemPort,omitempty"`
 }
 
 func (x *Body) Reset() {
@@ -353,20 +343,6 @@ func (x *Body) GetRelayServerIP() []byte {
 	return nil
 }
 
-func (x *Body) GetMlkemPayload() []byte {
-	if x != nil {
-		return x.MlkemPayload
-	}
-	return nil
-}
-
-func (x *Body) GetMlkemPort() uint32 {
-	if x != nil && x.MlkemPort != nil {
-		return *x.MlkemPort
-	}
-	return 0
-}
-
 // Mode indicates a connection mode
 type Mode struct {
 	state         protoimpl.MessageState
@@ -490,7 +466,7 @@ var file_signalexchange_proto_rawDesc = []byte{
 	0x52, 0x09, 0x72, 0x65, 0x6d, 0x6f, 0x74, 0x65, 0x4b, 0x65, 0x79, 0x12, 0x28, 0x0a, 0x04, 0x62,
 	0x6f, 0x64, 0x79, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x73, 0x69, 0x67, 0x6e,
 	0x61, 0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x42, 0x6f, 0x64, 0x79, 0x52,
-	0x04, 0x62, 0x6f, 0x64, 0x79, 0x22, 0xbd, 0x05, 0x0a, 0x04, 0x42, 0x6f, 0x64, 0x79, 0x12, 0x2d,
+	0x04, 0x62, 0x6f, 0x64, 0x79, 0x22, 0xd2, 0x04, 0x0a, 0x04, 0x42, 0x6f, 0x64, 0x79, 0x12, 0x2d,
 	0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x19, 0x2e, 0x73,
 	0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x42, 0x6f,
 	0x64, 0x79, 0x2e, 0x54, 0x79, 0x70, 0x65, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x12, 0x18, 0x0a,
@@ -518,46 +494,39 @@ var file_signalexchange_proto_rawDesc = []byte{
 	0x52, 0x09, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x88, 0x01, 0x01, 0x12, 0x29,
 	0x0a, 0x0d, 0x72, 0x65, 0x6c, 0x61, 0x79, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x49, 0x50, 0x18,
 	0x0b, 0x20, 0x01, 0x28, 0x0c, 0x48, 0x02, 0x52, 0x0d, 0x72, 0x65, 0x6c, 0x61, 0x79, 0x53, 0x65,
-	0x72, 0x76, 0x65, 0x72, 0x49, 0x50, 0x88, 0x01, 0x01, 0x12, 0x27, 0x0a, 0x0c, 0x6d, 0x6c, 0x6b,
-	0x65, 0x6d, 0x50, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64, 0x18, 0x0c, 0x20, 0x01, 0x28, 0x0c, 0x48,
-	0x03, 0x52, 0x0c, 0x6d, 0x6c, 0x6b, 0x65, 0x6d, 0x50, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64, 0x88,
-	0x01, 0x01, 0x12, 0x21, 0x0a, 0x09, 0x6d, 0x6c, 0x6b, 0x65, 0x6d, 0x50, 0x6f, 0x72, 0x74, 0x18,
-	0x0d, 0x20, 0x01, 0x28, 0x0d, 0x48, 0x04, 0x52, 0x09, 0x6d, 0x6c, 0x6b, 0x65, 0x6d, 0x50, 0x6f,
-	0x72, 0x74, 0x88, 0x01, 0x01, 0x22, 0x52, 0x0a, 0x04, 0x54, 0x79, 0x70, 0x65, 0x12, 0x09, 0x0a,
-	0x05, 0x4f, 0x46, 0x46, 0x45, 0x52, 0x10, 0x00, 0x12, 0x0a, 0x0a, 0x06, 0x41, 0x4e, 0x53, 0x57,
-	0x45, 0x52, 0x10, 0x01, 0x12, 0x0d, 0x0a, 0x09, 0x43, 0x41, 0x4e, 0x44, 0x49, 0x44, 0x41, 0x54,
-	0x45, 0x10, 0x02, 0x12, 0x08, 0x0a, 0x04, 0x4d, 0x4f, 0x44, 0x45, 0x10, 0x04, 0x12, 0x0b, 0x0a,
-	0x07, 0x47, 0x4f, 0x5f, 0x49, 0x44, 0x4c, 0x45, 0x10, 0x05, 0x12, 0x0d, 0x0a, 0x09, 0x48, 0x45,
-	0x41, 0x52, 0x54, 0x42, 0x45, 0x41, 0x54, 0x10, 0x06, 0x42, 0x15, 0x0a, 0x13, 0x5f, 0x72, 0x65,
-	0x6c, 0x61, 0x79, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x64, 0x64, 0x72, 0x65, 0x73, 0x73,
-	0x42, 0x0c, 0x0a, 0x0a, 0x5f, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x42, 0x10,
-	0x0a, 0x0e, 0x5f, 0x72, 0x65, 0x6c, 0x61, 0x79, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x49, 0x50,
-	0x42, 0x0f, 0x0a, 0x0d, 0x5f, 0x6d, 0x6c, 0x6b, 0x65, 0x6d, 0x50, 0x61, 0x79, 0x6c, 0x6f, 0x61,
-	0x64, 0x42, 0x0c, 0x0a, 0x0a, 0x5f, 0x6d, 0x6c, 0x6b, 0x65, 0x6d, 0x50, 0x6f, 0x72, 0x74, 0x4a,
-	0x04, 0x08, 0x09, 0x10, 0x0a, 0x22, 0x2e, 0x0a, 0x04, 0x4d, 0x6f, 0x64, 0x65, 0x12, 0x1b, 0x0a,
-	0x06, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x48, 0x00, 0x52,
-	0x06, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x88, 0x01, 0x01, 0x42, 0x09, 0x0a, 0x07, 0x5f, 0x64,
-	0x69, 0x72, 0x65, 0x63, 0x74, 0x22, 0x6d, 0x0a, 0x0f, 0x52, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61,
-	0x73, 0x73, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x28, 0x0a, 0x0f, 0x72, 0x6f, 0x73, 0x65,
-	0x6e, 0x70, 0x61, 0x73, 0x73, 0x50, 0x75, 0x62, 0x4b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28,
-	0x0c, 0x52, 0x0f, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x50, 0x75, 0x62, 0x4b,
-	0x65, 0x79, 0x12, 0x30, 0x0a, 0x13, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x53,
-	0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x64, 0x64, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52,
-	0x13, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72,
-	0x41, 0x64, 0x64, 0x72, 0x32, 0xb9, 0x01, 0x0a, 0x0e, 0x53, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x45,
-	0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x12, 0x4c, 0x0a, 0x04, 0x53, 0x65, 0x6e, 0x64, 0x12,
-	0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65,
-	0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67,
-	0x65, 0x1a, 0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e,
-	0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73,
-	0x61, 0x67, 0x65, 0x22, 0x00, 0x12, 0x59, 0x0a, 0x0d, 0x43, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74,
-	0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x12, 0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65,
+	0x72, 0x76, 0x65, 0x72, 0x49, 0x50, 0x88, 0x01, 0x01, 0x22, 0x52, 0x0a, 0x04, 0x54, 0x79, 0x70,
+	0x65, 0x12, 0x09, 0x0a, 0x05, 0x4f, 0x46, 0x46, 0x45, 0x52, 0x10, 0x00, 0x12, 0x0a, 0x0a, 0x06,
+	0x41, 0x4e, 0x53, 0x57, 0x45, 0x52, 0x10, 0x01, 0x12, 0x0d, 0x0a, 0x09, 0x43, 0x41, 0x4e, 0x44,
+	0x49, 0x44, 0x41, 0x54, 0x45, 0x10, 0x02, 0x12, 0x08, 0x0a, 0x04, 0x4d, 0x4f, 0x44, 0x45, 0x10,
+	0x04, 0x12, 0x0b, 0x0a, 0x07, 0x47, 0x4f, 0x5f, 0x49, 0x44, 0x4c, 0x45, 0x10, 0x05, 0x12, 0x0d,
+	0x0a, 0x09, 0x48, 0x45, 0x41, 0x52, 0x54, 0x42, 0x45, 0x41, 0x54, 0x10, 0x06, 0x42, 0x15, 0x0a,
+	0x13, 0x5f, 0x72, 0x65, 0x6c, 0x61, 0x79, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x64, 0x64,
+	0x72, 0x65, 0x73, 0x73, 0x42, 0x0c, 0x0a, 0x0a, 0x5f, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e,
+	0x49, 0x64, 0x42, 0x10, 0x0a, 0x0e, 0x5f, 0x72, 0x65, 0x6c, 0x61, 0x79, 0x53, 0x65, 0x72, 0x76,
+	0x65, 0x72, 0x49, 0x50, 0x4a, 0x04, 0x08, 0x09, 0x10, 0x0a, 0x22, 0x2e, 0x0a, 0x04, 0x4d, 0x6f,
+	0x64, 0x65, 0x12, 0x1b, 0x0a, 0x06, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x18, 0x01, 0x20, 0x01,
+	0x28, 0x08, 0x48, 0x00, 0x52, 0x06, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x88, 0x01, 0x01, 0x42,
+	0x09, 0x0a, 0x07, 0x5f, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x22, 0x6d, 0x0a, 0x0f, 0x52, 0x6f,
+	0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x28, 0x0a,
+	0x0f, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x50, 0x75, 0x62, 0x4b, 0x65, 0x79,
+	0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x0f, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73,
+	0x73, 0x50, 0x75, 0x62, 0x4b, 0x65, 0x79, 0x12, 0x30, 0x0a, 0x13, 0x72, 0x6f, 0x73, 0x65, 0x6e,
+	0x70, 0x61, 0x73, 0x73, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x64, 0x64, 0x72, 0x18, 0x02,
+	0x20, 0x01, 0x28, 0x09, 0x52, 0x13, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x53,
+	0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x64, 0x64, 0x72, 0x32, 0xb9, 0x01, 0x0a, 0x0e, 0x53, 0x69,
+	0x67, 0x6e, 0x61, 0x6c, 0x45, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x12, 0x4c, 0x0a, 0x04,
+	0x53, 0x65, 0x6e, 0x64, 0x12, 0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63,
+	0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d,
+	0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65,
 	0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65,
-	0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61,
-	0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70,
-	0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x28, 0x01, 0x30, 0x01,
-	0x42, 0x08, 0x5a, 0x06, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74,
-	0x6f, 0x33,
+	0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x12, 0x59, 0x0a, 0x0d, 0x43, 0x6f,
+	0x6e, 0x6e, 0x65, 0x63, 0x74, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x12, 0x20, 0x2e, 0x73, 0x69,
+	0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63,
+	0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x20, 0x2e,
+	0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45,
+	0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22,
+	0x00, 0x28, 0x01, 0x30, 0x01, 0x42, 0x08, 0x5a, 0x06, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62,
+	0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
 }
 
 var (

shared/signal/proto/signalexchange.proto

@@ -75,18 +75,6 @@ message Body {
   // fallback dial target when DNS resolution of relayServerAddress fails.
   // SNI/TLS verification still uses relayServerAddress.
   optional bytes relayServerIP = 11;
-
-  // mlkemPayload carries a post-quantum X25519MLKEM768 handshake message that
-  // seeds the WireGuard PSK, riding this Body's OFFER/ANSWER: on an OFFER it is
-  // the KEM offer, on an ANSWER the KEM answer. It is opaque to signal — the
-  // pqkem library frames and parses it. Absent when the sender does not run the
-  // ML-KEM PQ exchange; unknown to older clients, which ignore it.
-  optional bytes mlkemPayload = 12;
-
-  // mlkemPort is the UDP port of the sender's ML-KEM PQ service, bound on its
-  // WireGuard overlay IP. Peers send subsequent rekey messages there over the
-  // data path. Zero/absent when the ML-KEM PQ exchange is not running.
-  optional uint32 mlkemPort = 13;
 }
 
 // Mode indicates a connection mode