From bab5572a749d86adb8b9899d077fabd3b910d91a Mon Sep 17 00:00:00 2001 From: Maycon Santos Date: Tue, 28 Jul 2026 03:43:59 +0900 Subject: [PATCH] [management, proxy] scope agent-network model allowlist per policy/group and provider (#6905) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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. --- docs/agent-networks/01-end-to-end-flows.md | 17 +- .../modules/21-management-agentnetwork.md | 2 +- .../modules/31-proxy-middleware-builtin.md | 2 +- e2e/agentnetwork/guardrail_block_test.go | 209 +++++++++ .../guardrail_groupswitch_test.go | 205 +++++++++ .../guardrail_multipolicy_test.go | 201 +++++++++ .../guardrail_pergroup_providers_test.go | 422 ++++++++++++++++++ e2e/harness/proxy.go | 12 +- .../internals/modules/agentnetwork/manager.go | 7 +- .../modules/agentnetwork/policyselect.go | 108 +++++ .../agentnetwork/policyselect_model_test.go | 329 ++++++++++++++ .../modules/agentnetwork/synthesizer.go | 182 ++++---- .../synthesizer_guardrail_realstore_test.go | 6 +- .../synthesizer_provider_allowlist_test.go | 95 ++++ .../modules/agentnetwork/synthesizer_test.go | 8 +- management/internals/shared/grpc/proxy.go | 1 + .../grpc/proxy_llm_policy_limits_test.go | 138 ++++++ proxy/internal/auth/tunnel_cache.go | 39 +- proxy/internal/auth/tunnel_cache_test.go | 29 ++ .../builtin/llm_guardrail/factory.go | 37 +- .../builtin/llm_guardrail/middleware.go | 42 +- .../builtin/llm_guardrail/middleware_test.go | 143 +++++- .../builtin/llm_limit_check/middleware.go | 17 +- .../llm_limit_check/middleware_test.go | 40 ++ .../guardrail_allowlist_test.go | 11 +- 25 files changed, 2148 insertions(+), 154 deletions(-) create mode 100644 e2e/agentnetwork/guardrail_block_test.go create mode 100644 e2e/agentnetwork/guardrail_groupswitch_test.go create mode 100644 e2e/agentnetwork/guardrail_multipolicy_test.go create mode 100644 e2e/agentnetwork/guardrail_pergroup_providers_test.go create mode 100644 management/internals/modules/agentnetwork/policyselect_model_test.go create mode 100644 management/internals/modules/agentnetwork/synthesizer_provider_allowlist_test.go create mode 100644 management/internals/shared/grpc/proxy_llm_policy_limits_test.go diff --git a/docs/agent-networks/01-end-to-end-flows.md b/docs/agent-networks/01-end-to-end-flows.md index 7264f3768..b8891001b 100644 --- a/docs/agent-networks/01-end-to-end-flows.md +++ b/docs/agent-networks/01-end-to-end-flows.md @@ -109,7 +109,7 @@ sequenceDiagram Chk->>Inj: continue Inj->>Inj: inject NetBird identity headers per provider config Inj->>Grd: continue - Grd->>Grd: enforce model allowlist + Grd->>Grd: enforce per-provider allowlist (fail-closed backstop) Grd->>Up: forward (over WireGuard) Up-->>Resp: response (JSON or SSE stream) Resp->>Resp: parse usage tokens, completion @@ -135,6 +135,21 @@ sequenceDiagram (`redact_pii = settings.RedactPii`). Phones, emails, credit cards, PII names — see `redact.go` for the full set. See [`modules/31-proxy-middleware-builtin.md`](modules/31-proxy-middleware-builtin.md). +- The model allowlist is enforced in TWO places. `CheckLLMPolicyLimits` + is authoritative: it resolves the policy that governs this + (provider, caller-groups) and denies (`deny_code = llm_policy.model_blocked`) + when no applicable policy permits the model — so an allowlist scoped to + one group/provider never leaks to another, and an un-guardrailed policy + is genuinely unrestricted. `llm_guardrail` is a per-provider fail-closed + backstop: it only carries an allowlist for a provider every authorising + policy restricts, and blocks unknown/undetermined models even when + management is unreachable. Because that backstop allowlist is the UNION + of every restricting policy's models, per-group narrowing lives only in + the authoritative check: during a `CheckLLMPolicyLimits` outage + `llm_limit_check` fails open, so a caller can reach any model in the + provider's union — a group scoped to model A could reach model B if + another group restricts the same provider to B. This is the documented + fail-open trade-off; a future flag may switch it to fail-closed. - SSE streaming requires special handling on the response side; the parser must handle partial chunks without buffering the whole stream. See [`modules/32-proxy-llm-parsers.md`](modules/32-proxy-llm-parsers.md). diff --git a/docs/agent-networks/modules/21-management-agentnetwork.md b/docs/agent-networks/modules/21-management-agentnetwork.md index b64c1ba20..cc74206e9 100644 --- a/docs/agent-networks/modules/21-management-agentnetwork.md +++ b/docs/agent-networks/modules/21-management-agentnetwork.md @@ -122,7 +122,7 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest | on_request | 1 | `llm_router` | `{"providers":[{id, models[], upstream_*, auth_header_*, allowed_group_ids[]}]}` | **true** | | on_request | 2 | `llm_limit_check` | `{}` | – | | on_request | 3 | `llm_identity_inject` | `{"providers":[{provider_id, header_pair?, json_metadata?, extra_headers?}]}` | **true** | - | on_request | 4 | `llm_guardrail` | `{"model_allowlist"?, "prompt_capture":{enabled,redact_pii}}` | – | + | on_request | 4 | `llm_guardrail` | `{"provider_allowlists"?: {providerID: []model}, "prompt_capture":{enabled,redact_pii}}` | – | | on_response | 5 | `llm_limit_record` | `{}` (runs LAST at runtime) | – | | on_response | 6 | `cost_meter` | `{}` | – | | on_response | 7 | `llm_response_parser` | `{"capture_completion": , "redact_pii"?: true}` | – | diff --git a/docs/agent-networks/modules/31-proxy-middleware-builtin.md b/docs/agent-networks/modules/31-proxy-middleware-builtin.md index 904de6424..efe1bc4ce 100644 --- a/docs/agent-networks/modules/31-proxy-middleware-builtin.md +++ b/docs/agent-networks/modules/31-proxy-middleware-builtin.md @@ -244,7 +244,7 @@ no mocks. Tests: `TestChain_AllowPath_StampsAttributionAndRecordsCounter` | `llm_router` | `{providers: [{id, models, upstream_scheme, upstream_host, upstream_path?, auth_header_name, auth_header_value, allowed_group_ids}]}` | | `llm_limit_check` | `{}` — pulls `MgmtClient` from `FactoryContext` | | `llm_identity_inject` | `{providers: [{provider_id, header_pair?|json_metadata?, extra_headers?}]}` | -| `llm_guardrail` | `{model_allowlist: []string, prompt_capture: {enabled, redact_pii}}` | +| `llm_guardrail` | `{provider_allowlists: {providerID: []string}, prompt_capture: {enabled, redact_pii}}` — allowlist keyed by resolved provider id; a provider absent from the map is unrestricted (fail-closed backstop; authoritative per-policy/group check is management's `CheckLLMPolicyLimits`) | | `llm_response_parser` | `{redact_pii?, capture_completion?: *bool}` | | `cost_meter` | `{pricing_path?}` (basename inside data-dir; defaults `pricing.yaml`) | | `llm_limit_record` | `{}` — same pattern as `llm_limit_check` | diff --git a/e2e/agentnetwork/guardrail_block_test.go b/e2e/agentnetwork/guardrail_block_test.go new file mode 100644 index 000000000..c4b22ae25 --- /dev/null +++ b/e2e/agentnetwork/guardrail_block_test.go @@ -0,0 +1,209 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/e2e/harness" + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// pathRoutedGuardrailCase is one provider's self-contained scenario: its own +// provider, its own guardrail whose allowlist holds ONLY that provider's +// allowed model, and its own policy. Each case runs in isolation (its own +// proxy + client), so the guardrail the proxy enforces contains exactly this +// provider's model — never a mixed cross-provider list. +type pathRoutedGuardrailCase struct { + name string + catalogID string // agent-network catalog provider id + wire string // harness.WireVertex | harness.WireBedrock + allowEntry string // the single model id put on the guardrail allowlist + allowModel string // model id sent that MUST be served (200) + blockModel string // model id sent that MUST be denied (403 model_blocked) +} + +// TestGuardrailBlocksUnselectedModel_PathRouted is the end-to-end regression +// guard for the customer report that a model-allowlist guardrail attached to a +// policy has no effect for PATH-ROUTED providers — where the model travels in +// the URL, not the JSON body: Google Vertex (…/models/{model}:rawPredict) and +// AWS Bedrock (/model/{id}/invoke). +// +// Each provider is tested in isolation with a guardrail allowlisting a single +// model of its own: the allowed model (in the URL path) is served (200) and an +// unselected model (in the URL path) is denied 403 by the guardrail +// (llm_policy.model_blocked) before the upstream. The Vertex case mirrors the +// customer verbatim — allow Sonnet, and the unselected model is the exact +// claude-opus-4-6 they reported reaching the model unblocked. The Bedrock case +// sends a region-prefixed, versioned inference-profile id so URL-path model +// normalization is exercised too. +// +// The provider is catch-all (no models), so the router forwards any model and a +// 403 can only come from the guardrail, never model_not_routable. Only the +// upstream LLM is mocked (the vLLM nginx answers any path with 200); management +// synth/reconcile, the proxy middleware chain (URL-path model extraction, +// router, guardrail) and the tunnel are all real, and the guardrail denies +// before the upstream is dialed so the mock cannot influence the block. A +// static bearer api key is used so the router injects a static Authorization +// header instead of minting a GCP token — the only reason path-routed providers +// normally need live credentials — so the test runs with none and is always on. +func TestGuardrailBlocksUnselectedModel_PathRouted(t *testing.T) { + cases := []pathRoutedGuardrailCase{ + { + name: "vertex", + catalogID: "vertex_ai_api", + wire: harness.WireVertex, + allowEntry: "claude-sonnet-4-5", + allowModel: "claude-sonnet-4-5", + blockModel: "claude-opus-4-6", // the customer-reported model + }, + { + name: "bedrock", + catalogID: "bedrock_api", + wire: harness.WireBedrock, + allowEntry: "anthropic.claude-sonnet-4-5", // normalized catalog id + allowModel: "us.anthropic.claude-sonnet-4-5-v1:0", + blockModel: "us.anthropic.claude-opus-4-8-v1:0", + }, + } + + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + runPathRoutedGuardrailCase(t, tc) + }) + } +} + +func runPathRoutedGuardrailCase(t *testing.T, tc pathRoutedGuardrailCase) { + t.Helper() + + const ( + vertexProject = "e2e-project" + vertexRegion = "global" + ) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) + defer cancel() + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-" + tc.name}) + require.NoError(t, err, "create group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-guardrail-" + tc.name + "-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grp.Id}, + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + require.NotEmpty(t, sk.Key, "setup key plaintext") + + // Catch-all provider (no models) so the router forwards any model; a static + // bearer key means the router injects a static auth header instead of minting + // a GCP token. Bootstraps the cluster if it isn't already. + staticKey := "static-e2e-token" + prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: tc.name, + ProviderId: tc.catalogID, + UpstreamUrl: vllm.URL, + ApiKey: &staticKey, + Enabled: ptr(true), + BootstrapCluster: ptr(harness.AgentNetworkCluster), + }) + require.NoError(t, err, "create %s provider", tc.name) + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) }) + + // Guardrail allowlisting ONLY this provider's allowed model. + var gr api.AgentNetworkGuardrailRequest + gr.Name = "e2e-guardrail-" + tc.name + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = []string{tc.allowEntry} + guard, err := srv.CreateGuardrail(ctx, gr) + require.NoError(t, err, "create guardrail") + t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) }) + + enabled := true + pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-guardrail-" + tc.name, + Enabled: &enabled, + SourceGroups: []string{grp.Id}, + DestinationProviderIds: []string{prov.Id}, + GuardrailIds: &[]string{guard.Id}, + }) + require.NoError(t, err, "create policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) }) + + settings, err := srv.GetSettings(ctx) + require.NoError(t, err, "read settings") + require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned") + + proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-"+tc.name+"-proxy") + require.NoError(t, err, "mint proxy token") + px, err := harness.StartProxy(ctx, srv, proxyToken) + require.NoError(t, err, "start proxy") + t.Cleanup(func() { _ = px.Terminate(context.Background()) }) + + cl, err := harness.StartClient(ctx, srv, sk.Key) + require.NoError(t, err, "start client") + t.Cleanup(func() { _ = cl.Terminate(context.Background()) }) + + require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management") + // Probe first: the GET resolves the endpoint and its first packet wakes the + // lazy proxy peer, so WaitProxyPeer then observes it connected. + proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint) + require.NoError(t, err, "resolve endpoint to proxy IP") + if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil { + t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background())) + } + + send := func(model string) (int, string) { + var code int + var body string + var cerr error + switch tc.wire { + case harness.WireVertex: + code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "") + case harness.WireBedrock: + code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "") + default: + t.Fatalf("unsupported wire %q", tc.wire) + } + require.NoError(t, cerr, "request must reach the proxy for %s", tc.name) + return code, body + } + + // Allowed model (in the URL path) is served. Retry to absorb tunnel/DNS + // jitter on the first call over the freshly warmed tunnel. + var code int + var body string + deadline := time.Now().Add(90 * time.Second) + for time.Now().Before(deadline) { + code, body = send(tc.allowModel) + if code == 200 { + break + } + time.Sleep(5 * time.Second) + } + assert.Equal(t, 200, code, + "allowed %s model (URL path) must be served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background())) + + // Unselected model (in the URL path) must be blocked by the guardrail. + code, body = send(tc.blockModel) + assert.Equal(t, 403, code, + "unselected %s model (URL path) must be denied, not served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background())) + assert.Contains(t, body, "llm_policy.model_blocked", + "%s denial must come from the guardrail allowlist, not routing; body: %s", tc.name, body) +} diff --git a/e2e/agentnetwork/guardrail_groupswitch_test.go b/e2e/agentnetwork/guardrail_groupswitch_test.go new file mode 100644 index 000000000..5f2dce52e --- /dev/null +++ b/e2e/agentnetwork/guardrail_groupswitch_test.go @@ -0,0 +1,205 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/e2e/harness" + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// TestGuardrailGroupSwitchTakesEffectAfterTTL proves that moving a peer between +// groups flips its model-allowlist decision once the proxy's tunnel-peer cache +// expires. The peer's groups reach the guardrail via ValidateTunnelPeer, which +// the proxy caches; the switch is invisible until that cache expires. The proxy +// runs with a short NB_PROXY_TUNNEL_CACHE_TTL so the flip happens in seconds +// instead of the 5-minute default. +// +// Setup: one catch-all provider declaring modelA + modelB; polA (grpA -> allow +// modelA) and polB (grpB -> allow modelB). The client starts in grpA. modelA is +// served and modelB denied; after switching the client grpA -> grpB, modelB is +// served and modelA denied. The cross-group deny comes from management's +// per-policy/group CheckLLMPolicyLimits (the proxy backstop carries the union). +func TestGuardrailGroupSwitchTakesEffectAfterTTL(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) + defer cancel() + + const ( + modelA = "e2e-model-a" + modelB = "e2e-model-b" + // Short tunnel-cache TTL so a group switch propagates in seconds. + // Exercises the NB_PROXY_TUNNEL_CACHE_TTL override. + cacheTTL = 3 * time.Second + ) + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + grpA, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-a"}) + require.NoError(t, err, "create group A") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpA.Id) }) + + grpB, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-b"}) + require.NoError(t, err, "create group B") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpB.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-gswitch-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grpA.Id}, // client starts in group A + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + require.NotEmpty(t, sk.Key, "setup key plaintext") + + staticKey := "static-e2e-token" + prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "gswitch", + ProviderId: "openai_api", + UpstreamUrl: vllm.URL, + ApiKey: &staticKey, + Enabled: ptr(true), + Models: &[]api.AgentNetworkProviderModel{ + {Id: modelA, InputPer1k: 0.001, OutputPer1k: 0.001}, + {Id: modelB, InputPer1k: 0.001, OutputPer1k: 0.001}, + }, + BootstrapCluster: ptr(harness.AgentNetworkCluster), + }) + require.NoError(t, err, "create provider") + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) }) + + mkGuard := func(name, model string) api.AgentNetworkGuardrail { + var gr api.AgentNetworkGuardrailRequest + gr.Name = name + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = []string{model} + g, gerr := srv.CreateGuardrail(ctx, gr) + require.NoError(t, gerr, "create guardrail %s", name) + t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) }) + return g + } + gA := mkGuard("e2e-gswitch-a", modelA) + gB := mkGuard("e2e-gswitch-b", modelB) + + enabled := true + polA, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-gswitch-a", + Enabled: &enabled, + SourceGroups: []string{grpA.Id}, + DestinationProviderIds: []string{prov.Id}, + GuardrailIds: &[]string{gA.Id}, + }) + require.NoError(t, err, "create policy A") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polA.Id) }) + + polB, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-gswitch-b", + Enabled: &enabled, + SourceGroups: []string{grpB.Id}, + DestinationProviderIds: []string{prov.Id}, + GuardrailIds: &[]string{gB.Id}, + }) + require.NoError(t, err, "create policy B") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polB.Id) }) + + settings, err := srv.GetSettings(ctx) + require.NoError(t, err, "read settings") + require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned") + + proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-gswitch-proxy") + require.NoError(t, err, "mint proxy token") + px, err := harness.StartProxy(ctx, srv, proxyToken, map[string]string{ + "NB_PROXY_TUNNEL_CACHE_TTL": cacheTTL.String(), + }) + require.NoError(t, err, "start proxy") + t.Cleanup(func() { _ = px.Terminate(context.Background()) }) + + cl, err := harness.StartClient(ctx, srv, sk.Key) + require.NoError(t, err, "start client") + t.Cleanup(func() { _ = cl.Terminate(context.Background()) }) + + require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management") + proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint) + require.NoError(t, err, "resolve endpoint to proxy IP") + if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil { + t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background())) + } + + send := func(model string) (int, string) { + code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "") + require.NoError(t, cerr, "request must reach the proxy") + return code, body + } + sendUntil := func(model string, want int, timeout time.Duration) (int, string) { + var code int + var body string + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + code, body = send(model) + if code == want { + return code, body + } + time.Sleep(2 * time.Second) + } + return code, body + } + + // Phase 1 — client is in group A: modelA served, modelB denied. + code, body := sendUntil(modelA, 200, 90*time.Second) + assert.Equal(t, 200, code, + "group-A model must be served while the client is in group A; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background())) + code, body = send(modelB) + assert.Equal(t, 403, code, + "group-B model must be denied while the client is in group A; body: %s", body) + assert.Contains(t, body, "llm_policy.model_blocked") + + // Switch the client peer from group A to group B. + peerID := clientPeerInGroup(t, ctx, grpA.Id) + _, err = srv.API().Groups.Update(ctx, grpB.Id, api.PutApiGroupsGroupIdJSONRequestBody{ + Name: grpB.Name, + Peers: &[]string{peerID}, + }) + require.NoError(t, err, "add peer to group B") + _, err = srv.API().Groups.Update(ctx, grpA.Id, api.PutApiGroupsGroupIdJSONRequestBody{ + Name: grpA.Name, + Peers: &[]string{}, + }) + require.NoError(t, err, "remove peer from group A") + + // Phase 2 — after the short TTL expires the proxy re-validates the peer, + // sees group B, and the decision flips. Poll to absorb TTL + re-validation. + code, body = sendUntil(modelB, 200, 60*time.Second) + assert.Equal(t, 200, code, + "after the group switch + TTL, the group-B model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background())) + code, body = send(modelA) + assert.Equal(t, 403, code, + "after the switch, the old group-A model must be denied; body: %s", body) + assert.Contains(t, body, "llm_policy.model_blocked") +} + +// clientPeerInGroup returns the id of the single peer that is a member of the +// given group — the test client. The proxy peer is never added to test groups. +func clientPeerInGroup(t *testing.T, ctx context.Context, groupID string) string { + t.Helper() + peers, err := srv.API().Peers.List(ctx) + require.NoError(t, err, "list peers") + for _, p := range peers { + for _, g := range p.Groups { + if g.Id == groupID { + return p.Id + } + } + } + t.Fatalf("no peer found in group %s", groupID) + return "" +} diff --git a/e2e/agentnetwork/guardrail_multipolicy_test.go b/e2e/agentnetwork/guardrail_multipolicy_test.go new file mode 100644 index 000000000..1d664fb58 --- /dev/null +++ b/e2e/agentnetwork/guardrail_multipolicy_test.go @@ -0,0 +1,201 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/e2e/harness" + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// TestGuardrailMultiPolicyModelAllowlist: modelSelected served (200), grpOther's +// modelOther denied for the grpMain client (403 model_blocked, no cross-group +// leak), and openModel on the un-guardrailed policy's provider served (200). +func TestGuardrailMultiPolicyModelAllowlist(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) + defer cancel() + + const ( + modelSelected = "e2e-selected" + modelOther = "e2e-other" + openModel = "e2e-open" + ) + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-main"}) + require.NoError(t, err, "create main group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) }) + + grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-other"}) + require.NoError(t, err, "create other group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-guardrail-mp-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grpMain.Id}, // client joins grpMain only + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + require.NotEmpty(t, sk.Key, "setup key plaintext") + + staticKey := "static-e2e-token" + models := func(ids ...string) *[]api.AgentNetworkProviderModel { + out := make([]api.AgentNetworkProviderModel, 0, len(ids)) + for _, id := range ids { + out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001}) + } + return &out + } + + // pRestricted declares the two guardrailed models so routing is deterministic + // (model -> provider). Created first, so it carries the bootstrap cluster. + pRestricted, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "restricted", + ProviderId: "openai_api", + UpstreamUrl: vllm.URL, + ApiKey: &staticKey, + Enabled: ptr(true), + Models: models(modelSelected, modelOther), + BootstrapCluster: ptr(harness.AgentNetworkCluster), + }) + require.NoError(t, err, "create restricted provider") + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pRestricted.Id) }) + + pOpen, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "open", + ProviderId: "openai_api", + UpstreamUrl: vllm.URL, + ApiKey: &staticKey, + Enabled: ptr(true), + Models: models(openModel), + }) + require.NoError(t, err, "create open provider") + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pOpen.Id) }) + + mkGuardrail := func(name, model string) api.AgentNetworkGuardrail { + var gr api.AgentNetworkGuardrailRequest + gr.Name = name + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = []string{model} + g, gerr := srv.CreateGuardrail(ctx, gr) + require.NoError(t, gerr, "create guardrail %s", name) + t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) }) + return g + } + gMain := mkGuardrail("e2e-guardrail-mp-main", modelSelected) + gOther := mkGuardrail("e2e-guardrail-mp-other", modelOther) + + enabled := true + // polMain: grpMain restricted to modelSelected on pRestricted. + polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-guardrail-mp-main", + Enabled: &enabled, + SourceGroups: []string{grpMain.Id}, + DestinationProviderIds: []string{pRestricted.Id}, + GuardrailIds: &[]string{gMain.Id}, + }) + require.NoError(t, err, "create main policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) }) + + // polOther: grpOther restricted to modelOther on the SAME provider. The + // client is not in grpOther, so modelOther must never be usable by it. + polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-guardrail-mp-other", + Enabled: &enabled, + SourceGroups: []string{grpOther.Id}, + DestinationProviderIds: []string{pRestricted.Id}, + GuardrailIds: &[]string{gOther.Id}, + }) + require.NoError(t, err, "create other policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) }) + + // polOpen: grpMain on pOpen with NO guardrail — unrestricted. + polOpen, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-guardrail-mp-open", + Enabled: &enabled, + SourceGroups: []string{grpMain.Id}, + DestinationProviderIds: []string{pOpen.Id}, + }) + require.NoError(t, err, "create open policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOpen.Id) }) + + settings, err := srv.GetSettings(ctx) + require.NoError(t, err, "read settings") + require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned") + + proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-mp-proxy") + require.NoError(t, err, "mint proxy token") + px, err := harness.StartProxy(ctx, srv, proxyToken) + require.NoError(t, err, "start proxy") + t.Cleanup(func() { _ = px.Terminate(context.Background()) }) + + cl, err := harness.StartClient(ctx, srv, sk.Key) + require.NoError(t, err, "start client") + t.Cleanup(func() { _ = cl.Terminate(context.Background()) }) + + require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management") + proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint) + require.NoError(t, err, "resolve endpoint to proxy IP") + if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil { + t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background())) + } + + send := func(model string) (int, string) { + code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "") + require.NoError(t, cerr, "request must reach the proxy") + return code, body + } + // sendUntil200 absorbs first-call tunnel/DNS jitter on the freshly warmed tunnel. + sendUntil200 := func(model string) (int, string) { + var code int + var body string + deadline := time.Now().Add(90 * time.Second) + for time.Now().Before(deadline) { + code, body = send(model) + if code == 200 { + break + } + time.Sleep(5 * time.Second) + } + return code, body + } + + t.Run("selected model allowed for its group", func(t *testing.T) { + code, body := sendUntil200(modelSelected) + assert.Equal(t, 200, code, + "grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background())) + }) + + t.Run("other group's model does not leak", func(t *testing.T) { + // modelOther is allowlisted only for grpOther. The grpMain client must be + // denied by management's per-policy/group check — not waved through by an + // account-wide union. This is the security-critical wrong-ALLOW guard. + code, body := send(modelOther) + assert.Equal(t, 403, code, + "another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background())) + assert.Contains(t, body, "llm_policy.model_blocked", + "denial must be a model-allowlist decision; body: %s", body) + }) + + t.Run("unguarded policy leaves its provider unrestricted", func(t *testing.T) { + // polOpen carries no guardrail, so pOpen is unrestricted for grpMain. The + // old account-wide union would have blocked openModel (it is on no + // allowlist); it must now be served — the false-DENY guard. + code, body := sendUntil200(openModel) + assert.Equal(t, 200, code, + "an un-guardrailed policy's provider must not be blocked by another policy's allowlist; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background())) + }) +} diff --git a/e2e/agentnetwork/guardrail_pergroup_providers_test.go b/e2e/agentnetwork/guardrail_pergroup_providers_test.go new file mode 100644 index 000000000..eddae65c3 --- /dev/null +++ b/e2e/agentnetwork/guardrail_pergroup_providers_test.go @@ -0,0 +1,422 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/e2e/harness" + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// pergroupCase describes one provider surface for the per-group allowlist matrix. +// selectedReq/otherReq are the model identifiers as they travel in the request +// (URL path for Bedrock/Vertex, body "model" for chat/messages). selectedAllow/ +// otherAllow are the (normalized) forms the guardrail allowlist holds — for +// Bedrock these differ from the request form so path normalization is exercised. +type pergroupCase struct { + name string + catalogID string + wire string // "chat", "messages", "vertex", "bedrock" + models *[]api.AgentNetworkProviderModel + selectedReq string + selectedAllow string + otherReq string + otherAllow string + + providerID string // filled during setup +} + +// TestGuardrailPerGroupAllowlist_AllProviders proves the per-policy/group model +// allowlist end to end across every always-on provider surface, including the +// path-routed ones (Vertex, Bedrock) where the model travels in the URL. +// +// For each provider two policies target it: grpMain (the client) is allowed only +// selectedReq; grpOther (which the client is NOT in) is allowed only otherReq. +// The client must get selectedReq served (200) and otherReq denied (403, +// llm_policy.model_blocked) — the cross-group no-leak property. The deny is the +// authoritative per-policy/group decision from management (the proxy per-provider +// backstop carries the union of both models), so this also confirms management +// receives the correct normalized model for path-routed providers. +func TestGuardrailPerGroupAllowlist_AllProviders(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) + defer cancel() + + const ( + vertexProject = "e2e-project" + vertexRegion = "global" + ) + + priced := func(ids ...string) *[]api.AgentNetworkProviderModel { + out := make([]api.AgentNetworkProviderModel, 0, len(ids)) + for _, id := range ids { + out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001}) + } + return &out + } + + cases := []*pergroupCase{ + { + name: "openai", catalogID: "openai_api", wire: harness.WireChat, + models: priced("oai-model-a", "oai-model-b"), + selectedReq: "oai-model-a", selectedAllow: "oai-model-a", + otherReq: "oai-model-b", otherAllow: "oai-model-b", + }, + { + name: "anthropic", catalogID: "anthropic_api", wire: harness.WireMessages, + models: priced("ant-model-a", "ant-model-b"), + selectedReq: "ant-model-a", selectedAllow: "ant-model-a", + otherReq: "ant-model-b", otherAllow: "ant-model-b", + }, + { + // Vertex catalog ids travel bare in the rawPredict path. + name: "vertex", catalogID: "vertex_ai_api", wire: "vertex", + selectedReq: "claude-sonnet-4-5", selectedAllow: "claude-sonnet-4-5", + otherReq: "claude-opus-4-6", otherAllow: "claude-opus-4-6", + }, + { + // Bedrock request ids are region-prefixed/versioned; the parser + // normalizes them to the catalog key the allowlist holds. + name: "bedrock", catalogID: "bedrock_api", wire: "bedrock", + selectedReq: "us.anthropic.claude-sonnet-4-5-v1:0", selectedAllow: "anthropic.claude-sonnet-4-5", + otherReq: "us.anthropic.claude-opus-4-8-v1:0", otherAllow: "anthropic.claude-opus-4-8", + }, + } + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-main"}) + require.NoError(t, err, "create main group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) }) + + grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-other"}) + require.NoError(t, err, "create other group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-pergroup-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grpMain.Id}, + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + require.NotEmpty(t, sk.Key, "setup key plaintext") + + staticKey := "static-e2e-token" + enabled := true + + for i, c := range cases { + req := api.AgentNetworkProviderRequest{ + Name: "e2e-pergroup-" + c.name, + ProviderId: c.catalogID, + UpstreamUrl: vllm.URL, + ApiKey: &staticKey, + Enabled: ptr(true), + Models: c.models, + } + if i == 0 { + req.BootstrapCluster = ptr(harness.AgentNetworkCluster) + } + prov, perr := srv.CreateProvider(ctx, req) + require.NoError(t, perr, "create provider %s", c.name) + c.providerID = prov.Id + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) }) + + gSel := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-sel", c.selectedAllow) + gOth := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-oth", c.otherAllow) + + polMain, merr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-pergroup-" + c.name + "-main", + Enabled: &enabled, + SourceGroups: []string{grpMain.Id}, + DestinationProviderIds: []string{prov.Id}, + GuardrailIds: &[]string{gSel.Id}, + }) + require.NoError(t, merr, "create main policy %s", c.name) + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) }) + + polOther, oerr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-pergroup-" + c.name + "-other", + Enabled: &enabled, + SourceGroups: []string{grpOther.Id}, + DestinationProviderIds: []string{prov.Id}, + GuardrailIds: &[]string{gOth.Id}, + }) + require.NoError(t, oerr, "create other policy %s", c.name) + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) }) + } + + settings, err := srv.GetSettings(ctx) + require.NoError(t, err, "read settings") + require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned") + + proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-pergroup-proxy") + require.NoError(t, err, "mint proxy token") + px, err := harness.StartProxy(ctx, srv, proxyToken) + require.NoError(t, err, "start proxy") + t.Cleanup(func() { _ = px.Terminate(context.Background()) }) + + cl, err := harness.StartClient(ctx, srv, sk.Key) + require.NoError(t, err, "start client") + t.Cleanup(func() { _ = cl.Terminate(context.Background()) }) + + require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management") + proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint) + require.NoError(t, err, "resolve endpoint to proxy IP") + if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil { + t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background())) + } + + send := func(c *pergroupCase, model string) (int, string) { + var code int + var body string + var cerr error + switch c.wire { + case "vertex": + code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "") + case "bedrock": + code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "") + default: + code, body, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, c.wire, model, "Reply with exactly: pong", "") + } + require.NoError(t, cerr, "request must reach the proxy for %s", c.name) + return code, body + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + // grpMain's own model is served. Retry to absorb tunnel/DNS jitter on + // the first call over the freshly warmed tunnel. + var code int + var body string + deadline := time.Now().Add(90 * time.Second) + for time.Now().Before(deadline) { + code, body = send(c, c.selectedReq) + if code == 200 { + break + } + time.Sleep(5 * time.Second) + } + assert.Equal(t, 200, code, + "%s: grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background())) + + // grpOther's model must NOT leak to the grpMain client. + code, body = send(c, c.otherReq) + assert.Equal(t, 403, code, + "%s: another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background())) + assert.Contains(t, body, "llm_policy.model_blocked", + "%s: denial must be a model-allowlist decision, not routing; body: %s", c.name, body) + }) + } +} + +// TestGuardrailMultiGroupUser proves the per-policy/group decision for a caller +// that belongs to MULTIPLE groups at once. Two scenarios, one shared stack: +// +// - union across the user's groups: the client is in gUX and gUY, each with +// its own policy+guardrail on provider P1 (gUX->union-a, gUY->union-b). The +// client may use BOTH models (the union of its groups' allowlists) while a +// third, un-allowlisted model is denied. +// - an un-guardrailed group lifts the restriction: the client is in gMP and +// gMQ on provider P2, where gMP restricts to mix-a but gMQ's policy carries +// NO guardrail. Because one applicable policy is unrestricted, the client may +// use a model on no allowlist (mix-z) as well as mix-a. +func TestGuardrailMultiGroupUser(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) + defer cancel() + + const ( + unionA = "mg-union-a" + unionB = "mg-union-b" + unionC = "mg-union-c" // allowlisted by neither group + mixA = "mg-mix-a" + mixZ = "mg-mix-z" // on no allowlist; reachable only via the un-guardrailed policy + ) + + priced := func(ids ...string) *[]api.AgentNetworkProviderModel { + out := make([]api.AgentNetworkProviderModel, 0, len(ids)) + for _, id := range ids { + out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001}) + } + return &out + } + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + mkGroup := func(name string) *api.Group { + g, gerr := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: name}) + require.NoError(t, gerr, "create group %s", name) + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), g.Id) }) + return g + } + gUX := mkGroup("e2e-mg-union-x") + gUY := mkGroup("e2e-mg-union-y") + gMP := mkGroup("e2e-mg-mix-p") + gMQ := mkGroup("e2e-mg-mix-q") + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-mg-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{gUX.Id, gUY.Id, gMP.Id, gMQ.Id}, // client in all four groups + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + require.NotEmpty(t, sk.Key, "setup key plaintext") + + staticKey := "static-e2e-token" + enabled := true + + // P1 — union scenario: two restricting policies, one per group. + p1, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "e2e-mg-union", + ProviderId: "openai_api", + UpstreamUrl: vllm.URL, + ApiKey: &staticKey, + Enabled: ptr(true), + Models: priced(unionA, unionB, unionC), + BootstrapCluster: ptr(harness.AgentNetworkCluster), + }) + require.NoError(t, err, "create union provider") + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p1.Id) }) + + polUX, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-mg-union-x", + Enabled: &enabled, + SourceGroups: []string{gUX.Id}, + DestinationProviderIds: []string{p1.Id}, + GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-x", unionA).Id}, + }) + require.NoError(t, err, "create union policy X") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUX.Id) }) + + polUY, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-mg-union-y", + Enabled: &enabled, + SourceGroups: []string{gUY.Id}, + DestinationProviderIds: []string{p1.Id}, + GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-y", unionB).Id}, + }) + require.NoError(t, err, "create union policy Y") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUY.Id) }) + + // P2 — mixed scenario: one restricting policy + one un-guardrailed policy. + p2, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "e2e-mg-mix", + ProviderId: "openai_api", + UpstreamUrl: vllm.URL, + ApiKey: &staticKey, + Enabled: ptr(true), + Models: priced(mixA, mixZ), + }) + require.NoError(t, err, "create mix provider") + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p2.Id) }) + + polMP, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-mg-mix-p", + Enabled: &enabled, + SourceGroups: []string{gMP.Id}, + DestinationProviderIds: []string{p2.Id}, + GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-mix-p", mixA).Id}, + }) + require.NoError(t, err, "create mix policy P") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMP.Id) }) + + polMQ, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-mg-mix-q", + Enabled: &enabled, + SourceGroups: []string{gMQ.Id}, + DestinationProviderIds: []string{p2.Id}, // NO guardrail -> unrestricted + }) + require.NoError(t, err, "create mix policy Q") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMQ.Id) }) + + settings, err := srv.GetSettings(ctx) + require.NoError(t, err, "read settings") + require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned") + + proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-mg-proxy") + require.NoError(t, err, "mint proxy token") + px, err := harness.StartProxy(ctx, srv, proxyToken) + require.NoError(t, err, "start proxy") + t.Cleanup(func() { _ = px.Terminate(context.Background()) }) + + cl, err := harness.StartClient(ctx, srv, sk.Key) + require.NoError(t, err, "start client") + t.Cleanup(func() { _ = cl.Terminate(context.Background()) }) + + require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management") + proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint) + require.NoError(t, err, "resolve endpoint to proxy IP") + if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil { + t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background())) + } + + send := func(model string) (int, string) { + code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "") + require.NoError(t, cerr, "request must reach the proxy") + return code, body + } + sendUntil200 := func(model string) (int, string) { + var code int + var body string + deadline := time.Now().Add(90 * time.Second) + for time.Now().Before(deadline) { + code, body = send(model) + if code == 200 { + break + } + time.Sleep(5 * time.Second) + } + return code, body + } + + t.Run("union across the user's groups", func(t *testing.T) { + code, body := sendUntil200(unionA) + assert.Equal(t, 200, code, "model allowed by group X must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background())) + + code, body = sendUntil200(unionB) + assert.Equal(t, 200, code, "model allowed by group Y must also be served (union across the user's groups); body: %s", body) + + code, body = send(unionC) + assert.Equal(t, 403, code, "a model on neither group's allowlist must be denied; body: %s", body) + assert.Contains(t, body, "llm_policy.model_blocked") + }) + + t.Run("an un-guardrailed group lifts the restriction", func(t *testing.T) { + code, body := sendUntil200(mixA) + assert.Equal(t, 200, code, "the restricted group's model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background())) + + code, body = sendUntil200(mixZ) + assert.Equal(t, 200, code, + "a non-allowlisted model must be served because the user is also in a group whose policy has no guardrail; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background())) + }) +} + +// mkAllowGuardrail creates a guardrail whose model allowlist is enabled and holds +// exactly the given model, registering cleanup. +func mkAllowGuardrail(t *testing.T, ctx context.Context, name, model string) api.AgentNetworkGuardrail { + t.Helper() + var gr api.AgentNetworkGuardrailRequest + gr.Name = name + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = []string{model} + g, err := srv.CreateGuardrail(ctx, gr) + require.NoError(t, err, "create guardrail %s", name) + t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) }) + return g +} diff --git a/e2e/harness/proxy.go b/e2e/harness/proxy.go index 8db2c140f..85f3518d4 100644 --- a/e2e/harness/proxy.go +++ b/e2e/harness/proxy.go @@ -38,7 +38,11 @@ type Proxy struct { // network, registered via the given account proxy token and serving the // AgentNetworkCluster over a self-signed wildcard cert. It does not wait for // peer connectivity — callers poll management for the proxy peer. -func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, error) { +// StartProxy launches the reverse-proxy container. Optional envOverrides are +// merged into the container environment after the defaults, so callers can set +// or override any NB_PROXY_* var (e.g. NB_PROXY_TUNNEL_CACHE_TTL for tests that +// need a short authorization-cache window). +func StartProxy(ctx context.Context, c *Combined, proxyToken string, envOverrides ...map[string]string) (*Proxy, error) { root, err := repoRoot() if err != nil { return nil, err @@ -93,6 +97,12 @@ func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, er WaitingFor: wait.ForLog("Initial mapping sync complete").WithStartupTimeout(90 * time.Second), } + for _, ov := range envOverrides { + for k, v := range ov { + req.Env[k] = v + } + } + ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ ContainerRequest: req, Started: true, diff --git a/management/internals/modules/agentnetwork/manager.go b/management/internals/modules/agentnetwork/manager.go index d88e0d77c..77c77ce44 100644 --- a/management/internals/modules/agentnetwork/manager.go +++ b/management/internals/modules/agentnetwork/manager.go @@ -86,12 +86,17 @@ type Manager interface { // PolicySelectionInput is the per-request selection envelope. The // proxy populates it from CapturedData (account, user, groups) plus -// the provider llm_router resolved. +// the provider llm_router resolved and the model it extracted. type PolicySelectionInput struct { AccountID string UserID string GroupIDs []string ProviderID string + // Model is the already-normalised upstream model id the proxy extracted + // (parser strips Bedrock region/version, Vertex @version), so a + // case-insensitive compare suffices. Empty = undetermined → not permitted + // (fail closed). + Model string } // PolicySelectionResult names the policy that "pays" for this request diff --git a/management/internals/modules/agentnetwork/policyselect.go b/management/internals/modules/agentnetwork/policyselect.go index 9203a1910..9bb893b36 100644 --- a/management/internals/modules/agentnetwork/policyselect.go +++ b/management/internals/modules/agentnetwork/policyselect.go @@ -5,6 +5,7 @@ import ( "fmt" "math" "sort" + "strings" "time" "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" @@ -35,6 +36,10 @@ const ( denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded" //nolint:gosec // account deny code label, not a credential denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded" + // denyCodeModelBlocked is returned when policies govern the request's + // (provider, caller-groups) but none permits the model. Matches the proxy + // guardrail's code so both layers surface the same label. + denyCodeModelBlocked = "llm_policy.model_blocked" ) // consumptionCache holds the consumption counters prefetched for one @@ -159,6 +164,25 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec } candidates := filterApplicablePolicies(policies, in) + // Model-allowlist gate scoped to the matched policies: keep candidates whose + // guardrails permit the model (none enabled = unrestricted), deny when + // policies apply but none permits it. Skip the load when none has a guardrail. + if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) { + guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID) + if gErr != nil { + return nil, gErr + } + permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model) + if len(permitted) == 0 { + return &PolicySelectionResult{ + Allow: false, + DenyCode: denyCodeModelBlocked, + DenyReason: modelBlockedReason(in.Model), + }, nil + } + candidates = permitted + } + // Prefetch every consumption counter the ceiling + candidate policies will // read, in a single store round-trip, then score against the cache. cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now) @@ -250,6 +274,90 @@ func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput) return out } +// anyPolicyHasGuardrails reports whether any policy references at least one +// guardrail, so the selector can skip loading guardrails when none do. +func anyPolicyHasGuardrails(policies []*types.Policy) bool { + for _, p := range policies { + if p != nil && len(p.GuardrailIDs) > 0 { + return true + } + } + return false +} + +// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the +// model-allowlist gate to resolve each candidate policy's attached guardrails. +func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) { + guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID) + if err != nil { + return nil, fmt.Errorf("list account guardrails: %w", err) + } + byID := make(map[string]*types.Guardrail, len(guardrails)) + for _, g := range guardrails { + if g != nil { + byID[g.ID] = g + } + } + return byID, nil +} + +// filterModelPermittedPolicies returns the subset of policies whose guardrails +// permit the model. Order is preserved so downstream scoring is unaffected. +func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy { + out := make([]*types.Policy, 0, len(policies)) + for _, p := range policies { + if policyPermitsModel(p, byID, model) { + out = append(out, p) + } + } + return out +} + +// policyPermitsModel reports whether a policy permits the model. No +// allowlist-enabled guardrail = unrestricted (permits any, incl. empty); +// otherwise the model must be in the union of its allowlists, so an +// empty/undetermined model fails closed. +func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool { + if p == nil { + return false + } + wanted := normaliseModelID(model) + restricted := false + for _, gID := range p.GuardrailIDs { + g, ok := byID[gID] + if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled { + continue + } + restricted = true + if wanted == "" { + continue + } + for _, allowed := range g.Checks.ModelAllowlist.Models { + if normaliseModelID(allowed) == wanted { + return true + } + } + } + return !restricted +} + +// normaliseModelID lowercases and trims a model identifier so the allowlist +// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's +// normaliseModel so both layers agree on what "same model" means. +func normaliseModelID(model string) string { + return strings.ToLower(strings.TrimSpace(model)) +} + +// modelBlockedReason builds the human-readable deny reason for a model-allowlist +// rejection. The model is quoted when known; an undetermined model is reported +// as such so the access log distinguishes "wrong model" from "no model". +func modelBlockedReason(model string) string { + if normaliseModelID(model) == "" { + return "request model could not be determined for the policy allowlist" + } + return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model) +} + // candidate is the per-policy intermediate the selector ranks. A // policy that's been exhausted on any enabled cap never makes it // into this slice; the selector's deny envelope carries the latest diff --git a/management/internals/modules/agentnetwork/policyselect_model_test.go b/management/internals/modules/agentnetwork/policyselect_model_test.go new file mode 100644 index 000000000..c122cc36c --- /dev/null +++ b/management/internals/modules/agentnetwork/policyselect_model_test.go @@ -0,0 +1,329 @@ +package agentnetwork + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/golang/mock/gomock" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" +) + +// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups +// to reach providerID under the given guardrails. Uncapped keeps the selector's +// headroom scoring trivial so these tests isolate the model-allowlist gate. +func guardedPolicy(id, account string, sourceGroups []string, providerID string, guardrailIDs ...string) *types.Policy { + return &types.Policy{ + ID: id, + AccountID: account, + Enabled: true, + SourceGroups: sourceGroups, + DestinationProviderIDs: []string{providerID}, + GuardrailIDs: guardrailIDs, + CreatedAt: time.Now().UTC(), + } +} + +// allowlistGuardrail builds a guardrail whose model allowlist is enabled and +// carries the given models. +func allowlistGuardrail(id, account string, models ...string) *types.Guardrail { + return &types.Guardrail{ + ID: id, + AccountID: account, + Checks: types.GuardrailChecks{ + ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: models}, + }, + } +} + +func expectPolicies(mockStore *store.MockStore, account string, policies ...*types.Policy) { + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), account). + Return(policies, nil) +} + +func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...*types.Guardrail) { + mockStore.EXPECT(). + GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), account). + Return(guardrails, nil) +} + +// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist +// decision: a policy authorises the (provider, group) but restricts the model, +// and the requested model isn't on the list, so the request is denied. +func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1") + expectPolicies(mockStore, "acc-1", policy) + expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o")) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + UserID: "user-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + Model: "claude-opus-4", + }) + require.NoError(t, err) + assert.False(t, res.Allow, "a model outside the only applicable policy's allowlist must be denied") + assert.Equal(t, denyCodeModelBlocked, res.DenyCode, "deny code must be model_blocked") + assert.NotEmpty(t, res.DenyReason, "deny reason must be populated") +} + +// TestSelectPolicy_ModelAllowedByAllowlist is the allow counterpart: the model +// is on the applicable policy's allowlist, so selection proceeds normally. +func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1") + expectPolicies(mockStore, "acc-1", policy) + expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4")) + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + UserID: "user-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + Model: "claude-opus-4", + }) + require.NoError(t, err) + assert.True(t, res.Allow, "a model on the applicable policy's allowlist must be allowed") + assert.Equal(t, "pol-A", res.SelectedPolicyID) +} + +// TestSelectPolicy_CaseInsensitiveModelMatch proves the compare tolerates case +// and surrounding whitespace, matching the proxy guardrail's normalisation. +func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1") + expectPolicies(mockStore, "acc-1", policy) + expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", " GPT-4o ")) + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + Model: "gpt-4o", + }) + require.NoError(t, err) + assert.True(t, res.Allow, "case/whitespace variants must match the allowlist entry") +} + +// TestSelectPolicy_UnguardedPolicyIsUnrestricted is the false-deny fix: when two +// policies authorise the same (provider, group) and one has no guardrail, that +// policy makes the request unrestricted — not caught by the other's allowlist. +func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + restricted := guardedPolicy("pol-restricted", "acc-1", []string{"grp-eng"}, "prov-1", "g-1") + open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail + expectPolicies(mockStore, "acc-1", restricted, open) + expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o")) + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + Model: "claude-opus-4", + }) + require.NoError(t, err) + assert.True(t, res.Allow, "an un-guardrailed policy for the same (provider, group) must leave the request unrestricted") + assert.Equal(t, "pol-open", res.SelectedPolicyID, "the unrestricted policy must be the one that pays") +} + +// TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups is the false-allow fix: a +// model allowlisted only for grp-b must not be usable by a grp-a caller. The +// selector considers only policies applicable to the caller's groups. +func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + polA := guardedPolicy("pol-a", "acc-1", []string{"grp-a"}, "prov-1", "g-a") + polB := guardedPolicy("pol-b", "acc-1", []string{"grp-b"}, "prov-1", "g-b") + expectPolicies(mockStore, "acc-1", polA, polB) + expectGuardrails(mockStore, "acc-1", + allowlistGuardrail("g-a", "acc-1", "gpt-4o"), + allowlistGuardrail("g-b", "acc-1", "claude-opus-4"), + ) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-a"}, + ProviderID: "prov-1", + Model: "claude-opus-4", // only allowed for grp-b + }) + require.NoError(t, err) + assert.False(t, res.Allow, "grp-b's allowlisted model must not leak to a grp-a caller") + assert.Equal(t, denyCodeModelBlocked, res.DenyCode) +} + +// TestSelectPolicy_UndeterminedModelFailsClosed proves the fail-closed contract +// mirrors the proxy: with a restricted applicable policy and an empty model +// (e.g. a path-routed shape the parser couldn't map), the request is denied. +func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1") + expectPolicies(mockStore, "acc-1", policy) + expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o")) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + Model: "", // undetermined + }) + require.NoError(t, err) + assert.False(t, res.Allow, "an undetermined model must fail closed against a restricted policy") + assert.Equal(t, denyCodeModelBlocked, res.DenyCode) +} + +// TestSelectPolicy_DisabledAllowlistDoesNotRestrict proves a guardrail whose +// model allowlist is disabled imposes no model restriction, even though the +// policy references it. +func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1") + disabled := &types.Guardrail{ + ID: "g-1", + AccountID: "acc-1", + Checks: types.GuardrailChecks{ + ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}, + }, + } + expectPolicies(mockStore, "acc-1", policy) + expectGuardrails(mockStore, "acc-1", disabled) + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + Model: "anything-goes", + }) + require.NoError(t, err) + assert.True(t, res.Allow, "a disabled allowlist must not restrict the model") + assert.Equal(t, "pol-A", res.SelectedPolicyID) +} + +// TestSelectPolicy_UnionAcrossPolicyGuardrails proves a policy with multiple +// allowlist guardrails permits the union of their models (not just the first). +func TestSelectPolicy_UnionAcrossPolicyGuardrails(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1", "g-2") + expectPolicies(mockStore, "acc-1", policy) + expectGuardrails(mockStore, "acc-1", + allowlistGuardrail("g-1", "acc-1", "gpt-4o"), + allowlistGuardrail("g-2", "acc-1", "claude-opus-4"), + ) + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + Model: "claude-opus-4", // only in the second guardrail's list + }) + require.NoError(t, err) + assert.True(t, res.Allow, "a model in any of the policy's allowlist guardrails must be permitted") + assert.Equal(t, "pol-A", res.SelectedPolicyID) +} + +// TestSelectPolicy_GuardrailLookupErrorPropagates proves a store failure while +// resolving the candidate policies' guardrails surfaces as an error, not a +// silent allow/deny. +func TestSelectPolicy_GuardrailLookupErrorPropagates(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1") + expectPolicies(mockStore, "acc-1", policy) + mockStore.EXPECT(). + GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), "acc-1"). + Return(nil, errors.New("store unavailable")) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + Model: "gpt-4o", + }) + require.Error(t, err, "a guardrail-lookup failure must surface as an error") + assert.Nil(t, res) +} + +// TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted proves a +// policy referencing a guardrail ID absent from the account's set (a stale +// reference) imposes no model restriction — same as no guardrail. +func TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-missing") + expectPolicies(mockStore, "acc-1", policy) + expectGuardrails(mockStore, "acc-1") + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + Model: "anything-goes", + }) + require.NoError(t, err) + assert.True(t, res.Allow, "an orphaned guardrail reference must not restrict the model") + assert.Equal(t, "pol-A", res.SelectedPolicyID) +} + +// TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter proves the model +// gate narrows candidates before cap scoring: the permitting policy is selected +// even though the blocked one has a larger, more attractive cap. +func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + polBig := guardedPolicy("pol-big", "acc-1", []string{"grp-eng"}, "prov-1", "g-restrict") + polBig.Limits = types.PolicyLimits{ + TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 1_000_000, WindowSeconds: 3600}, + } + polSmall := guardedPolicy("pol-small", "acc-1", []string{"grp-eng"}, "prov-1", "g-permit") + polSmall.Limits = types.PolicyLimits{ + TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 100, WindowSeconds: 3600}, + } + expectPolicies(mockStore, "acc-1", polBig, polSmall) + expectGuardrails(mockStore, "acc-1", + allowlistGuardrail("g-restrict", "acc-1", "gpt-4o"), + allowlistGuardrail("g-permit", "acc-1", "claude-opus-4"), + ) + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + Model: "claude-opus-4", // only pol-small's guardrail permits this + }) + require.NoError(t, err) + assert.True(t, res.Allow) + assert.Equal(t, "pol-small", res.SelectedPolicyID, + "the model filter must exclude pol-big before cap scoring") +} diff --git a/management/internals/modules/agentnetwork/synthesizer.go b/management/internals/modules/agentnetwork/synthesizer.go index 95fe91773..169bdd4fd 100644 --- a/management/internals/modules/agentnetwork/synthesizer.go +++ b/management/internals/modules/agentnetwork/synthesizer.go @@ -235,7 +235,12 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([ mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID) applyAccountCollectionControls(&mergedGuardrails, settings) - guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails) + // The proxy guardrail is a per-provider fail-closed backstop; the + // authoritative per-policy/group decision is management's + // SelectPolicyForRequest. A provider lands in this map only when every + // authorising policy restricts models. + providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID) + guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture) if err != nil { return nil, err } @@ -780,10 +785,12 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt // guardrailConfig is the JSON shape the proxy-side llm_guardrail // middleware expects. Mirrors the proxy registration documented in -// the management→proxy contract. +// the management→proxy contract. provider_allowlists is keyed by the +// resolved provider id llm_router stamps; a provider absent from the map is +// unrestricted at the proxy layer. type guardrailConfig struct { - ModelAllowlist []string `json:"model_allowlist,omitempty"` - PromptCapture guardrailPromptCapture `json:"prompt_capture"` + ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"` + PromptCapture guardrailPromptCapture `json:"prompt_capture"` } type guardrailPromptCapture struct { @@ -828,13 +835,10 @@ func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Se merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii } -func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) { +func marshalGuardrailConfig(providerAllowlists map[string][]string, capture MergedPromptCapture) ([]byte, error) { cfg := guardrailConfig{ - ModelAllowlist: merged.ModelAllowlist, - PromptCapture: guardrailPromptCapture{ - Enabled: merged.PromptCapture.Enabled, - RedactPii: merged.PromptCapture.RedactPii, - }, + ProviderAllowlists: providerAllowlists, + PromptCapture: guardrailPromptCapture(capture), } out, err := json.Marshal(cfg) if err != nil { @@ -843,6 +847,74 @@ func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) { return out, nil } +// buildProviderAllowlists returns the proxy's per-provider backstop: a provider +// is included only when every authorising policy restricts models (their union); +// if any leaves it unrestricted it is omitted, so management decides per group. +func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string { + type providerAcc struct { + models map[string]struct{} + anyUnrestricted bool + } + accs := make(map[string]*providerAcc) + for _, p := range policies { + if p == nil { + continue + } + restricted, models := policyModelAllowlist(p, byID) + for _, providerID := range p.DestinationProviderIDs { + if providerID == "" { + continue + } + acc, ok := accs[providerID] + if !ok { + acc = &providerAcc{models: make(map[string]struct{})} + accs[providerID] = acc + } + if !restricted { + acc.anyUnrestricted = true + continue + } + for _, m := range models { + acc.models[m] = struct{}{} + } + } + } + out := make(map[string][]string, len(accs)) + for providerID, acc := range accs { + if acc.anyUnrestricted { + continue + } + models := make([]string, 0, len(acc.models)) + for m := range acc.models { + models = append(models, m) + } + sort.Strings(models) + out[providerID] = models + } + return out +} + +// policyModelAllowlist reports whether a policy restricts models (has an +// allowlist-enabled guardrail) and the union of allowed models. Models are +// verbatim; the proxy factory lowercases/trims them at decode time. +func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) { + restricted := false + var models []string + for _, gID := range p.GuardrailIDs { + g, ok := byID[gID] + if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled { + continue + } + restricted = true + for _, m := range g.Checks.ModelAllowlist.Models { + if m != "" { + models = append(models, m) + } + } + } + return restricted, models +} + // buildAccountService composes the per-account gateway Service. The // target carries the noop placeholder URL — the router middleware // rewrites every request to the matched provider's upstream before the @@ -986,38 +1058,11 @@ func unionSourceGroups(policies []*types.Policy) []string { return out } -// MergedGuardrails is the JSON shape passed to the proxy via the -// guardrail middleware's config_json. Mirrors the proxy-side -// expectations and is intentionally distinct from -// types.GuardrailChecks so we can evolve either side independently. +// MergedGuardrails is the synthesiser's fold target. Only prompt capture is +// merged here — the model allowlist is emitted per-provider, and +// token/budget/retention moved onto Policy.Limits and account Settings. type MergedGuardrails struct { - ModelAllowlist []string `json:"model_allowlist,omitempty"` - TokenLimits MergedTokenLimits `json:"token_limits"` - Budget MergedBudget `json:"budget"` - PromptCapture MergedPromptCapture `json:"prompt_capture"` - Retention MergedRetention `json:"retention"` -} - -type MergedTokenLimits struct { - Hourly *MergedTokenWindow `json:"hourly,omitempty"` - Daily *MergedTokenWindow `json:"daily,omitempty"` - Monthly *MergedTokenWindow `json:"monthly,omitempty"` -} - -type MergedTokenWindow struct { - MaxInputTokens int `json:"max_input_tokens,omitempty"` - MaxOutputTokens int `json:"max_output_tokens,omitempty"` -} - -type MergedBudget struct { - Hourly *MergedBudgetWindow `json:"hourly,omitempty"` - Daily *MergedBudgetWindow `json:"daily,omitempty"` - Monthly *MergedBudgetWindow `json:"monthly,omitempty"` -} - -type MergedBudgetWindow struct { - SoftCapUSD float64 `json:"soft_cap_usd,omitempty"` - HardCapUSD float64 `json:"hard_cap_usd,omitempty"` + PromptCapture MergedPromptCapture } type MergedPromptCapture struct { @@ -1025,64 +1070,31 @@ type MergedPromptCapture struct { RedactPii bool `json:"redact_pii"` } -type MergedRetention struct { - Enabled bool `json:"enabled"` - Days int `json:"days"` -} - -// mergeGuardrails computes the effective guardrail spec applied at the -// proxy, given the referencing policies and the account's guardrail -// catalogue. Policy enabled-ness is the caller's responsibility — only -// enabled policies should be passed in. +// mergeGuardrails folds the referencing policies' guardrails into the +// prompt-capture decision only. The model allowlist is enforced per-policy/group +// in management and shipped per-provider; token/budget/retention live off +// guardrails now. // -// Merge rules: -// - Model allowlist: union of allowlists across policies that enable it. -// - Token / Budget: most-restrictive (min of non-zero caps) per window. -// - Prompt capture: enabled if any policy enables it; redact_pii sticks -// if any enabling policy turns it on. -// - Retention: enabled if any enables it; smallest non-zero days wins. +// Merge rule — prompt capture: enabled if any policy enables it; redact_pii +// sticks if any enabling policy turns it on. func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails { merged := MergedGuardrails{} - allowlist := make(map[string]struct{}) - allowlistEnabled := false - for _, policy := range policies { for _, gID := range policy.GuardrailIDs { g, ok := byID[gID] if !ok || g == nil { continue } - mergeGuardrail(g, &merged, allowlist, &allowlistEnabled) + mergeGuardrail(g, &merged) } } - - if allowlistEnabled { - merged.ModelAllowlist = make([]string, 0, len(allowlist)) - for m := range allowlist { - merged.ModelAllowlist = append(merged.ModelAllowlist, m) - } - sort.Strings(merged.ModelAllowlist) - } return merged } -// mergeGuardrail folds a single guardrail's enabled checks into the -// running merge: model-allowlist models join the shared set (and flip -// allowlistEnabled), and prompt-capture / redact-pii stick once any -// enabling guardrail turns them on. -// -// TokenLimits, Budget, and Retention have moved off guardrails — token -// and budget caps now live on the Policy itself (Policy.Limits) and -// retention moves to account-level Settings — so they are not merged here. -func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) { - if g.Checks.ModelAllowlist.Enabled { - *allowlistEnabled = true - for _, m := range g.Checks.ModelAllowlist.Models { - if m != "" { - allowlist[m] = struct{}{} - } - } - } +// mergeGuardrail folds a single guardrail's prompt-capture settings into the +// running merge: enabled / redact-pii stick once any enabling guardrail turns +// them on. +func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) { if g.Checks.PromptCapture.Enabled { merged.PromptCapture.Enabled = true if g.Checks.PromptCapture.RedactPii { diff --git a/management/internals/modules/agentnetwork/synthesizer_guardrail_realstore_test.go b/management/internals/modules/agentnetwork/synthesizer_guardrail_realstore_test.go index 8ed4910da..8b13f78ac 100644 --- a/management/internals/modules/agentnetwork/synthesizer_guardrail_realstore_test.go +++ b/management/internals/modules/agentnetwork/synthesizer_guardrail_realstore_test.go @@ -81,8 +81,8 @@ func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testi require.Len(t, services, 1) cfg := decodeServiceGuardrailConfig(t, services[0]) - assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist, - "model allowlist is a pure policy guardrail and must always reach the config") + assert.Equal(t, map[string][]string{"prov-1": {"gpt-5.4"}}, cfg.ProviderAllowlists, + "model allowlist is a pure policy guardrail and must reach the per-provider config") assert.False(t, cfg.PromptCapture.Enabled, "prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail") } @@ -172,7 +172,7 @@ func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) { require.Len(t, services, 1, "exactly one synth service expected") cfg := decodeServiceGuardrailConfig(t, services[0]) - assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist") + assert.Empty(t, cfg.ProviderAllowlists, "no guardrail → provider unrestricted (absent from map)") assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default") assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default") } diff --git a/management/internals/modules/agentnetwork/synthesizer_provider_allowlist_test.go b/management/internals/modules/agentnetwork/synthesizer_provider_allowlist_test.go new file mode 100644 index 000000000..2cfc0db8c --- /dev/null +++ b/management/internals/modules/agentnetwork/synthesizer_provider_allowlist_test.go @@ -0,0 +1,95 @@ +package agentnetwork + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" +) + +// policyForProviders builds an enabled policy authorising the given providers +// under the given guardrails (both optional). Groups are irrelevant to +// buildProviderAllowlists, which keys purely on destination provider. +func policyForProviders(id string, guardrailIDs []string, providerIDs ...string) *types.Policy { + return &types.Policy{ + ID: id, + Enabled: true, + DestinationProviderIDs: providerIDs, + GuardrailIDs: guardrailIDs, + } +} + +func TestBuildProviderAllowlists(t *testing.T) { + byID := map[string]*types.Guardrail{ + "g-4o": allowlistGuardrail("g-4o", "acc-1", "gpt-4o"), + "g-opus": allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"), + "g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}}, + } + + t.Run("all authorising policies restrict yields per-provider union", func(t *testing.T) { + policies := []*types.Policy{ + policyForProviders("p1", []string{"g-4o"}, "prov-x"), + policyForProviders("p2", []string{"g-opus"}, "prov-x"), + } + got := buildProviderAllowlists(policies, byID) + assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got, + "a provider every policy restricts carries the sorted union of their models") + }) + + t.Run("any un-guardrailed policy leaves the provider unrestricted (omitted)", func(t *testing.T) { + policies := []*types.Policy{ + policyForProviders("p1", []string{"g-4o"}, "prov-x"), + policyForProviders("p2", nil, "prov-x"), // no guardrail + } + got := buildProviderAllowlists(policies, byID) + assert.NotContains(t, got, "prov-x", + "a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted") + }) + + t.Run("a disabled allowlist counts as unrestricted", func(t *testing.T) { + policies := []*types.Policy{ + policyForProviders("p1", []string{"g-disabled"}, "prov-x"), + } + got := buildProviderAllowlists(policies, byID) + assert.NotContains(t, got, "prov-x", + "a policy whose only guardrail has a disabled allowlist is unrestricted") + }) + + t.Run("providers are isolated from one another", func(t *testing.T) { + policies := []*types.Policy{ + policyForProviders("p1", []string{"g-4o"}, "prov-x"), + policyForProviders("p2", []string{"g-opus"}, "prov-y"), + } + got := buildProviderAllowlists(policies, byID) + assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model") + assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model") + }) + + t.Run("one policy authorising two providers restricts both", func(t *testing.T) { + policies := []*types.Policy{ + policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"), + } + got := buildProviderAllowlists(policies, byID) + assert.Equal(t, []string{"gpt-4o"}, got["prov-x"]) + assert.Equal(t, []string{"gpt-4o"}, got["prov-y"]) + }) + + t.Run("union across a single policy's guardrails", func(t *testing.T) { + policies := []*types.Policy{ + policyForProviders("p1", []string{"g-4o", "g-opus"}, "prov-x"), + } + got := buildProviderAllowlists(policies, byID) + assert.ElementsMatch(t, []string{"claude-opus-4", "gpt-4o"}, got["prov-x"], + "a policy's own multiple allowlist guardrails union together") + }) + + t.Run("an enabled allowlist with no models denies everything", func(t *testing.T) { + empty := map[string]*types.Guardrail{"g-empty": allowlistGuardrail("g-empty", "acc-1")} + got := buildProviderAllowlists([]*types.Policy{ + policyForProviders("p1", []string{"g-empty"}, "prov-x"), + }, empty) + assert.Equal(t, map[string][]string{"prov-x": {}}, got, + "an enabled-but-empty allowlist is restricted with an empty set, not unrestricted") + }) +} diff --git a/management/internals/modules/agentnetwork/synthesizer_test.go b/management/internals/modules/agentnetwork/synthesizer_test.go index 206d1d12a..8a18a9b59 100644 --- a/management/internals/modules/agentnetwork/synthesizer_test.go +++ b/management/internals/modules/agentnetwork/synthesizer_test.go @@ -1031,8 +1031,12 @@ func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *t var cfg guardrailConfig require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly") - assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist, - "model allowlist union must keep both models") + // Both policies restrict the same provider, so the per-provider backstop + // carries the union of their models — a coarse gate that management's + // per-policy/group check narrows; it only blocks models outside the union + // when management is down. + assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ProviderAllowlists["prov-1"], + "per-provider allowlist union must keep both models") } func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) { diff --git a/management/internals/shared/grpc/proxy.go b/management/internals/shared/grpc/proxy.go index b289f8c71..8f24de116 100644 --- a/management/internals/shared/grpc/proxy.go +++ b/management/internals/shared/grpc/proxy.go @@ -285,6 +285,7 @@ func (s *ProxyServiceServer) CheckLLMPolicyLimits(ctx context.Context, req *prot UserID: req.GetUserId(), GroupIDs: req.GetGroupIds(), ProviderID: req.GetProviderId(), + Model: req.GetModel(), }) if err != nil { log.WithContext(ctx).Errorf("select policy for request: %v", err) diff --git a/management/internals/shared/grpc/proxy_llm_policy_limits_test.go b/management/internals/shared/grpc/proxy_llm_policy_limits_test.go new file mode 100644 index 000000000..c293c39ea --- /dev/null +++ b/management/internals/shared/grpc/proxy_llm_policy_limits_test.go @@ -0,0 +1,138 @@ +package grpc + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + grpcstatus "google.golang.org/grpc/status" + + "github.com/netbirdio/netbird/management/internals/modules/agentnetwork" + "github.com/netbirdio/netbird/shared/management/proto" +) + +// fakeAgentNetworkLimits records the PolicySelectionInput it was invoked with +// and returns a pre-programmed result, so tests can assert what the handler +// forwards to the selector. +type fakeAgentNetworkLimits struct { + gotInput agentnetwork.PolicySelectionInput + result *agentnetwork.PolicySelectionResult + err error +} + +func (f *fakeAgentNetworkLimits) SelectPolicyForRequest(_ context.Context, in agentnetwork.PolicySelectionInput) (*agentnetwork.PolicySelectionResult, error) { + f.gotInput = in + if f.err != nil { + return nil, f.err + } + return f.result, nil +} + +func (f *fakeAgentNetworkLimits) RecordUsage(_ context.Context, _ agentnetwork.RecordUsageInput) error { + return nil +} + +// TestCheckLLMPolicyLimits_ForwardsModelToSelector proves the wiring added here: +// the model the proxy extracted must reach the selector's Model unchanged, +// alongside the account/user/group/provider fields. +func TestCheckLLMPolicyLimits_ForwardsModelToSelector(t *testing.T) { + fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true, SelectedPolicyID: "pol-1"}} + s := &ProxyServiceServer{} + s.SetAgentNetworkLimitsService(fake) + + req := &proto.CheckLLMPolicyLimitsRequest{ + AccountId: "acc-1", + UserId: "user-1", + GroupIds: []string{"grp-a", "grp-b"}, + ProviderId: "prov-1", + Model: "claude-opus-4", + } + + resp, err := s.CheckLLMPolicyLimits(context.Background(), req) + require.NoError(t, err) + require.NotNil(t, resp) + + assert.Equal(t, "acc-1", fake.gotInput.AccountID) + assert.Equal(t, "user-1", fake.gotInput.UserID) + assert.Equal(t, []string{"grp-a", "grp-b"}, fake.gotInput.GroupIDs) + assert.Equal(t, "prov-1", fake.gotInput.ProviderID) + assert.Equal(t, "claude-opus-4", fake.gotInput.Model, + "the request's model must be forwarded to the selector") +} + +// TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty proves an undetermined +// model (empty string) is forwarded as-is; the selector decides how to treat it. +func TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty(t *testing.T) { + fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true}} + s := &ProxyServiceServer{} + s.SetAgentNetworkLimitsService(fake) + + req := &proto.CheckLLMPolicyLimitsRequest{ + AccountId: "acc-1", + ProviderId: "prov-1", + } + + _, err := s.CheckLLMPolicyLimits(context.Background(), req) + require.NoError(t, err) + assert.Equal(t, "", fake.gotInput.Model, "an absent model must be forwarded as empty, not fabricated") +} + +// TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode proves the deny +// envelope surfaces the model-allowlist deny code + reason through the response. +func TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode(t *testing.T) { + fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{ + Allow: false, + DenyCode: "llm_policy.model_blocked", + DenyReason: `model "claude-opus-4" is not permitted by any applicable policy allowlist`, + }} + s := &ProxyServiceServer{} + s.SetAgentNetworkLimitsService(fake) + + resp, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{ + AccountId: "acc-1", + ProviderId: "prov-1", + Model: "claude-opus-4", + }) + require.NoError(t, err) + require.NotNil(t, resp) + assert.Equal(t, "deny", resp.Decision) + assert.Equal(t, "llm_policy.model_blocked", resp.DenyCode) + assert.NotEmpty(t, resp.DenyReason) + assert.Empty(t, resp.SelectedPolicyId, "a denied request must carry no selected policy") +} + +// TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal proves a selector +// failure surfaces as an Internal gRPC error rather than a silent allow. +func TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal(t *testing.T) { + fake := &fakeAgentNetworkLimits{err: errors.New("boom")} + s := &ProxyServiceServer{} + s.SetAgentNetworkLimitsService(fake) + + _, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{ + AccountId: "acc-1", + ProviderId: "prov-1", + Model: "gpt-4o", + }) + require.Error(t, err) + st, ok := grpcstatus.FromError(err) + require.True(t, ok) + assert.Equal(t, codes.Internal, st.Code(), "selector errors must never fail open on the hot path") +} + +// TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented locks the +// fallback: with no limits service wired the RPC returns Unimplemented. +func TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented(t *testing.T) { + s := &ProxyServiceServer{} + + _, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{ + AccountId: "acc-1", + ProviderId: "prov-1", + }) + require.Error(t, err) + st, ok := grpcstatus.FromError(err) + require.True(t, ok) + assert.Equal(t, codes.Unimplemented, st.Code()) +} diff --git a/proxy/internal/auth/tunnel_cache.go b/proxy/internal/auth/tunnel_cache.go index 10b671d82..185c53c62 100644 --- a/proxy/internal/auth/tunnel_cache.go +++ b/proxy/internal/auth/tunnel_cache.go @@ -3,20 +3,30 @@ package auth import ( "context" "net/netip" + "os" + "strings" "sync" "time" + log "github.com/sirupsen/logrus" "golang.org/x/sync/singleflight" "github.com/netbirdio/netbird/proxy/internal/types" "github.com/netbirdio/netbird/shared/management/proto" ) -// tunnelCacheTTL caps how long a positive ValidateTunnelPeer result is -// reused before re-fetching from management. 5 minutes balances freshness -// against management load on busy mesh networks. +// tunnelCacheTTL is the default cap on how long a positive ValidateTunnelPeer +// result is reused before re-fetching from management. 5 minutes balances +// freshness against management load on busy mesh networks. Override it with +// envTunnelCacheTTL when an account needs authorization changes to take effect +// sooner (at the cost of more ValidateTunnelPeer RPCs). const tunnelCacheTTL = 300 * time.Second +// envTunnelCacheTTL overrides tunnelCacheTTL. The value is a Go duration string +// (e.g. "30s", "2m"); an unset, unparseable, or non-positive value keeps the +// default. +const envTunnelCacheTTL = "NB_PROXY_TUNNEL_CACHE_TTL" + // tunnelCachePerAccount caps the number of cached identities per account. // Bounded eviction avoids memory growth in pathological cases (huge peer // roster, brief request bursts) while staying generous for normal use. @@ -60,16 +70,35 @@ type accountBucket struct { order []tunnelCacheKey } -// newTunnelValidationCache constructs a cache with default TTL and bounds. +// newTunnelValidationCache constructs a cache with the configured TTL +// (envTunnelCacheTTL override or default) and default bounds. func newTunnelValidationCache() *tunnelValidationCache { return &tunnelValidationCache{ entries: make(map[types.AccountID]*accountBucket), - ttl: tunnelCacheTTL, + ttl: tunnelCacheTTLFromEnv(), maxSize: tunnelCachePerAccount, now: time.Now, } } +// tunnelCacheTTLFromEnv returns the tunnel-cache TTL, honoring the +// envTunnelCacheTTL override. The override must be a positive Go duration +// string (e.g. "30s", "2m"); anything unset, unparseable, or non-positive +// falls back to tunnelCacheTTL. +func tunnelCacheTTLFromEnv() time.Duration { + raw := strings.TrimSpace(os.Getenv(envTunnelCacheTTL)) + if raw == "" { + return tunnelCacheTTL + } + d, err := time.ParseDuration(raw) + if err != nil || d <= 0 { + log.Warnf("ignoring invalid %s=%q (want a positive Go duration like 30s or 2m); using default %s", + envTunnelCacheTTL, raw, tunnelCacheTTL) + return tunnelCacheTTL + } + return d +} + // get returns a cached response for the key, or nil when missing or // expired. Expired entries are evicted lazily on read. func (c *tunnelValidationCache) get(key tunnelCacheKey) *proto.ValidateTunnelPeerResponse { diff --git a/proxy/internal/auth/tunnel_cache_test.go b/proxy/internal/auth/tunnel_cache_test.go index 1a63dc107..d663b91e7 100644 --- a/proxy/internal/auth/tunnel_cache_test.go +++ b/proxy/internal/auth/tunnel_cache_test.go @@ -169,3 +169,32 @@ func TestTunnelCache_BoundedSizeEvictsOldest(t *testing.T) { assert.NotNil(t, cache.get(keys[1]), "second-newest must remain cached") assert.NotNil(t, cache.get(keys[2]), "newest must remain cached") } + +func TestTunnelCacheTTLFromEnv(t *testing.T) { + t.Run("unset uses default", func(t *testing.T) { + t.Setenv(envTunnelCacheTTL, "") + assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv()) + }) + t.Run("valid duration overrides", func(t *testing.T) { + t.Setenv(envTunnelCacheTTL, "45s") + assert.Equal(t, 45*time.Second, tunnelCacheTTLFromEnv()) + }) + t.Run("whitespace trimmed", func(t *testing.T) { + t.Setenv(envTunnelCacheTTL, " 2m ") + assert.Equal(t, 2*time.Minute, tunnelCacheTTLFromEnv()) + }) + t.Run("unparseable uses default", func(t *testing.T) { + t.Setenv(envTunnelCacheTTL, "nonsense") + assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv()) + }) + t.Run("non-positive uses default", func(t *testing.T) { + t.Setenv(envTunnelCacheTTL, "0s") + assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv()) + t.Setenv(envTunnelCacheTTL, "-30s") + assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv()) + }) + t.Run("constructor honors override", func(t *testing.T) { + t.Setenv(envTunnelCacheTTL, "90s") + assert.Equal(t, 90*time.Second, newTunnelValidationCache().ttl) + }) +} diff --git a/proxy/internal/middleware/builtin/llm_guardrail/factory.go b/proxy/internal/middleware/builtin/llm_guardrail/factory.go index 6dd2a8e8d..9cec0dc68 100644 --- a/proxy/internal/middleware/builtin/llm_guardrail/factory.go +++ b/proxy/internal/middleware/builtin/llm_guardrail/factory.go @@ -10,11 +10,15 @@ import ( ) // Config is the JSON-decoded shape accepted by the factory. The -// runtime path consumes the normalised allowlist; raw config is not +// runtime path consumes the normalised allowlists; raw config is not // retained beyond construction. type Config struct { - ModelAllowlist []string `json:"model_allowlist"` - PromptCapture PromptCapture `json:"prompt_capture"` + // ProviderAllowlists maps a resolved provider id (KeyLLMResolvedProviderID) to + // its model allowlist. A provider present is restricted to those models; one + // absent is unrestricted. Kept per-provider so one provider's list can't leak + // onto another. + ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"` + PromptCapture PromptCapture `json:"prompt_capture"` } // PromptCapture toggles the optional prompt capture + redaction step @@ -54,21 +58,28 @@ func isEmptyJSON(raw []byte) bool { return false } -// normaliseConfig lowercases and trims allowlist entries so the runtime -// match is case-insensitive. Empty entries are dropped. +// normaliseConfig lowercases and trims allowlist entries for case-insensitive +// matching; empty entries drop. A provider whose entries all drop keeps an empty +// (non-nil) list — "deny every model" — distinct from an absent provider +// (unrestricted). func normaliseConfig(cfg Config) Config { - if len(cfg.ModelAllowlist) == 0 { + if len(cfg.ProviderAllowlists) == 0 { + cfg.ProviderAllowlists = nil return cfg } - cleaned := make([]string, 0, len(cfg.ModelAllowlist)) - for _, entry := range cfg.ModelAllowlist { - n := normaliseModel(entry) - if n == "" { - continue + cleaned := make(map[string][]string, len(cfg.ProviderAllowlists)) + for provider, models := range cfg.ProviderAllowlists { + list := make([]string, 0, len(models)) + for _, entry := range models { + n := normaliseModel(entry) + if n == "" { + continue + } + list = append(list, n) } - cleaned = append(cleaned, n) + cleaned[provider] = list } - cfg.ModelAllowlist = cleaned + cfg.ProviderAllowlists = cleaned return cfg } diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go index eded877ac..1863aff20 100644 --- a/proxy/internal/middleware/builtin/llm_guardrail/middleware.go +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go @@ -83,8 +83,9 @@ func (m *Middleware) MutationsSupported() bool { return false } // prompt capture only affects the metadata emitted alongside an allow. func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) { model, modelPresent := lookupMetadata(in.Metadata, middleware.KeyLLMModel) + providerID, _ := lookupMetadata(in.Metadata, middleware.KeyLLMResolvedProviderID) - if denial := m.evaluateAllowlist(model, modelPresent); denial != nil { + if denial := m.evaluateAllowlist(providerID, model, modelPresent); denial != nil { return denial, nil } @@ -110,20 +111,32 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar // is a no-op. func (m *Middleware) Close() error { return nil } -// evaluateAllowlist returns a deny Output when the configured allowlist -// rejects the model. A nil return means the request should proceed. -func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middleware.Output { - if len(m.cfg.ModelAllowlist) == 0 { +// evaluateAllowlist denies when the resolved provider's allowlist rejects the +// model; nil means proceed. Scoped to the provider llm_router resolved, so an +// unrestricted provider (absent from config) is never caught by another's list. +func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bool) *middleware.Output { + if len(m.cfg.ProviderAllowlists) == 0 { return nil } - // Fail closed: with an allowlist configured, a request whose model the - // upstream parser could not extract (absent or empty) must be denied rather - // than allowed. This is what enforces the allowlist for URL/path-routed - // providers (Bedrock, Vertex, ...) whose model lives outside the JSON body. + // Restrictions exist but the resolved provider is unknown, so we can't tell + // if this request targets a restricted provider — fail closed. llm_router + // normally stamps the provider first, so this is a defensive guard. + if providerID == "" { + return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown) + } + allowlist, restricted := m.cfg.ProviderAllowlists[providerID] + if !restricted { + // This provider has no allowlist (some authorising policy left it + // unrestricted); management owns any per-policy/group decision. + return nil + } + // Fail closed: with an allowlist in effect for this provider, a request whose + // model the parser couldn't extract (absent/empty) is denied. This enforces + // the allowlist for path-routed providers (Bedrock, Vertex) with no body model. if !modelPresent || normaliseModel(model) == "" { return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown) } - if m.modelInAllowlist(model) { + if modelInAllowlist(allowlist, model) { return nil } return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel) @@ -151,14 +164,15 @@ func denyModel(model, code, message, reason string) *middleware.Output { } } -// modelInAllowlist reports whether the model matches any allowlist -// entry under the case-insensitive, trim-tolerant comparison rule. -func (m *Middleware) modelInAllowlist(model string) bool { +// modelInAllowlist reports whether the model matches any entry in the supplied +// (already-normalised) allowlist under the case-insensitive, trim-tolerant +// comparison rule. +func modelInAllowlist(allowlist []string, model string) bool { normalised := normaliseModel(model) if normalised == "" { return false } - for _, allowed := range m.cfg.ModelAllowlist { + for _, allowed := range allowlist { if allowed == normalised { return true } diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go index cd7e256dd..5f35fefd3 100644 --- a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go @@ -26,6 +26,25 @@ func newInput(meta ...middleware.KV) *middleware.Input { return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: meta} } +const ( + testProvider = "prov-1" + otherProvider = "prov-2" +) + +// providerCfg builds a Config restricting testProvider to the given models. +func providerCfg(models ...string) Config { + return Config{ProviderAllowlists: map[string][]string{testProvider: models}} +} + +// newInputProvider builds an input that carries a resolved provider id (as +// llm_router would stamp) plus any extra metadata. +func newInputProvider(provider string, meta ...middleware.KV) *middleware.Input { + all := make([]middleware.KV, 0, len(meta)+1) + all = append(all, middleware.KV{Key: middleware.KeyLLMResolvedProviderID, Value: provider}) + all = append(all, meta...) + return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: all} +} + func TestMiddlewareIdentity(t *testing.T) { mw := New(Config{}) assert.Equal(t, ID, mw.ID(), "middleware ID must be llm_guardrail") @@ -47,12 +66,12 @@ func TestMiddlewareIdentity(t *testing.T) { func TestAllowlistEmptyAllowsAnyModel(t *testing.T) { mw := New(Config{}) - out, err := mw.Invoke(context.Background(), newInput( + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"}, )) require.NoError(t, err) require.NotNil(t, out) - assert.Equal(t, middleware.DecisionAllow, out.Decision, "empty allowlist must allow any model") + assert.Equal(t, middleware.DecisionAllow, out.Decision, "no provider allowlists must allow any model") v, ok := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision) require.True(t, ok, "decision metadata must be emitted") assert.Equal(t, "allow", v, "decision must be allow") @@ -62,8 +81,8 @@ func TestAllowlistEmptyAllowsAnyModel(t *testing.T) { } func TestAllowlistMatchAllows(t *testing.T) { - mw := New(Config{ModelAllowlist: []string{"gpt-4o", "claude-opus-4"}}) - out, err := mw.Invoke(context.Background(), newInput( + mw := New(providerCfg("gpt-4o", "claude-opus-4")) + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"}, )) require.NoError(t, err) @@ -71,8 +90,8 @@ func TestAllowlistMatchAllows(t *testing.T) { } func TestAllowlistMissDenies(t *testing.T) { - mw := New(Config{ModelAllowlist: []string{"gpt-4o"}}) - out, err := mw.Invoke(context.Background(), newInput( + mw := New(providerCfg("gpt-4o")) + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"}, )) require.NoError(t, err) @@ -91,10 +110,10 @@ func TestAllowlistMissDenies(t *testing.T) { } func TestAllowlistCaseInsensitive(t *testing.T) { - mw := New(Config{ModelAllowlist: []string{" GPT-4o ", "Claude-OPUS-4"}}) + mw := New(providerCfg(" GPT-4o ", "Claude-OPUS-4")) cases := []string{"gpt-4o", "GPT-4O", " claude-opus-4 "} for _, model := range cases { - out, err := mw.Invoke(context.Background(), newInput( + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, middleware.KV{Key: middleware.KeyLLMModel, Value: model}, )) require.NoError(t, err) @@ -103,14 +122,15 @@ func TestAllowlistCaseInsensitive(t *testing.T) { } func TestAllowlistMissingModelKeyDenies(t *testing.T) { - // Fail closed: with an allowlist configured, a request whose model the - // parser could not extract (URL/path-routed providers such as Bedrock or - // Vertex whose shape wasn't recognised) must be denied, not allowed. - mw := New(Config{ModelAllowlist: []string{"gpt-4o"}}) - out, err := mw.Invoke(context.Background(), newInput()) + // Fail closed: with an allowlist in effect for the resolved provider, a + // request whose model the parser could not extract (URL/path-routed + // providers such as Bedrock or Vertex whose shape wasn't recognised) must be + // denied, not allowed. + mw := New(providerCfg("gpt-4o")) + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider)) require.NoError(t, err) require.NotNil(t, out) - assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when an allowlist is set") + assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when the provider is restricted") assert.Equal(t, 403, out.DenyStatus, "deny status must be 403") require.NotNil(t, out.DenyReason, "deny reason must be populated") assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown") @@ -122,26 +142,101 @@ func TestAllowlistMissingModelKeyDenies(t *testing.T) { func TestAllowlistEmptyModelValueDenies(t *testing.T) { // A present-but-empty model is as undeterminable as an absent one. - mw := New(Config{ModelAllowlist: []string{"gpt-4o"}}) - out, err := mw.Invoke(context.Background(), newInput( + mw := New(providerCfg("gpt-4o")) + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, middleware.KV{Key: middleware.KeyLLMModel, Value: " "}, )) require.NoError(t, err) require.NotNil(t, out) - assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when an allowlist is set") + assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when the provider is restricted") require.NotNil(t, out.DenyReason, "deny reason must be populated") assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown") } func TestAllowlistEmptyListAllowsMissingModel(t *testing.T) { - // Without an allowlist there is nothing to enforce, so a missing model is - // still allowed — the fail-closed rule only applies when a list is set. + // Without any provider allowlists there is nothing to enforce, so a missing + // model is still allowed — the fail-closed rule only applies when a + // restriction is in effect. mw := New(Config{}) out, err := mw.Invoke(context.Background(), newInput()) require.NoError(t, err) assert.Equal(t, middleware.DecisionAllow, out.Decision, "no allowlist must allow even without a model") } +func TestUnrestrictedProviderAllowsAnyModel(t *testing.T) { + // The request resolved to otherProvider, which has no allowlist, so its + // traffic must not be caught by testProvider's restriction — the + // cross-provider-leak / false-deny guard. + mw := New(providerCfg("gpt-4o")) + out, err := mw.Invoke(context.Background(), newInputProvider(otherProvider, + middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"}, + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "an unrestricted provider must not inherit another provider's allowlist") +} + +func TestPerProviderAllowlistsAreIsolated(t *testing.T) { + // gpt-4o is allowed only on testProvider; claude-opus-4 only on + // otherProvider. A model allowlisted for one provider must not be usable on + // the other — the fail-closed layer never unions allowlists across providers. + mw := New(Config{ProviderAllowlists: map[string][]string{ + testProvider: {"gpt-4o"}, + otherProvider: {"claude-opus-4"}, + }}) + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, + middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"}, + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "claude-opus-4 is allowed only on otherProvider, not testProvider") + require.NotNil(t, out.DenyReason) + assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "cross-provider model must be blocked, not model_unknown") +} + +func TestRestrictionsButNoResolvedProviderFailsClosed(t *testing.T) { + // Restrictions exist for the account but the resolved provider id is absent, + // so the request cannot be scoped to a provider. Fail closed rather than + // wave it through. + mw := New(providerCfg("gpt-4o")) + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + )) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "missing resolved provider must fail closed when restrictions exist") + require.NotNil(t, out.DenyReason) + assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown") +} + +func TestEnabledButEmptyAllowlistDeniesEveryModel(t *testing.T) { + // An allowlist-enabled provider with zero models is distinct from an + // unrestricted (absent) provider: it must deny every model. + mw := New(providerCfg()) + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, + middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + )) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "an enabled-but-empty allowlist must deny every model") + require.NotNil(t, out.DenyReason) + assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked, not model_unknown") +} + +func TestFactoryAllEmptyEntriesDenyEveryModel(t *testing.T) { + // All the provider's entries are blank; they collapse to a non-nil empty + // list (deny everything for that provider), not "no restriction". + raw := []byte(`{"provider_allowlists":{"prov-1":[""," "]}}`) + mw, err := Factory{}.New(raw) + require.NoError(t, err) + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, + middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + )) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "all-blank allowlist entries must still restrict the provider") + require.NotNil(t, out.DenyReason) + assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked") +} + func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) { mw := New(Config{}) out, err := mw.Invoke(context.Background(), newInput( @@ -217,8 +312,8 @@ func TestFactoryAcceptsZeroConfigs(t *testing.T) { func TestFactoryDecodesValidConfig(t *testing.T) { cfg := Config{ - ModelAllowlist: []string{"gpt-4o"}, - PromptCapture: PromptCapture{Enabled: true, RedactPii: true}, + ProviderAllowlists: map[string][]string{testProvider: {"gpt-4o"}}, + PromptCapture: PromptCapture{Enabled: true, RedactPii: true}, } raw, err := json.Marshal(cfg) require.NoError(t, err, "marshalling test config must succeed") @@ -234,15 +329,15 @@ func TestFactoryRejectsMalformedJSON(t *testing.T) { } func TestFactoryNormalisesAllowlist(t *testing.T) { - raw := []byte(`{"model_allowlist":[" GPT-4o ","",""," Claude-3 "]}`) + raw := []byte(`{"provider_allowlists":{"prov-1":[" GPT-4o ","",""," Claude-3 "]}}`) mw, err := Factory{}.New(raw) require.NoError(t, err) - out, err := mw.Invoke(context.Background(), newInput( + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"}, )) require.NoError(t, err) assert.Equal(t, middleware.DecisionAllow, out.Decision, "factory must lowercase + trim allowlist entries") - out2, err := mw.Invoke(context.Background(), newInput( + out2, err := mw.Invoke(context.Background(), newInputProvider(testProvider, middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-3"}, )) require.NoError(t, err) diff --git a/proxy/internal/middleware/builtin/llm_limit_check/middleware.go b/proxy/internal/middleware/builtin/llm_limit_check/middleware.go index bebe4dca4..42ac56b9b 100644 --- a/proxy/internal/middleware/builtin/llm_limit_check/middleware.go +++ b/proxy/internal/middleware/builtin/llm_limit_check/middleware.go @@ -175,7 +175,7 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou DenyStatus: 403, DenyReason: &middleware.DenyReason{ Code: code, - Message: "LLM policy limit exceeded", + Message: denyMessageForCode(code), }, Metadata: []middleware.KV{ {Key: middleware.KeyLLMPolicyDecision, Value: "deny"}, @@ -184,6 +184,21 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou } } +// denyMessageForCode maps a management deny code to a public message. +// Model-allowlist rejections get a model-specific message matching the +// local guardrail; everything else keeps the generic quota wording. The +// message stays generic so it never leaks internal quota detail. +func denyMessageForCode(code string) string { + switch code { + case "llm_policy.model_blocked": + return "model is not in the policy allowlist" + case "llm_policy.model_unknown": + return "request model could not be determined for the policy allowlist" + default: + return "LLM policy limit exceeded" + } +} + // lookupKV returns the value associated with key, or the empty // string when absent. func lookupKV(kvs []middleware.KV, key string) string { diff --git a/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go b/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go index 2c26c2abe..87aa8e9e9 100644 --- a/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go @@ -115,6 +115,46 @@ func TestInvoke_DenyConvertsToProxyDeny(t *testing.T) { assert.NotContains(t, out.DenyReason.Message, "1000", "internal cap numbers must not reach the caller") } +// TestInvoke_ModelDenyMessages proves a model-allowlist rejection gets a +// model-specific public message rather than the generic quota wording, so a +// blocked or undetermined model reads consistently with the local guardrail. +func TestInvoke_ModelDenyMessages(t *testing.T) { + cases := []struct { + name string + code string + message string + }{ + {"blocked", "llm_policy.model_blocked", "model is not in the policy allowlist"}, + {"unknown", "llm_policy.model_unknown", "request model could not be determined for the policy allowlist"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + mgmt := &fakeMgmt{ + checkResp: &proto.CheckLLMPolicyLimitsResponse{ + Decision: "deny", + DenyCode: tc.code, + }, + } + m := New(mgmt, nil) + + out := runInvoke(t, m, &middleware.Input{ + AccountID: "acc-1", + UserGroups: []string{"grp-engineers"}, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"}, + {Key: middleware.KeyLLMModel, Value: "some-model"}, + }, + }) + + assert.Equal(t, middleware.DecisionDeny, out.Decision) + require.NotNil(t, out.DenyReason, "deny envelope must carry a reason payload") + assert.Equal(t, tc.code, out.DenyReason.Code, "canonical deny code surfaces to the caller") + assert.Equal(t, tc.message, out.DenyReason.Message, + "model denials must use a model-specific message, matching the local guardrail") + }) + } +} + // TestInvoke_NoMgmtClientPassesThrough proves the partial-wiring // safety: a middleware constructed without a management client // allows every request without attribution. This makes a half-set-up diff --git a/proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go b/proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go index 0074411cc..1eae26d73 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go @@ -25,10 +25,17 @@ func runParserGuardrail(t *testing.T, url string, body []byte, allowlist []strin }) require.NoError(t, err, "parser must not error") - guard := llm_guardrail.New(llm_guardrail.Config{ModelAllowlist: allowlist}) + const providerID = "prov-under-test" + guard := llm_guardrail.New(llm_guardrail.Config{ + ProviderAllowlists: map[string][]string{providerID: allowlist}, + }) + // The real chain has llm_router stamp the resolved provider id before the + // guardrail runs; the parser doesn't, so add it here so the guardrail can + // scope the allowlist to this provider. + meta := append([]middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: providerID}}, parsed.Metadata...) out, err := guard.Invoke(context.Background(), &middleware.Input{ Slot: middleware.SlotOnRequest, - Metadata: parsed.Metadata, + Metadata: meta, }) require.NoError(t, err, "guardrail must not error") require.NotNil(t, out, "guardrail must return an output")