diff --git a/client/cmd/login.go b/client/cmd/login.go index ee32a3727..a53cb6d5f 100644 --- a/client/cmd/login.go +++ b/client/cmd/login.go @@ -17,7 +17,9 @@ import ( "github.com/netbirdio/netbird/client/internal" "github.com/netbirdio/netbird/client/internal/auth" "github.com/netbirdio/netbird/client/internal/profilemanager" + nbnet "github.com/netbirdio/netbird/client/net" "github.com/netbirdio/netbird/client/proto" + "github.com/netbirdio/netbird/client/server" "github.com/netbirdio/netbird/client/system" "github.com/netbirdio/netbird/util" ) @@ -331,6 +333,14 @@ func doForegroundLogin(ctx context.Context, cmd *cobra.Command, setupKey string, return fmt.Errorf("read config file %s: %v", configFilePath, err) } + // Mirror runInForegroundMode: recover residual state (DNS, firewall, + // ssh config, legacy routing) from a previous unclean shutdown and + // enable advanced routing before dialing management. + if err := server.RestoreResidualState(ctx, profilemanager.NewServiceManager(configFilePath).GetStatePath()); err != nil { + log.Warnf("failed to restore residual state: %v", err) + } + nbnet.Init() + err = foregroundLogin(ctx, cmd, config, setupKey, activeProf.ID) if err != nil { return fmt.Errorf("foreground login failed: %v", err) diff --git a/client/cmd/up.go b/client/cmd/up.go index 8b3de3c66..2d9731f26 100644 --- a/client/cmd/up.go +++ b/client/cmd/up.go @@ -22,6 +22,8 @@ import ( "github.com/netbirdio/netbird/client/internal/peer" "github.com/netbirdio/netbird/client/internal/profilemanager" "github.com/netbirdio/netbird/client/proto" + nbnet "github.com/netbirdio/netbird/client/net" + "github.com/netbirdio/netbird/client/server" "github.com/netbirdio/netbird/client/system" "github.com/netbirdio/netbird/shared/management/domain" "github.com/netbirdio/netbird/util" @@ -229,6 +231,24 @@ func runInForegroundMode(ctx context.Context, cmd *cobra.Command, activeProf *pr _, _ = profilemanager.UpdateOldManagementURL(ctx, config, configFilePath) + // Restore residual state left by a previous run that did not shut down + // cleanly, mirroring what the daemon does before connecting: it recovers + // DNS config (a stale resolv.conf takeover can make the management + // hostname unresolvable), firewall rules, ssh config and legacy routing. + // Route cleanup itself happens at engine start; nbnet.Init() below lets + // the management dial bypass a leftover fwmark rule until then. + // Foreground mode is particularly exposed in containers: a crashed + // container restarts inside the same (pod) network namespace, so stale + // state survives while the process does not. + if err := server.RestoreResidualState(ctx, profilemanager.NewServiceManager(configPath).GetStatePath()); err != nil { + log.Warnf("failed to restore residual state: %v", err) + } + + // Enable advanced routing (as the daemon does on startup) so the + // management dial bypasses a leftover fwmark rule instead of being + // shunted into a stale routing table. + nbnet.Init() + err = foregroundLogin(ctx, cmd, config, providedSetupKey, activeProf.ID) if err != nil { return fmt.Errorf("foreground login failed: %v", err) diff --git a/client/server/server.go b/client/server/server.go index 46f9a6055..2b919d58d 100644 --- a/client/server/server.go +++ b/client/server/server.go @@ -181,7 +181,7 @@ func (s *Server) Start() error { log.Warnf("failed to redirect stderr: %v", err) } - if err := restoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil { + if err := RestoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil { log.Warnf(errRestoreResidualState, err) } @@ -551,7 +551,7 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro s.actCancel = cancel s.mutex.Unlock() - if err := restoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil { + if err := RestoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil { log.Warnf(errRestoreResidualState, err) } @@ -858,7 +858,7 @@ func (s *Server) Up(callerCtx context.Context, msg *proto.UpRequest) (*proto.UpR return s.waitForUp(callerCtx) } - if err := restoreResidualState(callerCtx, s.profileManager.GetStatePath()); err != nil { + if err := RestoreResidualState(callerCtx, s.profileManager.GetStatePath()); err != nil { log.Warnf(errRestoreResidualState, err) } diff --git a/client/server/state.go b/client/server/state.go index f2d823465..a4e91468e 100644 --- a/client/server/state.go +++ b/client/server/state.go @@ -46,7 +46,7 @@ func (s *Server) CleanState(ctx context.Context, req *proto.CleanStateRequest) ( if req.All { // Reuse existing cleanup logic for all states - if err := restoreResidualState(ctx, statePath); err != nil { + if err := RestoreResidualState(ctx, statePath); err != nil { return nil, status.Errorf(codes.Internal, "failed to clean all states: %v", err) } @@ -113,9 +113,9 @@ func (s *Server) DeleteState(ctx context.Context, req *proto.DeleteStateRequest) }, nil } -// restoreResidualState checks if the client was not shut down in a clean way and restores residual if required. +// RestoreResidualState checks if the client was not shut down in a clean way and restores residual if required. // Otherwise, we might not be able to connect to the management server to retrieve new config. -func restoreResidualState(ctx context.Context, statePath string) error { +func RestoreResidualState(ctx context.Context, statePath string) error { if statePath == "" { return nil } 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/main.go b/management/main.go index ff8482f97..a19b741a9 100644 --- a/management/main.go +++ b/management/main.go @@ -1,19 +1,24 @@ package main import ( - "log" "net/http" // nolint:gosec _ "net/http/pprof" "os" + log "github.com/sirupsen/logrus" + "github.com/netbirdio/netbird/management/cmd" ) func main() { - go func() { - log.Println(http.ListenAndServe("localhost:6060", nil)) - }() + if pprofAddr := os.Getenv("NB_PPROF_ADDR"); pprofAddr != "" { + log.Infof("pprof enabled, listening on: %s", pprofAddr) + go func() { + log.Println(http.ListenAndServe(pprofAddr, nil)) + }() + } + if err := cmd.Execute(); err != nil { os.Exit(1) } 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) + } + }) + } +}