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