From afb34f3139f692db848daa69a929939cdc3ee843 Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Sun, 26 Jul 2026 17:49:56 +0000 Subject: [PATCH] e2e: add group-switch scenario using a short tunnel-cache TTL StartProxy now accepts env overrides so a test can set a short NB_PROXY_TUNNEL_CACHE_TTL and observe authorization changes propagate in seconds instead of the 5-minute default. Add TestGuardrailGroupSwitchTakesEffectAfterTTL: a client starts in group A (allowed modelA) with a second policy granting group B modelB. After moving the peer A -> B via the API, once the short tunnel-cache TTL expires the proxy re-validates, the guardrail decision flips (modelB served, modelA denied), and the cross-group deny is enforced by management's per-policy/group check. --- .../guardrail_groupswitch_test.go | 205 ++++++++++++++++++ e2e/harness/proxy.go | 12 +- 2 files changed, 216 insertions(+), 1 deletion(-) create mode 100644 e2e/agentnetwork/guardrail_groupswitch_test.go 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/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,