mirror of
https://github.com/netbirdio/netbird.git
synced 2026-10-25 13:59:08 +01:00
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.
This commit is contained in:
@@ -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 ""
|
||||||
|
}
|
||||||
+11
-1
@@ -38,7 +38,11 @@ type Proxy struct {
|
|||||||
// network, registered via the given account proxy token and serving the
|
// network, registered via the given account proxy token and serving the
|
||||||
// AgentNetworkCluster over a self-signed wildcard cert. It does not wait for
|
// AgentNetworkCluster over a self-signed wildcard cert. It does not wait for
|
||||||
// peer connectivity — callers poll management for the proxy peer.
|
// 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()
|
root, err := repoRoot()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
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),
|
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{
|
ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
|
||||||
ContainerRequest: req,
|
ContainerRequest: req,
|
||||||
Started: true,
|
Started: true,
|
||||||
|
|||||||
Reference in New Issue
Block a user