mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-25 01:41:29 +02:00
Compare commits
9 Commits
agent-netw
...
revert/com
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7ed3737cda | ||
|
|
b65ec8b68a | ||
|
|
e13bcdbd44 | ||
|
|
46568f7af8 | ||
|
|
3358138ccc | ||
|
|
6d15d0729a | ||
|
|
0936918d24 | ||
|
|
d4a4418969 | ||
|
|
31ed241a1a |
8
.github/workflows/agent-network-e2e.yml
vendored
8
.github/workflows/agent-network-e2e.yml
vendored
@@ -5,11 +5,6 @@ on:
|
||||
schedule:
|
||||
- cron: "0 3 * * *"
|
||||
workflow_dispatch:
|
||||
# TEMPORARY: run on every push to the per-policy model-allowlist branch while
|
||||
# it is under review. Remove this push trigger before merge.
|
||||
push:
|
||||
branches:
|
||||
- agent-network-per-policy-model-allowlist
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.ref }}
|
||||
@@ -56,6 +51,9 @@ jobs:
|
||||
# token (and URL, for gateways) is unset, so partial coverage is fine.
|
||||
OPENAI_TOKEN: ${{ secrets.E2E_OPENAI_TOKEN }}
|
||||
ANTHROPIC_TOKEN: ${{ secrets.E2E_ANTHROPIC_TOKEN }}
|
||||
# Moonshot AI platform key (platform.kimi.ai); drives both Kimi wire
|
||||
# shapes (OpenAI /v1 and Anthropic /anthropic) through kimi_api.
|
||||
KIMI_TOKEN: ${{ secrets.E2E_KIMI_TOKEN }}
|
||||
VERCEL_URL: ${{ secrets.E2E_VERCEL_URL }}
|
||||
VERCEL_TOKEN: ${{ secrets.E2E_VERCEL_TOKEN }}
|
||||
OPENROUTER_URL: ${{ secrets.E2E_OPENROUTER_URL }}
|
||||
|
||||
@@ -569,17 +569,17 @@ func Test_ConnectPeers(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// todo: investigate why in some tests execution we need 30s
|
||||
// The peers use userspace WireGuard (stdnet transport). A tight busy-loop
|
||||
// here starves the wireguard-go goroutines that process the handshake, so
|
||||
// poll on a ticker instead and yield the CPU between checks. WireGuard also
|
||||
// only retries a lost handshake initiation every REKEY_TIMEOUT (5s), which
|
||||
// is why the overall wait can occasionally stretch to tens of seconds.
|
||||
timeout := 30 * time.Second
|
||||
timeoutChannel := time.After(timeout)
|
||||
ticker := time.NewTicker(500 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-timeoutChannel:
|
||||
t.Fatalf("waiting for peer handshake timeout after %s", timeout.String())
|
||||
default:
|
||||
}
|
||||
|
||||
peer, gpErr := getPeer(peer1ifaceName, peer2Key.PublicKey().String())
|
||||
if gpErr != nil {
|
||||
t.Fatal(gpErr)
|
||||
@@ -588,6 +588,12 @@ func Test_ConnectPeers(t *testing.T) {
|
||||
t.Log("peers successfully handshake")
|
||||
break
|
||||
}
|
||||
|
||||
select {
|
||||
case <-timeoutChannel:
|
||||
t.Fatalf("waiting for peer handshake timeout after %s", timeout.String())
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -49,11 +49,21 @@ type ConnMgr struct {
|
||||
// engine.syncMsgMux; all other reads stay under engine.syncMsgMux only.
|
||||
lazyConnMgrMu sync.RWMutex
|
||||
|
||||
// reconcileRoutedIPs re-applies a peer's routed allowed IPs after its lazy wake endpoint is
|
||||
// (re)armed (Mode A at arm time). Injected by the engine; nil disables the reconcile.
|
||||
reconcileRoutedIPs func(peerKey string) error
|
||||
|
||||
wg sync.WaitGroup
|
||||
lazyCtx context.Context
|
||||
lazyCtxCancel context.CancelFunc
|
||||
}
|
||||
|
||||
// SetRoutedIPsReconciler injects the callback used to re-apply a peer's routed allowed IPs when
|
||||
// its lazy wake endpoint is (re)armed. Must be called before the lazy manager starts.
|
||||
func (e *ConnMgr) SetRoutedIPsReconciler(fn func(peerKey string) error) {
|
||||
e.reconcileRoutedIPs = fn
|
||||
}
|
||||
|
||||
func NewConnMgr(engineConfig *EngineConfig, statusRecorder *peer.Status, peerStore *peerstore.Store, iface lazyconn.WGIface) *ConnMgr {
|
||||
e := &ConnMgr{
|
||||
peerStore: peerStore,
|
||||
@@ -291,6 +301,7 @@ func (e *ConnMgr) Close() {
|
||||
func (e *ConnMgr) initLazyManager(engineCtx context.Context) {
|
||||
cfg := manager.Config{
|
||||
InactivityThreshold: inactivityThresholdEnv(),
|
||||
ReconcileAllowedIPs: e.reconcileRoutedIPs,
|
||||
}
|
||||
|
||||
e.lazyConnMgrMu.Lock()
|
||||
|
||||
Binary file not shown.
Binary file not shown.
@@ -52,11 +52,14 @@ int xdp_dns_fwd(struct iphdr *ip, struct udphdr *udp) {
|
||||
|
||||
if (udp->dest == GENERAL_DNS_PORT && ip->daddr == dns_ip) {
|
||||
udp->dest = dns_port;
|
||||
// Clear the now-stale checksum; zero means "not computed" for IPv4.
|
||||
udp->check = 0;
|
||||
return XDP_PASS;
|
||||
}
|
||||
|
||||
if (udp->source == dns_port && ip->saddr == dns_ip) {
|
||||
udp->source = GENERAL_DNS_PORT;
|
||||
udp->check = 0;
|
||||
return XDP_PASS;
|
||||
}
|
||||
|
||||
|
||||
@@ -50,5 +50,11 @@ int xdp_wg_proxy(struct iphdr *ip, struct udphdr *udp) {
|
||||
__be16 new_dst_port = htons(proxy_port);
|
||||
udp->dest = new_dst_port;
|
||||
udp->source = new_src_port;
|
||||
|
||||
// The ports are covered by the UDP checksum. This is an IPv4 loopback hop
|
||||
// and the payload is already integrity-protected, so clear the checksum (a
|
||||
// zero UDP checksum means "not computed" for IPv4) rather than leave a
|
||||
// stale value the kernel would drop as UDP_CSUM.
|
||||
udp->check = 0;
|
||||
return XDP_PASS;
|
||||
}
|
||||
|
||||
@@ -663,6 +663,12 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
|
||||
iceCfg := e.createICEConfig()
|
||||
|
||||
e.connMgr = NewConnMgr(e.config, e.statusRecorder, e.peerStore, wgIface)
|
||||
e.connMgr.SetRoutedIPsReconciler(func(peerKey string) error {
|
||||
if e.routeManager == nil {
|
||||
return nil
|
||||
}
|
||||
return e.routeManager.ReconcilePeerAllowedIPs(peerKey)
|
||||
})
|
||||
e.connMgr.Start(e.ctx)
|
||||
|
||||
// Wire DNS-time lazy-connection warm-up now that the connection manager
|
||||
|
||||
@@ -29,6 +29,11 @@ type managedPeer struct {
|
||||
|
||||
type Config struct {
|
||||
InactivityThreshold *time.Duration
|
||||
// ReconcileAllowedIPs re-applies a peer's routed allowed IPs after its wake endpoint is
|
||||
// armed. The activity listener creates the wake peer with the overlay /32 only; without the
|
||||
// routed prefixes WireGuard would not steer subnet-bound traffic to the wake endpoint, so an
|
||||
// idle routing peer could never be woken by that traffic. Optional; nil disables the reconcile.
|
||||
ReconcileAllowedIPs func(peerKey string) error
|
||||
}
|
||||
|
||||
// Manager manages lazy connections
|
||||
@@ -56,6 +61,9 @@ type Manager struct {
|
||||
peerToHAGroups map[string][]route.HAUniqueID // peer ID -> HA groups they belong to
|
||||
haGroupToPeers map[route.HAUniqueID][]string // HA group -> peer IDs in the group
|
||||
routesMu sync.RWMutex
|
||||
|
||||
// reconcileAllowedIPs re-applies a peer's routed allowed IPs after its wake endpoint is armed.
|
||||
reconcileAllowedIPs func(peerKey string) error
|
||||
}
|
||||
|
||||
// NewManager creates a new lazy connection manager
|
||||
@@ -73,6 +81,7 @@ func NewManager(config Config, engineCtx context.Context, peerStore *peerstore.S
|
||||
activityManager: activity.NewManager(wgIface),
|
||||
peerToHAGroups: make(map[string][]route.HAUniqueID),
|
||||
haGroupToPeers: make(map[route.HAUniqueID][]string),
|
||||
reconcileAllowedIPs: config.ReconcileAllowedIPs,
|
||||
}
|
||||
|
||||
if wgIface.IsUserspaceBind() {
|
||||
@@ -201,7 +210,7 @@ func (m *Manager) AddPeer(peerCfg lazyconn.PeerConfig) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil {
|
||||
if err := m.armActivityListener(peerCfg); err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
@@ -288,7 +297,7 @@ func (m *Manager) DeactivatePeer(peerID peerid.ConnID) {
|
||||
|
||||
m.inactivityManager.RemovePeer(mp.peerCfg.PublicKey)
|
||||
|
||||
if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
|
||||
if err := m.armActivityListener(*mp.peerCfg); err != nil {
|
||||
mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
|
||||
return
|
||||
}
|
||||
@@ -465,6 +474,31 @@ func (m *Manager) close() {
|
||||
}
|
||||
|
||||
// shouldDeferIdleForHA checks if peer should stay connected due to HA group requirements
|
||||
// armRoutedAllowedIPs re-applies the peer's routed allowed IPs onto its freshly armed wake
|
||||
// endpoint. The activity listener creates the wake peer with the overlay /32 only, so without
|
||||
// this the routed prefixes would be missing and traffic to a routed subnet could not wake the
|
||||
// idle routing peer. It is a no-op when no reconciler is configured.
|
||||
// armActivityListener (re)arms the peer's wake endpoint via the activity manager and then
|
||||
// re-applies its routed allowed IPs, so traffic to a routed subnet can wake an idle routing
|
||||
// peer. The routed prefixes must be re-applied after the wake endpoint exists because the
|
||||
// listener creates it with the overlay /32 only.
|
||||
func (m *Manager) armActivityListener(peerCfg lazyconn.PeerConfig) error {
|
||||
if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil {
|
||||
return err
|
||||
}
|
||||
m.armRoutedAllowedIPs(&peerCfg)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) armRoutedAllowedIPs(peerCfg *lazyconn.PeerConfig) {
|
||||
if m.reconcileAllowedIPs == nil {
|
||||
return
|
||||
}
|
||||
if err := m.reconcileAllowedIPs(peerCfg.PublicKey); err != nil {
|
||||
peerCfg.Log.Errorf("failed to reconcile routed allowed IPs on wake endpoint: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) shouldDeferIdleForHA(inactivePeers map[string]struct{}, peerID string) bool {
|
||||
m.routesMu.RLock()
|
||||
defer m.routesMu.RUnlock()
|
||||
@@ -577,7 +611,7 @@ func (m *Manager) onPeerInactivityTimedOut(peerIDs map[string]struct{}) {
|
||||
|
||||
mp.peerCfg.Log.Infof("start activity monitor")
|
||||
|
||||
if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
|
||||
if err := m.armActivityListener(*mp.peerCfg); err != nil {
|
||||
mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -61,6 +61,7 @@ type Manager interface {
|
||||
InitialRouteRange() []string
|
||||
SetFirewall(firewall.Manager) error
|
||||
SetDNSForwarderPort(port uint16)
|
||||
ReconcilePeerAllowedIPs(peerKey string) error
|
||||
Stop(stateManager *statemanager.Manager)
|
||||
}
|
||||
|
||||
@@ -232,6 +233,30 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) {
|
||||
)
|
||||
}
|
||||
|
||||
// ReconcilePeerAllowedIPs re-applies every routed allowed IP currently tracked for the peer
|
||||
// onto the WireGuard device. The allowed-IP refcounter only calls its AddFunc (which pushes to
|
||||
// the device) on a prefix's 0->1 transition, so a peer whose device entry was rebuilt without a
|
||||
// matching refcounter change — e.g. a lazy connection cycling through idle->wake, which recreates
|
||||
// the WireGuard peer with the overlay /32 only — ends up missing routed prefixes the refcounter
|
||||
// still considers installed, and nothing retries. Calling this when the peer's WireGuard entry is
|
||||
// (re)created restores convergence. It is add-only and idempotent: AddAllowedIP is update-only, so
|
||||
// prefixes are re-added to an existing peer and an absent peer is left untouched.
|
||||
func (m *DefaultManager) ReconcilePeerAllowedIPs(peerKey string) error {
|
||||
if m.allowedIPsRefCounter == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
return m.allowedIPsRefCounter.ReapplyMatching(
|
||||
func(out string) bool { return out == peerKey },
|
||||
func(prefix netip.Prefix) error {
|
||||
if err := m.wgInterface.AddAllowedIP(peerKey, prefix); err != nil {
|
||||
return fmt.Errorf("add allowed IP %s for peer %s: %w", prefix, peerKey, err)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// Init sets up the routing
|
||||
func (m *DefaultManager) Init() error {
|
||||
m.routeSelector = m.initSelector()
|
||||
|
||||
@@ -112,6 +112,11 @@ func (m *MockManager) SetFirewall(firewall.Manager) error {
|
||||
func (m *MockManager) SetDNSForwarderPort(port uint16) {
|
||||
}
|
||||
|
||||
// ReconcilePeerAllowedIPs mock implementation of ReconcilePeerAllowedIPs from Manager interface
|
||||
func (m *MockManager) ReconcilePeerAllowedIPs(peerKey string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop mock implementation of Stop from Manager interface
|
||||
func (m *MockManager) Stop(stateManager *statemanager.Manager) {
|
||||
if m.StopFunc != nil {
|
||||
|
||||
90
client/internal/routemanager/reconcile_test.go
Normal file
90
client/internal/routemanager/reconcile_test.go
Normal file
@@ -0,0 +1,90 @@
|
||||
//go:build !windows
|
||||
|
||||
package routemanager
|
||||
|
||||
import (
|
||||
"net"
|
||||
"net/netip"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.zx2c4.com/wireguard/tun/netstack"
|
||||
|
||||
"github.com/netbirdio/netbird/client/iface/device"
|
||||
"github.com/netbirdio/netbird/client/iface/wgaddr"
|
||||
"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
|
||||
)
|
||||
|
||||
// reconcileWGMock is a minimal iface.WGIface that only records AddAllowedIP calls; every other
|
||||
// method is an inert stub because ReconcilePeerAllowedIPs exercises none of them.
|
||||
type reconcileWGMock struct {
|
||||
mu sync.Mutex
|
||||
adds map[string][]netip.Prefix
|
||||
}
|
||||
|
||||
func (m *reconcileWGMock) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if m.adds == nil {
|
||||
m.adds = map[string][]netip.Prefix{}
|
||||
}
|
||||
m.adds[peerKey] = append(m.adds[peerKey], allowedIP)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *reconcileWGMock) added(peerKey string) []netip.Prefix {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return m.adds[peerKey]
|
||||
}
|
||||
|
||||
func (m *reconcileWGMock) RemoveAllowedIP(string, netip.Prefix) error { return nil }
|
||||
func (m *reconcileWGMock) Name() string { return "utun-test" }
|
||||
func (m *reconcileWGMock) Address() wgaddr.Address { return wgaddr.Address{} }
|
||||
func (m *reconcileWGMock) ToInterface() *net.Interface { return nil }
|
||||
func (m *reconcileWGMock) IsUserspaceBind() bool { return false }
|
||||
func (m *reconcileWGMock) GetFilter() device.PacketFilter { return nil }
|
||||
func (m *reconcileWGMock) GetDevice() *device.FilteredDevice { return nil }
|
||||
func (m *reconcileWGMock) GetNet() *netstack.Net { return nil }
|
||||
|
||||
// TestReconcilePeerAllowedIPs verifies the declarative reconcile re-applies every routed prefix
|
||||
// tracked for the peer (self-heal, independent of refcount level) and stays scoped to that peer.
|
||||
func TestReconcilePeerAllowedIPs(t *testing.T) {
|
||||
wg := &reconcileWGMock{}
|
||||
m := &DefaultManager{wgInterface: wg}
|
||||
m.allowedIPsRefCounter = refcounter.New[netip.Prefix, string, string](
|
||||
func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil },
|
||||
func(netip.Prefix, string) error { return nil },
|
||||
)
|
||||
|
||||
peerA1 := netip.MustParsePrefix("10.0.0.0/24")
|
||||
peerA2 := netip.MustParsePrefix("10.1.0.0/24")
|
||||
peerB1 := netip.MustParsePrefix("10.2.0.0/24")
|
||||
|
||||
for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} {
|
||||
_, err := m.allowedIPsRefCounter.Increment(prefix, peer)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
// Extra reference: reconcile must still re-apply the prefix even though its refcount never
|
||||
// hit 0 again (the exact case the plain incremental path skips).
|
||||
_, err := m.allowedIPsRefCounter.Increment(peerA1, "peerA")
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NoError(t, m.ReconcilePeerAllowedIPs("peerA"))
|
||||
|
||||
assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, wg.added("peerA"),
|
||||
"reconcile must re-apply all routed prefixes of the peer")
|
||||
assert.Empty(t, wg.added("peerB"), "reconcile must not touch another peer's prefixes")
|
||||
}
|
||||
|
||||
// TestReconcilePeerAllowedIPsNoCounter verifies reconcile is a safe no-op before the refcounter is
|
||||
// set up.
|
||||
func TestReconcilePeerAllowedIPsNoCounter(t *testing.T) {
|
||||
wg := &reconcileWGMock{}
|
||||
m := &DefaultManager{wgInterface: wg}
|
||||
|
||||
require.NoError(t, m.ReconcilePeerAllowedIPs("peerA"))
|
||||
assert.Empty(t, wg.added("peerA"))
|
||||
}
|
||||
@@ -94,6 +94,26 @@ func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) {
|
||||
return ref, ok
|
||||
}
|
||||
|
||||
// ReapplyMatching calls apply for every key whose stored Out satisfies pred, holding the
|
||||
// counter lock for the whole pass. Running apply under the lock keeps it atomic with respect
|
||||
// to Increment/Decrement: a prefix dropped to zero is removed from the map (and had its
|
||||
// RemoveFunc called) before this pass observes it, so a stale key can never be re-applied.
|
||||
// pred and apply are invoked under the lock, so they must not call back into the counter.
|
||||
func (rm *Counter[Key, I, O]) ReapplyMatching(pred func(out O) bool, apply func(key Key) error) error {
|
||||
rm.mu.Lock()
|
||||
defer rm.mu.Unlock()
|
||||
|
||||
var merr *multierror.Error
|
||||
for key, ref := range rm.refCountMap {
|
||||
if pred(ref.Out) {
|
||||
if err := apply(key); err != nil {
|
||||
merr = multierror.Append(merr, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nberrors.FormatErrorOrNil(merr)
|
||||
}
|
||||
|
||||
// Increment increments the reference count for the given key.
|
||||
// If this is the first reference to the key, the AddFunc is called.
|
||||
func (rm *Counter[Key, I, O]) Increment(key Key, in I) (Ref[O], error) {
|
||||
|
||||
47
client/internal/routemanager/refcounter/refcounter_test.go
Normal file
47
client/internal/routemanager/refcounter/refcounter_test.go
Normal file
@@ -0,0 +1,47 @@
|
||||
package refcounter
|
||||
|
||||
import (
|
||||
"net/netip"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestReapplyMatching verifies ReapplyMatching invokes apply for exactly the keys whose stored
|
||||
// Out satisfies the predicate (no duplicates for multiply-referenced keys) — the primitive
|
||||
// ReconcilePeerAllowedIPs relies on to re-apply a single peer's routed prefixes.
|
||||
func TestReapplyMatching(t *testing.T) {
|
||||
rc := New[netip.Prefix, string, string](
|
||||
func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil },
|
||||
func(netip.Prefix, string) error { return nil },
|
||||
)
|
||||
|
||||
peerA1 := netip.MustParsePrefix("10.0.0.0/24")
|
||||
peerA2 := netip.MustParsePrefix("10.1.0.0/24")
|
||||
peerB1 := netip.MustParsePrefix("10.2.0.0/24")
|
||||
|
||||
for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} {
|
||||
_, err := rc.Increment(prefix, peer)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
// a second reference must not make the key applied twice
|
||||
_, err := rc.Increment(peerA1, "peerA")
|
||||
require.NoError(t, err)
|
||||
|
||||
var applied []netip.Prefix
|
||||
err = rc.ReapplyMatching(
|
||||
func(out string) bool { return out == "peerA" },
|
||||
func(key netip.Prefix) error { applied = append(applied, key); return nil },
|
||||
)
|
||||
require.NoError(t, err)
|
||||
assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, applied)
|
||||
|
||||
var none []netip.Prefix
|
||||
err = rc.ReapplyMatching(
|
||||
func(out string) bool { return out == "missing" },
|
||||
func(key netip.Prefix) error { none = append(none, key); return nil },
|
||||
)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, none)
|
||||
}
|
||||
@@ -109,7 +109,7 @@ sequenceDiagram
|
||||
Chk->>Inj: continue
|
||||
Inj->>Inj: inject NetBird identity headers per provider config
|
||||
Inj->>Grd: continue
|
||||
Grd->>Grd: enforce per-provider allowlist (fail-closed backstop)
|
||||
Grd->>Grd: enforce model allowlist
|
||||
Grd->>Up: forward (over WireGuard)
|
||||
Up-->>Resp: response (JSON or SSE stream)
|
||||
Resp->>Resp: parse usage tokens, completion
|
||||
@@ -135,15 +135,6 @@ sequenceDiagram
|
||||
(`redact_pii = settings.RedactPii`). Phones, emails, credit cards,
|
||||
PII names — see `redact.go` for the full set. See
|
||||
[`modules/31-proxy-middleware-builtin.md`](modules/31-proxy-middleware-builtin.md).
|
||||
- The model allowlist is enforced in TWO places. `CheckLLMPolicyLimits`
|
||||
is authoritative: it resolves the policy that governs this
|
||||
(provider, caller-groups) and denies (`deny_code = llm_policy.model_blocked`)
|
||||
when no applicable policy permits the model — so an allowlist scoped to
|
||||
one group/provider never leaks to another, and an un-guardrailed policy
|
||||
is genuinely unrestricted. `llm_guardrail` is a per-provider fail-closed
|
||||
backstop: it only carries an allowlist for a provider every authorising
|
||||
policy restricts, and blocks unknown/undetermined models even when
|
||||
management is unreachable.
|
||||
- SSE streaming requires special handling on the response side; the
|
||||
parser must handle partial chunks without buffering the whole
|
||||
stream. See [`modules/32-proxy-llm-parsers.md`](modules/32-proxy-llm-parsers.md).
|
||||
|
||||
@@ -122,7 +122,7 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
|
||||
| on_request | 1 | `llm_router` | `{"providers":[{id, models[], upstream_*, auth_header_*, allowed_group_ids[]}]}` | **true** |
|
||||
| on_request | 2 | `llm_limit_check` | `{}` | – |
|
||||
| on_request | 3 | `llm_identity_inject` | `{"providers":[{provider_id, header_pair?, json_metadata?, extra_headers?}]}` | **true** |
|
||||
| on_request | 4 | `llm_guardrail` | `{"provider_allowlists"?: {providerID: []model}, "prompt_capture":{enabled,redact_pii}}` | – |
|
||||
| on_request | 4 | `llm_guardrail` | `{"model_allowlist"?, "prompt_capture":{enabled,redact_pii}}` | – |
|
||||
| on_response | 5 | `llm_limit_record` | `{}` (runs LAST at runtime) | – |
|
||||
| on_response | 6 | `cost_meter` | `{}` | – |
|
||||
| on_response | 7 | `llm_response_parser` | `{"capture_completion": <bool>, "redact_pii"?: true}` | – |
|
||||
|
||||
@@ -244,7 +244,7 @@ no mocks. Tests: `TestChain_AllowPath_StampsAttributionAndRecordsCounter`
|
||||
| `llm_router` | `{providers: [{id, models, upstream_scheme, upstream_host, upstream_path?, auth_header_name, auth_header_value, allowed_group_ids}]}` |
|
||||
| `llm_limit_check` | `{}` — pulls `MgmtClient` from `FactoryContext` |
|
||||
| `llm_identity_inject` | `{providers: [{provider_id, header_pair?|json_metadata?, extra_headers?}]}` |
|
||||
| `llm_guardrail` | `{provider_allowlists: {providerID: []string}, prompt_capture: {enabled, redact_pii}}` — allowlist keyed by resolved provider id; a provider absent from the map is unrestricted (fail-closed backstop; authoritative per-policy/group check is management's `CheckLLMPolicyLimits`) |
|
||||
| `llm_guardrail` | `{model_allowlist: []string, prompt_capture: {enabled, redact_pii}}` |
|
||||
| `llm_response_parser` | `{redact_pii?, capture_completion?: *bool}` |
|
||||
| `cost_meter` | `{pricing_path?}` (basename inside data-dir; defaults `pricing.yaml`) |
|
||||
| `llm_limit_record` | `{}` — same pattern as `llm_limit_check` |
|
||||
|
||||
@@ -20,14 +20,15 @@ import (
|
||||
// covers whatever credentials are present (source ~/.llm-keys locally / set the
|
||||
// Actions secrets in CI).
|
||||
type providerCase struct {
|
||||
name string
|
||||
catalogID string
|
||||
upstream string
|
||||
apiKey string
|
||||
model string // body model (chat/messages) or path model@version (vertex)
|
||||
kind string // harness.WireChat, harness.WireMessages, or harness.WireVertex
|
||||
project string // vertex only: GCP project for the rawPredict path
|
||||
region string // vertex only: GCP region for the rawPredict path
|
||||
name string
|
||||
catalogID string
|
||||
upstream string
|
||||
apiKey string
|
||||
model string // body model (chat/messages) or path model@version (vertex)
|
||||
kind string // harness.WireChat, harness.WireMessages, or harness.WireVertex
|
||||
project string // vertex only: GCP project for the rawPredict path
|
||||
region string // vertex only: GCP region for the rawPredict path
|
||||
pathPrefix string // base-URL path prefix the agent carries (e.g. "/anthropic" for Kimi)
|
||||
}
|
||||
|
||||
// availableProviders builds the matrix from the provider env vars that are set.
|
||||
@@ -39,6 +40,24 @@ func availableProviders() []providerCase {
|
||||
if k := os.Getenv("ANTHROPIC_TOKEN"); k != "" {
|
||||
ps = append(ps, providerCase{name: "anthropic", catalogID: "anthropic_api", upstream: "https://api.anthropic.com", apiKey: k, model: "claude-haiku-4-5", kind: harness.WireMessages})
|
||||
}
|
||||
if k := os.Getenv("KIMI_TOKEN"); k != "" {
|
||||
// Kimi (Moonshot AI) serves two body shapes from the same key: OpenAI
|
||||
// Chat Completions on the bare host (/v1/...) and the Anthropic
|
||||
// Messages API under the /anthropic path prefix (the endpoint
|
||||
// Moonshot's Claude Code guide uses). The provider keeps the bare
|
||||
// default upstream and the AGENT carries the /anthropic prefix in
|
||||
// its base URL — exactly the documented Claude Code / Kimi CLI
|
||||
// setup (ANTHROPIC_BASE_URL=https://<endpoint>/anthropic) — so one
|
||||
// provider serves both shapes and the prefix rides through to
|
||||
// Moonshot. Run the Anthropic shape, the flagship Claude Code path;
|
||||
// the OpenAI wire shape is covered live by the other chat-shaped
|
||||
// matrix providers, and Kimi-over-chat passed with kimi-k3 before
|
||||
// the single-model constraint surfaced (run #73 on the kimi feature
|
||||
// branch). The platform serves this account exactly ONE model —
|
||||
// kimi-k3 (kimi-k2-thinking and even kimi-latest return
|
||||
// resource_not_found_error on both surfaces).
|
||||
ps = append(ps, providerCase{name: "kimi", catalogID: "kimi_api", upstream: "https://api.moonshot.ai", apiKey: k, model: "kimi-k3", kind: harness.WireMessages, pathPrefix: "/anthropic"})
|
||||
}
|
||||
if k, u := os.Getenv("VERCEL_TOKEN"), os.Getenv("VERCEL_URL"); k != "" && u != "" {
|
||||
ps = append(ps, providerCase{name: "vercel", catalogID: "vercel_ai_gateway", upstream: u, apiKey: k, model: "openai/gpt-4o-mini", kind: harness.WireChat})
|
||||
}
|
||||
@@ -84,12 +103,18 @@ func availableProviders() []providerCase {
|
||||
}
|
||||
}
|
||||
|
||||
// Bedrock: path-routed, bearer auth. Model is a cross-region inference
|
||||
// profile id (distinct string from the first-party Anthropic case).
|
||||
// Bedrock: path-routed, bearer auth. Model is the FULL cross-region
|
||||
// inference-profile id exactly as AWS issues it — region-family prefix
|
||||
// plus the date/version suffix. A bare or wrong-region id makes Bedrock
|
||||
// reject the request with "The provided model identifier is invalid"
|
||||
// before any inference runs. The proxy normalizes this id to the catalog
|
||||
// key (anthropic.claude-haiku-4-5) for routing/pricing/allowlists.
|
||||
// Defaults pair eu-central-1 with the eu.* profile; AWS_REGION overrides
|
||||
// the region and the prefix follows its family.
|
||||
if k := os.Getenv("AWS_BEARER_TOKEN_BEDROCK"); k != "" {
|
||||
region := os.Getenv("AWS_REGION")
|
||||
if region == "" {
|
||||
region = "us-east-1"
|
||||
region = "eu-central-1"
|
||||
}
|
||||
// A valid Bedrock inference-profile id (region prefix + date + version),
|
||||
// overridable per account. `global.` profiles can be invoked from any
|
||||
@@ -246,7 +271,7 @@ func TestProvidersMatrix(t *testing.T) {
|
||||
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)
|
||||
c, b, cerr = cl.ChatPrefixed(ctx, settings.Endpoint, proxyIP, pc.pathPrefix, pc.kind, pc.model, "Reply with exactly: pong", sessionID)
|
||||
}
|
||||
if cerr == nil {
|
||||
code, body = c, b
|
||||
|
||||
@@ -1,212 +0,0 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// TestGuardrailBlocksUnselectedModel_PathRouted is the end-to-end regression
|
||||
// guard for the customer report that a model-allowlist guardrail attached to a
|
||||
// policy "has no effect" for PATH-ROUTED providers — where the model travels in
|
||||
// the URL, not the JSON body (Google Vertex `…/models/{model}:rawPredict`, AWS
|
||||
// Bedrock `/model/{id}/invoke`). This is the shape the customer actually hits.
|
||||
//
|
||||
// Through the real management API it creates a Vertex and a Bedrock provider
|
||||
// (both catch-all, so the router forwards any model — a 403 can only come from
|
||||
// the guardrail, never model_not_routable) and one guardrail whose allowlist
|
||||
// selects a single model per provider, attached to the policy. Over the tunnel,
|
||||
// for each provider:
|
||||
//
|
||||
// - the SELECTED model (in the URL path) is served (200), and
|
||||
// - an UNSELECTED model (in the URL path) is denied 403 by the guardrail
|
||||
// (llm_policy.model_blocked) before it reaches the upstream.
|
||||
//
|
||||
// Only the upstream LLM is mocked (the vLLM nginx answers any path with 200);
|
||||
// management synth/reconcile, the proxy middleware chain (URL-path model
|
||||
// extraction in llm_request_parser, llm_router, llm_guardrail), and the tunnel
|
||||
// are all real. The guardrail denies before the upstream is dialed, so the mock
|
||||
// cannot influence the block. Providers use a static bearer api key (not
|
||||
// "keyfile::…") so the router injects a static Authorization header instead of
|
||||
// minting a GCP OAuth token — that is the only reason path-routed providers
|
||||
// normally need live credentials, so this test runs with none and is always on.
|
||||
//
|
||||
// If the guardrail were inert for path-routed requests (its middleware dropped
|
||||
// by the proxy, the allowlist not synthesized, or the URL-path model not
|
||||
// extracted so the allowlist can't match), the unselected model would reach the
|
||||
// mock and return 200 — the fail-open this test catches.
|
||||
func TestGuardrailBlocksUnselectedModel_PathRouted(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
vertexProject = "e2e-project"
|
||||
vertexRegion = "global"
|
||||
|
||||
// Vertex catalog ids travel bare in the rawPredict path.
|
||||
vertexSelected = "claude-sonnet-4-5"
|
||||
vertexUnselected = "claude-opus-4-6"
|
||||
|
||||
// Bedrock request ids travel as region-prefixed, versioned
|
||||
// inference-profile ids; the parser normalizes them to the catalog key
|
||||
// the allowlist holds — so this also exercises Bedrock normalization.
|
||||
bedrockSelectedPath = "us.anthropic.claude-sonnet-4-5-v1:0"
|
||||
bedrockSelectedCatalog = "anthropic.claude-sonnet-4-5"
|
||||
bedrockUnselectedPath = "us.anthropic.claude-opus-4-8-v1:0"
|
||||
)
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-path"})
|
||||
require.NoError(t, err, "create group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-guardrail-path-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
// Static bearer (not "keyfile::…") so the router injects a static auth header
|
||||
// instead of minting a GCP token. Both providers are catch-all (no Models)
|
||||
// and point at the mock upstream. Create both up front, enabled, before the
|
||||
// proxy starts (provider changes after connect don't reconcile to the proxy).
|
||||
staticKey := "static-e2e-token"
|
||||
|
||||
vertexProv, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "vertex",
|
||||
ProviderId: "vertex_ai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create vertex provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), vertexProv.Id) })
|
||||
|
||||
bedrockProv, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "bedrock",
|
||||
ProviderId: "bedrock_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
})
|
||||
require.NoError(t, err, "create bedrock provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), bedrockProv.Id) })
|
||||
|
||||
// One guardrail allowlisting exactly the selected model per provider.
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-guardrail-path"
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{vertexSelected, bedrockSelectedCatalog}
|
||||
guard, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail")
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
|
||||
|
||||
enabled := true
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-path",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: []string{vertexProv.Id, bedrockProv.Id},
|
||||
GuardrailIds: &[]string{guard.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-path-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
// Probe first: the GET resolves the endpoint and its first packet wakes the
|
||||
// lazy proxy peer, so WaitProxyPeer then observes it connected.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
// vertexSend / bedrockSend drive one path-routed request and return status+body.
|
||||
vertexSend := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "vertex request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
bedrockSend := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "bedrock request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
|
||||
t.Run("vertex", func(t *testing.T) {
|
||||
// Selected model (in the URL path) is served. Retry to absorb tunnel/DNS
|
||||
// jitter on the first call over the freshly warmed tunnel.
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = vertexSend(vertexSelected)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
assert.Equal(t, 200, code,
|
||||
"selected Vertex model (URL path) must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
|
||||
// Unselected model (in the URL path) must be blocked by the guardrail.
|
||||
code, body = vertexSend(vertexUnselected)
|
||||
assert.Equal(t, 403, code,
|
||||
"unselected Vertex model (URL path) must be denied, not served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"Vertex denial must come from the guardrail allowlist, not routing; body: %s", body)
|
||||
})
|
||||
|
||||
t.Run("bedrock", func(t *testing.T) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = bedrockSend(bedrockSelectedPath)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
assert.Equal(t, 200, code,
|
||||
"selected Bedrock model (URL path, normalized) must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
|
||||
code, body = bedrockSend(bedrockUnselectedPath)
|
||||
assert.Equal(t, 403, code,
|
||||
"unselected Bedrock model (URL path) must be denied, not served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"Bedrock denial must come from the guardrail allowlist, not routing; body: %s", body)
|
||||
})
|
||||
}
|
||||
@@ -1,232 +0,0 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// TestGuardrailMultiPolicyModelAllowlist is the end-to-end regression guard for
|
||||
// the multi-policy interactions of the model-allowlist guardrail: what happens
|
||||
// when several enabled policies govern an account and some carry a guardrail
|
||||
// while others don't. The earlier implementation merged every policy's allowlist
|
||||
// into ONE account-wide union enforced flat on every request, which produced two
|
||||
// defects this test pins:
|
||||
//
|
||||
// - false-ALLOW (cross-group leak): a model allowlisted only for another
|
||||
// group's policy became usable by any caller, because the union ignored
|
||||
// which policy/group actually authorised the request; and
|
||||
// - false-DENY: a policy with NO guardrail (intended unrestricted) still had
|
||||
// its traffic blocked by some other policy's allowlist.
|
||||
//
|
||||
// The fix scopes enforcement to the matched policy/group in management
|
||||
// (SelectPolicyForRequest) with a per-provider fail-closed backstop at the
|
||||
// proxy. The account here has one client in grpMain and three policies over two
|
||||
// catch-all upstreams (the mock vLLM answers any model 200, so a 403 can only be
|
||||
// a policy decision, never a routing miss):
|
||||
//
|
||||
// - polMain : grpMain -> pRestricted, guardrail allowlisting modelSelected
|
||||
// - polOther: grpOther -> pRestricted, guardrail allowlisting modelOther
|
||||
// - polOpen : grpMain -> pOpen, NO guardrail (unrestricted)
|
||||
//
|
||||
// Providers declare their models so routing is deterministic. Over the tunnel,
|
||||
// as the grpMain client:
|
||||
//
|
||||
// - modelSelected on pRestricted is served (200) — allowed by grpMain's policy;
|
||||
// - modelOther on pRestricted is denied 403 (llm_policy.model_blocked) — it is
|
||||
// allowlisted only for grpOther and must NOT leak to grpMain; and
|
||||
// - openModel on pOpen is served (200) — the un-guardrailed policy leaves that
|
||||
// provider unrestricted and must NOT be blocked by another policy's list.
|
||||
//
|
||||
// Under the old account-wide union the middle case returned 200 (the leak) and
|
||||
// the last case returned 403 (the false-deny); both are inverted here.
|
||||
func TestGuardrailMultiPolicyModelAllowlist(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
modelSelected = "e2e-selected"
|
||||
modelOther = "e2e-other"
|
||||
openModel = "e2e-open"
|
||||
)
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-main"})
|
||||
require.NoError(t, err, "create main group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
|
||||
|
||||
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-other"})
|
||||
require.NoError(t, err, "create other group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-guardrail-mp-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpMain.Id}, // client joins grpMain only
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
models := func(ids ...string) *[]api.AgentNetworkProviderModel {
|
||||
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
// pRestricted declares the two guardrailed models so routing is deterministic
|
||||
// (model -> provider). Created first, so it carries the bootstrap cluster.
|
||||
pRestricted, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "restricted",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: models(modelSelected, modelOther),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create restricted provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pRestricted.Id) })
|
||||
|
||||
pOpen, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "open",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: models(openModel),
|
||||
})
|
||||
require.NoError(t, err, "create open provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pOpen.Id) })
|
||||
|
||||
mkGuardrail := func(name, model string) api.AgentNetworkGuardrail {
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, gerr := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, gerr, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gMain := mkGuardrail("e2e-guardrail-mp-main", modelSelected)
|
||||
gOther := mkGuardrail("e2e-guardrail-mp-other", modelOther)
|
||||
|
||||
enabled := true
|
||||
// polMain: grpMain restricted to modelSelected on pRestricted.
|
||||
polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-main",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{pRestricted.Id},
|
||||
GuardrailIds: &[]string{gMain.Id},
|
||||
})
|
||||
require.NoError(t, err, "create main policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
|
||||
|
||||
// polOther: grpOther restricted to modelOther on the SAME provider. The
|
||||
// client is not in grpOther, so modelOther must never be usable by it.
|
||||
polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-other",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpOther.Id},
|
||||
DestinationProviderIds: []string{pRestricted.Id},
|
||||
GuardrailIds: &[]string{gOther.Id},
|
||||
})
|
||||
require.NoError(t, err, "create other policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
|
||||
|
||||
// polOpen: grpMain on pOpen with NO guardrail — unrestricted.
|
||||
polOpen, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-open",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{pOpen.Id},
|
||||
})
|
||||
require.NoError(t, err, "create open policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOpen.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-mp-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
// sendUntil200 absorbs first-call tunnel/DNS jitter on the freshly warmed tunnel.
|
||||
sendUntil200 := func(model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(model)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
t.Run("selected model allowed for its group", func(t *testing.T) {
|
||||
code, body := sendUntil200(modelSelected)
|
||||
assert.Equal(t, 200, code,
|
||||
"grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
|
||||
t.Run("other group's model does not leak", func(t *testing.T) {
|
||||
// modelOther is allowlisted only for grpOther. The grpMain client must be
|
||||
// denied by management's per-policy/group check — not waved through by an
|
||||
// account-wide union. This is the security-critical wrong-ALLOW guard.
|
||||
code, body := send(modelOther)
|
||||
assert.Equal(t, 403, code,
|
||||
"another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"denial must be a model-allowlist decision; body: %s", body)
|
||||
})
|
||||
|
||||
t.Run("unguarded policy leaves its provider unrestricted", func(t *testing.T) {
|
||||
// polOpen carries no guardrail, so pOpen is unrestricted for grpMain. The
|
||||
// old account-wide union would have blocked openModel (it is on no
|
||||
// allowlist); it must now be served — the false-DENY guard.
|
||||
code, body := sendUntil200(openModel)
|
||||
assert.Equal(t, 200, code,
|
||||
"an un-guardrailed policy's provider must not be blocked by another policy's allowlist; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
}
|
||||
@@ -52,7 +52,9 @@ func catalogModel(pc providerCase) string {
|
||||
func disallowedModel(pc providerCase) string {
|
||||
switch pc.kind {
|
||||
case harness.WireBedrock:
|
||||
return "us.anthropic.claude-opus-4-8"
|
||||
// Same profile prefix as the allowed model so only the model name
|
||||
// differs; the guardrail must deny it before it reaches AWS.
|
||||
return strings.SplitN(pc.model, ".", 2)[0] + ".anthropic.claude-opus-4-8"
|
||||
case harness.WireVertex:
|
||||
return "claude-opus-4-8@20250101"
|
||||
default:
|
||||
@@ -72,7 +74,7 @@ func sendModel(ctx context.Context, t *testing.T, cl *harness.Client, endpoint,
|
||||
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", "")
|
||||
code, _, err = cl.ChatPrefixed(ctx, endpoint, proxyIP, pc.pathPrefix, pc.kind, model, "Reply with exactly: pong", "")
|
||||
}
|
||||
require.NoError(t, err, "request must reach the proxy for %s", pc.name)
|
||||
return code
|
||||
|
||||
@@ -239,6 +239,17 @@ const (
|
||||
// the wire shape: WireChat (OpenAI) or WireMessages (Anthropic). A non-empty
|
||||
// sessionID is sent as the universal x-session-id header the proxy records.
|
||||
func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prompt, sessionID string) (int, string, error) {
|
||||
return cl.ChatPrefixed(ctx, endpoint, proxyIP, "", kind, model, prompt, sessionID)
|
||||
}
|
||||
|
||||
// ChatPrefixed is Chat with a base-URL path prefix prepended to the wire
|
||||
// path, mirroring agents whose base URL carries a shape-selecting prefix that
|
||||
// rides through to the upstream — e.g. Claude Code against a Kimi provider
|
||||
// sets ANTHROPIC_BASE_URL=https://<endpoint>/anthropic so the proxy forwards
|
||||
// /anthropic/v1/messages to Moonshot's Anthropic surface while the provider's
|
||||
// upstream URL stays the bare https://api.moonshot.ai. Empty prefix is plain
|
||||
// Chat.
|
||||
func (cl *Client) ChatPrefixed(ctx context.Context, endpoint, proxyIP, pathPrefix, kind, model, prompt, sessionID string) (int, string, error) {
|
||||
var path, body string
|
||||
var headers []string
|
||||
switch kind {
|
||||
@@ -250,7 +261,7 @@ func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prom
|
||||
path = "/v1/chat/completions"
|
||||
body = fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":%q}]}`, model, prompt)
|
||||
}
|
||||
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(headers, sessionID))
|
||||
return cl.post(ctx, endpoint, proxyIP, pathPrefix+path, body, withSessionID(headers, sessionID))
|
||||
}
|
||||
|
||||
// Vertex issues an Anthropic-on-Vertex rawPredict POST over the tunnel. Unlike
|
||||
|
||||
@@ -234,9 +234,6 @@ init_environment() {
|
||||
|
||||
NETBIRD_LICENSE_KEY=$(read_secret "Enter license key (input hidden)")
|
||||
|
||||
GHCR_USERNAME="netbirdExtAccess1"
|
||||
GHCR_TOKEN=$(read_secret "Enter GHCR token (input hidden)")
|
||||
|
||||
POSTGRES_USER="netbird"
|
||||
POSTGRES_DB="netbird"
|
||||
POSTGRES_PASSWORD=$(rand_secret)
|
||||
@@ -263,10 +260,6 @@ init_environment() {
|
||||
install -m 600 /dev/null config.yaml
|
||||
render_config_yaml >> config.yaml
|
||||
|
||||
echo "Logging in to ghcr.io ..."
|
||||
printf '%s' "$GHCR_TOKEN" | docker login ghcr.io -u "$GHCR_USERNAME" --password-stdin
|
||||
unset GHCR_TOKEN
|
||||
|
||||
echo ""
|
||||
echo "Pulling images ..."
|
||||
$DOCKER_COMPOSE_COMMAND pull
|
||||
|
||||
@@ -490,8 +490,6 @@ init_migration() {
|
||||
echo ""
|
||||
echo "Step 1: Image swap (community → Enterprise). License key required."
|
||||
NB_LICENSE_KEY=$(read_secret " License key")
|
||||
GHCR_USERNAME="netbirdExtAccess1"
|
||||
GHCR_TOKEN=$(read_secret " GHCR token (input hidden)")
|
||||
|
||||
# Step 2 — optional
|
||||
echo ""
|
||||
@@ -588,11 +586,6 @@ apply_changes() {
|
||||
fi
|
||||
} >> "$ENV_FILE"
|
||||
|
||||
echo ""
|
||||
echo "Logging in to ghcr.io ..."
|
||||
printf '%s' "$GHCR_TOKEN" | docker login ghcr.io -u "$GHCR_USERNAME" --password-stdin
|
||||
unset GHCR_TOKEN
|
||||
|
||||
echo ""
|
||||
echo "Pulling enterprise images ..."
|
||||
$DOCKER_COMPOSE_COMMAND pull
|
||||
|
||||
@@ -420,6 +420,47 @@ var providers = []Provider{
|
||||
{ID: "mistral-embed", Label: "Mistral Embed", InputPer1k: 0.0001, OutputPer1k: 0, ContextWindow: 8192},
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: "kimi_api",
|
||||
Kind: KindProvider,
|
||||
Name: "Kimi (Moonshot AI) API",
|
||||
Description: "Kimi K3 / K2 models via the Moonshot AI platform",
|
||||
DefaultHost: "api.moonshot.ai",
|
||||
AuthHeaderName: "Authorization",
|
||||
AuthHeaderTemplate: "Bearer ${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#1A1A2E",
|
||||
// ParserID empty on purpose: Moonshot serves two body shapes on
|
||||
// the same host and key, and the proxy's URL sniffer dispatches
|
||||
// both (same pattern as Bifrost). /v1/chat/completions matches
|
||||
// OpenAIParser; the Anthropic-compatible endpoint the official
|
||||
// Claude Code guide uses (/anthropic/v1/messages) contains
|
||||
// "/v1/messages" and matches AnthropicParser. Pinning "openai"
|
||||
// here would misparse the Claude Code path — the primary way
|
||||
// teams consume Kimi for coding today. Both endpoints accept the
|
||||
// same Moonshot key via Authorization: Bearer (Claude Code's
|
||||
// ANTHROPIC_AUTH_TOKEN rides that header too).
|
||||
//
|
||||
// api.moonshot.ai is the international platform; mainland-China
|
||||
// accounts live on api.moonshot.cn with separate billing —
|
||||
// operators there override the host on the provider record. The
|
||||
// kimi.com subscription coding endpoint (api.kimi.com/coding,
|
||||
// model id "k3") is account-bound seat licensing rather than a
|
||||
// meterable platform key, so it's deliberately not the default.
|
||||
ParserID: "",
|
||||
// Pricing per Moonshot's platform rates at K3 launch (July 2026):
|
||||
// $3/$15 per MTok with $0.30 cached input, flat across the 1M-token
|
||||
// window. kimi-k3 is the ONLY model the platform serves newer
|
||||
// accounts — K2-era ids (kimi-k2-thinking) and even the kimi-latest
|
||||
// alias return resource_not_found_error, verified live 2026-07-21 —
|
||||
// so it's the only catalog entry. Grandfathered accounts with K2
|
||||
// access can still type those ids on the provider's model rows.
|
||||
// The consumer app's "K3 Swarm Max" mode is not an API SKU, so it
|
||||
// doesn't appear here.
|
||||
Models: []Model{
|
||||
{ID: "kimi-k3", Label: "Kimi K3", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000},
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: "litellm_proxy",
|
||||
Kind: KindGateway,
|
||||
|
||||
@@ -86,20 +86,12 @@ type Manager interface {
|
||||
|
||||
// PolicySelectionInput is the per-request selection envelope. The
|
||||
// proxy populates it from CapturedData (account, user, groups) plus
|
||||
// the provider llm_router resolved and the model it extracted.
|
||||
// the provider llm_router resolved.
|
||||
type PolicySelectionInput struct {
|
||||
AccountID string
|
||||
UserID string
|
||||
GroupIDs []string
|
||||
ProviderID string
|
||||
// Model is the (already-normalised) upstream model identifier the proxy
|
||||
// extracted from the request. The proxy's request parser strips
|
||||
// path-routed provider decoration (Bedrock region/version, Vertex
|
||||
// @version) before emitting it, so a plain case-insensitive compare
|
||||
// against a guardrail allowlist is sufficient here. Empty when the model
|
||||
// could not be determined — treated as "not permitted" by any policy that
|
||||
// restricts models (fail closed), mirroring the proxy guardrail.
|
||||
Model string
|
||||
}
|
||||
|
||||
// PolicySelectionResult names the policy that "pays" for this request
|
||||
|
||||
@@ -5,7 +5,6 @@ import (
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
@@ -36,11 +35,6 @@ const (
|
||||
denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded"
|
||||
//nolint:gosec // account deny code label, not a credential
|
||||
denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded"
|
||||
// denyCodeModelBlocked is returned when policies govern the request's
|
||||
// (provider, caller-groups) but none of them permits the requested model
|
||||
// under its guardrail allowlist. Matches the proxy guardrail's code so the
|
||||
// two enforcement layers surface the same label.
|
||||
denyCodeModelBlocked = "llm_policy.model_blocked"
|
||||
)
|
||||
|
||||
// consumptionCache holds the consumption counters prefetched for one
|
||||
@@ -165,34 +159,6 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec
|
||||
}
|
||||
candidates := filterApplicablePolicies(policies, in)
|
||||
|
||||
// Model-allowlist gate (per-policy, per-group). Among the policies that
|
||||
// authorise this (provider, caller-groups), keep only those whose guardrails
|
||||
// permit the requested model; a policy with no allowlist-enabled guardrail is
|
||||
// unrestricted and always permits. When policies govern this request but none
|
||||
// permits the model, deny. This is the authoritative allowlist decision:
|
||||
// because it is scoped to the matched policies, a model allowlisted for one
|
||||
// group/provider never leaks to another (no account-wide union), and a policy
|
||||
// that carries no guardrail is genuinely unrestricted rather than being caught
|
||||
// by some other policy's allowlist.
|
||||
// Only consult guardrails when at least one candidate policy references one;
|
||||
// a policy with no guardrail is unrestricted, so when none of them carries a
|
||||
// guardrail the gate is a no-op and we skip the store read entirely.
|
||||
if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) {
|
||||
guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID)
|
||||
if gErr != nil {
|
||||
return nil, gErr
|
||||
}
|
||||
permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model)
|
||||
if len(permitted) == 0 {
|
||||
return &PolicySelectionResult{
|
||||
Allow: false,
|
||||
DenyCode: denyCodeModelBlocked,
|
||||
DenyReason: modelBlockedReason(in.Model),
|
||||
}, nil
|
||||
}
|
||||
candidates = permitted
|
||||
}
|
||||
|
||||
// Prefetch every consumption counter the ceiling + candidate policies will
|
||||
// read, in a single store round-trip, then score against the cache.
|
||||
cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now)
|
||||
@@ -284,93 +250,6 @@ func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput)
|
||||
return out
|
||||
}
|
||||
|
||||
// anyPolicyHasGuardrails reports whether any policy references at least one
|
||||
// guardrail, so the selector can skip loading guardrails when none do.
|
||||
func anyPolicyHasGuardrails(policies []*types.Policy) bool {
|
||||
for _, p := range policies {
|
||||
if p != nil && len(p.GuardrailIDs) > 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
|
||||
// model-allowlist gate to resolve each candidate policy's attached guardrails.
|
||||
func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
|
||||
guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list account guardrails: %w", err)
|
||||
}
|
||||
byID := make(map[string]*types.Guardrail, len(guardrails))
|
||||
for _, g := range guardrails {
|
||||
if g != nil {
|
||||
byID[g.ID] = g
|
||||
}
|
||||
}
|
||||
return byID, nil
|
||||
}
|
||||
|
||||
// filterModelPermittedPolicies returns the subset of policies whose guardrails
|
||||
// permit the model. Order is preserved so downstream scoring is unaffected.
|
||||
func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy {
|
||||
out := make([]*types.Policy, 0, len(policies))
|
||||
for _, p := range policies {
|
||||
if policyPermitsModel(p, byID, model) {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// policyPermitsModel reports whether a policy permits the requested model under
|
||||
// its attached guardrails. A policy with no allowlist-enabled guardrail is
|
||||
// unrestricted (permits any model, including an empty/undetermined one). A
|
||||
// policy with one or more allowlist-enabled guardrails permits the model only
|
||||
// when it appears in the union of those allowlists; an empty/undetermined model
|
||||
// never matches a non-empty allowlist, so such a policy fails closed — the same
|
||||
// contract the proxy guardrail enforces for path-routed providers.
|
||||
func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool {
|
||||
if p == nil {
|
||||
return false
|
||||
}
|
||||
wanted := normaliseModelID(model)
|
||||
restricted := false
|
||||
for _, gID := range p.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
|
||||
continue
|
||||
}
|
||||
restricted = true
|
||||
if wanted == "" {
|
||||
continue
|
||||
}
|
||||
for _, allowed := range g.Checks.ModelAllowlist.Models {
|
||||
if normaliseModelID(allowed) == wanted {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return !restricted
|
||||
}
|
||||
|
||||
// normaliseModelID lowercases and trims a model identifier so the allowlist
|
||||
// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's
|
||||
// normaliseModel so both layers agree on what "same model" means.
|
||||
func normaliseModelID(model string) string {
|
||||
return strings.ToLower(strings.TrimSpace(model))
|
||||
}
|
||||
|
||||
// modelBlockedReason builds the human-readable deny reason for a model-allowlist
|
||||
// rejection. The model is quoted when known; an undetermined model is reported
|
||||
// as such so the access log distinguishes "wrong model" from "no model".
|
||||
func modelBlockedReason(model string) string {
|
||||
if normaliseModelID(model) == "" {
|
||||
return "request model could not be determined for the policy allowlist"
|
||||
}
|
||||
return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model)
|
||||
}
|
||||
|
||||
// candidate is the per-policy intermediate the selector ranks. A
|
||||
// policy that's been exhausted on any enabled cap never makes it
|
||||
// into this slice; the selector's deny envelope carries the latest
|
||||
|
||||
@@ -1,226 +0,0 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/golang/mock/gomock"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
)
|
||||
|
||||
// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups
|
||||
// to reach providerID under the given guardrails. Uncapped keeps the selector's
|
||||
// headroom scoring trivial so these tests isolate the model-allowlist gate.
|
||||
func guardedPolicy(id, account string, sourceGroups []string, providerID string, guardrailIDs ...string) *types.Policy {
|
||||
return &types.Policy{
|
||||
ID: id,
|
||||
AccountID: account,
|
||||
Enabled: true,
|
||||
SourceGroups: sourceGroups,
|
||||
DestinationProviderIDs: []string{providerID},
|
||||
GuardrailIDs: guardrailIDs,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
}
|
||||
|
||||
// allowlistGuardrail builds a guardrail whose model allowlist is enabled and
|
||||
// carries the given models.
|
||||
func allowlistGuardrail(id, account string, models ...string) *types.Guardrail {
|
||||
return &types.Guardrail{
|
||||
ID: id,
|
||||
AccountID: account,
|
||||
Checks: types.GuardrailChecks{
|
||||
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: models},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func expectPolicies(mockStore *store.MockStore, account string, policies ...*types.Policy) {
|
||||
mockStore.EXPECT().
|
||||
GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), account).
|
||||
Return(policies, nil)
|
||||
}
|
||||
|
||||
func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...*types.Guardrail) {
|
||||
mockStore.EXPECT().
|
||||
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), account).
|
||||
Return(guardrails, nil)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist
|
||||
// decision: a policy authorises the (provider, group) but restricts the model,
|
||||
// and the requested model isn't on the list, so the request is denied.
|
||||
func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "a model outside the only applicable policy's allowlist must be denied")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode, "deny code must be model_blocked")
|
||||
assert.NotEmpty(t, res.DenyReason, "deny reason must be populated")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_ModelAllowedByAllowlist is the allow counterpart: the model
|
||||
// is on the applicable policy's allowlist, so selection proceeds normally.
|
||||
func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4"))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "a model on the applicable policy's allowlist must be allowed")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_CaseInsensitiveModelMatch proves the compare tolerates case
|
||||
// and surrounding whitespace, matching the proxy guardrail's normalisation.
|
||||
func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", " GPT-4o "))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "gpt-4o",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "case/whitespace variants must match the allowlist entry")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_UnguardedPolicyIsUnrestricted is the false-deny fix: when
|
||||
// two policies authorise the same (provider, group) and one carries no
|
||||
// guardrail, that un-guardrailed policy makes the request unrestricted — it is
|
||||
// NOT caught by the other policy's allowlist.
|
||||
func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
restricted := guardedPolicy("pol-restricted", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail
|
||||
expectPolicies(mockStore, "acc-1", restricted, open)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "an un-guardrailed policy for the same (provider, group) must leave the request unrestricted")
|
||||
assert.Equal(t, "pol-open", res.SelectedPolicyID, "the unrestricted policy must be the one that pays")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups is the false-allow fix: a
|
||||
// model allowlisted only for grp-b must not be usable by a caller in grp-a,
|
||||
// even though both groups' policies target the same provider. The selector only
|
||||
// considers policies applicable to the caller's groups, so grp-b's allowlist
|
||||
// never enters grp-a's decision.
|
||||
func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
polA := guardedPolicy("pol-a", "acc-1", []string{"grp-a"}, "prov-1", "g-a")
|
||||
polB := guardedPolicy("pol-b", "acc-1", []string{"grp-b"}, "prov-1", "g-b")
|
||||
expectPolicies(mockStore, "acc-1", polA, polB)
|
||||
expectGuardrails(mockStore, "acc-1",
|
||||
allowlistGuardrail("g-a", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-b", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-a"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4", // only allowed for grp-b
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "grp-b's allowlisted model must not leak to a grp-a caller")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_UndeterminedModelFailsClosed proves the fail-closed contract
|
||||
// mirrors the proxy: with a restricted applicable policy and an empty model
|
||||
// (e.g. a path-routed shape the parser couldn't map), the request is denied.
|
||||
func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "", // undetermined
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "an undetermined model must fail closed against a restricted policy")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_DisabledAllowlistDoesNotRestrict proves a guardrail whose
|
||||
// model allowlist is disabled imposes no model restriction, even though the
|
||||
// policy references it.
|
||||
func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
disabled := &types.Guardrail{
|
||||
ID: "g-1",
|
||||
AccountID: "acc-1",
|
||||
Checks: types.GuardrailChecks{
|
||||
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}},
|
||||
},
|
||||
}
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", disabled)
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "anything-goes",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "a disabled allowlist must not restrict the model")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
@@ -235,14 +235,7 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
|
||||
mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
|
||||
applyAccountCollectionControls(&mergedGuardrails, settings)
|
||||
// The proxy guardrail is the fail-closed backstop, scoped per provider so it
|
||||
// never applies one policy's allowlist to another provider's traffic. The
|
||||
// authoritative per-policy/group model decision lives in management's
|
||||
// SelectPolicyForRequest; a provider only lands in this map when EVERY policy
|
||||
// authorising it restricts models, so an un-guardrailed policy leaves its
|
||||
// providers unrestricted here.
|
||||
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
|
||||
guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture)
|
||||
guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -787,12 +780,10 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt
|
||||
|
||||
// guardrailConfig is the JSON shape the proxy-side llm_guardrail
|
||||
// middleware expects. Mirrors the proxy registration documented in
|
||||
// the management→proxy contract. provider_allowlists is keyed by the
|
||||
// resolved provider id llm_router stamps; a provider absent from the map is
|
||||
// unrestricted at the proxy layer.
|
||||
// the management→proxy contract.
|
||||
type guardrailConfig struct {
|
||||
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
|
||||
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
|
||||
ModelAllowlist []string `json:"model_allowlist,omitempty"`
|
||||
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
|
||||
}
|
||||
|
||||
type guardrailPromptCapture struct {
|
||||
@@ -837,12 +828,12 @@ func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Se
|
||||
merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii
|
||||
}
|
||||
|
||||
func marshalGuardrailConfig(providerAllowlists map[string][]string, capture MergedPromptCapture) ([]byte, error) {
|
||||
func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
|
||||
cfg := guardrailConfig{
|
||||
ProviderAllowlists: providerAllowlists,
|
||||
ModelAllowlist: merged.ModelAllowlist,
|
||||
PromptCapture: guardrailPromptCapture{
|
||||
Enabled: capture.Enabled,
|
||||
RedactPii: capture.RedactPii,
|
||||
Enabled: merged.PromptCapture.Enabled,
|
||||
RedactPii: merged.PromptCapture.RedactPii,
|
||||
},
|
||||
}
|
||||
out, err := json.Marshal(cfg)
|
||||
@@ -852,81 +843,6 @@ func marshalGuardrailConfig(providerAllowlists map[string][]string, capture Merg
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// buildProviderAllowlists computes the per-provider model allowlist the proxy
|
||||
// guardrail enforces as a fail-closed backstop. A provider is included ONLY when
|
||||
// every enabled policy that authorises it restricts models (has at least one
|
||||
// allowlist-enabled guardrail); its list is then the union of those policies'
|
||||
// allowed models. If any authorising policy leaves models unrestricted, the
|
||||
// provider is omitted entirely — the proxy treats an absent provider as
|
||||
// unrestricted and defers to management's per-policy/group decision, so a
|
||||
// mixed set of policies can neither leak a model across providers nor wrongly
|
||||
// block an un-guardrailed policy's traffic.
|
||||
func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string {
|
||||
type providerAcc struct {
|
||||
models map[string]struct{}
|
||||
anyUnrestricted bool
|
||||
}
|
||||
accs := make(map[string]*providerAcc)
|
||||
for _, p := range policies {
|
||||
if p == nil {
|
||||
continue
|
||||
}
|
||||
restricted, models := policyModelAllowlist(p, byID)
|
||||
for _, providerID := range p.DestinationProviderIDs {
|
||||
if providerID == "" {
|
||||
continue
|
||||
}
|
||||
acc, ok := accs[providerID]
|
||||
if !ok {
|
||||
acc = &providerAcc{models: make(map[string]struct{})}
|
||||
accs[providerID] = acc
|
||||
}
|
||||
if !restricted {
|
||||
acc.anyUnrestricted = true
|
||||
continue
|
||||
}
|
||||
for _, m := range models {
|
||||
acc.models[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
out := make(map[string][]string, len(accs))
|
||||
for providerID, acc := range accs {
|
||||
if acc.anyUnrestricted {
|
||||
continue
|
||||
}
|
||||
models := make([]string, 0, len(acc.models))
|
||||
for m := range acc.models {
|
||||
models = append(models, m)
|
||||
}
|
||||
sort.Strings(models)
|
||||
out[providerID] = models
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// policyModelAllowlist reports whether a policy restricts models (has at least
|
||||
// one allowlist-enabled guardrail) and the union of allowed models across those
|
||||
// guardrails. Models are returned verbatim; the proxy factory lowercases and
|
||||
// trims them at decode time, matching the runtime compare.
|
||||
func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) {
|
||||
restricted := false
|
||||
var models []string
|
||||
for _, gID := range p.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
|
||||
continue
|
||||
}
|
||||
restricted = true
|
||||
for _, m := range g.Checks.ModelAllowlist.Models {
|
||||
if m != "" {
|
||||
models = append(models, m)
|
||||
}
|
||||
}
|
||||
}
|
||||
return restricted, models
|
||||
}
|
||||
|
||||
// buildAccountService composes the per-account gateway Service. The
|
||||
// target carries the noop placeholder URL — the router middleware
|
||||
// rewrites every request to the matched provider's upstream before the
|
||||
@@ -1075,10 +991,11 @@ func unionSourceGroups(policies []*types.Policy) []string {
|
||||
// expectations and is intentionally distinct from
|
||||
// types.GuardrailChecks so we can evolve either side independently.
|
||||
type MergedGuardrails struct {
|
||||
TokenLimits MergedTokenLimits `json:"token_limits"`
|
||||
Budget MergedBudget `json:"budget"`
|
||||
PromptCapture MergedPromptCapture `json:"prompt_capture"`
|
||||
Retention MergedRetention `json:"retention"`
|
||||
ModelAllowlist []string `json:"model_allowlist,omitempty"`
|
||||
TokenLimits MergedTokenLimits `json:"token_limits"`
|
||||
Budget MergedBudget `json:"budget"`
|
||||
PromptCapture MergedPromptCapture `json:"prompt_capture"`
|
||||
Retention MergedRetention `json:"retention"`
|
||||
}
|
||||
|
||||
type MergedTokenLimits struct {
|
||||
@@ -1113,38 +1030,59 @@ type MergedRetention struct {
|
||||
Days int `json:"days"`
|
||||
}
|
||||
|
||||
// mergeGuardrails computes the prompt-capture portion of the effective
|
||||
// guardrail spec, given the referencing policies and the account's guardrail
|
||||
// catalogue. Policy enabled-ness is the caller's responsibility — only enabled
|
||||
// policies should be passed in.
|
||||
// mergeGuardrails computes the effective guardrail spec applied at the
|
||||
// proxy, given the referencing policies and the account's guardrail
|
||||
// catalogue. Policy enabled-ness is the caller's responsibility — only
|
||||
// enabled policies should be passed in.
|
||||
//
|
||||
// The model allowlist is NOT merged here: it is enforced per-policy/group in
|
||||
// management (SelectPolicyForRequest) and shipped to the proxy per-provider via
|
||||
// buildProviderAllowlists, so a single account-wide union would be both
|
||||
// incorrect (leaks a model across groups/providers) and redundant. Token,
|
||||
// budget, and retention likewise live off guardrails now (Policy.Limits and
|
||||
// account Settings), so prompt capture is all that remains to merge.
|
||||
//
|
||||
// Merge rule — prompt capture: enabled if any policy enables it; redact_pii
|
||||
// sticks if any enabling policy turns it on.
|
||||
// Merge rules:
|
||||
// - Model allowlist: union of allowlists across policies that enable it.
|
||||
// - Token / Budget: most-restrictive (min of non-zero caps) per window.
|
||||
// - Prompt capture: enabled if any policy enables it; redact_pii sticks
|
||||
// if any enabling policy turns it on.
|
||||
// - Retention: enabled if any enables it; smallest non-zero days wins.
|
||||
func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails {
|
||||
merged := MergedGuardrails{}
|
||||
allowlist := make(map[string]struct{})
|
||||
allowlistEnabled := false
|
||||
|
||||
for _, policy := range policies {
|
||||
for _, gID := range policy.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil {
|
||||
continue
|
||||
}
|
||||
mergeGuardrail(g, &merged)
|
||||
mergeGuardrail(g, &merged, allowlist, &allowlistEnabled)
|
||||
}
|
||||
}
|
||||
|
||||
if allowlistEnabled {
|
||||
merged.ModelAllowlist = make([]string, 0, len(allowlist))
|
||||
for m := range allowlist {
|
||||
merged.ModelAllowlist = append(merged.ModelAllowlist, m)
|
||||
}
|
||||
sort.Strings(merged.ModelAllowlist)
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
// mergeGuardrail folds a single guardrail's prompt-capture settings into the
|
||||
// running merge: enabled / redact-pii stick once any enabling guardrail turns
|
||||
// them on.
|
||||
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
|
||||
// mergeGuardrail folds a single guardrail's enabled checks into the
|
||||
// running merge: model-allowlist models join the shared set (and flip
|
||||
// allowlistEnabled), and prompt-capture / redact-pii stick once any
|
||||
// enabling guardrail turns them on.
|
||||
//
|
||||
// TokenLimits, Budget, and Retention have moved off guardrails — token
|
||||
// and budget caps now live on the Policy itself (Policy.Limits) and
|
||||
// retention moves to account-level Settings — so they are not merged here.
|
||||
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) {
|
||||
if g.Checks.ModelAllowlist.Enabled {
|
||||
*allowlistEnabled = true
|
||||
for _, m := range g.Checks.ModelAllowlist.Models {
|
||||
if m != "" {
|
||||
allowlist[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
if g.Checks.PromptCapture.Enabled {
|
||||
merged.PromptCapture.Enabled = true
|
||||
if g.Checks.PromptCapture.RedactPii {
|
||||
|
||||
@@ -81,8 +81,8 @@ func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testi
|
||||
require.Len(t, services, 1)
|
||||
|
||||
cfg := decodeServiceGuardrailConfig(t, services[0])
|
||||
assert.Equal(t, map[string][]string{"prov-1": {"gpt-5.4"}}, cfg.ProviderAllowlists,
|
||||
"model allowlist is a pure policy guardrail and must reach the per-provider config")
|
||||
assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist,
|
||||
"model allowlist is a pure policy guardrail and must always reach the config")
|
||||
assert.False(t, cfg.PromptCapture.Enabled,
|
||||
"prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail")
|
||||
}
|
||||
@@ -172,7 +172,7 @@ func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) {
|
||||
require.Len(t, services, 1, "exactly one synth service expected")
|
||||
|
||||
cfg := decodeServiceGuardrailConfig(t, services[0])
|
||||
assert.Empty(t, cfg.ProviderAllowlists, "no guardrail → provider unrestricted (absent from map)")
|
||||
assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist")
|
||||
assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default")
|
||||
assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default")
|
||||
}
|
||||
|
||||
@@ -1,77 +0,0 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
)
|
||||
|
||||
// policyForProviders builds an enabled policy authorising the given providers
|
||||
// under the given guardrails (both optional). Groups are irrelevant to
|
||||
// buildProviderAllowlists, which keys purely on destination provider.
|
||||
func policyForProviders(id string, guardrailIDs []string, providerIDs ...string) *types.Policy {
|
||||
return &types.Policy{
|
||||
ID: id,
|
||||
Enabled: true,
|
||||
DestinationProviderIDs: providerIDs,
|
||||
GuardrailIDs: guardrailIDs,
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildProviderAllowlists(t *testing.T) {
|
||||
byID := map[string]*types.Guardrail{
|
||||
"g-4o": allowlistGuardrail("g-4o", "acc-1", "gpt-4o"),
|
||||
"g-opus": allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"),
|
||||
"g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}},
|
||||
}
|
||||
|
||||
t.Run("all authorising policies restrict yields per-provider union", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", []string{"g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got,
|
||||
"a provider every policy restricts carries the sorted union of their models")
|
||||
})
|
||||
|
||||
t.Run("any un-guardrailed policy leaves the provider unrestricted (omitted)", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", nil, "prov-x"), // no guardrail
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.NotContains(t, got, "prov-x",
|
||||
"a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted")
|
||||
})
|
||||
|
||||
t.Run("a disabled allowlist counts as unrestricted", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-disabled"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.NotContains(t, got, "prov-x",
|
||||
"a policy whose only guardrail has a disabled allowlist is unrestricted")
|
||||
})
|
||||
|
||||
t.Run("providers are isolated from one another", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", []string{"g-opus"}, "prov-y"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model")
|
||||
assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model")
|
||||
})
|
||||
|
||||
t.Run("one policy authorising two providers restricts both", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"])
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-y"])
|
||||
})
|
||||
}
|
||||
@@ -1031,14 +1031,8 @@ func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *t
|
||||
|
||||
var cfg guardrailConfig
|
||||
require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly")
|
||||
// Both policies restrict the same provider, so the proxy-side per-provider
|
||||
// backstop carries the union of their models. This coarse union is
|
||||
// deliberate: it can still let grp-a reach grp-b's model at the proxy layer,
|
||||
// which management's per-policy/group check (SelectPolicyForRequest) is the
|
||||
// one that rejects. The backstop only guarantees nothing outside the
|
||||
// provider's union slips through when management is unreachable.
|
||||
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ProviderAllowlists["prov-1"],
|
||||
"per-provider allowlist union must keep both models")
|
||||
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist,
|
||||
"model allowlist union must keep both models")
|
||||
}
|
||||
|
||||
func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) {
|
||||
|
||||
@@ -285,7 +285,6 @@ func (s *ProxyServiceServer) CheckLLMPolicyLimits(ctx context.Context, req *prot
|
||||
UserID: req.GetUserId(),
|
||||
GroupIDs: req.GetGroupIds(),
|
||||
ProviderID: req.GetProviderId(),
|
||||
Model: req.GetModel(),
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("select policy for request: %v", err)
|
||||
|
||||
@@ -283,8 +283,8 @@ func (a *Account) SynthesizePrivateServiceZones(peerID string) []nbdns.CustomZon
|
||||
// it, adding a single private service would black-hole every
|
||||
// other name under the zone apex.
|
||||
zone = &nbdns.CustomZone{
|
||||
Domain: dns.Fqdn(serviceDomainZone),
|
||||
Records: []nbdns.SimpleRecord{},
|
||||
Domain: dns.Fqdn(serviceDomainZone),
|
||||
Records: []nbdns.SimpleRecord{},
|
||||
NonAuthoritative: true,
|
||||
SearchDomainDisabled: true,
|
||||
}
|
||||
|
||||
@@ -161,6 +161,16 @@ openai:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0
|
||||
|
||||
# Kimi / Moonshot AI (kimi_api) — OpenAI-compatible /v1 endpoint. Moonshot
|
||||
# reports cache hits OpenAI-style when present; cached input is 10% of
|
||||
# input ($0.30 vs $3.00 per MTok). kimi-k3 is the only model the platform
|
||||
# serves newer accounts (K2-era ids and kimi-latest 404), matching the
|
||||
# management catalog.
|
||||
kimi-k3:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cached_input_per_1k: 0.0003
|
||||
|
||||
anthropic:
|
||||
# Claude 4.x family — cache reads ≈10% of input, cache writes ≈125% of input.
|
||||
# Pricing source: Anthropic's current published rates per million tokens,
|
||||
@@ -206,6 +216,20 @@ anthropic:
|
||||
cache_read_per_1k: 0.0001
|
||||
cache_creation_per_1k: 0.00125
|
||||
|
||||
# Kimi / Moonshot AI (kimi_api) via the Anthropic-compatible endpoint
|
||||
# (/anthropic/v1/messages — the official Claude Code setup). Same rates
|
||||
# as the OpenAI-shape entry above. "kimi-k3[1m]" is the model id some
|
||||
# Claude Code guides set for the 1M-context alias; priced identically so
|
||||
# cost metering doesn't silently skip those requests.
|
||||
kimi-k3:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
"kimi-k3[1m]":
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
|
||||
bedrock:
|
||||
# AWS Bedrock model ids, normalised by the request parser (cross-region
|
||||
# inference-profile prefix + version/throughput suffix stripped), e.g.
|
||||
|
||||
@@ -10,22 +10,11 @@ import (
|
||||
)
|
||||
|
||||
// Config is the JSON-decoded shape accepted by the factory. The
|
||||
// runtime path consumes the normalised allowlists; raw config is not
|
||||
// runtime path consumes the normalised allowlist; raw config is not
|
||||
// retained beyond construction.
|
||||
type Config struct {
|
||||
// ProviderAllowlists maps a resolved provider id (the value llm_router
|
||||
// stamps as KeyLLMResolvedProviderID) to that provider's model allowlist. A
|
||||
// provider present here is restricted to the listed models; a provider
|
||||
// absent is unrestricted. The synthesiser only lists a provider when EVERY
|
||||
// policy authorising it enables an allowlist, so a provider reachable by any
|
||||
// un-guardrailed policy is intentionally absent (unrestricted) here and the
|
||||
// precise per-policy/group decision is left to management. Keeping the gate
|
||||
// per-provider — rather than one account-wide union — is what stops a model
|
||||
// allowlisted for one provider from leaking onto another and stops an
|
||||
// un-guardrailed policy's traffic from being blocked by an unrelated
|
||||
// policy's allowlist.
|
||||
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
|
||||
PromptCapture PromptCapture `json:"prompt_capture"`
|
||||
ModelAllowlist []string `json:"model_allowlist"`
|
||||
PromptCapture PromptCapture `json:"prompt_capture"`
|
||||
}
|
||||
|
||||
// PromptCapture toggles the optional prompt capture + redaction step
|
||||
@@ -66,28 +55,20 @@ func isEmptyJSON(raw []byte) bool {
|
||||
}
|
||||
|
||||
// normaliseConfig lowercases and trims allowlist entries so the runtime
|
||||
// match is case-insensitive. Empty entries are dropped. A provider whose
|
||||
// entries all drop out keeps an empty (non-nil) list — an allowlist that
|
||||
// permits nothing — which is the intended "deny every model" for that
|
||||
// provider, distinct from the provider being absent (unrestricted).
|
||||
// match is case-insensitive. Empty entries are dropped.
|
||||
func normaliseConfig(cfg Config) Config {
|
||||
if len(cfg.ProviderAllowlists) == 0 {
|
||||
cfg.ProviderAllowlists = nil
|
||||
if len(cfg.ModelAllowlist) == 0 {
|
||||
return cfg
|
||||
}
|
||||
cleaned := make(map[string][]string, len(cfg.ProviderAllowlists))
|
||||
for provider, models := range cfg.ProviderAllowlists {
|
||||
list := make([]string, 0, len(models))
|
||||
for _, entry := range models {
|
||||
n := normaliseModel(entry)
|
||||
if n == "" {
|
||||
continue
|
||||
}
|
||||
list = append(list, n)
|
||||
cleaned := make([]string, 0, len(cfg.ModelAllowlist))
|
||||
for _, entry := range cfg.ModelAllowlist {
|
||||
n := normaliseModel(entry)
|
||||
if n == "" {
|
||||
continue
|
||||
}
|
||||
cleaned[provider] = list
|
||||
cleaned = append(cleaned, n)
|
||||
}
|
||||
cfg.ProviderAllowlists = cleaned
|
||||
cfg.ModelAllowlist = cleaned
|
||||
return cfg
|
||||
}
|
||||
|
||||
|
||||
@@ -83,9 +83,8 @@ func (m *Middleware) MutationsSupported() bool { return false }
|
||||
// prompt capture only affects the metadata emitted alongside an allow.
|
||||
func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
|
||||
model, modelPresent := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
|
||||
providerID, _ := lookupMetadata(in.Metadata, middleware.KeyLLMResolvedProviderID)
|
||||
|
||||
if denial := m.evaluateAllowlist(providerID, model, modelPresent); denial != nil {
|
||||
if denial := m.evaluateAllowlist(model, modelPresent); denial != nil {
|
||||
return denial, nil
|
||||
}
|
||||
|
||||
@@ -111,38 +110,20 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
|
||||
// is a no-op.
|
||||
func (m *Middleware) Close() error { return nil }
|
||||
|
||||
// evaluateAllowlist returns a deny Output when the allowlist for the request's
|
||||
// resolved provider rejects the model. A nil return means the request should
|
||||
// proceed. The allowlist is scoped to the provider llm_router resolved: a
|
||||
// provider absent from the config is unrestricted, so an un-guardrailed policy's
|
||||
// traffic is never caught by an unrelated provider's allowlist.
|
||||
func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bool) *middleware.Output {
|
||||
if len(m.cfg.ProviderAllowlists) == 0 {
|
||||
// evaluateAllowlist returns a deny Output when the configured allowlist
|
||||
// rejects the model. A nil return means the request should proceed.
|
||||
func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middleware.Output {
|
||||
if len(m.cfg.ModelAllowlist) == 0 {
|
||||
return nil
|
||||
}
|
||||
// Restrictions exist for this account but the resolved provider is unknown,
|
||||
// so we cannot tell whether this request is bound for a restricted provider.
|
||||
// Fail closed rather than wave it through. In practice llm_router always
|
||||
// stamps the resolved provider before the guardrail runs (or denies first),
|
||||
// so this is a defensive guard, not the common path.
|
||||
if providerID == "" {
|
||||
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
|
||||
}
|
||||
allowlist, restricted := m.cfg.ProviderAllowlists[providerID]
|
||||
if !restricted {
|
||||
// This provider has no allowlist (some authorising policy left it
|
||||
// unrestricted); management owns any per-policy/group decision.
|
||||
return nil
|
||||
}
|
||||
// Fail closed: with an allowlist in effect for this provider, a request whose
|
||||
// model the 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.
|
||||
// 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 modelInAllowlist(allowlist, model) {
|
||||
if m.modelInAllowlist(model) {
|
||||
return nil
|
||||
}
|
||||
return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
|
||||
@@ -170,15 +151,14 @@ func denyModel(model, code, message, reason string) *middleware.Output {
|
||||
}
|
||||
}
|
||||
|
||||
// modelInAllowlist reports whether the model matches any entry in the supplied
|
||||
// (already-normalised) allowlist under the case-insensitive, trim-tolerant
|
||||
// comparison rule.
|
||||
func modelInAllowlist(allowlist []string, model string) bool {
|
||||
// modelInAllowlist reports whether the model matches any allowlist
|
||||
// entry under the case-insensitive, trim-tolerant comparison rule.
|
||||
func (m *Middleware) modelInAllowlist(model string) bool {
|
||||
normalised := normaliseModel(model)
|
||||
if normalised == "" {
|
||||
return false
|
||||
}
|
||||
for _, allowed := range allowlist {
|
||||
for _, allowed := range m.cfg.ModelAllowlist {
|
||||
if allowed == normalised {
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -26,25 +26,6 @@ func newInput(meta ...middleware.KV) *middleware.Input {
|
||||
return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: meta}
|
||||
}
|
||||
|
||||
const (
|
||||
testProvider = "prov-1"
|
||||
otherProvider = "prov-2"
|
||||
)
|
||||
|
||||
// providerCfg builds a Config restricting testProvider to the given models.
|
||||
func providerCfg(models ...string) Config {
|
||||
return Config{ProviderAllowlists: map[string][]string{testProvider: models}}
|
||||
}
|
||||
|
||||
// newInputProvider builds an input that carries a resolved provider id (as
|
||||
// llm_router would stamp) plus any extra metadata.
|
||||
func newInputProvider(provider string, meta ...middleware.KV) *middleware.Input {
|
||||
all := make([]middleware.KV, 0, len(meta)+1)
|
||||
all = append(all, middleware.KV{Key: middleware.KeyLLMResolvedProviderID, Value: provider})
|
||||
all = append(all, meta...)
|
||||
return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: all}
|
||||
}
|
||||
|
||||
func TestMiddlewareIdentity(t *testing.T) {
|
||||
mw := New(Config{})
|
||||
assert.Equal(t, ID, mw.ID(), "middleware ID must be llm_guardrail")
|
||||
@@ -66,12 +47,12 @@ func TestMiddlewareIdentity(t *testing.T) {
|
||||
|
||||
func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
|
||||
mw := New(Config{})
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "no provider allowlists must allow any model")
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "empty allowlist must allow any model")
|
||||
v, ok := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision)
|
||||
require.True(t, ok, "decision metadata must be emitted")
|
||||
assert.Equal(t, "allow", v, "decision must be allow")
|
||||
@@ -81,8 +62,8 @@ func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAllowlistMatchAllows(t *testing.T) {
|
||||
mw := New(providerCfg("gpt-4o", "claude-opus-4"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
mw := New(Config{ModelAllowlist: []string{"gpt-4o", "claude-opus-4"}})
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
@@ -90,8 +71,8 @@ func TestAllowlistMatchAllows(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAllowlistMissDenies(t *testing.T) {
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
@@ -110,10 +91,10 @@ func TestAllowlistMissDenies(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAllowlistCaseInsensitive(t *testing.T) {
|
||||
mw := New(providerCfg(" GPT-4o ", "Claude-OPUS-4"))
|
||||
mw := New(Config{ModelAllowlist: []string{" GPT-4o ", "Claude-OPUS-4"}})
|
||||
cases := []string{"gpt-4o", "GPT-4O", " claude-opus-4 "}
|
||||
for _, model := range cases {
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: model},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
@@ -122,15 +103,14 @@ func TestAllowlistCaseInsensitive(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAllowlistMissingModelKeyDenies(t *testing.T) {
|
||||
// Fail closed: with an allowlist in effect for the resolved provider, a
|
||||
// request whose model the parser could not extract (URL/path-routed
|
||||
// providers such as Bedrock or Vertex whose shape wasn't recognised) must be
|
||||
// denied, not allowed.
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider))
|
||||
// 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)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when the provider is restricted")
|
||||
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")
|
||||
@@ -142,72 +122,26 @@ func TestAllowlistMissingModelKeyDenies(t *testing.T) {
|
||||
|
||||
func TestAllowlistEmptyModelValueDenies(t *testing.T) {
|
||||
// A present-but-empty model is as undeterminable as an absent one.
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
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 the provider is restricted")
|
||||
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 any provider allowlists there is nothing to enforce, so a missing
|
||||
// model is still allowed — the fail-closed rule only applies when a
|
||||
// restriction is in effect.
|
||||
// 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 TestUnrestrictedProviderAllowsAnyModel(t *testing.T) {
|
||||
// testProvider is restricted, but the request resolved to otherProvider,
|
||||
// which carries no allowlist. Its traffic must not be caught by the other
|
||||
// provider's restriction — this is the cross-provider-leak / false-deny
|
||||
// guard. Management owns any per-policy/group decision for otherProvider.
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(otherProvider,
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "an unrestricted provider must not inherit another provider's allowlist")
|
||||
}
|
||||
|
||||
func TestPerProviderAllowlistsAreIsolated(t *testing.T) {
|
||||
// gpt-4o is allowed only on testProvider; claude-opus-4 only on
|
||||
// otherProvider. A model allowlisted for one provider must not be usable on
|
||||
// the other — the fail-closed layer never unions allowlists across providers.
|
||||
mw := New(Config{ProviderAllowlists: map[string][]string{
|
||||
testProvider: {"gpt-4o"},
|
||||
otherProvider: {"claude-opus-4"},
|
||||
}})
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "claude-opus-4 is allowed only on otherProvider, not testProvider")
|
||||
require.NotNil(t, out.DenyReason)
|
||||
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "cross-provider model must be blocked, not model_unknown")
|
||||
}
|
||||
|
||||
func TestRestrictionsButNoResolvedProviderFailsClosed(t *testing.T) {
|
||||
// Restrictions exist for the account but the resolved provider id is absent,
|
||||
// so the request cannot be scoped to a provider. Fail closed rather than
|
||||
// wave it through.
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "missing resolved provider must fail closed when restrictions exist")
|
||||
require.NotNil(t, out.DenyReason)
|
||||
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
|
||||
}
|
||||
|
||||
func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) {
|
||||
mw := New(Config{})
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
@@ -283,8 +217,8 @@ func TestFactoryAcceptsZeroConfigs(t *testing.T) {
|
||||
|
||||
func TestFactoryDecodesValidConfig(t *testing.T) {
|
||||
cfg := Config{
|
||||
ProviderAllowlists: map[string][]string{testProvider: {"gpt-4o"}},
|
||||
PromptCapture: PromptCapture{Enabled: true, RedactPii: true},
|
||||
ModelAllowlist: []string{"gpt-4o"},
|
||||
PromptCapture: PromptCapture{Enabled: true, RedactPii: true},
|
||||
}
|
||||
raw, err := json.Marshal(cfg)
|
||||
require.NoError(t, err, "marshalling test config must succeed")
|
||||
@@ -300,15 +234,15 @@ func TestFactoryRejectsMalformedJSON(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestFactoryNormalisesAllowlist(t *testing.T) {
|
||||
raw := []byte(`{"provider_allowlists":{"prov-1":[" GPT-4o ","",""," Claude-3 "]}}`)
|
||||
raw := []byte(`{"model_allowlist":[" GPT-4o ","",""," Claude-3 "]}`)
|
||||
mw, err := Factory{}.New(raw)
|
||||
require.NoError(t, err)
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "factory must lowercase + trim allowlist entries")
|
||||
out2, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
out2, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-3"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -25,17 +25,10 @@ func runParserGuardrail(t *testing.T, url string, body []byte, allowlist []strin
|
||||
})
|
||||
require.NoError(t, err, "parser must not error")
|
||||
|
||||
const providerID = "prov-under-test"
|
||||
guard := llm_guardrail.New(llm_guardrail.Config{
|
||||
ProviderAllowlists: map[string][]string{providerID: allowlist},
|
||||
})
|
||||
// The real chain has llm_router stamp the resolved provider id before the
|
||||
// guardrail runs; the parser doesn't, so add it here so the guardrail can
|
||||
// scope the allowlist to this provider.
|
||||
meta := append([]middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: providerID}}, parsed.Metadata...)
|
||||
guard := llm_guardrail.New(llm_guardrail.Config{ModelAllowlist: allowlist})
|
||||
out, err := guard.Invoke(context.Background(), &middleware.Input{
|
||||
Slot: middleware.SlotOnRequest,
|
||||
Metadata: meta,
|
||||
Metadata: parsed.Metadata,
|
||||
})
|
||||
require.NoError(t, err, "guardrail must not error")
|
||||
require.NotNil(t, out, "guardrail must return an output")
|
||||
|
||||
@@ -3,8 +3,8 @@
|
||||
# FreeBSD Port Diff Generator for NetBird
|
||||
#
|
||||
# This script generates the diff file required for submitting a FreeBSD port update.
|
||||
# It works on macOS, Linux, and FreeBSD by fetching files from FreeBSD cgit and
|
||||
# computing checksums from the Go module proxy.
|
||||
# It works on macOS, Linux, and FreeBSD by fetching files from the FreeBSD ports
|
||||
# GitHub mirror and computing checksums from the Go module proxy.
|
||||
#
|
||||
# Usage: ./freebsd-port-diff.sh [new_version]
|
||||
# Example: ./freebsd-port-diff.sh 0.60.7
|
||||
@@ -14,7 +14,7 @@
|
||||
set -e
|
||||
|
||||
GITHUB_REPO="netbirdio/netbird"
|
||||
PORTS_CGIT_BASE="https://cgit.freebsd.org/ports/plain/security/netbird"
|
||||
PORTS_MIRROR_BASE="https://raw.githubusercontent.com/freebsd/freebsd-ports/main/security/netbird"
|
||||
GO_PROXY="https://proxy.golang.org/github.com/netbirdio/netbird/@v"
|
||||
OUTPUT_DIR="${OUTPUT_DIR:-.}"
|
||||
AWK_FIRST_FIELD='{print $1}'
|
||||
@@ -30,10 +30,17 @@ fetch_all_tags() {
|
||||
|
||||
fetch_current_ports_version() {
|
||||
echo "Fetching current version from FreeBSD ports..." >&2
|
||||
curl -sL "${PORTS_CGIT_BASE}/Makefile" 2>/dev/null | \
|
||||
local makefile version
|
||||
makefile=$(fetch_ports_file "Makefile") || return 1
|
||||
version=$(echo "$makefile" | \
|
||||
grep -E "^DISTVERSION=" | \
|
||||
sed 's/DISTVERSION=[[:space:]]*//' | \
|
||||
tr -d '\t '
|
||||
tr -d '\t ')
|
||||
if [[ -z "$version" ]]; then
|
||||
echo "Error: Could not extract DISTVERSION from ports Makefile" >&2
|
||||
return 1
|
||||
fi
|
||||
echo "$version"
|
||||
return 0
|
||||
}
|
||||
|
||||
@@ -45,7 +52,16 @@ fetch_latest_github_release() {
|
||||
|
||||
fetch_ports_file() {
|
||||
local filename="$1"
|
||||
curl -sL "${PORTS_CGIT_BASE}/${filename}" 2>/dev/null
|
||||
local content
|
||||
if ! content=$(curl -fsL --proto '=https' --proto-redir '=https' --retry 3 "${PORTS_MIRROR_BASE}/${filename}" 2>/dev/null); then
|
||||
echo "Error: Could not fetch ${filename} from ${PORTS_MIRROR_BASE}" >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ "$content" == \<* ]]; then
|
||||
echo "Error: Received HTML instead of ${filename} from ${PORTS_MIRROR_BASE}" >&2
|
||||
return 1
|
||||
fi
|
||||
printf '%s' "$content"
|
||||
return 0
|
||||
}
|
||||
|
||||
|
||||
@@ -9,18 +9,22 @@
|
||||
# Example: ./freebsd-port-issue-body.sh 0.56.0 0.59.1
|
||||
#
|
||||
# If no versions are provided, the script will:
|
||||
# - Fetch OLD version from FreeBSD ports cgit (current version in ports tree)
|
||||
# - Fetch OLD version from the FreeBSD ports GitHub mirror (current version in ports tree)
|
||||
# - Fetch NEW version from latest NetBird GitHub release tag
|
||||
|
||||
set -e
|
||||
|
||||
GITHUB_REPO="netbirdio/netbird"
|
||||
PORTS_CGIT_URL="https://cgit.freebsd.org/ports/plain/security/netbird/Makefile"
|
||||
PORTS_MAKEFILE_URL="https://raw.githubusercontent.com/freebsd/freebsd-ports/main/security/netbird/Makefile"
|
||||
|
||||
fetch_current_ports_version() {
|
||||
echo "Fetching current version from FreeBSD ports..." >&2
|
||||
local makefile_content
|
||||
makefile_content=$(curl -sL "$PORTS_CGIT_URL" 2>/dev/null)
|
||||
makefile_content=$(curl -fsL --proto '=https' --proto-redir '=https' --retry 3 "$PORTS_MAKEFILE_URL" 2>/dev/null) || makefile_content=""
|
||||
if [[ "$makefile_content" == \<* ]]; then
|
||||
echo "Error: Received HTML instead of Makefile from ${PORTS_MAKEFILE_URL}" >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ -z "$makefile_content" ]]; then
|
||||
echo "Error: Could not fetch Makefile from FreeBSD ports" >&2
|
||||
return 1
|
||||
|
||||
@@ -187,16 +187,16 @@ func (c *GrpcClient) ready() bool {
|
||||
// Sync wraps the real client's Sync endpoint call and takes care of retries and encryption/decryption of messages
|
||||
// Blocking request. The result will be sent via msgHandler callback function
|
||||
func (c *GrpcClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
|
||||
return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key) error {
|
||||
return c.handleSyncStream(ctx, serverPubKey, sysInfo, msgHandler)
|
||||
return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key, backOff backoff.BackOff) error {
|
||||
return c.handleSyncStream(ctx, serverPubKey, sysInfo, msgHandler, backOff)
|
||||
})
|
||||
}
|
||||
|
||||
// Job wraps the real client's Job endpoint call and takes care of retries and encryption/decryption of messages
|
||||
// Blocking request. The result will be sent via msgHandler callback function
|
||||
func (c *GrpcClient) Job(ctx context.Context, msgHandler func(msg *proto.JobRequest) *proto.JobResponse) error {
|
||||
return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key) error {
|
||||
return c.handleJobStream(ctx, serverPubKey, msgHandler)
|
||||
return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key, backOff backoff.BackOff) error {
|
||||
return c.handleJobStream(ctx, serverPubKey, msgHandler, backOff)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -204,7 +204,7 @@ func (c *GrpcClient) Job(ctx context.Context, msgHandler func(msg *proto.JobRequ
|
||||
// It takes care of retries, connection readiness, and fetching server public key.
|
||||
func (c *GrpcClient) withMgmtStream(
|
||||
ctx context.Context,
|
||||
handler func(ctx context.Context, serverPubKey wgtypes.Key) error,
|
||||
handler func(ctx context.Context, serverPubKey wgtypes.Key, backOff backoff.BackOff) error,
|
||||
) error {
|
||||
backOff := defaultBackoff(ctx)
|
||||
operation := func() error {
|
||||
@@ -224,7 +224,7 @@ func (c *GrpcClient) withMgmtStream(
|
||||
return err
|
||||
}
|
||||
|
||||
return handler(ctx, *serverPubKey)
|
||||
return handler(ctx, *serverPubKey, backOff)
|
||||
}
|
||||
|
||||
err := backoff.Retry(operation, backOff)
|
||||
@@ -239,6 +239,7 @@ func (c *GrpcClient) handleJobStream(
|
||||
ctx context.Context,
|
||||
serverPubKey wgtypes.Key,
|
||||
msgHandler func(msg *proto.JobRequest) *proto.JobResponse,
|
||||
backOff backoff.BackOff,
|
||||
) error {
|
||||
ctx, cancelStream := context.WithCancel(ctx)
|
||||
defer cancelStream()
|
||||
@@ -256,6 +257,19 @@ func (c *GrpcClient) handleJobStream(
|
||||
|
||||
log.Debug("job stream handshake sent successfully")
|
||||
|
||||
// The stream is up, so reset the backoff. This matters for two reasons,
|
||||
// both caused by the backoff lib not resetting its state on a successful
|
||||
// connection:
|
||||
// 1. Without a reset, after a connect followed by an error the next retry
|
||||
// starts from the accumulated (large) interval instead of retrying
|
||||
// promptly, delaying reconnection.
|
||||
// 2. Worse, once the accumulated elapsed time exceeds MaxElapsedTime, the
|
||||
// next stream error makes NextBackOff() return Stop, so the retry loop
|
||||
// exits immediately. That error is then mislabeled unrecoverable and
|
||||
// bubbles up to trigger a full engine restart / data-plane teardown
|
||||
// instead of a silent reconnection.
|
||||
backOff.Reset()
|
||||
|
||||
// Main loop: receive, process, respond
|
||||
for {
|
||||
jobReq, err := c.receiveJobRequest(ctx, stream, serverPubKey)
|
||||
@@ -371,7 +385,7 @@ func (c *GrpcClient) sendJobResponse(
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.Key, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
|
||||
func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.Key, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error, backOff backoff.BackOff) error {
|
||||
ctx, cancelStream := context.WithCancel(ctx)
|
||||
defer cancelStream()
|
||||
|
||||
@@ -390,6 +404,19 @@ func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.
|
||||
c.notifyConnected()
|
||||
c.setSyncStreamConnected()
|
||||
|
||||
// The stream is up, so reset the backoff. This matters for two reasons,
|
||||
// both caused by the backoff lib not resetting its state on a successful
|
||||
// connection:
|
||||
// 1. Without a reset, after a connect followed by an error the next retry
|
||||
// starts from the accumulated (large) interval instead of retrying
|
||||
// promptly, delaying reconnection.
|
||||
// 2. Worse, once the accumulated elapsed time exceeds MaxElapsedTime, the
|
||||
// next stream error makes NextBackOff() return Stop, so the retry loop
|
||||
// exits immediately. That error is then mislabeled unrecoverable and
|
||||
// bubbles up to trigger a full engine restart / data-plane teardown
|
||||
// instead of a silent reconnection.
|
||||
backOff.Reset()
|
||||
|
||||
// blocking until error
|
||||
err = c.receiveUpdatesEvents(stream, serverPubKey, msgHandler)
|
||||
if err != nil {
|
||||
|
||||
@@ -21,9 +21,9 @@ import (
|
||||
"net/netip"
|
||||
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/shared/management/types"
|
||||
nbroute "github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
"github.com/netbirdio/netbird/shared/management/types"
|
||||
"github.com/netbirdio/netbird/shared/netiputil"
|
||||
"github.com/netbirdio/netbird/shared/sshauth"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user