mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-23 07:09:08 +02:00
[management, proxy] scope agent-network model allowlist per policy/group and provider (#6905)
## Describe your changes The model-allowlist guardrail was merged into one account-wide union and enforced flat on every request, ignoring which policy/group/provider authorised it. With multiple policies — especially a mix of guardrailed and un-guardrailed ones — this caused: - **false-allow**: a model allowlisted for one group/provider leaked to any caller; and - **false-deny**: an un-guardrailed policy (intended unrestricted) was blocked by another policy's allowlist. Enforcement is now policy/group-aware, mirroring `llm_limit_check`: - **Management (`SelectPolicyForRequest`) is authoritative.** It uses the request model (already carried in `CheckLLMPolicyLimitsRequest.model`, previously ignored) to keep only applicable policies whose guardrails permit the model; no allowlist-enabled guardrail = unrestricted. Denies `llm_policy.model_blocked` when policies govern the (provider, groups) but none permits the model. - **Proxy `llm_guardrail` becomes a per-provider fail-closed backstop.** The synthesiser emits an allowlist only for providers every authorising policy restricts; the middleware keys off the resolved provider id and keeps unknown-model fail-closed.
This commit is contained in:
Maycon Santos
@@ -0,0 +1,209 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// pathRoutedGuardrailCase is one provider's self-contained scenario: its own
|
||||
// provider, its own guardrail whose allowlist holds ONLY that provider's
|
||||
// allowed model, and its own policy. Each case runs in isolation (its own
|
||||
// proxy + client), so the guardrail the proxy enforces contains exactly this
|
||||
// provider's model — never a mixed cross-provider list.
|
||||
type pathRoutedGuardrailCase struct {
|
||||
name string
|
||||
catalogID string // agent-network catalog provider id
|
||||
wire string // harness.WireVertex | harness.WireBedrock
|
||||
allowEntry string // the single model id put on the guardrail allowlist
|
||||
allowModel string // model id sent that MUST be served (200)
|
||||
blockModel string // model id sent that MUST be denied (403 model_blocked)
|
||||
}
|
||||
|
||||
// TestGuardrailBlocksUnselectedModel_PathRouted is the end-to-end regression
|
||||
// guard for the customer report that a model-allowlist guardrail attached to a
|
||||
// policy has no effect for PATH-ROUTED providers — where the model travels in
|
||||
// the URL, not the JSON body: Google Vertex (…/models/{model}:rawPredict) and
|
||||
// AWS Bedrock (/model/{id}/invoke).
|
||||
//
|
||||
// Each provider is tested in isolation with a guardrail allowlisting a single
|
||||
// model of its own: the allowed model (in the URL path) is served (200) and an
|
||||
// unselected model (in the URL path) is denied 403 by the guardrail
|
||||
// (llm_policy.model_blocked) before the upstream. The Vertex case mirrors the
|
||||
// customer verbatim — allow Sonnet, and the unselected model is the exact
|
||||
// claude-opus-4-6 they reported reaching the model unblocked. The Bedrock case
|
||||
// sends a region-prefixed, versioned inference-profile id so URL-path model
|
||||
// normalization is exercised too.
|
||||
//
|
||||
// The provider is catch-all (no models), so the router forwards any model and a
|
||||
// 403 can only come from the guardrail, never model_not_routable. Only the
|
||||
// upstream LLM is mocked (the vLLM nginx answers any path with 200); management
|
||||
// synth/reconcile, the proxy middleware chain (URL-path model extraction,
|
||||
// router, guardrail) and the tunnel are all real, and the guardrail denies
|
||||
// before the upstream is dialed so the mock cannot influence the block. A
|
||||
// static bearer api key is used so the router injects a static Authorization
|
||||
// header instead of minting a GCP token — the only reason path-routed providers
|
||||
// normally need live credentials — so the test runs with none and is always on.
|
||||
func TestGuardrailBlocksUnselectedModel_PathRouted(t *testing.T) {
|
||||
cases := []pathRoutedGuardrailCase{
|
||||
{
|
||||
name: "vertex",
|
||||
catalogID: "vertex_ai_api",
|
||||
wire: harness.WireVertex,
|
||||
allowEntry: "claude-sonnet-4-5",
|
||||
allowModel: "claude-sonnet-4-5",
|
||||
blockModel: "claude-opus-4-6", // the customer-reported model
|
||||
},
|
||||
{
|
||||
name: "bedrock",
|
||||
catalogID: "bedrock_api",
|
||||
wire: harness.WireBedrock,
|
||||
allowEntry: "anthropic.claude-sonnet-4-5", // normalized catalog id
|
||||
allowModel: "us.anthropic.claude-sonnet-4-5-v1:0",
|
||||
blockModel: "us.anthropic.claude-opus-4-8-v1:0",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
runPathRoutedGuardrailCase(t, tc)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func runPathRoutedGuardrailCase(t *testing.T, tc pathRoutedGuardrailCase) {
|
||||
t.Helper()
|
||||
|
||||
const (
|
||||
vertexProject = "e2e-project"
|
||||
vertexRegion = "global"
|
||||
)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-" + tc.name})
|
||||
require.NoError(t, err, "create group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-guardrail-" + tc.name + "-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
// Catch-all provider (no models) so the router forwards any model; a static
|
||||
// bearer key means the router injects a static auth header instead of minting
|
||||
// a GCP token. Bootstraps the cluster if it isn't already.
|
||||
staticKey := "static-e2e-token"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: tc.name,
|
||||
ProviderId: tc.catalogID,
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create %s provider", tc.name)
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
// Guardrail allowlisting ONLY this provider's allowed model.
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-guardrail-" + tc.name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{tc.allowEntry}
|
||||
guard, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail")
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
|
||||
|
||||
enabled := true
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-" + tc.name,
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{guard.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-"+tc.name+"-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
// Probe first: the GET resolves the endpoint and its first packet wakes the
|
||||
// lazy proxy peer, so WaitProxyPeer then observes it connected.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
var cerr error
|
||||
switch tc.wire {
|
||||
case harness.WireVertex:
|
||||
code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
|
||||
case harness.WireBedrock:
|
||||
code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
|
||||
default:
|
||||
t.Fatalf("unsupported wire %q", tc.wire)
|
||||
}
|
||||
require.NoError(t, cerr, "request must reach the proxy for %s", tc.name)
|
||||
return code, body
|
||||
}
|
||||
|
||||
// Allowed model (in the URL path) is served. Retry to absorb tunnel/DNS
|
||||
// jitter on the first call over the freshly warmed tunnel.
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(tc.allowModel)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
assert.Equal(t, 200, code,
|
||||
"allowed %s model (URL path) must be served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
|
||||
|
||||
// Unselected model (in the URL path) must be blocked by the guardrail.
|
||||
code, body = send(tc.blockModel)
|
||||
assert.Equal(t, 403, code,
|
||||
"unselected %s model (URL path) must be denied, not served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"%s denial must come from the guardrail allowlist, not routing; body: %s", tc.name, body)
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// TestGuardrailGroupSwitchTakesEffectAfterTTL proves that moving a peer between
|
||||
// groups flips its model-allowlist decision once the proxy's tunnel-peer cache
|
||||
// expires. The peer's groups reach the guardrail via ValidateTunnelPeer, which
|
||||
// the proxy caches; the switch is invisible until that cache expires. The proxy
|
||||
// runs with a short NB_PROXY_TUNNEL_CACHE_TTL so the flip happens in seconds
|
||||
// instead of the 5-minute default.
|
||||
//
|
||||
// Setup: one catch-all provider declaring modelA + modelB; polA (grpA -> allow
|
||||
// modelA) and polB (grpB -> allow modelB). The client starts in grpA. modelA is
|
||||
// served and modelB denied; after switching the client grpA -> grpB, modelB is
|
||||
// served and modelA denied. The cross-group deny comes from management's
|
||||
// per-policy/group CheckLLMPolicyLimits (the proxy backstop carries the union).
|
||||
func TestGuardrailGroupSwitchTakesEffectAfterTTL(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
modelA = "e2e-model-a"
|
||||
modelB = "e2e-model-b"
|
||||
// Short tunnel-cache TTL so a group switch propagates in seconds.
|
||||
// Exercises the NB_PROXY_TUNNEL_CACHE_TTL override.
|
||||
cacheTTL = 3 * time.Second
|
||||
)
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpA, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-a"})
|
||||
require.NoError(t, err, "create group A")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpA.Id) })
|
||||
|
||||
grpB, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-b"})
|
||||
require.NoError(t, err, "create group B")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpB.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-gswitch-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpA.Id}, // client starts in group A
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "gswitch",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: &[]api.AgentNetworkProviderModel{
|
||||
{Id: modelA, InputPer1k: 0.001, OutputPer1k: 0.001},
|
||||
{Id: modelB, InputPer1k: 0.001, OutputPer1k: 0.001},
|
||||
},
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
mkGuard := func(name, model string) api.AgentNetworkGuardrail {
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, gerr := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, gerr, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gA := mkGuard("e2e-gswitch-a", modelA)
|
||||
gB := mkGuard("e2e-gswitch-b", modelB)
|
||||
|
||||
enabled := true
|
||||
polA, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-gswitch-a",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpA.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gA.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy A")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polA.Id) })
|
||||
|
||||
polB, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-gswitch-b",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpB.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gB.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy B")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polB.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-gswitch-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken, map[string]string{
|
||||
"NB_PROXY_TUNNEL_CACHE_TTL": cacheTTL.String(),
|
||||
})
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
sendUntil := func(model string, want int, timeout time.Duration) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(model)
|
||||
if code == want {
|
||||
return code, body
|
||||
}
|
||||
time.Sleep(2 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
// Phase 1 — client is in group A: modelA served, modelB denied.
|
||||
code, body := sendUntil(modelA, 200, 90*time.Second)
|
||||
assert.Equal(t, 200, code,
|
||||
"group-A model must be served while the client is in group A; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
code, body = send(modelB)
|
||||
assert.Equal(t, 403, code,
|
||||
"group-B model must be denied while the client is in group A; body: %s", body)
|
||||
assert.Contains(t, body, "llm_policy.model_blocked")
|
||||
|
||||
// Switch the client peer from group A to group B.
|
||||
peerID := clientPeerInGroup(t, ctx, grpA.Id)
|
||||
_, err = srv.API().Groups.Update(ctx, grpB.Id, api.PutApiGroupsGroupIdJSONRequestBody{
|
||||
Name: grpB.Name,
|
||||
Peers: &[]string{peerID},
|
||||
})
|
||||
require.NoError(t, err, "add peer to group B")
|
||||
_, err = srv.API().Groups.Update(ctx, grpA.Id, api.PutApiGroupsGroupIdJSONRequestBody{
|
||||
Name: grpA.Name,
|
||||
Peers: &[]string{},
|
||||
})
|
||||
require.NoError(t, err, "remove peer from group A")
|
||||
|
||||
// Phase 2 — after the short TTL expires the proxy re-validates the peer,
|
||||
// sees group B, and the decision flips. Poll to absorb TTL + re-validation.
|
||||
code, body = sendUntil(modelB, 200, 60*time.Second)
|
||||
assert.Equal(t, 200, code,
|
||||
"after the group switch + TTL, the group-B model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
code, body = send(modelA)
|
||||
assert.Equal(t, 403, code,
|
||||
"after the switch, the old group-A model must be denied; body: %s", body)
|
||||
assert.Contains(t, body, "llm_policy.model_blocked")
|
||||
}
|
||||
|
||||
// clientPeerInGroup returns the id of the single peer that is a member of the
|
||||
// given group — the test client. The proxy peer is never added to test groups.
|
||||
func clientPeerInGroup(t *testing.T, ctx context.Context, groupID string) string {
|
||||
t.Helper()
|
||||
peers, err := srv.API().Peers.List(ctx)
|
||||
require.NoError(t, err, "list peers")
|
||||
for _, p := range peers {
|
||||
for _, g := range p.Groups {
|
||||
if g.Id == groupID {
|
||||
return p.Id
|
||||
}
|
||||
}
|
||||
}
|
||||
t.Fatalf("no peer found in group %s", groupID)
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// TestGuardrailMultiPolicyModelAllowlist: modelSelected served (200), grpOther's
|
||||
// modelOther denied for the grpMain client (403 model_blocked, no cross-group
|
||||
// leak), and openModel on the un-guardrailed policy's provider served (200).
|
||||
func TestGuardrailMultiPolicyModelAllowlist(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
modelSelected = "e2e-selected"
|
||||
modelOther = "e2e-other"
|
||||
openModel = "e2e-open"
|
||||
)
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-main"})
|
||||
require.NoError(t, err, "create main group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
|
||||
|
||||
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-other"})
|
||||
require.NoError(t, err, "create other group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-guardrail-mp-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpMain.Id}, // client joins grpMain only
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
models := func(ids ...string) *[]api.AgentNetworkProviderModel {
|
||||
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
// pRestricted declares the two guardrailed models so routing is deterministic
|
||||
// (model -> provider). Created first, so it carries the bootstrap cluster.
|
||||
pRestricted, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "restricted",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: models(modelSelected, modelOther),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create restricted provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pRestricted.Id) })
|
||||
|
||||
pOpen, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "open",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: models(openModel),
|
||||
})
|
||||
require.NoError(t, err, "create open provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pOpen.Id) })
|
||||
|
||||
mkGuardrail := func(name, model string) api.AgentNetworkGuardrail {
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, gerr := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, gerr, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gMain := mkGuardrail("e2e-guardrail-mp-main", modelSelected)
|
||||
gOther := mkGuardrail("e2e-guardrail-mp-other", modelOther)
|
||||
|
||||
enabled := true
|
||||
// polMain: grpMain restricted to modelSelected on pRestricted.
|
||||
polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-main",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{pRestricted.Id},
|
||||
GuardrailIds: &[]string{gMain.Id},
|
||||
})
|
||||
require.NoError(t, err, "create main policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
|
||||
|
||||
// polOther: grpOther restricted to modelOther on the SAME provider. The
|
||||
// client is not in grpOther, so modelOther must never be usable by it.
|
||||
polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-other",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpOther.Id},
|
||||
DestinationProviderIds: []string{pRestricted.Id},
|
||||
GuardrailIds: &[]string{gOther.Id},
|
||||
})
|
||||
require.NoError(t, err, "create other policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
|
||||
|
||||
// polOpen: grpMain on pOpen with NO guardrail — unrestricted.
|
||||
polOpen, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-open",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{pOpen.Id},
|
||||
})
|
||||
require.NoError(t, err, "create open policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOpen.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-mp-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
// sendUntil200 absorbs first-call tunnel/DNS jitter on the freshly warmed tunnel.
|
||||
sendUntil200 := func(model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(model)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
t.Run("selected model allowed for its group", func(t *testing.T) {
|
||||
code, body := sendUntil200(modelSelected)
|
||||
assert.Equal(t, 200, code,
|
||||
"grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
|
||||
t.Run("other group's model does not leak", func(t *testing.T) {
|
||||
// modelOther is allowlisted only for grpOther. The grpMain client must be
|
||||
// denied by management's per-policy/group check — not waved through by an
|
||||
// account-wide union. This is the security-critical wrong-ALLOW guard.
|
||||
code, body := send(modelOther)
|
||||
assert.Equal(t, 403, code,
|
||||
"another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"denial must be a model-allowlist decision; body: %s", body)
|
||||
})
|
||||
|
||||
t.Run("unguarded policy leaves its provider unrestricted", func(t *testing.T) {
|
||||
// polOpen carries no guardrail, so pOpen is unrestricted for grpMain. The
|
||||
// old account-wide union would have blocked openModel (it is on no
|
||||
// allowlist); it must now be served — the false-DENY guard.
|
||||
code, body := sendUntil200(openModel)
|
||||
assert.Equal(t, 200, code,
|
||||
"an un-guardrailed policy's provider must not be blocked by another policy's allowlist; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,422 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// pergroupCase describes one provider surface for the per-group allowlist matrix.
|
||||
// selectedReq/otherReq are the model identifiers as they travel in the request
|
||||
// (URL path for Bedrock/Vertex, body "model" for chat/messages). selectedAllow/
|
||||
// otherAllow are the (normalized) forms the guardrail allowlist holds — for
|
||||
// Bedrock these differ from the request form so path normalization is exercised.
|
||||
type pergroupCase struct {
|
||||
name string
|
||||
catalogID string
|
||||
wire string // "chat", "messages", "vertex", "bedrock"
|
||||
models *[]api.AgentNetworkProviderModel
|
||||
selectedReq string
|
||||
selectedAllow string
|
||||
otherReq string
|
||||
otherAllow string
|
||||
|
||||
providerID string // filled during setup
|
||||
}
|
||||
|
||||
// TestGuardrailPerGroupAllowlist_AllProviders proves the per-policy/group model
|
||||
// allowlist end to end across every always-on provider surface, including the
|
||||
// path-routed ones (Vertex, Bedrock) where the model travels in the URL.
|
||||
//
|
||||
// For each provider two policies target it: grpMain (the client) is allowed only
|
||||
// selectedReq; grpOther (which the client is NOT in) is allowed only otherReq.
|
||||
// The client must get selectedReq served (200) and otherReq denied (403,
|
||||
// llm_policy.model_blocked) — the cross-group no-leak property. The deny is the
|
||||
// authoritative per-policy/group decision from management (the proxy per-provider
|
||||
// backstop carries the union of both models), so this also confirms management
|
||||
// receives the correct normalized model for path-routed providers.
|
||||
func TestGuardrailPerGroupAllowlist_AllProviders(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
vertexProject = "e2e-project"
|
||||
vertexRegion = "global"
|
||||
)
|
||||
|
||||
priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
|
||||
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
cases := []*pergroupCase{
|
||||
{
|
||||
name: "openai", catalogID: "openai_api", wire: harness.WireChat,
|
||||
models: priced("oai-model-a", "oai-model-b"),
|
||||
selectedReq: "oai-model-a", selectedAllow: "oai-model-a",
|
||||
otherReq: "oai-model-b", otherAllow: "oai-model-b",
|
||||
},
|
||||
{
|
||||
name: "anthropic", catalogID: "anthropic_api", wire: harness.WireMessages,
|
||||
models: priced("ant-model-a", "ant-model-b"),
|
||||
selectedReq: "ant-model-a", selectedAllow: "ant-model-a",
|
||||
otherReq: "ant-model-b", otherAllow: "ant-model-b",
|
||||
},
|
||||
{
|
||||
// Vertex catalog ids travel bare in the rawPredict path.
|
||||
name: "vertex", catalogID: "vertex_ai_api", wire: "vertex",
|
||||
selectedReq: "claude-sonnet-4-5", selectedAllow: "claude-sonnet-4-5",
|
||||
otherReq: "claude-opus-4-6", otherAllow: "claude-opus-4-6",
|
||||
},
|
||||
{
|
||||
// Bedrock request ids are region-prefixed/versioned; the parser
|
||||
// normalizes them to the catalog key the allowlist holds.
|
||||
name: "bedrock", catalogID: "bedrock_api", wire: "bedrock",
|
||||
selectedReq: "us.anthropic.claude-sonnet-4-5-v1:0", selectedAllow: "anthropic.claude-sonnet-4-5",
|
||||
otherReq: "us.anthropic.claude-opus-4-8-v1:0", otherAllow: "anthropic.claude-opus-4-8",
|
||||
},
|
||||
}
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-main"})
|
||||
require.NoError(t, err, "create main group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
|
||||
|
||||
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-other"})
|
||||
require.NoError(t, err, "create other group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-pergroup-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpMain.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
enabled := true
|
||||
|
||||
for i, c := range cases {
|
||||
req := api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-pergroup-" + c.name,
|
||||
ProviderId: c.catalogID,
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: c.models,
|
||||
}
|
||||
if i == 0 {
|
||||
req.BootstrapCluster = ptr(harness.AgentNetworkCluster)
|
||||
}
|
||||
prov, perr := srv.CreateProvider(ctx, req)
|
||||
require.NoError(t, perr, "create provider %s", c.name)
|
||||
c.providerID = prov.Id
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
gSel := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-sel", c.selectedAllow)
|
||||
gOth := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-oth", c.otherAllow)
|
||||
|
||||
polMain, merr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-pergroup-" + c.name + "-main",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gSel.Id},
|
||||
})
|
||||
require.NoError(t, merr, "create main policy %s", c.name)
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
|
||||
|
||||
polOther, oerr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-pergroup-" + c.name + "-other",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpOther.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gOth.Id},
|
||||
})
|
||||
require.NoError(t, oerr, "create other policy %s", c.name)
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
|
||||
}
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-pergroup-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(c *pergroupCase, model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
var cerr error
|
||||
switch c.wire {
|
||||
case "vertex":
|
||||
code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
|
||||
case "bedrock":
|
||||
code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
|
||||
default:
|
||||
code, body, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, c.wire, model, "Reply with exactly: pong", "")
|
||||
}
|
||||
require.NoError(t, cerr, "request must reach the proxy for %s", c.name)
|
||||
return code, body
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
// grpMain's own model is served. Retry to absorb tunnel/DNS jitter on
|
||||
// the first call over the freshly warmed tunnel.
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(c, c.selectedReq)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
assert.Equal(t, 200, code,
|
||||
"%s: grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
|
||||
|
||||
// grpOther's model must NOT leak to the grpMain client.
|
||||
code, body = send(c, c.otherReq)
|
||||
assert.Equal(t, 403, code,
|
||||
"%s: another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"%s: denial must be a model-allowlist decision, not routing; body: %s", c.name, body)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGuardrailMultiGroupUser proves the per-policy/group decision for a caller
|
||||
// that belongs to MULTIPLE groups at once. Two scenarios, one shared stack:
|
||||
//
|
||||
// - union across the user's groups: the client is in gUX and gUY, each with
|
||||
// its own policy+guardrail on provider P1 (gUX->union-a, gUY->union-b). The
|
||||
// client may use BOTH models (the union of its groups' allowlists) while a
|
||||
// third, un-allowlisted model is denied.
|
||||
// - an un-guardrailed group lifts the restriction: the client is in gMP and
|
||||
// gMQ on provider P2, where gMP restricts to mix-a but gMQ's policy carries
|
||||
// NO guardrail. Because one applicable policy is unrestricted, the client may
|
||||
// use a model on no allowlist (mix-z) as well as mix-a.
|
||||
func TestGuardrailMultiGroupUser(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
unionA = "mg-union-a"
|
||||
unionB = "mg-union-b"
|
||||
unionC = "mg-union-c" // allowlisted by neither group
|
||||
mixA = "mg-mix-a"
|
||||
mixZ = "mg-mix-z" // on no allowlist; reachable only via the un-guardrailed policy
|
||||
)
|
||||
|
||||
priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
|
||||
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
mkGroup := func(name string) *api.Group {
|
||||
g, gerr := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: name})
|
||||
require.NoError(t, gerr, "create group %s", name)
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gUX := mkGroup("e2e-mg-union-x")
|
||||
gUY := mkGroup("e2e-mg-union-y")
|
||||
gMP := mkGroup("e2e-mg-mix-p")
|
||||
gMQ := mkGroup("e2e-mg-mix-q")
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-mg-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{gUX.Id, gUY.Id, gMP.Id, gMQ.Id}, // client in all four groups
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
enabled := true
|
||||
|
||||
// P1 — union scenario: two restricting policies, one per group.
|
||||
p1, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-mg-union",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: priced(unionA, unionB, unionC),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create union provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p1.Id) })
|
||||
|
||||
polUX, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-union-x",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gUX.Id},
|
||||
DestinationProviderIds: []string{p1.Id},
|
||||
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-x", unionA).Id},
|
||||
})
|
||||
require.NoError(t, err, "create union policy X")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUX.Id) })
|
||||
|
||||
polUY, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-union-y",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gUY.Id},
|
||||
DestinationProviderIds: []string{p1.Id},
|
||||
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-y", unionB).Id},
|
||||
})
|
||||
require.NoError(t, err, "create union policy Y")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUY.Id) })
|
||||
|
||||
// P2 — mixed scenario: one restricting policy + one un-guardrailed policy.
|
||||
p2, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-mg-mix",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: priced(mixA, mixZ),
|
||||
})
|
||||
require.NoError(t, err, "create mix provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p2.Id) })
|
||||
|
||||
polMP, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-mix-p",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gMP.Id},
|
||||
DestinationProviderIds: []string{p2.Id},
|
||||
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-mix-p", mixA).Id},
|
||||
})
|
||||
require.NoError(t, err, "create mix policy P")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMP.Id) })
|
||||
|
||||
polMQ, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-mix-q",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gMQ.Id},
|
||||
DestinationProviderIds: []string{p2.Id}, // NO guardrail -> unrestricted
|
||||
})
|
||||
require.NoError(t, err, "create mix policy Q")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMQ.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-mg-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
sendUntil200 := func(model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(model)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
t.Run("union across the user's groups", func(t *testing.T) {
|
||||
code, body := sendUntil200(unionA)
|
||||
assert.Equal(t, 200, code, "model allowed by group X must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
|
||||
code, body = sendUntil200(unionB)
|
||||
assert.Equal(t, 200, code, "model allowed by group Y must also be served (union across the user's groups); body: %s", body)
|
||||
|
||||
code, body = send(unionC)
|
||||
assert.Equal(t, 403, code, "a model on neither group's allowlist must be denied; body: %s", body)
|
||||
assert.Contains(t, body, "llm_policy.model_blocked")
|
||||
})
|
||||
|
||||
t.Run("an un-guardrailed group lifts the restriction", func(t *testing.T) {
|
||||
code, body := sendUntil200(mixA)
|
||||
assert.Equal(t, 200, code, "the restricted group's model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
|
||||
code, body = sendUntil200(mixZ)
|
||||
assert.Equal(t, 200, code,
|
||||
"a non-allowlisted model must be served because the user is also in a group whose policy has no guardrail; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
}
|
||||
|
||||
// mkAllowGuardrail creates a guardrail whose model allowlist is enabled and holds
|
||||
// exactly the given model, registering cleanup.
|
||||
func mkAllowGuardrail(t *testing.T, ctx context.Context, name, model string) api.AgentNetworkGuardrail {
|
||||
t.Helper()
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
+11
-1
@@ -38,7 +38,11 @@ type Proxy struct {
|
||||
// network, registered via the given account proxy token and serving the
|
||||
// AgentNetworkCluster over a self-signed wildcard cert. It does not wait for
|
||||
// peer connectivity — callers poll management for the proxy peer.
|
||||
func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, error) {
|
||||
// StartProxy launches the reverse-proxy container. Optional envOverrides are
|
||||
// merged into the container environment after the defaults, so callers can set
|
||||
// or override any NB_PROXY_* var (e.g. NB_PROXY_TUNNEL_CACHE_TTL for tests that
|
||||
// need a short authorization-cache window).
|
||||
func StartProxy(ctx context.Context, c *Combined, proxyToken string, envOverrides ...map[string]string) (*Proxy, error) {
|
||||
root, err := repoRoot()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -93,6 +97,12 @@ func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, er
|
||||
WaitingFor: wait.ForLog("Initial mapping sync complete").WithStartupTimeout(90 * time.Second),
|
||||
}
|
||||
|
||||
for _, ov := range envOverrides {
|
||||
for k, v := range ov {
|
||||
req.Env[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
|
||||
ContainerRequest: req,
|
||||
Started: true,
|
||||
|
||||
Reference in New Issue
Block a user