From 5353cab54f99adbac7119ca2bfc7365f1f62a9f8 Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Tue, 11 Aug 2026 03:00:32 +0000 Subject: [PATCH] [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. --- .../builtin/llm_router/middleware.go | 22 ++- .../builtin/llm_router/middleware_test.go | 45 +++++ proxy/internal/middleware/types.go | 6 + proxy/internal/proxy/discovery_filter.go | 154 +++++++++++++++++ proxy/internal/proxy/discovery_filter_test.go | 155 ++++++++++++++++++ proxy/internal/proxy/reverseproxy.go | 3 + 6 files changed, 381 insertions(+), 4 deletions(-) create mode 100644 proxy/internal/proxy/discovery_filter.go create mode 100644 proxy/internal/proxy/discovery_filter_test.go diff --git a/proxy/internal/middleware/builtin/llm_router/middleware.go b/proxy/internal/middleware/builtin/llm_router/middleware.go index 88d7d6cbd..f76d24a4c 100644 --- a/proxy/internal/middleware/builtin/llm_router/middleware.go +++ b/proxy/internal/middleware/builtin/llm_router/middleware.go @@ -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 } diff --git a/proxy/internal/middleware/builtin/llm_router/middleware_test.go b/proxy/internal/middleware/builtin/llm_router/middleware_test.go index ec67ec05b..f1b0e4583 100644 --- a/proxy/internal/middleware/builtin/llm_router/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_router/middleware_test.go @@ -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") + }) +} diff --git a/proxy/internal/middleware/types.go b/proxy/internal/middleware/types.go index f81892fc8..3c0ac0ab6 100644 --- a/proxy/internal/middleware/types.go +++ b/proxy/internal/middleware/types.go @@ -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 diff --git a/proxy/internal/proxy/discovery_filter.go b/proxy/internal/proxy/discovery_filter.go new file mode 100644 index 000000000..b4b0e92e0 --- /dev/null +++ b/proxy/internal/proxy/discovery_filter.go @@ -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))) +} diff --git a/proxy/internal/proxy/discovery_filter_test.go b/proxy/internal/proxy/discovery_filter_test.go new file mode 100644 index 000000000..ab61a0f3c --- /dev/null +++ b/proxy/internal/proxy/discovery_filter_test.go @@ -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") +} diff --git a/proxy/internal/proxy/reverseproxy.go b/proxy/internal/proxy/reverseproxy.go index 9150c0329..7c9e21261 100644 --- a/proxy/internal/proxy/reverseproxy.go +++ b/proxy/internal/proxy/reverseproxy.go @@ -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)) }