Compare commits

..

5 Commits

Author SHA1 Message Date
Maycon Santos
c6bf5fbbfb [management,client] 0.74.5 branch sync (#6769)
## Describe your changes
* [proxy] enforce model allowlist for URL-routed providers
(Bedrock/Vertex) by @mlsmaycon in
https://github.com/netbirdio/netbird/pull/6764
* [management] Remove proxy peer stale deduplication logic by @mlsmaycon
in https://github.com/netbirdio/netbird/pull/6768
## Issue ticket number and link

## Stack

<!-- branch-stack -->

### Checklist
- [ ] Is it a bug fix
- [ ] Is a typo/documentation fix
- [ ] Is a feature enhancement
- [ ] It is a refactor
- [ ] Created tests that fail without the change (if possible)
- [ ] This change does **not** modify the public API, gRPC protocols,
functionality behavior, CLI / service flags, or introduce a new feature
— **OR** I have discussed it with the NetBird team beforehand (link the
issue / Slack thread in the description). See
[CONTRIBUTING.md](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTING.md#discuss-changes-with-the-netbird-team-first).

> By submitting this pull request, you confirm that you have read and
agree to the terms of the [Contributor License
Agreement](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTOR_LICENSE_AGREEMENT.md).

## Documentation
Select exactly one:

- [ ] I added/updated documentation for this change
- [x] Documentation is **not needed** for this change (explain why)

### Docs PR URL (required if "docs added" is checked)
Paste the PR link from https://github.com/netbirdio/docs here:

https://github.com/netbirdio/docs/pull/__


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Added model-allowlist guardrails for path-routed providers, including
Bedrock and Vertex.
  - Added Bedrock request support for chat interactions.
  - Added guardrail management capabilities.

- **Bug Fixes**
- Requests with missing or blank model identifiers are now denied when a
model allowlist is configured, improving fail-closed protection.
- Corrected provider-specific request handling and session tracking for
Bedrock interactions.

- **Tests**
- Expanded coverage for allowlist enforcement and provider routing
scenarios.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Theodor Midtlien <theodor@midtlien.com>
Co-authored-by: blaugrau90 <61945343+blaugrau90@users.noreply.github.com>
Co-authored-by: Viktor Liu <17948409+lixmal@users.noreply.github.com>
2026-07-14 21:22:40 +02:00
David Fry
e70a69bbcf [client] Restore residual state in foreground mode before login (#6707)
* Improved residual state restoration during foreground startup and
foreground login, ensuring consistent recovery with stale states.
* Foreground flows now initialize advanced routing so stale routes 
are bypassed during login.
2026-07-14 17:43:59 +02:00
Riccardo Manfrin
a48618c074 [client] Fix forwarder peers never excluded from lazy connections (#6674)
* [client] Extract peerRoutesAddr helper in toExcludedLazyPeers

Refactor: pull the AllowedIPs match into a named
peerRoutesAddr helper and document why forward-target peers are excluded
from lazy connections. No behavior change; the existing address match is
preserved as-is.

* [client] Add failing test for lazy-conn forward-target exclusion

toExcludedLazyPeers compares AllowedIPs (CIDR) against the unmasked
TranslatedAddress, so forward-target peers are never excluded. This test
asserts the peer is excluded and fails on the current behavior; the fix
follows.

* [client] Fix lazy-conn exclusion for ingress forward peers

peerRoutesAddr compared AllowedIPs (CIDR, e.g. a peer's overlay IP as /32)
against the unmasked TranslatedAddress string, so the match never fired and
forward-target peers were never excluded from lazy connections. Use prefix
containment so a routed address matches the peer's AllowedIP

* [client] Reuse parsed AllowedIPs from peerStore in lazy exclusion

Instead of re-parsing the network map AllowedIPs strings, look up the
already-parsed []netip.Prefix from peerStore.AllowedIPs (the same typed
value the lazy manager itself consumes). A down/lazy peer still has its
conn in the store, so exclusion is unaffected by connection state. Extract
a pure prefixesContain helper and unit-test it.
2026-07-14 12:12:37 +02:00
Riccardo Manfrin
39193396f5 [client] Fix WGWatcher silently failing to restart on fast disconnect/reconnect (#6664)
* Stick new watcher creation to actual existence of af the conn

and its removal to the removal of such same conn.
Avoid debouncing and cross lock dead locking

* Discriminate not updated from timeout handshakes

* [Recheck watcher ctx cancellation under conn.mu in onWGDisconnected

onWGDisconnected only checked conn.ctx (the engine-scoped context), never
the watcher's own context. disableWgWatcherIfNeeded cancels the wgWatcherCtx,
not conn.ctx, so a disabled watcher's timeout callback did not see the
cancellation.

handshakeCheck runs lock-free, so between the ctx check in periodicHandshakeCheck
and acquiring conn.mu a fast disconnect/reconnect can slip in: the stale watcher
then acquires the lock and tears down the *new*, healthy connection based on the
old timeout, forcing the guard into an unnecessary reconnect (flap).

Recheck watcherCtx.Err() under conn.mu so a superseded watcher exits without
touching the connection that replaced it.

* Remove verbose comments

* Fixup merge conflict leftovers

* Fixup context brought by onWGDisconnected
2026-07-14 10:21:59 +02:00
Riccardo Manfrin
5343402385 [client, relay] Increase early-message buffer cap to 10000 to avoid dropping relayed handshakes 2026-07-13 22:46:03 +02:00
20 changed files with 566 additions and 327 deletions

View File

@@ -17,7 +17,9 @@ 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"
)
@@ -331,6 +333,14 @@ 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,6 +22,8 @@ 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"
@@ -229,6 +231,24 @@ 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

@@ -2605,13 +2605,14 @@ func (e *Engine) updateForwardRules(rules []*mgmProto.ForwardingRule) ([]firewal
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 {
ip := r.TranslatedAddress
for _, p := range peers {
for _, allowedIP := range p.GetAllowedIps() {
if allowedIP != ip.String() {
continue
}
if e.peerRoutesAddr(p, r.TranslatedAddress) {
log.Infof("exclude forwarder peer from lazy connection: %s", p.GetWgPubKey())
excludedPeers[p.GetWgPubKey()] = true
}
@@ -2621,6 +2622,27 @@ func (e *Engine) toExcludedLazyPeers(rules []firewallManager.ForwardRule, peers
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
}
return prefixesContain(prefixes, addr)
}
// prefixesContain reports whether addr falls within any of the prefixes.
func prefixesContain(prefixes []netip.Prefix, addr netip.Addr) bool {
for _, prefix := range prefixes {
if prefix.Contains(addr) {
return true
}
}
return false
}
// isChecksEqual checks if two slices of checks are equal.
func isChecksEqual(checks1, checks2 []*mgmProto.Checks) bool {
normalize := func(checks []*mgmProto.Checks) []string {

View File

@@ -0,0 +1,87 @@
package internal
import (
"net/netip"
"testing"
"github.com/stretchr/testify/require"
firewallManager "github.com/netbirdio/netbird/client/firewall/manager"
"github.com/netbirdio/netbird/client/internal/peer"
"github.com/netbirdio/netbird/client/internal/peerstore"
mgmProto "github.com/netbirdio/netbird/shared/management/proto"
)
func TestPrefixesContain(t *testing.T) {
tests := []struct {
name string
prefixes []string
addr string
want bool
}{
{name: "own overlay /32 matches", prefixes: []string{"100.110.8.145/32"}, addr: "100.110.8.145", want: true},
{name: "addr inside routed subnet", prefixes: []string{"10.121.0.0/16"}, addr: "10.121.208.4", want: true},
{name: "addr outside subnet", prefixes: []string{"10.121.0.0/16"}, addr: "10.122.0.1", want: false},
{name: "different /32", prefixes: []string{"100.110.8.145/32"}, addr: "100.110.8.146", want: false},
{name: "ipv6 /128 matches", prefixes: []string{"fd00::1/128"}, addr: "fd00::1", want: true},
{name: "no prefixes", prefixes: nil, addr: "10.121.208.4", want: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
prefixes := make([]netip.Prefix, 0, len(tt.prefixes))
for _, p := range tt.prefixes {
prefixes = append(prefixes, netip.MustParsePrefix(p))
}
require.Equal(t, tt.want, prefixesContain(prefixes, netip.MustParseAddr(tt.addr)))
})
}
}
// TestToExcludedLazyPeers_ForwardTarget guards a regression: the forward-target
// peer (the peer routing a ForwardRule.TranslatedAddress) must be excluded from
// lazy connections, matched via the peer's already-parsed AllowedIPs.
func TestToExcludedLazyPeers_ForwardTarget(t *testing.T) {
const targetPeerKey = "cccccccccccccccccccccccccccccccccccccccccc0="
const otherPeerKey = "dddddddddddddddddddddddddddddddddddddddddd0="
store := peerstore.NewConnStore()
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}
peers := []*mgmProto.RemotePeerConfig{
{WgPubKey: targetPeerKey, AllowedIps: []string{"100.110.8.145/32"}},
{WgPubKey: otherPeerKey, AllowedIps: []string{"100.110.9.10/32"}},
}
rules := []firewallManager.ForwardRule{
{TranslatedAddress: netip.MustParseAddr("100.110.8.145")},
}
excluded := e.toExcludedLazyPeers(rules, peers)
require.True(t, excluded[targetPeerKey], "forward-target peer must be excluded from lazy connections")
require.False(t, excluded[otherPeerKey], "non-target peer must not be excluded")
require.Len(t, excluded, 1)
}
func TestToExcludedLazyPeers_NoRules(t *testing.T) {
e := &Engine{peerStore: peerstore.NewConnStore()}
peers := []*mgmProto.RemotePeerConfig{
{WgPubKey: "peer-a", AllowedIps: []string{"100.110.8.145/32"}},
}
require.Empty(t, e.toExcludedLazyPeers(nil, peers))
}
func newTestConn(t *testing.T, key, allowedIP string) *peer.Conn {
t.Helper()
conn, err := peer.NewConn(peer.ConnConfig{
Key: key,
WgConfig: peer.WgConfig{AllowedIps: []netip.Prefix{netip.MustParsePrefix(allowedIP)}},
}, peer.ServiceDependencies{})
require.NoError(t, err)
return conn
}

View File

@@ -203,7 +203,6 @@ func NewConn(config ConnConfig, services ServiceDependencies) (*Conn, error) {
statusICE: worker.NewAtomicStatus(),
dumpState: dumpState,
endpointUpdater: NewEndpointUpdater(connLog, config.WgConfig, isController(config)),
wgWatcher: NewWGWatcher(connLog, config.WgConfig.WgInterface, config.Key, dumpState),
metricsRecorder: services.MetricsRecorder,
}
@@ -671,11 +670,12 @@ func (conn *Conn) onGuardEvent() {
}
}
func (conn *Conn) onWGDisconnected() {
func (conn *Conn) onWGDisconnected(watcherCtx context.Context) {
conn.mu.Lock()
defer conn.mu.Unlock()
if conn.ctx.Err() != nil {
// watcherCtx guards against a stale watcher tearing down a connection that already superseded it.
if conn.ctx.Err() != nil || watcherCtx.Err() != nil {
return
}
@@ -833,25 +833,39 @@ func (conn *Conn) isConnectedOnAllWay() (status guard.ConnStatus) {
})
}
// enableWgWatcherIfNeeded starts a fresh watcher instance per connection attempt, so its
// lifecycle stays bound to conn.mu and enable/disable can't race an old goroutine's shutdown.
// Caller must hold conn.mu.
func (conn *Conn) enableWgWatcherIfNeeded(enabledTime time.Time) {
if !conn.wgWatcher.PrepareInitialHandshake() {
if conn.wgWatcher != nil {
return
}
watcher := NewWGWatcher(conn.Log, conn.config.WgConfig.WgInterface, conn.config.Key, conn.dumpState)
watcher.PrepareInitialHandshake()
wgWatcherCtx, wgWatcherCancel := context.WithCancel(conn.ctx)
conn.wgWatcher = watcher
conn.wgWatcherCancel = wgWatcherCancel
conn.wgWatcherWg.Add(1)
go func() {
defer conn.wgWatcherWg.Done()
conn.wgWatcher.EnableWgWatcher(wgWatcherCtx, enabledTime, conn.onWGDisconnected, conn.onWGHandshakeSuccess, conn.onWGCheckSuccess)
onDisconnected := func() { conn.onWGDisconnected(wgWatcherCtx) }
watcher.EnableWgWatcher(wgWatcherCtx, enabledTime, onDisconnected, conn.onWGHandshakeSuccess, conn.onWGCheckSuccess)
}()
}
// disableWgWatcherIfNeeded cancels and drops the watcher once no transport is active. It never
// waits for the goroutine: the timeout path reentrantly calls back here under conn.mu, so
// blocking would deadlock. Caller must hold conn.mu.
func (conn *Conn) disableWgWatcherIfNeeded() {
if conn.currentConnPriority == conntype.None && conn.wgWatcherCancel != nil {
conn.wgWatcherCancel()
conn.wgWatcherCancel = nil
if conn.currentConnPriority != conntype.None || conn.wgWatcher == nil {
return
}
conn.wgWatcherCancel()
conn.wgWatcher = nil
conn.wgWatcherCancel = nil
}
func (conn *Conn) newProxy(remoteConn net.Conn) (wgproxy.Proxy, error) {
@@ -874,7 +888,9 @@ func (conn *Conn) resetEndpoint() {
return
}
conn.Log.Infof("reset wg endpoint")
conn.wgWatcher.Reset()
if conn.wgWatcher != nil {
conn.wgWatcher.Reset()
}
if err := conn.endpointUpdater.RemoveEndpointAddress(); err != nil {
conn.Log.Warnf("failed to remove endpoint address before update: %v", err)
}

View File

@@ -339,20 +339,20 @@ func TestConn_onWGDisconnected_EscalatesToRosenpassReset(t *testing.T) {
conn := newWGTimeoutTestConn(true, &disconnected)
for i := 0; i < wgTimeoutEscalationThreshold-1; i++ {
conn.onWGDisconnected()
conn.onWGDisconnected(conn.ctx)
}
assert.Empty(t, disconnected, "escalation must not fire below the threshold")
conn.onWGDisconnected()
conn.onWGDisconnected(conn.ctx)
assert.Equal(t, []string{conn.config.WgConfig.RemoteKey}, disconnected,
"reaching the threshold must report the peer disconnected once")
for i := 0; i < wgTimeoutEscalationThreshold-1; i++ {
conn.onWGDisconnected()
conn.onWGDisconnected(conn.ctx)
}
assert.Len(t, disconnected, 1, "escalation must restart counting after firing")
conn.onWGDisconnected()
conn.onWGDisconnected(conn.ctx)
assert.Len(t, disconnected, 2, "continued timeouts must escalate again")
}
@@ -364,12 +364,12 @@ func TestConn_onWGDisconnected_CheckSuccessResetsEscalation(t *testing.T) {
conn := newWGTimeoutTestConn(true, &disconnected)
for i := 0; i < wgTimeoutEscalationThreshold-1; i++ {
conn.onWGDisconnected()
conn.onWGDisconnected(conn.ctx)
}
conn.onWGCheckSuccess()
for i := 0; i < wgTimeoutEscalationThreshold-1; i++ {
conn.onWGDisconnected()
conn.onWGDisconnected(conn.ctx)
}
assert.Empty(t, disconnected, "handshake success must reset the timeout count")
}
@@ -382,7 +382,7 @@ func TestConn_onWGDisconnected_NoEscalationWithoutRosenpass(t *testing.T) {
conn := newWGTimeoutTestConn(false, &disconnected)
for i := 0; i < wgTimeoutEscalationThreshold*3; i++ {
conn.onWGDisconnected()
conn.onWGDisconnected(conn.ctx)
}
assert.Empty(t, disconnected, "escalation must be limited to rosenpass connections")
}

View File

@@ -3,7 +3,6 @@ package peer
import (
"context"
"fmt"
"sync"
"time"
log "github.com/sirupsen/logrus"
@@ -24,14 +23,14 @@ type WGInterfaceStater interface {
GetStats() (map[string]configurer.WGStats, error)
}
// WGWatcher is single-shot: one instance per connection attempt, run once, then discarded.
// Lifecycle is owned by Conn under conn.mu, so it keeps no "enabled" state to go stale.
type WGWatcher struct {
log *log.Entry
wgIfaceStater WGInterfaceStater
peerKey string
stateDump *stateDump
enabled bool
muEnabled sync.Mutex
// initialHandshake is not thread-safe; never call PrepareInitialHandshake and EnableWgWatcher concurrently.
initialHandshake time.Time
@@ -48,25 +47,14 @@ func NewWGWatcher(log *log.Entry, wgIfaceStater WGInterfaceStater, peerKey strin
}
}
// PrepareInitialHandshake reserves the watcher and reads the peer's current WireGuard
// handshake time. It must be called before the peer is (re)configured on the WireGuard
// interface, so the captured baseline reflects the state prior to this connection attempt
// instead of racing with that configuration. Returns ok=false if the watcher is already
// running, in which case EnableWgWatcher must not be called.
func (w *WGWatcher) PrepareInitialHandshake() (ok bool) {
w.muEnabled.Lock()
if w.enabled {
w.muEnabled.Unlock()
return false
}
// PrepareInitialHandshake reads the peer's current WireGuard handshake time. It must be
// called before the peer is (re)configured on the WireGuard interface, so the captured
// baseline reflects the state prior to this connection attempt instead of racing with
// that configuration.
func (w *WGWatcher) PrepareInitialHandshake() {
w.log.Debugf("enable WireGuard watcher")
w.enabled = true
w.muEnabled.Unlock()
handshake, _ := w.wgState()
w.initialHandshake = handshake
return true
}
// EnableWgWatcher runs the WireGuard watcher loop using the handshake baseline captured by
@@ -76,10 +64,6 @@ func (w *WGWatcher) PrepareInitialHandshake() (ok bool) {
// handshake, including the first.
func (w *WGWatcher) EnableWgWatcher(ctx context.Context, enabledTime time.Time, onDisconnectedFn func(), onHandshakeSuccessFn func(when time.Time), onCheckSuccessFn func()) {
w.periodicHandshakeCheck(ctx, onDisconnectedFn, onHandshakeSuccessFn, onCheckSuccessFn, enabledTime, w.initialHandshake)
w.muEnabled.Lock()
w.enabled = false
w.muEnabled.Unlock()
}
// Reset signals the watcher that the WireGuard peer has been reset and a new
@@ -105,6 +89,7 @@ func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, onDisconnectedFn
case <-timer.C:
handshake, ok := w.handshakeCheck(lastHandshake)
if !ok {
// early ctx cancel check return
if ctx.Err() != nil {
return
}
@@ -153,9 +138,9 @@ func (w *WGWatcher) handshakeCheck(lastHandshake time.Time) (*time.Time, bool) {
w.log.Tracef("previous handshake, handshake: %v, %v", lastHandshake, handshake)
// the current know handshake did not change
// the current known handshake did not change
if handshake.Equal(lastHandshake) {
w.log.Warnf("WireGuard handshake timed out: %v", handshake)
w.log.Warnf("WireGuard handshake not updated: %v", handshake)
return nil, false
}

View File

@@ -7,7 +7,6 @@ import (
"time"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/client/iface/configurer"
)
@@ -62,7 +61,7 @@ func TestWGWatcher_CheckSuccessCallback(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
require.True(t, watcher.PrepareInitialHandshake())
watcher.PrepareInitialHandshake()
firstHandshake := make(chan struct{}, 1)
checkSuccess := make(chan struct{}, 1)
@@ -101,8 +100,7 @@ func TestWGWatcher_EnableWgWatcher(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
ok := watcher.PrepareInitialHandshake()
require.True(t, ok, "watcher should not be enabled yet")
watcher.PrepareInitialHandshake()
onDisconnected := make(chan struct{}, 1)
go watcher.EnableWgWatcher(ctx, time.Now(), func() {
@@ -132,8 +130,7 @@ func TestWGWatcher_ReEnable(t *testing.T) {
watcher := NewWGWatcher(mlog, mocWgIface, "", newStateDump("peer", mlog, &Status{}))
ctx, cancel := context.WithCancel(context.Background())
ok := watcher.PrepareInitialHandshake()
require.True(t, ok, "watcher should not be enabled yet")
watcher.PrepareInitialHandshake()
wg := &sync.WaitGroup{}
wg.Add(1)
@@ -149,8 +146,7 @@ func TestWGWatcher_ReEnable(t *testing.T) {
ctx, cancel = context.WithCancel(context.Background())
defer cancel()
ok = watcher.PrepareInitialHandshake()
require.True(t, ok, "watcher should be re-enabled after the previous run stopped")
watcher.PrepareInitialHandshake()
onDisconnected := make(chan struct{}, 1)
go watcher.EnableWgWatcher(ctx, time.Now(), func() {

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.WireMessages})
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})
}
return ps
}
@@ -224,9 +224,12 @@ func TestProvidersMatrix(t *testing.T) {
var c int
var b string
var cerr error
if pc.kind == harness.WireVertex {
switch pc.kind {
case harness.WireVertex:
c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, "Reply with exactly: pong", sessionID)
} else {
case harness.WireBedrock:
c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, "Reply with exactly: pong", sessionID)
default:
c, b, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, pc.kind, pc.model, "Reply with exactly: pong", sessionID)
}
if cerr == nil {

View File

@@ -0,0 +1,168 @@
//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,6 +107,17 @@ 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,6 +194,11 @@ 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
@@ -226,6 +231,17 @@ 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,31 +226,6 @@ 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).
log.WithContext(ctx).Debugf("removing %d stale embedded proxy peer records [%s]", len(staleIDs), staleIDs)
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,
@@ -276,33 +251,3 @@ 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) {
peers, err := m.store.GetAccountPeers(ctx, store.LockingStrengthNone, accountID, "", "")
if err != nil {
return nil, err
}
var stale []string
for _, p := range peers {
if p == nil || !p.ProxyMeta.Embedded {
continue
}
if p.ProxyMeta.Cluster != cluster {
continue
}
if p.Key == newKey {
continue
}
stale = append(stale, p.ID)
}
if len(stale) > 0 {
log.WithContext(ctx).Tracef("found stale embedded proxy peer(s) with cluster: %s, and pubkey: %s but voided", cluster, newKey)
}
// returning empty for validating
return []string{}, nil
}

View File

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

@@ -25,6 +25,14 @@ 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
@@ -108,23 +116,37 @@ func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middlew
if len(m.cfg.ModelAllowlist) == 0 {
return nil
}
if !modelPresent {
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 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: denyCodeModel,
Message: denyMessageModel,
Details: map[string]string{"model": model},
Code: code,
Message: message,
Details: details,
},
Metadata: []middleware.KV{
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
{Key: middleware.KeyLLMPolicyReason, Value: denyReasonModel},
{Key: middleware.KeyLLMPolicyReason, Value: reason},
},
}
}

View File

@@ -102,13 +102,44 @@ func TestAllowlistCaseInsensitive(t *testing.T) {
}
}
func TestAllowlistMissingModelKeyAllows(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())
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "missing model key must allow even with non-empty allowlist")
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")
dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision)
assert.Equal(t, "allow", dec, "decision must be allow when model key is absent")
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")
}
func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) {

View File

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

View File

@@ -10,7 +10,7 @@ import (
const (
earlyMsgTTL = 5 * time.Second
earlyMsgCapacity = 1000
earlyMsgCapacity = 10000
)
// earlyMsgBuffer buffers transport messages that arrive before the corresponding