Compare commits

..

5 Commits

30 changed files with 1113 additions and 1528 deletions

View File

@@ -17,9 +17,7 @@ import (
"github.com/netbirdio/netbird/client/internal"
"github.com/netbirdio/netbird/client/internal/auth"
"github.com/netbirdio/netbird/client/internal/profilemanager"
nbnet "github.com/netbirdio/netbird/client/net"
"github.com/netbirdio/netbird/client/proto"
"github.com/netbirdio/netbird/client/server"
"github.com/netbirdio/netbird/client/system"
"github.com/netbirdio/netbird/util"
)
@@ -333,14 +331,6 @@ func doForegroundLogin(ctx context.Context, cmd *cobra.Command, setupKey string,
return fmt.Errorf("read config file %s: %v", configFilePath, err)
}
// Mirror runInForegroundMode: recover residual state (DNS, firewall,
// ssh config, legacy routing) from a previous unclean shutdown and
// enable advanced routing before dialing management.
if err := server.RestoreResidualState(ctx, profilemanager.NewServiceManager(configFilePath).GetStatePath()); err != nil {
log.Warnf("failed to restore residual state: %v", err)
}
nbnet.Init()
err = foregroundLogin(ctx, cmd, config, setupKey, activeProf.ID)
if err != nil {
return fmt.Errorf("foreground login failed: %v", err)

View File

@@ -22,8 +22,6 @@ import (
"github.com/netbirdio/netbird/client/internal/peer"
"github.com/netbirdio/netbird/client/internal/profilemanager"
"github.com/netbirdio/netbird/client/proto"
nbnet "github.com/netbirdio/netbird/client/net"
"github.com/netbirdio/netbird/client/server"
"github.com/netbirdio/netbird/client/system"
"github.com/netbirdio/netbird/shared/management/domain"
"github.com/netbirdio/netbird/util"
@@ -231,24 +229,6 @@ func runInForegroundMode(ctx context.Context, cmd *cobra.Command, activeProf *pr
_, _ = profilemanager.UpdateOldManagementURL(ctx, config, configFilePath)
// Restore residual state left by a previous run that did not shut down
// cleanly, mirroring what the daemon does before connecting: it recovers
// DNS config (a stale resolv.conf takeover can make the management
// hostname unresolvable), firewall rules, ssh config and legacy routing.
// Route cleanup itself happens at engine start; nbnet.Init() below lets
// the management dial bypass a leftover fwmark rule until then.
// Foreground mode is particularly exposed in containers: a crashed
// container restarts inside the same (pod) network namespace, so stale
// state survives while the process does not.
if err := server.RestoreResidualState(ctx, profilemanager.NewServiceManager(configPath).GetStatePath()); err != nil {
log.Warnf("failed to restore residual state: %v", err)
}
// Enable advanced routing (as the daemon does on startup) so the
// management dial bypasses a leftover fwmark rule instead of being
// shunted into a stale routing table.
nbnet.Init()
err = foregroundLogin(ctx, cmd, config, providedSetupKey, activeProf.ID)
if err != nil {
return fmt.Errorf("foreground login failed: %v", err)

View File

@@ -2,6 +2,7 @@ package internal
import (
"context"
"maps"
"os"
"strconv"
"sync"
@@ -14,6 +15,7 @@ import (
"github.com/netbirdio/netbird/client/internal/peer"
"github.com/netbirdio/netbird/client/internal/peerstore"
"github.com/netbirdio/netbird/route"
mgmProto "github.com/netbirdio/netbird/shared/management/proto"
)
// lazyForce is the resolved local decision for lazy connections, layered above the
@@ -40,8 +42,14 @@ type ConnMgr struct {
iface lazyconn.WGIface
force lazyForce
rosenpassEnabled bool
// remoteLazyEnabled caches the account-wide lazy feature flag from management.
// It is the default for peers that do not carry a per-peer lazy hint.
remoteLazyEnabled bool
lazyConnMgr *manager.Manager
// appliedExcludeList is the exclude set last handed to the lazy manager, kept so an
// unchanged set on the next sync skips the O(n) reconciliation.
appliedExcludeList map[string]bool
wg sync.WaitGroup
lazyCtx context.Context
@@ -59,69 +67,56 @@ func NewConnMgr(engineConfig *EngineConfig, statusRecorder *peer.Status, peerSto
return e
}
// Start initializes the connection manager. It starts the lazy connection manager when a
// local override forces it on; with no local override it waits for the management feature flag.
// Start initializes the connection manager. The lazy connection manager always runs so that
// per-peer lazy defaults (e.g. proxy peers) work even when the account flag is off; the
// account flag and the local override decide the default lazy state per peer (see
// PeerLazyDefault). Rosenpass is the only condition that disables it.
func (e *ConnMgr) Start(ctx context.Context) {
if e.lazyConnMgr != nil {
log.Errorf("lazy connection manager is already started")
return
}
switch e.force {
case lazyForceOff:
log.Infof("lazy connection manager is disabled by local override (%s or MDM policy)", lazyconn.EnvLazyConn)
e.statusRecorder.UpdateLazyConnection(false)
return
case lazyForceNone:
log.Infof("lazy connection manager is managed by the management feature flag")
e.statusRecorder.UpdateLazyConnection(false)
return
}
if e.rosenpassEnabled {
log.Warnf("rosenpass connection manager is enabled, lazy connection manager will not be started")
log.Warnf("rosenpass is enabled, lazy connection manager will not be started")
e.statusRecorder.UpdateLazyConnection(false)
return
}
e.initLazyManager(ctx)
e.statusRecorder.UpdateLazyConnection(true)
e.statusRecorder.UpdateLazyConnection(e.PeerLazyDefault(mgmProto.LazyState_LazyStateDefault))
}
// UpdatedRemoteFeatureFlag is called when the remote feature flag is updated.
// If enabled, it initializes the lazy connection manager and start it. Do not need to call Start() again.
// If disabled, then it closes the lazy connection manager and open the connections to all peers.
func (e *ConnMgr) UpdatedRemoteFeatureFlag(ctx context.Context, enabled bool) error {
// a local override (NB_LAZY_CONN or local config) takes precedence over management
if e.force != lazyForceNone {
return nil
// UpdatedRemoteFeatureFlag caches the account-wide lazy feature flag. The manager itself is
// not started or stopped here; the per-sync exclude-list reconciliation moves normal peers
// between the lazy and always-active sets when the flag flips.
func (e *ConnMgr) UpdatedRemoteFeatureFlag(_ context.Context, enabled bool) error {
e.remoteLazyEnabled = enabled
if e.isStartedWithLazyMgr() {
e.statusRecorder.UpdateLazyConnection(e.PeerLazyDefault(mgmProto.LazyState_LazyStateDefault))
}
return nil
}
// PeerLazyDefault reports whether a peer should be lazy. The local override
// (NB_LAZY_CONN/MDM) wins over everything; without a local override the
// management per-peer state applies (LazyStateLazy/Eager force the decision),
// and LazyStateDefault follows the account-wide flag.
func (e *ConnMgr) PeerLazyDefault(state mgmProto.LazyState) bool {
switch e.force {
case lazyForceOn:
return true
case lazyForceOff:
return false
}
if enabled {
// if the lazy connection manager is already started, do not start it again
if e.lazyConnMgr != nil {
return nil
}
if e.rosenpassEnabled {
log.Infof("rosenpass connection manager is enabled, lazy connection manager will not be started")
e.statusRecorder.UpdateLazyConnection(false)
return nil
}
log.Infof("lazy connection manager is enabled by the management feature flag")
e.initLazyManager(ctx)
e.statusRecorder.UpdateLazyConnection(true)
return e.addPeersToLazyConnManager()
} else {
if e.lazyConnMgr == nil {
e.statusRecorder.UpdateLazyConnection(false)
return nil
}
log.Infof("lazy connection manager is disabled by management feature flag")
e.closeManager(ctx)
e.statusRecorder.UpdateLazyConnection(false)
return nil
switch state {
case mgmProto.LazyState_LazyStateLazy:
return true
case mgmProto.LazyState_LazyStateEager:
return false
default:
return e.remoteLazyEnabled
}
}
@@ -141,6 +136,13 @@ func (e *ConnMgr) SetExcludeList(ctx context.Context, peerIDs map[string]bool) {
return
}
// The exclude set is recomputed every sync but rarely changes; skip the O(n)
// store lookups and reconciliation when it matches what was already applied.
if maps.Equal(peerIDs, e.appliedExcludeList) {
return
}
e.appliedExcludeList = maps.Clone(peerIDs)
excludedPeers := make([]lazyconn.PeerConfig, 0, len(peerIDs))
for peerID := range peerIDs {
@@ -176,12 +178,16 @@ func (e *ConnMgr) SetExcludeList(ctx context.Context, peerIDs map[string]bool) {
}
}
func (e *ConnMgr) AddPeerConn(ctx context.Context, peerKey string, conn *peer.Conn) (exists bool) {
// AddPeerConn registers a peer connection. permanent requests an always-active connection
// (the peer belongs to the exclude set: a forwarder, or a peer that is not lazy by policy).
// Non-permanent peers are handed to the lazy manager. The subsequent SetExcludeList call
// reconciles membership for existing peers across flag flips.
func (e *ConnMgr) AddPeerConn(ctx context.Context, peerKey string, conn *peer.Conn, permanent bool) (exists bool) {
if success := e.peerStore.AddPeerConn(peerKey, conn); !success {
return true
}
if !e.isStartedWithLazyMgr() {
if !e.isStartedWithLazyMgr() || permanent {
if err := conn.Open(ctx); err != nil {
conn.Log.Errorf("failed to open connection: %v", err)
}
@@ -269,6 +275,7 @@ func (e *ConnMgr) Close() {
e.lazyCtxCancel()
e.wg.Wait()
e.lazyConnMgr = nil
e.appliedExcludeList = nil
}
func (e *ConnMgr) initLazyManager(engineCtx context.Context) {
@@ -276,6 +283,7 @@ func (e *ConnMgr) initLazyManager(engineCtx context.Context) {
InactivityThreshold: inactivityThresholdEnv(),
}
e.lazyConnMgr = manager.NewManager(cfg, engineCtx, e.peerStore, e.iface)
e.appliedExcludeList = nil
e.lazyCtx, e.lazyCtxCancel = context.WithCancel(engineCtx)
@@ -286,43 +294,6 @@ func (e *ConnMgr) initLazyManager(engineCtx context.Context) {
}()
}
func (e *ConnMgr) addPeersToLazyConnManager() error {
peers := e.peerStore.PeersPubKey()
lazyPeerCfgs := make([]lazyconn.PeerConfig, 0, len(peers))
for _, peerID := range peers {
var peerConn *peer.Conn
var exists bool
if peerConn, exists = e.peerStore.PeerConn(peerID); !exists {
log.Warnf("failed to find peer conn for peerID: %s", peerID)
continue
}
lazyPeerCfg := lazyconn.PeerConfig{
PublicKey: peerID,
AllowedIPs: peerConn.WgConfig().AllowedIps,
PeerConnID: peerConn.ConnID(),
Log: peerConn.Log,
}
lazyPeerCfgs = append(lazyPeerCfgs, lazyPeerCfg)
}
return e.lazyConnMgr.AddActivePeers(lazyPeerCfgs)
}
func (e *ConnMgr) closeManager(ctx context.Context) {
if e.lazyConnMgr == nil {
return
}
e.lazyCtxCancel()
e.wg.Wait()
e.lazyConnMgr = nil
for _, peerID := range e.peerStore.PeersPubKey() {
e.peerStore.PeerConnOpen(ctx, peerID)
}
}
func (e *ConnMgr) isStartedWithLazyMgr() bool {
return e.lazyConnMgr != nil && e.lazyCtxCancel != nil
}

View File

@@ -5,6 +5,7 @@ import (
"testing"
"github.com/netbirdio/netbird/client/internal/lazyconn"
mgmProto "github.com/netbirdio/netbird/shared/management/proto"
)
func TestResolveLazyForce(t *testing.T) {
@@ -38,3 +39,90 @@ func TestResolveLazyForce(t *testing.T) {
})
}
}
func TestPeerLazyDefault(t *testing.T) {
tests := []struct {
name string
force lazyForce
remoteEnabled bool
state mgmProto.LazyState
want bool
}{
{name: "force on wins over eager state", force: lazyForceOn, state: mgmProto.LazyState_LazyStateEager, want: true},
{name: "force off wins over lazy state", force: lazyForceOff, remoteEnabled: true, state: mgmProto.LazyState_LazyStateLazy, want: false},
{name: "none, default, account off -> active", force: lazyForceNone, state: mgmProto.LazyState_LazyStateDefault, want: false},
{name: "none, default, account on -> lazy", force: lazyForceNone, remoteEnabled: true, state: mgmProto.LazyState_LazyStateDefault, want: true},
{name: "none, lazy state, account off -> lazy", force: lazyForceNone, state: mgmProto.LazyState_LazyStateLazy, want: true},
{name: "none, eager state, account on -> active", force: lazyForceNone, remoteEnabled: true, state: mgmProto.LazyState_LazyStateEager, want: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
e := &ConnMgr{force: tt.force, remoteLazyEnabled: tt.remoteEnabled}
if got := e.PeerLazyDefault(tt.state); got != tt.want {
t.Fatalf("PeerLazyDefault(%v) = %v, want %v", tt.state, got, tt.want)
}
})
}
}
// TestToExcludedLazyPeers covers the per-peer lazy classification (proxy vs
// normal, across the force/account-flag matrix). Forwarder-target exclusion is
// covered by TestToExcludedLazyPeers_ForwardTarget.
func TestToExcludedLazyPeers(t *testing.T) {
const (
normalKey = "normal"
lazyKey = "lazy-state"
eagerKey = "eager-state"
)
peers := []*mgmProto.RemotePeerConfig{
{WgPubKey: normalKey, AllowedIps: []string{"100.64.0.1/32"}},
{WgPubKey: lazyKey, AllowedIps: []string{"100.64.0.2/32"}, LazyState: mgmProto.LazyState_LazyStateLazy},
{WgPubKey: eagerKey, AllowedIps: []string{"100.64.0.3/32"}, LazyState: mgmProto.LazyState_LazyStateEager},
}
tests := []struct {
name string
force lazyForce
remoteEnabled bool
want map[string]bool
}{
{
name: "account off: lazy-state peer lazy, normal + eager active",
force: lazyForceNone, remoteEnabled: false,
want: map[string]bool{normalKey: true, eagerKey: true},
},
{
name: "account on: only eager-state peer active",
force: lazyForceNone, remoteEnabled: true,
want: map[string]bool{eagerKey: true},
},
{
name: "force off: everything active",
force: lazyForceOff, remoteEnabled: true,
want: map[string]bool{normalKey: true, lazyKey: true, eagerKey: true},
},
{
name: "force on: nothing active",
force: lazyForceOn, remoteEnabled: false,
want: map[string]bool{},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
e := &Engine{connMgr: &ConnMgr{force: tt.force, remoteLazyEnabled: tt.remoteEnabled}}
got := e.toExcludedLazyPeers(nil, peers)
if len(got) != len(tt.want) {
t.Fatalf("toExcludedLazyPeers() = %v, want %v", got, tt.want)
}
for k := range tt.want {
if !got[k] {
t.Fatalf("expected peer %s excluded, got %v", k, got)
}
}
})
}
}

View File

@@ -803,7 +803,7 @@ func (e *Engine) blockLanAccess() {
// modifyPeers updates peers that have been modified (e.g. IP address has been changed).
// It closes the existing connection, removes it from the peerConns map, and creates a new one.
func (e *Engine) modifyPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
func (e *Engine) modifyPeers(peersUpdate []*mgmProto.RemotePeerConfig, forwardingRules []firewallManager.ForwardRule) error {
// first, check if peers have been modified
var modified []*mgmProto.RemotePeerConfig
@@ -842,7 +842,7 @@ func (e *Engine) modifyPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
}
// third, add the peer connections again
for _, p := range modified {
err := e.addNewPeer(p)
err := e.addNewPeer(p, forwardingRules)
if err != nil {
return err
}
@@ -1403,8 +1403,12 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
return nil
}
if err := e.connMgr.UpdatedRemoteFeatureFlag(e.ctx, networkMap.GetPeerConfig().GetLazyConnectionEnabled()); err != nil {
log.Errorf("failed to update lazy connection feature flag: %v", err)
// Only update the flag when the sync carries a peer config; a nil peer config
// (e.g. a partial update) must not reset the cached flag to false.
if peerConfig := networkMap.GetPeerConfig(); peerConfig != nil {
if err := e.connMgr.UpdatedRemoteFeatureFlag(e.ctx, peerConfig.GetLazyConnectionEnabled()); err != nil {
log.Errorf("failed to update lazy connection feature flag: %v", err)
}
}
if e.firewall != nil {
@@ -1482,15 +1486,14 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
e.updateOfflinePeers(networkMap.GetOfflinePeers())
done()
remotePeers, err := e.reconcilePeers(networkMap)
remotePeers, err := e.reconcilePeers(networkMap, forwardingRules)
if err != nil {
return err
}
// must set the exclude list after the peers are added. Without it the manager can not figure out the peers parameters from the store
done = e.phase("lazy_exclude")
excludedLazyPeers := e.toExcludedLazyPeers(forwardingRules, remotePeers)
e.connMgr.SetExcludeList(e.ctx, excludedLazyPeers)
e.connMgr.SetExcludeList(e.ctx, e.toExcludedLazyPeers(forwardingRules, remotePeers))
done()
e.networkSerial = serial
@@ -1500,8 +1503,10 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
// reconcilePeers applies the remote peer list from the network map (removing,
// modifying and adding peers, then updating SSH config) and returns the remote
// peers with our own peer filtered out, for use by later sync steps.
func (e *Engine) reconcilePeers(networkMap *mgmProto.NetworkMap) ([]*mgmProto.RemotePeerConfig, error) {
// peers with our own peer filtered out, for use by later sync steps. The
// forwarding rules are used to decide whether a newly added peer needs an
// always-active connection.
func (e *Engine) reconcilePeers(networkMap *mgmProto.NetworkMap, forwardingRules []firewallManager.ForwardRule) ([]*mgmProto.RemotePeerConfig, error) {
// Filter out own peer from the remote peers list
localPubKey := e.config.WgPrivateKey.PublicKey().String()
remotePeers := make([]*mgmProto.RemotePeerConfig, 0, len(networkMap.GetRemotePeers()))
@@ -1529,14 +1534,14 @@ func (e *Engine) reconcilePeers(networkMap *mgmProto.NetworkMap) ([]*mgmProto.Re
}
done = e.phase("modified_peers")
err = e.modifyPeers(remotePeers)
err = e.modifyPeers(remotePeers, forwardingRules)
done()
if err != nil {
return nil, err
}
done = e.phase("added_peers")
err = e.addNewPeers(remotePeers)
err = e.addNewPeers(remotePeers, forwardingRules)
done()
if err != nil {
return nil, err
@@ -1732,9 +1737,9 @@ func addrToString(addr netip.Addr) string {
}
// addNewPeers adds peers that were not know before but arrived from the Management service with the update
func (e *Engine) addNewPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
func (e *Engine) addNewPeers(peersUpdate []*mgmProto.RemotePeerConfig, forwardingRules []firewallManager.ForwardRule) error {
for _, p := range peersUpdate {
err := e.addNewPeer(p)
err := e.addNewPeer(p, forwardingRules)
if err != nil {
return err
}
@@ -1742,8 +1747,9 @@ func (e *Engine) addNewPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
return nil
}
// addNewPeer add peer if connection doesn't exist
func (e *Engine) addNewPeer(peerConfig *mgmProto.RemotePeerConfig) error {
// addNewPeer add peer if connection doesn't exist. A peer that is not lazy by
// policy (or is a forwarder) gets an always-active connection instead.
func (e *Engine) addNewPeer(peerConfig *mgmProto.RemotePeerConfig, forwardingRules []firewallManager.ForwardRule) error {
peerKey := peerConfig.GetWgPubKey()
peerIPs := make([]netip.Prefix, 0, len(peerConfig.GetAllowedIps()))
if _, ok := e.peerStore.PeerConn(peerKey); ok {
@@ -1777,7 +1783,7 @@ func (e *Engine) addNewPeer(peerConfig *mgmProto.RemotePeerConfig) error {
log.Warnf("error adding peer %s to status recorder, got error: %v", peerKey, err)
}
if exists := e.connMgr.AddPeerConn(e.ctx, peerKey, conn); exists {
if exists := e.connMgr.AddPeerConn(e.ctx, peerKey, conn, e.isPermanentPeer(peerConfig, forwardingRules)); exists {
conn.Close(false)
return fmt.Errorf("peer already exists: %s", peerKey)
}
@@ -2603,34 +2609,44 @@ func (e *Engine) updateForwardRules(rules []*mgmProto.ForwardingRule) ([]firewal
return forwardingRules, nberrors.FormatErrorOrNil(merr)
}
// toExcludedLazyPeers returns the peers that must have an always-active
// connection, so the caller can reconcile the lazy manager's exclude list.
func (e *Engine) toExcludedLazyPeers(rules []firewallManager.ForwardRule, peers []*mgmProto.RemotePeerConfig) map[string]bool {
excludedPeers := make(map[string]bool)
// Ingress forward targets: inbound forwarded traffic is initiated remotely and
// cannot wake a lazy connection, so the peer routing the target must stay
// permanently connected. AllowedIPs are already parsed on the peer conn, so
// reuse those typed prefixes instead of re-parsing the network map strings.
for _, r := range rules {
for _, p := range peers {
if e.peerRoutesAddr(p, r.TranslatedAddress) {
log.Infof("exclude forwarder peer from lazy connection: %s", p.GetWgPubKey())
excludedPeers[p.GetWgPubKey()] = true
}
for _, p := range peers {
if e.isPermanentPeer(p, rules) {
excludedPeers[p.GetWgPubKey()] = true
}
}
return excludedPeers
}
// peerRoutesAddr reports whether the peer is a router for addr, matched against
// the peer's already-parsed AllowedIPs from the store (the same typed value the
// lazy manager consumes) rather than re-parsing the network map strings.
func (e *Engine) peerRoutesAddr(p *mgmProto.RemotePeerConfig, addr netip.Addr) bool {
prefixes, ok := e.peerStore.AllowedIPs(p.GetWgPubKey())
if !ok {
return false
// isPermanentPeer reports whether a peer needs an always-active connection: it
// is not lazy by policy (the per-peer lazy hint or account flag, subject to the
// local override), or it is an ingress forward target. Inbound forwarded traffic
// is initiated remotely and cannot wake a lazy connection, so the peer routing
// the target must stay permanently connected.
func (e *Engine) isPermanentPeer(p *mgmProto.RemotePeerConfig, rules []firewallManager.ForwardRule) bool {
if !e.connMgr.PeerLazyDefault(p.GetLazyState()) {
return true
}
return prefixesContain(prefixes, addr)
// Match against the incoming config's AllowedIPs rather than the peer store:
// isPermanentPeer runs in addNewPeer before the peer is in the store, so a
// store lookup would miss a forward target and register it as lazy.
prefixes := make([]netip.Prefix, 0, len(p.GetAllowedIps()))
for _, ipStr := range p.GetAllowedIps() {
if prefix, err := netip.ParsePrefix(ipStr); err == nil {
prefixes = append(prefixes, prefix)
}
}
for _, r := range rules {
if prefixesContain(prefixes, r.TranslatedAddress) {
log.Infof("exclude forwarder peer from lazy connection: %s", p.GetWgPubKey())
return true
}
}
return false
}
// prefixesContain reports whether addr falls within any of the prefixes.

View File

@@ -49,7 +49,8 @@ func TestToExcludedLazyPeers_ForwardTarget(t *testing.T) {
store.AddPeerConn(targetPeerKey, newTestConn(t, targetPeerKey, "100.110.8.145/32"))
store.AddPeerConn(otherPeerKey, newTestConn(t, otherPeerKey, "100.110.9.10/32"))
e := &Engine{peerStore: store}
// Lazy on for normal peers, so the only exclusion under test is the forward target.
e := &Engine{peerStore: store, connMgr: &ConnMgr{force: lazyForceOn}}
peers := []*mgmProto.RemotePeerConfig{
{WgPubKey: targetPeerKey, AllowedIps: []string{"100.110.8.145/32"}},
@@ -67,7 +68,8 @@ func TestToExcludedLazyPeers_ForwardTarget(t *testing.T) {
}
func TestToExcludedLazyPeers_NoRules(t *testing.T) {
e := &Engine{peerStore: peerstore.NewConnStore()}
// Lazy on for normal peers and no forward rules, so nothing is excluded.
e := &Engine{peerStore: peerstore.NewConnStore(), connMgr: &ConnMgr{force: lazyForceOn}}
peers := []*mgmProto.RemotePeerConfig{
{WgPubKey: "peer-a", AllowedIps: []string{"100.110.8.145/32"}},

View File

@@ -279,7 +279,8 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
}, MobileDependency{})
wgIface := &MockWGIface{
NameFunc: func() string { return "utun102" },
NameFunc: func() string { return "utun102" },
IsUserspaceBindFunc: func() bool { return true },
RemovePeerFunc: func(peerKey string) error {
return nil
},

View File

@@ -95,7 +95,7 @@ func (d *DnsInterceptor) RemoveRoute() error {
// AllowedIPs should use real IPs
if d.currentPeerKey != "" {
if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
if _, err := d.allowedIPsRefcounter.Decrement(prefix); 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, d.currentPeerKey); err != nil {
if _, err := d.allowedIPsRefcounter.Decrement(realPrefix); 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, d.currentPeerKey); err != nil {
if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
}
}

View File

@@ -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, r.currentPeerKey); err != nil {
if _, err := r.allowedIPsRefcounter.Decrement(prefix); 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, r.currentPeerKey); err != nil {
if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
}
}

View File

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

View File

@@ -1,185 +0,0 @@
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
}

View File

@@ -1,241 +0,0 @@
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])
}
}

View File

@@ -5,7 +5,5 @@ 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 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.
// 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]

View File

@@ -15,11 +15,6 @@ 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 {
@@ -57,15 +52,12 @@ func (r *Route) AddAllowedIPs(peerKey string) error {
ref.Out,
)
}
r.currentPeerKey = peerKey
return nil
}
func (r *Route) RemoveAllowedIPs() error {
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)
if _, err := r.allowedIPsRefcounter.Decrement(r.route.Network); err != nil {
return err
}
r.currentPeerKey = ""
return err
return nil
}

View File

@@ -181,7 +181,7 @@ func (s *Server) Start() error {
log.Warnf("failed to redirect stderr: %v", err)
}
if err := RestoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil {
if err := restoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil {
log.Warnf(errRestoreResidualState, err)
}
@@ -551,7 +551,7 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro
s.actCancel = cancel
s.mutex.Unlock()
if err := RestoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil {
if err := restoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil {
log.Warnf(errRestoreResidualState, err)
}
@@ -858,7 +858,7 @@ func (s *Server) Up(callerCtx context.Context, msg *proto.UpRequest) (*proto.UpR
return s.waitForUp(callerCtx)
}
if err := RestoreResidualState(callerCtx, s.profileManager.GetStatePath()); err != nil {
if err := restoreResidualState(callerCtx, s.profileManager.GetStatePath()); err != nil {
log.Warnf(errRestoreResidualState, err)
}

View File

@@ -46,7 +46,7 @@ func (s *Server) CleanState(ctx context.Context, req *proto.CleanStateRequest) (
if req.All {
// Reuse existing cleanup logic for all states
if err := RestoreResidualState(ctx, statePath); err != nil {
if err := restoreResidualState(ctx, statePath); err != nil {
return nil, status.Errorf(codes.Internal, "failed to clean all states: %v", err)
}
@@ -113,9 +113,9 @@ func (s *Server) DeleteState(ctx context.Context, req *proto.DeleteStateRequest)
}, nil
}
// RestoreResidualState checks if the client was not shut down in a clean way and restores residual if required.
// restoreResidualState checks if the client was not shut down in a clean way and restores residual if required.
// Otherwise, we might not be able to connect to the management server to retrieve new config.
func RestoreResidualState(ctx context.Context, statePath string) error {
func restoreResidualState(ctx context.Context, statePath string) error {
if statePath == "" {
return nil
}

View File

@@ -91,7 +91,7 @@ func availableProviders() []providerCase {
if region == "" {
region = "us-east-1"
}
ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: "us.anthropic.claude-haiku-4-5", kind: harness.WireBedrock})
ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: "us.anthropic.claude-haiku-4-5", kind: harness.WireMessages})
}
return ps
}
@@ -224,12 +224,9 @@ func TestProvidersMatrix(t *testing.T) {
var c int
var b string
var cerr error
switch pc.kind {
case harness.WireVertex:
if pc.kind == harness.WireVertex {
c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, "Reply with exactly: pong", sessionID)
case harness.WireBedrock:
c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, "Reply with exactly: pong", sessionID)
default:
} else {
c, b, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, pc.kind, pc.model, "Reply with exactly: pong", sessionID)
}
if cerr == nil {

View File

@@ -1,168 +0,0 @@
//go:build e2e
package agentnetwork
import (
"context"
"strings"
"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"
)
// catalogModel returns the normalized catalog id the proxy stamps for a
// path-routed provider's configured model — the form the guardrail allowlist is
// compared against (region prefix / @version stripped).
func catalogModel(pc providerCase) string {
switch pc.kind {
case harness.WireBedrock:
return strings.TrimPrefix(pc.model, "us.")
case harness.WireVertex:
return strings.SplitN(pc.model, "@", 2)[0]
default:
return pc.model
}
}
// disallowedModel returns a valid-shaped model id for the provider that is NOT
// the configured/allowed one, so the guardrail must reject it before the
// request ever reaches the upstream.
func disallowedModel(pc providerCase) string {
switch pc.kind {
case harness.WireBedrock:
return "us.anthropic.claude-opus-4-8"
case harness.WireVertex:
return "claude-opus-4-8@20250101"
default:
return "unlisted-model"
}
}
// sendModel drives one request for the given model through the provider's native
// wire shape and returns the HTTP status.
func sendModel(ctx context.Context, t *testing.T, cl *harness.Client, endpoint, proxyIP string, pc providerCase, model string) int {
t.Helper()
var code int
var err error
switch pc.kind {
case harness.WireBedrock:
code, _, err = cl.Bedrock(ctx, endpoint, proxyIP, model, "Reply with exactly: pong", "")
case harness.WireVertex:
code, _, err = cl.Vertex(ctx, endpoint, proxyIP, pc.project, pc.region, model, "Reply with exactly: pong", "")
default:
code, _, err = cl.Chat(ctx, endpoint, proxyIP, pc.kind, model, "Reply with exactly: pong", "")
}
require.NoError(t, err, "request must reach the proxy for %s", pc.name)
return code
}
// TestModelAllowlistEnforced provisions a Model Allowlist guardrail limiting each
// path-routed provider (Bedrock, Vertex) to its configured model, then drives
// requests over the tunnel: the allowed model returns 200 while a model outside
// the allowlist is denied 403 by the guardrail before it reaches the upstream.
// This is the coverage missing for #6751 — the model for these providers travels
// in the URL path, and the allowlist must be enforced there.
func TestModelAllowlistEnforced(t *testing.T) {
var providers []providerCase
for _, pc := range availableProviders() {
if pc.kind == harness.WireBedrock || pc.kind == harness.WireVertex {
providers = append(providers, pc)
}
}
if len(providers) == 0 {
t.Skip("no path-routed provider keys set (AWS_BEARER_TOKEN_BEDROCK / GOOGLE_VERTEX_*); source ~/.llm-keys")
}
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
defer cancel()
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-allowlist"})
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-allowlist-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{grp.Id},
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
// Providers with their configured (allowed) models; the first bootstraps the cluster.
ids := make([]string, 0, len(providers))
allowed := make([]string, 0, len(providers))
for i, pc := range providers {
req := providerRequest(pc)
if i == 0 {
req.BootstrapCluster = ptr(harness.AgentNetworkCluster)
}
prov, perr := srv.CreateProvider(ctx, req)
require.NoError(t, perr, "create provider %s", pc.name)
id := prov.Id
ids = append(ids, id)
allowed = append(allowed, catalogModel(pc))
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), id) })
}
// Guardrail allowlisting exactly the configured models.
var gr api.AgentNetworkGuardrailRequest
gr.Name = "e2e-allowlist"
gr.Checks.ModelAllowlist.Enabled = true
gr.Checks.ModelAllowlist.Models = allowed
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-allowlist",
Enabled: &enabled,
SourceGroups: []string{grp.Id},
DestinationProviderIds: ids,
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 for endpoint")
require.NotEmpty(t, settings.Endpoint, "agent-network endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-proxy-allowlist")
require.NoError(t, err, "mint proxy token via CLI")
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")
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()))
}
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
for _, pc := range providers {
pc := pc
t.Run(pc.name, func(t *testing.T) {
// The admin's allowlisted model is served end to end.
assert.Equal(t, 200, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, pc.model),
"allowlisted model must be permitted for %s", pc.name)
// A model outside the allowlist is rejected by the guardrail (before
// the upstream), regardless of whether it is a real catalog model.
assert.Equal(t, 403, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, disallowedModel(pc)),
"model outside the allowlist must be denied for %s", pc.name)
})
}
}

View File

@@ -107,17 +107,6 @@ func (c *Combined) DeletePolicy(ctx context.Context, id string) error {
return anDelete(ctx, c, "/api/agent-network/policies/"+id)
}
// CreateGuardrail creates an agent-network guardrail (e.g. a model allowlist)
// that can then be attached to a policy via its GuardrailIds.
func (c *Combined) CreateGuardrail(ctx context.Context, req api.AgentNetworkGuardrailRequest) (api.AgentNetworkGuardrail, error) {
return anRequest[api.AgentNetworkGuardrail](ctx, c, http.MethodPost, "/api/agent-network/guardrails", req)
}
// DeleteGuardrail removes a guardrail by id.
func (c *Combined) DeleteGuardrail(ctx context.Context, id string) error {
return anDelete(ctx, c, "/api/agent-network/guardrails/"+id)
}
// GetSettings returns the account's agent-network settings row. It exists only
// after the first provider create bootstraps it.
func (c *Combined) GetSettings(ctx context.Context) (api.AgentNetworkSettings, error) {

View File

@@ -194,11 +194,6 @@ const (
// WireVertex is the Anthropic-on-Vertex rawPredict shape: the client posts
// the full Vertex model path and the proxy mints the SA OAuth token.
WireVertex = "vertex"
// WireBedrock is the native AWS Bedrock InvokeModel shape: the model id
// travels in the URL path (/model/{id}/invoke), not the body, so the proxy
// routes by path. This is what a Bedrock SDK client sends and the shape the
// model-allowlist guardrail must enforce.
WireBedrock = "bedrock"
)
// Chat issues a chat-completion POST to the agent-network endpoint over the
@@ -231,17 +226,6 @@ func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
}
// Bedrock issues a native AWS Bedrock InvokeModel POST over the tunnel. The
// model id is carried in the request path (/model/{id}/invoke), so the proxy
// routes by path; the body uses the bedrock anthropic_version rather than a
// model field. A non-empty sessionID is sent as the universal x-session-id
// header the proxy records.
func (cl *Client) Bedrock(ctx context.Context, endpoint, proxyIP, model, prompt, sessionID string) (int, string, error) {
path := "/model/" + model + "/invoke"
body := fmt.Sprintf(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":64,"messages":[{"role":"user","content":%q}]}`, prompt)
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
}
// withSessionID appends the x-session-id header when sessionID is non-empty.
func withSessionID(headers []string, sessionID string) []string {
if sessionID == "" {

View File

@@ -226,6 +226,30 @@ func (m *managerImpl) CreateProxyPeer(ctx context.Context, accountID string, pee
return nil
}
// Dedupe stale embedded peer records for the same (account, cluster).
// The proxy generates a fresh WireGuard keypair on every startup
// (proxy/internal/roundtrip/netbird.go), so without this sweep the
// prior embedded peer would linger forever — holding its CGNAT IP
// allocation, polluting other peers' rosters, and (most visibly)
// leaving the synth DNS pointing at the dead address. The
// (account, cluster) tuple identifies "the embedded peer for this
// proxy instance at this cluster"; any record matching that tuple
// with a different pubkey is by definition stale and must go.
staleIDs, err := m.findStaleEmbeddedProxyPeers(ctx, accountID, cluster, peerKey)
if err != nil {
return fmt.Errorf("scan for stale embedded proxy peers: %w", err)
}
if len(staleIDs) > 0 {
// userID="" + checkConnected=false: the deletion is initiated
// by management itself on behalf of the freshly-registering
// proxy, not by an end user; the stale peer may still be
// marked Connected from its prior session, but its session is
// dead by definition (its key no longer exists).
if err := m.DeletePeers(ctx, accountID, staleIDs, "", false); err != nil {
return fmt.Errorf("delete stale embedded proxy peers %v: %w", staleIDs, err)
}
}
name := fmt.Sprintf("proxy-%s", xid.New().String())
newPeer := &peer.Peer{
Ephemeral: true,
@@ -251,3 +275,29 @@ func (m *managerImpl) CreateProxyPeer(ctx context.Context, accountID string, pee
return nil
}
// findStaleEmbeddedProxyPeers returns the peer IDs of embedded proxy peer
// records in accountID that target the same cluster but carry a different
// WireGuard pubkey than the freshly-registering one. Used by CreateProxyPeer
// to garbage-collect stale records left behind when the proxy restarts with a
// regenerated keypair.
func (m *managerImpl) findStaleEmbeddedProxyPeers(ctx context.Context, accountID, cluster, newKey string) ([]string, error) {
account, err := m.store.GetAccount(ctx, accountID)
if err != nil {
return nil, err
}
var stale []string
for _, p := range account.Peers {
if p == nil || !p.ProxyMeta.Embedded {
continue
}
if p.ProxyMeta.Cluster != cluster {
continue
}
if p.Key == newKey {
continue
}
stale = append(stale, p.ID)
}
return stale, nil
}

View File

@@ -164,6 +164,7 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
// filtered at the source (network map builder).
includeIPv6 := peer.SupportsIPv6() && peer.IPv6.IsValid()
useSourcePrefixes := peer.SupportsSourcePrefixes()
localIsProxy := peer.ProxyMeta.Embedded
response := &proto.SyncResponse{
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH),
@@ -183,7 +184,7 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
response.NetworkMap.PeerConfig = response.PeerConfig
remotePeers := make([]*proto.RemotePeerConfig, 0, len(networkMap.Peers)+len(networkMap.OfflinePeers))
remotePeers = appendRemotePeerConfig(remotePeers, networkMap.Peers, dnsName, includeIPv6)
remotePeers = appendRemotePeerConfig(remotePeers, networkMap.Peers, dnsName, includeIPv6, localIsProxy)
if !shouldSkipSendingDeprecatedRemotePeers(peer.Meta.WtVersion) {
response.RemotePeers = remotePeers
@@ -193,7 +194,7 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
response.RemotePeersIsEmpty = len(remotePeers) == 0
response.NetworkMap.RemotePeersIsEmpty = response.RemotePeersIsEmpty
response.NetworkMap.OfflinePeers = appendRemotePeerConfig(nil, networkMap.OfflinePeers, dnsName, includeIPv6)
response.NetworkMap.OfflinePeers = appendRemotePeerConfig(nil, networkMap.OfflinePeers, dnsName, includeIPv6, localIsProxy)
firewallRules := toProtocolFirewallRules(networkMap.FirewallRules, includeIPv6, useSourcePrefixes)
response.NetworkMap.FirewallRules = firewallRules
@@ -292,7 +293,7 @@ func shouldSkipSendingDeprecatedRemotePeers(peerVersion string) bool {
return precomputedDeprecatedRemotePeersConstraint.Check(peerNBVersion)
}
func appendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*nbpeer.Peer, dnsName string, includeIPv6 bool) []*proto.RemotePeerConfig {
func appendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*nbpeer.Peer, dnsName string, includeIPv6 bool, localIsProxy bool) []*proto.RemotePeerConfig {
for _, rPeer := range peers {
allowedIPs := []string{rPeer.IP.String() + "/32"}
if includeIPv6 && rPeer.IPv6.IsValid() {
@@ -304,11 +305,24 @@ func appendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*nbpeer.Peer,
SshConfig: &proto.SSHConfig{SshPubKey: []byte(rPeer.SSHKey)},
Fqdn: rPeer.FQDN(dnsName),
AgentVersion: rPeer.Meta.WtVersion,
LazyState: lazyStateFor(localIsProxy, rPeer),
})
}
return dst
}
// lazyStateFor returns the per-peer lazy override for a remote peer. Connections
// involving an ephemeral proxy peer on either endpoint default to lazy so shared
// proxy infrastructure is not kept permanently connected to every peer. All
// other peers follow the account-wide flag. A future admin-facing per-peer
// setting can return LazyStateEager here to force a peer always-active.
func lazyStateFor(localIsProxy bool, rPeer *nbpeer.Peer) proto.LazyState {
if localIsProxy || rPeer.ProxyMeta.Embedded {
return proto.LazyState_LazyStateLazy
}
return proto.LazyState_LazyStateDefault
}
// toProtocolDNSConfig converts nbdns.Config to proto.DNSConfig using the cache
func toProtocolDNSConfig(update nbdns.Config, cache *cache.DNSConfigCache, forwardPort int64) *proto.DNSConfig {
protoUpdate := &proto.DNSConfig{

View File

@@ -1,24 +1,19 @@
package main
import (
"log"
"net/http"
// nolint:gosec
_ "net/http/pprof"
"os"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/management/cmd"
)
func main() {
if pprofAddr := os.Getenv("NB_PPROF_ADDR"); pprofAddr != "" {
log.Infof("pprof enabled, listening on: %s", pprofAddr)
go func() {
log.Println(http.ListenAndServe(pprofAddr, nil))
}()
}
go func() {
log.Println(http.ListenAndServe("localhost:6060", nil))
}()
if err := cmd.Execute(); err != nil {
os.Exit(1)
}

View File

@@ -0,0 +1,199 @@
package server
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/peers"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agenttypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/permissions"
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/management/server/types"
)
// TestAgentNetwork_ProxyRestart_PropagatesNewPeerAndDropsStale is the no-mock
// regression guard for the bug the user reported: restarting the proxy creates
// a fresh embedded peer with a NEW WireGuard public key (the proxy generates
// the keypair on every startup at proxy/internal/roundtrip/netbird.go:312).
// The PRIOR embedded peer record is never deleted on management, so the
// account accumulates a stale peer holding a stale CGNAT IP. Other peers
// in the account either keep routing to the dead IP, or — if synth DNS
// picks the wrong record — never see the new IP at all.
//
// What this test exercises (no mocks):
// - real SQLite test store
// - real DefaultAccountManager, network-map controller, peer-update channels
// - real peers.Manager.CreateProxyPeer path (the very method the proxy
// invokes over gRPC on every startup)
// - real agentnetwork.Manager + synth chain so the client receives a
// concrete DNS record that must point at the LATEST proxy peer.
//
// Pre-fix expected behavior (red): two embedded peers exist after the
// "restart"; the synth DNS record points at the stale one; the client
// receives an update reflecting the new peer but the old one lingers.
// Post-fix expected behavior (green): exactly one embedded peer exists
// after restart (with the new key) AND the client's network map carries
// the synth DNS pointing at that new peer's CGNAT IP.
func TestAgentNetwork_ProxyRestart_PropagatesNewPeerAndDropsStale(t *testing.T) {
am, updateManager, err := createManager(t)
require.NoError(t, err, "createManager must succeed")
ctx := context.Background()
const (
accountID = "an-restart-acct"
adminUserID = "an-restart-admin"
groupAID = "an-restart-grp-A"
clusterAddr = "eu.proxy.netbird.io"
clientKey = "BhRPtynAAYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8="
// Two different proxy pubkeys — the "before" and "after" of a
// proxy-process restart with fresh-keypair generation.
proxyKey1 = "Aaaaa1aaaaYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8="
proxyKey2 = "Bbbbb2bbbbYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8="
)
// --- Account scaffold ---
account := newAccountWithId(ctx, accountID, adminUserID, "an-restart.test", "", "", false)
require.NoError(t, am.Store.SaveAccount(ctx, account))
clientPeer := &nbpeer.Peer{
Key: clientKey,
Name: "an-restart-client",
DNSLabel: "an-restart-client",
Meta: nbpeer.PeerSystemMeta{Hostname: "an-restart-client", GoOS: "linux", WtVersion: "development"},
}
addedClient, _, _, _, err := am.AddPeer(ctx, "", "", adminUserID, clientPeer, false)
require.NoError(t, err, "AddPeer for client must succeed")
require.NoError(t, am.MarkPeerConnected(ctx, clientKey, accountID, time.Now().UnixNano(), &types.NetworkMap{}),
"MarkPeerConnected for the client peer must succeed (affected-peer fan-out skips disconnected peers)")
// Place the client in group A so the synth policy reaches it.
account, err = am.Store.GetAccount(ctx, accountID)
require.NoError(t, err)
account.Groups[groupAID] = &types.Group{ID: groupAID, Name: "groupA", Peers: []string{addedClient.ID}}
require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must persist group A")
// --- Real peers + agent-network managers ---
permMgr := permissions.NewManager(am.Store)
peersMgr := peers.NewManager(am.Store, permMgr)
peersMgr.SetAccountManager(am)
peersMgr.SetNetworkMapController(am.networkMapController)
agentMgr := agentnetwork.NewManager(am.Store, permMgr, am, nil)
// Subscribe BEFORE any state-mutating call so we don't lose the update
// that contains the synth DNS record.
clientCh := updateManager.CreateChannel(ctx, addedClient.ID)
t.Cleanup(func() { updateManager.CloseChannel(ctx, addedClient.ID) })
drain(clientCh)
// --- First proxy startup: register peer key K1, then mark it
// connected. In production the proxy follows CreateProxyPeer with the
// regular sync stream which lands on MarkPeerConnected; the synth DNS
// path filters out peers that aren't Connected (types/account.go:323),
// so without this step no DNS record would be emitted.
require.NoError(t, peersMgr.CreateProxyPeer(ctx, accountID, proxyKey1, clusterAddr),
"first CreateProxyPeer (proxy startup) must succeed")
peer1ID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey1)
require.NoError(t, err, "proxy peer for K1 must be persisted after CreateProxyPeer")
require.NotEmpty(t, peer1ID)
require.NoError(t, am.MarkPeerConnected(ctx, proxyKey1, accountID, time.Now().UnixNano(), &types.NetworkMap{}),
"MarkPeerConnected for K1 must succeed")
account, err = am.Store.GetAccount(ctx, accountID)
require.NoError(t, err)
proxyIP1 := account.Peers[peer1ID].IP.String()
require.NotEmpty(t, proxyIP1, "K1 must have an assigned overlay IP")
// --- Provider + policy. CreateProvider / CreatePolicy trigger the
// agentnetwork reconcile which runs UpdateAccountPeers; the resulting
// NetworkMap delivered to the client carries the synth DNS record
// pointing at K1's IP. ---
provider, err := agentMgr.CreateProvider(ctx, adminUserID, &agenttypes.Provider{
AccountID: accountID,
ProviderID: "openai_api",
Name: "openai-test",
UpstreamURL: "https://api.openai.com",
APIKey: "sk-test-key",
Enabled: true,
Models: []agenttypes.ProviderModel{{ID: "gpt-5.4"}},
}, clusterAddr)
require.NoError(t, err, "CreateProvider must succeed")
_, err = agentMgr.CreatePolicy(ctx, adminUserID, &agenttypes.Policy{
AccountID: accountID,
Name: "p1",
Enabled: true,
SourceGroups: []string{groupAID},
DestinationProviderIDs: []string{provider.ID},
})
require.NoError(t, err, "CreatePolicy must succeed")
settings, err := am.Store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID)
require.NoError(t, err)
fqdn := settings.Endpoint()
rdata1 := awaitZoneRData(clientCh, clusterAddr, fqdn, true)
require.Equal(t, proxyIP1, rdata1,
"client must receive a synth DNS record pointing at K1's overlay IP after the synth path runs")
drain(clientCh)
// --- Proxy restart: NEW keypair K2, same account, same cluster ---
require.NoError(t, peersMgr.CreateProxyPeer(ctx, accountID, proxyKey2, clusterAddr),
"second CreateProxyPeer (proxy restart with fresh keypair) must succeed")
peer2ID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey2)
require.NoError(t, err, "proxy peer for K2 must be persisted after restart")
require.NotEmpty(t, peer2ID)
require.NoError(t, am.MarkPeerConnected(ctx, proxyKey2, accountID, time.Now().UnixNano(), &types.NetworkMap{}),
"MarkPeerConnected for K2 must succeed")
// In production the agent's sync stream pulls a fresh NetworkMap as
// part of its normal reconcile cadence; in this isolated test
// MarkPeerConnected's affected-peer fan-out can race the channel-side
// buffer in a way that swallows the synth-DNS-bearing update before
// our await reads it. Trigger an explicit account-wide fan-out so the
// assertion below tests what production actually delivers, not the
// in-test buffer race.
am.UpdateAccountPeers(ctx, accountID, types.UpdateReason{Resource: types.UpdateResourcePeer, Operation: types.UpdateOperationUpdate})
account, err = am.Store.GetAccount(ctx, accountID)
require.NoError(t, err)
proxyIP2 := account.Peers[peer2ID].IP.String()
require.NotEmpty(t, proxyIP2, "K2 must have an assigned overlay IP")
require.NotEqual(t, proxyIP1, proxyIP2, "K2 must get a different overlay IP than K1 (sanity)")
// CRITICAL ASSERTION 1: K1 must no longer be in the store. The SqlStore
// returns ("", nil) for a missing key rather than NotFound, so assert
// on the returned ID being empty.
staleID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey1)
require.NoError(t, err, "GetPeerIDByKey for a missing peer must not error")
assert.Empty(t, staleID,
"stale embedded proxy peer K1 must be removed when a new embedded peer registers for the same (account, cluster); pre-fix this assertion fails because management never cleans up the prior peer record")
// CRITICAL ASSERTION 2: exactly one embedded proxy peer remains, and it
// is K2.
account, err = am.Store.GetAccount(ctx, accountID)
require.NoError(t, err)
embeddedKeys := []string{}
for _, p := range account.Peers {
if p.ProxyMeta.Embedded {
embeddedKeys = append(embeddedKeys, p.Key)
}
}
assert.Equal(t, []string{proxyKey2}, embeddedKeys,
"after a proxy restart exactly one embedded proxy peer should remain — the one with the new key K2")
// CRITICAL ASSERTION 3: the synth DNS record the client receives now
// points at K2's IP, not K1's.
rdata2 := awaitZoneRData(clientCh, clusterAddr, fqdn, true)
assert.Equal(t, proxyIP2, rdata2,
"after proxy restart, the client's synth DNS record must point at the NEW embedded peer's IP, not the stale K1 IP")
}

View File

@@ -7,7 +7,6 @@ import (
"slices"
"strconv"
"strings"
"sync"
"time"
"github.com/netbirdio/netbird/client/ssh/auth"
@@ -43,14 +42,6 @@ type NetworkMapComponents struct {
PostureFailedPeers map[string]map[string]struct{}
RouterPeers map[string]*nbpeer.Peer
routesByPeerOnce sync.Once
routesByPeerIdx map[string][]routeIndexEntry
}
type routeIndexEntry struct {
route *route.Route
viaGroup bool
}
type AccountSettingsInfo struct {
@@ -539,43 +530,33 @@ func (c *NetworkMapComponents) getRoutingPeerRoutes(peerID string) (enabledRoute
disabledRoutes = append(disabledRoutes, r)
}
for _, entry := range c.routesByPeer()[peerID] {
if entry.viaGroup {
newPeerRoute := entry.route.Copy()
newPeerRoute.PeerGroups = nil
newPeerRoute.ID = route.ID(string(entry.route.ID) + ":" + peerID)
takeRoute(newPeerRoute)
continue
for _, r := range c.Routes {
for _, groupID := range r.PeerGroups {
group := c.GetGroupInfo(groupID)
if group == nil {
continue
}
for _, id := range group.Peers {
if id != peerID {
continue
}
newPeerRoute := r.Copy()
newPeerRoute.Peer = id
newPeerRoute.PeerGroups = nil
newPeerRoute.ID = route.ID(string(r.ID) + ":" + id)
takeRoute(newPeerRoute)
break
}
}
if r.Peer == peerID {
takeRoute(r.Copy())
}
takeRoute(entry.route.Copy())
}
return enabledRoutes, disabledRoutes
}
func (c *NetworkMapComponents) routesByPeer() map[string][]routeIndexEntry {
c.routesByPeerOnce.Do(func() {
idx := make(map[string][]routeIndexEntry)
for _, r := range c.Routes {
for _, groupID := range r.PeerGroups {
group := c.GetGroupInfo(groupID)
if group == nil {
continue
}
for _, id := range group.Peers {
idx[id] = append(idx[id], routeIndexEntry{route: r, viaGroup: true})
}
}
if r.Peer != "" {
idx[r.Peer] = append(idx[r.Peer], routeIndexEntry{route: r})
}
}
c.routesByPeerIdx = idx
})
return c.routesByPeerIdx
}
func (c *NetworkMapComponents) filterRoutesByGroups(routes []*route.Route, groupListMap LookupMap) []*route.Route {
var filteredRoutes []*route.Route
for _, r := range routes {

View File

@@ -25,14 +25,6 @@ const (
denyCodeModel = "llm_policy.model_blocked"
denyReasonModel = "model_blocked"
denyMessageModel = "model is not in the policy allowlist"
// Deny reason used when an allowlist is configured but the request model
// could not be determined. URL/path-routed providers (AWS Bedrock, Google
// Vertex, ...) carry the model outside the JSON body, so a request shape the
// parser does not recognise reaches the guardrail with no model. Such a
// request must be denied (fail closed), never waved through.
denyCodeModelUnknown = "llm_policy.model_unknown"
denyReasonModelUnknown = "model_unknown"
denyMessageModelUnknown = "request model could not be determined for the policy allowlist"
)
// Middleware enforces the model allowlist and optionally captures the
@@ -116,37 +108,23 @@ func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middlew
if len(m.cfg.ModelAllowlist) == 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.
if !modelPresent || normaliseModel(model) == "" {
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
if !modelPresent {
return nil
}
if m.modelInAllowlist(model) {
return nil
}
return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
}
// denyModel builds a 403 deny Output for a model-allowlist rejection. model is
// included in the details only when non-empty.
func denyModel(model, code, message, reason string) *middleware.Output {
details := map[string]string{}
if model != "" {
details["model"] = model
}
return &middleware.Output{
Decision: middleware.DecisionDeny,
DenyStatus: 403,
DenyReason: &middleware.DenyReason{
Code: code,
Message: message,
Details: details,
Code: denyCodeModel,
Message: denyMessageModel,
Details: map[string]string{"model": model},
},
Metadata: []middleware.KV{
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
{Key: middleware.KeyLLMPolicyReason, Value: reason},
{Key: middleware.KeyLLMPolicyReason, Value: denyReasonModel},
},
}
}

View File

@@ -102,44 +102,13 @@ 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.
func TestAllowlistMissingModelKeyAllows(t *testing.T) {
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
out, err := mw.Invoke(context.Background(), newInput())
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, 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")
assert.Equal(t, middleware.DecisionAllow, out.Decision, "missing model key must allow even with non-empty allowlist")
dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision)
assert.Equal(t, "deny", dec, "decision must be deny when model key is absent")
reason, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyReason)
assert.Equal(t, "model_unknown", reason, "reason metadata must be model_unknown")
}
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(
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")
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.
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")
assert.Equal(t, "allow", dec, "decision must be allow when model key is absent")
}
func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) {

View File

@@ -1,106 +0,0 @@
package llm_request_parser
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_guardrail"
)
// runParserGuardrail runs the request parser then the model-allowlist guardrail
// in SlotOnRequest order, threading the parser's metadata into the guardrail the
// same way the real chain does. It returns the guardrail decision so tests can
// assert allowlist enforcement for URL/path-routed providers end to end.
func runParserGuardrail(t *testing.T, url string, body []byte, allowlist []string) *middleware.Output {
t.Helper()
parser := newMiddleware(t)
parsed, err := parser.Invoke(context.Background(), &middleware.Input{
Slot: middleware.SlotOnRequest,
URL: url,
Body: body,
})
require.NoError(t, err, "parser must not error")
guard := llm_guardrail.New(llm_guardrail.Config{ModelAllowlist: allowlist})
out, err := guard.Invoke(context.Background(), &middleware.Input{
Slot: middleware.SlotOnRequest,
Metadata: parsed.Metadata,
})
require.NoError(t, err, "guardrail must not error")
require.NotNil(t, out, "guardrail must return an output")
return out
}
// TestModelAllowlist_URLRoutedProviders validates that the model allowlist is
// enforced for providers whose model travels in the URL path (AWS Bedrock,
// Google Vertex) rather than the JSON body. The "unknown action" case is the
// regression guard for #6751: a Bedrock request shape the parser cannot map to a
// model must fail closed under an allowlist instead of bypassing it.
func TestModelAllowlist_URLRoutedProviders(t *testing.T) {
const bedrockBody = `{"anthropic_version":"bedrock-2023-05-31","messages":[{"role":"user","content":"hi"}]}`
const vertexBody = `{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"hi"}]}`
tests := []struct {
name string
url string
body string
allowlist []string
decision middleware.Decision
denyCode string
}{
{
name: "bedrock allowed model passes",
url: "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-haiku-4-5-v1:0/invoke",
body: bedrockBody,
allowlist: []string{"anthropic.claude-haiku-4-5"},
decision: middleware.DecisionAllow,
},
{
name: "bedrock disallowed model denied",
url: "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-opus-4-8-v1:0/invoke",
body: bedrockBody,
allowlist: []string{"anthropic.claude-haiku-4-5"},
decision: middleware.DecisionDeny,
denyCode: "llm_policy.model_blocked",
},
{
name: "bedrock unknown action fails closed",
url: "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-opus-4-8-v1:0/some-future-action",
body: bedrockBody,
allowlist: []string{"anthropic.claude-haiku-4-5"},
decision: middleware.DecisionDeny,
denyCode: "llm_policy.model_unknown",
},
{
name: "vertex disallowed model denied",
url: "/v1/projects/p/locations/global/publishers/anthropic/models/claude-opus-4-8@20250101:rawPredict",
body: vertexBody,
allowlist: []string{"claude-haiku-4-5"},
decision: middleware.DecisionDeny,
denyCode: "llm_policy.model_blocked",
},
{
name: "vertex allowed model passes",
url: "/v1/projects/p/locations/global/publishers/anthropic/models/claude-haiku-4-5@20250101:rawPredict",
body: vertexBody,
allowlist: []string{"claude-haiku-4-5"},
decision: middleware.DecisionAllow,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
out := runParserGuardrail(t, tt.url, []byte(tt.body), tt.allowlist)
assert.Equal(t, tt.decision, out.Decision, "unexpected decision for %s", tt.name)
if tt.decision == middleware.DecisionDeny {
require.NotNil(t, out.DenyReason, "deny reason must be set for %s", tt.name)
assert.Equal(t, 403, out.DenyStatus, "deny status must be 403 for %s", tt.name)
assert.Equal(t, tt.denyCode, out.DenyReason.Code, "deny code for %s", tt.name)
}
})
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -486,6 +486,22 @@ message RemotePeerConfig {
string fqdn = 4;
string agentVersion = 5;
// lazyState is the management per-peer override for lazy (on-demand)
// connections to this remote peer. LazyStateDefault follows the account-wide
// flag; LazyStateLazy forces lazy; LazyStateEager forces an always-active
// connection. A local NB_LAZY_CONN/MDM override still wins over this.
LazyState lazyState = 6;
}
// LazyState is the management per-peer override for lazy connections.
enum LazyState {
// Follow the account-wide lazy connection flag.
LazyStateDefault = 0;
// Force a lazy (on-demand) connection regardless of the account flag.
LazyStateLazy = 1;
// Force an always-active connection regardless of the account flag.
LazyStateEager = 2;
}
// SSHConfig represents SSH configurations of a peer.