diff --git a/core/services/failover/prober.go b/core/services/failover/prober.go index c6d412030..9a407e8a0 100644 --- a/core/services/failover/prober.go +++ b/core/services/failover/prober.go @@ -37,7 +37,12 @@ type DefaultProber struct { } func NewProber(loaded LoadedFunc) *DefaultProber { - return &DefaultProber{HTTP: &http.Client{}, Loaded: loaded} + return &DefaultProber{HTTP: &http.Client{ + // A redirect is a failed probe, not something to follow: Go resends + // custom headers such as x-api-key to any host, and the target's + // API key must reach only the configured upstream. + CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }, + }, Loaded: loaded} } func (p *DefaultProber) Liveness(ctx context.Context, cfg config.ModelConfig, kind Kind, warm bool) error { diff --git a/core/services/failover/prober_test.go b/core/services/failover/prober_test.go index b557c4f96..842934876 100644 --- a/core/services/failover/prober_test.go +++ b/core/services/failover/prober_test.go @@ -125,6 +125,23 @@ var _ = Describe("DefaultProber", func() { Expect(up.auth).To(Equal("Bearer sekret")) }) + It("does not follow a redirect, so the API key never leaves the upstream", func() { + GinkgoT().Setenv("FAILOVER_PROBE_KEY", "sekret") + other := newFakeUpstream() + DeferCleanup(other.srv.Close) + other.models = []string{"argus-llm"} + redirect := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, other.srv.URL+r.URL.Path, http.StatusFound) + })) + DeferCleanup(redirect.Close) + c := proxied("argus-llm", "") + c.Proxy.UpstreamURL = redirect.URL + "/v1/chat/completions" + c.Proxy.APIKeyEnv = "FAILOVER_PROBE_KEY" + c.Proxy.Provider = config.ProxyProviderAnthropic + Expect(p.Liveness(ctx, c, KindRemote, false)).To(MatchError(ContainSubstring("302"))) + Expect(other.paths).To(BeEmpty()) + }) + DescribeTable("remote inference hits the usecase endpoint", func(usecase, path string) { Expect(p.Inference(ctx, proxied("m", "", usecase), KindRemote, false)).To(Succeed())