[management] Close four review findings on live model discovery

Four points from the review of #7246, all confirmed against the code.

**regionFromUpstream panicked on Bedrock's regionless endpoint.** The
template is "bedrock-runtime.<region>.amazonaws.com", so the two fixed
halves are "bedrock-runtime." and ".amazonaws.com". The regionless host
"bedrock-runtime.amazonaws.com" carries both at once, with the halves
overlapping rather than sandwiching a region — it satisfied HasPrefix and
HasSuffix, then sliced host[16:15]:

    panic: runtime error: slice bounds out of range [16:15]

That is reachable from any operator who types that host into upstream_url on
a Bedrock record. The length check makes the overlap read as "no region
here", which is what it is.

**A DNS-rebinding window sat between the guard and the dial.**
checkPublicHost resolved the host and the transport resolved it again to
dial, and the name's owner picks both answers. Public to the first lookup,
127.0.0.1 to the second, and the request reached loopback carrying the
operator's provider credential. The dialer now re-checks at the socket
through net.Dialer.Control, which runs post-resolution and pre-connect for
each address tried, so it sees what the second lookup actually returned.
The transport is cloned from http.DefaultTransport to keep its proxy and
TLS behaviour, and shared package-wide so the connection pool survives.

**Caller-input failures answered 500.** An unknown provider, an unusable
upstream, a region that cannot be read and a missing key are all reachable
from a well-formed request with a bad field value, and the OpenAPI document
already declares 400 for this endpoint. They now carry ErrInvalidRequest and
the handler branches on the sentinel rather than on message text. The
non-public-address refusal is included: that is the caller's own URL.

**The catalog id was trimmed for the emptiness test and then discarded.** A
padded " openai_api " cleared the check and reached catalog.Lookup with its
spaces, so the operator was told their provider was unknown.
This commit is contained in:
mlsmaycon
2026-08-22 11:29:08 +00:00
parent 73e3a2bacf
commit 52373f87f4
4 changed files with 203 additions and 9 deletions

View File

@@ -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": `{`,

View File

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

View File

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

View File

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