mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-07 15:31:30 +02:00
agentnetwork: scope model allowlist per policy/group and provider
The model-allowlist guardrail was merged into one account-wide union
and enforced flat on every request, independent of which policy,
group, or provider actually authorised it. With multiple policies —
especially a mix of guardrailed and un-guardrailed ones — this
produced two defects:
- false-allow: a model allowlisted for one group/provider leaked to
any caller, because the flat union ignored which policy authorised
the request; and
- false-deny: a policy with no guardrail (intended unrestricted) had
its traffic blocked by an unrelated policy's allowlist.
Make enforcement policy/group-aware, mirroring how llm_limit_check
already resolves caps:
- Management (SelectPolicyForRequest) is authoritative. It now uses
the request model (already carried on the wire in
CheckLLMPolicyLimitsRequest.model, previously ignored) to keep only
the applicable policies whose guardrails permit the model; a policy
with no allowlist-enabled guardrail is unrestricted. When policies
govern the (provider, groups) but none permits the model it denies
with llm_policy.model_blocked. No proto change — the field already
existed and the proxy already sends it.
- The proxy llm_guardrail becomes a per-provider fail-closed backstop.
The synthesiser emits a per-provider allowlist only when every
policy authorising that provider restricts models; a provider
reachable by any un-guardrailed policy is omitted (unrestricted).
The middleware keys the allowlist by the resolved provider id and
keeps the unknown/undetermined-model fail-closed behaviour so
path-routed providers stay enforced when management is unreachable.
Unit tests cover the selector gate (block, allow, cross-group no-leak,
un-guardrailed policy unrestricted, undetermined-model fail-closed) and
the per-provider synthesis; an e2e proves all three properties end to
end over the tunnel. Docs updated for the new config shape and the
two-layer enforcement.
This commit is contained in:
@@ -86,12 +86,20 @@ type Manager interface {
|
||||
|
||||
// PolicySelectionInput is the per-request selection envelope. The
|
||||
// proxy populates it from CapturedData (account, user, groups) plus
|
||||
// the provider llm_router resolved.
|
||||
// the provider llm_router resolved and the model it extracted.
|
||||
type PolicySelectionInput struct {
|
||||
AccountID string
|
||||
UserID string
|
||||
GroupIDs []string
|
||||
ProviderID string
|
||||
// Model is the (already-normalised) upstream model identifier the proxy
|
||||
// extracted from the request. The proxy's request parser strips
|
||||
// path-routed provider decoration (Bedrock region/version, Vertex
|
||||
// @version) before emitting it, so a plain case-insensitive compare
|
||||
// against a guardrail allowlist is sufficient here. Empty when the model
|
||||
// could not be determined — treated as "not permitted" by any policy that
|
||||
// restricts models (fail closed), mirroring the proxy guardrail.
|
||||
Model string
|
||||
}
|
||||
|
||||
// PolicySelectionResult names the policy that "pays" for this request
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
@@ -35,6 +36,11 @@ const (
|
||||
denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded"
|
||||
//nolint:gosec // account deny code label, not a credential
|
||||
denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded"
|
||||
// denyCodeModelBlocked is returned when policies govern the request's
|
||||
// (provider, caller-groups) but none of them permits the requested model
|
||||
// under its guardrail allowlist. Matches the proxy guardrail's code so the
|
||||
// two enforcement layers surface the same label.
|
||||
denyCodeModelBlocked = "llm_policy.model_blocked"
|
||||
)
|
||||
|
||||
// consumptionCache holds the consumption counters prefetched for one
|
||||
@@ -159,6 +165,34 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec
|
||||
}
|
||||
candidates := filterApplicablePolicies(policies, in)
|
||||
|
||||
// Model-allowlist gate (per-policy, per-group). Among the policies that
|
||||
// authorise this (provider, caller-groups), keep only those whose guardrails
|
||||
// permit the requested model; a policy with no allowlist-enabled guardrail is
|
||||
// unrestricted and always permits. When policies govern this request but none
|
||||
// permits the model, deny. This is the authoritative allowlist decision:
|
||||
// because it is scoped to the matched policies, a model allowlisted for one
|
||||
// group/provider never leaks to another (no account-wide union), and a policy
|
||||
// that carries no guardrail is genuinely unrestricted rather than being caught
|
||||
// by some other policy's allowlist.
|
||||
// Only consult guardrails when at least one candidate policy references one;
|
||||
// a policy with no guardrail is unrestricted, so when none of them carries a
|
||||
// guardrail the gate is a no-op and we skip the store read entirely.
|
||||
if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) {
|
||||
guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID)
|
||||
if gErr != nil {
|
||||
return nil, gErr
|
||||
}
|
||||
permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model)
|
||||
if len(permitted) == 0 {
|
||||
return &PolicySelectionResult{
|
||||
Allow: false,
|
||||
DenyCode: denyCodeModelBlocked,
|
||||
DenyReason: modelBlockedReason(in.Model),
|
||||
}, nil
|
||||
}
|
||||
candidates = permitted
|
||||
}
|
||||
|
||||
// Prefetch every consumption counter the ceiling + candidate policies will
|
||||
// read, in a single store round-trip, then score against the cache.
|
||||
cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now)
|
||||
@@ -250,6 +284,93 @@ func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput)
|
||||
return out
|
||||
}
|
||||
|
||||
// anyPolicyHasGuardrails reports whether any policy references at least one
|
||||
// guardrail, so the selector can skip loading guardrails when none do.
|
||||
func anyPolicyHasGuardrails(policies []*types.Policy) bool {
|
||||
for _, p := range policies {
|
||||
if p != nil && len(p.GuardrailIDs) > 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
|
||||
// model-allowlist gate to resolve each candidate policy's attached guardrails.
|
||||
func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
|
||||
guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list account guardrails: %w", err)
|
||||
}
|
||||
byID := make(map[string]*types.Guardrail, len(guardrails))
|
||||
for _, g := range guardrails {
|
||||
if g != nil {
|
||||
byID[g.ID] = g
|
||||
}
|
||||
}
|
||||
return byID, nil
|
||||
}
|
||||
|
||||
// filterModelPermittedPolicies returns the subset of policies whose guardrails
|
||||
// permit the model. Order is preserved so downstream scoring is unaffected.
|
||||
func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy {
|
||||
out := make([]*types.Policy, 0, len(policies))
|
||||
for _, p := range policies {
|
||||
if policyPermitsModel(p, byID, model) {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// policyPermitsModel reports whether a policy permits the requested model under
|
||||
// its attached guardrails. A policy with no allowlist-enabled guardrail is
|
||||
// unrestricted (permits any model, including an empty/undetermined one). A
|
||||
// policy with one or more allowlist-enabled guardrails permits the model only
|
||||
// when it appears in the union of those allowlists; an empty/undetermined model
|
||||
// never matches a non-empty allowlist, so such a policy fails closed — the same
|
||||
// contract the proxy guardrail enforces for path-routed providers.
|
||||
func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool {
|
||||
if p == nil {
|
||||
return false
|
||||
}
|
||||
wanted := normaliseModelID(model)
|
||||
restricted := false
|
||||
for _, gID := range p.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
|
||||
continue
|
||||
}
|
||||
restricted = true
|
||||
if wanted == "" {
|
||||
continue
|
||||
}
|
||||
for _, allowed := range g.Checks.ModelAllowlist.Models {
|
||||
if normaliseModelID(allowed) == wanted {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return !restricted
|
||||
}
|
||||
|
||||
// normaliseModelID lowercases and trims a model identifier so the allowlist
|
||||
// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's
|
||||
// normaliseModel so both layers agree on what "same model" means.
|
||||
func normaliseModelID(model string) string {
|
||||
return strings.ToLower(strings.TrimSpace(model))
|
||||
}
|
||||
|
||||
// modelBlockedReason builds the human-readable deny reason for a model-allowlist
|
||||
// rejection. The model is quoted when known; an undetermined model is reported
|
||||
// as such so the access log distinguishes "wrong model" from "no model".
|
||||
func modelBlockedReason(model string) string {
|
||||
if normaliseModelID(model) == "" {
|
||||
return "request model could not be determined for the policy allowlist"
|
||||
}
|
||||
return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model)
|
||||
}
|
||||
|
||||
// candidate is the per-policy intermediate the selector ranks. A
|
||||
// policy that's been exhausted on any enabled cap never makes it
|
||||
// into this slice; the selector's deny envelope carries the latest
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/golang/mock/gomock"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
)
|
||||
|
||||
// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups
|
||||
// to reach providerID under the given guardrails. Uncapped keeps the selector's
|
||||
// headroom scoring trivial so these tests isolate the model-allowlist gate.
|
||||
func guardedPolicy(id, account string, sourceGroups []string, providerID string, guardrailIDs ...string) *types.Policy {
|
||||
return &types.Policy{
|
||||
ID: id,
|
||||
AccountID: account,
|
||||
Enabled: true,
|
||||
SourceGroups: sourceGroups,
|
||||
DestinationProviderIDs: []string{providerID},
|
||||
GuardrailIDs: guardrailIDs,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
}
|
||||
|
||||
// allowlistGuardrail builds a guardrail whose model allowlist is enabled and
|
||||
// carries the given models.
|
||||
func allowlistGuardrail(id, account string, models ...string) *types.Guardrail {
|
||||
return &types.Guardrail{
|
||||
ID: id,
|
||||
AccountID: account,
|
||||
Checks: types.GuardrailChecks{
|
||||
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: models},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func expectPolicies(mockStore *store.MockStore, account string, policies ...*types.Policy) {
|
||||
mockStore.EXPECT().
|
||||
GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), account).
|
||||
Return(policies, nil)
|
||||
}
|
||||
|
||||
func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...*types.Guardrail) {
|
||||
mockStore.EXPECT().
|
||||
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), account).
|
||||
Return(guardrails, nil)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist
|
||||
// decision: a policy authorises the (provider, group) but restricts the model,
|
||||
// and the requested model isn't on the list, so the request is denied.
|
||||
func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "a model outside the only applicable policy's allowlist must be denied")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode, "deny code must be model_blocked")
|
||||
assert.NotEmpty(t, res.DenyReason, "deny reason must be populated")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_ModelAllowedByAllowlist is the allow counterpart: the model
|
||||
// is on the applicable policy's allowlist, so selection proceeds normally.
|
||||
func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4"))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "a model on the applicable policy's allowlist must be allowed")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_CaseInsensitiveModelMatch proves the compare tolerates case
|
||||
// and surrounding whitespace, matching the proxy guardrail's normalisation.
|
||||
func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", " GPT-4o "))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "gpt-4o",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "case/whitespace variants must match the allowlist entry")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_UnguardedPolicyIsUnrestricted is the false-deny fix: when
|
||||
// two policies authorise the same (provider, group) and one carries no
|
||||
// guardrail, that un-guardrailed policy makes the request unrestricted — it is
|
||||
// NOT caught by the other policy's allowlist.
|
||||
func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
restricted := guardedPolicy("pol-restricted", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail
|
||||
expectPolicies(mockStore, "acc-1", restricted, open)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "an un-guardrailed policy for the same (provider, group) must leave the request unrestricted")
|
||||
assert.Equal(t, "pol-open", res.SelectedPolicyID, "the unrestricted policy must be the one that pays")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups is the false-allow fix: a
|
||||
// model allowlisted only for grp-b must not be usable by a caller in grp-a,
|
||||
// even though both groups' policies target the same provider. The selector only
|
||||
// considers policies applicable to the caller's groups, so grp-b's allowlist
|
||||
// never enters grp-a's decision.
|
||||
func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
polA := guardedPolicy("pol-a", "acc-1", []string{"grp-a"}, "prov-1", "g-a")
|
||||
polB := guardedPolicy("pol-b", "acc-1", []string{"grp-b"}, "prov-1", "g-b")
|
||||
expectPolicies(mockStore, "acc-1", polA, polB)
|
||||
expectGuardrails(mockStore, "acc-1",
|
||||
allowlistGuardrail("g-a", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-b", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-a"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4", // only allowed for grp-b
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "grp-b's allowlisted model must not leak to a grp-a caller")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_UndeterminedModelFailsClosed proves the fail-closed contract
|
||||
// mirrors the proxy: with a restricted applicable policy and an empty model
|
||||
// (e.g. a path-routed shape the parser couldn't map), the request is denied.
|
||||
func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "", // undetermined
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "an undetermined model must fail closed against a restricted policy")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_DisabledAllowlistDoesNotRestrict proves a guardrail whose
|
||||
// model allowlist is disabled imposes no model restriction, even though the
|
||||
// policy references it.
|
||||
func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
disabled := &types.Guardrail{
|
||||
ID: "g-1",
|
||||
AccountID: "acc-1",
|
||||
Checks: types.GuardrailChecks{
|
||||
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}},
|
||||
},
|
||||
}
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", disabled)
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "anything-goes",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "a disabled allowlist must not restrict the model")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
@@ -235,7 +235,14 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
|
||||
mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
|
||||
applyAccountCollectionControls(&mergedGuardrails, settings)
|
||||
guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails)
|
||||
// The proxy guardrail is the fail-closed backstop, scoped per provider so it
|
||||
// never applies one policy's allowlist to another provider's traffic. The
|
||||
// authoritative per-policy/group model decision lives in management's
|
||||
// SelectPolicyForRequest; a provider only lands in this map when EVERY policy
|
||||
// authorising it restricts models, so an un-guardrailed policy leaves its
|
||||
// providers unrestricted here.
|
||||
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
|
||||
guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -780,10 +787,12 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt
|
||||
|
||||
// guardrailConfig is the JSON shape the proxy-side llm_guardrail
|
||||
// middleware expects. Mirrors the proxy registration documented in
|
||||
// the management→proxy contract.
|
||||
// the management→proxy contract. provider_allowlists is keyed by the
|
||||
// resolved provider id llm_router stamps; a provider absent from the map is
|
||||
// unrestricted at the proxy layer.
|
||||
type guardrailConfig struct {
|
||||
ModelAllowlist []string `json:"model_allowlist,omitempty"`
|
||||
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
|
||||
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
|
||||
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
|
||||
}
|
||||
|
||||
type guardrailPromptCapture struct {
|
||||
@@ -828,12 +837,12 @@ func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Se
|
||||
merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii
|
||||
}
|
||||
|
||||
func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
|
||||
func marshalGuardrailConfig(providerAllowlists map[string][]string, capture MergedPromptCapture) ([]byte, error) {
|
||||
cfg := guardrailConfig{
|
||||
ModelAllowlist: merged.ModelAllowlist,
|
||||
ProviderAllowlists: providerAllowlists,
|
||||
PromptCapture: guardrailPromptCapture{
|
||||
Enabled: merged.PromptCapture.Enabled,
|
||||
RedactPii: merged.PromptCapture.RedactPii,
|
||||
Enabled: capture.Enabled,
|
||||
RedactPii: capture.RedactPii,
|
||||
},
|
||||
}
|
||||
out, err := json.Marshal(cfg)
|
||||
@@ -843,6 +852,81 @@ func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// buildProviderAllowlists computes the per-provider model allowlist the proxy
|
||||
// guardrail enforces as a fail-closed backstop. A provider is included ONLY when
|
||||
// every enabled policy that authorises it restricts models (has at least one
|
||||
// allowlist-enabled guardrail); its list is then the union of those policies'
|
||||
// allowed models. If any authorising policy leaves models unrestricted, the
|
||||
// provider is omitted entirely — the proxy treats an absent provider as
|
||||
// unrestricted and defers to management's per-policy/group decision, so a
|
||||
// mixed set of policies can neither leak a model across providers nor wrongly
|
||||
// block an un-guardrailed policy's traffic.
|
||||
func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string {
|
||||
type providerAcc struct {
|
||||
models map[string]struct{}
|
||||
anyUnrestricted bool
|
||||
}
|
||||
accs := make(map[string]*providerAcc)
|
||||
for _, p := range policies {
|
||||
if p == nil {
|
||||
continue
|
||||
}
|
||||
restricted, models := policyModelAllowlist(p, byID)
|
||||
for _, providerID := range p.DestinationProviderIDs {
|
||||
if providerID == "" {
|
||||
continue
|
||||
}
|
||||
acc, ok := accs[providerID]
|
||||
if !ok {
|
||||
acc = &providerAcc{models: make(map[string]struct{})}
|
||||
accs[providerID] = acc
|
||||
}
|
||||
if !restricted {
|
||||
acc.anyUnrestricted = true
|
||||
continue
|
||||
}
|
||||
for _, m := range models {
|
||||
acc.models[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
out := make(map[string][]string, len(accs))
|
||||
for providerID, acc := range accs {
|
||||
if acc.anyUnrestricted {
|
||||
continue
|
||||
}
|
||||
models := make([]string, 0, len(acc.models))
|
||||
for m := range acc.models {
|
||||
models = append(models, m)
|
||||
}
|
||||
sort.Strings(models)
|
||||
out[providerID] = models
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// policyModelAllowlist reports whether a policy restricts models (has at least
|
||||
// one allowlist-enabled guardrail) and the union of allowed models across those
|
||||
// guardrails. Models are returned verbatim; the proxy factory lowercases and
|
||||
// trims them at decode time, matching the runtime compare.
|
||||
func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) {
|
||||
restricted := false
|
||||
var models []string
|
||||
for _, gID := range p.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
|
||||
continue
|
||||
}
|
||||
restricted = true
|
||||
for _, m := range g.Checks.ModelAllowlist.Models {
|
||||
if m != "" {
|
||||
models = append(models, m)
|
||||
}
|
||||
}
|
||||
}
|
||||
return restricted, models
|
||||
}
|
||||
|
||||
// buildAccountService composes the per-account gateway Service. The
|
||||
// target carries the noop placeholder URL — the router middleware
|
||||
// rewrites every request to the matched provider's upstream before the
|
||||
@@ -991,11 +1075,10 @@ func unionSourceGroups(policies []*types.Policy) []string {
|
||||
// expectations and is intentionally distinct from
|
||||
// types.GuardrailChecks so we can evolve either side independently.
|
||||
type MergedGuardrails struct {
|
||||
ModelAllowlist []string `json:"model_allowlist,omitempty"`
|
||||
TokenLimits MergedTokenLimits `json:"token_limits"`
|
||||
Budget MergedBudget `json:"budget"`
|
||||
PromptCapture MergedPromptCapture `json:"prompt_capture"`
|
||||
Retention MergedRetention `json:"retention"`
|
||||
TokenLimits MergedTokenLimits `json:"token_limits"`
|
||||
Budget MergedBudget `json:"budget"`
|
||||
PromptCapture MergedPromptCapture `json:"prompt_capture"`
|
||||
Retention MergedRetention `json:"retention"`
|
||||
}
|
||||
|
||||
type MergedTokenLimits struct {
|
||||
@@ -1030,59 +1113,38 @@ type MergedRetention struct {
|
||||
Days int `json:"days"`
|
||||
}
|
||||
|
||||
// mergeGuardrails computes the effective guardrail spec applied at the
|
||||
// proxy, given the referencing policies and the account's guardrail
|
||||
// catalogue. Policy enabled-ness is the caller's responsibility — only
|
||||
// enabled policies should be passed in.
|
||||
// mergeGuardrails computes the prompt-capture portion of the effective
|
||||
// guardrail spec, given the referencing policies and the account's guardrail
|
||||
// catalogue. Policy enabled-ness is the caller's responsibility — only enabled
|
||||
// policies should be passed in.
|
||||
//
|
||||
// Merge rules:
|
||||
// - Model allowlist: union of allowlists across policies that enable it.
|
||||
// - Token / Budget: most-restrictive (min of non-zero caps) per window.
|
||||
// - Prompt capture: enabled if any policy enables it; redact_pii sticks
|
||||
// if any enabling policy turns it on.
|
||||
// - Retention: enabled if any enables it; smallest non-zero days wins.
|
||||
// The model allowlist is NOT merged here: it is enforced per-policy/group in
|
||||
// management (SelectPolicyForRequest) and shipped to the proxy per-provider via
|
||||
// buildProviderAllowlists, so a single account-wide union would be both
|
||||
// incorrect (leaks a model across groups/providers) and redundant. Token,
|
||||
// budget, and retention likewise live off guardrails now (Policy.Limits and
|
||||
// account Settings), so prompt capture is all that remains to merge.
|
||||
//
|
||||
// Merge rule — prompt capture: enabled if any policy enables it; redact_pii
|
||||
// sticks if any enabling policy turns it on.
|
||||
func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails {
|
||||
merged := MergedGuardrails{}
|
||||
allowlist := make(map[string]struct{})
|
||||
allowlistEnabled := false
|
||||
|
||||
for _, policy := range policies {
|
||||
for _, gID := range policy.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil {
|
||||
continue
|
||||
}
|
||||
mergeGuardrail(g, &merged, allowlist, &allowlistEnabled)
|
||||
mergeGuardrail(g, &merged)
|
||||
}
|
||||
}
|
||||
|
||||
if allowlistEnabled {
|
||||
merged.ModelAllowlist = make([]string, 0, len(allowlist))
|
||||
for m := range allowlist {
|
||||
merged.ModelAllowlist = append(merged.ModelAllowlist, m)
|
||||
}
|
||||
sort.Strings(merged.ModelAllowlist)
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
// mergeGuardrail folds a single guardrail's enabled checks into the
|
||||
// running merge: model-allowlist models join the shared set (and flip
|
||||
// allowlistEnabled), and prompt-capture / redact-pii stick once any
|
||||
// enabling guardrail turns them on.
|
||||
//
|
||||
// TokenLimits, Budget, and Retention have moved off guardrails — token
|
||||
// and budget caps now live on the Policy itself (Policy.Limits) and
|
||||
// retention moves to account-level Settings — so they are not merged here.
|
||||
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) {
|
||||
if g.Checks.ModelAllowlist.Enabled {
|
||||
*allowlistEnabled = true
|
||||
for _, m := range g.Checks.ModelAllowlist.Models {
|
||||
if m != "" {
|
||||
allowlist[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
// mergeGuardrail folds a single guardrail's prompt-capture settings into the
|
||||
// running merge: enabled / redact-pii stick once any enabling guardrail turns
|
||||
// them on.
|
||||
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
|
||||
if g.Checks.PromptCapture.Enabled {
|
||||
merged.PromptCapture.Enabled = true
|
||||
if g.Checks.PromptCapture.RedactPii {
|
||||
|
||||
@@ -81,8 +81,8 @@ func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testi
|
||||
require.Len(t, services, 1)
|
||||
|
||||
cfg := decodeServiceGuardrailConfig(t, services[0])
|
||||
assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist,
|
||||
"model allowlist is a pure policy guardrail and must always reach the config")
|
||||
assert.Equal(t, map[string][]string{"prov-1": {"gpt-5.4"}}, cfg.ProviderAllowlists,
|
||||
"model allowlist is a pure policy guardrail and must reach the per-provider config")
|
||||
assert.False(t, cfg.PromptCapture.Enabled,
|
||||
"prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail")
|
||||
}
|
||||
@@ -172,7 +172,7 @@ func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) {
|
||||
require.Len(t, services, 1, "exactly one synth service expected")
|
||||
|
||||
cfg := decodeServiceGuardrailConfig(t, services[0])
|
||||
assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist")
|
||||
assert.Empty(t, cfg.ProviderAllowlists, "no guardrail → provider unrestricted (absent from map)")
|
||||
assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default")
|
||||
assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
)
|
||||
|
||||
// policyForProviders builds an enabled policy authorising the given providers
|
||||
// under the given guardrails (both optional). Groups are irrelevant to
|
||||
// buildProviderAllowlists, which keys purely on destination provider.
|
||||
func policyForProviders(id string, guardrailIDs []string, providerIDs ...string) *types.Policy {
|
||||
return &types.Policy{
|
||||
ID: id,
|
||||
Enabled: true,
|
||||
DestinationProviderIDs: providerIDs,
|
||||
GuardrailIDs: guardrailIDs,
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildProviderAllowlists(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("all authorising policies restrict yields per-provider union", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", []string{"g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got,
|
||||
"a provider every policy restricts carries the sorted union of their models")
|
||||
})
|
||||
|
||||
t.Run("any un-guardrailed policy leaves the provider unrestricted (omitted)", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", nil, "prov-x"), // no guardrail
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.NotContains(t, got, "prov-x",
|
||||
"a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted")
|
||||
})
|
||||
|
||||
t.Run("a disabled allowlist counts as unrestricted", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-disabled"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.NotContains(t, got, "prov-x",
|
||||
"a policy whose only guardrail has a disabled allowlist is unrestricted")
|
||||
})
|
||||
|
||||
t.Run("providers are isolated from one another", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", []string{"g-opus"}, "prov-y"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model")
|
||||
assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model")
|
||||
})
|
||||
|
||||
t.Run("one policy authorising two providers restricts both", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"])
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-y"])
|
||||
})
|
||||
}
|
||||
@@ -1031,8 +1031,14 @@ func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *t
|
||||
|
||||
var cfg guardrailConfig
|
||||
require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly")
|
||||
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist,
|
||||
"model allowlist union must keep both models")
|
||||
// Both policies restrict the same provider, so the proxy-side per-provider
|
||||
// backstop carries the union of their models. This coarse union is
|
||||
// deliberate: it can still let grp-a reach grp-b's model at the proxy layer,
|
||||
// which management's per-policy/group check (SelectPolicyForRequest) is the
|
||||
// one that rejects. The backstop only guarantees nothing outside the
|
||||
// provider's union slips through when management is unreachable.
|
||||
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ProviderAllowlists["prov-1"],
|
||||
"per-provider allowlist union must keep both models")
|
||||
}
|
||||
|
||||
func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) {
|
||||
|
||||
@@ -285,6 +285,7 @@ func (s *ProxyServiceServer) CheckLLMPolicyLimits(ctx context.Context, req *prot
|
||||
UserID: req.GetUserId(),
|
||||
GroupIDs: req.GetGroupIds(),
|
||||
ProviderID: req.GetProviderId(),
|
||||
Model: req.GetModel(),
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("select policy for request: %v", err)
|
||||
|
||||
Reference in New Issue
Block a user