[proxy] Bound the model-listing response to what policy authorises

Discovery proxies the upstream's full list, so the picker offers every
model the shared provider key can reach and each one outside the policy is
a request the chain denies a moment later. Restricting models is the point
of the product, and the client had no way to see the restriction.

Carry the resolved route's model list on the upstream rewrite and drop the
rest from the listing response. Only a route that enumerates its models
bounds anything: a catch-all claims every model, so its list passes
through. Anything the filter cannot safely rewrite, including a compressed
or oversized body, reaches the client untouched.
This commit is contained in:
mlsmaycon
2026-08-11 03:00:32 +00:00
parent a64a417ea5
commit 5353cab54f
6 changed files with 381 additions and 4 deletions

View File

@@ -199,8 +199,16 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
case matchOutcomeFound:
out := m.allowWithRoute(route, surface, in.UserGroups)
out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"})
if _, hadPrefix := splitBedrockNamespace(reqPath); hadPrefix && out.Mutations != nil && out.Mutations.RewriteUpstream != nil {
out.Mutations.RewriteUpstream.StripPathPrefix = bedrockNamespacePrefix
if out.Mutations != nil && out.Mutations.RewriteUpstream != nil {
if _, hadPrefix := splitBedrockNamespace(reqPath); hadPrefix {
out.Mutations.RewriteUpstream.StripPathPrefix = bedrockNamespacePrefix
}
// 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.RewriteUpstream.DiscoveryModels = append([]string(nil), route.Models...)
}
}
return out, nil
case matchOutcomeUnauthorised:
@@ -318,9 +326,15 @@ const connectionWarmPath = "/api/hello"
// that legitimately carries no model in its request (model listing and the
// connection-warming probe). These must route to an upstream rather than
// deny, so model enumeration works end to end.
// 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"
func isModelLessPath(reqPath string) bool {
return reqPath == "/v1/models" ||
strings.HasPrefix(reqPath, "/v1/models/") ||
return reqPath == modelListingPath ||
strings.HasPrefix(reqPath, modelListingPath+"/") ||
reqPath == connectionWarmPath
}

View File

@@ -935,3 +935,48 @@ func TestRouter_ConnectionWarmProbeRoutes(t *testing.T) {
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")
})
}

View File

@@ -253,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

View File

@@ -0,0 +1,154 @@
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
}
body, err := io.ReadAll(io.LimitReader(resp.Body, maxDiscoveryBodyBytes+1))
closeErr := resp.Body.Close()
if err != nil {
return err
}
if closeErr != nil {
return closeErr
}
if len(body) > maxDiscoveryBodyBytes {
restoreBody(resp, body)
return nil
}
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 _, ok := permitted[entryModelID(entry)]; ok {
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
}
// entryModelID returns the entry's model id in the form the policy stores
// it, or "" when the entry carries no usable id. A provider-prefixed id
// ("bedrock/anthropic.claude-sonnet-5") keeps only its last segment, which
// is what the operator registers.
func entryModelID(entry map[string]json.RawMessage) string {
raw, ok := entry["id"]
if !ok {
return ""
}
var id string
if err := json.Unmarshal(raw, &id); err != nil {
return ""
}
if slash := strings.LastIndex(id, "/"); slash >= 0 {
id = id[slash+1:]
}
return sharedllm.NormalizeAnthropicModel(id)
}
// restoreBody puts body back on the response and fixes the length headers
// so the client reads exactly what is there.
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)))
}

View File

@@ -0,0 +1,155 @@
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")
}

View File

@@ -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))
}