From c6bf5fbbfb8324bfd66f787585d6670bc1b5a3f3 Mon Sep 17 00:00:00 2001 From: Maycon Santos Date: Tue, 14 Jul 2026 21:22:40 +0200 Subject: [PATCH] [management,client] 0.74.5 branch sync (#6769) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Describe your changes * [proxy] enforce model allowlist for URL-routed providers (Bedrock/Vertex) by @mlsmaycon in https://github.com/netbirdio/netbird/pull/6764 * [management] Remove proxy peer stale deduplication logic by @mlsmaycon in https://github.com/netbirdio/netbird/pull/6768 ## Issue ticket number and link ## Stack ### Checklist - [ ] Is it a bug fix - [ ] Is a typo/documentation fix - [ ] Is a feature enhancement - [ ] It is a refactor - [ ] Created tests that fail without the change (if possible) - [ ] This change does **not** modify the public API, gRPC protocols, functionality behavior, CLI / service flags, or introduce a new feature — **OR** I have discussed it with the NetBird team beforehand (link the issue / Slack thread in the description). See [CONTRIBUTING.md](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTING.md#discuss-changes-with-the-netbird-team-first). > By submitting this pull request, you confirm that you have read and agree to the terms of the [Contributor License Agreement](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTOR_LICENSE_AGREEMENT.md). ## Documentation Select exactly one: - [ ] I added/updated documentation for this change - [x] Documentation is **not needed** for this change (explain why) ### Docs PR URL (required if "docs added" is checked) Paste the PR link from https://github.com/netbirdio/docs here: https://github.com/netbirdio/docs/pull/__ ## Summary by CodeRabbit - **New Features** - Added model-allowlist guardrails for path-routed providers, including Bedrock and Vertex. - Added Bedrock request support for chat interactions. - Added guardrail management capabilities. - **Bug Fixes** - Requests with missing or blank model identifiers are now denied when a model allowlist is configured, improving fail-closed protection. - Corrected provider-specific request handling and session tracking for Bedrock interactions. - **Tests** - Expanded coverage for allowlist enforcement and provider routing scenarios. --------- Co-authored-by: Theodor Midtlien Co-authored-by: blaugrau90 <61945343+blaugrau90@users.noreply.github.com> Co-authored-by: Viktor Liu <17948409+lixmal@users.noreply.github.com> --- e2e/agentnetwork/chat_test.go | 9 +- e2e/agentnetwork/guardrail_test.go | 168 +++++++++++++++ e2e/harness/agentnetwork.go | 11 + e2e/harness/client.go | 16 ++ management/internals/modules/peers/manager.go | 50 ----- .../agentnetwork_proxypeer_restart_test.go | 199 ------------------ .../builtin/llm_guardrail/middleware.go | 34 ++- .../builtin/llm_guardrail/middleware_test.go | 37 +++- .../guardrail_allowlist_test.go | 106 ++++++++++ 9 files changed, 369 insertions(+), 261 deletions(-) create mode 100644 e2e/agentnetwork/guardrail_test.go delete mode 100644 management/server/agentnetwork_proxypeer_restart_test.go create mode 100644 proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go diff --git a/e2e/agentnetwork/chat_test.go b/e2e/agentnetwork/chat_test.go index 5e3a79273..487aa3cea 100644 --- a/e2e/agentnetwork/chat_test.go +++ b/e2e/agentnetwork/chat_test.go @@ -91,7 +91,7 @@ func availableProviders() []providerCase { if region == "" { region = "us-east-1" } - ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: "us.anthropic.claude-haiku-4-5", kind: harness.WireMessages}) + ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: "us.anthropic.claude-haiku-4-5", kind: harness.WireBedrock}) } return ps } @@ -224,9 +224,12 @@ func TestProvidersMatrix(t *testing.T) { var c int var b string var cerr error - if pc.kind == harness.WireVertex { + switch pc.kind { + case harness.WireVertex: c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, "Reply with exactly: pong", sessionID) - } else { + case harness.WireBedrock: + c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, "Reply with exactly: pong", sessionID) + default: c, b, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, pc.kind, pc.model, "Reply with exactly: pong", sessionID) } if cerr == nil { diff --git a/e2e/agentnetwork/guardrail_test.go b/e2e/agentnetwork/guardrail_test.go new file mode 100644 index 000000000..bb952044f --- /dev/null +++ b/e2e/agentnetwork/guardrail_test.go @@ -0,0 +1,168 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/e2e/harness" + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// catalogModel returns the normalized catalog id the proxy stamps for a +// path-routed provider's configured model — the form the guardrail allowlist is +// compared against (region prefix / @version stripped). +func catalogModel(pc providerCase) string { + switch pc.kind { + case harness.WireBedrock: + return strings.TrimPrefix(pc.model, "us.") + case harness.WireVertex: + return strings.SplitN(pc.model, "@", 2)[0] + default: + return pc.model + } +} + +// disallowedModel returns a valid-shaped model id for the provider that is NOT +// the configured/allowed one, so the guardrail must reject it before the +// request ever reaches the upstream. +func disallowedModel(pc providerCase) string { + switch pc.kind { + case harness.WireBedrock: + return "us.anthropic.claude-opus-4-8" + case harness.WireVertex: + return "claude-opus-4-8@20250101" + default: + return "unlisted-model" + } +} + +// sendModel drives one request for the given model through the provider's native +// wire shape and returns the HTTP status. +func sendModel(ctx context.Context, t *testing.T, cl *harness.Client, endpoint, proxyIP string, pc providerCase, model string) int { + t.Helper() + var code int + var err error + switch pc.kind { + case harness.WireBedrock: + code, _, err = cl.Bedrock(ctx, endpoint, proxyIP, model, "Reply with exactly: pong", "") + case harness.WireVertex: + code, _, err = cl.Vertex(ctx, endpoint, proxyIP, pc.project, pc.region, model, "Reply with exactly: pong", "") + default: + code, _, err = cl.Chat(ctx, endpoint, proxyIP, pc.kind, model, "Reply with exactly: pong", "") + } + require.NoError(t, err, "request must reach the proxy for %s", pc.name) + return code +} + +// TestModelAllowlistEnforced provisions a Model Allowlist guardrail limiting each +// path-routed provider (Bedrock, Vertex) to its configured model, then drives +// requests over the tunnel: the allowed model returns 200 while a model outside +// the allowlist is denied 403 by the guardrail before it reaches the upstream. +// This is the coverage missing for #6751 — the model for these providers travels +// in the URL path, and the allowlist must be enforced there. +func TestModelAllowlistEnforced(t *testing.T) { + var providers []providerCase + for _, pc := range availableProviders() { + if pc.kind == harness.WireBedrock || pc.kind == harness.WireVertex { + providers = append(providers, pc) + } + } + if len(providers) == 0 { + t.Skip("no path-routed provider keys set (AWS_BEARER_TOKEN_BEDROCK / GOOGLE_VERTEX_*); source ~/.llm-keys") + } + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) + defer cancel() + + grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-allowlist"}) + require.NoError(t, err, "create group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-allowlist-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grp.Id}, + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + + // Providers with their configured (allowed) models; the first bootstraps the cluster. + ids := make([]string, 0, len(providers)) + allowed := make([]string, 0, len(providers)) + for i, pc := range providers { + req := providerRequest(pc) + if i == 0 { + req.BootstrapCluster = ptr(harness.AgentNetworkCluster) + } + prov, perr := srv.CreateProvider(ctx, req) + require.NoError(t, perr, "create provider %s", pc.name) + id := prov.Id + ids = append(ids, id) + allowed = append(allowed, catalogModel(pc)) + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), id) }) + } + + // Guardrail allowlisting exactly the configured models. + var gr api.AgentNetworkGuardrailRequest + gr.Name = "e2e-allowlist" + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = allowed + guard, err := srv.CreateGuardrail(ctx, gr) + require.NoError(t, err, "create guardrail") + t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) }) + + enabled := true + pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-allowlist", + Enabled: &enabled, + SourceGroups: []string{grp.Id}, + DestinationProviderIds: ids, + GuardrailIds: &[]string{guard.Id}, + }) + require.NoError(t, err, "create policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) }) + + settings, err := srv.GetSettings(ctx) + require.NoError(t, err, "read settings for endpoint") + require.NotEmpty(t, settings.Endpoint, "agent-network endpoint must be assigned") + + proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-proxy-allowlist") + require.NoError(t, err, "mint proxy token via CLI") + px, err := harness.StartProxy(ctx, srv, proxyToken) + require.NoError(t, err, "start proxy") + t.Cleanup(func() { _ = px.Terminate(context.Background()) }) + + cl, err := harness.StartClient(ctx, srv, sk.Key) + require.NoError(t, err, "start client") + t.Cleanup(func() { _ = cl.Terminate(context.Background()) }) + + require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management") + if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil { + t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background())) + } + proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint) + require.NoError(t, err, "resolve agent-network endpoint to proxy IP") + + for _, pc := range providers { + pc := pc + t.Run(pc.name, func(t *testing.T) { + // The admin's allowlisted model is served end to end. + assert.Equal(t, 200, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, pc.model), + "allowlisted model must be permitted for %s", pc.name) + // A model outside the allowlist is rejected by the guardrail (before + // the upstream), regardless of whether it is a real catalog model. + assert.Equal(t, 403, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, disallowedModel(pc)), + "model outside the allowlist must be denied for %s", pc.name) + }) + } +} diff --git a/e2e/harness/agentnetwork.go b/e2e/harness/agentnetwork.go index 192385ab1..53aa8e342 100644 --- a/e2e/harness/agentnetwork.go +++ b/e2e/harness/agentnetwork.go @@ -107,6 +107,17 @@ func (c *Combined) DeletePolicy(ctx context.Context, id string) error { return anDelete(ctx, c, "/api/agent-network/policies/"+id) } +// CreateGuardrail creates an agent-network guardrail (e.g. a model allowlist) +// that can then be attached to a policy via its GuardrailIds. +func (c *Combined) CreateGuardrail(ctx context.Context, req api.AgentNetworkGuardrailRequest) (api.AgentNetworkGuardrail, error) { + return anRequest[api.AgentNetworkGuardrail](ctx, c, http.MethodPost, "/api/agent-network/guardrails", req) +} + +// DeleteGuardrail removes a guardrail by id. +func (c *Combined) DeleteGuardrail(ctx context.Context, id string) error { + return anDelete(ctx, c, "/api/agent-network/guardrails/"+id) +} + // GetSettings returns the account's agent-network settings row. It exists only // after the first provider create bootstraps it. func (c *Combined) GetSettings(ctx context.Context) (api.AgentNetworkSettings, error) { diff --git a/e2e/harness/client.go b/e2e/harness/client.go index 1ce8c0f6e..19210349f 100644 --- a/e2e/harness/client.go +++ b/e2e/harness/client.go @@ -194,6 +194,11 @@ const ( // WireVertex is the Anthropic-on-Vertex rawPredict shape: the client posts // the full Vertex model path and the proxy mints the SA OAuth token. WireVertex = "vertex" + // WireBedrock is the native AWS Bedrock InvokeModel shape: the model id + // travels in the URL path (/model/{id}/invoke), not the body, so the proxy + // routes by path. This is what a Bedrock SDK client sends and the shape the + // model-allowlist guardrail must enforce. + WireBedrock = "bedrock" ) // Chat issues a chat-completion POST to the agent-network endpoint over the @@ -226,6 +231,17 @@ func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID)) } +// Bedrock issues a native AWS Bedrock InvokeModel POST over the tunnel. The +// model id is carried in the request path (/model/{id}/invoke), so the proxy +// routes by path; the body uses the bedrock anthropic_version rather than a +// model field. A non-empty sessionID is sent as the universal x-session-id +// header the proxy records. +func (cl *Client) Bedrock(ctx context.Context, endpoint, proxyIP, model, prompt, sessionID string) (int, string, error) { + path := "/model/" + model + "/invoke" + body := fmt.Sprintf(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":64,"messages":[{"role":"user","content":%q}]}`, prompt) + return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID)) +} + // withSessionID appends the x-session-id header when sessionID is non-empty. func withSessionID(headers []string, sessionID string) []string { if sessionID == "" { diff --git a/management/internals/modules/peers/manager.go b/management/internals/modules/peers/manager.go index 5e4538d08..6f292f6ed 100644 --- a/management/internals/modules/peers/manager.go +++ b/management/internals/modules/peers/manager.go @@ -226,30 +226,6 @@ func (m *managerImpl) CreateProxyPeer(ctx context.Context, accountID string, pee return nil } - // Dedupe stale embedded peer records for the same (account, cluster). - // The proxy generates a fresh WireGuard keypair on every startup - // (proxy/internal/roundtrip/netbird.go), so without this sweep the - // prior embedded peer would linger forever — holding its CGNAT IP - // allocation, polluting other peers' rosters, and (most visibly) - // leaving the synth DNS pointing at the dead address. The - // (account, cluster) tuple identifies "the embedded peer for this - // proxy instance at this cluster"; any record matching that tuple - // with a different pubkey is by definition stale and must go. - staleIDs, err := m.findStaleEmbeddedProxyPeers(ctx, accountID, cluster, peerKey) - if err != nil { - return fmt.Errorf("scan for stale embedded proxy peers: %w", err) - } - if len(staleIDs) > 0 { - // userID="" + checkConnected=false: the deletion is initiated - // by management itself on behalf of the freshly-registering - // proxy, not by an end user; the stale peer may still be - // marked Connected from its prior session, but its session is - // dead by definition (its key no longer exists). - if err := m.DeletePeers(ctx, accountID, staleIDs, "", false); err != nil { - return fmt.Errorf("delete stale embedded proxy peers %v: %w", staleIDs, err) - } - } - name := fmt.Sprintf("proxy-%s", xid.New().String()) newPeer := &peer.Peer{ Ephemeral: true, @@ -275,29 +251,3 @@ func (m *managerImpl) CreateProxyPeer(ctx context.Context, accountID string, pee return nil } - -// findStaleEmbeddedProxyPeers returns the peer IDs of embedded proxy peer -// records in accountID that target the same cluster but carry a different -// WireGuard pubkey than the freshly-registering one. Used by CreateProxyPeer -// to garbage-collect stale records left behind when the proxy restarts with a -// regenerated keypair. -func (m *managerImpl) findStaleEmbeddedProxyPeers(ctx context.Context, accountID, cluster, newKey string) ([]string, error) { - account, err := m.store.GetAccount(ctx, accountID) - if err != nil { - return nil, err - } - var stale []string - for _, p := range account.Peers { - if p == nil || !p.ProxyMeta.Embedded { - continue - } - if p.ProxyMeta.Cluster != cluster { - continue - } - if p.Key == newKey { - continue - } - stale = append(stale, p.ID) - } - return stale, nil -} diff --git a/management/server/agentnetwork_proxypeer_restart_test.go b/management/server/agentnetwork_proxypeer_restart_test.go deleted file mode 100644 index 1e4b8d016..000000000 --- a/management/server/agentnetwork_proxypeer_restart_test.go +++ /dev/null @@ -1,199 +0,0 @@ -package server - -import ( - "context" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/netbirdio/netbird/management/internals/modules/peers" - "github.com/netbirdio/netbird/management/internals/modules/agentnetwork" - agenttypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" - nbpeer "github.com/netbirdio/netbird/management/server/peer" - "github.com/netbirdio/netbird/management/server/permissions" - "github.com/netbirdio/netbird/management/server/store" - "github.com/netbirdio/netbird/management/server/types" -) - -// TestAgentNetwork_ProxyRestart_PropagatesNewPeerAndDropsStale is the no-mock -// regression guard for the bug the user reported: restarting the proxy creates -// a fresh embedded peer with a NEW WireGuard public key (the proxy generates -// the keypair on every startup at proxy/internal/roundtrip/netbird.go:312). -// The PRIOR embedded peer record is never deleted on management, so the -// account accumulates a stale peer holding a stale CGNAT IP. Other peers -// in the account either keep routing to the dead IP, or — if synth DNS -// picks the wrong record — never see the new IP at all. -// -// What this test exercises (no mocks): -// - real SQLite test store -// - real DefaultAccountManager, network-map controller, peer-update channels -// - real peers.Manager.CreateProxyPeer path (the very method the proxy -// invokes over gRPC on every startup) -// - real agentnetwork.Manager + synth chain so the client receives a -// concrete DNS record that must point at the LATEST proxy peer. -// -// Pre-fix expected behavior (red): two embedded peers exist after the -// "restart"; the synth DNS record points at the stale one; the client -// receives an update reflecting the new peer but the old one lingers. -// Post-fix expected behavior (green): exactly one embedded peer exists -// after restart (with the new key) AND the client's network map carries -// the synth DNS pointing at that new peer's CGNAT IP. -func TestAgentNetwork_ProxyRestart_PropagatesNewPeerAndDropsStale(t *testing.T) { - am, updateManager, err := createManager(t) - require.NoError(t, err, "createManager must succeed") - ctx := context.Background() - - const ( - accountID = "an-restart-acct" - adminUserID = "an-restart-admin" - groupAID = "an-restart-grp-A" - clusterAddr = "eu.proxy.netbird.io" - clientKey = "BhRPtynAAYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8=" - // Two different proxy pubkeys — the "before" and "after" of a - // proxy-process restart with fresh-keypair generation. - proxyKey1 = "Aaaaa1aaaaYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8=" - proxyKey2 = "Bbbbb2bbbbYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8=" - ) - - // --- Account scaffold --- - account := newAccountWithId(ctx, accountID, adminUserID, "an-restart.test", "", "", false) - require.NoError(t, am.Store.SaveAccount(ctx, account)) - - clientPeer := &nbpeer.Peer{ - Key: clientKey, - Name: "an-restart-client", - DNSLabel: "an-restart-client", - Meta: nbpeer.PeerSystemMeta{Hostname: "an-restart-client", GoOS: "linux", WtVersion: "development"}, - } - addedClient, _, _, _, err := am.AddPeer(ctx, "", "", adminUserID, clientPeer, false) - require.NoError(t, err, "AddPeer for client must succeed") - require.NoError(t, am.MarkPeerConnected(ctx, clientKey, accountID, time.Now().UnixNano(), &types.NetworkMap{}), - "MarkPeerConnected for the client peer must succeed (affected-peer fan-out skips disconnected peers)") - - // Place the client in group A so the synth policy reaches it. - account, err = am.Store.GetAccount(ctx, accountID) - require.NoError(t, err) - account.Groups[groupAID] = &types.Group{ID: groupAID, Name: "groupA", Peers: []string{addedClient.ID}} - require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must persist group A") - - // --- Real peers + agent-network managers --- - permMgr := permissions.NewManager(am.Store) - peersMgr := peers.NewManager(am.Store, permMgr) - peersMgr.SetAccountManager(am) - peersMgr.SetNetworkMapController(am.networkMapController) - agentMgr := agentnetwork.NewManager(am.Store, permMgr, am, nil) - - // Subscribe BEFORE any state-mutating call so we don't lose the update - // that contains the synth DNS record. - clientCh := updateManager.CreateChannel(ctx, addedClient.ID) - t.Cleanup(func() { updateManager.CloseChannel(ctx, addedClient.ID) }) - drain(clientCh) - - // --- First proxy startup: register peer key K1, then mark it - // connected. In production the proxy follows CreateProxyPeer with the - // regular sync stream which lands on MarkPeerConnected; the synth DNS - // path filters out peers that aren't Connected (types/account.go:323), - // so without this step no DNS record would be emitted. - require.NoError(t, peersMgr.CreateProxyPeer(ctx, accountID, proxyKey1, clusterAddr), - "first CreateProxyPeer (proxy startup) must succeed") - - peer1ID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey1) - require.NoError(t, err, "proxy peer for K1 must be persisted after CreateProxyPeer") - require.NotEmpty(t, peer1ID) - - require.NoError(t, am.MarkPeerConnected(ctx, proxyKey1, accountID, time.Now().UnixNano(), &types.NetworkMap{}), - "MarkPeerConnected for K1 must succeed") - - account, err = am.Store.GetAccount(ctx, accountID) - require.NoError(t, err) - proxyIP1 := account.Peers[peer1ID].IP.String() - require.NotEmpty(t, proxyIP1, "K1 must have an assigned overlay IP") - - // --- Provider + policy. CreateProvider / CreatePolicy trigger the - // agentnetwork reconcile which runs UpdateAccountPeers; the resulting - // NetworkMap delivered to the client carries the synth DNS record - // pointing at K1's IP. --- - provider, err := agentMgr.CreateProvider(ctx, adminUserID, &agenttypes.Provider{ - AccountID: accountID, - ProviderID: "openai_api", - Name: "openai-test", - UpstreamURL: "https://api.openai.com", - APIKey: "sk-test-key", - Enabled: true, - Models: []agenttypes.ProviderModel{{ID: "gpt-5.4"}}, - }, clusterAddr) - require.NoError(t, err, "CreateProvider must succeed") - - _, err = agentMgr.CreatePolicy(ctx, adminUserID, &agenttypes.Policy{ - AccountID: accountID, - Name: "p1", - Enabled: true, - SourceGroups: []string{groupAID}, - DestinationProviderIDs: []string{provider.ID}, - }) - require.NoError(t, err, "CreatePolicy must succeed") - - settings, err := am.Store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) - require.NoError(t, err) - fqdn := settings.Endpoint() - - rdata1 := awaitZoneRData(clientCh, clusterAddr, fqdn, true) - require.Equal(t, proxyIP1, rdata1, - "client must receive a synth DNS record pointing at K1's overlay IP after the synth path runs") - drain(clientCh) - - // --- Proxy restart: NEW keypair K2, same account, same cluster --- - require.NoError(t, peersMgr.CreateProxyPeer(ctx, accountID, proxyKey2, clusterAddr), - "second CreateProxyPeer (proxy restart with fresh keypair) must succeed") - - peer2ID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey2) - require.NoError(t, err, "proxy peer for K2 must be persisted after restart") - require.NotEmpty(t, peer2ID) - - require.NoError(t, am.MarkPeerConnected(ctx, proxyKey2, accountID, time.Now().UnixNano(), &types.NetworkMap{}), - "MarkPeerConnected for K2 must succeed") - - // In production the agent's sync stream pulls a fresh NetworkMap as - // part of its normal reconcile cadence; in this isolated test - // MarkPeerConnected's affected-peer fan-out can race the channel-side - // buffer in a way that swallows the synth-DNS-bearing update before - // our await reads it. Trigger an explicit account-wide fan-out so the - // assertion below tests what production actually delivers, not the - // in-test buffer race. - am.UpdateAccountPeers(ctx, accountID, types.UpdateReason{Resource: types.UpdateResourcePeer, Operation: types.UpdateOperationUpdate}) - - account, err = am.Store.GetAccount(ctx, accountID) - require.NoError(t, err) - proxyIP2 := account.Peers[peer2ID].IP.String() - require.NotEmpty(t, proxyIP2, "K2 must have an assigned overlay IP") - require.NotEqual(t, proxyIP1, proxyIP2, "K2 must get a different overlay IP than K1 (sanity)") - - // CRITICAL ASSERTION 1: K1 must no longer be in the store. The SqlStore - // returns ("", nil) for a missing key rather than NotFound, so assert - // on the returned ID being empty. - staleID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey1) - require.NoError(t, err, "GetPeerIDByKey for a missing peer must not error") - assert.Empty(t, staleID, - "stale embedded proxy peer K1 must be removed when a new embedded peer registers for the same (account, cluster); pre-fix this assertion fails because management never cleans up the prior peer record") - - // CRITICAL ASSERTION 2: exactly one embedded proxy peer remains, and it - // is K2. - account, err = am.Store.GetAccount(ctx, accountID) - require.NoError(t, err) - embeddedKeys := []string{} - for _, p := range account.Peers { - if p.ProxyMeta.Embedded { - embeddedKeys = append(embeddedKeys, p.Key) - } - } - assert.Equal(t, []string{proxyKey2}, embeddedKeys, - "after a proxy restart exactly one embedded proxy peer should remain — the one with the new key K2") - - // CRITICAL ASSERTION 3: the synth DNS record the client receives now - // points at K2's IP, not K1's. - rdata2 := awaitZoneRData(clientCh, clusterAddr, fqdn, true) - assert.Equal(t, proxyIP2, rdata2, - "after proxy restart, the client's synth DNS record must point at the NEW embedded peer's IP, not the stale K1 IP") -} diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go index e6259f06f..eded877ac 100644 --- a/proxy/internal/middleware/builtin/llm_guardrail/middleware.go +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go @@ -25,6 +25,14 @@ const ( denyCodeModel = "llm_policy.model_blocked" denyReasonModel = "model_blocked" denyMessageModel = "model is not in the policy allowlist" + // Deny reason used when an allowlist is configured but the request model + // could not be determined. URL/path-routed providers (AWS Bedrock, Google + // Vertex, ...) carry the model outside the JSON body, so a request shape the + // parser does not recognise reaches the guardrail with no model. Such a + // request must be denied (fail closed), never waved through. + denyCodeModelUnknown = "llm_policy.model_unknown" + denyReasonModelUnknown = "model_unknown" + denyMessageModelUnknown = "request model could not be determined for the policy allowlist" ) // Middleware enforces the model allowlist and optionally captures the @@ -108,23 +116,37 @@ func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middlew if len(m.cfg.ModelAllowlist) == 0 { return nil } - if !modelPresent { - return nil + // Fail closed: with an allowlist configured, a request whose model the + // upstream parser could not extract (absent or empty) must be denied rather + // than allowed. This is what enforces the allowlist for URL/path-routed + // providers (Bedrock, Vertex, ...) whose model lives outside the JSON body. + if !modelPresent || normaliseModel(model) == "" { + return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown) } if m.modelInAllowlist(model) { return nil } + return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel) +} + +// denyModel builds a 403 deny Output for a model-allowlist rejection. model is +// included in the details only when non-empty. +func denyModel(model, code, message, reason string) *middleware.Output { + details := map[string]string{} + if model != "" { + details["model"] = model + } return &middleware.Output{ Decision: middleware.DecisionDeny, DenyStatus: 403, DenyReason: &middleware.DenyReason{ - Code: denyCodeModel, - Message: denyMessageModel, - Details: map[string]string{"model": model}, + Code: code, + Message: message, + Details: details, }, Metadata: []middleware.KV{ {Key: middleware.KeyLLMPolicyDecision, Value: "deny"}, - {Key: middleware.KeyLLMPolicyReason, Value: denyReasonModel}, + {Key: middleware.KeyLLMPolicyReason, Value: reason}, }, } } diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go index 865dc07af..cd7e256dd 100644 --- a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go @@ -102,13 +102,44 @@ func TestAllowlistCaseInsensitive(t *testing.T) { } } -func TestAllowlistMissingModelKeyAllows(t *testing.T) { +func TestAllowlistMissingModelKeyDenies(t *testing.T) { + // Fail closed: with an allowlist configured, a request whose model the + // parser could not extract (URL/path-routed providers such as Bedrock or + // Vertex whose shape wasn't recognised) must be denied, not allowed. mw := New(Config{ModelAllowlist: []string{"gpt-4o"}}) out, err := mw.Invoke(context.Background(), newInput()) require.NoError(t, err) - assert.Equal(t, middleware.DecisionAllow, out.Decision, "missing model key must allow even with non-empty allowlist") + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when an allowlist is set") + assert.Equal(t, 403, out.DenyStatus, "deny status must be 403") + require.NotNil(t, out.DenyReason, "deny reason must be populated") + assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown") dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision) - assert.Equal(t, "allow", dec, "decision must be allow when model key is absent") + assert.Equal(t, "deny", dec, "decision must be deny when model key is absent") + reason, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyReason) + assert.Equal(t, "model_unknown", reason, "reason metadata must be model_unknown") +} + +func TestAllowlistEmptyModelValueDenies(t *testing.T) { + // A present-but-empty model is as undeterminable as an absent one. + mw := New(Config{ModelAllowlist: []string{"gpt-4o"}}) + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMModel, Value: " "}, + )) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when an allowlist is set") + require.NotNil(t, out.DenyReason, "deny reason must be populated") + assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown") +} + +func TestAllowlistEmptyListAllowsMissingModel(t *testing.T) { + // Without an allowlist there is nothing to enforce, so a missing model is + // still allowed — the fail-closed rule only applies when a list is set. + mw := New(Config{}) + out, err := mw.Invoke(context.Background(), newInput()) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "no allowlist must allow even without a model") } func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) { 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 new file mode 100644 index 000000000..0074411cc --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go @@ -0,0 +1,106 @@ +package llm_request_parser + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_guardrail" +) + +// runParserGuardrail runs the request parser then the model-allowlist guardrail +// in SlotOnRequest order, threading the parser's metadata into the guardrail the +// same way the real chain does. It returns the guardrail decision so tests can +// assert allowlist enforcement for URL/path-routed providers end to end. +func runParserGuardrail(t *testing.T, url string, body []byte, allowlist []string) *middleware.Output { + t.Helper() + parser := newMiddleware(t) + parsed, err := parser.Invoke(context.Background(), &middleware.Input{ + Slot: middleware.SlotOnRequest, + URL: url, + Body: body, + }) + require.NoError(t, err, "parser must not error") + + guard := llm_guardrail.New(llm_guardrail.Config{ModelAllowlist: allowlist}) + out, err := guard.Invoke(context.Background(), &middleware.Input{ + Slot: middleware.SlotOnRequest, + Metadata: parsed.Metadata, + }) + require.NoError(t, err, "guardrail must not error") + require.NotNil(t, out, "guardrail must return an output") + return out +} + +// TestModelAllowlist_URLRoutedProviders validates that the model allowlist is +// enforced for providers whose model travels in the URL path (AWS Bedrock, +// Google Vertex) rather than the JSON body. The "unknown action" case is the +// regression guard for #6751: a Bedrock request shape the parser cannot map to a +// model must fail closed under an allowlist instead of bypassing it. +func TestModelAllowlist_URLRoutedProviders(t *testing.T) { + const bedrockBody = `{"anthropic_version":"bedrock-2023-05-31","messages":[{"role":"user","content":"hi"}]}` + const vertexBody = `{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"hi"}]}` + + tests := []struct { + name string + url string + body string + allowlist []string + decision middleware.Decision + denyCode string + }{ + { + name: "bedrock allowed model passes", + url: "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-haiku-4-5-v1:0/invoke", + body: bedrockBody, + allowlist: []string{"anthropic.claude-haiku-4-5"}, + decision: middleware.DecisionAllow, + }, + { + name: "bedrock disallowed model denied", + url: "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-opus-4-8-v1:0/invoke", + body: bedrockBody, + allowlist: []string{"anthropic.claude-haiku-4-5"}, + decision: middleware.DecisionDeny, + denyCode: "llm_policy.model_blocked", + }, + { + name: "bedrock unknown action fails closed", + url: "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-opus-4-8-v1:0/some-future-action", + body: bedrockBody, + allowlist: []string{"anthropic.claude-haiku-4-5"}, + decision: middleware.DecisionDeny, + denyCode: "llm_policy.model_unknown", + }, + { + name: "vertex disallowed model denied", + url: "/v1/projects/p/locations/global/publishers/anthropic/models/claude-opus-4-8@20250101:rawPredict", + body: vertexBody, + allowlist: []string{"claude-haiku-4-5"}, + decision: middleware.DecisionDeny, + denyCode: "llm_policy.model_blocked", + }, + { + name: "vertex allowed model passes", + url: "/v1/projects/p/locations/global/publishers/anthropic/models/claude-haiku-4-5@20250101:rawPredict", + body: vertexBody, + allowlist: []string{"claude-haiku-4-5"}, + decision: middleware.DecisionAllow, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + out := runParserGuardrail(t, tt.url, []byte(tt.body), tt.allowlist) + assert.Equal(t, tt.decision, out.Decision, "unexpected decision for %s", tt.name) + if tt.decision == middleware.DecisionDeny { + require.NotNil(t, out.DenyReason, "deny reason must be set for %s", tt.name) + assert.Equal(t, 403, out.DenyStatus, "deny status must be 403 for %s", tt.name) + assert.Equal(t, tt.denyCode, out.DenyReason.Code, "deny code for %s", tt.name) + } + }) + } +}