mirror of
https://github.com/netbirdio/netbird.git
synced 2026-08-24 16:41:30 +02:00
[proxy,management] Bound model discovery to the caller's own policies
The listing was narrowed by the provider record's enumerated models, which is the right bound only while one policy reaches a provider. Where two teams share a provider under different allowlists, every caller was offered the union: each model outside their own policy is a request the guardrail refuses a moment later, which is the empty-or-wrong picker this endpoint exists to avoid, moved one level up. A gateway record enumerating nothing was worse still — it offered the upstream's entire catalogue however narrow the policy. The synthesiser already knows which policies authorise a provider and which groups each binds, so the router can answer this at request time where it knows the caller's groups. Each route now carries one rule per authorising policy — its source groups and the models it permits — and the listing is bounded to the union across the rules matching the caller, intersected with what the provider serves. This is deliberately finer than the guardrail's own per-provider allowlist, which stays as it is: that list is a fail-closed backstop and cannot tell who is asking, so discovery is now narrower than the backstop rather than wider. A policy setting no allowlist lifts the restriction for the groups it binds, so nil and empty model lists stay distinct end to end — collapsing them would let a listing that should offer nothing fall open to everything.
This commit is contained in:
122
e2e/agentnetwork/discovery_multipolicy_test.go
Normal file
122
e2e/agentnetwork/discovery_multipolicy_test.go
Normal file
@@ -0,0 +1,122 @@
|
||||
//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"
|
||||
)
|
||||
|
||||
// TestDiscoveryBoundToCallersPolicies covers a model listing on a provider two
|
||||
// teams reach under different allowlists.
|
||||
//
|
||||
// Bounding the listing by the provider's enumerated models alone is not enough
|
||||
// once more than one policy is in play: the caller would be offered every model
|
||||
// any team may use, and each one outside their own policy is a request the
|
||||
// guardrail refuses a moment later — the empty-or-wrong picker this endpoint
|
||||
// exists to avoid, just moved one level up.
|
||||
//
|
||||
// The client joins the main group only. Both models are enumerated by the same
|
||||
// provider and both are advertised by the upstream, so a listing that leaked
|
||||
// the other team's model would visibly contain it.
|
||||
func TestDiscoveryBoundToCallersPolicies(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-disc-mp-main"})
|
||||
require.NoError(t, err, "create main group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
|
||||
|
||||
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-disc-mp-other"})
|
||||
require.NoError(t, err, "create other group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-disc-mp-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpMain.Id}, // the client joins grpMain only
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
// One provider enumerating both models the upstream advertises, so the
|
||||
// listing is narrowed by policy rather than by what the provider serves.
|
||||
staticKey := "static-e2e-token"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-disc-mp",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: &[]api.AgentNetworkProviderModel{
|
||||
{Id: harness.VLLMModel, InputPer1k: 0.001, OutputPer1k: 0.001},
|
||||
{Id: harness.VLLMUnlistedModel, InputPer1k: 0.001, OutputPer1k: 0.001},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "create provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
mkGuardrail := func(name, model string) api.AgentNetworkGuardrail {
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, gerr := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, gerr, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gMain := mkGuardrail("e2e-disc-mp-main", harness.VLLMModel)
|
||||
gOther := mkGuardrail("e2e-disc-mp-other", harness.VLLMUnlistedModel)
|
||||
|
||||
enabled := true
|
||||
polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-disc-mp-main",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gMain.Id},
|
||||
})
|
||||
require.NoError(t, err, "create main policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
|
||||
|
||||
// The other team's policy, on the same provider, permitting the model the
|
||||
// client must never be offered.
|
||||
polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-disc-mp-other",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpOther.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gOther.Id},
|
||||
})
|
||||
require.NoError(t, err, "create other policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
|
||||
|
||||
endpoint, proxyIP, cl, px := connectClient(t, ctx, "disc-mp", sk.Key)
|
||||
_ = px
|
||||
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return cl.Get(ctx, endpoint, proxyIP, "/v1/models?limit=1000", nil)
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "discovery must be served; body: %s", body)
|
||||
|
||||
assert.Contains(t, body, harness.VLLMModel,
|
||||
"the model the caller's own policy permits must reach the picker")
|
||||
assert.NotContains(t, body, harness.VLLMUnlistedModel,
|
||||
"a model only another group's policy permits must not be offered to this caller")
|
||||
}
|
||||
@@ -211,7 +211,19 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
|
||||
groupIndex := indexProviderGroups(enabledPolicies)
|
||||
|
||||
routerCfgJSON, err := buildRouterConfigJSON(enabledProviders, groupIndex)
|
||||
// The proxy guardrail is a per-provider fail-closed backstop; the
|
||||
// authoritative per-policy/group decision is management's
|
||||
// SelectPolicyForRequest. A provider lands in that map only when every
|
||||
// authorising policy restricts models.
|
||||
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
|
||||
|
||||
// Discovery gets the finer view: per policy rather than flattened per
|
||||
// provider, so a listing can be bounded to what the calling groups may
|
||||
// actually use instead of the union across everyone who reaches the
|
||||
// provider.
|
||||
modelPolicies := buildModelPolicies(enabledPolicies, guardrailsByID)
|
||||
|
||||
routerCfgJSON, err := buildRouterConfigJSON(enabledProviders, groupIndex, modelPolicies)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -228,11 +240,6 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
|
||||
mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
|
||||
applyAccountCollectionControls(&mergedGuardrails, settings)
|
||||
// The proxy guardrail is a per-provider fail-closed backstop; the
|
||||
// authoritative per-policy/group decision is management's
|
||||
// SelectPolicyForRequest. A provider lands in this map only when every
|
||||
// authorising policy restricts models.
|
||||
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
|
||||
guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -351,6 +358,11 @@ type routerProviderRoute struct {
|
||||
AuthHeaderName string `json:"auth_header_name"`
|
||||
AuthHeaderValue string `json:"auth_header_value"`
|
||||
AllowedGroupIDs []string `json:"allowed_group_ids,omitempty"`
|
||||
// ModelPolicies is one entry per enabled policy authorising this provider,
|
||||
// carrying that policy's source groups and the models it permits. The
|
||||
// router bounds a model listing with it, so a provider two groups reach
|
||||
// under different allowlists offers each only its own.
|
||||
ModelPolicies []routerModelPolicy `json:"model_policies,omitempty"`
|
||||
// Vertex marks a Google Vertex AI provider, whose requests carry the
|
||||
// model in the URL path. The router selects it by path, bypassing the
|
||||
// model/vendor table.
|
||||
@@ -422,7 +434,7 @@ func indexProviderGroups(policies []*types.Policy) map[string][]string {
|
||||
// path-prefix tiebreak. Providers no enabled policy authorises
|
||||
// (orphans) are intentionally OMITTED so the router never observes a
|
||||
// route with an empty ACL.
|
||||
func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]string) ([]byte, error) {
|
||||
func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]string, modelPolicies map[string][]routerModelPolicy) ([]byte, error) {
|
||||
cfg := routerConfig{Providers: make([]routerProviderRoute, 0, len(providers))}
|
||||
for _, p := range providers {
|
||||
groups, hasPolicy := groupIndex[p.ID]
|
||||
@@ -449,6 +461,7 @@ func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]
|
||||
AuthHeaderName: headerName,
|
||||
AuthHeaderValue: headerValue,
|
||||
AllowedGroupIDs: groups,
|
||||
ModelPolicies: modelPolicies[p.ID],
|
||||
Vertex: catalog.IsVertexPathStyle(p.ProviderID),
|
||||
Bedrock: catalog.IsBedrockPathStyle(p.ProviderID),
|
||||
GCPServiceAccountKeyB64: gcpSAKeyB64,
|
||||
@@ -1098,3 +1111,46 @@ func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// routerModelPolicy mirrors the router's ModelPolicyRule: one authorising
|
||||
// policy's source groups plus the models it permits. Models is nil for a
|
||||
// policy that sets no model allowlist, which lifts the restriction for the
|
||||
// groups it binds — so nil and empty must survive the round trip distinctly.
|
||||
type routerModelPolicy struct {
|
||||
GroupIDs []string `json:"group_ids"`
|
||||
Models []string `json:"models"`
|
||||
}
|
||||
|
||||
// buildModelPolicies indexes, per provider, one rule for each enabled policy
|
||||
// authorising it: the policy's source groups and the models its guardrail
|
||||
// permits.
|
||||
//
|
||||
// This is deliberately finer than buildProviderAllowlists, which flattens the
|
||||
// same inputs into one list per provider for the proxy's fail-closed guardrail.
|
||||
// A flattened list cannot answer "what may THIS caller see", so a provider two
|
||||
// teams reach under different allowlists would offer each team the other's
|
||||
// models — a picker full of entries the next request refuses. Keeping the
|
||||
// source groups alongside the models lets the router answer it at request time,
|
||||
// where it knows the caller's groups.
|
||||
func buildModelPolicies(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]routerModelPolicy {
|
||||
out := make(map[string][]routerModelPolicy)
|
||||
for _, p := range policies {
|
||||
if p == nil || len(p.SourceGroups) == 0 {
|
||||
continue
|
||||
}
|
||||
restricted, models := policyModelAllowlist(p, byID)
|
||||
rule := routerModelPolicy{GroupIDs: append([]string(nil), p.SourceGroups...)}
|
||||
if restricted {
|
||||
// Never nil when restricted: an allowlist permitting nothing must
|
||||
// stay distinguishable from no allowlist at all.
|
||||
rule.Models = append([]string{}, models...)
|
||||
}
|
||||
for _, providerID := range p.DestinationProviderIDs {
|
||||
if providerID == "" {
|
||||
continue
|
||||
}
|
||||
out[providerID] = append(out[providerID], rule)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
)
|
||||
@@ -93,3 +94,75 @@ func TestBuildProviderAllowlists(t *testing.T) {
|
||||
"an enabled-but-empty allowlist is restricted with an empty set, not unrestricted")
|
||||
})
|
||||
}
|
||||
|
||||
// policyForGroups builds an enabled policy binding the given source groups to
|
||||
// the given providers under an optional guardrail.
|
||||
func policyForGroups(id string, groups []string, guardrailIDs []string, providerIDs ...string) *types.Policy {
|
||||
return &types.Policy{
|
||||
ID: id,
|
||||
Enabled: true,
|
||||
SourceGroups: groups,
|
||||
DestinationProviderIDs: providerIDs,
|
||||
GuardrailIDs: guardrailIDs,
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildModelPolicies covers the finer index discovery needs. Where
|
||||
// buildProviderAllowlists flattens every authorising policy into one list per
|
||||
// provider — enough for a fail-closed backstop, but blind to who is asking —
|
||||
// this keeps each policy's source groups beside its models so the router can
|
||||
// bound a listing to the calling groups.
|
||||
func TestBuildModelPolicies(t *testing.T) {
|
||||
byID := map[string]*types.Guardrail{
|
||||
"g-4o": allowlistGuardrail("g-4o", "acc-1", "gpt-4o"),
|
||||
"g-opus": allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"),
|
||||
"g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}},
|
||||
}
|
||||
|
||||
t.Run("each policy keeps its own groups and models", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForGroups("p1", []string{"grp-eng"}, []string{"g-4o"}, "prov-x"),
|
||||
policyForGroups("p2", []string{"grp-sales"}, []string{"g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildModelPolicies(policies, byID)
|
||||
assert.Equal(t, []routerModelPolicy{
|
||||
{GroupIDs: []string{"grp-eng"}, Models: []string{"gpt-4o"}},
|
||||
{GroupIDs: []string{"grp-sales"}, Models: []string{"claude-opus-4"}},
|
||||
}, got["prov-x"],
|
||||
"the two policies must stay separable so neither group is offered the other's models")
|
||||
})
|
||||
|
||||
t.Run("an unrestricted policy carries nil models", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForGroups("p1", []string{"grp-eng"}, []string{"g-4o"}, "prov-x"),
|
||||
policyForGroups("p2", []string{"grp-admin"}, nil, "prov-x"),
|
||||
}
|
||||
got := buildModelPolicies(policies, byID)
|
||||
assert.Nil(t, got["prov-x"][1].Models,
|
||||
"no allowlist must reach the router as nil, which lifts the restriction for its groups")
|
||||
})
|
||||
|
||||
t.Run("a disabled allowlist is not a restriction", func(t *testing.T) {
|
||||
policies := []*types.Policy{policyForGroups("p1", []string{"grp-eng"}, []string{"g-disabled"}, "prov-x")}
|
||||
got := buildModelPolicies(policies, byID)
|
||||
assert.Nil(t, got["prov-x"][0].Models,
|
||||
"a guardrail with the allowlist check off restricts nothing")
|
||||
})
|
||||
|
||||
t.Run("an enabled allowlist with no models permits nothing", func(t *testing.T) {
|
||||
byIDEmpty := map[string]*types.Guardrail{
|
||||
"g-empty": {ID: "g-empty", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true}}},
|
||||
}
|
||||
policies := []*types.Policy{policyForGroups("p1", []string{"grp-eng"}, []string{"g-empty"}, "prov-x")}
|
||||
got := buildModelPolicies(policies, byIDEmpty)
|
||||
require.NotNil(t, got["prov-x"][0].Models,
|
||||
"an empty allowlist must not arrive as nil — that would read as unrestricted")
|
||||
assert.Empty(t, got["prov-x"][0].Models)
|
||||
})
|
||||
|
||||
t.Run("a policy binding no groups is skipped", func(t *testing.T) {
|
||||
policies := []*types.Policy{policyForGroups("p1", nil, []string{"g-4o"}, "prov-x")}
|
||||
assert.Empty(t, buildModelPolicies(policies, byID),
|
||||
"a policy with no source groups authorises nobody, so it bounds nobody's listing")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -44,6 +44,12 @@ type ProviderRoute struct {
|
||||
AuthHeaderName string `json:"auth_header_name"`
|
||||
AuthHeaderValue string `json:"auth_header_value"`
|
||||
AllowedGroupIDs []string `json:"allowed_group_ids"`
|
||||
// ModelPolicies carries, per authorising policy, the source groups it
|
||||
// binds and the models it permits. The router uses it to bound a model
|
||||
// listing to what THIS caller may use: a provider reachable by two groups
|
||||
// under different allowlists must not offer either group the other's
|
||||
// models. Empty means no policy restricts models on this route.
|
||||
ModelPolicies []ModelPolicyRule `json:"model_policies,omitempty"`
|
||||
// Vertex marks a Google Vertex AI provider. Vertex requests carry the
|
||||
// model in the URL path, so the router selects this route by path
|
||||
// (isVertexPath) and bypasses the model/vendor table entirely.
|
||||
@@ -65,6 +71,18 @@ type ProviderRoute struct {
|
||||
SkipTLSVerify bool `json:"skip_tls_verify,omitempty"`
|
||||
}
|
||||
|
||||
// ModelPolicyRule is one authorising policy's contribution to what a caller
|
||||
// may use on a route: the source groups it binds, and the models it permits.
|
||||
//
|
||||
// Models is nil when the policy sets no model allowlist — an unrestricted
|
||||
// policy, which lifts the restriction for the groups it binds. That is why
|
||||
// nil and empty must stay distinct: an empty list is a guardrail that permits
|
||||
// nothing, and collapsing the two would let a listing fail open.
|
||||
type ModelPolicyRule struct {
|
||||
GroupIDs []string `json:"group_ids"`
|
||||
Models []string `json:"models"`
|
||||
}
|
||||
|
||||
// Config is the on-wire configuration accepted by the factory. An
|
||||
// empty Providers slice yields a router that denies every request as
|
||||
// not-routable; the synthesiser is responsible for stamping the
|
||||
|
||||
@@ -242,12 +242,13 @@ func (m *Middleware) routeModelless(reqPath, surface, method string, userGroups
|
||||
if _, hadPrefix := splitBedrockNamespace(reqPath); hadPrefix {
|
||||
stripBedrockNamespace(out)
|
||||
}
|
||||
// A route that enumerates its models bounds what the caller may use,
|
||||
// so the picker must not offer the rest: every entry outside the list
|
||||
// is a request the chain will deny.
|
||||
if reqPath == modelListingPath && len(route.Models) > 0 &&
|
||||
out.Mutations != nil && out.Mutations.RewriteUpstream != nil {
|
||||
out.Mutations.RewriteUpstream.DiscoveryModels = append([]string(nil), route.Models...)
|
||||
// What the caller may actually use bounds what the picker may offer:
|
||||
// every entry outside it is a request the chain will deny a moment
|
||||
// later.
|
||||
if reqPath == modelListingPath && out.Mutations != nil && out.Mutations.RewriteUpstream != nil {
|
||||
if models, bounded := discoverableModels(route, userGroups); bounded {
|
||||
out.Mutations.RewriteUpstream.DiscoveryModels = models
|
||||
}
|
||||
}
|
||||
return out
|
||||
case matchOutcomeUnauthorised:
|
||||
@@ -271,6 +272,96 @@ func isNonInferenceMethod(method string) bool {
|
||||
return method == http.MethodGet || method == http.MethodHead
|
||||
}
|
||||
|
||||
// discoverableModels returns the model ids a caller in userGroups may actually
|
||||
// use on this route, and whether the listing should be bounded to them at all.
|
||||
//
|
||||
// Two things narrow a listing, and both must apply or the picker offers models
|
||||
// the very next request refuses:
|
||||
//
|
||||
// - the provider's own enumerated models, when it lists any (a gateway record
|
||||
// enumerates nothing and claims everything);
|
||||
// - the model allowlists of the policies that authorise THIS caller. A
|
||||
// provider reachable by two groups under different allowlists must not
|
||||
// offer either group the other's models, which is why the rules carry their
|
||||
// source groups rather than arriving pre-flattened.
|
||||
//
|
||||
// A policy that sets no allowlist lifts the restriction for the groups it
|
||||
// binds, so a caller holding one unrestricted policy sees the provider's full
|
||||
// list. bounded is false when nothing narrows the listing — an unrestricted
|
||||
// caller on a route that enumerates nothing — in which case the upstream's own
|
||||
// answer passes through untouched.
|
||||
func discoverableModels(route ProviderRoute, userGroups []string) ([]string, bool) {
|
||||
permitted, restricted := policyPermittedModels(route, userGroups)
|
||||
|
||||
switch {
|
||||
case !restricted && len(route.Models) == 0:
|
||||
return nil, false
|
||||
case !restricted:
|
||||
return append([]string(nil), route.Models...), true
|
||||
case len(route.Models) == 0:
|
||||
// A gateway record enumerates nothing, so the allowlist is the whole
|
||||
// bound — previously such a record offered the upstream's entire
|
||||
// catalogue however narrow the policy was.
|
||||
return sortedModels(permitted), true
|
||||
}
|
||||
|
||||
// Both bound: only what the provider serves and the policy permits.
|
||||
intersection := make(map[string]struct{}, len(route.Models))
|
||||
for _, m := range route.Models {
|
||||
if _, ok := permitted[m]; ok {
|
||||
intersection[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
return sortedModels(intersection), true
|
||||
}
|
||||
|
||||
// policyPermittedModels folds the rules whose groups intersect the caller's
|
||||
// into the set of models they permit. restricted is false when the caller
|
||||
// holds at least one authorising policy that sets no allowlist, or when no
|
||||
// rule binds them at all.
|
||||
func policyPermittedModels(route ProviderRoute, userGroups []string) (map[string]struct{}, bool) {
|
||||
permitted := make(map[string]struct{})
|
||||
restricted := false
|
||||
for _, rule := range route.ModelPolicies {
|
||||
if !groupsIntersect(rule.GroupIDs, userGroups) {
|
||||
continue
|
||||
}
|
||||
if rule.Models == nil {
|
||||
// An unrestricted policy the caller holds lifts the restriction
|
||||
// entirely, whatever the others say.
|
||||
return nil, false
|
||||
}
|
||||
restricted = true
|
||||
for _, m := range rule.Models {
|
||||
permitted[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
return permitted, restricted
|
||||
}
|
||||
|
||||
// groupsIntersect reports whether the two group-id sets share a member.
|
||||
func groupsIntersect(a, b []string) bool {
|
||||
for _, x := range a {
|
||||
for _, y := range b {
|
||||
if x == y {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// sortedModels flattens a model set into a stable slice so the bound the proxy
|
||||
// applies — and any test asserting on it — does not depend on map order.
|
||||
func sortedModels(set map[string]struct{}) []string {
|
||||
out := make([]string, 0, len(set))
|
||||
for m := range set {
|
||||
out = append(out, m)
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// markNonInference tags an allow as a request that spends no tokens, so the
|
||||
// limit check skips the management pre-flight it would charge nothing against.
|
||||
func markNonInference(out *middleware.Output) {
|
||||
|
||||
@@ -1145,3 +1145,144 @@ func TestRouter_PinnedDatedModelStaysDistinct(t *testing.T) {
|
||||
"declaration order must not decide between two deliberately pinned builds")
|
||||
})
|
||||
}
|
||||
|
||||
// TestRouter_DiscoveryBoundToCallersPolicies pins that a model listing is
|
||||
// bounded by the policies that authorise the caller, not by the union across
|
||||
// everyone who can reach the provider. Two teams sharing one provider record
|
||||
// under different allowlists is the case that makes the difference visible: a
|
||||
// flattened per-provider list would offer each team the other's models, and
|
||||
// every one of those entries is a request the guardrail then refuses.
|
||||
func TestRouter_DiscoveryBoundToCallersPolicies(t *testing.T) {
|
||||
const (
|
||||
eng = "grp-eng"
|
||||
sales = "grp-sales"
|
||||
)
|
||||
route := ProviderRoute{
|
||||
ID: "shared-gateway",
|
||||
Models: []string{"claude-sonnet-5", "claude-haiku-4-5", "gpt-4o"},
|
||||
AllowedGroupIDs: []string{eng, sales},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "gateway.example.com",
|
||||
ModelPolicies: []ModelPolicyRule{
|
||||
{GroupIDs: []string{eng}, Models: []string{"claude-sonnet-5"}},
|
||||
{GroupIDs: []string{sales}, Models: []string{"gpt-4o"}},
|
||||
},
|
||||
}
|
||||
|
||||
listingFor := func(t *testing.T, group string) []string {
|
||||
t.Helper()
|
||||
mw := New(Config{Providers: []ProviderRoute{route}})
|
||||
in := newModellessInput(modelListingPath)
|
||||
in.UserGroups = []string{group}
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, middleware.DecisionAllow, out.Decision)
|
||||
require.NotNil(t, out.Mutations)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
return out.Mutations.RewriteUpstream.DiscoveryModels
|
||||
}
|
||||
|
||||
t.Run("each group sees only its own policy's models", func(t *testing.T) {
|
||||
assert.Equal(t, []string{"claude-sonnet-5"}, listingFor(t, eng),
|
||||
"engineering must not be offered the model only sales may use")
|
||||
assert.Equal(t, []string{"gpt-4o"}, listingFor(t, sales),
|
||||
"sales must not be offered the model only engineering may use")
|
||||
})
|
||||
|
||||
t.Run("a model no policy allows is offered to nobody", func(t *testing.T) {
|
||||
for _, group := range []string{eng, sales} {
|
||||
assert.NotContains(t, listingFor(t, group), "claude-haiku-4-5",
|
||||
"the provider serves it, but no policy permits it")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestRouter_DiscoveryUnrestrictedPolicy covers the lifting rule: a caller
|
||||
// holding one policy without a model allowlist sees everything the provider
|
||||
// enumerates, whatever the other policies say.
|
||||
func TestRouter_DiscoveryUnrestrictedPolicy(t *testing.T) {
|
||||
const (
|
||||
eng = "grp-eng"
|
||||
admin = "grp-admin"
|
||||
)
|
||||
route := ProviderRoute{
|
||||
ID: "shared-gateway",
|
||||
Models: []string{"claude-sonnet-5", "gpt-4o"},
|
||||
AllowedGroupIDs: []string{eng, admin},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "gateway.example.com",
|
||||
ModelPolicies: []ModelPolicyRule{
|
||||
{GroupIDs: []string{eng}, Models: []string{"claude-sonnet-5"}},
|
||||
// nil Models: a policy that sets no allowlist at all.
|
||||
{GroupIDs: []string{admin}},
|
||||
},
|
||||
}
|
||||
mw := New(Config{Providers: []ProviderRoute{route}})
|
||||
|
||||
in := newModellessInput(modelListingPath)
|
||||
in.UserGroups = []string{eng, admin}
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.ElementsMatch(t, []string{"claude-sonnet-5", "gpt-4o"},
|
||||
out.Mutations.RewriteUpstream.DiscoveryModels,
|
||||
"an unrestricted policy the caller holds lifts the restriction")
|
||||
}
|
||||
|
||||
// TestRouter_DiscoveryOnGatewayRecord covers a record that enumerates no
|
||||
// models. It previously offered the upstream's whole catalogue however narrow
|
||||
// the policy was, because there was nothing to intersect against; the policy
|
||||
// allowlist is now the bound on its own.
|
||||
func TestRouter_DiscoveryOnGatewayRecord(t *testing.T) {
|
||||
const eng = "grp-eng"
|
||||
base := ProviderRoute{
|
||||
ID: "litellm",
|
||||
AllowedGroupIDs: []string{eng},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "litellm.internal",
|
||||
}
|
||||
|
||||
t.Run("a policy allowlist bounds it", func(t *testing.T) {
|
||||
route := base
|
||||
route.ModelPolicies = []ModelPolicyRule{{GroupIDs: []string{eng}, Models: []string{"gpt-4o"}}}
|
||||
mw := New(Config{Providers: []ProviderRoute{route}})
|
||||
|
||||
in := newModellessInput(modelListingPath)
|
||||
in.UserGroups = []string{eng}
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []string{"gpt-4o"}, out.Mutations.RewriteUpstream.DiscoveryModels,
|
||||
"a catch-all record must still be bounded by what policy permits")
|
||||
})
|
||||
|
||||
t.Run("an allowlist permitting nothing offers nothing", func(t *testing.T) {
|
||||
route := base
|
||||
route.ModelPolicies = []ModelPolicyRule{{GroupIDs: []string{eng}, Models: []string{}}}
|
||||
mw := New(Config{Providers: []ProviderRoute{route}})
|
||||
|
||||
in := newModellessInput(modelListingPath)
|
||||
in.UserGroups = []string{eng}
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Empty(t, out.Mutations.RewriteUpstream.DiscoveryModels,
|
||||
"an empty allowlist permits nothing, and must not be read as unrestricted")
|
||||
})
|
||||
|
||||
t.Run("no policy restriction leaves the listing alone", func(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{base}})
|
||||
|
||||
in := newModellessInput(modelListingPath)
|
||||
in.UserGroups = []string{eng}
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Nil(t, out.Mutations.RewriteUpstream.DiscoveryModels,
|
||||
"nothing narrows the listing, so the upstream's own answer passes through")
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user