mirror of
https://github.com/netbirdio/netbird.git
synced 2026-08-12 18:51:28 +02:00
Compare commits
23 Commits
agent-netw
...
agent-netw
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dc91325ac1 | ||
|
|
1796b2a1d8 | ||
|
|
dd87760f3d | ||
|
|
63d16a7cd4 | ||
|
|
fa691721d0 | ||
|
|
d46e2574a7 | ||
|
|
649a867cd3 | ||
|
|
a39b3c4af4 | ||
|
|
b1337f09d0 | ||
|
|
08e187ccb7 | ||
|
|
93cdf64a19 | ||
|
|
c04fb1c388 | ||
|
|
5353cab54f | ||
|
|
a64a417ea5 | ||
|
|
652c5c8b68 | ||
|
|
4d2b8b407b | ||
|
|
03e02c86ce | ||
|
|
789d416215 | ||
|
|
6415215126 | ||
|
|
1ae352a08d | ||
|
|
d928bcb630 | ||
|
|
47b2667653 | ||
|
|
875dda1708 |
@@ -40,6 +40,35 @@ You can then use this private endpoint to configure your AI agents, whether that
|
||||
Full step-by-step setup:
|
||||
**https://docs.netbird.io/agent-network/quickstart**
|
||||
|
||||
## Client settings that don't follow the endpoint
|
||||
|
||||
Most of an agent's traffic follows the base URL you hand it, but a few
|
||||
client-side checks call their vendor directly and never reach the proxy. On a
|
||||
network that blocks direct egress they fail even though inference works, so
|
||||
they are worth setting once when you roll the endpoint out.
|
||||
|
||||
For Claude Code:
|
||||
|
||||
- **Fast mode** checks availability against `api.anthropic.com` rather than the
|
||||
configured base URL. Set `CLAUDE_CODE_SKIP_FAST_MODE_ORG_CHECK=1` when the
|
||||
agent authenticates with `ANTHROPIC_AUTH_TOKEN` alone (the usual shape when
|
||||
the proxy injects the real provider key) or when a TLS-inspecting proxy
|
||||
answers the check itself. Set
|
||||
`CLAUDE_CODE_SKIP_FAST_MODE_NETWORK_ERRORS=1` when the network refuses the
|
||||
connection outright. Fast mode is an Anthropic-API feature, so it is
|
||||
unavailable on a Bedrock- or Vertex-backed endpoint whatever you set.
|
||||
- **Model discovery** is off by default. Set
|
||||
`CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` for the picker to list the
|
||||
models your policies authorise; the proxy filters the response to that set.
|
||||
The client gives discovery a three-second budget and treats any redirect as
|
||||
a failure, so the endpoint must serve `/v1/models` directly.
|
||||
- **The WebFetch domain safety check** also calls `api.anthropic.com` directly
|
||||
and is unaffected by the variables above.
|
||||
|
||||
Allowing direct egress to `api.anthropic.com` covers the network cases but not
|
||||
the credential one, where the check reaches Anthropic and is rejected because
|
||||
the agent presents a proxy-issued key.
|
||||
|
||||
## Architecture
|
||||
|
||||
Agent Network is built on two existing NetBird capabilities:
|
||||
|
||||
@@ -23,9 +23,10 @@ import (
|
||||
// model the client asks for. The proxy prices off the REQUEST model, not the
|
||||
// upstream response model, so a made-up model id billed at operator rates lets
|
||||
// these tests assert exact costs without a real vendor key.
|
||||
// Sourced from the harness so the counts can't drift from the mock's config.
|
||||
const (
|
||||
vllmPromptTokens = 11
|
||||
vllmCompletionTokens = 2
|
||||
vllmPromptTokens = harness.VLLMChatInputTokens
|
||||
vllmCompletionTokens = harness.VLLMChatOutputTokens
|
||||
)
|
||||
|
||||
// pricedEnv is a connected single-provider agent-network deployment pointed at
|
||||
|
||||
455
e2e/agentnetwork/gateway_protocol_test.go
Normal file
455
e2e/agentnetwork/gateway_protocol_test.go
Normal file
@@ -0,0 +1,455 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"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"
|
||||
)
|
||||
|
||||
// Models each catalog surface is registered with in the matrix below. They
|
||||
// differ per provider so the router's choice is unambiguous: a request that
|
||||
// lands on the wrong provider record fails the surface assertion instead of
|
||||
// passing by coincidence.
|
||||
const (
|
||||
matrixAnthropicModel = "claude-sonnet-5"
|
||||
matrixBedrockModel = "anthropic.claude-sonnet-5"
|
||||
// matrixBedrockPathModel is what a Bedrock SDK client puts in the URL: a
|
||||
// cross-region inference profile with a release date and version suffix.
|
||||
// The proxy must normalise it back to matrixBedrockModel to route and price.
|
||||
matrixBedrockPathModel = "us.anthropic.claude-sonnet-5-20250101-v1:0"
|
||||
// matrixVertexModel differs from the Anthropic record's model on purpose:
|
||||
// a shared id would leave two routes claiming it and make which one serves
|
||||
// /v1/messages depend on declaration order.
|
||||
matrixVertexModel = "claude-haiku-4-5"
|
||||
matrixVertexProject = "e2e-project"
|
||||
matrixVertexRegion = "us-east5"
|
||||
)
|
||||
|
||||
// gatewayEnv is a connected client plus a set of provider records, all pointed
|
||||
// at one mock upstream, so several wire shapes can be driven over a single
|
||||
// tunnel.
|
||||
type gatewayEnv struct {
|
||||
endpoint string
|
||||
proxyIP string
|
||||
client *harness.Client
|
||||
proxy *harness.Proxy
|
||||
vllm *harness.VLLM
|
||||
// providerIDs maps the catalog id to the created provider record id.
|
||||
providerIDs map[string]string
|
||||
}
|
||||
|
||||
// provisionGatewayMatrix brings up one mock upstream and one provider record
|
||||
// per catalog surface, all authorised for the same group by a single policy.
|
||||
// Sharing one proxy and client keeps the wire-shape cases to one tunnel setup;
|
||||
// each case still creates its own session id so its access-log row is findable.
|
||||
func provisionGatewayMatrix(t *testing.T, ctx context.Context) gatewayEnv {
|
||||
t.Helper()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gw-matrix"})
|
||||
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-gw-matrix-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
// The mock ignores auth, so a dummy credential satisfies each catalog
|
||||
// entry's auth template. Vertex is the exception: its api_key is a GCP
|
||||
// service-account keyfile the proxy mints an OAuth token from, and a dummy
|
||||
// one cannot mint. That is deliberate — the Vertex case below asserts on
|
||||
// routing, which happens before the token mint.
|
||||
dummyKey := "sk-gw-e2e"
|
||||
dummyKeyfile := "keyfile::" + "e2e-not-a-real-service-account-key"
|
||||
|
||||
specs := []struct {
|
||||
name string
|
||||
catalogID string
|
||||
apiKey string
|
||||
models []api.AgentNetworkProviderModel
|
||||
}{
|
||||
{
|
||||
name: "openai", catalogID: "openai_api", apiKey: dummyKey,
|
||||
models: []api.AgentNetworkProviderModel{{Id: harness.VLLMModel, InputPer1k: 0.001, OutputPer1k: 0.002}},
|
||||
},
|
||||
{
|
||||
name: "anthropic", catalogID: "anthropic_api", apiKey: dummyKey,
|
||||
models: []api.AgentNetworkProviderModel{{Id: matrixAnthropicModel, InputPer1k: 0.003, OutputPer1k: 0.015}},
|
||||
},
|
||||
{
|
||||
name: "bedrock", catalogID: "bedrock_api", apiKey: dummyKey,
|
||||
models: []api.AgentNetworkProviderModel{{Id: matrixBedrockModel, InputPer1k: 0.003, OutputPer1k: 0.015}},
|
||||
},
|
||||
{
|
||||
name: "vertex", catalogID: "vertex_ai_api", apiKey: dummyKeyfile,
|
||||
models: []api.AgentNetworkProviderModel{{Id: matrixVertexModel, InputPer1k: 0.001, OutputPer1k: 0.005}},
|
||||
},
|
||||
}
|
||||
|
||||
providerIDs := make(map[string]string, len(specs))
|
||||
ids := make([]string, 0, len(specs))
|
||||
for _, spec := range specs {
|
||||
key := spec.apiKey
|
||||
models := spec.models
|
||||
prov, perr := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-gw-" + spec.name,
|
||||
ProviderId: spec.catalogID,
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &key,
|
||||
Enabled: ptr(true),
|
||||
Models: &models,
|
||||
})
|
||||
require.NoError(t, perr, "create %s provider", spec.name)
|
||||
id := prov.Id
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), id) })
|
||||
providerIDs[spec.catalogID] = id
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
// Uncapped token limit: never blocks the handful of tokens driven here, but
|
||||
// switches on usage metering so consumption and cost land in the row.
|
||||
enabled := true
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-gw-matrix",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: ids,
|
||||
Limits: &api.AgentNetworkPolicyLimits{
|
||||
TokenLimit: api.AgentNetworkPolicyTokenLimit{
|
||||
Enabled: true,
|
||||
GroupCap: 10_000_000,
|
||||
UserCap: 10_000_000,
|
||||
WindowSeconds: 60,
|
||||
},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
endpoint, proxyIP, cl, px := connectClient(t, ctx, "gw-matrix", sk.Key)
|
||||
return gatewayEnv{
|
||||
endpoint: endpoint,
|
||||
proxyIP: proxyIP,
|
||||
client: cl,
|
||||
proxy: px,
|
||||
vllm: vllm,
|
||||
providerIDs: providerIDs,
|
||||
}
|
||||
}
|
||||
|
||||
// connectClient starts a proxy and a tunnel client for the shared account and
|
||||
// waits until the client can reach the proxy peer, returning the endpoint and
|
||||
// the proxy's tunnel IP to pin requests to.
|
||||
func connectClient(t *testing.T, ctx context.Context, name, setupKey string) (string, string, *harness.Client, *harness.Proxy) {
|
||||
t.Helper()
|
||||
|
||||
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-"+name+"-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
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, setupKey)
|
||||
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")
|
||||
// The probe resolves the endpoint and its first packet wakes the lazy proxy
|
||||
// peer, so WaitProxyPeer then observes it connected.
|
||||
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()))
|
||||
}
|
||||
return settings.Endpoint, proxyIP, cl, px
|
||||
}
|
||||
|
||||
// callUntil retries an HTTP call through the tunnel until it returns one of the
|
||||
// wanted statuses or the deadline passes, absorbing the DNS and tunnel jitter
|
||||
// the first call through a fresh tunnel can hit. The last status and body are
|
||||
// returned either way so the caller can assert with real detail.
|
||||
func callUntil(t *testing.T, call func() (int, string, error), want ...int) (int, string) {
|
||||
t.Helper()
|
||||
wanted := make(map[int]struct{}, len(want))
|
||||
for _, w := range want {
|
||||
wanted[w] = struct{}{}
|
||||
}
|
||||
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
c, b, err := call()
|
||||
if err == nil {
|
||||
code, body = c, b
|
||||
if _, ok := wanted[code]; ok {
|
||||
return code, body
|
||||
}
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
// TestGatewayProtocolProviderMatrix drives one request per wire shape over a
|
||||
// single tunnel, with a provider record per catalog surface behind it. It is
|
||||
// the regression net for the routing and parser-selection changes: each case
|
||||
// asserts the surface the request was metered under and the token counts that
|
||||
// surface's own usage block carries, so a request parsed by the wrong provider's
|
||||
// parser meters zero and fails rather than passing on a coincidence.
|
||||
func TestGatewayProtocolProviderMatrix(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
env := provisionGatewayMatrix(t, ctx)
|
||||
diag := func() string {
|
||||
return fmt.Sprintf("\n=== upstream logs ===\n%s\n=== proxy logs ===\n%s",
|
||||
env.vllm.Logs(context.Background()), env.proxy.Logs(context.Background()))
|
||||
}
|
||||
|
||||
t.Run("openai chat completions", func(t *testing.T) {
|
||||
session := "e2e-gw-openai"
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireChat, harness.VLLMModel, "ping", session)
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "openai chat must succeed; body: %s%s", body, diag())
|
||||
require.Contains(t, body, "chat.completion", "body must be an OpenAI completion; got: %s", body)
|
||||
|
||||
row := findAccessLogBySession(t, ctx, session)
|
||||
require.NotNil(t, row.Provider)
|
||||
assert.Equal(t, "openai", *row.Provider, "the OpenAI chat path must meter under the openai surface")
|
||||
assert.Equal(t, int64(harness.VLLMChatInputTokens), row.InputTokens, "OpenAI usage block must be read")
|
||||
assert.Equal(t, int64(harness.VLLMChatOutputTokens), row.OutputTokens)
|
||||
})
|
||||
|
||||
t.Run("anthropic messages", func(t *testing.T) {
|
||||
session := "e2e-gw-anthropic"
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, matrixAnthropicModel, "ping", session)
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "anthropic messages must succeed; body: %s%s", body, diag())
|
||||
|
||||
row := findAccessLogBySession(t, ctx, session)
|
||||
require.NotNil(t, row.Provider)
|
||||
assert.Equal(t, "anthropic", *row.Provider, "the /v1/messages path must meter under the anthropic surface")
|
||||
// These counts only appear if the Anthropic parser read the response:
|
||||
// its usage fields are named differently from the OpenAI block.
|
||||
assert.Equal(t, int64(harness.VLLMMessagesInputTokens), row.InputTokens,
|
||||
"Anthropic input_tokens must be read; zero here means the wrong parser ran")
|
||||
assert.Equal(t, int64(harness.VLLMMessagesOutputTokens), row.OutputTokens)
|
||||
assert.Positive(t, row.CachedInputTokens, "the Anthropic cache-read bucket must be recorded")
|
||||
assert.Positive(t, row.CostUsd, "a metered request must carry a cost")
|
||||
require.NotNil(t, row.ResolvedProviderId)
|
||||
assert.Equal(t, env.providerIDs["anthropic_api"], *row.ResolvedProviderId,
|
||||
"a vendor-tagged request must not cross to another provider's record")
|
||||
})
|
||||
|
||||
t.Run("bedrock invoke normalises the path model", func(t *testing.T) {
|
||||
session := "e2e-gw-bedrock"
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Bedrock(ctx, env.endpoint, env.proxyIP, matrixBedrockPathModel, "ping", session)
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "bedrock invoke must succeed; body: %s%s", body, diag())
|
||||
|
||||
row := findAccessLogBySession(t, ctx, session)
|
||||
require.NotNil(t, row.Provider)
|
||||
assert.Equal(t, "bedrock", *row.Provider, "a native Bedrock path must meter under the bedrock surface")
|
||||
require.NotNil(t, row.Model)
|
||||
assert.Equal(t, matrixBedrockModel, *row.Model,
|
||||
"the inference-profile prefix, release date and version suffix must be normalised away")
|
||||
assert.Equal(t, int64(harness.VLLMMessagesInputTokens), row.InputTokens)
|
||||
})
|
||||
|
||||
t.Run("anthropic token counting", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.PostJSON(ctx, env.endpoint, env.proxyIP, "/v1/messages/count_tokens",
|
||||
fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"ping"}]}`, matrixAnthropicModel),
|
||||
[]string{"anthropic-version: 2023-06-01"})
|
||||
}, 200)
|
||||
assert.Equal(t, 200, code, "token counting must route rather than deny; body: %s%s", body, diag())
|
||||
})
|
||||
|
||||
t.Run("bedrock token counting", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.PostJSON(ctx, env.endpoint, env.proxyIP,
|
||||
"/model/"+matrixBedrockPathModel+"/count-tokens",
|
||||
`{"input":{"converse":{"messages":[{"role":"user","content":[{"text":"ping"}]}]}}}`, nil)
|
||||
}, 200)
|
||||
assert.Equal(t, 200, code,
|
||||
"the Bedrock count-tokens action must route; denying it pushes counting onto the billable inference path; body: %s%s",
|
||||
body, diag())
|
||||
})
|
||||
|
||||
t.Run("vertex token counting reaches its provider", func(t *testing.T) {
|
||||
// The dummy service-account key cannot mint an OAuth token, so the
|
||||
// request stops at the upstream credential. Both outcomes render as
|
||||
// 403, so the deny code is what distinguishes them: upstream_auth_failed
|
||||
// means the path resolved to the Vertex route and only the credential
|
||||
// failed, while model_not_routable would mean the method segment was
|
||||
// swallowed into the model id and no route ever claimed it.
|
||||
path := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/anthropic/models/%s/count-tokens:rawPredict",
|
||||
matrixVertexProject, matrixVertexRegion, matrixVertexModel)
|
||||
_, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.PostJSON(ctx, env.endpoint, env.proxyIP, path,
|
||||
`{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"ping"}]}`, nil)
|
||||
}, 403)
|
||||
assert.NotContains(t, body, "model_not_routable",
|
||||
"the count-tokens method segment must not be parsed as part of the model id; body: %s%s", body, diag())
|
||||
assert.Contains(t, body, "llm_policy.upstream_auth_failed",
|
||||
"the request must reach the Vertex route and fail only at the credential; body: %s%s", body, diag())
|
||||
})
|
||||
|
||||
t.Run("connection warming probe", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Get(ctx, env.endpoint, env.proxyIP, "/api/hello", nil)
|
||||
}, 200)
|
||||
assert.NotEqual(t, 403, code,
|
||||
"the warm-up probe carries no model and must not be refused as unroutable; body: %s%s", body, diag())
|
||||
})
|
||||
|
||||
t.Run("unknown model denies in the caller's error shape", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages,
|
||||
"claude-not-a-real-model-9", "ping", "e2e-gw-unknown")
|
||||
}, 403)
|
||||
require.Equal(t, 403, code, "a model no provider claims must still be refused; body: %s%s", body, diag())
|
||||
|
||||
// The NetBird fields stay where they were for existing consumers.
|
||||
assert.Contains(t, body, "llm_policy.model_not_routable", "the deny code must be preserved")
|
||||
// And the vendor's own envelope rides alongside, so the client can show
|
||||
// the reason instead of an unexplained API error.
|
||||
assert.Contains(t, body, `"type":"error"`, "an Anthropic caller must get the Anthropic error envelope")
|
||||
assert.Contains(t, body, "permission_error", "403 must map to the vendor's permission error type")
|
||||
})
|
||||
}
|
||||
|
||||
// TestModelDiscoveryWithModelAllowlist covers gateway model discovery on an
|
||||
// account that restricts models, which is the configuration that broke: the
|
||||
// listing carries no model, and the per-model allowlist fails closed on an
|
||||
// undetermined one, so discovery denied for exactly the accounts using the
|
||||
// feature. It also asserts the allowlist still refuses a model outside it, so
|
||||
// the exemption cannot be read as a way around the gate.
|
||||
func TestModelDiscoveryWithModelAllowlist(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gw-discovery"})
|
||||
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-gw-discovery-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
// One provider enumerating a single model, while the upstream's own listing
|
||||
// advertises two. The proxy must serve the shorter list.
|
||||
dummyKey := "sk-discovery-e2e"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-gw-discovery",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &dummyKey,
|
||||
Enabled: ptr(true),
|
||||
Models: &[]api.AgentNetworkProviderModel{
|
||||
{Id: harness.VLLMModel, InputPer1k: 0.001, OutputPer1k: 0.002},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "create provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
// The model allowlist is what makes this a regression test: without a
|
||||
// guardrail enabled, discovery was never gated in the first place.
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-gw-discovery-allowlist"
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{harness.VLLMModel}
|
||||
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-gw-discovery",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{guard.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
endpoint, proxyIP, cl, px := connectClient(t, ctx, "gw-discovery", sk.Key)
|
||||
diag := func() string {
|
||||
return fmt.Sprintf("\n=== upstream logs ===\n%s\n=== proxy logs ===\n%s",
|
||||
vllm.Logs(context.Background()), px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
t.Run("listing is served and bounded by policy", func(t *testing.T) {
|
||||
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 not be refused because the request carries no model; body: %s%s", body, diag())
|
||||
|
||||
assert.Contains(t, body, harness.VLLMModel, "the authorised model must reach the picker")
|
||||
assert.NotContains(t, body, harness.VLLMUnlistedModel,
|
||||
"a model the policy does not authorise must not be offered; body: %s", body)
|
||||
})
|
||||
|
||||
t.Run("allowlist still refuses a model outside it", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return cl.Chat(ctx, endpoint, proxyIP, harness.WireChat,
|
||||
harness.VLLMUnlistedModel, "ping", "e2e-gw-discovery-blocked")
|
||||
}, 403)
|
||||
require.Equal(t, 403, code,
|
||||
"exempting model-less endpoints must not exempt inference; body: %s%s", body, diag())
|
||||
assert.True(t,
|
||||
strings.Contains(body, "llm_policy.model_blocked") || strings.Contains(body, "llm_policy.model_not_routable"),
|
||||
"the refusal must name a model policy code; body: %s", body)
|
||||
})
|
||||
|
||||
t.Run("allowlisted model still routes", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return cl.Chat(ctx, endpoint, proxyIP, harness.WireChat,
|
||||
harness.VLLMModel, "ping", "e2e-gw-discovery-allowed")
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "the allowlisted model must still be served; body: %s%s", body, diag())
|
||||
})
|
||||
}
|
||||
@@ -11,7 +11,6 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/client/rest"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
@@ -178,89 +177,3 @@ func TestSettingsBootstrapSelfAddressed(t *testing.T) {
|
||||
require.NoError(t, err, "bootstrap after delete must succeed")
|
||||
assert.Equal(t, "gw2.e2e.netbird.selfhosted", recreated.Endpoint, "the fresh bootstrap claims the new hostname")
|
||||
}
|
||||
|
||||
// TestSettingsConditionalWrites covers the lost-update guard end to end, over
|
||||
// the same REST client the Terraform provider uses: read the settings, take
|
||||
// the entity-tag, and have a write refused when the row moved underneath it.
|
||||
//
|
||||
// The scenario is the one that motivates the feature. A client reads the
|
||||
// settings and computes an update. An operator turns PII redaction on in the
|
||||
// dashboard in the meantime. Without a precondition the client's write puts
|
||||
// redaction straight back off — no error, no drift warning, a
|
||||
// compliance-relevant control silently disabled. With one, the write is
|
||||
// refused and the client can read again.
|
||||
func TestSettingsConditionalWrites(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
fresh, err := harnessStartFresh(ctx, t)
|
||||
require.NoError(t, err, "start dedicated combined server")
|
||||
|
||||
const cluster = "eu.e2e.netbird.selfhosted"
|
||||
bootstrapped, err := fresh.CreateSettings(ctx, api.AgentNetworkSettingsCreateRequest{
|
||||
ProxyAddress: ptr(cluster),
|
||||
})
|
||||
require.NoError(t, err, "bootstrap must succeed")
|
||||
|
||||
// What the client plans against.
|
||||
planned, etag, err := fresh.GetSettingsWithETag(ctx)
|
||||
require.NoError(t, err, "read must succeed")
|
||||
require.NotEmpty(t, etag, "the read must carry a validator")
|
||||
assert.Equal(t, bootstrapped.Endpoint, planned.Endpoint)
|
||||
|
||||
_, again, err := fresh.GetSettingsWithETag(ctx)
|
||||
require.NoError(t, err, "second read must succeed")
|
||||
assert.Equal(t, etag, again, "an unchanged row must read as the same validator")
|
||||
|
||||
update := func(redactPii bool, retention int) api.AgentNetworkSettingsRequest {
|
||||
return api.AgentNetworkSettingsRequest{
|
||||
Endpoint: planned.Endpoint,
|
||||
ProxyAddress: planned.ProxyAddress,
|
||||
EnableLogCollection: true,
|
||||
EnablePromptCollection: true,
|
||||
RedactPii: redactPii,
|
||||
AccessLogRetentionDays: retention,
|
||||
}
|
||||
}
|
||||
|
||||
// The operator's change, which the planning client never saw.
|
||||
_, err = fresh.UpdateSettings(ctx, update(true, 21))
|
||||
require.NoError(t, err, "the intervening update must succeed")
|
||||
|
||||
// The client's write, planned against the earlier read, would have turned
|
||||
// redaction back off. It is refused instead.
|
||||
_, _, err = fresh.UpdateSettingsIfMatch(ctx, update(false, 7), etag)
|
||||
require.Error(t, err, "a stale precondition must be refused")
|
||||
require.True(t, rest.IsPreconditionFailed(err),
|
||||
"the refusal must be a precondition failure, got: %v", err)
|
||||
|
||||
intact, current, err := fresh.GetSettingsWithETag(ctx)
|
||||
require.NoError(t, err, "read after the refusal must succeed")
|
||||
assert.True(t, intact.RedactPii, "the refused write must not have turned redaction off")
|
||||
require.NotNil(t, intact.AccessLogRetentionDays)
|
||||
assert.Equal(t, 21, *intact.AccessLogRetentionDays, "the refused write must not have changed retention")
|
||||
assert.NotEqual(t, etag, current, "the validator must have moved with the intervening update")
|
||||
|
||||
// Retrying against the current validator goes through, and hands back the
|
||||
// validator for the write after it.
|
||||
updated, next, err := fresh.UpdateSettingsIfMatch(ctx, update(true, 7), current)
|
||||
require.NoError(t, err, "a matching precondition must be honoured")
|
||||
require.NotNil(t, updated.AccessLogRetentionDays)
|
||||
assert.Equal(t, 7, *updated.AccessLogRetentionDays, "the conditional write must apply")
|
||||
assert.NotEmpty(t, next, "the write must return a validator")
|
||||
assert.NotEqual(t, current, next, "the write must move the validator")
|
||||
|
||||
// The delete is conditional too, and refusing a stale one leaves the
|
||||
// endpoint claimed.
|
||||
err = fresh.DeleteSettingsIfMatch(ctx, etag)
|
||||
require.Error(t, err, "a stale precondition must refuse the delete")
|
||||
require.True(t, rest.IsPreconditionFailed(err),
|
||||
"the delete must be refused for staleness rather than for a state guard or a server error, got: %v", err)
|
||||
stillThere, err := fresh.GetSettings(ctx)
|
||||
require.NoError(t, err, "read after the refused delete must succeed")
|
||||
assert.Equal(t, planned.Endpoint, stillThere.Endpoint, "the refused delete must leave the endpoint claimed")
|
||||
|
||||
require.NoError(t, fresh.DeleteSettingsIfMatch(ctx, next), "a matching precondition must be honoured")
|
||||
gone, err := fresh.GetSettings(ctx)
|
||||
require.NoError(t, err, "read after the delete must succeed")
|
||||
assert.Empty(t, gone.Endpoint, "the row must be gone")
|
||||
}
|
||||
|
||||
@@ -153,37 +153,6 @@ func (c *Combined) DeleteSettings(ctx context.Context) error {
|
||||
return anDelete(ctx, c, "/api/agent-network/settings")
|
||||
}
|
||||
|
||||
// The conditional-request wrappers go through the typed REST client rather
|
||||
// than anRequest, so the e2e run exercises the client's own header handling —
|
||||
// the quoting on the way out and the unquoting on the way back — against a
|
||||
// real server, which is the path the Terraform provider takes.
|
||||
|
||||
// GetSettingsWithETag reads the settings along with the entity-tag that makes
|
||||
// a following write conditional.
|
||||
func (c *Combined) GetSettingsWithETag(ctx context.Context) (api.AgentNetworkSettings, string, error) {
|
||||
settings, etag, err := c.api.AgentNetwork.GetSettingsWithETag(ctx)
|
||||
if err != nil {
|
||||
return api.AgentNetworkSettings{}, "", err
|
||||
}
|
||||
return *settings, etag, nil
|
||||
}
|
||||
|
||||
// UpdateSettingsIfMatch applies the update only if etag is still current,
|
||||
// returning the entity-tag of the row it wrote.
|
||||
func (c *Combined) UpdateSettingsIfMatch(ctx context.Context, req api.AgentNetworkSettingsRequest, etag string) (api.AgentNetworkSettings, string, error) {
|
||||
settings, newETag, err := c.api.AgentNetwork.UpdateSettingsIfMatch(ctx, req, etag)
|
||||
if err != nil {
|
||||
return api.AgentNetworkSettings{}, "", err
|
||||
}
|
||||
return *settings, newETag, nil
|
||||
}
|
||||
|
||||
// DeleteSettingsIfMatch deletes the settings row only if etag is still
|
||||
// current.
|
||||
func (c *Combined) DeleteSettingsIfMatch(ctx context.Context, etag string) error {
|
||||
return c.api.AgentNetwork.DeleteSettingsIfMatch(ctx, etag)
|
||||
}
|
||||
|
||||
// ListConsumption returns the account's consumption rows (possibly empty).
|
||||
func (c *Combined) ListConsumption(ctx context.Context) ([]api.AgentNetworkConsumption, error) {
|
||||
return anRequest[[]api.AgentNetworkConsumption](ctx, c, http.MethodGet, "/api/agent-network/consumption", nil)
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -294,10 +295,29 @@ func withSessionID(headers []string, sessionID string) []string {
|
||||
return append(headers, "x-session-id: "+sessionID)
|
||||
}
|
||||
|
||||
// post runs curl in a throwaway container sharing the client's network
|
||||
// namespace so the request traverses the WireGuard tunnel, pinning the endpoint
|
||||
// to the proxy IP. It returns the HTTP status and response body.
|
||||
// Get issues a GET to the agent-network endpoint over the client's tunnel.
|
||||
// Model discovery and the connection-warming probe are read-only endpoints
|
||||
// that carry no body, so they can't go through the chat helpers.
|
||||
func (cl *Client) Get(ctx context.Context, endpoint, proxyIP, path string, extraHeaders []string) (int, string, error) {
|
||||
return cl.do(ctx, http.MethodGet, endpoint, proxyIP, path, "", extraHeaders)
|
||||
}
|
||||
|
||||
// PostJSON issues an arbitrary JSON POST over the client's tunnel, for wire
|
||||
// shapes the typed helpers don't cover (token counting, say).
|
||||
func (cl *Client) PostJSON(ctx context.Context, endpoint, proxyIP, path, body string, extraHeaders []string) (int, string, error) {
|
||||
return cl.do(ctx, http.MethodPost, endpoint, proxyIP, path, body, extraHeaders)
|
||||
}
|
||||
|
||||
// post issues a JSON POST. Retained as the shorthand the chat helpers use.
|
||||
func (cl *Client) post(ctx context.Context, endpoint, proxyIP, path, body string, extraHeaders []string) (int, string, error) {
|
||||
return cl.do(ctx, http.MethodPost, endpoint, proxyIP, path, body, extraHeaders)
|
||||
}
|
||||
|
||||
// do runs curl in a throwaway container sharing the client's network
|
||||
// namespace so the request traverses the WireGuard tunnel, pinning the endpoint
|
||||
// to the proxy IP. It returns the HTTP status and response body. An empty body
|
||||
// sends no payload, which is what a GET needs.
|
||||
func (cl *Client) do(ctx context.Context, method, endpoint, proxyIP, path, body string, extraHeaders []string) (int, string, error) {
|
||||
url := "https://" + endpoint + path
|
||||
args := []string{
|
||||
"run", "--rm",
|
||||
@@ -306,13 +326,15 @@ func (cl *Client) post(ctx context.Context, endpoint, proxyIP, path, body string
|
||||
"-sk", "--connect-timeout", "5", "--max-time", "90",
|
||||
"--resolve", endpoint + ":443:" + proxyIP,
|
||||
"-o", "/dev/stderr", "-w", "%{http_code}",
|
||||
"-X", "POST", url,
|
||||
"-X", method, url,
|
||||
"-H", "Content-Type: application/json",
|
||||
}
|
||||
for _, h := range extraHeaders {
|
||||
args = append(args, "-H", h)
|
||||
}
|
||||
args = append(args, "--data", body)
|
||||
if body != "" {
|
||||
args = append(args, "--data", body)
|
||||
}
|
||||
cmd := exec.CommandContext(ctx, "docker", args...)
|
||||
// -w writes the status code to stdout; -o /dev/stderr writes the body to
|
||||
// stderr so we can capture both separately.
|
||||
|
||||
@@ -22,14 +22,42 @@ const (
|
||||
// matches a real small model commonly served by vLLM so the provider's
|
||||
// enumerated model and the client's request line up.
|
||||
VLLMModel = "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
// VLLMUnlistedModel is a second id the mock's model listing advertises but
|
||||
// no test provider enumerates, so a filtered listing is observably shorter
|
||||
// than the upstream's own.
|
||||
VLLMUnlistedModel = "Qwen/Qwen2.5-7B-Instruct"
|
||||
)
|
||||
|
||||
// Token counts the mock reports per wire shape. Tests assert on these rather
|
||||
// than on "> 0" so a response parsed with the wrong provider's parser (which
|
||||
// would read a different field, or none) fails loudly instead of passing on
|
||||
// a coincidental non-zero.
|
||||
const (
|
||||
// VLLMChatInputTokens / VLLMChatOutputTokens ride the OpenAI usage block.
|
||||
VLLMChatInputTokens = 11
|
||||
VLLMChatOutputTokens = 2
|
||||
// VLLMMessagesInputTokens / VLLMMessagesOutputTokens ride the Anthropic
|
||||
// usage block, whose field names the OpenAI parser cannot read.
|
||||
VLLMMessagesInputTokens = 17
|
||||
VLLMMessagesOutputTokens = 3
|
||||
)
|
||||
|
||||
// vllmNginxConf emulates a vLLM OpenAI-compatible server over plain HTTP (vLLM's
|
||||
// default: no TLS, port 8000). It answers /v1/models with a one-model list and
|
||||
// any chat/completions path with a canned OpenAI-shaped chat completion carrying
|
||||
// a non-zero usage block, so the proxy's OpenAI parser records real token
|
||||
// consumption. Running actual vLLM in CI is infeasible (GPU + multi-GB model
|
||||
// default: no TLS, port 8000), and additionally answers the wire shapes the
|
||||
// other catalog surfaces speak so one mock can stand in for every provider the
|
||||
// proxy routes to. Running actual vLLM in CI is infeasible (GPU + multi-GB model
|
||||
// download), so this stands in for the wire contract the proxy depends on.
|
||||
//
|
||||
// Each shape answers with its own vendor's usage block, so a response parsed
|
||||
// under the wrong surface meters zero rather than passing by accident:
|
||||
//
|
||||
// - /v1/chat/completions (and any unmatched path): OpenAI chat completion.
|
||||
// - /v1/messages: Anthropic Messages, snake_case usage plus a cache bucket.
|
||||
// - /model/{id}/invoke: Bedrock InvokeModel, which carries the Anthropic body.
|
||||
// - the token-counting endpoints: a count, with no usage block at all.
|
||||
//
|
||||
// The model listing advertises two models so a policy that authorises one
|
||||
// produces an observably shorter list than the upstream's own.
|
||||
const vllmNginxConf = `pid /tmp/nginx.pid;
|
||||
events {}
|
||||
http {
|
||||
@@ -37,7 +65,26 @@ http {
|
||||
listen 8000;
|
||||
location = /v1/models {
|
||||
default_type application/json;
|
||||
return 200 '{"object":"list","data":[{"id":"Qwen/Qwen2.5-0.5B-Instruct","object":"model","owned_by":"vllm"}]}';
|
||||
return 200 '{"object":"list","data":[{"id":"Qwen/Qwen2.5-0.5B-Instruct","object":"model","owned_by":"vllm"},{"id":"Qwen/Qwen2.5-7B-Instruct","object":"model","owned_by":"vllm"}]}';
|
||||
}
|
||||
location = /v1/messages {
|
||||
default_type application/json;
|
||||
return 200 '{"id":"msg_e2e","type":"message","role":"assistant","model":"claude-sonnet-5","content":[{"type":"text","text":"pong"}],"stop_reason":"end_turn","usage":{"input_tokens":17,"output_tokens":3,"cache_read_input_tokens":5}}';
|
||||
}
|
||||
location = /v1/messages/count_tokens {
|
||||
default_type application/json;
|
||||
return 200 '{"input_tokens":7}';
|
||||
}
|
||||
location ~ ^/model/.+/invoke$ {
|
||||
default_type application/json;
|
||||
return 200 '{"id":"msg_e2e_bedrock","type":"message","role":"assistant","content":[{"type":"text","text":"pong"}],"stop_reason":"end_turn","usage":{"input_tokens":17,"output_tokens":3,"cache_read_input_tokens":5}}';
|
||||
}
|
||||
location ~ ^/model/.+/count-tokens$ {
|
||||
default_type application/json;
|
||||
return 200 '{"inputTokens":9}';
|
||||
}
|
||||
location = /api/hello {
|
||||
return 200;
|
||||
}
|
||||
location / {
|
||||
default_type application/json;
|
||||
|
||||
@@ -296,6 +296,8 @@ var providers = []Provider{
|
||||
// account to be on >= 30-day data retention or all requests
|
||||
// 400.
|
||||
Models: []Model{
|
||||
{ID: "claude-opus-5", Label: "Claude Opus 5", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-sonnet-5", Label: "Claude Sonnet 5", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000},
|
||||
{ID: "claude-fable-5", Label: "Claude Fable 5", InputPer1k: 0.010, OutputPer1k: 0.050, CacheReadPer1k: 0.001, CacheCreationPer1k: 0.0125, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-8", Label: "Claude Opus 4.8", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-7", Label: "Claude Opus 4.7", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
@@ -355,6 +357,8 @@ var providers = []Provider{
|
||||
// Llama 3.3 70B entry kept unchanged — LiteLLM tracks only
|
||||
// per-region Llama 3 entries; standalone 3.3 not yet listed.
|
||||
Models: []Model{
|
||||
{ID: "anthropic.claude-opus-5", Label: "Claude Opus 5 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-sonnet-5", Label: "Claude Sonnet 5 (Bedrock)", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-opus-4-8", Label: "Claude Opus 4.8 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-opus-4-7", Label: "Claude Opus 4.7 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-opus-4-6", Label: "Claude Opus 4.6 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
@@ -406,6 +410,8 @@ var providers = []Provider{
|
||||
// exists — the router denies unmeterable publishers rather than forward
|
||||
// them uncounted.
|
||||
Models: []Model{
|
||||
{ID: "claude-opus-5", Label: "Claude Opus 5 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-sonnet-5", Label: "Claude Sonnet 5 (Vertex)", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000},
|
||||
{ID: "claude-fable-5", Label: "Claude Fable 5 (Vertex)", InputPer1k: 0.010, OutputPer1k: 0.050, CacheReadPer1k: 0.001, CacheCreationPer1k: 0.0125, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-8", Label: "Claude Opus 4.8 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-7", Label: "Claude Opus 4.7 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
package catalog
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestClaudeLineupSelectable pins the models Claude Code resolves to by
|
||||
// default. A model absent from the lineup can't be ticked on a provider
|
||||
// record, so llm_router denies it as not-routable and the operator has no
|
||||
// way to authorise the client's own default.
|
||||
func TestClaudeLineupSelectable(t *testing.T) {
|
||||
for providerID, wanted := range map[string][]string{
|
||||
"anthropic_api": {"claude-opus-5", "claude-sonnet-5", "claude-haiku-4-5"},
|
||||
"bedrock_api": {"anthropic.claude-opus-5", "anthropic.claude-sonnet-5", "anthropic.claude-haiku-4-5"},
|
||||
"vertex_ai_api": {"claude-opus-5", "claude-sonnet-5", "claude-haiku-4-5"},
|
||||
} {
|
||||
provider, ok := Lookup(providerID)
|
||||
require.True(t, ok, "catalog must define %s", providerID)
|
||||
|
||||
selectable := make(map[string]Model, len(provider.Models))
|
||||
for _, m := range provider.Models {
|
||||
selectable[m.ID] = m
|
||||
}
|
||||
for _, id := range wanted {
|
||||
model, found := selectable[id]
|
||||
require.True(t, found, "%s must offer %s", providerID, id)
|
||||
assert.NotEmpty(t, model.Label, "%s/%s needs a label for the picker", providerID, id)
|
||||
assert.Positive(t, model.InputPer1k, "%s/%s needs an input rate", providerID, id)
|
||||
assert.Positive(t, model.OutputPer1k, "%s/%s needs an output rate", providerID, id)
|
||||
assert.Positive(t, model.ContextWindow, "%s/%s needs a context window", providerID, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -92,13 +92,6 @@ func newAgentNetworkHandlerFixture(t *testing.T) *agentNetworkHandlerFixture {
|
||||
}
|
||||
|
||||
func (f *agentNetworkHandlerFixture) do(t *testing.T, method, path, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
return f.doWithHeaders(t, method, path, body, nil)
|
||||
}
|
||||
|
||||
// doWithHeaders is do with request headers, for the cases where the header is
|
||||
// the thing under test (conditional requests).
|
||||
func (f *agentNetworkHandlerFixture) doWithHeaders(t *testing.T, method, path, body string, headers map[string]string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
var reader io.Reader
|
||||
if body != "" {
|
||||
@@ -108,9 +101,6 @@ func (f *agentNetworkHandlerFixture) doWithHeaders(t *testing.T, method, path, b
|
||||
if body != "" {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
for name, value := range headers {
|
||||
req.Header.Set(name, value)
|
||||
}
|
||||
req = nbcontext.SetUserAuthInRequest(req, auth.UserAuth{
|
||||
UserId: testUserID,
|
||||
AccountId: testAccountID,
|
||||
|
||||
@@ -60,20 +60,12 @@ func (h *handler) createSettings(w http.ResponseWriter, r *http.Request) {
|
||||
util.WriteError(r.Context(), err, w)
|
||||
return
|
||||
}
|
||||
// Emitting the validator here lets a client that just bootstrapped issue a
|
||||
// conditional PUT without an intervening GET.
|
||||
util.SetETag(w, created.ETag())
|
||||
util.WriteJSONObject(r.Context(), w, created.ToAPIResponse())
|
||||
}
|
||||
|
||||
// updateSettings replaces the mutable settings fields on the account's row.
|
||||
// A request carrying a cluster bootstraps the row when the account doesn't
|
||||
// have one yet.
|
||||
//
|
||||
// An If-Match header makes the update conditional: it is honoured against the
|
||||
// stored row inside the write's transaction, and a stale validator is refused
|
||||
// with 412 rather than overwriting what changed since the client read. Omitting
|
||||
// the header keeps the pre-existing last-write-wins behaviour.
|
||||
func (h *handler) updateSettings(w http.ResponseWriter, r *http.Request) {
|
||||
userAuth, err := nbcontext.GetUserAuthFromContext(r.Context())
|
||||
if err != nil {
|
||||
@@ -90,12 +82,11 @@ func (h *handler) updateSettings(w http.ResponseWriter, r *http.Request) {
|
||||
settings := &types.Settings{AccountID: userAuth.AccountId}
|
||||
settings.FromAPIRequest(&req)
|
||||
|
||||
updated, err := h.manager.UpdateSettings(r.Context(), userAuth.UserId, settings, util.IfMatch(r))
|
||||
updated, err := h.manager.UpdateSettings(r.Context(), userAuth.UserId, settings)
|
||||
if err != nil {
|
||||
util.WriteError(r.Context(), err, w)
|
||||
return
|
||||
}
|
||||
util.SetETag(w, updated.ETag())
|
||||
util.WriteJSONObject(r.Context(), w, updated.ToAPIResponse())
|
||||
}
|
||||
|
||||
@@ -103,11 +94,6 @@ func (h *handler) updateSettings(w http.ResponseWriter, r *http.Request) {
|
||||
// The manager refuses (412) while providers exist or a proxy is actively
|
||||
// serving the endpoint; a later POST bootstraps fresh, allocating a new
|
||||
// endpoint.
|
||||
//
|
||||
// An If-Match header makes the delete conditional, and is worth sending here
|
||||
// even more than on update: both existing guards are about state rather than
|
||||
// staleness, so nothing else stops a client from deleting a row that was
|
||||
// replaced since it read one.
|
||||
func (h *handler) deleteSettings(w http.ResponseWriter, r *http.Request) {
|
||||
userAuth, err := nbcontext.GetUserAuthFromContext(r.Context())
|
||||
if err != nil {
|
||||
@@ -115,7 +101,7 @@ func (h *handler) deleteSettings(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.manager.DeleteSettings(r.Context(), userAuth.AccountId, userAuth.UserId, util.IfMatch(r)); err != nil {
|
||||
if err := h.manager.DeleteSettings(r.Context(), userAuth.AccountId, userAuth.UserId); err != nil {
|
||||
util.WriteError(r.Context(), err, w)
|
||||
return
|
||||
}
|
||||
@@ -137,9 +123,5 @@ func (h *handler) getSettings(w http.ResponseWriter, r *http.Request) {
|
||||
util.WriteError(r.Context(), err, w)
|
||||
return
|
||||
}
|
||||
// The pre-bootstrap defaults are a representation like any other and carry
|
||||
// a validator too, so an If-Match taken before bootstrap cannot silently
|
||||
// match the row that appeared since.
|
||||
util.SetETag(w, settings.ETag())
|
||||
util.WriteJSONObject(r.Context(), w, settings.ToAPIResponse())
|
||||
}
|
||||
|
||||
@@ -393,202 +393,3 @@ func TestSettingsHandler_DeleteReleasesEndpointForFreshBootstrap(t *testing.T) {
|
||||
"the fresh row must carry bootstrap defaults, not the deleted row's toggles")
|
||||
assert.NotNil(t, second.CreatedAt, "the fresh row is persisted and carries timestamps")
|
||||
}
|
||||
|
||||
// bootstrapForETag bootstraps a settings row and returns the response body
|
||||
// alongside the validator the bootstrap emitted, which is what a client would
|
||||
// carry into its first conditional write.
|
||||
func bootstrapForETag(t *testing.T, f *agentNetworkHandlerFixture) (api.AgentNetworkSettings, string) {
|
||||
t.Helper()
|
||||
|
||||
rec := f.do(t, http.MethodPost, "/agent-network/settings",
|
||||
`{"proxy_address": "eu.proxy.netbird.io", "enable_prompt_collection": true, "redact_pii": true, "access_log_retention_days": 14}`)
|
||||
require.Equal(t, http.StatusOK, rec.Code, "bootstrap POST must succeed: %s", rec.Body.String())
|
||||
|
||||
var settings api.AgentNetworkSettings
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &settings))
|
||||
|
||||
etag := rec.Header().Get("ETag")
|
||||
require.NotEmpty(t, etag, "bootstrap must emit a validator so a client can PUT without an intervening GET")
|
||||
return settings, etag
|
||||
}
|
||||
|
||||
// putBody renders a complete settings update — every field, with the identity
|
||||
// echo the endpoint requires — so the conditional-request tests differ only in
|
||||
// their headers.
|
||||
func putBody(settings api.AgentNetworkSettings, redactPii bool, retention int) string {
|
||||
return fmt.Sprintf(
|
||||
`{"endpoint": %q, "proxy_address": %q, "enable_log_collection": true, "enable_prompt_collection": true, "redact_pii": %t, "access_log_retention_days": %d}`,
|
||||
settings.Endpoint, settings.ProxyAddress, redactPii, retention)
|
||||
}
|
||||
|
||||
// TestSettingsHandler_EmitsETag pins that every read and every write hands the
|
||||
// client back a validator, quoted as a strong entity-tag. Without one on the
|
||||
// write responses a client would have to re-GET after every update to stay
|
||||
// able to make the next one conditional.
|
||||
func TestSettingsHandler_EmitsETag(t *testing.T) {
|
||||
f := newAgentNetworkHandlerFixture(t)
|
||||
|
||||
// The pre-bootstrap defaults are a representation too, and validate like
|
||||
// one — an If-Match taken here must not match the row that appears later.
|
||||
rec := f.do(t, http.MethodGet, "/agent-network/settings", "")
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
defaultsETag := rec.Header().Get("ETag")
|
||||
assert.NotEmpty(t, defaultsETag, "the unbootstrapped view must carry a validator")
|
||||
|
||||
settings, bootstrapETag := bootstrapForETag(t, f)
|
||||
assert.Regexp(t, `^"[0-9a-f]+"$`, bootstrapETag, "the validator must be a quoted strong entity-tag")
|
||||
assert.NotEqual(t, defaultsETag, bootstrapETag, "bootstrapping must move the validator")
|
||||
|
||||
rec = f.do(t, http.MethodGet, "/agent-network/settings", "")
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
assert.Equal(t, bootstrapETag, rec.Header().Get("ETag"),
|
||||
"reading an unchanged row must derive the same validator the bootstrap returned")
|
||||
|
||||
rec = f.do(t, http.MethodPut, "/agent-network/settings", putBody(settings, false, 7))
|
||||
require.Equal(t, http.StatusOK, rec.Code, "update must succeed: %s", rec.Body.String())
|
||||
assert.NotEqual(t, bootstrapETag, rec.Header().Get("ETag"),
|
||||
"an update that changed the representation must return a different validator")
|
||||
}
|
||||
|
||||
// TestSettingsHandler_PutIfMatch walks the conditional-update contract. The
|
||||
// stale case is the one the feature exists for: a client that planned against
|
||||
// an earlier read must be refused rather than silently reverting whatever
|
||||
// changed in between — RedactPii above all, where a silent revert turns a
|
||||
// compliance control off with no error and no drift warning.
|
||||
func TestSettingsHandler_PutIfMatch(t *testing.T) {
|
||||
t.Run("matching validator succeeds", func(t *testing.T) {
|
||||
f := newAgentNetworkHandlerFixture(t)
|
||||
settings, etag := bootstrapForETag(t, f)
|
||||
|
||||
rec := f.doWithHeaders(t, http.MethodPut, "/agent-network/settings",
|
||||
putBody(settings, false, 7), map[string]string{"If-Match": etag})
|
||||
require.Equal(t, http.StatusOK, rec.Code,
|
||||
"a matching precondition must be honoured: got %d body=%s", rec.Code, rec.Body.String())
|
||||
assert.NotEqual(t, etag, rec.Header().Get("ETag"),
|
||||
"the response must carry the new validator, not the one that was matched")
|
||||
})
|
||||
|
||||
t.Run("stale validator is refused and changes nothing", func(t *testing.T) {
|
||||
f := newAgentNetworkHandlerFixture(t)
|
||||
settings, stale := bootstrapForETag(t, f)
|
||||
|
||||
// Someone else writes in between — the dashboard operator enabling
|
||||
// something the planning client never saw.
|
||||
rec := f.do(t, http.MethodPut, "/agent-network/settings", putBody(settings, true, 21))
|
||||
require.Equal(t, http.StatusOK, rec.Code, "the intervening update must succeed: %s", rec.Body.String())
|
||||
var intervened api.AgentNetworkSettings
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &intervened))
|
||||
|
||||
rec = f.doWithHeaders(t, http.MethodPut, "/agent-network/settings",
|
||||
putBody(settings, false, 7), map[string]string{"If-Match": stale})
|
||||
require.Equal(t, http.StatusPreconditionFailed, rec.Code,
|
||||
"a stale precondition must be refused: got %d body=%s", rec.Code, rec.Body.String())
|
||||
|
||||
// Asserting the state, not just the status: a partial write would pass
|
||||
// a status-only check.
|
||||
rec = f.do(t, http.MethodGet, "/agent-network/settings", "")
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var after api.AgentNetworkSettings
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &after))
|
||||
assert.Equal(t, intervened, after, "the refused update must leave the row byte-identical")
|
||||
})
|
||||
|
||||
t.Run("star matches the existing row", func(t *testing.T) {
|
||||
f := newAgentNetworkHandlerFixture(t)
|
||||
settings, _ := bootstrapForETag(t, f)
|
||||
|
||||
rec := f.doWithHeaders(t, http.MethodPut, "/agent-network/settings",
|
||||
putBody(settings, false, 7), map[string]string{"If-Match": "*"})
|
||||
assert.Equal(t, http.StatusOK, rec.Code,
|
||||
"* must match any current representation: got %d body=%s", rec.Code, rec.Body.String())
|
||||
})
|
||||
|
||||
t.Run("no precondition still succeeds", func(t *testing.T) {
|
||||
f := newAgentNetworkHandlerFixture(t)
|
||||
settings, _ := bootstrapForETag(t, f)
|
||||
|
||||
// The back-compatibility guarantee: clients that predate conditional
|
||||
// requests — the dashboard among them — keep last-write-wins.
|
||||
rec := f.do(t, http.MethodPut, "/agent-network/settings", putBody(settings, false, 7))
|
||||
assert.Equal(t, http.StatusOK, rec.Code,
|
||||
"an unconditional update must keep working: got %d body=%s", rec.Code, rec.Body.String())
|
||||
})
|
||||
|
||||
t.Run("precondition is checked before the immutability echo", func(t *testing.T) {
|
||||
f := newAgentNetworkHandlerFixture(t)
|
||||
settings, stale := bootstrapForETag(t, f)
|
||||
|
||||
rec := f.do(t, http.MethodPut, "/agent-network/settings", putBody(settings, true, 21))
|
||||
require.Equal(t, http.StatusOK, rec.Code, "the intervening update must succeed: %s", rec.Body.String())
|
||||
|
||||
// A client stale enough to hold an old validator may be stale in its
|
||||
// identity echo too. Answering 412 tells it the useful thing — go and
|
||||
// read again — where 422 would send it hunting an immutability bug.
|
||||
body := fmt.Sprintf(
|
||||
`{"endpoint": "other.gateway.example.com", "proxy_address": %q, "enable_log_collection": true, "enable_prompt_collection": true, "redact_pii": false, "access_log_retention_days": 7}`,
|
||||
settings.ProxyAddress)
|
||||
rec = f.doWithHeaders(t, http.MethodPut, "/agent-network/settings", body,
|
||||
map[string]string{"If-Match": stale})
|
||||
assert.Equal(t, http.StatusPreconditionFailed, rec.Code,
|
||||
"staleness must be reported ahead of the identity mismatch: got %d body=%s", rec.Code, rec.Body.String())
|
||||
})
|
||||
}
|
||||
|
||||
// TestSettingsHandler_DeleteIfMatch covers the conditional delete, which
|
||||
// carries more weight than the conditional update: both existing delete guards
|
||||
// are about state — no providers, no serving proxy — so nothing else stops a
|
||||
// client from deleting a row that was replaced since it read one.
|
||||
func TestSettingsHandler_DeleteIfMatch(t *testing.T) {
|
||||
t.Run("stale validator is refused and the row survives", func(t *testing.T) {
|
||||
f := newAgentNetworkHandlerFixture(t)
|
||||
settings, stale := bootstrapForETag(t, f)
|
||||
|
||||
rec := f.do(t, http.MethodPut, "/agent-network/settings", putBody(settings, true, 21))
|
||||
require.Equal(t, http.StatusOK, rec.Code, "the intervening update must succeed: %s", rec.Body.String())
|
||||
|
||||
rec = f.doWithHeaders(t, http.MethodDelete, "/agent-network/settings", "",
|
||||
map[string]string{"If-Match": stale})
|
||||
require.Equal(t, http.StatusPreconditionFailed, rec.Code,
|
||||
"a stale precondition must refuse the delete: got %d body=%s", rec.Code, rec.Body.String())
|
||||
|
||||
rec = f.do(t, http.MethodGet, "/agent-network/settings", "")
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var after api.AgentNetworkSettings
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &after))
|
||||
assert.Equal(t, settings.Endpoint, after.Endpoint, "the refused delete must leave the row in place")
|
||||
})
|
||||
|
||||
t.Run("matching validator deletes", func(t *testing.T) {
|
||||
f := newAgentNetworkHandlerFixture(t)
|
||||
_, etag := bootstrapForETag(t, f)
|
||||
|
||||
rec := f.doWithHeaders(t, http.MethodDelete, "/agent-network/settings", "",
|
||||
map[string]string{"If-Match": etag})
|
||||
require.Equal(t, http.StatusOK, rec.Code,
|
||||
"a matching precondition must be honoured: got %d body=%s", rec.Code, rec.Body.String())
|
||||
|
||||
rec = f.do(t, http.MethodGet, "/agent-network/settings", "")
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var after api.AgentNetworkSettings
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &after))
|
||||
assert.Empty(t, after.Endpoint, "the row must be gone")
|
||||
})
|
||||
|
||||
t.Run("precondition is checked before the state guards", func(t *testing.T) {
|
||||
f := newAgentNetworkHandlerFixture(t)
|
||||
settings, stale := bootstrapForETag(t, f)
|
||||
|
||||
rec := f.do(t, http.MethodPut, "/agent-network/settings", putBody(settings, true, 21))
|
||||
require.Equal(t, http.StatusOK, rec.Code, "the intervening update must succeed: %s", rec.Body.String())
|
||||
f.seedProvider(t, "prov-precondition")
|
||||
|
||||
// Both refusals are 412, so the status cannot tell them apart — the
|
||||
// message must, or a stale client is sent to delete providers it may
|
||||
// not even know about.
|
||||
rec = f.doWithHeaders(t, http.MethodDelete, "/agent-network/settings", "",
|
||||
map[string]string{"If-Match": stale})
|
||||
require.Equal(t, http.StatusPreconditionFailed, rec.Code, "the delete must be refused: %s", rec.Body.String())
|
||||
assert.Contains(t, rec.Body.String(), "if-match",
|
||||
"staleness must be reported ahead of the provider guard: %s", rec.Body.String())
|
||||
})
|
||||
}
|
||||
|
||||
@@ -22,7 +22,6 @@ import (
|
||||
"github.com/netbirdio/netbird/management/server/permissions/modules"
|
||||
"github.com/netbirdio/netbird/management/server/permissions/operations"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
httputil "github.com/netbirdio/netbird/shared/management/http/util"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
@@ -72,8 +71,8 @@ type Manager interface {
|
||||
|
||||
GetSettings(ctx context.Context, accountID, userID string) (*types.Settings, error)
|
||||
CreateSettings(ctx context.Context, userID string, settings *types.Settings, proxyAddress, endpoint string) (*types.Settings, error)
|
||||
UpdateSettings(ctx context.Context, userID string, settings *types.Settings, precondition *httputil.Precondition) (*types.Settings, error)
|
||||
DeleteSettings(ctx context.Context, accountID, userID string, precondition *httputil.Precondition) error
|
||||
UpdateSettings(ctx context.Context, userID string, settings *types.Settings) (*types.Settings, error)
|
||||
DeleteSettings(ctx context.Context, accountID, userID string) error
|
||||
|
||||
ListConsumption(ctx context.Context, accountID, userID string) ([]*types.Consumption, error)
|
||||
ListAccessLogs(ctx context.Context, accountID, userID string, filter types.AgentNetworkAccessLogFilter) ([]*types.AgentNetworkAccessLog, int64, error)
|
||||
@@ -545,13 +544,6 @@ func (m *managerImpl) DeleteBudgetRule(ctx context.Context, accountID, userID, r
|
||||
return nil
|
||||
}
|
||||
|
||||
// stalePreconditionMsg is the refusal both conditional settings writes return.
|
||||
// Shared so the two cannot drift: DeleteSettings answers 412 for its state
|
||||
// guards as well, so the message is the only thing telling a client that it is
|
||||
// working from an old read rather than tripping over providers or a serving
|
||||
// proxy.
|
||||
const stalePreconditionMsg = "if-match precondition failed: the settings have changed since they were read; GET them again and retry"
|
||||
|
||||
// UpdateSettings replaces the mutable account-level settings — the collection
|
||||
// toggles and retention — on the account's row. The identity fields (Domain,
|
||||
// ProxyAddress) are assigned at bootstrap (CreateSettings) and immutable: the
|
||||
@@ -562,11 +554,7 @@ const stalePreconditionMsg = "if-match precondition failed: the settings have ch
|
||||
// Because the collection toggles change the synthesised service config
|
||||
// (prompt-capture gating, access-log emission), a reconcile is triggered so
|
||||
// the proxy and peer network maps converge on the new state.
|
||||
//
|
||||
// precondition carries the caller's If-Match, and is nil for an unconditional
|
||||
// update — last write wins, which is what the dashboard wants and what every
|
||||
// client that predates conditional requests gets.
|
||||
func (m *managerImpl) UpdateSettings(ctx context.Context, userID string, settings *types.Settings, precondition *httputil.Precondition) (*types.Settings, error) {
|
||||
func (m *managerImpl) UpdateSettings(ctx context.Context, userID string, settings *types.Settings) (*types.Settings, error) {
|
||||
if err := m.requirePermission(ctx, settings.AccountID, userID, modules.AgentNetworkSettings, operations.Update); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -585,20 +573,6 @@ func (m *managerImpl) UpdateSettings(ctx context.Context, userID string, setting
|
||||
return fmt.Errorf("get agent network settings: %w", err)
|
||||
}
|
||||
|
||||
// Evaluated here, under the row lock and inside the write's own
|
||||
// transaction, rather than in the handler: comparing before the
|
||||
// transaction only narrows the race, since two requests can both pass
|
||||
// the check before either writes. Locking the row first makes it a
|
||||
// genuine compare-and-set.
|
||||
//
|
||||
// It comes before the identity comparison because a client holding a
|
||||
// stale validator is stale in its identity echo too, and "you are
|
||||
// working from an old read" is the more accurate answer than "the
|
||||
// endpoint is immutable".
|
||||
if !precondition.Matches(existing.ETag()) {
|
||||
return status.Errorf(status.PreconditionFailed, "%s", stalePreconditionMsg)
|
||||
}
|
||||
|
||||
// The identity echo is compared leniently (trimmed, case-insensitive):
|
||||
// the stored values are normalized lowercase, and a client replaying a
|
||||
// GET response must never be rejected over casing it didn't choose.
|
||||
@@ -661,12 +635,7 @@ func hostnamesEquivalent(supplied, stored string) bool {
|
||||
// is not reserved. That full-reset semantic is what gives clients that model
|
||||
// immutability as replace-on-change (e.g. Terraform's RequiresReplace) a real
|
||||
// path: tear down providers, delete, re-create.
|
||||
//
|
||||
// precondition carries the caller's If-Match, and is nil for an unconditional
|
||||
// delete. It matters more here than on update: the two guards above are about
|
||||
// state rather than staleness, so without it nothing stops a client from
|
||||
// deleting a row that was replaced since it last read one.
|
||||
func (m *managerImpl) DeleteSettings(ctx context.Context, accountID, userID string, precondition *httputil.Precondition) error {
|
||||
func (m *managerImpl) DeleteSettings(ctx context.Context, accountID, userID string) error {
|
||||
if err := m.requirePermission(ctx, accountID, userID, modules.AgentNetworkSettings, operations.Delete); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -682,13 +651,6 @@ func (m *managerImpl) DeleteSettings(ctx context.Context, accountID, userID stri
|
||||
return fmt.Errorf("get agent network settings: %w", err)
|
||||
}
|
||||
|
||||
// Under the row lock, for the same reason as in UpdateSettings, and
|
||||
// before the state guards: a caller working from an old read should
|
||||
// learn that first, not be told about providers it may not know exist.
|
||||
if !precondition.Matches(existing.ETag()) {
|
||||
return status.Errorf(status.PreconditionFailed, "%s", stalePreconditionMsg)
|
||||
}
|
||||
|
||||
providers, err := tx.GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, accountID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("get agent network providers: %w", err)
|
||||
@@ -1138,13 +1100,11 @@ func (*mockManager) CreateSettings(_ context.Context, _ string, s *types.Setting
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (*mockManager) UpdateSettings(_ context.Context, _ string, s *types.Settings, _ *httputil.Precondition) (*types.Settings, error) {
|
||||
func (*mockManager) UpdateSettings(_ context.Context, _ string, s *types.Settings) (*types.Settings, error) {
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (*mockManager) DeleteSettings(_ context.Context, _, _ string, _ *httputil.Precondition) error {
|
||||
return nil
|
||||
}
|
||||
func (*mockManager) DeleteSettings(_ context.Context, _, _ string) error { return nil }
|
||||
|
||||
func (*mockManager) ListConsumption(_ context.Context, _, _ string) ([]*types.Consumption, error) {
|
||||
return nil, nil
|
||||
|
||||
@@ -47,17 +47,11 @@ var supplementalDefaults = map[string]map[string]Entry{
|
||||
"gpt-5-nano": {InputPer1k: 0.00005, OutputPer1k: 0.0004, CachedInputPer1k: 0.000005},
|
||||
},
|
||||
"anthropic": {
|
||||
// claude-opus-5 is not yet in the catalog lineup but gateway /
|
||||
// grandfathered traffic uses it; priced so it isn't skipped.
|
||||
"claude-opus-5": {InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625},
|
||||
// "kimi-k3[1m]" is the 1M-context alias some Claude Code guides
|
||||
// configure against Moonshot's Anthropic-compatible endpoint;
|
||||
// priced identically to kimi-k3 so those requests aren't skipped.
|
||||
"kimi-k3[1m]": {InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003},
|
||||
},
|
||||
"bedrock": {
|
||||
"anthropic.claude-opus-5": {InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625},
|
||||
},
|
||||
}
|
||||
|
||||
var (
|
||||
|
||||
@@ -82,6 +82,11 @@ anthropic:
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
claude-sonnet-5:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
kimi-k3:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
@@ -145,6 +150,11 @@ bedrock:
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
anthropic.claude-sonnet-5:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
meta.llama3-3-70b-instruct:
|
||||
input_per_1k: 0.00072
|
||||
output_per_1k: 0.00072
|
||||
|
||||
@@ -116,11 +116,13 @@ func TestDefaultTable_PinnedRates(t *testing.T) {
|
||||
assert.InDelta(t, 0.010, fable.InputPer1k, 1e-9, "claude-fable-5 input")
|
||||
assert.InDelta(t, 0.0125, fable.CacheCreationPer1k, 1e-9, "claude-fable-5 cache creation")
|
||||
|
||||
// Supplementals present on their surfaces.
|
||||
// Every id below must stay priced whichever source provides it: the
|
||||
// catalog lineup for the current Claude 5 family, supplementalDefaults
|
||||
// for the ids the dashboard deliberately doesn't offer.
|
||||
for surface, ids := range map[string][]string{
|
||||
"openai": {"gpt-5", "gpt-5-mini", "gpt-5-nano"},
|
||||
"anthropic": {"claude-opus-5", "kimi-k3[1m]", "kimi-k3"},
|
||||
"bedrock": {"anthropic.claude-opus-5"},
|
||||
"anthropic": {"claude-opus-5", "claude-sonnet-5", "kimi-k3[1m]", "kimi-k3"},
|
||||
"bedrock": {"anthropic.claude-opus-5", "anthropic.claude-sonnet-5"},
|
||||
} {
|
||||
for _, id := range ids {
|
||||
_, ok := table[surface][id]
|
||||
|
||||
@@ -1,199 +0,0 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"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/permissions/modules"
|
||||
"github.com/netbirdio/netbird/management/server/permissions/operations"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
httputil "github.com/netbirdio/netbird/shared/management/http/util"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// ifMatch builds the precondition a client sending this validator would
|
||||
// produce, by going through the same header parse the handler uses rather than
|
||||
// reaching past it.
|
||||
func ifMatch(t *testing.T, etag string) *httputil.Precondition {
|
||||
t.Helper()
|
||||
|
||||
r := httptest.NewRequest(http.MethodPut, "/", nil)
|
||||
r.Header.Set("If-Match", strconv.Quote(etag))
|
||||
return httputil.IfMatch(r)
|
||||
}
|
||||
|
||||
// updateFor renders a complete update for the given row, echoing the identity
|
||||
// fields the endpoint requires and setting retention to tell writers apart.
|
||||
func updateFor(settings *types.Settings, retention int) *types.Settings {
|
||||
return &types.Settings{
|
||||
AccountID: settings.AccountID,
|
||||
Domain: settings.Domain,
|
||||
ProxyAddress: settings.ProxyAddress,
|
||||
EnableLogCollection: true,
|
||||
EnablePromptCollection: true,
|
||||
RedactPii: true,
|
||||
AccessLogRetentionDays: retention,
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpdateSettingsPreconditionSerializesConcurrentWriters is the test the
|
||||
// design rests on. Two writers start from the same validator and race; exactly
|
||||
// one may win.
|
||||
//
|
||||
// An implementation that compares the validator before opening the write
|
||||
// transaction passes every sequential test in this suite and fails here: both
|
||||
// writers read the same row, both find their precondition satisfied, and both
|
||||
// then write — which is the lost update the feature exists to prevent, merely
|
||||
// narrowed to a smaller window. Holding the row under LockingStrengthUpdate
|
||||
// and comparing inside the write's own transaction is what makes it a genuine
|
||||
// compare-and-set.
|
||||
//
|
||||
// The test store is sqlite, which serializes writers of its own accord, so
|
||||
// what this pins directly is the outcome — exactly one success — rather than
|
||||
// the mechanism. It still has teeth against the check-before-transaction
|
||||
// shape, whose two reads interleave freely before either write. Running it
|
||||
// against postgres (NB_STORE_ENGINE_POSTGRES_DSN) exercises real concurrent
|
||||
// transactions.
|
||||
func TestUpdateSettingsPreconditionSerializesConcurrentWriters(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
f := newBootstrapFixture(t)
|
||||
|
||||
const accountID, userID = "account1", "user1"
|
||||
f.expectPermission(accountID, userID, modules.AgentNetworkSettings, operations.Create, true)
|
||||
f.expectPermission(accountID, userID, modules.AgentNetworkSettings, operations.Update, true)
|
||||
f.expectPermission(accountID, userID, modules.AgentNetworkSettings, operations.Update, true)
|
||||
|
||||
created, err := f.createSettings(ctx, accountID, userID, "cluster1.example.com", "")
|
||||
require.NoError(t, err, "bootstrap must succeed")
|
||||
|
||||
// Both writers plan against this one read, as a client that read, computed
|
||||
// a diff and is about to write the whole object back would.
|
||||
shared := created.ETag()
|
||||
|
||||
// noWrite is a retention value neither writer sends and the API would
|
||||
// never store, so an assertion that lands on it is a test bug rather than
|
||||
// a silently satisfied comparison. Zero would not do: the API documents 0
|
||||
// as "keep indefinitely", so it is a value the row could legitimately hold.
|
||||
const noWrite = -1
|
||||
|
||||
var (
|
||||
wg sync.WaitGroup
|
||||
start = make(chan struct{})
|
||||
errs = make([]error, 2)
|
||||
wrote = []int{7, 21}
|
||||
returned = []int{noWrite, noWrite}
|
||||
)
|
||||
for i := range 2 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
updated, err := f.manager.UpdateSettings(ctx, userID, updateFor(created, wrote[i]), ifMatch(t, shared))
|
||||
errs[i] = err
|
||||
if err == nil {
|
||||
returned[i] = updated.AccessLogRetentionDays
|
||||
}
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
|
||||
succeeded, winner := 0, noWrite
|
||||
for i, err := range errs {
|
||||
if err == nil {
|
||||
succeeded++
|
||||
winner = wrote[i]
|
||||
assert.Equal(t, wrote[i], returned[i], "the winner's response must carry what it sent")
|
||||
continue
|
||||
}
|
||||
assert.Truef(t, isPreconditionFailed(err),
|
||||
"the losing writer must be refused for staleness, got: %v (writer %d)", err, i)
|
||||
}
|
||||
require.Equal(t, 1, succeeded, "exactly one writer may win: %v", errs)
|
||||
|
||||
// The row must carry the winner's value and nothing blended.
|
||||
stored, err := f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID)
|
||||
require.NoError(t, err, "the row must survive the race")
|
||||
assert.Equal(t, winner, stored.AccessLogRetentionDays,
|
||||
"the stored row must be exactly what the winning writer sent")
|
||||
assert.NotEqual(t, shared, stored.ETag(), "the surviving row must derive a new validator")
|
||||
}
|
||||
|
||||
// TestUpdateSettingsUnconditionalIgnoresStaleness pins the back-compatibility
|
||||
// half: without a precondition the manager keeps last-write-wins, which is
|
||||
// what the dashboard relies on and what any client that predates conditional
|
||||
// requests does.
|
||||
func TestUpdateSettingsUnconditionalIgnoresStaleness(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
f := newBootstrapFixture(t)
|
||||
|
||||
const accountID, userID = "account1", "user1"
|
||||
f.expectPermission(accountID, userID, modules.AgentNetworkSettings, operations.Create, true)
|
||||
f.expectPermission(accountID, userID, modules.AgentNetworkSettings, operations.Update, true)
|
||||
f.expectPermission(accountID, userID, modules.AgentNetworkSettings, operations.Update, true)
|
||||
|
||||
created, err := f.createSettings(ctx, accountID, userID, "cluster1.example.com", "")
|
||||
require.NoError(t, err, "bootstrap must succeed")
|
||||
|
||||
_, err = f.manager.UpdateSettings(ctx, userID, updateFor(created, 21), nil)
|
||||
require.NoError(t, err, "the first unconditional update must succeed")
|
||||
|
||||
// The second writer is working from a read that is now stale, and with no
|
||||
// precondition it overwrites regardless.
|
||||
updated, err := f.manager.UpdateSettings(ctx, userID, updateFor(created, 7), nil)
|
||||
require.NoError(t, err, "an unconditional update must not be refused for staleness")
|
||||
assert.Equal(t, 7, updated.AccessLogRetentionDays, "last write wins without a precondition")
|
||||
}
|
||||
|
||||
// TestDeleteSettingsPreconditionRefusesStale pins the conditional delete at
|
||||
// the manager level: a stale validator refuses, and the row is still there
|
||||
// afterwards. Deletion is the destructive operation and its two other guards
|
||||
// are about state rather than staleness, so this is the only thing standing
|
||||
// between a client working from an old read and a released endpoint.
|
||||
func TestDeleteSettingsPreconditionRefusesStale(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
f := newBootstrapFixture(t)
|
||||
|
||||
const accountID, userID = "account1", "user1"
|
||||
f.expectPermission(accountID, userID, modules.AgentNetworkSettings, operations.Create, true)
|
||||
f.expectPermission(accountID, userID, modules.AgentNetworkSettings, operations.Update, true)
|
||||
f.expectPermission(accountID, userID, modules.AgentNetworkSettings, operations.Delete, true)
|
||||
f.expectPermission(accountID, userID, modules.AgentNetworkSettings, operations.Delete, true)
|
||||
|
||||
created, err := f.createSettings(ctx, accountID, userID, "cluster1.example.com", "")
|
||||
require.NoError(t, err, "bootstrap must succeed")
|
||||
stale := created.ETag()
|
||||
|
||||
updated, err := f.manager.UpdateSettings(ctx, userID, updateFor(created, 21), nil)
|
||||
require.NoError(t, err, "the intervening update must succeed")
|
||||
|
||||
err = f.manager.DeleteSettings(ctx, accountID, userID, ifMatch(t, stale))
|
||||
require.Error(t, err, "a stale precondition must refuse the delete")
|
||||
assert.True(t, isPreconditionFailed(err), "the refusal must be a precondition failure, got: %v", err)
|
||||
|
||||
stored, err := f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID)
|
||||
require.NoError(t, err, "the refused delete must leave the row in place")
|
||||
assert.Equal(t, created.Domain, stored.Domain, "the endpoint must not have been released")
|
||||
|
||||
// The validator the intervening update returned is the current one, and
|
||||
// deleting with it goes through.
|
||||
require.NoError(t, f.manager.DeleteSettings(ctx, accountID, userID, ifMatch(t, updated.ETag())),
|
||||
"a matching precondition must be honoured")
|
||||
_, err = f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID)
|
||||
assert.Error(t, err, "the row must be gone")
|
||||
}
|
||||
|
||||
// isPreconditionFailed reports whether err is the 412-mapped status error.
|
||||
func isPreconditionFailed(err error) bool {
|
||||
var sErr *status.Error
|
||||
return errors.As(err, &sErr) && sErr.Type() == status.PreconditionFailed
|
||||
}
|
||||
@@ -1,8 +1,6 @@
|
||||
package types
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -69,64 +67,6 @@ func DefaultSettings(accountID string) *Settings {
|
||||
}
|
||||
}
|
||||
|
||||
// etagLength is how much of the hash the validator carries. 16 hex characters
|
||||
// — 64 bits — is far more than enough to make an accidental collision between
|
||||
// two representations of one account's settings unreachable, and keeps the
|
||||
// header short enough to read in a log line.
|
||||
const etagLength = 16
|
||||
|
||||
// ETag returns a strong validator over the settings representation, for
|
||||
// conditional requests (RFC 9110 If-Match). The value is unquoted; applying
|
||||
// the quoting is the transport layer's job.
|
||||
//
|
||||
// The hash covers an explicit field tuple rather than the marshalled API
|
||||
// representation: field ordering in the generated API types is not a contract,
|
||||
// so hashing serialized output would make the validator churn with codegen.
|
||||
// Two exclusions are deliberate:
|
||||
//
|
||||
// - AccountID identifies the resource — it is the URL, not the
|
||||
// representation. Including it would make the validator differ between
|
||||
// accounts whose settings are genuinely identical, which no client can
|
||||
// observe and no precondition needs.
|
||||
// - UpdatedAt is excluded so that equal representations always yield equal
|
||||
// validators. A write that changes nothing must not invalidate a
|
||||
// precondition another client is holding.
|
||||
//
|
||||
// Everything else is in, including the identity fields and CreatedAt. A
|
||||
// validator that covered only the mutable toggles would survive a delete
|
||||
// followed by a fresh bootstrap onto the same toggle values, and an If-Match
|
||||
// held across that gap would then authorize a write against what is really a
|
||||
// different resource. CreatedAt is what distinguishes the re-bootstrapped row.
|
||||
//
|
||||
// CreatedAt is hashed at whole-second precision because the validator has to
|
||||
// agree across a store round-trip. A freshly bootstrapped row derives its
|
||||
// validator in memory, from a time.Time carrying nanoseconds, while every
|
||||
// later comparison derives it from a row read back out of the store — and the
|
||||
// engines truncate: PostgreSQL to microseconds, MySQL DATETIME to whole
|
||||
// seconds without an fsp. At nanosecond precision the two never agree again,
|
||||
// so the validator a bootstrap hands out is permanently unusable. Seconds is
|
||||
// the floor every supported engine preserves. The cost is that a delete and
|
||||
// re-bootstrap within the same second, onto the same endpoint and the same
|
||||
// toggles, derives the same validator; a labeled bootstrap draws a fresh
|
||||
// random label, so that needs a self-addressed endpoint reclaimed inside one
|
||||
// second.
|
||||
//
|
||||
// Adding a field to Settings means deciding whether it belongs here; the
|
||||
// field-count guard in the tests is what forces that decision.
|
||||
func (s *Settings) ETag() string {
|
||||
h := sha256.New()
|
||||
fmt.Fprintf(h, "%s\x00%s\x00%t\x00%t\x00%t\x00%d\x00%d",
|
||||
s.Domain,
|
||||
s.ProxyAddress,
|
||||
s.EnableLogCollection,
|
||||
s.EnablePromptCollection,
|
||||
s.RedactPii,
|
||||
s.AccessLogRetentionDays,
|
||||
s.CreatedAt.Unix(),
|
||||
)
|
||||
return hex.EncodeToString(h.Sum(nil))[:etagLength]
|
||||
}
|
||||
|
||||
// Endpoint returns the bare hostname agents reach this account at — the
|
||||
// Domain column. Empty until the row is bootstrapped.
|
||||
func (s *Settings) Endpoint() string { return s.Domain }
|
||||
|
||||
@@ -1,184 +0,0 @@
|
||||
package types
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// etagSettings is a fully populated settings row — every hashed field set to a
|
||||
// distinctive value — so a mutation test can flip exactly one thing at a time.
|
||||
// The timestamp carries sub-second precision on purpose: a whole-second value
|
||||
// would make the precision test below pass without proving anything.
|
||||
func etagSettings() *Settings {
|
||||
created := time.Date(2026, 8, 11, 9, 30, 0, 123456789, time.UTC)
|
||||
return &Settings{
|
||||
AccountID: "acc-1",
|
||||
Domain: "cool-otter.eu.proxy.netbird.io",
|
||||
ProxyAddress: "eu.proxy.netbird.io",
|
||||
EnableLogCollection: true,
|
||||
EnablePromptCollection: true,
|
||||
RedactPii: true,
|
||||
AccessLogRetentionDays: 30,
|
||||
CreatedAt: created,
|
||||
UpdatedAt: created,
|
||||
}
|
||||
}
|
||||
|
||||
// TestSettings_ETagShape pins the wire shape of the validator: a bare
|
||||
// lowercase hex string of the documented length, with no quoting — quoting is
|
||||
// the transport layer's job, and a validator that arrived pre-quoted would be
|
||||
// double-quoted on the way out.
|
||||
func TestSettings_ETagShape(t *testing.T) {
|
||||
etag := etagSettings().ETag()
|
||||
|
||||
assert.Len(t, etag, etagLength, "the validator must be exactly etagLength characters")
|
||||
assert.NotContains(t, etag, `"`, "the derived validator must not carry its own quoting")
|
||||
for _, r := range etag {
|
||||
require.Truef(t, (r >= '0' && r <= '9') || (r >= 'a' && r <= 'f'),
|
||||
"the validator must be lowercase hex, got %q in %q", r, etag)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSettings_ETagIsStable covers the guarantee every conditional request
|
||||
// rests on: an unchanged row derives the same validator every time, including
|
||||
// across a fresh struct built from the same values. A validator that varied
|
||||
// per derivation would fail every If-Match and make the feature unusable.
|
||||
func TestSettings_ETagIsStable(t *testing.T) {
|
||||
s := etagSettings()
|
||||
|
||||
first := s.ETag()
|
||||
assert.Equal(t, first, s.ETag(), "repeated derivation from one value must agree")
|
||||
assert.Equal(t, first, etagSettings().ETag(), "an equal row must derive an equal validator")
|
||||
}
|
||||
|
||||
// TestSettings_ETagSensitivity is the other half of the contract: every field
|
||||
// the validator covers must actually move it. The cases are also what makes
|
||||
// the field-count guard meaningful — a new field that belongs in the tuple but
|
||||
// is missing from it has no case here, and the guard is what catches that.
|
||||
//
|
||||
// The mutations are checked to be pairwise distinct, not merely different from
|
||||
// the baseline: that is what catches an ambiguous concatenation, where moving
|
||||
// a character across a field boundary would hash identically without the
|
||||
// delimiter.
|
||||
func TestSettings_ETagSensitivity(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
mutate func(*Settings)
|
||||
}{
|
||||
{"domain", func(s *Settings) { s.Domain = "brave-otter.eu.proxy.netbird.io" }},
|
||||
{"proxy address", func(s *Settings) { s.ProxyAddress = "us.proxy.netbird.io" }},
|
||||
{"log collection", func(s *Settings) { s.EnableLogCollection = false }},
|
||||
{"prompt collection", func(s *Settings) { s.EnablePromptCollection = false }},
|
||||
{"redact pii", func(s *Settings) { s.RedactPii = false }},
|
||||
{"retention", func(s *Settings) { s.AccessLogRetentionDays = 14 }},
|
||||
{"created at", func(s *Settings) { s.CreatedAt = s.CreatedAt.Add(time.Second) }},
|
||||
// Moving characters across the Domain/ProxyAddress boundary leaves
|
||||
// the two fields' concatenation byte-identical, so this case passes
|
||||
// only because the tuple is delimited.
|
||||
{"identity boundary shifted", func(s *Settings) {
|
||||
joined := s.Domain + s.ProxyAddress
|
||||
split := len(s.Domain) - 3
|
||||
s.Domain, s.ProxyAddress = joined[:split], joined[split:]
|
||||
}},
|
||||
}
|
||||
|
||||
baseline := etagSettings().ETag()
|
||||
seen := map[string]string{"baseline": baseline}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := etagSettings()
|
||||
tc.mutate(s)
|
||||
|
||||
etag := s.ETag()
|
||||
assert.NotEqual(t, baseline, etag, "changing %s must change the validator", tc.name)
|
||||
|
||||
if other, clash := seen[etag]; clash {
|
||||
t.Fatalf("changing %s derives the same validator as %s (%s) — the field tuple is ambiguous", tc.name, other, etag)
|
||||
}
|
||||
seen[etag] = tc.name
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestSettings_ETagExclusions pins the two deliberate omissions. AccountID is
|
||||
// the resource's identity rather than its representation. UpdatedAt is left
|
||||
// out so that a write which changes nothing observable does not invalidate a
|
||||
// precondition another client is holding — equal representations must always
|
||||
// derive equal validators.
|
||||
func TestSettings_ETagExclusions(t *testing.T) {
|
||||
baseline := etagSettings().ETag()
|
||||
|
||||
other := etagSettings()
|
||||
other.AccountID = "acc-2"
|
||||
assert.Equal(t, baseline, other.ETag(), "the account id must not reach the validator")
|
||||
|
||||
touched := etagSettings()
|
||||
touched.UpdatedAt = touched.UpdatedAt.Add(time.Hour)
|
||||
assert.Equal(t, baseline, touched.ETag(), "a write that changed nothing must not move the validator")
|
||||
}
|
||||
|
||||
// TestSettings_ETagSurvivesTimestampTruncation pins the store round-trip the
|
||||
// validator has to survive. A freshly bootstrapped row derives its validator
|
||||
// in memory, from a time.Time carrying nanoseconds; every later comparison
|
||||
// derives it from a row read back out of the store, and the engines truncate
|
||||
// on the way through — PostgreSQL to microseconds, MySQL DATETIME to whole
|
||||
// seconds without an fsp. If the hash is sensitive below its coarsest engine's
|
||||
// precision, the validator a bootstrap hands out never matches again and the
|
||||
// documented "conditional PUT without an intervening GET" is a permanent 412.
|
||||
//
|
||||
// Asserted on the type rather than through a store, so it holds without running
|
||||
// the suite against every engine. The sqlite test store preserves nanoseconds,
|
||||
// so a sqlite-only suite cannot observe the truncation at all.
|
||||
func TestSettings_ETagSurvivesTimestampTruncation(t *testing.T) {
|
||||
inMemory := etagSettings()
|
||||
require.NotZero(t, inMemory.CreatedAt.Nanosecond(), "the fixture must carry sub-second precision to prove anything")
|
||||
|
||||
for name, truncation := range map[string]time.Duration{
|
||||
"postgres (microseconds)": time.Microsecond,
|
||||
"mysql (milliseconds)": time.Millisecond,
|
||||
"mysql datetime (seconds)": time.Second,
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
roundTripped := etagSettings()
|
||||
roundTripped.CreatedAt = roundTripped.CreatedAt.Truncate(truncation)
|
||||
|
||||
assert.Equal(t, inMemory.ETag(), roundTripped.ETag(),
|
||||
"a validator derived before the write must still match one derived after reading the row back")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestSettings_ETagOfDefaults covers the pre-bootstrap view, which GET serves
|
||||
// as a real representation and therefore validates like one. It must derive
|
||||
// without panicking on the zero CreatedAt, and it must not collide with a
|
||||
// bootstrapped row — otherwise an If-Match taken before bootstrap would
|
||||
// authorize a write against the row that appeared since.
|
||||
func TestSettings_ETagOfDefaults(t *testing.T) {
|
||||
defaults := DefaultSettings("acc-1").ETag()
|
||||
|
||||
assert.Len(t, defaults, etagLength, "the default view must derive a well-formed validator")
|
||||
assert.NotEqual(t, etagSettings().ETag(), defaults,
|
||||
"the unbootstrapped view must not validate as a bootstrapped row")
|
||||
}
|
||||
|
||||
// etagFieldCount is the number of fields Settings carries. ETag hashes an
|
||||
// explicit tuple rather than the struct, so a field added here is silently
|
||||
// outside the validator until someone decides otherwise — the worst kind of
|
||||
// gap, because the mechanism looks present and works for every other field.
|
||||
//
|
||||
// If this constant needs updating, that is the decision point: either add the
|
||||
// new field to ETag and give it a case in TestSettings_ETagSensitivity, or
|
||||
// record here why it stays out.
|
||||
const etagFieldCount = 9
|
||||
|
||||
// TestSettings_ETagFieldCountGuard fails when a field is added to or removed
|
||||
// from Settings, forcing the question of whether it belongs in the validator.
|
||||
func TestSettings_ETagFieldCountGuard(t *testing.T) {
|
||||
assert.Equal(t, etagFieldCount, reflect.TypeFor[Settings]().NumField(),
|
||||
"Settings gained or lost a field: decide whether it belongs in ETag(), then update etagFieldCount")
|
||||
}
|
||||
@@ -117,7 +117,7 @@ func TestAgentNetwork_UpdateSettings_PreservesImmutableAndTogglesCollection(t *t
|
||||
EnablePromptCollection: true,
|
||||
RedactPii: true,
|
||||
AccessLogRetentionDays: before.AccessLogRetentionDays,
|
||||
}, nil)
|
||||
})
|
||||
require.NoError(t, err, "UpdateSettings must succeed")
|
||||
assert.Equal(t, before.Domain, updated.Domain, "domain is immutable and must be preserved")
|
||||
assert.Equal(t, before.ProxyAddress, updated.ProxyAddress, "proxy address is immutable and must be preserved")
|
||||
@@ -147,7 +147,7 @@ func TestAgentNetwork_UpdateSettings_PreservesImmutableAndTogglesCollection(t *t
|
||||
EnablePromptCollection: false,
|
||||
RedactPii: false,
|
||||
AccessLogRetentionDays: before.AccessLogRetentionDays,
|
||||
}, nil)
|
||||
})
|
||||
assert.Error(t, err, "a mismatched identity echo must be rejected")
|
||||
assert.ErrorContains(t, err, "immutable", "the rejection must name the immutability rule")
|
||||
})
|
||||
|
||||
@@ -98,7 +98,7 @@ func NewAPIHandler(ctx context.Context, router *mux.Router, accountManager accou
|
||||
isValidChildAccount,
|
||||
)
|
||||
|
||||
corsMiddleware := newCORSMiddleware()
|
||||
corsMiddleware := cors.AllowAll()
|
||||
|
||||
metricsMiddleware := appMetrics.HTTPMiddleware()
|
||||
|
||||
@@ -145,32 +145,3 @@ func NewAPIHandler(ctx context.Context, router *mux.Router, accountManager accou
|
||||
|
||||
return router, nil
|
||||
}
|
||||
|
||||
// newCORSMiddleware builds the API's CORS policy: cors.AllowAll() plus ETag in
|
||||
// ExposedHeaders.
|
||||
//
|
||||
// The addition is what makes conditional requests usable from a browser. A
|
||||
// response header that is not CORS-safelisted is invisible to JavaScript
|
||||
// unless it is named in Access-Control-Expose-Headers, and ETag is not on that
|
||||
// list — so without this the server can hand a browser client a validator it
|
||||
// has no way to read, leaving conditional requests to non-browser clients
|
||||
// only. If-Match needs nothing further, since AllowedHeaders is already "*".
|
||||
//
|
||||
// Everything else mirrors cors.AllowAll() exactly. It is spelled out rather
|
||||
// than called because the library offers no way to extend it.
|
||||
func newCORSMiddleware() *cors.Cors {
|
||||
return cors.New(cors.Options{
|
||||
AllowedOrigins: []string{"*"},
|
||||
AllowedMethods: []string{
|
||||
http.MethodHead,
|
||||
http.MethodGet,
|
||||
http.MethodPost,
|
||||
http.MethodPut,
|
||||
http.MethodPatch,
|
||||
http.MethodDelete,
|
||||
},
|
||||
AllowedHeaders: []string{"*"},
|
||||
ExposedHeaders: []string{"ETag"},
|
||||
AllowCredentials: false,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1,88 +0,0 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestCORSExposesETag pins the reason this policy is spelled out instead of
|
||||
// being cors.AllowAll(). ETag is not a CORS-safelisted response header, so
|
||||
// without it named in Access-Control-Expose-Headers a browser client is handed
|
||||
// a validator it cannot read — conditional requests would work for the CLI,
|
||||
// the REST client and Terraform, and silently not for the dashboard.
|
||||
//
|
||||
// Collapsing this back to cors.AllowAll() is exactly the simplification that
|
||||
// would reintroduce that, which is what this test is here to catch.
|
||||
func TestCORSExposesETag(t *testing.T) {
|
||||
handler := newCORSMiddleware().Handler(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("ETag", `"9f86d081884c7d65"`)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/agent-network/settings", nil)
|
||||
req.Header.Set("Origin", "https://app.netbird.io")
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
handler.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
// Compared canonicalized: the library normalizes the name it echoes, so
|
||||
// this reads "Etag" rather than "ETag". Browsers match the exposed-header
|
||||
// list case-insensitively, so the spelling does not matter — but asserting
|
||||
// it byte-exactly would fail for a reason that has nothing to do with the
|
||||
// behaviour being pinned.
|
||||
assert.Equal(t, http.CanonicalHeaderKey("ETag"),
|
||||
http.CanonicalHeaderKey(rec.Header().Get("Access-Control-Expose-Headers")),
|
||||
"browser clients must be allowed to read the validator they are sent")
|
||||
}
|
||||
|
||||
// TestCORSAllowsIfMatchPreflight covers the request half. It needs nothing
|
||||
// beyond the wildcard AllowedHeaders that was already there, so this is a
|
||||
// regression guard rather than a new grant: narrowing AllowedHeaders to a list
|
||||
// later must not drop If-Match and leave writes readable but not conditional.
|
||||
func TestCORSAllowsIfMatchPreflight(t *testing.T) {
|
||||
handler := newCORSMiddleware().Handler(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
req := httptest.NewRequest(http.MethodOptions, "/api/agent-network/settings", nil)
|
||||
req.Header.Set("Origin", "https://app.netbird.io")
|
||||
req.Header.Set("Access-Control-Request-Method", http.MethodPut)
|
||||
req.Header.Set("Access-Control-Request-Headers", "If-Match")
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
handler.ServeHTTP(rec, req)
|
||||
|
||||
assert.Contains(t, rec.Header().Get("Access-Control-Allow-Headers"), "If-Match",
|
||||
"a conditional write must survive preflight")
|
||||
assert.Contains(t, rec.Header().Get("Access-Control-Allow-Methods"), http.MethodPut,
|
||||
"the conditional write's method must survive preflight")
|
||||
}
|
||||
|
||||
// TestCORSMatchesAllowAllOtherwise pins the rest of the policy, which is a
|
||||
// verbatim copy of cors.AllowAll(). Spelling the options out is what let ETag
|
||||
// be added; it also means a change to the library's defaults no longer reaches
|
||||
// this API, so the settings that matter are asserted here rather than assumed.
|
||||
func TestCORSMatchesAllowAllOtherwise(t *testing.T) {
|
||||
handler := newCORSMiddleware().Handler(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
req := httptest.NewRequest(http.MethodOptions, "/api/peers", nil)
|
||||
req.Header.Set("Origin", "https://anywhere.example.com")
|
||||
req.Header.Set("Access-Control-Request-Method", http.MethodDelete)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
handler.ServeHTTP(rec, req)
|
||||
|
||||
assert.Equal(t, "*", rec.Header().Get("Access-Control-Allow-Origin"), "any origin must still be allowed")
|
||||
assert.Empty(t, rec.Header().Get("Access-Control-Allow-Credentials"),
|
||||
"credentials must stay disallowed — allowing them alongside a wildcard origin would be a real weakening")
|
||||
assert.Contains(t, rec.Header().Get("Access-Control-Allow-Methods"), http.MethodDelete,
|
||||
"the full method set must still be allowed")
|
||||
}
|
||||
@@ -13,6 +13,14 @@ func NormalizeBedrockModel(modelID string) string {
|
||||
return sharedllm.NormalizeBedrockModel(modelID)
|
||||
}
|
||||
|
||||
// NormalizeAnthropicModel strips the trailing "-YYYYMMDD" release-date suffix
|
||||
// from an Anthropic model id so a dated id a client pins matches the undated
|
||||
// one the operator registered. Thin delegate to shared/llm for the same
|
||||
// contract reason as the two below.
|
||||
func NormalizeAnthropicModel(modelID string) string {
|
||||
return sharedllm.NormalizeAnthropicModel(modelID)
|
||||
}
|
||||
|
||||
// NormalizeVertexModel strips the "@version" suffix from a Vertex AI model id
|
||||
// so it matches the catalog/pricing key. Thin delegate to shared/llm, kept
|
||||
// beside NormalizeBedrockModel for the same contract reason.
|
||||
|
||||
@@ -10,6 +10,8 @@ package pricing
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
sharedllm "github.com/netbirdio/netbird/shared/llm"
|
||||
)
|
||||
|
||||
// Entry is a single model's input and output pricing, expressed in USD per
|
||||
@@ -92,7 +94,10 @@ func NewTable(raw map[string]map[string]EntryJSON) (*Table, error) {
|
||||
return &Table{entries: entries}, nil
|
||||
}
|
||||
|
||||
// Lookup returns the entry for the given provider surface and model.
|
||||
// Lookup returns the entry for the given provider surface and model. A
|
||||
// dated Anthropic id falls back to its undated form, so a client pinning
|
||||
// "claude-sonnet-4-5-20250929" bills at the registered "claude-sonnet-4-5"
|
||||
// rate instead of recording no cost at all.
|
||||
func (t *Table) Lookup(provider, model string) (Entry, bool) {
|
||||
if t == nil {
|
||||
return Entry{}, false
|
||||
@@ -101,7 +106,14 @@ func (t *Table) Lookup(provider, model string) (Entry, bool) {
|
||||
if !ok {
|
||||
return Entry{}, false
|
||||
}
|
||||
e, ok := byModel[model]
|
||||
if e, found := byModel[model]; found {
|
||||
return e, true
|
||||
}
|
||||
undated := sharedllm.NormalizeAnthropicModel(model)
|
||||
if undated == model {
|
||||
return Entry{}, false
|
||||
}
|
||||
e, ok := byModel[undated]
|
||||
return e, ok
|
||||
}
|
||||
|
||||
|
||||
@@ -175,3 +175,22 @@ func TestNewTable_NilAndEmpty(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, entries, "nil in, empty (never-matching) map out for the per-record map")
|
||||
}
|
||||
|
||||
// TestLookup_DatedAnthropicIDFallsBackToUndated covers a client pinning a
|
||||
// release date on a model priced under its undated id. Without the
|
||||
// fallback the request records no cost at all.
|
||||
func TestLookup_DatedAnthropicIDFallsBackToUndated(t *testing.T) {
|
||||
table, err := NewTable(map[string]map[string]EntryJSON{
|
||||
"anthropic": {
|
||||
"claude-sonnet-4-5": {InputPer1K: 0.003, OutputPer1K: 0.015},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "table must build from a valid defaults map")
|
||||
|
||||
entry, ok := table.Lookup("anthropic", "claude-sonnet-4-5-20250929")
|
||||
require.True(t, ok, "a dated id must resolve to the undated entry")
|
||||
assert.InDelta(t, 0.003, entry.InputPer1K, 1e-9, "dated id must bill at the registered rate")
|
||||
|
||||
_, ok = table.Lookup("anthropic", "claude-sonnet-9-9-20250929")
|
||||
assert.False(t, ok, "an unknown family must stay unpriced")
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
|
||||
"github.com/netbirdio/netbird/proxy/internal/llm"
|
||||
"github.com/netbirdio/netbird/proxy/internal/llm/pricing"
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware"
|
||||
)
|
||||
@@ -175,13 +176,28 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
|
||||
// Anthropic route still bills its cache buckets additively.
|
||||
func (m *Middleware) lookupCosts(md []middleware.KV, surface, model string, inTokens, outTokens, cachedTokens, cacheCreationTokens int64) (pricing.Costs, bool) {
|
||||
if recordID := lookupKV(md, middleware.KeyLLMResolvedProviderID); recordID != "" {
|
||||
if entry, ok := m.perRecord[recordID][model]; ok {
|
||||
if entry, ok := perRecordEntry(m.perRecord[recordID], model); ok {
|
||||
return pricing.EntryCosts(entry, surface, inTokens, outTokens, cachedTokens, cacheCreationTokens), true
|
||||
}
|
||||
}
|
||||
return m.defaults.Costs(surface, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
|
||||
}
|
||||
|
||||
// perRecordEntry resolves the operator's stored price for a model on one
|
||||
// provider record, falling back to the undated form of a dated Anthropic id
|
||||
// so a client that pins a release date still bills at the registered rate.
|
||||
func perRecordEntry(byModel map[string]pricing.Entry, model string) (pricing.Entry, bool) {
|
||||
if entry, ok := byModel[model]; ok {
|
||||
return entry, true
|
||||
}
|
||||
undated := llm.NormalizeAnthropicModel(model)
|
||||
if undated == model {
|
||||
return pricing.Entry{}, false
|
||||
}
|
||||
entry, ok := byModel[undated]
|
||||
return entry, ok
|
||||
}
|
||||
|
||||
// usd renders a cost as the fixed-precision string every cost.usd_* key
|
||||
// carries, so the per-bucket values and the aggregates round identically.
|
||||
//
|
||||
|
||||
@@ -84,8 +84,10 @@ func (m *Middleware) MutationsSupported() bool { return false }
|
||||
func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
|
||||
model, modelPresent := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
|
||||
providerID, _ := lookupMetadata(in.Metadata, middleware.KeyLLMResolvedProviderID)
|
||||
surface, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider)
|
||||
nonInference, _ := lookupMetadata(in.Metadata, middleware.KeyLLMNonInference)
|
||||
|
||||
if denial := m.evaluateAllowlist(providerID, model, modelPresent); denial != nil {
|
||||
if denial := m.evaluateAllowlist(providerID, surface, model, modelPresent, nonInference == "true"); denial != nil {
|
||||
return denial, nil
|
||||
}
|
||||
|
||||
@@ -114,7 +116,7 @@ func (m *Middleware) Close() error { return nil }
|
||||
// evaluateAllowlist denies when the resolved provider's allowlist rejects the
|
||||
// model; nil means proceed. Scoped to the provider llm_router resolved, so an
|
||||
// unrestricted provider (absent from config) is never caught by another's list.
|
||||
func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bool) *middleware.Output {
|
||||
func (m *Middleware) evaluateAllowlist(providerID, surface, model string, modelPresent, nonInference bool) *middleware.Output {
|
||||
if len(m.cfg.ProviderAllowlists) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -122,7 +124,7 @@ func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bo
|
||||
// if this request targets a restricted provider — fail closed. llm_router
|
||||
// normally stamps the provider first, so this is a defensive guard.
|
||||
if providerID == "" {
|
||||
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
|
||||
return denyModel(surface, "", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
|
||||
}
|
||||
allowlist, restricted := m.cfg.ProviderAllowlists[providerID]
|
||||
if !restricted {
|
||||
@@ -133,18 +135,29 @@ func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bo
|
||||
// Fail closed: with an allowlist in effect for this provider, a request whose
|
||||
// model the parser couldn't extract (absent/empty) is denied. This enforces
|
||||
// the allowlist for path-routed providers (Bedrock, Vertex) with no body model.
|
||||
//
|
||||
// The exception is a non-inference endpoint the router already authorised.
|
||||
// The model listing and the connection-warming probe name no model
|
||||
// anywhere — not in a body, not in the path — so failing closed here
|
||||
// rejected model discovery for exactly the accounts that configured an
|
||||
// allowlist, which is the outage this endpoint is meant to avoid. The
|
||||
// per-model lookup does name one (the router stamps it from the path), so
|
||||
// it still falls through to the allowlist check below.
|
||||
if !modelPresent || normaliseModel(model) == "" {
|
||||
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
|
||||
if nonInference {
|
||||
return nil
|
||||
}
|
||||
return denyModel(surface, "", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
|
||||
}
|
||||
if modelInAllowlist(allowlist, model) {
|
||||
return nil
|
||||
}
|
||||
return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
|
||||
return denyModel(surface, 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 {
|
||||
func denyModel(surface, model, code, message, reason string) *middleware.Output {
|
||||
details := map[string]string{}
|
||||
if model != "" {
|
||||
details["model"] = model
|
||||
@@ -156,6 +169,7 @@ func denyModel(model, code, message, reason string) *middleware.Output {
|
||||
Code: code,
|
||||
Message: message,
|
||||
Details: details,
|
||||
Surface: surface,
|
||||
},
|
||||
Metadata: []middleware.KV{
|
||||
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
|
||||
|
||||
@@ -343,3 +343,52 @@ func TestFactoryNormalisesAllowlist(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out2.Decision, "trimmed entry must still match")
|
||||
}
|
||||
|
||||
// TestAllowlistSkipsNonInferenceWithoutModel covers the reported regression:
|
||||
// GET /v1/models carries no model anywhere, so the fail-closed rule above
|
||||
// denied model discovery for exactly the accounts that configured a provider
|
||||
// allowlist — the clients that read a 403 here render an empty model picker.
|
||||
// The router authorises those endpoints by path before the guardrail sees
|
||||
// them, so an absent model there is expected rather than undeterminable.
|
||||
func TestAllowlistSkipsNonInferenceWithoutModel(t *testing.T) {
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision,
|
||||
"model discovery must not be refused because it names no model")
|
||||
}
|
||||
|
||||
// TestAllowlistStillAppliesToNonInferenceWithModel pins that the exemption is
|
||||
// scoped to requests that genuinely name nothing. The per-model lookup
|
||||
// (GET /v1/models/{id}) is non-inference too, but the router stamps the model
|
||||
// from its path, so the allowlist must still decide it — otherwise the
|
||||
// exemption becomes a way to confirm a model the policy blocks.
|
||||
func TestAllowlistStillAppliesToNonInferenceWithModel(t *testing.T) {
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
|
||||
t.Run("model in the allowlist", func(t *testing.T) {
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"},
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision,
|
||||
"an allowlisted model must stay reachable")
|
||||
})
|
||||
|
||||
t.Run("model outside the allowlist", func(t *testing.T) {
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"},
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-5"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision,
|
||||
"non-inference must not become a way past the allowlist")
|
||||
require.NotNil(t, out.DenyReason)
|
||||
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code,
|
||||
"a named but blocked model is blocked, not unknown")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -217,6 +217,32 @@ func applyHeaderPair(rule *HeaderPairRule, in *middleware.Input) *middleware.Mut
|
||||
return mutations
|
||||
}
|
||||
|
||||
// bodyInjectableSurfaces are the request-body dialects that accept the
|
||||
// OpenAI-standard identity fields this middleware writes. A surface
|
||||
// outside this set gets header-only stamping: "user" and "metadata.tags"
|
||||
// are not part of the Anthropic Messages schema, which rejects unknown
|
||||
// top-level fields and permits only "user_id" under metadata, so writing
|
||||
// them into an Anthropic-shaped body turns a working request into a 400.
|
||||
// Claude Code speaks that shape through gateway records pinned to the
|
||||
// OpenAI parser, so the check keys on the detected surface rather than
|
||||
// on the provider record.
|
||||
var bodyInjectableSurfaces = map[string]struct{}{
|
||||
"openai": {},
|
||||
// An empty surface means no parser claimed the path (a custom gateway
|
||||
// base). Those upstreams are OpenAI-compatible by convention, so keep
|
||||
// the long-standing behaviour rather than silently dropping identity.
|
||||
"": {},
|
||||
}
|
||||
|
||||
// bodyAcceptsOpenAIIdentity reports whether the request body may carry the
|
||||
// OpenAI-standard identity fields, read from the surface llm_request_parser
|
||||
// resolved from the request path.
|
||||
func bodyAcceptsOpenAIIdentity(in *middleware.Input) bool {
|
||||
surface, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider)
|
||||
_, ok := bodyInjectableSurfaces[surface]
|
||||
return ok
|
||||
}
|
||||
|
||||
// injectIntoBody parses the request body and writes the supplied
|
||||
// identity dimensions into it. Tags land at metadata.tags (creating
|
||||
// the metadata object when absent); the user identity lands at the
|
||||
@@ -225,6 +251,8 @@ func applyHeaderPair(rule *HeaderPairRule, in *middleware.Input) *middleware.Mut
|
||||
// was written. Returns ok=false (no mutation) when:
|
||||
//
|
||||
// - both inputs are empty (nothing to write);
|
||||
// - the body speaks a dialect without these fields (see
|
||||
// bodyInjectableSurfaces);
|
||||
// - the body is empty or truncated (we don't have the full document
|
||||
// to safely round-trip);
|
||||
// - the body isn't a JSON object (skip silently — this middleware
|
||||
@@ -245,6 +273,9 @@ func injectIntoBody(in *middleware.Input, tags []string, userID string) ([]byte,
|
||||
if in == nil || len(in.Body) == 0 || in.BodyTruncated {
|
||||
return nil, false
|
||||
}
|
||||
if !bodyAcceptsOpenAIIdentity(in) {
|
||||
return nil, false
|
||||
}
|
||||
var doc map[string]any
|
||||
if err := json.Unmarshal(in.Body, &doc); err != nil {
|
||||
return nil, false
|
||||
|
||||
@@ -704,3 +704,57 @@ func TestInject_ExtraHeaders_EmptyValueSkipped(t *testing.T) {
|
||||
"empty extra value must not be stamped")
|
||||
}
|
||||
}
|
||||
|
||||
// TestInject_AnthropicBodyIsNotRewritten pins the shape gate. Claude Code
|
||||
// reaches a LiteLLM record on /v1/messages, where "user" is not a
|
||||
// permitted top-level field and metadata accepts only "user_id", so
|
||||
// writing the OpenAI-standard fields would turn a working request into a
|
||||
// 400 naming a field the client never sent. Header stamping still runs, so
|
||||
// spend tracking and per-end-user budgets keep working.
|
||||
func TestInject_AnthropicBodyIsNotRewritten(t *testing.T) {
|
||||
rule := liteLLMRuleWithBody()
|
||||
rule.HeaderPair.EndUserIDInBody = true
|
||||
mw := New(Config{Providers: []ProviderInjection{rule}})
|
||||
|
||||
in := newInput(litellmProvider, "alice", []string{"grp-eng"})
|
||||
in.UserEmail = "alice@example.com"
|
||||
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
|
||||
in.Body = []byte(`{"model":"claude-sonnet-5","messages":[]}`)
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out.Mutations)
|
||||
assert.Empty(t, out.Mutations.BodyReplace,
|
||||
"an Anthropic-shaped body must reach the upstream unmodified")
|
||||
|
||||
var endUser string
|
||||
for _, kv := range out.Mutations.HeadersAdd {
|
||||
if kv.Key == "x-litellm-end-user-id" {
|
||||
endUser = kv.Value
|
||||
}
|
||||
}
|
||||
assert.Equal(t, "alice@example.com", endUser,
|
||||
"header stamping must still carry identity when body inject is skipped")
|
||||
}
|
||||
|
||||
// TestInject_OpenAIBodyStillRewritten guards the gate against
|
||||
// over-reaching: the OpenAI surface must keep its body-level identity,
|
||||
// which is the only path LiteLLM's tag-budget check reads.
|
||||
func TestInject_OpenAIBodyStillRewritten(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}})
|
||||
|
||||
in := newInput(litellmProvider, "alice", []string{"grp-eng"})
|
||||
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "openai"})
|
||||
in.Body = []byte(`{"model":"gpt-4o-mini","messages":[]}`)
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out.Mutations)
|
||||
require.NotEmpty(t, out.Mutations.BodyReplace, "the OpenAI surface still gets body tags")
|
||||
|
||||
var doc map[string]any
|
||||
require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc))
|
||||
meta, ok := doc["metadata"].(map[string]any)
|
||||
require.True(t, ok, "metadata must be an object")
|
||||
assert.NotEmpty(t, meta["tags"], "metadata.tags must still be written")
|
||||
}
|
||||
|
||||
@@ -84,6 +84,15 @@ func (m *Middleware) Invoke(ctx context.Context, in *middleware.Input) (*middlew
|
||||
return allowNoAttribution(), nil
|
||||
}
|
||||
|
||||
// Model-listing and other non-inference endpoints carry no model, and
|
||||
// management's per-model allowlist fails closed on an empty one. The
|
||||
// router has already authorised the route against the caller's groups
|
||||
// and the request consumes no tokens, so gating it on a model that
|
||||
// cannot exist would only break gateway model discovery.
|
||||
if lookupKV(in.Metadata, middleware.KeyLLMNonInference) == "true" {
|
||||
return allowNoAttribution(), nil
|
||||
}
|
||||
|
||||
providerID := lookupKV(in.Metadata, middleware.KeyLLMResolvedProviderID)
|
||||
if providerID == "" {
|
||||
// llm_router didn't emit a resolved provider id — usually
|
||||
@@ -117,7 +126,7 @@ func (m *Middleware) Invoke(ctx context.Context, in *middleware.Input) (*middlew
|
||||
}
|
||||
|
||||
if resp.GetDecision() == "deny" {
|
||||
return denyFromManagement(resp), nil
|
||||
return denyFromManagement(resp, lookupKV(in.Metadata, middleware.KeyLLMProvider)), nil
|
||||
}
|
||||
return allowFromManagement(resp), nil
|
||||
}
|
||||
@@ -161,7 +170,7 @@ func allowFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.O
|
||||
// envelope. The deny code surfaces verbatim through the framework's
|
||||
// fixed JSON template; arbitrary middleware bytes can't reach the
|
||||
// wire.
|
||||
func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Output {
|
||||
func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse, surface string) *middleware.Output {
|
||||
code := resp.GetDenyCode()
|
||||
if code == "" {
|
||||
code = "llm_policy.cap_exceeded"
|
||||
@@ -176,6 +185,7 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou
|
||||
DenyReason: &middleware.DenyReason{
|
||||
Code: code,
|
||||
Message: denyMessageForCode(code),
|
||||
Surface: surface,
|
||||
},
|
||||
Metadata: []middleware.KV{
|
||||
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
|
||||
|
||||
@@ -224,3 +224,35 @@ func TestMetadataKeys_Allowlist(t *testing.T) {
|
||||
}
|
||||
assert.ElementsMatch(t, want, keys)
|
||||
}
|
||||
|
||||
// TestInvoke_NonInferenceSkipsPreflight covers gateway model discovery:
|
||||
// GET /v1/models carries no model, and management's per-model allowlist
|
||||
// fails closed on an empty one, so a pre-flight would deny discovery for
|
||||
// exactly the accounts that use the model allowlist. The router marks the
|
||||
// request non-inference after authorising the route, and the gate must
|
||||
// then allow without calling management at all.
|
||||
func TestInvoke_NonInferenceSkipsPreflight(t *testing.T) {
|
||||
mgmt := &fakeMgmt{
|
||||
checkResp: &proto.CheckLLMPolicyLimitsResponse{
|
||||
Decision: "deny",
|
||||
DenyCode: "llm_policy.model_blocked",
|
||||
},
|
||||
}
|
||||
m := New(mgmt, nil)
|
||||
|
||||
out := runInvoke(t, m, &middleware.Input{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-bob",
|
||||
UserGroups: []string{"grp-engineers"},
|
||||
Metadata: []middleware.KV{
|
||||
{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"},
|
||||
{Key: middleware.KeyLLMNonInference, Value: "true"},
|
||||
},
|
||||
})
|
||||
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "model-less endpoints must not be gated on a model")
|
||||
assert.Nil(t, mgmt.checkReq, "no pre-flight may be sent for a non-inference request")
|
||||
|
||||
assert.Empty(t, lookupKV(out.Metadata, middleware.KeyLLMSelectedPolicyID),
|
||||
"no policy is attributed when nothing was metered")
|
||||
}
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
package llm_request_parser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware"
|
||||
)
|
||||
|
||||
func TestParseBedrockPath(t *testing.T) {
|
||||
@@ -36,3 +40,25 @@ func TestParseBedrockPath(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestInvoke_BedrockCountTokens covers the dedicated token-counting
|
||||
// endpoint. Denying it does not break the client, it just pushes context
|
||||
// counting back onto the inference endpoint, which is billable.
|
||||
func TestInvoke_BedrockCountTokens(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/count-tokens",
|
||||
Body: []byte(`{"input":{"converse":{"messages":[]}}}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision)
|
||||
|
||||
model, ok := metaValue(t, out.Metadata, middleware.KeyLLMModel)
|
||||
require.True(t, ok, "count-tokens carries a model in the path and must emit it")
|
||||
assert.Equal(t, "anthropic.claude-sonnet-4-5", model, "model must be normalized like any other action")
|
||||
|
||||
stream, _ := metaValue(t, out.Metadata, middleware.KeyLLMStream)
|
||||
assert.Equal(t, "false", stream, "count-tokens never streams")
|
||||
}
|
||||
|
||||
@@ -61,6 +61,8 @@ func (middlewareImpl) MetadataKeys() []string {
|
||||
middleware.KeyLLMRequestPromptRaw,
|
||||
middleware.KeyLLMCaptureTruncated,
|
||||
middleware.KeyLLMSessionID,
|
||||
middleware.KeyLLMAgentID,
|
||||
middleware.KeyLLMParentAgentID,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -72,9 +74,9 @@ func (middlewareImpl) Close() error { return nil }
|
||||
|
||||
// Invoke detects the LLM provider, parses request facts, and emits
|
||||
// metadata. Always returns DecisionAllow; never errors. Provider
|
||||
// selection prefers the configured providerID (synthesiser-stamped on
|
||||
// agent-network targets) so requests routed to a custom upstream URL
|
||||
// still resolve. Falls back to URL sniffing when no providerID is set.
|
||||
// selection prefers the request path, falling back to the configured
|
||||
// providerID (synthesiser-stamped on agent-network targets) so requests
|
||||
// routed to a custom upstream URL still resolve.
|
||||
func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
|
||||
out := &middleware.Output{Decision: middleware.DecisionAllow}
|
||||
if in == nil {
|
||||
@@ -92,9 +94,14 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle
|
||||
return m.invokeBedrock(in, br), nil
|
||||
}
|
||||
|
||||
parser, ok := llm.ParserByName(m.providerID)
|
||||
// A path that names an API surface wins over the configured providerID:
|
||||
// a gateway record pinned to "openai" still serves Claude Code on
|
||||
// /v1/messages, and reading that body with the OpenAI parser loses the
|
||||
// Anthropic usage block and prices the request on the wrong surface.
|
||||
// providerID stays the fallback for upstreams whose path says nothing.
|
||||
parser, ok := llm.DetectParser(extractPath(in.URL))
|
||||
if !ok {
|
||||
parser, ok = llm.DetectParser(extractPath(in.URL))
|
||||
parser, ok = llm.ParserByName(m.providerID)
|
||||
}
|
||||
if !ok {
|
||||
return out, nil
|
||||
@@ -116,9 +123,9 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle
|
||||
}
|
||||
appendSessionID := func(md []middleware.KV) []middleware.KV {
|
||||
if sessionID != "" {
|
||||
return append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
|
||||
}
|
||||
return md
|
||||
return appendAgentIDs(md, in.Headers)
|
||||
}
|
||||
|
||||
facts, err := parser.ParseRequest(in.Body)
|
||||
@@ -160,6 +167,41 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// agentIDHeader and parentAgentIDHeader carry sub-agent attribution: a
|
||||
// coding agent that spawns helpers stamps the spawned agent's id, plus the
|
||||
// spawning agent's when that helper is itself nested. Both are opaque
|
||||
// identifiers rather than content, so they're emitted regardless of the
|
||||
// prompt-collection toggle, the same way the session id is.
|
||||
const (
|
||||
agentIDHeader = "x-claude-code-agent-id"
|
||||
parentAgentIDHeader = "x-claude-code-parent-agent-id"
|
||||
)
|
||||
|
||||
// appendAgentIDs stamps the sub-agent attribution headers onto the metadata
|
||||
// bag, skipping either one the request doesn't carry.
|
||||
func appendAgentIDs(md []middleware.KV, headers []middleware.KV) []middleware.KV {
|
||||
for _, pair := range []struct{ key, header string }{
|
||||
{middleware.KeyLLMAgentID, agentIDHeader},
|
||||
{middleware.KeyLLMParentAgentID, parentAgentIDHeader},
|
||||
} {
|
||||
if v := headerValue(headers, pair.header); v != "" {
|
||||
md = append(md, middleware.KV{Key: pair.key, Value: v})
|
||||
}
|
||||
}
|
||||
return md
|
||||
}
|
||||
|
||||
// headerValue returns the first non-empty value for the named header.
|
||||
// Headers arrive in canonical form, so the match is case-insensitive.
|
||||
func headerValue(headers []middleware.KV, want string) string {
|
||||
for _, kv := range headers {
|
||||
if strings.EqualFold(kv.Key, want) && kv.Value != "" {
|
||||
return kv.Value
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// sessionIDHeaders are request header names that may carry a client
|
||||
// session identifier, checked in order, case-insensitively. Matching is
|
||||
// against Go's canonical header form, so use the hyphenated names the
|
||||
@@ -173,10 +215,8 @@ var sessionIDHeaders = []string{"x-claude-code-session-id", "session-id", "x-ses
|
||||
// canonical form, so the match is case-insensitive.
|
||||
func sessionIDFromHeaders(headers []middleware.KV) string {
|
||||
for _, want := range sessionIDHeaders {
|
||||
for _, kv := range headers {
|
||||
if strings.EqualFold(kv.Key, want) && kv.Value != "" {
|
||||
return kv.Value
|
||||
}
|
||||
if v := headerValue(headers, want); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
@@ -252,6 +292,12 @@ func parseVertexPath(reqPath string) (vertexRequest, bool) {
|
||||
if c := strings.LastIndex(rest, ":"); c >= 0 {
|
||||
model, action = rest[:c], rest[c+1:]
|
||||
}
|
||||
// Token counting hangs off the model as its own path segment
|
||||
// (".../models/{model}/count-tokens:rawPredict"), so anything past the
|
||||
// first "/" belongs to the method rather than the model id.
|
||||
if slash := strings.Index(model, "/"); slash >= 0 {
|
||||
model = model[:slash]
|
||||
}
|
||||
model = llm.NormalizeVertexModel(model)
|
||||
if model == "" {
|
||||
return vertexRequest{}, false
|
||||
@@ -298,6 +344,7 @@ func (m middlewareImpl) invokeVertex(in *middleware.Input, vx vertexRequest) *mi
|
||||
if sessionID != "" {
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
|
||||
}
|
||||
md = appendAgentIDs(md, in.Headers)
|
||||
|
||||
promptTruncated := false
|
||||
if parser != nil && m.capturePrompt {
|
||||
@@ -345,7 +392,9 @@ func trimBedrockNamespace(reqPath string) string {
|
||||
//
|
||||
// /model/{modelId}/{action}
|
||||
//
|
||||
// action ∈ {invoke, invoke-with-response-stream, converse, converse-stream}.
|
||||
// action ∈ {invoke, invoke-with-response-stream, converse, converse-stream,
|
||||
// count-tokens}. Token counting carries a model and no usage, so it routes
|
||||
// like any other action and meters to zero.
|
||||
// The modelId may be URL-encoded and may carry a cross-region inference-profile
|
||||
// prefix and a version suffix; normalizeBedrockModel strips both so the model
|
||||
// matches catalog pricing.
|
||||
@@ -369,7 +418,7 @@ func parseBedrockPath(reqPath string) (bedrockRequest, bool) {
|
||||
return bedrockRequest{}, false
|
||||
}
|
||||
switch action {
|
||||
case "invoke", "converse":
|
||||
case "invoke", "converse", "count-tokens":
|
||||
return bedrockRequest{model: model}, true
|
||||
case "invoke-with-response-stream", "converse-stream":
|
||||
return bedrockRequest{model: model, stream: true}, true
|
||||
@@ -397,6 +446,7 @@ func (m middlewareImpl) invokeBedrock(in *middleware.Input, br bedrockRequest) *
|
||||
if sessionID != "" {
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
|
||||
}
|
||||
md = appendAgentIDs(md, in.Headers)
|
||||
|
||||
promptTruncated := false
|
||||
if parser != nil && m.capturePrompt {
|
||||
|
||||
@@ -45,6 +45,8 @@ func TestMiddleware_StaticSurface(t *testing.T) {
|
||||
middleware.KeyLLMRequestPromptRaw,
|
||||
middleware.KeyLLMCaptureTruncated,
|
||||
middleware.KeyLLMSessionID,
|
||||
middleware.KeyLLMAgentID,
|
||||
middleware.KeyLLMParentAgentID,
|
||||
}
|
||||
assert.Equal(t, expected, keys, "metadata key allowlist must match the spec")
|
||||
}
|
||||
@@ -230,6 +232,31 @@ func TestInvoke_ProviderIDConfigBypassesURLSniff(t *testing.T) {
|
||||
assert.Equal(t, "gpt-4o-mini", model)
|
||||
}
|
||||
|
||||
func TestInvoke_PathSurfaceBeatsProviderIDConfig(t *testing.T) {
|
||||
// Gateway records (LiteLLM, Portkey, OpenRouter) pin provider_id
|
||||
// "openai", but the same record serves Claude Code on /v1/messages.
|
||||
// Parsing that body as OpenAI reads no usage off the Anthropic
|
||||
// response and prices the request on a surface where no claude-*
|
||||
// model exists, so the path has to win.
|
||||
mw, err := Factory{}.New([]byte(`{"provider_id":"openai"}`))
|
||||
require.NoError(t, err, "factory must accept provider_id config")
|
||||
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/messages",
|
||||
Body: []byte(`{"model":"claude-sonnet-5","stream":true,"messages":[{"role":"user","content":"Hi"}]}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
|
||||
provider, ok := metaValue(t, out.Metadata, middleware.KeyLLMProvider)
|
||||
require.True(t, ok, "provider must be emitted")
|
||||
assert.Equal(t, "anthropic", provider, "the /v1/messages path selects the Anthropic surface")
|
||||
|
||||
model, ok := metaValue(t, out.Metadata, middleware.KeyLLMModel)
|
||||
require.True(t, ok, "model must be extracted")
|
||||
assert.Equal(t, "claude-sonnet-5", model)
|
||||
}
|
||||
|
||||
func TestInvoke_UnknownProviderIDFallsBackToURL(t *testing.T) {
|
||||
mw, err := Factory{}.New([]byte(`{"provider_id":"not-a-real-parser"}`))
|
||||
require.NoError(t, err, "factory must accept any provider_id string")
|
||||
@@ -416,3 +443,81 @@ func TestInvoke_NilInputAllows(t *testing.T) {
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "nil input still allows")
|
||||
assert.Empty(t, out.Metadata, "nil input emits no metadata")
|
||||
}
|
||||
|
||||
// TestParseVertexPath_CountTokensKeepsModel covers Vertex token counting,
|
||||
// where the method hangs off the model as its own path segment. Splitting
|
||||
// only on the final colon swallowed "/count-tokens" into the model id, so
|
||||
// the router saw a model no route could claim.
|
||||
func TestParseVertexPath_CountTokensKeepsModel(t *testing.T) {
|
||||
cases := map[string]struct {
|
||||
model string
|
||||
stream bool
|
||||
}{
|
||||
"/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5:rawPredict": {model: "claude-sonnet-5"},
|
||||
"/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5:streamRawPredict": {model: "claude-sonnet-5", stream: true},
|
||||
"/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5/count-tokens:rawPredict": {model: "claude-sonnet-5"},
|
||||
"/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5@20250929/count-tokens:rawPredict": {model: "claude-sonnet-5"},
|
||||
}
|
||||
for path, want := range cases {
|
||||
vx, ok := parseVertexPath(path)
|
||||
require.True(t, ok, "must parse %q", path)
|
||||
assert.Equal(t, want.model, vx.model, "model for %q", path)
|
||||
assert.Equal(t, want.stream, vx.stream, "stream flag for %q", path)
|
||||
assert.Equal(t, "anthropic", vx.publisher, "publisher for %q", path)
|
||||
}
|
||||
}
|
||||
|
||||
// TestInvoke_EmitsAgentIDs covers sub-agent attribution: several agents run
|
||||
// in parallel inside one session, and without their ids every request in
|
||||
// the session attributes to the session alone.
|
||||
func TestInvoke_EmitsAgentIDs(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
|
||||
t.Run("spawned agent", func(t *testing.T) {
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/messages",
|
||||
Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`),
|
||||
Headers: []middleware.KV{
|
||||
{Key: "X-Claude-Code-Session-Id", Value: "sess-1"},
|
||||
{Key: "X-Claude-Code-Agent-Id", Value: "agent-7"},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
agent, ok := metaValue(t, out.Metadata, middleware.KeyLLMAgentID)
|
||||
require.True(t, ok, "the spawned agent's id must be emitted")
|
||||
assert.Equal(t, "agent-7", agent)
|
||||
|
||||
_, ok = metaValue(t, out.Metadata, middleware.KeyLLMParentAgentID)
|
||||
assert.False(t, ok, "a top-level agent has no parent to emit")
|
||||
})
|
||||
|
||||
t.Run("nested agent", func(t *testing.T) {
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/messages",
|
||||
Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`),
|
||||
Headers: []middleware.KV{
|
||||
{Key: "X-Claude-Code-Agent-Id", Value: "agent-9"},
|
||||
{Key: "X-Claude-Code-Parent-Agent-Id", Value: "agent-7"},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
agent, _ := metaValue(t, out.Metadata, middleware.KeyLLMAgentID)
|
||||
assert.Equal(t, "agent-9", agent)
|
||||
parent, ok := metaValue(t, out.Metadata, middleware.KeyLLMParentAgentID)
|
||||
require.True(t, ok, "a nested agent must carry the spawning agent's id")
|
||||
assert.Equal(t, "agent-7", parent)
|
||||
})
|
||||
|
||||
t.Run("absent on a plain request", func(t *testing.T) {
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/messages",
|
||||
Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, ok := metaValue(t, out.Metadata, middleware.KeyLLMAgentID)
|
||||
assert.False(t, ok, "no key is emitted when the client sends no agent id")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
package llm_router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware"
|
||||
)
|
||||
|
||||
// TestRouteClaimsModel_BedrockNormalizesCandidate guards the fix for the native
|
||||
@@ -28,3 +32,86 @@ func TestRouteClaimsModel_BedrockNormalizesCandidate(t *testing.T) {
|
||||
assert.False(t, routeClaimsModel(openai, "us.gpt-4o"),
|
||||
"non-Bedrock routes must not strip a us. prefix")
|
||||
}
|
||||
|
||||
// TestRouter_BedrockCountTokensRoutes pins that the token-counting action
|
||||
// reaches the Bedrock route instead of denying as not-routable.
|
||||
func TestRouter_BedrockCountTokensRoutes(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{{
|
||||
ID: "bedrock-prod",
|
||||
Bedrock: true,
|
||||
Models: []string{"anthropic.claude-sonnet-4-5"},
|
||||
AllowedGroupIDs: []string{defaultTestGroup},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com",
|
||||
}}})
|
||||
|
||||
in := newInputWithModelAndURL("anthropic.claude-sonnet-4-5",
|
||||
"/model/anthropic.claude-sonnet-4-5-20250929-v1:0/count-tokens")
|
||||
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "bedrock"})
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "count-tokens must route, not deny")
|
||||
require.NotNil(t, out.Mutations)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Equal(t, "bedrock-runtime.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host)
|
||||
}
|
||||
|
||||
// TestRouter_BedrockInferenceProfilesRoutes covers the startup lookups a
|
||||
// client makes to resolve a configured inference profile. They carry no
|
||||
// model, so before they were recognised they denied and wrote a policy
|
||||
// rejection into the access log on every session start.
|
||||
func TestRouter_BedrockInferenceProfilesRoutes(t *testing.T) {
|
||||
bedrock := ProviderRoute{
|
||||
ID: "bedrock-prod",
|
||||
Bedrock: true,
|
||||
AllowedGroupIDs: []string{defaultTestGroup},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com",
|
||||
}
|
||||
openai := ProviderRoute{
|
||||
ID: "openai-prod",
|
||||
Models: []string{"gpt-4o"},
|
||||
AllowedGroupIDs: []string{defaultTestGroup},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "api.openai.com",
|
||||
}
|
||||
mw := New(Config{Providers: []ProviderRoute{openai, bedrock}})
|
||||
|
||||
for _, path := range []string{
|
||||
"/inference-profiles?type=SYSTEM_DEFINED",
|
||||
"/inference-profiles/us.anthropic.claude-sonnet-5",
|
||||
} {
|
||||
out, err := mw.Invoke(context.Background(), newModellessInput(path))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "%s must route", path)
|
||||
require.NotNil(t, out.Mutations)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Equal(t, "bedrock-runtime.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host,
|
||||
"%s must reach the Bedrock provider, not the first authorised one", path)
|
||||
|
||||
nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
|
||||
assert.Equal(t, "true", nonInference, "%s carries no model to gate on", path)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRouter_BedrockNamespacedInferenceProfilesStripsPrefix pins that the
|
||||
// optional gateway namespace is removed before the request goes upstream.
|
||||
func TestRouter_BedrockNamespacedInferenceProfilesStripsPrefix(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{{
|
||||
ID: "bedrock-prod",
|
||||
Bedrock: true,
|
||||
AllowedGroupIDs: []string{defaultTestGroup},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com",
|
||||
}}})
|
||||
|
||||
out, err := mw.Invoke(context.Background(), newModellessInput("/bedrock/inference-profiles"))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out.Mutations)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Equal(t, "/bedrock", out.Mutations.RewriteUpstream.StripPathPrefix,
|
||||
"the namespace prefix must not reach the real Bedrock endpoint")
|
||||
}
|
||||
|
||||
@@ -109,6 +109,10 @@ func (m *Middleware) MetadataKeys() []string {
|
||||
middleware.KeyLLMAuthorisingGroups,
|
||||
middleware.KeyLLMPolicyDecision,
|
||||
middleware.KeyLLMPolicyReason,
|
||||
middleware.KeyLLMNonInference,
|
||||
// Emitted only for the per-model lookup, whose model lives in the path
|
||||
// rather than a body the parser could read.
|
||||
middleware.KeyLLMModel,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -137,29 +141,26 @@ const (
|
||||
// known to a provider that no policy authorises for the caller deny
|
||||
// with no_authorised_provider.
|
||||
func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
|
||||
reqPath := requestPath(in.URL)
|
||||
// The caller's API dialect, used to mirror a denial in the vendor's own
|
||||
// error shape so the client can explain it to the user.
|
||||
surface, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider)
|
||||
model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
|
||||
|
||||
// Vertex AI carries the model in the URL path, not the body, and is
|
||||
// selected by path rather than by the model/vendor table. Route it before
|
||||
// the model lookup so a model the parser extracted from the path can't be
|
||||
// claimed by a same-vendor direct provider (e.g. claude-* on api.anthropic.com).
|
||||
reqPath := requestPath(in.URL)
|
||||
if isVertexPath(reqPath) {
|
||||
model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
|
||||
// The request parser emits no llm.provider for a Vertex publisher it
|
||||
// can't parse (e.g. google/gemini). Forwarding such a request would
|
||||
// bypass token/budget metering, so deny it rather than serve it
|
||||
// unmetered.
|
||||
if vendor, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider); vendor == "" {
|
||||
return denyUnmeterable(), nil
|
||||
if surface == "" {
|
||||
return denyUnmeterable(surface), nil
|
||||
}
|
||||
route, outcome := m.matchVertex(reqPath, model, in.UserGroups)
|
||||
switch outcome {
|
||||
case matchOutcomeFound:
|
||||
return m.allowWithRoute(route, in.UserGroups), nil
|
||||
case matchOutcomeUnauthorised:
|
||||
return denyNoAuthorisedRoute(model), nil
|
||||
default:
|
||||
return denyUnknownModel(model), nil
|
||||
}
|
||||
return m.decide(route, outcome, surface, model, in.UserGroups, nil), nil
|
||||
}
|
||||
|
||||
// Bedrock likewise carries the model in the URL path (/model/{id}/{action}),
|
||||
@@ -167,52 +168,120 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
|
||||
// before the model lookup; when the prefix is present, strip it from the
|
||||
// forwarded path so the real Bedrock endpoint receives its native path.
|
||||
if isBedrockPath(reqPath) {
|
||||
model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
|
||||
native, hadPrefix := splitBedrockNamespace(reqPath)
|
||||
route, outcome := m.matchBedrock(native, model, in.UserGroups)
|
||||
switch outcome {
|
||||
case matchOutcomeFound:
|
||||
out := m.allowWithRoute(route, in.UserGroups)
|
||||
if hadPrefix && out.Mutations != nil && out.Mutations.RewriteUpstream != nil {
|
||||
out.Mutations.RewriteUpstream.StripPathPrefix = bedrockNamespacePrefix
|
||||
return m.decide(route, outcome, surface, model, in.UserGroups, func(out *middleware.Output) {
|
||||
if hadPrefix {
|
||||
stripBedrockNamespace(out)
|
||||
}
|
||||
return out, nil
|
||||
case matchOutcomeUnauthorised:
|
||||
return denyNoAuthorisedRoute(model), nil
|
||||
default:
|
||||
return denyUnknownModel(model), nil
|
||||
}
|
||||
}), nil
|
||||
}
|
||||
|
||||
model, ok := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
|
||||
if !ok || model == "" {
|
||||
// Non-inference endpoints (model listing) carry no model but still
|
||||
// need rewriting from the synth placeholder to a real upstream;
|
||||
// clients such as Codex call GET /v1/models at startup to enumerate
|
||||
// availability and read a 403 as "model unavailable".
|
||||
route, outcome := m.matchModelless(requestPath(in.URL), in.UserGroups)
|
||||
switch outcome {
|
||||
case matchOutcomeFound:
|
||||
return m.allowWithRoute(route, in.UserGroups), nil
|
||||
case matchOutcomeUnauthorised:
|
||||
// A recognised model-less endpoint exists but no provider
|
||||
// authorises the caller — deny as an authorisation failure
|
||||
// rather than masking it as a missing model.
|
||||
return denyNoAuthorisedRoute(model), nil
|
||||
default:
|
||||
return denyMissingModel(), nil
|
||||
}
|
||||
// GET /v1/models/{id} carries no body, so no model reaches the router in
|
||||
// metadata — but the path names one, and answering it confirms a model
|
||||
// exists and is reachable. Authorise it against the model table like any
|
||||
// other per-model request, then mark it non-inference so it still skips
|
||||
// the token pre-flight it would otherwise charge nothing against.
|
||||
if detail, isDetail := modelDetailID(reqPath); isDetail && isNonInferenceMethod(in.Method) {
|
||||
route, outcome := m.matchRoute(detail, surface, reqPath, in.UserGroups)
|
||||
return m.decide(route, outcome, surface, detail, in.UserGroups, func(out *middleware.Output) {
|
||||
markNonInference(out)
|
||||
// The parser reads models from JSON bodies only, and this request
|
||||
// has none, so stamp the one the path names. Without it the
|
||||
// guardrail's own allowlist — a separate, possibly narrower list
|
||||
// than the route's — never sees a model to check.
|
||||
out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMModel, Value: detail})
|
||||
}), nil
|
||||
}
|
||||
|
||||
vendor, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider)
|
||||
route, outcome := m.matchRoute(model, vendor, requestPath(in.URL), in.UserGroups)
|
||||
if model == "" {
|
||||
return m.routeModelless(reqPath, surface, in.Method, in.UserGroups), nil
|
||||
}
|
||||
|
||||
route, outcome := m.matchRoute(model, surface, reqPath, in.UserGroups)
|
||||
return m.decide(route, outcome, surface, model, in.UserGroups, nil), nil
|
||||
}
|
||||
|
||||
// decide turns a per-model match result into the middleware's decision. Every
|
||||
// surface that routes by model shares the same two denial arms — a model no
|
||||
// route claims is not routable, one that some route claims but none authorises
|
||||
// for this caller is an authorisation failure — so they live here once.
|
||||
// decorate, when non-nil, adjusts the allow with whatever that surface needs.
|
||||
func (m *Middleware) decide(
|
||||
route ProviderRoute,
|
||||
outcome matchOutcome,
|
||||
surface, model string,
|
||||
userGroups []string,
|
||||
decorate func(*middleware.Output),
|
||||
) *middleware.Output {
|
||||
switch outcome {
|
||||
case matchOutcomeFound:
|
||||
return m.allowWithRoute(route, in.UserGroups), nil
|
||||
out := m.allowWithRoute(route, surface, userGroups)
|
||||
if decorate != nil {
|
||||
decorate(out)
|
||||
}
|
||||
return out
|
||||
case matchOutcomeUnauthorised:
|
||||
return denyNoAuthorisedRoute(model), nil
|
||||
return denyNoAuthorisedRoute(surface, model)
|
||||
default:
|
||||
return denyUnknownModel(model), nil
|
||||
return denyUnknownModel(surface, model)
|
||||
}
|
||||
}
|
||||
|
||||
// routeModelless serves the endpoints that name no model at all: the model
|
||||
// listing, the connection-warming probe, and the Bedrock inference-profile
|
||||
// lookup. They still need rewriting from the synth placeholder to a real
|
||||
// upstream — clients such as Codex call GET /v1/models at startup to enumerate
|
||||
// availability and read a 403 as "model unavailable".
|
||||
func (m *Middleware) routeModelless(reqPath, surface, method string, userGroups []string) *middleware.Output {
|
||||
route, outcome := m.matchModelless(reqPath, method, userGroups)
|
||||
switch outcome {
|
||||
case matchOutcomeFound:
|
||||
out := m.allowWithRoute(route, surface, userGroups)
|
||||
markNonInference(out)
|
||||
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...)
|
||||
}
|
||||
return out
|
||||
case matchOutcomeUnauthorised:
|
||||
// A recognised model-less endpoint exists but no provider authorises
|
||||
// the caller — deny as an authorisation failure rather than masking it
|
||||
// as a missing model.
|
||||
return denyNoAuthorisedRoute(surface, "")
|
||||
default:
|
||||
return denyMissingModel(surface)
|
||||
}
|
||||
}
|
||||
|
||||
// isNonInferenceMethod reports whether a request method is one the
|
||||
// non-inference endpoints actually use: the listing and the per-model lookup
|
||||
// are GET, the connection-warming probe is HEAD or GET. The method is the only
|
||||
// thing separating "GET /v1/models/{id}" from a POST to the same path carrying
|
||||
// an inference body, and the non-inference mark exempts a request from the
|
||||
// token pre-flight — so anything else falls through to normal per-model
|
||||
// routing, which denies when the request names no model.
|
||||
func isNonInferenceMethod(method string) bool {
|
||||
return method == http.MethodGet || method == http.MethodHead
|
||||
}
|
||||
|
||||
// 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) {
|
||||
out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"})
|
||||
}
|
||||
|
||||
// stripBedrockNamespace tells the rewrite to drop the optional "/bedrock"
|
||||
// gateway namespace so the upstream receives its native Bedrock path.
|
||||
func stripBedrockNamespace(out *middleware.Output) {
|
||||
if out.Mutations != nil && out.Mutations.RewriteUpstream != nil {
|
||||
out.Mutations.RewriteUpstream.StripPathPrefix = bedrockNamespacePrefix
|
||||
}
|
||||
}
|
||||
|
||||
@@ -300,12 +369,60 @@ func (m *Middleware) matchRoute(model, vendor, reqPath string, userGroups []stri
|
||||
return best, matchOutcomeFound
|
||||
}
|
||||
|
||||
// isModelLessPath reports whether reqPath is a known OpenAI-shaped
|
||||
// non-inference endpoint that legitimately carries no model in its
|
||||
// request (the model-listing endpoints). These must route to an upstream
|
||||
// rather than deny, so model enumeration works end to end.
|
||||
// connectionWarmPath is the probe Anthropic clients send before their first
|
||||
// inference request to open the upstream connection early. Forwarding it
|
||||
// warms the connection the request will actually use; denying it only fills
|
||||
// the access log with rejections at every session start.
|
||||
const connectionWarmPath = "/api/hello"
|
||||
|
||||
// modelListingPath is the endpoint clients read at startup to populate
|
||||
// their model picker. Its response is a list the proxy can bound; the
|
||||
// per-model "/v1/models/{id}" lookup returns a single object and is left
|
||||
// alone.
|
||||
const modelListingPath = "/v1/models"
|
||||
|
||||
// isModelLessPath reports whether reqPath is a known non-inference endpoint
|
||||
// that legitimately carries no model at all: the model listing and the
|
||||
// connection-warming probe. These must route to an upstream rather than
|
||||
// deny, so model enumeration works end to end. The per-model
|
||||
// "/v1/models/{id}" lookup is deliberately excluded — it names a model, so
|
||||
// it is authorised against the model table instead (see modelDetailID).
|
||||
func isModelLessPath(reqPath string) bool {
|
||||
return reqPath == "/v1/models" || strings.HasPrefix(reqPath, "/v1/models/")
|
||||
return reqPath == modelListingPath || reqPath == connectionWarmPath
|
||||
}
|
||||
|
||||
// modelDetailID returns the model id named by a "/v1/models/{id}" lookup.
|
||||
// reqPath comes from url.URL.Path, which is already percent-decoded, so an
|
||||
// id carrying a "/" (a self-hosted "Qwen/Qwen2.5-0.5B-Instruct" sent as
|
||||
// "Qwen%2FQwen2.5-...") arrives whole and everything after the prefix is the
|
||||
// id, separators included.
|
||||
func modelDetailID(reqPath string) (string, bool) {
|
||||
if !strings.HasPrefix(reqPath, modelListingPath+"/") {
|
||||
return "", false
|
||||
}
|
||||
id := strings.TrimPrefix(reqPath, modelListingPath+"/")
|
||||
if id == "" {
|
||||
return "", false
|
||||
}
|
||||
return id, true
|
||||
}
|
||||
|
||||
// isBedrockModelLessPath reports whether reqPath is a Bedrock
|
||||
// inference-profile lookup, optionally behind the "/bedrock" gateway
|
||||
// namespace. Clients read these at startup to resolve a configured profile
|
||||
// to its underlying model. They carry no model of their own, so they route
|
||||
// by path to a Bedrock provider rather than through the model table.
|
||||
//
|
||||
// On native AWS these live on the control plane ("bedrock.<region>") while a
|
||||
// provider's upstream is normally the runtime host ("bedrock-runtime.<region>"),
|
||||
// so forwarding yields a 404 there. That is deliberate: a client has one base
|
||||
// URL, so pointing it straight at the runtime host 404s identically, and
|
||||
// forwarding keeps the proxy transparent instead of inventing a policy denial
|
||||
// the client would never otherwise see. Operators whose Bedrock upstream is a
|
||||
// gateway that does serve the lookup get a working answer.
|
||||
func isBedrockModelLessPath(reqPath string) bool {
|
||||
native, _ := splitBedrockNamespace(reqPath)
|
||||
return native == "/inference-profiles" || strings.HasPrefix(native, "/inference-profiles/")
|
||||
}
|
||||
|
||||
// isVertexPath reports whether reqPath is a Google Vertex AI publisher
|
||||
@@ -332,20 +449,33 @@ func splitBedrockNamespace(reqPath string) (string, bool) {
|
||||
return reqPath, false
|
||||
}
|
||||
|
||||
// bedrockActions are the runtime actions that follow the model id in a
|
||||
// Bedrock path. count-tokens is here so a client can price its context
|
||||
// against the dedicated endpoint; denying it pushes that work back onto
|
||||
// the inference endpoint, which bills for it.
|
||||
var bedrockActions = []string{
|
||||
"/invoke",
|
||||
"/invoke-with-response-stream",
|
||||
"/converse",
|
||||
"/converse-stream",
|
||||
"/count-tokens",
|
||||
}
|
||||
|
||||
// isBedrockPath reports whether reqPath is an AWS Bedrock runtime model
|
||||
// endpoint: /model/{modelId}/{action} where action is invoke,
|
||||
// invoke-with-response-stream, converse, or converse-stream — optionally behind
|
||||
// a "/bedrock" gateway-namespace prefix. The model lives in the path, so these
|
||||
// requests are routed by path to the Bedrock provider.
|
||||
// endpoint: /model/{modelId}/{action} — optionally behind a "/bedrock"
|
||||
// gateway-namespace prefix. The model lives in the path, so these requests
|
||||
// are routed by path to the Bedrock provider.
|
||||
func isBedrockPath(reqPath string) bool {
|
||||
native, _ := splitBedrockNamespace(reqPath)
|
||||
if !strings.HasPrefix(native, "/model/") {
|
||||
return false
|
||||
}
|
||||
return strings.HasSuffix(native, "/invoke") ||
|
||||
strings.HasSuffix(native, "/invoke-with-response-stream") ||
|
||||
strings.HasSuffix(native, "/converse") ||
|
||||
strings.HasSuffix(native, "/converse-stream")
|
||||
for _, action := range bedrockActions {
|
||||
if strings.HasSuffix(native, action) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// matchVertex selects the Vertex provider authorised for the caller's groups
|
||||
@@ -425,19 +555,26 @@ func (m *Middleware) matchPathRoute(reqPath, model string, userGroups []string,
|
||||
// declaration order), matchOutcomeUnauthorised when no provider authorises
|
||||
// the caller, or matchOutcomeUnknownModel when the path isn't a recognised
|
||||
// model-less endpoint.
|
||||
func (m *Middleware) matchModelless(reqPath string, userGroups []string) (ProviderRoute, matchOutcome) {
|
||||
if !isModelLessPath(reqPath) {
|
||||
func (m *Middleware) matchModelless(reqPath, method string, userGroups []string) (ProviderRoute, matchOutcome) {
|
||||
if !isNonInferenceMethod(method) {
|
||||
return ProviderRoute{}, matchOutcomeUnknownModel
|
||||
}
|
||||
var candidates []ProviderRoute
|
||||
for _, route := range m.cfg.Providers {
|
||||
var eligible func(ProviderRoute) bool
|
||||
switch {
|
||||
case isBedrockModelLessPath(reqPath):
|
||||
eligible = func(r ProviderRoute) bool { return r.Bedrock }
|
||||
case isModelLessPath(reqPath):
|
||||
// Vertex/Bedrock are path-routed and don't serve OpenAI-style
|
||||
// model-listing endpoints; including them here could rewrite a
|
||||
// GET /v1/models to an upstream that 404s it.
|
||||
if route.Vertex || route.Bedrock {
|
||||
continue
|
||||
}
|
||||
if routeAuthorisesGroups(route, userGroups) {
|
||||
eligible = func(r ProviderRoute) bool { return !r.Vertex && !r.Bedrock }
|
||||
default:
|
||||
return ProviderRoute{}, matchOutcomeUnknownModel
|
||||
}
|
||||
|
||||
var candidates []ProviderRoute
|
||||
for _, route := range m.cfg.Providers {
|
||||
if eligible(route) && routeAuthorisesGroups(route, userGroups) {
|
||||
candidates = append(candidates, route)
|
||||
}
|
||||
}
|
||||
@@ -564,6 +701,16 @@ func routeClaimsModel(route ProviderRoute, model string) bool {
|
||||
if route.Bedrock && llm.NormalizeBedrockModel(candidate) == model {
|
||||
return true
|
||||
}
|
||||
// A client may pin a dated Anthropic id ("claude-sonnet-4-5-20250929")
|
||||
// where the operator registered the undated one. Only an undated
|
||||
// registration absorbs a dated request: normalising both sides would
|
||||
// let a route pinned to one dated release claim a different one, so an
|
||||
// operator who deliberately pinned a build would silently serve
|
||||
// another — and with several such routes, ordering would decide which.
|
||||
if candidate == llm.NormalizeAnthropicModel(candidate) &&
|
||||
candidate == llm.NormalizeAnthropicModel(model) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -612,7 +759,7 @@ func requestPath(raw string) string {
|
||||
// provider id so identity-stamping middlewares (llm_identity_inject)
|
||||
// tag the request with ONLY the groups that authorised this specific
|
||||
// route — not every group the peer happens to be in.
|
||||
func (m *Middleware) allowWithRoute(route ProviderRoute, userGroups []string) *middleware.Output {
|
||||
func (m *Middleware) allowWithRoute(route ProviderRoute, surface string, userGroups []string) *middleware.Output {
|
||||
rewrite := &middleware.UpstreamRewrite{
|
||||
Scheme: route.UpstreamScheme,
|
||||
Host: route.UpstreamHost,
|
||||
@@ -634,7 +781,7 @@ func (m *Middleware) allowWithRoute(route ProviderRoute, userGroups []string) *m
|
||||
// request time (cached + auto-refreshed) instead of a static value.
|
||||
bearer, err := m.gcpBearer(route.GCPServiceAccountKeyB64)
|
||||
if err != nil {
|
||||
return denyUpstreamAuth()
|
||||
return denyUpstreamAuth(surface)
|
||||
}
|
||||
authValue = bearer
|
||||
}
|
||||
@@ -704,11 +851,12 @@ func (m *Middleware) gcpTokenSource(saKeyB64 string) (oauth2.TokenSource, error)
|
||||
// denyUpstreamAuth is returned when the router cannot obtain the upstream
|
||||
// credential (e.g. a malformed service-account key or an unreachable token
|
||||
// endpoint). It surfaces as a 502 — an upstream problem, not a policy denial.
|
||||
func denyUpstreamAuth() *middleware.Output {
|
||||
func denyUpstreamAuth(surface string) *middleware.Output {
|
||||
return &middleware.Output{
|
||||
Decision: middleware.DecisionDeny,
|
||||
DenyStatus: 502,
|
||||
DenyReason: &middleware.DenyReason{
|
||||
Surface: surface,
|
||||
Code: denyCodeUpstreamAuth,
|
||||
Message: "could not obtain upstream credential",
|
||||
},
|
||||
@@ -722,11 +870,12 @@ func denyUpstreamAuth() *middleware.Output {
|
||||
// denyUnmeterable returns the deny envelope for a path-routed request whose
|
||||
// publisher has no parser surface, so its usage can't be metered. Serving it
|
||||
// would bypass token/budget caps, so it is rejected with a 403.
|
||||
func denyUnmeterable() *middleware.Output {
|
||||
func denyUnmeterable(surface string) *middleware.Output {
|
||||
return &middleware.Output{
|
||||
Decision: middleware.DecisionDeny,
|
||||
DenyStatus: 403,
|
||||
DenyReason: &middleware.DenyReason{
|
||||
Surface: surface,
|
||||
Code: denyCodeUnmeterable,
|
||||
Message: "request publisher is not supported for metering",
|
||||
},
|
||||
@@ -739,11 +888,12 @@ func denyUnmeterable() *middleware.Output {
|
||||
|
||||
// denyMissingModel returns the deny envelope for a request whose
|
||||
// envelope has no llm.model metadata.
|
||||
func denyMissingModel() *middleware.Output {
|
||||
func denyMissingModel(surface string) *middleware.Output {
|
||||
return &middleware.Output{
|
||||
Decision: middleware.DecisionDeny,
|
||||
DenyStatus: 403,
|
||||
DenyReason: &middleware.DenyReason{
|
||||
Surface: surface,
|
||||
Code: denyCodeNotRoutable,
|
||||
Message: "missing llm.model on request envelope",
|
||||
},
|
||||
@@ -756,11 +906,12 @@ func denyMissingModel() *middleware.Output {
|
||||
|
||||
// denyUnknownModel returns the deny envelope for a model that no
|
||||
// configured provider claims.
|
||||
func denyUnknownModel(model string) *middleware.Output {
|
||||
func denyUnknownModel(surface, model string) *middleware.Output {
|
||||
return &middleware.Output{
|
||||
Decision: middleware.DecisionDeny,
|
||||
DenyStatus: 403,
|
||||
DenyReason: &middleware.DenyReason{
|
||||
Surface: surface,
|
||||
Code: denyCodeNotRoutable,
|
||||
Message: fmt.Sprintf("no provider configured for model %s", model),
|
||||
Details: map[string]string{"model": model},
|
||||
@@ -775,11 +926,12 @@ func denyUnknownModel(model string) *middleware.Output {
|
||||
// denyNoAuthorisedRoute returns the deny envelope for a model that one
|
||||
// or more providers claim, but where no policy authorises the caller's
|
||||
// groups for any of those providers.
|
||||
func denyNoAuthorisedRoute(model string) *middleware.Output {
|
||||
func denyNoAuthorisedRoute(surface, model string) *middleware.Output {
|
||||
return &middleware.Output{
|
||||
Decision: middleware.DecisionDeny,
|
||||
DenyStatus: 403,
|
||||
DenyReason: &middleware.DenyReason{
|
||||
Surface: surface,
|
||||
Code: denyCodeNoAuthorisedRoute,
|
||||
Message: fmt.Sprintf("no policy authorises model %s for the caller's groups", model),
|
||||
Details: map[string]string{"model": model},
|
||||
|
||||
@@ -2,6 +2,7 @@ package llm_router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -60,6 +61,8 @@ func TestMiddlewareIdentity(t *testing.T) {
|
||||
[]string{
|
||||
middleware.KeyLLMResolvedProviderID,
|
||||
middleware.KeyLLMAuthorisingGroups,
|
||||
middleware.KeyLLMNonInference,
|
||||
middleware.KeyLLMModel,
|
||||
middleware.KeyLLMPolicyDecision,
|
||||
middleware.KeyLLMPolicyReason,
|
||||
},
|
||||
@@ -171,8 +174,12 @@ func TestRouter_MissingModel(t *testing.T) {
|
||||
// from which a model could be parsed). UserGroups matches defaultTestGroup.
|
||||
func newModellessInput(reqURL string) *middleware.Input {
|
||||
return &middleware.Input{
|
||||
Slot: middleware.SlotOnRequest,
|
||||
URL: reqURL,
|
||||
Slot: middleware.SlotOnRequest,
|
||||
URL: reqURL,
|
||||
// The non-inference endpoints are read requests; the method is what
|
||||
// separates them from an inference body posted to the same path, so
|
||||
// state it rather than leaning on the zero value.
|
||||
Method: http.MethodGet,
|
||||
UserGroups: []string{defaultTestGroup},
|
||||
}
|
||||
}
|
||||
@@ -197,6 +204,12 @@ func TestRouter_ModelLessPath_RoutesToAuthorisedProvider(t *testing.T) {
|
||||
|
||||
provider, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID)
|
||||
assert.Equal(t, "openai-prod", provider, "resolved provider must be the authorised route")
|
||||
|
||||
// The limits gate reads this to tell "no model applies here" from
|
||||
// "the model could not be determined", which fails closed.
|
||||
nonInference, ok := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
|
||||
require.True(t, ok, "model-less allow must mark the request non-inference")
|
||||
assert.Equal(t, "true", nonInference)
|
||||
}
|
||||
|
||||
func TestRouter_ModelLessPath_MultiProviderDeclarationOrder(t *testing.T) {
|
||||
@@ -873,3 +886,262 @@ func TestRouter_EmptyModelsClaimsAnyModel(t *testing.T) {
|
||||
resolved, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID)
|
||||
assert.Equal(t, "litellm", resolved)
|
||||
}
|
||||
|
||||
// TestRouter_DatedAnthropicModelRoutes covers a client pinning a release
|
||||
// date on a model the operator registered undated. Exact matches still win,
|
||||
// so an operator who registers both dated releases keeps them distinct.
|
||||
func TestRouter_DatedAnthropicModelRoutes(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{{
|
||||
ID: "anthropic-prod",
|
||||
Vendor: "anthropic",
|
||||
Models: []string{"claude-sonnet-4-5"},
|
||||
AllowedGroupIDs: []string{defaultTestGroup},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "api.anthropic.com",
|
||||
}}})
|
||||
|
||||
in := newInputWithModelAndURL("claude-sonnet-4-5-20250929", "/v1/messages")
|
||||
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "a dated id must route to the undated registration")
|
||||
require.NotNil(t, out.Mutations)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host)
|
||||
}
|
||||
|
||||
// TestRouter_ConnectionWarmProbeRoutes covers the HEAD /api/hello probe an
|
||||
// Anthropic client sends before its first request. Forwarding it warms the
|
||||
// connection that request will use; denying it only wrote a rejection into
|
||||
// the access log at every session start.
|
||||
func TestRouter_ConnectionWarmProbeRoutes(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{{
|
||||
ID: "anthropic-prod",
|
||||
Vendor: "anthropic",
|
||||
Models: []string{"claude-sonnet-5"},
|
||||
AllowedGroupIDs: []string{defaultTestGroup},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "api.anthropic.com",
|
||||
}}})
|
||||
|
||||
in := newModellessInput("/api/hello")
|
||||
in.Method = http.MethodHead
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "the warm-up probe must reach the upstream")
|
||||
require.NotNil(t, out.Mutations)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host)
|
||||
|
||||
nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
|
||||
assert.Equal(t, "true", nonInference, "the probe carries no model to gate on")
|
||||
}
|
||||
|
||||
// TestRouter_ModelListingCarriesAuthorisedModels pins the list the proxy
|
||||
// bounds the discovery response with. A catch-all route enumerates nothing,
|
||||
// so it must not bound the upstream's list at all.
|
||||
func TestRouter_ModelListingCarriesAuthorisedModels(t *testing.T) {
|
||||
enumerated := ProviderRoute{
|
||||
ID: "anthropic-prod",
|
||||
Models: []string{"claude-sonnet-5", "claude-haiku-4-5"},
|
||||
AllowedGroupIDs: []string{defaultTestGroup},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "api.anthropic.com",
|
||||
}
|
||||
|
||||
t.Run("enumerated route bounds the listing", func(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{enumerated}})
|
||||
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models"))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out.Mutations)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Equal(t, []string{"claude-sonnet-5", "claude-haiku-4-5"},
|
||||
out.Mutations.RewriteUpstream.DiscoveryModels,
|
||||
"the picker must be bounded by what the route authorises")
|
||||
})
|
||||
|
||||
t.Run("catch-all route leaves the listing alone", func(t *testing.T) {
|
||||
catchAll := enumerated
|
||||
catchAll.Models = nil
|
||||
mw := New(Config{Providers: []ProviderRoute{catchAll}})
|
||||
|
||||
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models"))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Empty(t, out.Mutations.RewriteUpstream.DiscoveryModels,
|
||||
"a route that claims every model cannot bound the upstream's list")
|
||||
})
|
||||
|
||||
t.Run("per-model lookup is not a listing", func(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{enumerated}})
|
||||
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-sonnet-5"))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Empty(t, out.Mutations.RewriteUpstream.DiscoveryModels,
|
||||
"the single-object lookup has no data array to filter")
|
||||
})
|
||||
}
|
||||
|
||||
// TestRouter_ModelDetailHonoursAllowlist pins that GET /v1/models/{id} is
|
||||
// authorised against the model table. It carries no body model, so treating
|
||||
// it as a model-less endpoint would let a caller confirm a model the route
|
||||
// does not list — the listing itself is bounded to the allowlist, so the
|
||||
// detail lookup must be too.
|
||||
func TestRouter_ModelDetailHonoursAllowlist(t *testing.T) {
|
||||
enumerated := ProviderRoute{
|
||||
ID: "anthropic-prod",
|
||||
Models: []string{"claude-sonnet-5"},
|
||||
AllowedGroupIDs: []string{defaultTestGroup},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "api.anthropic.com",
|
||||
}
|
||||
|
||||
t.Run("allowlisted model routes and skips metering", func(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{enumerated}})
|
||||
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-sonnet-5"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision)
|
||||
require.NotNil(t, out.Mutations)
|
||||
require.NotNil(t, out.Mutations.RewriteUpstream)
|
||||
assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host)
|
||||
|
||||
nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
|
||||
assert.Equal(t, "true", nonInference, "a detail lookup spends no tokens")
|
||||
})
|
||||
|
||||
t.Run("model outside the allowlist denies", func(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{enumerated}})
|
||||
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-opus-5"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision,
|
||||
"a model no route lists must not be confirmed by the detail lookup")
|
||||
})
|
||||
|
||||
t.Run("dated id matches its undated registration", func(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{enumerated}})
|
||||
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-sonnet-5-20250929"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision,
|
||||
"a pinned release of an allowlisted family stays reachable")
|
||||
})
|
||||
|
||||
t.Run("catch-all route still answers every lookup", func(t *testing.T) {
|
||||
catchAll := enumerated
|
||||
catchAll.Models = nil
|
||||
mw := New(Config{Providers: []ProviderRoute{catchAll}})
|
||||
|
||||
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/anything-at-all"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision,
|
||||
"a gateway that enumerates nothing cannot refuse a lookup")
|
||||
})
|
||||
}
|
||||
|
||||
// TestRouter_NonInferenceRequiresReadMethod pins that the non-inference mark —
|
||||
// which exempts a request from the token pre-flight — is reachable only by the
|
||||
// read methods these endpoints actually use. A POST to the same path could
|
||||
// carry an inference body, so it must not buy the exemption; it falls through
|
||||
// to normal per-model routing instead, which denies when no model is named.
|
||||
func TestRouter_NonInferenceRequiresReadMethod(t *testing.T) {
|
||||
route := ProviderRoute{
|
||||
ID: "gateway",
|
||||
AllowedGroupIDs: []string{defaultTestGroup},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "gateway.example.com",
|
||||
}
|
||||
|
||||
for _, path := range []string{"/v1/models", "/v1/models/claude-sonnet-5", "/api/hello"} {
|
||||
t.Run("POST "+path, func(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{route}})
|
||||
|
||||
in := newModellessInput(path)
|
||||
in.Method = http.MethodPost
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision,
|
||||
"a write to a non-inference path must not route unmetered")
|
||||
|
||||
nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
|
||||
assert.NotEqual(t, "true", nonInference,
|
||||
"only a read method may skip the token pre-flight")
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("HEAD keeps the warm probe working", func(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{route}})
|
||||
|
||||
in := newModellessInput(connectionWarmPath)
|
||||
in.Method = http.MethodHead
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision,
|
||||
"the HEAD warm probe must still reach the upstream")
|
||||
|
||||
nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
|
||||
assert.Equal(t, "true", nonInference,
|
||||
"the HEAD warm probe carries no model to meter")
|
||||
})
|
||||
}
|
||||
|
||||
// TestRouter_PinnedDatedModelStaysDistinct pins that a route registered
|
||||
// against one dated Anthropic release does not claim another. Normalising
|
||||
// both sides of the comparison made every dated build of a family
|
||||
// interchangeable, so an operator who deliberately pinned a build would have
|
||||
// served a different one — and with several such routes, declaration or path
|
||||
// order would have decided which.
|
||||
func TestRouter_PinnedDatedModelStaysDistinct(t *testing.T) {
|
||||
pinned := ProviderRoute{
|
||||
ID: "anthropic-pinned",
|
||||
Vendor: "anthropic",
|
||||
Models: []string{"claude-sonnet-4-5-20250101"},
|
||||
AllowedGroupIDs: []string{defaultTestGroup},
|
||||
UpstreamScheme: "https",
|
||||
UpstreamHost: "pinned.example.com",
|
||||
}
|
||||
|
||||
t.Run("a different dated release is not claimed", func(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{pinned}})
|
||||
in := newInputWithModelAndURL("claude-sonnet-4-5-20250202", "/v1/messages")
|
||||
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision,
|
||||
"a route pinned to one dated build must not serve another")
|
||||
})
|
||||
|
||||
t.Run("its own dated release still routes", func(t *testing.T) {
|
||||
mw := New(Config{Providers: []ProviderRoute{pinned}})
|
||||
in := newInputWithModelAndURL("claude-sonnet-4-5-20250101", "/v1/messages")
|
||||
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
|
||||
|
||||
out, err := mw.Invoke(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "the exact match must still route")
|
||||
})
|
||||
|
||||
t.Run("two pinned builds each route to their own provider", func(t *testing.T) {
|
||||
other := pinned
|
||||
other.ID = "anthropic-pinned-newer"
|
||||
other.Models = []string{"claude-sonnet-4-5-20250202"}
|
||||
other.UpstreamHost = "newer.example.com"
|
||||
mw := New(Config{Providers: []ProviderRoute{pinned, other}})
|
||||
|
||||
in := newInputWithModelAndURL("claude-sonnet-4-5-20250202", "/v1/messages")
|
||||
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
|
||||
|
||||
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)
|
||||
assert.Equal(t, "newer.example.com", out.Mutations.RewriteUpstream.Host,
|
||||
"declaration order must not decide between two deliberately pinned builds")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -11,11 +11,78 @@ var codeRegex = regexp.MustCompile(`^[a-z][a-z0-9._-]{0,63}$`)
|
||||
// denyResponse is the on-wire shape rendered by RenderDenyResponse.
|
||||
// Keeping this as a typed struct ensures we never leak
|
||||
// middleware-supplied bytes outside known fields.
|
||||
//
|
||||
// Type and Error mirror the denial in the vendor's own error shape when
|
||||
// the request reached a known LLM surface. LLM clients only parse their
|
||||
// provider's envelope, so without the mirror a budget stop reaches the
|
||||
// user as an unexplained API error. The NetBird fields stay where they
|
||||
// were, so the body is a superset and existing consumers are unaffected.
|
||||
type denyResponse struct {
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message,omitempty"`
|
||||
Details map[string]string `json:"details,omitempty"`
|
||||
Middleware string `json:"middleware,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Error *providerError `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// providerError is the nested error object both vendor envelopes carry.
|
||||
type providerError struct {
|
||||
Type string `json:"type"`
|
||||
Message string `json:"message,omitempty"`
|
||||
Code string `json:"code,omitempty"`
|
||||
}
|
||||
|
||||
// Vendor error types keyed by HTTP status, per each provider's published
|
||||
// error reference.
|
||||
const (
|
||||
anthropicErrInvalidRequest = "invalid_request_error"
|
||||
anthropicErrPermission = "permission_error"
|
||||
anthropicErrRateLimit = "rate_limit_error"
|
||||
anthropicErrAPI = "api_error"
|
||||
openAIErrInvalidRequest = "invalid_request_error"
|
||||
openAIErrRateLimit = "rate_limit_error"
|
||||
)
|
||||
|
||||
// providerEnvelope returns the vendor-shaped mirror for a denial on the
|
||||
// given surface, or nil when the surface has no envelope we can speak.
|
||||
// message is the already-redacted public message.
|
||||
func providerEnvelope(surface, code, message string, status int) (string, *providerError) {
|
||||
switch surface {
|
||||
case "anthropic":
|
||||
return "error", &providerError{
|
||||
Type: anthropicErrorType(status),
|
||||
Message: message,
|
||||
}
|
||||
case "openai":
|
||||
return "", &providerError{
|
||||
Type: openAIErrorType(status),
|
||||
Message: message,
|
||||
Code: code,
|
||||
}
|
||||
default:
|
||||
return "", nil
|
||||
}
|
||||
}
|
||||
|
||||
func anthropicErrorType(status int) string {
|
||||
switch status {
|
||||
case http.StatusForbidden:
|
||||
return anthropicErrPermission
|
||||
case http.StatusTooManyRequests:
|
||||
return anthropicErrRateLimit
|
||||
case http.StatusBadRequest:
|
||||
return anthropicErrInvalidRequest
|
||||
default:
|
||||
return anthropicErrAPI
|
||||
}
|
||||
}
|
||||
|
||||
func openAIErrorType(status int) string {
|
||||
if status == http.StatusTooManyRequests {
|
||||
return openAIErrRateLimit
|
||||
}
|
||||
return openAIErrInvalidRequest
|
||||
}
|
||||
|
||||
// RenderDenyResponse writes a structured JSON deny body. Status is
|
||||
@@ -36,6 +103,7 @@ func RenderDenyResponse(w http.ResponseWriter, middlewareID string, reason *Deny
|
||||
Message: truncate(Scan(reason.Message), 256),
|
||||
Middleware: truncate(Scan(middlewareID), 64),
|
||||
}
|
||||
resp.Type, resp.Error = providerEnvelope(reason.Surface, resp.Code, resp.Message, status)
|
||||
if n := len(reason.Details); n > 0 {
|
||||
resp.Details = make(map[string]string, min(n, 8))
|
||||
for k, v := range reason.Details {
|
||||
|
||||
92
proxy/internal/middleware/decision_test.go
Normal file
92
proxy/internal/middleware/decision_test.go
Normal file
@@ -0,0 +1,92 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// decodeDeny renders a denial and returns the parsed body plus the status.
|
||||
func decodeDeny(t *testing.T, reason *DenyReason, status int) (map[string]any, int) {
|
||||
t.Helper()
|
||||
rec := httptest.NewRecorder()
|
||||
RenderDenyResponse(rec, "llm_limit_check", reason, status)
|
||||
|
||||
var body map[string]any
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body), "deny body must be valid JSON")
|
||||
return body, rec.Code
|
||||
}
|
||||
|
||||
// TestRenderDeny_AnthropicSurfaceMirrorsVendorShape covers a budget stop
|
||||
// reaching Claude Code. The client only parses the Anthropic envelope, so
|
||||
// without the mirror the user sees an unexplained API error instead of the
|
||||
// reason their request was refused.
|
||||
func TestRenderDeny_AnthropicSurfaceMirrorsVendorShape(t *testing.T) {
|
||||
body, status := decodeDeny(t, &DenyReason{
|
||||
Code: "llm_policy.budget_cap_exceeded",
|
||||
Message: "LLM policy limit exceeded",
|
||||
Surface: "anthropic",
|
||||
}, http.StatusForbidden)
|
||||
|
||||
assert.Equal(t, http.StatusForbidden, status)
|
||||
assert.Equal(t, "error", body["type"], "Anthropic errors carry type=error at the top level")
|
||||
|
||||
errObj, ok := body["error"].(map[string]any)
|
||||
require.True(t, ok, "error must be an object")
|
||||
assert.Equal(t, "permission_error", errObj["type"], "403 maps to permission_error")
|
||||
assert.Equal(t, "LLM policy limit exceeded", errObj["message"])
|
||||
|
||||
// The NetBird fields stay put so existing consumers keep working.
|
||||
assert.Equal(t, "llm_policy.budget_cap_exceeded", body["code"])
|
||||
assert.Equal(t, "LLM policy limit exceeded", body["message"])
|
||||
assert.Equal(t, "llm_limit_check", body["middleware"])
|
||||
}
|
||||
|
||||
// TestRenderDeny_OpenAISurfaceMirrorsVendorShape pins the OpenAI envelope,
|
||||
// which nests the code and carries no top-level type.
|
||||
func TestRenderDeny_OpenAISurfaceMirrorsVendorShape(t *testing.T) {
|
||||
body, _ := decodeDeny(t, &DenyReason{
|
||||
Code: "llm_policy.model_blocked",
|
||||
Message: "model is not in the policy allowlist",
|
||||
Surface: "openai",
|
||||
}, http.StatusForbidden)
|
||||
|
||||
assert.NotContains(t, body, "type", "OpenAI errors have no top-level type")
|
||||
|
||||
errObj, ok := body["error"].(map[string]any)
|
||||
require.True(t, ok, "error must be an object")
|
||||
assert.Equal(t, "invalid_request_error", errObj["type"])
|
||||
assert.Equal(t, "llm_policy.model_blocked", errObj["code"], "the NetBird code rides in the vendor code field")
|
||||
assert.Equal(t, "model is not in the policy allowlist", errObj["message"])
|
||||
}
|
||||
|
||||
// TestRenderDeny_RateLimitStatusMapsToVendorRateLimit pins the mapping a
|
||||
// client's backoff keys on.
|
||||
func TestRenderDeny_RateLimitStatusMapsToVendorRateLimit(t *testing.T) {
|
||||
body, status := decodeDeny(t, &DenyReason{
|
||||
Code: "llm_policy.token_cap_exceeded",
|
||||
Message: "LLM policy limit exceeded",
|
||||
Surface: "anthropic",
|
||||
}, http.StatusTooManyRequests)
|
||||
|
||||
assert.Equal(t, http.StatusTooManyRequests, status, "429 must survive the status clamp")
|
||||
errObj := body["error"].(map[string]any)
|
||||
assert.Equal(t, "rate_limit_error", errObj["type"])
|
||||
}
|
||||
|
||||
// TestRenderDeny_NoSurfaceKeepsLegacyShape guards non-LLM middlewares and
|
||||
// denials raised before a surface is known.
|
||||
func TestRenderDeny_NoSurfaceKeepsLegacyShape(t *testing.T) {
|
||||
body, _ := decodeDeny(t, &DenyReason{
|
||||
Code: "llm_policy.model_not_routable",
|
||||
Message: "no provider configured for model x",
|
||||
}, http.StatusForbidden)
|
||||
|
||||
assert.NotContains(t, body, "type", "no surface means no vendor mirror")
|
||||
assert.NotContains(t, body, "error", "no surface means no vendor mirror")
|
||||
assert.Equal(t, "llm_policy.model_not_routable", body["code"])
|
||||
}
|
||||
@@ -22,6 +22,15 @@ const (
|
||||
// body. Empty for clients that don't send one.
|
||||
KeyLLMSessionID = "llm.session_id"
|
||||
|
||||
// Sub-agent attribution (emitted by llm_request_parser from the
|
||||
// client's request headers). A coding agent that spawns helpers
|
||||
// stamps the spawned agent's id, and the spawning agent's id when
|
||||
// the helper is itself nested, so cost within one session can be
|
||||
// split across the agents that ran in parallel. These identify an
|
||||
// agent, not a person or a device: never treat them as a user id.
|
||||
KeyLLMAgentID = "llm.agent_id"
|
||||
KeyLLMParentAgentID = "llm.parent_agent_id"
|
||||
|
||||
// LLM response-side metadata (emitted by llm_response_parser).
|
||||
//nolint:gosec // metadata key name, not a credential
|
||||
KeyLLMInputTokens = "llm.input_tokens"
|
||||
@@ -66,6 +75,14 @@ const (
|
||||
// downstream gateways' spend logs.
|
||||
KeyLLMAuthorisingGroups = "llm.authorising_groups"
|
||||
|
||||
// LLM non-inference marker (emitted by llm_router on the allow path
|
||||
// for endpoints that legitimately carry no model, such as model
|
||||
// listing). The router still authorises these against the caller's
|
||||
// groups; the marker only tells the limits gate that a per-model
|
||||
// allowlist has nothing to evaluate, so an empty model must not be
|
||||
// read as an undetermined one. Never derived from client input.
|
||||
KeyLLMNonInference = "llm.non_inference"
|
||||
|
||||
// LLM policy attribution (emitted by llm_limit_check on the allow
|
||||
// path). Names the policy that paid for this request and the
|
||||
// dimension counters the post-flight llm_limit_record middleware
|
||||
|
||||
@@ -179,6 +179,12 @@ type DenyReason struct {
|
||||
Code string
|
||||
Message string
|
||||
Details map[string]string
|
||||
// Surface names the LLM API dialect the caller speaks (the
|
||||
// llm.provider value), so the rendered body can mirror the denial in
|
||||
// that vendor's error shape alongside the NetBird fields. Empty for
|
||||
// non-LLM middlewares and for denials raised before a surface was
|
||||
// resolved; the body then carries the NetBird fields alone.
|
||||
Surface string
|
||||
}
|
||||
|
||||
// Output is the value each middleware returns to the dispatcher. The
|
||||
@@ -247,6 +253,12 @@ type UpstreamRewrite struct {
|
||||
// without verifying its TLS certificate. Set by llm_router from the
|
||||
// provider's skip_tls_verification for self-hosted / internal gateways.
|
||||
SkipTLSVerify bool
|
||||
// DiscoveryModels, when non-empty, is the set of model ids the resolved
|
||||
// route authorises, and the proxy drops everything else from the
|
||||
// model-listing response. Empty leaves the upstream's list untouched,
|
||||
// which is what a route that claims every model wants. Set by
|
||||
// llm_router on a model-listing request only.
|
||||
DiscoveryModels []string
|
||||
}
|
||||
|
||||
// AuthHeader is a single name/value pair the proxy injects on the
|
||||
|
||||
191
proxy/internal/proxy/discovery_filter.go
Normal file
191
proxy/internal/proxy/discovery_filter.go
Normal file
@@ -0,0 +1,191 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
sharedllm "github.com/netbirdio/netbird/shared/llm"
|
||||
)
|
||||
|
||||
// maxDiscoveryBodyBytes bounds the model-listing response the filter will
|
||||
// buffer. A listing is a few kilobytes of ids; anything larger is not a
|
||||
// listing we recognise, and buffering it to rewrite would cost more than
|
||||
// the filtering is worth.
|
||||
const maxDiscoveryBodyBytes = 1 << 20
|
||||
|
||||
// modelDiscoveryFilter returns a ModifyResponse hook that drops models the
|
||||
// caller's policy does not authorise from a model-listing response, then
|
||||
// delegates to next (which may be nil).
|
||||
//
|
||||
// Clients populate their model picker from this endpoint, so an unfiltered
|
||||
// list offers models the very next request denies. The filter is
|
||||
// best-effort: a response it cannot safely rewrite passes through
|
||||
// untouched rather than reaching the client corrupted.
|
||||
func modelDiscoveryFilter(allowed []string, next func(*http.Response) error) func(*http.Response) error {
|
||||
permitted := make(map[string]struct{}, len(allowed)*2)
|
||||
for _, id := range allowed {
|
||||
permitted[id] = struct{}{}
|
||||
permitted[sharedllm.NormalizeAnthropicModel(id)] = struct{}{}
|
||||
}
|
||||
|
||||
return func(resp *http.Response) error {
|
||||
if err := filterModelListing(resp, permitted); err != nil {
|
||||
return err
|
||||
}
|
||||
if next == nil {
|
||||
return nil
|
||||
}
|
||||
return next(resp)
|
||||
}
|
||||
}
|
||||
|
||||
// filterModelListing rewrites the response body in place, keeping only the
|
||||
// entries whose id the policy authorises. Responses that are not a plain
|
||||
// JSON listing are left alone.
|
||||
func filterModelListing(resp *http.Response, permitted map[string]struct{}) error {
|
||||
if !isPlainJSONListing(resp) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// One byte past the cap, so an oversized body is detectable without
|
||||
// buffering all of it.
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, maxDiscoveryBodyBytes+1))
|
||||
if err != nil {
|
||||
_ = resp.Body.Close()
|
||||
return err
|
||||
}
|
||||
if len(body) > maxDiscoveryBodyBytes {
|
||||
// Too large to filter. Put the bytes already read back in front of the
|
||||
// unread remainder and forward the response exactly as the upstream
|
||||
// sent it, headers included. Buffering what was read and closing here
|
||||
// would truncate the body at the cap and hand the client a short,
|
||||
// invalid listing — worse than not filtering at all.
|
||||
resp.Body = spliceBody(body, resp.Body)
|
||||
return nil
|
||||
}
|
||||
if err := resp.Body.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
filtered, ok := filterListingBody(body, permitted)
|
||||
if !ok {
|
||||
restoreBody(resp, body)
|
||||
return nil
|
||||
}
|
||||
restoreBody(resp, filtered)
|
||||
return nil
|
||||
}
|
||||
|
||||
// isPlainJSONListing reports whether the response is a JSON body the filter
|
||||
// can parse. A content-encoded body is skipped: the transport only
|
||||
// transparently decompresses what it negotiated itself, and the client
|
||||
// negotiates its own encoding on this request.
|
||||
func isPlainJSONListing(resp *http.Response) bool {
|
||||
if resp == nil || resp.Body == nil {
|
||||
return false
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return false
|
||||
}
|
||||
if enc := resp.Header.Get("Content-Encoding"); enc != "" && !strings.EqualFold(enc, "identity") {
|
||||
return false
|
||||
}
|
||||
return strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "application/json")
|
||||
}
|
||||
|
||||
// filterListingBody returns the listing with unauthorised entries removed.
|
||||
// ok is false when the body is not a listing shape, in which case the
|
||||
// caller must forward the original bytes.
|
||||
func filterListingBody(body []byte, permitted map[string]struct{}) ([]byte, bool) {
|
||||
var doc map[string]json.RawMessage
|
||||
if err := json.Unmarshal(body, &doc); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
raw, present := doc["data"]
|
||||
if !present {
|
||||
return nil, false
|
||||
}
|
||||
var entries []map[string]json.RawMessage
|
||||
if err := json.Unmarshal(raw, &entries); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
kept := make([]map[string]json.RawMessage, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
if entryPermitted(entry, permitted) {
|
||||
kept = append(kept, entry)
|
||||
}
|
||||
}
|
||||
|
||||
encoded, err := json.Marshal(kept)
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
doc["data"] = encoded
|
||||
out, err := json.Marshal(doc)
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
return out, true
|
||||
}
|
||||
|
||||
// entryPermitted reports whether a listing entry names a model the policy
|
||||
// authorises, trying every form the same model is written in.
|
||||
func entryPermitted(entry map[string]json.RawMessage, permitted map[string]struct{}) bool {
|
||||
raw, ok := entry["id"]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
var id string
|
||||
if err := json.Unmarshal(raw, &id); err != nil {
|
||||
return false
|
||||
}
|
||||
for _, candidate := range modelIDForms(id) {
|
||||
if _, ok := permitted[candidate]; ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// modelIDForms returns the forms a single model id may be written in: the id
|
||||
// itself, its undated form, and the same two with a gateway's provider
|
||||
// prefix removed ("vertex_ai/claude-sonnet-5"). The bare id is tried first,
|
||||
// because a self-hosted id can legitimately contain a slash of its own
|
||||
// ("Qwen/Qwen2.5-0.5B-Instruct") and must not be cut down to its tail.
|
||||
func modelIDForms(id string) []string {
|
||||
if id == "" {
|
||||
return nil
|
||||
}
|
||||
forms := []string{id, sharedllm.NormalizeAnthropicModel(id)}
|
||||
if slash := strings.LastIndex(id, "/"); slash >= 0 {
|
||||
tail := id[slash+1:]
|
||||
forms = append(forms, tail, sharedllm.NormalizeAnthropicModel(tail))
|
||||
}
|
||||
return forms
|
||||
}
|
||||
|
||||
// restoreBody puts body back on the response and fixes the length headers
|
||||
// so the client reads exactly what is there.
|
||||
// spliceBody returns a ReadCloser that yields prefix followed by whatever is
|
||||
// left in rest, closing rest when closed. It lets the filter put back bytes it
|
||||
// consumed while deciding, without owning the rest of the stream.
|
||||
func spliceBody(prefix []byte, rest io.ReadCloser) io.ReadCloser {
|
||||
return struct {
|
||||
io.Reader
|
||||
io.Closer
|
||||
}{
|
||||
Reader: io.MultiReader(bytes.NewReader(prefix), rest),
|
||||
Closer: rest,
|
||||
}
|
||||
}
|
||||
|
||||
func restoreBody(resp *http.Response, body []byte) {
|
||||
resp.Body = io.NopCloser(bytes.NewReader(body))
|
||||
resp.ContentLength = int64(len(body))
|
||||
resp.Header.Set("Content-Length", strconv.Itoa(len(body)))
|
||||
}
|
||||
217
proxy/internal/proxy/discovery_filter_test.go
Normal file
217
proxy/internal/proxy/discovery_filter_test.go
Normal file
@@ -0,0 +1,217 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// jsonListingResponse builds a 200 model-listing response with the given
|
||||
// body, as an upstream would return it.
|
||||
func jsonListingResponse(body string) *http.Response {
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
ContentLength: int64(len(body)),
|
||||
}
|
||||
resp.Header.Set("Content-Type", "application/json")
|
||||
return resp
|
||||
}
|
||||
|
||||
// listedIDs runs the filter and returns the ids left in the response.
|
||||
func listedIDs(t *testing.T, allowed []string, body string) []string {
|
||||
t.Helper()
|
||||
resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body, replaced by the filter
|
||||
require.NoError(t, modelDiscoveryFilter(allowed, nil)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter
|
||||
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
|
||||
var doc struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(raw, &doc), "filtered body must stay valid JSON")
|
||||
|
||||
ids := make([]string, 0, len(doc.Data))
|
||||
for _, entry := range doc.Data {
|
||||
ids = append(ids, entry.ID)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
// TestModelDiscoveryFilter_KeepsOnlyAuthorisedModels covers the picker a
|
||||
// developer sees: an unfiltered upstream list offers every model the shared
|
||||
// key can reach, and each one the policy excludes is a request the chain
|
||||
// denies a moment later.
|
||||
func TestModelDiscoveryFilter_KeepsOnlyAuthorisedModels(t *testing.T) {
|
||||
ids := listedIDs(t, []string{"claude-sonnet-5", "claude-haiku-4-5"}, `{
|
||||
"data": [
|
||||
{"id": "claude-opus-5", "display_name": "Claude Opus 5"},
|
||||
{"id": "claude-sonnet-5", "display_name": "Claude Sonnet 5"},
|
||||
{"id": "claude-haiku-4-5"}
|
||||
],
|
||||
"has_more": false
|
||||
}`)
|
||||
|
||||
assert.Equal(t, []string{"claude-sonnet-5", "claude-haiku-4-5"}, ids,
|
||||
"only the models the route authorises may reach the picker")
|
||||
}
|
||||
|
||||
// TestModelDiscoveryFilter_MatchesDatedAndPrefixedIDs pins the two id forms
|
||||
// a gateway returns for a model the operator registered plainly.
|
||||
func TestModelDiscoveryFilter_MatchesDatedAndPrefixedIDs(t *testing.T) {
|
||||
ids := listedIDs(t, []string{"claude-sonnet-4-5", "anthropic.claude-opus-5"}, `{
|
||||
"data": [
|
||||
{"id": "claude-sonnet-4-5-20250929"},
|
||||
{"id": "bedrock/anthropic.claude-opus-5"},
|
||||
{"id": "gpt-4o"}
|
||||
]
|
||||
}`)
|
||||
|
||||
assert.Equal(t, []string{"claude-sonnet-4-5-20250929", "bedrock/anthropic.claude-opus-5"}, ids,
|
||||
"a dated or provider-prefixed id must match its registered form")
|
||||
}
|
||||
|
||||
// TestModelDiscoveryFilter_PreservesEnvelopeFields guards the rest of the
|
||||
// document: clients read paging fields alongside data.
|
||||
func TestModelDiscoveryFilter_PreservesEnvelopeFields(t *testing.T) {
|
||||
resp := jsonListingResponse(`{"data":[{"id":"claude-sonnet-5"}],"has_more":true,"first_id":"x"}`) //nolint:bodyclose // in-memory body, replaced by the filter
|
||||
require.NoError(t, modelDiscoveryFilter([]string{"claude-sonnet-5"}, nil)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter
|
||||
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
|
||||
var doc map[string]any
|
||||
require.NoError(t, json.Unmarshal(raw, &doc))
|
||||
assert.Equal(t, true, doc["has_more"], "paging fields must survive the rewrite")
|
||||
assert.Equal(t, "x", doc["first_id"])
|
||||
assert.Equal(t, strconv.Itoa(len(raw)), resp.Header.Get("Content-Length"),
|
||||
"Content-Length must match the rewritten body")
|
||||
}
|
||||
|
||||
// TestModelDiscoveryFilter_PassesThroughUnfilterable covers the responses
|
||||
// the filter must not touch: a compressed body it cannot parse, a non-JSON
|
||||
// body, an error status, and a document with no data array.
|
||||
func TestModelDiscoveryFilter_PassesThroughUnfilterable(t *testing.T) {
|
||||
cases := map[string]func() *http.Response{
|
||||
"compressed": func() *http.Response {
|
||||
resp := jsonListingResponse(`{"data":[{"id":"gpt-4o"}]}`)
|
||||
resp.Header.Set("Content-Encoding", "gzip")
|
||||
return resp
|
||||
},
|
||||
"not json": func() *http.Response {
|
||||
resp := jsonListingResponse(`{"data":[{"id":"gpt-4o"}]}`)
|
||||
resp.Header.Set("Content-Type", "text/html")
|
||||
return resp
|
||||
},
|
||||
"error status": func() *http.Response {
|
||||
resp := jsonListingResponse(`{"data":[{"id":"gpt-4o"}]}`)
|
||||
resp.StatusCode = http.StatusInternalServerError
|
||||
return resp
|
||||
},
|
||||
"no data array": func() *http.Response {
|
||||
return jsonListingResponse(`{"object":"list"}`)
|
||||
},
|
||||
}
|
||||
|
||||
for name, build := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
resp := build() //nolint:bodyclose // in-memory body, replaced by the filter
|
||||
original, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
resp.Body = io.NopCloser(bytes.NewReader(original))
|
||||
|
||||
require.NoError(t, modelDiscoveryFilter([]string{"claude-sonnet-5"}, nil)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter
|
||||
|
||||
got, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, string(original), string(got), "an unfilterable response must reach the client unchanged")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelDiscoveryFilter_RunsNextHook pins that an existing
|
||||
// ModifyResponse hook still runs after filtering.
|
||||
func TestModelDiscoveryFilter_RunsNextHook(t *testing.T) {
|
||||
called := false
|
||||
next := func(*http.Response) error {
|
||||
called = true
|
||||
return nil
|
||||
}
|
||||
|
||||
resp := jsonListingResponse(`{"data":[{"id":"claude-sonnet-5"}]}`) //nolint:bodyclose // in-memory body, replaced by the filter
|
||||
require.NoError(t, modelDiscoveryFilter([]string{"claude-sonnet-5"}, next)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter
|
||||
assert.True(t, called, "the chained hook must still run")
|
||||
}
|
||||
|
||||
// TestModelDiscoveryFilter_KeepsSlashBearingIDs covers self-hosted backends
|
||||
// whose model ids carry a slash of their own. Treating the slash as a
|
||||
// gateway prefix and keeping only the tail dropped every such model from
|
||||
// the picker even though the policy named it exactly.
|
||||
func TestModelDiscoveryFilter_KeepsSlashBearingIDs(t *testing.T) {
|
||||
ids := listedIDs(t, []string{"Qwen/Qwen2.5-0.5B-Instruct"}, `{
|
||||
"object": "list",
|
||||
"data": [
|
||||
{"id": "Qwen/Qwen2.5-0.5B-Instruct"},
|
||||
{"id": "Qwen/Qwen2.5-7B-Instruct"}
|
||||
]
|
||||
}`)
|
||||
|
||||
assert.Equal(t, []string{"Qwen/Qwen2.5-0.5B-Instruct"}, ids,
|
||||
"a slash inside the model id is part of the id, not a provider prefix")
|
||||
}
|
||||
|
||||
// TestModelDiscoveryFilter_ForwardsOversizedBodyIntact covers a listing past
|
||||
// the buffering cap. The filter reads one byte beyond the cap to detect the
|
||||
// size; forwarding only what it read would hand the client a body truncated
|
||||
// at exactly 1 MiB — valid-looking, short, and unparseable as JSON. The bytes
|
||||
// already read must be spliced back in front of the unread remainder so the
|
||||
// response reaches the client exactly as the upstream sent it.
|
||||
func TestModelDiscoveryFilter_ForwardsOversizedBodyIntact(t *testing.T) {
|
||||
// A well-formed listing whose single entry pads the body past the cap.
|
||||
padding := strings.Repeat("x", maxDiscoveryBodyBytes)
|
||||
body := `{"object":"list","data":[{"id":"gpt-4o","note":"` + padding + `"}]}`
|
||||
require.Greater(t, len(body), maxDiscoveryBodyBytes+1,
|
||||
"the fixture must exceed the cap by more than the one-byte probe")
|
||||
|
||||
resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body
|
||||
require.NoError(t, modelDiscoveryFilter(nil, nil)(resp)) //nolint:bodyclose // in-memory body
|
||||
|
||||
got, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, len(body), len(got),
|
||||
"an oversized listing must reach the client whole, not truncated at the cap")
|
||||
assert.Equal(t, body, string(got), "the forwarded bytes must be the upstream's own")
|
||||
|
||||
var doc map[string]json.RawMessage
|
||||
assert.NoError(t, json.Unmarshal(got, &doc),
|
||||
"the forwarded body must still parse as JSON")
|
||||
}
|
||||
|
||||
// TestModelDiscoveryFilter_OversizedBodyKeepsUpstreamHeaders pins that the
|
||||
// oversized path leaves the response metadata alone. Rewriting Content-Length
|
||||
// to the truncated prefix is what made the corruption invisible to the client
|
||||
// until it tried to parse.
|
||||
func TestModelDiscoveryFilter_OversizedBodyKeepsUpstreamHeaders(t *testing.T) {
|
||||
padding := strings.Repeat("x", maxDiscoveryBodyBytes)
|
||||
body := `{"object":"list","data":[{"id":"gpt-4o","note":"` + padding + `"}]}`
|
||||
|
||||
resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body
|
||||
resp.Header.Set("Content-Length", strconv.Itoa(len(body)))
|
||||
require.NoError(t, modelDiscoveryFilter(nil, nil)(resp)) //nolint:bodyclose // in-memory body
|
||||
|
||||
assert.Equal(t, int64(len(body)), resp.ContentLength,
|
||||
"ContentLength must keep describing the body the client receives")
|
||||
assert.Equal(t, strconv.Itoa(len(body)), resp.Header.Get("Content-Length"),
|
||||
"the Content-Length header must not be rewritten to the truncated prefix")
|
||||
}
|
||||
@@ -363,6 +363,9 @@ func (p *ReverseProxy) forwardUpstream(respWriter http.ResponseWriter, r *http.R
|
||||
if result.rewriteRedirects {
|
||||
rp.ModifyResponse = p.rewriteLocationFunc(effectiveURL, rewriteMatchedPath, r) //nolint:bodyclose
|
||||
}
|
||||
if upstreamRewrite != nil && len(upstreamRewrite.DiscoveryModels) > 0 {
|
||||
rp.ModifyResponse = modelDiscoveryFilter(upstreamRewrite.DiscoveryModels, rp.ModifyResponse) //nolint:bodyclose // the hook replaces the body and closes the original
|
||||
}
|
||||
rp.ServeHTTP(respWriter, r.WithContext(ctx))
|
||||
}
|
||||
|
||||
|
||||
@@ -46,6 +46,27 @@ func NormalizeBedrockModel(modelID string) string {
|
||||
return bedrockVersionSuffix.ReplaceAllString(m, "")
|
||||
}
|
||||
|
||||
// anthropicDatedModel matches a Claude model id carrying the trailing
|
||||
// "-YYYYMMDD" release-date suffix Anthropic appends to a pinned release,
|
||||
// capturing the id without it. The "claude" anchor is load-bearing: pricing
|
||||
// looks every model up through this helper regardless of surface, and an
|
||||
// operator may register a custom id with any shape at all, so an unanchored
|
||||
// "-\d{8}$" would let "internal-llm-20250101" silently inherit the rate
|
||||
// registered for "internal-llm". The anchor also covers the vendor-prefixed
|
||||
// forms ("anthropic.claude-...", "us.anthropic.claude-...").
|
||||
var anthropicDatedModel = regexp.MustCompile(`(?i)^(.*claude.*)-\d{8}$`)
|
||||
|
||||
// NormalizeAnthropicModel strips the trailing release-date suffix from a
|
||||
// Claude model id, e.g. "claude-sonnet-4-5-20250929" -> "claude-sonnet-4-5",
|
||||
// so a dated id a client pins matches the undated one the operator
|
||||
// registered. Ids that are not Claude-family are returned untouched.
|
||||
// Callers try the verbatim id first and fall back to this, so two dated
|
||||
// releases of the same family stay distinct wherever both are registered
|
||||
// explicitly.
|
||||
func NormalizeAnthropicModel(modelID string) string {
|
||||
return anthropicDatedModel.ReplaceAllString(modelID, "$1")
|
||||
}
|
||||
|
||||
// NormalizeVertexModel strips the "@version" suffix from a Vertex AI model id
|
||||
// (e.g. "claude-sonnet-4-5@20250929" -> "claude-sonnet-4-5") so it matches
|
||||
// the catalog/pricing key. Vertex publisher models are priced under their
|
||||
|
||||
@@ -34,3 +34,29 @@ func TestNormalizeVertexModel(t *testing.T) {
|
||||
require.Equal(t, want, NormalizeVertexModel(in), "normalize %q", in)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeAnthropicModel(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"claude-sonnet-4-5-20250929": "claude-sonnet-4-5",
|
||||
"claude-3-5-haiku-20241022": "claude-3-5-haiku",
|
||||
"claude-sonnet-5": "claude-sonnet-5",
|
||||
"claude-opus-4-8": "claude-opus-4-8",
|
||||
"anthropic.claude-haiku-4-5": "anthropic.claude-haiku-4-5",
|
||||
"anthropic.claude-sonnet-4-5-20250929": "anthropic.claude-sonnet-4-5",
|
||||
"us.anthropic.claude-opus-4-8-20250101": "us.anthropic.claude-opus-4-8",
|
||||
// Non-Claude ids must survive untouched even when they end in eight
|
||||
// consecutive digits: an operator can register a custom model under
|
||||
// any id, and pricing looks every one of them up through this helper.
|
||||
"gpt-4o": "gpt-4o",
|
||||
"gpt-4o-2024-08-06": "gpt-4o-2024-08-06",
|
||||
"gpt-4o-20240806": "gpt-4o-20240806",
|
||||
"internal-llm-20250101": "internal-llm-20250101",
|
||||
"deepseek-r1-20250120": "deepseek-r1-20250120",
|
||||
"Qwen/Qwen2.5-20250101": "Qwen/Qwen2.5-20250101",
|
||||
"gemini-2-5-pro-20250101": "gemini-2-5-pro-20250101",
|
||||
"": "",
|
||||
}
|
||||
for in, want := range cases {
|
||||
require.Equal(t, want, NormalizeAnthropicModel(in), "normalize %q", in)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,8 +6,6 @@ import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
@@ -338,35 +336,25 @@ func (a *AgentNetworkAPI) DeleteBudgetRule(ctx context.Context, ruleID string) e
|
||||
// to an APIError matchable via IsNotFound rather than fabricating defaults
|
||||
// the server never stated.
|
||||
func (a *AgentNetworkAPI) GetSettings(ctx context.Context) (*api.AgentNetworkSettings, error) {
|
||||
settings, _, err := a.GetSettingsWithETag(ctx)
|
||||
return settings, err
|
||||
}
|
||||
|
||||
// GetSettingsWithETag is GetSettings, additionally returning the entity-tag
|
||||
// the server derived for the settings it returned. Hand that validator to
|
||||
// UpdateSettingsIfMatch or DeleteSettingsIfMatch to make the write conditional
|
||||
// on nothing having changed in between — the read-modify-write cycle that
|
||||
// otherwise silently reverts a concurrent change.
|
||||
func (a *AgentNetworkAPI) GetSettingsWithETag(ctx context.Context) (*api.AgentNetworkSettings, string, error) {
|
||||
resp, err := a.c.NewRequest(ctx, "GET", "/api/agent-network/settings", nil, nil)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
return nil, err
|
||||
}
|
||||
if resp.Body != nil {
|
||||
defer resp.Body.Close()
|
||||
}
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
return nil, err
|
||||
}
|
||||
if trimmed := bytes.TrimSpace(body); len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) {
|
||||
return nil, "", &APIError{StatusCode: http.StatusNotFound, Message: "agent network settings not found"}
|
||||
return nil, &APIError{StatusCode: http.StatusNotFound, Message: "agent network settings not found"}
|
||||
}
|
||||
var ret api.AgentNetworkSettings
|
||||
if err := json.Unmarshal(body, &ret); err != nil {
|
||||
return nil, "", err
|
||||
return nil, err
|
||||
}
|
||||
return &ret, etagFrom(resp), nil
|
||||
return &ret, nil
|
||||
}
|
||||
|
||||
// CreateSettings bootstraps the account's Agent Network settings row,
|
||||
@@ -375,30 +363,19 @@ func (a *AgentNetworkAPI) GetSettingsWithETag(ctx context.Context) (*api.AgentNe
|
||||
// request.Endpoint (self-addressed dedicated endpoint, claimed verbatim) must
|
||||
// be set. Returns a conflict when the account already has a settings row.
|
||||
func (a *AgentNetworkAPI) CreateSettings(ctx context.Context, request api.PostApiAgentNetworkSettingsJSONRequestBody) (*api.AgentNetworkSettings, error) {
|
||||
settings, _, err := a.CreateSettingsWithETag(ctx, request)
|
||||
return settings, err
|
||||
}
|
||||
|
||||
// CreateSettingsWithETag is CreateSettings, additionally returning the
|
||||
// entity-tag of the row it bootstrapped, so a client can follow the bootstrap
|
||||
// with a conditional write without an intervening read.
|
||||
func (a *AgentNetworkAPI) CreateSettingsWithETag(ctx context.Context, request api.PostApiAgentNetworkSettingsJSONRequestBody) (*api.AgentNetworkSettings, string, error) {
|
||||
requestBytes, err := json.Marshal(request)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
return nil, err
|
||||
}
|
||||
resp, err := a.c.NewRequest(ctx, "POST", "/api/agent-network/settings", bytes.NewReader(requestBytes), nil)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
return nil, err
|
||||
}
|
||||
if resp.Body != nil {
|
||||
defer resp.Body.Close()
|
||||
}
|
||||
ret, err := parseResponse[api.AgentNetworkSettings](resp)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return &ret, etagFrom(resp), nil
|
||||
return &ret, err
|
||||
}
|
||||
|
||||
// UpdateSettings updates the account's Agent Network settings; the request
|
||||
@@ -408,35 +385,19 @@ func (a *AgentNetworkAPI) CreateSettingsWithETag(ctx context.Context, request ap
|
||||
// a request carrying different values is rejected. Returns not-found until
|
||||
// the account is bootstrapped.
|
||||
func (a *AgentNetworkAPI) UpdateSettings(ctx context.Context, request api.PutApiAgentNetworkSettingsJSONRequestBody) (*api.AgentNetworkSettings, error) {
|
||||
settings, _, err := a.UpdateSettingsIfMatch(ctx, request, "")
|
||||
return settings, err
|
||||
}
|
||||
|
||||
// UpdateSettingsIfMatch is UpdateSettings made conditional on etag — the
|
||||
// validator from an earlier read — still being current, and returns the
|
||||
// validator of the row it wrote. This is what closes the read-modify-write
|
||||
// window: a settings change made between the read and this write makes the
|
||||
// request fail with a precondition-failed APIError instead of reverting it.
|
||||
//
|
||||
// An empty etag sends no precondition and updates unconditionally, which is
|
||||
// what UpdateSettings does.
|
||||
func (a *AgentNetworkAPI) UpdateSettingsIfMatch(ctx context.Context, request api.PutApiAgentNetworkSettingsJSONRequestBody, etag string) (*api.AgentNetworkSettings, string, error) {
|
||||
requestBytes, err := json.Marshal(request)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
return nil, err
|
||||
}
|
||||
resp, err := a.c.newRequest(ctx, "PUT", "/api/agent-network/settings", bytes.NewReader(requestBytes), nil, ifMatchHeader(etag))
|
||||
resp, err := a.c.NewRequest(ctx, "PUT", "/api/agent-network/settings", bytes.NewReader(requestBytes), nil)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
return nil, err
|
||||
}
|
||||
if resp.Body != nil {
|
||||
defer resp.Body.Close()
|
||||
}
|
||||
ret, err := parseResponse[api.AgentNetworkSettings](resp)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return &ret, etagFrom(resp), nil
|
||||
return &ret, err
|
||||
}
|
||||
|
||||
// DeleteSettings deletes the account's Agent Network settings row, releasing
|
||||
@@ -444,20 +405,7 @@ func (a *AgentNetworkAPI) UpdateSettingsIfMatch(ctx context.Context, request api
|
||||
// exists for the account or while a proxy is actively serving the endpoint.
|
||||
// Bootstrapping again afterwards allocates a new endpoint.
|
||||
func (a *AgentNetworkAPI) DeleteSettings(ctx context.Context) error {
|
||||
return a.DeleteSettingsIfMatch(ctx, "")
|
||||
}
|
||||
|
||||
// DeleteSettingsIfMatch is DeleteSettings made conditional on etag — the
|
||||
// validator from an earlier read — still being current. Sending it matters
|
||||
// more here than on update: the server's other two refusals are about state
|
||||
// (no providers, no serving proxy), so this is the only thing that stops a
|
||||
// client working from an old read of one row from releasing the endpoint of
|
||||
// the row that replaced it.
|
||||
//
|
||||
// An empty etag sends no precondition and deletes unconditionally, which is
|
||||
// what DeleteSettings does.
|
||||
func (a *AgentNetworkAPI) DeleteSettingsIfMatch(ctx context.Context, etag string) error {
|
||||
resp, err := a.c.newRequest(ctx, "DELETE", "/api/agent-network/settings", nil, nil, ifMatchHeader(etag))
|
||||
resp, err := a.c.NewRequest(ctx, "DELETE", "/api/agent-network/settings", nil, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -467,20 +415,3 @@ func (a *AgentNetworkAPI) DeleteSettingsIfMatch(ctx context.Context, etag string
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// etagFrom returns the bare validator from a response, with the transport's
|
||||
// quoting stripped so a caller can hand it straight back to an If-Match
|
||||
// parameter without knowing the wire syntax.
|
||||
func etagFrom(resp *http.Response) string {
|
||||
return strings.Trim(resp.Header.Get("ETag"), `"`)
|
||||
}
|
||||
|
||||
// ifMatchHeader renders the precondition headers for a bare validator,
|
||||
// re-applying the quoting etagFrom stripped. An empty validator yields no
|
||||
// headers at all — an unconditional request.
|
||||
func ifMatchHeader(etag string) map[string]string {
|
||||
if etag == "" {
|
||||
return nil
|
||||
}
|
||||
return map[string]string{"If-Match": strconv.Quote(etag)}
|
||||
}
|
||||
|
||||
@@ -559,123 +559,3 @@ func TestAgentNetwork_DeleteSettings_Guarded(t *testing.T) {
|
||||
assert.Contains(t, err.Error(), "cannot be deleted")
|
||||
})
|
||||
}
|
||||
|
||||
// TestAgentNetwork_GetSettings_ETag pins that the validator surfaces to the
|
||||
// caller with the transport's quoting stripped, so it can be handed straight
|
||||
// back to a conditional write without the caller knowing the wire syntax.
|
||||
func TestAgentNetwork_GetSettings_ETag(t *testing.T) {
|
||||
withMockClient(func(c *rest.Client, mux *http.ServeMux) {
|
||||
mux.HandleFunc("/api/agent-network/settings", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("ETag", `"9f86d081884c7d65"`)
|
||||
retBytes, _ := json.Marshal(testAgentNetworkSettings)
|
||||
_, err := w.Write(retBytes)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
ret, etag, err := c.AgentNetwork.GetSettingsWithETag(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, testAgentNetworkSettings, *ret)
|
||||
assert.Equal(t, "9f86d081884c7d65", etag, "the validator must arrive unquoted")
|
||||
})
|
||||
}
|
||||
|
||||
// TestAgentNetwork_UpdateSettings_IfMatch covers the round trip that makes the
|
||||
// whole feature usable: a validator taken from a read goes back out quoted on
|
||||
// the write, and the write's own validator comes back for the next one.
|
||||
func TestAgentNetwork_UpdateSettings_IfMatch(t *testing.T) {
|
||||
withMockClient(func(c *rest.Client, mux *http.ServeMux) {
|
||||
mux.HandleFunc("/api/agent-network/settings", func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, `"9f86d081884c7d65"`, r.Header.Get("If-Match"),
|
||||
"the precondition must go out quoted as a strong entity-tag")
|
||||
w.Header().Set("ETag", `"0011223344556677"`)
|
||||
retBytes, _ := json.Marshal(testAgentNetworkSettings)
|
||||
_, err := w.Write(retBytes)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
_, etag, err := c.AgentNetwork.UpdateSettingsIfMatch(context.Background(), api.PutApiAgentNetworkSettingsJSONRequestBody{
|
||||
Endpoint: "brave-otter.eu.proxy.netbird.io",
|
||||
ProxyAddress: "eu.proxy.netbird.io",
|
||||
EnableLogCollection: true,
|
||||
}, "9f86d081884c7d65")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "0011223344556677", etag, "the write must return the new validator")
|
||||
})
|
||||
}
|
||||
|
||||
// TestAgentNetwork_UpdateSettings_NoPrecondition pins the back-compatible
|
||||
// path: the plain method sends no If-Match at all, rather than an empty or
|
||||
// wildcard one, so it stays the unconditional update it has always been.
|
||||
func TestAgentNetwork_UpdateSettings_NoPrecondition(t *testing.T) {
|
||||
withMockClient(func(c *rest.Client, mux *http.ServeMux) {
|
||||
mux.HandleFunc("/api/agent-network/settings", func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Empty(t, r.Header.Values("If-Match"), "an unconditional update must send no precondition")
|
||||
retBytes, _ := json.Marshal(testAgentNetworkSettings)
|
||||
_, err := w.Write(retBytes)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
_, err := c.AgentNetwork.UpdateSettings(context.Background(), api.PutApiAgentNetworkSettingsJSONRequestBody{
|
||||
Endpoint: "brave-otter.eu.proxy.netbird.io",
|
||||
ProxyAddress: "eu.proxy.netbird.io",
|
||||
EnableLogCollection: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
// TestAgentNetwork_UpdateSettings_StalePrecondition pins how a refused write
|
||||
// reaches the caller: as an APIError a client can recognise as staleness and
|
||||
// answer by reading again, rather than as an opaque failure.
|
||||
func TestAgentNetwork_UpdateSettings_StalePrecondition(t *testing.T) {
|
||||
withMockClient(func(c *rest.Client, mux *http.ServeMux) {
|
||||
mux.HandleFunc("/api/agent-network/settings", func(w http.ResponseWriter, r *http.Request) {
|
||||
retBytes, _ := json.Marshal(util.ErrorResponse{Message: "if-match precondition failed: the settings have changed since they were read; get them again and retry", Code: 412})
|
||||
w.WriteHeader(412)
|
||||
_, err := w.Write(retBytes)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
_, _, err := c.AgentNetwork.UpdateSettingsIfMatch(context.Background(), api.PutApiAgentNetworkSettingsJSONRequestBody{
|
||||
Endpoint: "brave-otter.eu.proxy.netbird.io",
|
||||
ProxyAddress: "eu.proxy.netbird.io",
|
||||
EnableLogCollection: true,
|
||||
}, "9f86d081884c7d65")
|
||||
require.Error(t, err)
|
||||
assert.True(t, rest.IsPreconditionFailed(err), "a refused precondition must be recognisable as one")
|
||||
assert.False(t, rest.IsNotFound(err), "it must not be confused with an unbootstrapped account")
|
||||
})
|
||||
}
|
||||
|
||||
// TestAgentNetwork_CreateSettings_ETag pins that the bootstrap hands back a
|
||||
// validator, which is what lets a client follow it with a conditional write
|
||||
// without an intervening read.
|
||||
func TestAgentNetwork_CreateSettings_ETag(t *testing.T) {
|
||||
withMockClient(func(c *rest.Client, mux *http.ServeMux) {
|
||||
mux.HandleFunc("/api/agent-network/settings", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("ETag", `"9f86d081884c7d65"`)
|
||||
retBytes, _ := json.Marshal(testAgentNetworkSettings)
|
||||
_, err := w.Write(retBytes)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
_, etag, err := c.AgentNetwork.CreateSettingsWithETag(context.Background(), api.PostApiAgentNetworkSettingsJSONRequestBody{
|
||||
ProxyAddress: ptr("eu.proxy.netbird.io"),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "9f86d081884c7d65", etag, "the bootstrap must return a validator")
|
||||
})
|
||||
}
|
||||
|
||||
// TestAgentNetwork_DeleteSettings_IfMatch covers the conditional delete on the
|
||||
// wire, and that the plain method still sends nothing.
|
||||
func TestAgentNetwork_DeleteSettings_IfMatch(t *testing.T) {
|
||||
withMockClient(func(c *rest.Client, mux *http.ServeMux) {
|
||||
var seen []string
|
||||
mux.HandleFunc("/api/agent-network/settings", func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "DELETE", r.Method)
|
||||
seen = append(seen, r.Header.Get("If-Match"))
|
||||
_, err := w.Write([]byte("{}"))
|
||||
require.NoError(t, err)
|
||||
})
|
||||
require.NoError(t, c.AgentNetwork.DeleteSettingsIfMatch(context.Background(), "9f86d081884c7d65"))
|
||||
require.NoError(t, c.AgentNetwork.DeleteSettings(context.Background()))
|
||||
assert.Equal(t, []string{`"9f86d081884c7d65"`, ""}, seen,
|
||||
"the conditional delete must carry the quoted validator and the plain one must carry nothing")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -31,19 +31,6 @@ func IsNotFound(err error) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// IsPreconditionFailed returns true if the error represents a 412 Precondition
|
||||
// Failed response — an If-Match the server refused, or an endpoint's own
|
||||
// precondition. A caller that sent a conditional request can use this to tell
|
||||
// "someone else changed it, read again and retry" apart from a real failure;
|
||||
// the message distinguishes it from an endpoint's other 412s.
|
||||
func IsPreconditionFailed(err error) bool {
|
||||
var apiErr *APIError
|
||||
if ok := errors.As(err, &apiErr); ok {
|
||||
return apiErr.StatusCode == http.StatusPreconditionFailed
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Client Management service HTTP REST API Client
|
||||
type Client struct {
|
||||
managementURL string
|
||||
@@ -231,12 +218,6 @@ func (c *Client) initialize() {
|
||||
|
||||
// NewRequest creates and executes new management API request
|
||||
func (c *Client) NewRequest(ctx context.Context, method, path string, body io.Reader, query map[string]string) (*http.Response, error) {
|
||||
return c.newRequest(ctx, method, path, body, query, nil)
|
||||
}
|
||||
|
||||
// newRequest is NewRequest with request headers, for the endpoints whose
|
||||
// contract includes one — conditional requests carrying If-Match.
|
||||
func (c *Client) newRequest(ctx context.Context, method, path string, body io.Reader, query, headers map[string]string) (*http.Response, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, method, c.managementURL+path, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -250,9 +231,6 @@ func (c *Client) newRequest(ctx context.Context, method, path string, body io.Re
|
||||
if c.userAgent != "" {
|
||||
req.Header.Set("User-Agent", c.userAgent)
|
||||
}
|
||||
for name, value := range headers {
|
||||
req.Header.Set(name, value)
|
||||
}
|
||||
|
||||
if len(query) != 0 {
|
||||
q := req.URL.Query()
|
||||
|
||||
@@ -6438,15 +6438,6 @@ components:
|
||||
schema:
|
||||
type: string
|
||||
example: cot7r4n3l3vh3qj4qveg
|
||||
ETag:
|
||||
description: |
|
||||
Strong entity-tag identifying the returned representation. Send it back
|
||||
in `If-Match` on a subsequent write to make that write conditional, so
|
||||
a change made between the read and the write is refused with `412`
|
||||
rather than silently overwritten.
|
||||
schema:
|
||||
type: string
|
||||
example: '"9f86d081884c7d65"'
|
||||
securitySchemes:
|
||||
BearerAuth:
|
||||
type: http
|
||||
@@ -13742,9 +13733,6 @@ paths:
|
||||
responses:
|
||||
'200':
|
||||
description: Agent Network settings for the account
|
||||
headers:
|
||||
ETag:
|
||||
$ref: '#/components/headers/ETag'
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
@@ -13772,9 +13760,6 @@ paths:
|
||||
responses:
|
||||
'200':
|
||||
description: The freshly bootstrapped Agent Network settings
|
||||
headers:
|
||||
ETag:
|
||||
$ref: '#/components/headers/ETag'
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
@@ -13793,25 +13778,11 @@ paths:
|
||||
"$ref": "#/components/responses/internal_error"
|
||||
put:
|
||||
summary: Update Agent Network settings
|
||||
description: Updates the account-level Agent Network settings; the request carries every field, replacing the mutable ones (collection toggles and retention). Returns 404 when the account has no settings row yet — bootstrap it with POST first. The endpoint and proxy address are assigned at bootstrap and immutable; the request must carry them unchanged, and a request carrying different values is rejected. Supply `If-Match` to make the update conditional; without it the update is unconditional and the last write wins.
|
||||
description: Updates the account-level Agent Network settings; the request carries every field, replacing the mutable ones (collection toggles and retention). Returns 404 when the account has no settings row yet — bootstrap it with POST first. The endpoint and proxy address are assigned at bootstrap and immutable; the request must carry them unchanged, and a request carrying different values is rejected.
|
||||
tags: [ Agent Network ]
|
||||
security:
|
||||
- BearerAuth: [ ]
|
||||
- TokenAuth: [ ]
|
||||
parameters:
|
||||
- name: If-Match
|
||||
in: header
|
||||
required: false
|
||||
description: |
|
||||
Makes the update conditional on the settings not having changed since
|
||||
they were read. Send the `ETag` from an earlier `GET`, `POST` or `PUT`,
|
||||
or `*` to require only that a settings row exists. The precondition is
|
||||
evaluated against the stored row inside the update's own transaction,
|
||||
so two clients starting from the same `ETag` cannot both succeed.
|
||||
Omitting the header leaves the update unconditional.
|
||||
schema:
|
||||
type: string
|
||||
example: '"9f86d081884c7d65"'
|
||||
requestBody:
|
||||
description: Settings update request
|
||||
content:
|
||||
@@ -13821,9 +13792,6 @@ paths:
|
||||
responses:
|
||||
'200':
|
||||
description: Updated Agent Network settings
|
||||
headers:
|
||||
ETag:
|
||||
$ref: '#/components/headers/ETag'
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
@@ -13836,34 +13804,17 @@ paths:
|
||||
"$ref": "#/components/responses/forbidden"
|
||||
'404':
|
||||
"$ref": "#/components/responses/not_found"
|
||||
'412':
|
||||
description: The `If-Match` precondition failed — the settings changed since they were read. The stored settings are unmodified; read them again and retry.
|
||||
content: { }
|
||||
'422':
|
||||
"$ref": "#/components/responses/validation_failed"
|
||||
'500':
|
||||
"$ref": "#/components/responses/internal_error"
|
||||
delete:
|
||||
summary: Delete Agent Network settings
|
||||
description: Deletes the account's Agent Network settings row, releasing the endpoint. Guarded — the delete is refused with 412 while any Agent Network provider exists for the account or while a proxy is actively serving the endpoint. Bootstrapping again after a delete allocates a new endpoint; the released hostname is not reserved. Supply `If-Match` to make the delete conditional, which is worth doing here even more than on update — the other two guards are about state rather than staleness, so nothing else stops a client from deleting a row that was replaced since it read one.
|
||||
description: Deletes the account's Agent Network settings row, releasing the endpoint. Guarded — the delete is refused with 412 while any Agent Network provider exists for the account or while a proxy is actively serving the endpoint. Bootstrapping again after a delete allocates a new endpoint; the released hostname is not reserved.
|
||||
tags: [ Agent Network ]
|
||||
security:
|
||||
- BearerAuth: [ ]
|
||||
- TokenAuth: [ ]
|
||||
parameters:
|
||||
- name: If-Match
|
||||
in: header
|
||||
required: false
|
||||
description: |
|
||||
Makes the delete conditional on the settings not having changed since
|
||||
they were read. Send the `ETag` from an earlier `GET`, `POST` or `PUT`,
|
||||
or `*` to require only that a settings row exists. The precondition is
|
||||
evaluated inside the delete's own transaction, ahead of the provider
|
||||
and serving-proxy guards. Omitting the header leaves the delete
|
||||
unconditional.
|
||||
schema:
|
||||
type: string
|
||||
example: '"9f86d081884c7d65"'
|
||||
responses:
|
||||
'200':
|
||||
description: Settings deleted
|
||||
@@ -13874,7 +13825,7 @@ paths:
|
||||
'404':
|
||||
"$ref": "#/components/responses/not_found"
|
||||
'412':
|
||||
description: Delete refused — the `If-Match` precondition failed, or Agent Network providers still exist for the account, or a proxy is actively serving the endpoint. The stored settings are unmodified in every case; the response message distinguishes them.
|
||||
description: Delete refused — Agent Network providers still exist for the account, or a proxy is actively serving the endpoint
|
||||
content: { }
|
||||
'500':
|
||||
"$ref": "#/components/responses/internal_error"
|
||||
|
||||
@@ -5939,28 +5939,6 @@ type GetApiAgentNetworkAccessLogsParamsSortBy string
|
||||
// GetApiAgentNetworkAccessLogsParamsSortOrder defines parameters for GetApiAgentNetworkAccessLogs.
|
||||
type GetApiAgentNetworkAccessLogsParamsSortOrder string
|
||||
|
||||
// DeleteApiAgentNetworkSettingsParams defines parameters for DeleteApiAgentNetworkSettings.
|
||||
type DeleteApiAgentNetworkSettingsParams struct {
|
||||
// IfMatch Makes the delete conditional on the settings not having changed since
|
||||
// they were read. Send the `ETag` from an earlier `GET`, `POST` or `PUT`,
|
||||
// or `*` to require only that a settings row exists. The precondition is
|
||||
// evaluated inside the delete's own transaction, ahead of the provider
|
||||
// and serving-proxy guards. Omitting the header leaves the delete
|
||||
// unconditional.
|
||||
IfMatch *string `json:"If-Match,omitempty"`
|
||||
}
|
||||
|
||||
// PutApiAgentNetworkSettingsParams defines parameters for PutApiAgentNetworkSettings.
|
||||
type PutApiAgentNetworkSettingsParams struct {
|
||||
// IfMatch Makes the update conditional on the settings not having changed since
|
||||
// they were read. Send the `ETag` from an earlier `GET`, `POST` or `PUT`,
|
||||
// or `*` to require only that a settings row exists. The precondition is
|
||||
// evaluated against the stored row inside the update's own transaction,
|
||||
// so two clients starting from the same `ETag` cannot both succeed.
|
||||
// Omitting the header leaves the update unconditional.
|
||||
IfMatch *string `json:"If-Match,omitempty"`
|
||||
}
|
||||
|
||||
// GetApiAgentNetworkUsageOverviewParams defines parameters for GetApiAgentNetworkUsageOverview.
|
||||
type GetApiAgentNetworkUsageOverviewParams struct {
|
||||
// Granularity Time bucket width. Defaults to day.
|
||||
|
||||
@@ -1,103 +0,0 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
etagHeader = "ETag"
|
||||
ifMatchHeader = "If-Match"
|
||||
|
||||
// matchAny is the If-Match value that matches any current representation
|
||||
// of the resource (RFC 9110 §13.1.1).
|
||||
matchAny = "*"
|
||||
|
||||
// weakPrefix marks a weak validator. If-Match is defined in terms of the
|
||||
// strong comparison function, under which a weak validator never matches.
|
||||
weakPrefix = "W/"
|
||||
)
|
||||
|
||||
// SetETag writes etag as a strong ETag response header, quoted per RFC 9110.
|
||||
// The value passed in is the bare validator — callers derive it (typically
|
||||
// from the type being served) and this applies the wire syntax, so the quoting
|
||||
// is decided in one place rather than at every handler.
|
||||
//
|
||||
// Call it before writing the body: once the response is committed the header
|
||||
// no longer reaches the client. An empty etag writes no header at all, so a
|
||||
// caller with nothing to validate against does not have to special-case it.
|
||||
func SetETag(w http.ResponseWriter, etag string) {
|
||||
if etag == "" {
|
||||
return
|
||||
}
|
||||
w.Header().Set(etagHeader, strconv.Quote(etag))
|
||||
}
|
||||
|
||||
// Precondition is a parsed If-Match request precondition. The zero value
|
||||
// matches nothing; a nil *Precondition is an unconditional request and matches
|
||||
// everything, so a handler can pass the result of IfMatch straight through
|
||||
// without a presence check.
|
||||
type Precondition struct {
|
||||
// tags are the strong entity-tags the client will accept, unquoted.
|
||||
tags []string
|
||||
|
||||
// any records the "*" form, which matches any current representation.
|
||||
any bool
|
||||
}
|
||||
|
||||
// IfMatch parses the request's If-Match precondition. It returns nil when the
|
||||
// header is absent — an unconditional request, which is the back-compatible
|
||||
// default: clients that know nothing of conditional requests keep working.
|
||||
//
|
||||
// A header that is present but carries nothing usable — empty, or nothing but
|
||||
// weak validators — yields a precondition that matches nothing rather than
|
||||
// nil. Failing closed is the only safe direction: a client that meant to send
|
||||
// a precondition must not have it silently dropped and its write let through
|
||||
// unguarded.
|
||||
func IfMatch(r *http.Request) *Precondition {
|
||||
values := r.Header.Values(ifMatchHeader)
|
||||
if len(values) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
p := &Precondition{}
|
||||
for _, value := range values {
|
||||
for raw := range strings.SplitSeq(value, ",") {
|
||||
candidate := strings.TrimSpace(raw)
|
||||
switch {
|
||||
case candidate == "":
|
||||
// Tolerated rather than rejected: a stray comma changes
|
||||
// nothing about what the client is willing to accept.
|
||||
case candidate == matchAny:
|
||||
p.any = true
|
||||
case strings.HasPrefix(candidate, weakPrefix):
|
||||
// Dropped, not unwrapped. If-Match uses strong comparison, so
|
||||
// a weak validator cannot satisfy it — and unwrapping one into
|
||||
// a strong tag would quietly grant the match the client's own
|
||||
// header said it could not have.
|
||||
default:
|
||||
p.tags = append(p.tags, strings.Trim(candidate, `"`))
|
||||
}
|
||||
}
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// Matches reports whether etag — the bare validator of the resource as it
|
||||
// currently stands — satisfies the precondition. A nil precondition matches
|
||||
// everything.
|
||||
//
|
||||
// Callers must establish that the resource exists before consulting this: the
|
||||
// "*" form asks whether there is any current representation, a question only
|
||||
// the caller can answer, and this reports true for it.
|
||||
func (p *Precondition) Matches(etag string) bool {
|
||||
if p == nil {
|
||||
return true
|
||||
}
|
||||
if p.any {
|
||||
return true
|
||||
}
|
||||
return slices.Contains(p.tags, etag)
|
||||
}
|
||||
@@ -1,161 +0,0 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestSetETag covers the wire syntax: the bare validator goes in, a quoted
|
||||
// strong entity-tag comes out. Handlers pass what the type derived, so the
|
||||
// quoting has to happen here or every handler re-decides it.
|
||||
func TestSetETag(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
SetETag(rec, "9f86d081884c7d65")
|
||||
|
||||
assert.Equal(t, `"9f86d081884c7d65"`, rec.Header().Get("ETag"),
|
||||
"the validator must be emitted quoted")
|
||||
}
|
||||
|
||||
// TestSetETagEmpty pins the no-op: a caller with nothing to validate against
|
||||
// must not emit an empty entity-tag, which would be a validator that every
|
||||
// later request could match.
|
||||
func TestSetETagEmpty(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
SetETag(rec, "")
|
||||
|
||||
assert.Empty(t, rec.Header().Values("ETag"), "an empty validator must write no header")
|
||||
}
|
||||
|
||||
// TestSetETagRoundTrip closes the loop between the two halves of the helper:
|
||||
// what SetETag emits is what IfMatch accepts back. A client echoing the header
|
||||
// it was given must match, or conditional requests never succeed in practice.
|
||||
func TestSetETagRoundTrip(t *testing.T) {
|
||||
const etag = "9f86d081884c7d65"
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
SetETag(rec, etag)
|
||||
|
||||
r := httptest.NewRequest(http.MethodPut, "/", nil)
|
||||
r.Header.Set("If-Match", rec.Header().Get("ETag"))
|
||||
|
||||
assert.True(t, IfMatch(r).Matches(etag), "an echoed ETag header must satisfy the precondition")
|
||||
}
|
||||
|
||||
// TestIfMatchAbsent pins the back-compatibility guarantee: a request with no
|
||||
// If-Match is unconditional, and the nil precondition it yields matches
|
||||
// anything so handlers need no presence check.
|
||||
func TestIfMatchAbsent(t *testing.T) {
|
||||
p := IfMatch(httptest.NewRequest(http.MethodPut, "/", nil))
|
||||
|
||||
require.Nil(t, p, "an absent header must yield no precondition")
|
||||
assert.True(t, p.Matches("9f86d081884c7d65"), "a nil precondition must match anything")
|
||||
assert.True(t, p.Matches(""), "a nil precondition must not depend on the validator")
|
||||
}
|
||||
|
||||
// TestIfMatchParsing walks the header forms a client can send. The weak and
|
||||
// unusable cases are the ones that matter: each must yield a precondition that
|
||||
// exists and refuses, never one that is absent and waves the write through.
|
||||
func TestIfMatchParsing(t *testing.T) {
|
||||
const current = "9f86d081884c7d65"
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
header string
|
||||
match bool
|
||||
reason string
|
||||
}{
|
||||
{
|
||||
name: "quoted current validator",
|
||||
header: `"9f86d081884c7d65"`,
|
||||
match: true,
|
||||
reason: "the ordinary conditional request must be honoured",
|
||||
},
|
||||
{
|
||||
name: "unquoted current validator",
|
||||
header: "9f86d081884c7d65",
|
||||
match: true,
|
||||
reason: "a client that omits the quoting means the same thing, and only an exact value can match",
|
||||
},
|
||||
{
|
||||
name: "stale validator",
|
||||
header: `"0000000000000000"`,
|
||||
match: false,
|
||||
reason: "a validator from an earlier read must not match",
|
||||
},
|
||||
{
|
||||
name: "star",
|
||||
header: "*",
|
||||
match: true,
|
||||
reason: "* matches any current representation",
|
||||
},
|
||||
{
|
||||
name: "list containing the current validator",
|
||||
header: `"0000000000000000", "9f86d081884c7d65"`,
|
||||
match: true,
|
||||
reason: "If-Match is a list; any member matching is a match",
|
||||
},
|
||||
{
|
||||
name: "list of stale validators",
|
||||
header: `"0000000000000000", "1111111111111111"`,
|
||||
match: false,
|
||||
reason: "a list none of whose members match must not match",
|
||||
},
|
||||
{
|
||||
name: "surrounding whitespace",
|
||||
header: ` "9f86d081884c7d65" `,
|
||||
match: true,
|
||||
reason: "list whitespace is not part of the entity-tag",
|
||||
},
|
||||
{
|
||||
name: "stray comma",
|
||||
header: `"9f86d081884c7d65", `,
|
||||
match: true,
|
||||
reason: "an empty list element says nothing about what the client accepts",
|
||||
},
|
||||
{
|
||||
name: "weak validator of the current representation",
|
||||
header: `W/"9f86d081884c7d65"`,
|
||||
match: false,
|
||||
reason: "If-Match uses strong comparison, so a weak validator never satisfies it",
|
||||
},
|
||||
{
|
||||
name: "weak validator alongside a strong one",
|
||||
header: `W/"0000000000000000", "9f86d081884c7d65"`,
|
||||
match: true,
|
||||
reason: "dropping the weak member must not discard the rest of the list",
|
||||
},
|
||||
{
|
||||
name: "empty header",
|
||||
header: "",
|
||||
match: false,
|
||||
reason: "a precondition the server cannot make sense of must fail closed, not vanish",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
r := httptest.NewRequest(http.MethodPut, "/", nil)
|
||||
r.Header.Set("If-Match", tc.header)
|
||||
|
||||
p := IfMatch(r)
|
||||
require.NotNil(t, p, "a header that was sent must yield a precondition: %s", tc.reason)
|
||||
assert.Equal(t, tc.match, p.Matches(current), tc.reason)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestIfMatchRepeatedHeader covers the same list split across header lines,
|
||||
// which is semantically identical to the comma form and which a proxy is free
|
||||
// to produce.
|
||||
func TestIfMatchRepeatedHeader(t *testing.T) {
|
||||
r := httptest.NewRequest(http.MethodPut, "/", nil)
|
||||
r.Header.Add("If-Match", `"0000000000000000"`)
|
||||
r.Header.Add("If-Match", `"9f86d081884c7d65"`)
|
||||
|
||||
assert.True(t, IfMatch(r).Matches("9f86d081884c7d65"),
|
||||
"entity-tags split across header lines must be read as one list")
|
||||
}
|
||||
Reference in New Issue
Block a user