diff --git a/management/internals/modules/agentnetwork/handlers/model_discovery_handler_test.go b/management/internals/modules/agentnetwork/handlers/model_discovery_handler_test.go index df53612c9..9552d7202 100644 --- a/management/internals/modules/agentnetwork/handlers/model_discovery_handler_test.go +++ b/management/internals/modules/agentnetwork/handlers/model_discovery_handler_test.go @@ -3,6 +3,7 @@ package handlers import ( "context" "encoding/json" + "fmt" "net/http" "net/http/httptest" "strings" @@ -122,6 +123,32 @@ func TestDiscoverModelsReportsNoDiscoveryDistinctly(t *testing.T) { assert.Equal(t, http.StatusUnprocessableEntity, rec.Code) } +// TestDiscoverModelsTrimsTheCatalogID pins that the id the emptiness check +// accepts is the id the manager receives. A padded value that clears the check +// but reaches the catalog untrimmed misses the lookup, and the operator is told +// their provider does not exist. +func TestDiscoverModelsTrimsTheCatalogID(t *testing.T) { + stub := &discoveryManagerStub{} + + rec := postDiscovery(t, stub, `{"catalog_provider_id":" openai_api ","api_key":"sk"}`) + require.Equal(t, http.StatusOK, rec.Code, "body: %s", rec.Body.String()) + assert.Equal(t, "openai_api", stub.gotReq.CatalogID) +} + +// TestDiscoverModelsReportsCallerInputAsBadRequest covers the other half of the +// error mapping. These failures are all reachable from a well-formed request +// with a bad field value, so answering 500 both misinforms the operator and +// puts their typo into the server's error rate. +func TestDiscoverModelsReportsCallerInputAsBadRequest(t *testing.T) { + stub := &discoveryManagerStub{ + err: fmt.Errorf("%w: unknown catalog provider %q", modeldiscovery.ErrInvalidRequest, "nope"), + } + + rec := postDiscovery(t, stub, `{"catalog_provider_id":"nope","api_key":"sk"}`) + assert.Equal(t, http.StatusBadRequest, rec.Code) + assert.Contains(t, rec.Body.String(), "unknown catalog provider") +} + func TestDiscoverModelsRejectsMalformedRequests(t *testing.T) { for name, body := range map[string]string{ "not json": `{`, diff --git a/management/internals/modules/agentnetwork/handlers/providers_handler.go b/management/internals/modules/agentnetwork/handlers/providers_handler.go index 54b317442..d76885f2a 100644 --- a/management/internals/modules/agentnetwork/handlers/providers_handler.go +++ b/management/internals/modules/agentnetwork/handlers/providers_handler.go @@ -79,14 +79,18 @@ func (h *handler) discoverProviderModels(w http.ResponseWriter, r *http.Request) util.WriteErrorResponse("invalid json", http.StatusBadRequest, w) return } - if strings.TrimSpace(body.CatalogProviderId) == "" { + // Trimmed once and carried, not trimmed for the emptiness test and then + // discarded: a padded " openai_api " would clear the check here and miss + // the catalog lookup, reporting the provider as unknown. + catalogID := strings.TrimSpace(body.CatalogProviderId) + if catalogID == "" { util.WriteErrorResponse("catalog_provider_id is required", http.StatusBadRequest, w) return } recordID := strValue(body.ProviderId) req := modeldiscovery.Request{ - CatalogID: body.CatalogProviderId, + CatalogID: catalogID, UpstreamURL: strValue(body.UpstreamUrl), APIKey: strValue(body.ApiKey), } @@ -107,6 +111,14 @@ func (h *handler) discoverProviderModels(w http.ResponseWriter, r *http.Request) util.WriteErrorResponse(err.Error(), http.StatusUnprocessableEntity, w) return } + // An unknown provider, an unusable upstream, a missing region or a + // missing key are all things the caller sent, reachable from a + // well-formed request. Reporting them as 500 tells the operator the + // server broke and buries genuine faults in the error rate. + if errors.Is(err, modeldiscovery.ErrInvalidRequest) { + util.WriteErrorResponse(err.Error(), http.StatusBadRequest, w) + return + } util.WriteError(r.Context(), err, w) return } diff --git a/management/internals/modules/agentnetwork/modeldiscovery/discovery.go b/management/internals/modules/agentnetwork/modeldiscovery/discovery.go index 7fdef7c6f..00d59d238 100644 --- a/management/internals/modules/agentnetwork/modeldiscovery/discovery.go +++ b/management/internals/modules/agentnetwork/modeldiscovery/discovery.go @@ -25,6 +25,7 @@ import ( "net/netip" "net/url" "strings" + "syscall" "time" "golang.org/x/oauth2/google" @@ -55,6 +56,13 @@ const ( // a failure. var ErrNoDiscovery = errors.New("provider has no model-discovery endpoint") +// ErrInvalidRequest marks a discovery failure caused by the caller's own input +// rather than by the vendor or by this server. Every one of these is reachable +// from a well-formed request carrying a bad field value, so the handler owes +// the caller a 400 — a 500 would both misinform them and bury real server +// faults in the error rate. +var ErrInvalidRequest = errors.New("invalid discovery request") + // Model is one discovered model. type Model struct { // ID is the identifier to register on the provider record, in the form the @@ -98,7 +106,7 @@ type Client struct { func (c *Client) Fetch(ctx context.Context, req Request) ([]Model, error) { entry, ok := catalog.Lookup(req.CatalogID) if !ok { - return nil, fmt.Errorf("unknown catalog provider %q", req.CatalogID) + return nil, fmt.Errorf("%w: unknown catalog provider %q", ErrInvalidRequest, req.CatalogID) } if entry.Discovery == nil { return nil, ErrNoDiscovery @@ -161,7 +169,7 @@ func (c *Client) discoveryURL(entry catalog.Provider, req Request) (string, erro if host == "" { parsed, err := url.Parse(strings.TrimSpace(req.UpstreamURL)) if err != nil || parsed.Host == "" { - return "", fmt.Errorf("provider upstream %q is not a usable URL", req.UpstreamURL) + return "", fmt.Errorf("%w: provider upstream %q is not a usable URL", ErrInvalidRequest, req.UpstreamURL) } host = parsed.Host } @@ -174,7 +182,8 @@ func (c *Client) discoveryURL(entry catalog.Provider, req Request) (string, erro region = regionFromUpstream(entry, req.UpstreamURL) } if region == "" { - return "", fmt.Errorf("%s discovery needs a region, and none could be read from the provider upstream", entry.Name) + return "", fmt.Errorf("%w: %s discovery needs a region, and none could be read from the provider upstream", + ErrInvalidRequest, entry.Name) } host = strings.ReplaceAll(host, catalog.RegionPlaceholder, region) } @@ -206,7 +215,12 @@ func regionFromUpstream(entry catalog.Provider, upstreamURL string) string { // A bare host with no scheme parses as a path, not a host. host = strings.TrimSpace(upstreamURL) } - if !strings.HasPrefix(host, prefix) || !strings.HasSuffix(host, suffix) { + // The two halves must not overlap. "bedrock-runtime.amazonaws.com" carries + // both of Bedrock's — it is the regionless endpoint — and satisfies both + // checks above while leaving nothing between them, so slicing it would + // panic on an inverted range rather than report "no region here". + if !strings.HasPrefix(host, prefix) || !strings.HasSuffix(host, suffix) || + len(host) < len(prefix)+len(suffix) { return "" } region := host[len(prefix) : len(host)-len(suffix)] @@ -240,7 +254,7 @@ func (c *Client) checkPublicHost(host string) error { // loopback address is still a way to reach loopback. for _, addr := range addrs { if !isPublic(addr) { - return fmt.Errorf("discovery host %q resolves to a non-public address", host) + return fmt.Errorf("%w: discovery host %q resolves to a non-public address", ErrInvalidRequest, host) } } return nil @@ -277,7 +291,7 @@ func isPublic(addr netip.Addr) bool { func applyAuth(req *http.Request, entry catalog.Provider, apiKey string) error { key := strings.TrimSpace(apiKey) if key == "" { - return fmt.Errorf("%s discovery needs an API key", entry.Name) + return fmt.Errorf("%w: %s discovery needs an API key", ErrInvalidRequest, entry.Name) } if rest, ok := strings.CutPrefix(key, vertexKeyfilePrefix); ok { token, err := mintGCPToken(req.Context(), rest) @@ -347,8 +361,13 @@ func (c *Client) httpClient() *http.Client { if c.HTTPClient != nil { return c.HTTPClient } + transport := guardedTransport + if c.AllowPrivateHosts { + transport = http.DefaultTransport + } return &http.Client{ - Timeout: fetchTimeout, + Timeout: fetchTimeout, + Transport: transport, // A redirect is a way to move the request to a host the guard above // never checked, so none are followed. CheckRedirect: func(*http.Request, []*http.Request) error { @@ -356,3 +375,59 @@ func (c *Client) httpClient() *http.Client { }, } } + +// guardedTransport dials only addresses isPublic accepts. +// +// checkPublicHost resolves the host itself, and the transport then resolves it +// again when it dials — two lookups of a name whose owner chooses the answers. +// A record that returns a public address to the first and 127.0.0.1 to the +// second passes the guard and reaches loopback anyway, which is the whole of +// DNS rebinding. Re-checking at the socket closes that window: whatever the +// second lookup returned is what Control is handed, and an address the guard +// refuses never gets connected. +// +// Shared package-wide rather than built per Fetch so connections and their +// pool survive between calls; the guard holds no state. +var guardedTransport = newGuardedTransport() + +func newGuardedTransport() http.RoundTripper { + base, ok := http.DefaultTransport.(*http.Transport) + if !ok { + // Something replaced the default transport. Fall back to it rather + // than dropping its behaviour, and rely on checkPublicHost alone. + return http.DefaultTransport + } + // Cloned so proxy settings, TLS defaults and timeouts come from the + // standard transport rather than being restated here. + transport := base.Clone() + dialer := &net.Dialer{ + Timeout: fetchTimeout, + KeepAlive: 30 * time.Second, + Control: func(_, address string, _ syscall.RawConn) error { + return guardDialAddress(address) + }, + } + transport.DialContext = dialer.DialContext + return transport +} + +// guardDialAddress refuses a resolved socket address the discovery client has +// no business connecting to. Control hands it over post-resolution and +// pre-connect, once per address the dialer tries, so a name with several A +// records is checked at each one. +func guardDialAddress(address string) error { + host, _, err := net.SplitHostPort(address) + if err != nil { + return fmt.Errorf("discovery dial address %q is unreadable", address) + } + addr, err := netip.ParseAddr(host) + if err != nil { + // Control is documented to receive a resolved address; anything else + // is a state we cannot vet, so it does not get dialled. + return fmt.Errorf("discovery dial address %q is not an IP", host) + } + if !isPublic(addr) { + return fmt.Errorf("discovery refused to dial non-public address %s", addr) + } + return nil +} diff --git a/management/internals/modules/agentnetwork/modeldiscovery/discovery_test.go b/management/internals/modules/agentnetwork/modeldiscovery/discovery_test.go index 23827bbf6..4ae0c1594 100644 --- a/management/internals/modules/agentnetwork/modeldiscovery/discovery_test.go +++ b/management/internals/modules/agentnetwork/modeldiscovery/discovery_test.go @@ -255,6 +255,81 @@ func TestHostGuardResolvesAndRejectsLocalhost(t *testing.T) { assert.Contains(t, err.Error(), "non-public") } +// TestDialGuardRejectsRebindingToANonPublicAddress covers the window between +// the two DNS lookups. checkPublicHost resolves the host, then the transport +// resolves it again to dial; a name whose owner answers the first with a public +// address and the second with 127.0.0.1 would otherwise pass the guard and +// still reach loopback. The dial-time check sees whatever the second lookup +// actually returned. +func TestDialGuardRejectsRebindingToANonPublicAddress(t *testing.T) { + for _, tc := range []struct { + name string + address string + wantErr string + }{ + {"loopback", "127.0.0.1:443", "non-public"}, + {"cloud metadata", "169.254.169.254:80", "non-public"}, + {"rfc1918", "10.1.2.3:443", "non-public"}, + {"netbird overlay", "100.90.1.2:443", "non-public"}, + {"loopback v6", "[::1]:443", "non-public"}, + {"unresolved name", "evil.example.com:443", "not an IP"}, + {"no port", "1.1.1.1", "unreadable"}, + } { + t.Run(tc.name, func(t *testing.T) { + err := guardDialAddress(tc.address) + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantErr) + }) + } + + assert.NoError(t, guardDialAddress("1.1.1.1:443"), "a public address must still be dialled") + assert.NoError(t, guardDialAddress("[2606:4700:4700::1111]:443")) +} + +// TestDialGuardIsInstalledOnTheDefaultClient pins the wiring rather than the +// guard: a correct guard nothing calls protects nothing. +func TestDialGuardIsInstalledOnTheDefaultClient(t *testing.T) { + cl := &Client{} + transport, ok := cl.httpClient().Transport.(*http.Transport) + require.True(t, ok, "the default discovery client must carry the guarded transport") + require.NotNil(t, transport.DialContext, "the guarded transport must dial through the guard") + + _, err := transport.DialContext(context.Background(), "tcp", "127.0.0.1:9") + require.Error(t, err, "the guard must refuse loopback even when the caller dials it directly") + assert.Contains(t, err.Error(), "non-public") + + // Tests point the client at a loopback server on purpose, so the opt-out + // has to reach the dialer too. + relaxed := &Client{AllowPrivateHosts: true} + assert.Equal(t, http.DefaultTransport, relaxed.httpClient().Transport) +} + +// TestCallerInputFailuresAreMarkedInvalid keeps the handler's 400 mapping +// honest: it branches on this sentinel, so an unmarked caller-input failure +// silently becomes a 500. +func TestCallerInputFailuresAreMarkedInvalid(t *testing.T) { + for _, tc := range []struct { + name string + req Request + }{ + {"unknown provider", Request{CatalogID: "not_a_provider", APIKey: "k"}}, + {"unusable upstream", Request{CatalogID: "openai_api", UpstreamURL: "://", APIKey: "k"}}, + {"missing api key", Request{CatalogID: "openai_api", UpstreamURL: "https://api.openai.com"}}, + {"no region to read", Request{ + CatalogID: "bedrock_api", + UpstreamURL: "https://bedrock-runtime.amazonaws.com", + APIKey: "aws-bearer", + }}, + } { + t.Run(tc.name, func(t *testing.T) { + cl, _ := newStubClient(http.StatusOK, openAIListing) + _, err := cl.Fetch(context.Background(), tc.req) + require.Error(t, err) + assert.ErrorIs(t, err, ErrInvalidRequest) + }) + } +} + // TestEveryDiscoveryEntryHasAParser keeps the catalog and the parser table from // drifting: adding a Discovery block with a shape nothing parses would fail // only at runtime, in front of an operator. @@ -313,6 +388,11 @@ func TestRegionFromUpstream(t *testing.T) { // a region from it would build a URL pointing somewhere arbitrary. {"unrelated upstream", bedrock, "https://llm.internal.example.com", ""}, {"vertex global host has no region segment", vertex, "https://aiplatform.googleapis.com", ""}, + // Bedrock's regionless endpoint carries both halves of the template at + // once, with nothing between them. It has to read as "no region here" + // rather than as an inverted slice range. + {"bedrock regionless endpoint", bedrock, "https://bedrock-runtime.amazonaws.com", ""}, + {"bedrock regionless without scheme", bedrock, "bedrock-runtime.amazonaws.com", ""}, } { t.Run(tc.name, func(t *testing.T) { assert.Equal(t, tc.want, regionFromUpstream(tc.entry, tc.upstream))