From a9f5ed76444357faa1b9c74cf5b4f074cb5e42ac Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Thu, 28 May 2026 13:22:38 +0200 Subject: [PATCH] fix: re-validate provider per request and classify reloads (#25766) Refactors the `aibridgeproxyd` provider reload mechanism which was unnecessarily complex. Also ensures that providers are evaluated on each CONNECT request to prevent interception of requests to (newly) disabled providers; in this case the requests will passthrough unencrypted, by design. --- enterprise/aibridgeproxyd/aibridgeproxyd.go | 95 ++- .../aibridgeproxyd/aibridgeproxyd_test.go | 706 ++++++++---------- enterprise/aibridgeproxyd/reload.go | 123 +-- .../aibridgeproxyd/reload_internal_test.go | 106 +-- enterprise/aibridgeproxyd/reload_test.go | 288 +++++-- enterprise/cli/aibridgeproxyd.go | 62 +- .../cli/aibridgeproxyd_internal_test.go | 105 +++ 7 files changed, 856 insertions(+), 629 deletions(-) create mode 100644 enterprise/cli/aibridgeproxyd_internal_test.go diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index 19d05ca511..cfcb2071c4 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -31,18 +31,6 @@ import ( agplaibridge "github.com/coder/coder/v2/coderd/aibridge" ) -// ProviderRoute is the routing entry for a single AI provider: the -// instance name (the routing key) and the upstream base URL (the -// source of the MITM allowlist host). -type ProviderRoute struct { - Name string - BaseURL string -} - -// RefreshProvidersFunc returns the live provider set used by Reload to -// rebuild the proxy's routing snapshot. -type RefreshProvidersFunc func(ctx context.Context) ([]ProviderRoute, error) - // Known AI provider hosts. const ( HostAnthropic = "api.anthropic.com" @@ -161,7 +149,7 @@ type Server struct { // providerRouter keeps CONNECT matching and provider lookup in sync. type providerRouter struct { - mitmHosts []string // host:port allowlist for the goproxy condition. + mitmHosts []string // host:port set the goproxy condition matches against. nameByHost map[string]string // lowercase hostname -> provider name. } @@ -218,15 +206,8 @@ type Options struct { // CertStore is an optional certificate cache for MITM. If nil, a default // cache is created. Exposed for testing. CertStore goproxy.CertStorage - // DomainAllowlist seeds the boot-time MITM allowlist. Production - // callers should leave this empty and rely on RefreshProviders; - // tests use it to skip the refresh round-trip. - DomainAllowlist []string - // AIBridgeProviderFromHost seeds the boot-time host -> provider - // name mapping. Required iff DomainAllowlist is non-empty. - AIBridgeProviderFromHost func(host string) string // UpstreamProxy is the URL of an upstream HTTP proxy to chain tunneled - // (non-allowlisted) requests through. If empty, tunneled requests connect + // (non-provider-host) requests through. If empty, tunneled requests connect // directly to their destinations. // Format: http://[user:pass@]host:port or https://[user:pass@]host:port UpstreamProxy string @@ -249,7 +230,7 @@ type Options struct { // If nil, metrics will not be recorded. Metrics *Metrics // RefreshProviders, when set, is invoked by Server.Reload to fetch - // the live provider snapshot used to derive the MITM allowlist and + // the live provider snapshot used to derive the MITM host set and // host -> provider-name routing. Nil disables hot-reload. RefreshProviders RefreshProvidersFunc } @@ -296,14 +277,6 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) allowedPorts = []string{"80", "443"} } - // Build the boot-time router from DomainAllowlist + the lookup fn. - // Both empty is fine: the server fails closed (no MITM until - // Reload populates the router from the database). - bootRouter, err := buildBootRouter(opts.DomainAllowlist, opts.AIBridgeProviderFromHost, allowedPorts) - if err != nil { - return nil, err - } - // Parse configured exceptions to the blocked IP ranges. allowedPrivateRanges := make([]net.IPNet, 0, len(opts.AllowedPrivateCIDRs)) for _, cidr := range opts.AllowedPrivateCIDRs { @@ -352,13 +325,13 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) newDumper: opts.NewDumper, metrics: opts.Metrics, } - // Seed the boot-time router from the constructor inputs so the - // proxy can serve immediately. Reload may swap this snapshot at any - // point after construction. - srv.providerRouter.Store(bootRouter) + // Start with an empty router; the first Reload populates it from + // the configured provider source. The proxy fails closed (no MITM) + // until that happens. + srv.providerRouter.Store(emptyProviderRouter) - // Configure upstream proxy for tunneled (non-allowlisted) CONNECT requests. - // Allowlisted domains are MITM'd and forwarded to aibridge directly, + // Configure upstream proxy for tunneled (non-provider-host) CONNECT requests. + // Provider-host domains are MITM'd and forwarded to aibridge directly, // bypassing the upstream proxy. if opts.UpstreamProxy != "" { upstreamURL, err := url.Parse(opts.UpstreamProxy) @@ -443,7 +416,7 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) // Reject CONNECT requests to non-standard ports. proxy.OnRequest().HandleConnectFunc(srv.portMiddleware(allowedPorts)) - // Apply MITM with authentication only to allowlisted hosts. The host + // Apply MITM with authentication only to provider hosts. The host // list is loaded from the atomic router on every CONNECT so a // Reload while inflight requests are in progress takes effect on // the next CONNECT without touching the already-MITM'd ones. @@ -452,9 +425,9 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) srv.authMiddleware, ) - // Tunnel CONNECT requests for non-allowlisted domains directly to their destination. + // Tunnel CONNECT requests for non-provider-host domains directly to their destination. // goproxy calls handlers in registration order: this must come after the MITM handler - // so it only handles requests that weren't matched by the allowlist. + // so it only handles requests that weren't matched as provider hosts. proxy.OnRequest().HandleConnectFunc(srv.tunneledMiddleware) // Handle decrypted requests: route to aibridged for known AI providers, or tunnel to original destination. @@ -495,7 +468,6 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) slog.F("listen_addr", listener.Addr().String()), slog.F("tls_listener_enabled", srv.tlsEnabled), slog.F("coder_access_url", coderAccessURL.String()), - slog.F("domain_allowlist", bootRouter.mitmHosts), slog.F("upstream_proxy", opts.UpstreamProxy), slog.F("allowed_private_cidrs", opts.AllowedPrivateCIDRs), slog.F("api_dump_enabled", opts.NewDumper != nil), @@ -810,7 +782,7 @@ func newProxyAuthRequiredResponse(req *http.Request) *http.Response { } } -// tunneledMiddleware is a CONNECT middleware that handles tunneled (non-allowlisted) +// tunneledMiddleware is a CONNECT middleware that handles tunneled (non-provider-host) // connections. These connections are not MITM'd and are tunneled directly to their // destination. This middleware records metrics for tunneled CONNECT sessions. func (s *Server) tunneledMiddleware(host string, _ *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) { @@ -946,16 +918,28 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. return req, resp } - if reqCtx.Provider == "" { - // A concurrent Reload can remove the provider after CONNECT - // authentication. The request is MITM'd (decrypted), but without a - // mapping there is no known route to aibridge. Log and forward - // to the original destination as a fallback. - s.logger.Warn(s.ctx, "decrypted request has no provider mapping, passing through", + // Re-validate the CONNECT-time provider against the live router. + // A long-lived CONNECT tunnel can outlive a provider being disabled, + // removed, or renamed: the captured reqCtx.Provider is stale, but + // subsequent decrypted requests would still route to aibridged if we + // trusted it. Look up the provider for the current request's host + // and pass through if the mapping is gone or has changed. + host := req.URL.Hostname() + if host == "" { + host = req.Host + if h, _, splitErr := net.SplitHostPort(host); splitErr == nil { + host = h + } + } + liveProvider := s.loadProviderRouter().providerFromHost(host) + if liveProvider == "" || liveProvider != reqCtx.Provider { + s.logger.Warn(s.ctx, "provider mapping changed or removed since CONNECT, passing through", slog.F("connect_id", reqCtx.ConnectSessionID.String()), slog.F("host", req.Host), slog.F("method", req.Method), slog.F("path", originalPath), + slog.F("connect_provider", reqCtx.Provider), + slog.F("live_provider", liveProvider), ) return req, nil } @@ -1053,8 +1037,13 @@ func injectBYOKHeaderIfNeeded(header http.Header, coderToken string) { } // handleResponse handles responses received from aibridged. -// This is only called for MITM'd requests (allowlisted domains routed through aibridged). -// Tunneled requests (non-allowlisted domains) bypass this handler entirely. +// This is called for every MITM'd request, including the pass-through +// path where handleRequest re-validated the CONNECT-time provider and +// forwarded the request to the original upstream instead of aibridged. +// Pass-through responses are identified by reqCtx.RequestID == uuid.Nil +// (set only when handleRequest routes to aibridged) and are skipped here +// to avoid mislabeled logs and corrupting MITM metrics. +// Tunneled requests (non-provider-host domains) bypass this handler entirely. func (s *Server) handleResponse(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response { if resp == nil { return nil @@ -1077,6 +1066,14 @@ func (s *Server) handleResponse(resp *http.Response, ctx *goproxy.ProxyCtx) *htt slog.F("status", resp.StatusCode), ) + // Pass-through responses (handleRequest returned without routing to + // aibridged) come from the real upstream. The aibridged-specific log + // and metrics do not apply; the pass-through itself is already logged + // in handleRequest. + if requestID == uuid.Nil { + return resp + } + switch { case resp.StatusCode >= http.StatusInternalServerError: logger.Error(s.ctx, "received error response from aibridged") diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index 6b843d8b14..fbf77956a2 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -3,6 +3,7 @@ package aibridgeproxyd_test import ( "bufio" "bytes" + "context" "crypto/rand" "crypto/rsa" "crypto/tls" @@ -145,20 +146,19 @@ func generateListenerCert(t *testing.T) (certFile, keyFile string) { } type testProxyConfig struct { - listenAddr string - tlsCertFile string - tlsKeyFile string - coderAccessURL string - allowedPorts []string - certStore *aibridgeproxyd.CertCache - domainAllowlist []string - aibridgeProviderFromHost func(string) string - upstreamProxy string - upstreamProxyCA string - allowedPrivateCIDRs []string - newDumper func(string, string) aibridgeproxyd.RoundTripDumper - metrics *aibridgeproxyd.Metrics - refreshProviders aibridgeproxyd.RefreshProvidersFunc + listenAddr string + tlsCertFile string + tlsKeyFile string + coderAccessURL string + allowedPorts []string + certStore *aibridgeproxyd.CertCache + providers []aibridgeproxyd.ReloadedProvider + upstreamProxy string + upstreamProxyCA string + allowedPrivateCIDRs []string + newDumper func(string, string) aibridgeproxyd.RoundTripDumper + metrics *aibridgeproxyd.Metrics + refreshProviders aibridgeproxyd.RefreshProvidersFunc } type testProxyOption func(*testProxyConfig) @@ -181,15 +181,41 @@ func withCertStore(store *aibridgeproxyd.CertCache) testProxyOption { } } -func withDomainAllowlist(domains ...string) testProxyOption { +// withProviders configures the proxy with the given classified provider +// set. The reload helper synthesizes a RefreshProvidersFunc and the +// router is populated synchronously during newTestProxy before the +// server begins serving. +func withProviders(providers ...aibridgeproxyd.ReloadedProvider) testProxyOption { return func(cfg *testProxyConfig) { - cfg.domainAllowlist = domains + cfg.providers = providers } } -func withAIBridgeProviderFromHost(fn func(string) string) testProxyOption { +// withProviderHosts is a convenience that builds enabled +// ReloadedProvider entries from each host, looking up the well-known +// provider name via testProviderFromHost and falling back to +// "test-provider" for hosts without a well-known mapping. Equivalent +// to passing each entry individually to withProviders. +func withProviderHosts(hosts ...string) testProxyOption { return func(cfg *testProxyConfig) { - cfg.aibridgeProviderFromHost = fn + providers := make([]aibridgeproxyd.ReloadedProvider, 0, len(hosts)) + for _, h := range hosts { + name := testProviderFromHost(h) + if name == "" { + name = "test-provider" + } + host, _, splitErr := net.SplitHostPort(h) + if splitErr != nil { + host = h + } + providers = append(providers, aibridgeproxyd.ReloadedProvider{ + Name: name, + Type: "openai", + Host: strings.ToLower(host), + Status: aibridgeproxyd.ProviderStatusEnabled, + }) + } + cfg.providers = providers } } @@ -264,39 +290,48 @@ func newTestProxy(t *testing.T, opts ...testProxyOption) *aibridgeproxyd.Server t.Helper() cfg := &testProxyConfig{ - listenAddr: "127.0.0.1:0", - coderAccessURL: "http://localhost:3000", - domainAllowlist: []string{"127.0.0.1", "localhost"}, + listenAddr: "127.0.0.1:0", + coderAccessURL: "http://localhost:3000", // Allow 127.0.0.1 by default so test servers, which always listen on // loopback, are reachable. Tests that verify IP blocking override this. allowedPrivateCIDRs: []string{"127.0.0.1/32"}, - aibridgeProviderFromHost: func(host string) string { - return "test-provider" + providers: []aibridgeproxyd.ReloadedProvider{ + {Name: "test-provider", Type: "openai", Host: "127.0.0.1", Status: aibridgeproxyd.ProviderStatusEnabled}, + {Name: "test-provider", Type: "openai", Host: "localhost", Status: aibridgeproxyd.ProviderStatusEnabled}, }, } for _, opt := range opts { opt(cfg) } + // If the test did not supply a RefreshProviders, synthesize one + // that returns the configured providers verbatim. This populates + // the router synchronously below, mirroring how production starts + // up after the first reload completes. + if cfg.refreshProviders == nil { + providers := cfg.providers + cfg.refreshProviders = func(context.Context) (aibridgeproxyd.ProviderReload, error) { + return aibridgeproxyd.ProviderReload{Providers: providers}, nil + } + } + mitmCertFile, mitmKeyFile := getSharedTestMITMCert(t) logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug) aibridgeOpts := aibridgeproxyd.Options{ - ListenAddr: cfg.listenAddr, - TLSCertFile: cfg.tlsCertFile, - TLSKeyFile: cfg.tlsKeyFile, - CoderAccessURL: cfg.coderAccessURL, - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - AllowedPorts: cfg.allowedPorts, - DomainAllowlist: cfg.domainAllowlist, - AIBridgeProviderFromHost: cfg.aibridgeProviderFromHost, - UpstreamProxy: cfg.upstreamProxy, - UpstreamProxyCA: cfg.upstreamProxyCA, - AllowedPrivateCIDRs: cfg.allowedPrivateCIDRs, - NewDumper: cfg.newDumper, - Metrics: cfg.metrics, - RefreshProviders: cfg.refreshProviders, + ListenAddr: cfg.listenAddr, + TLSCertFile: cfg.tlsCertFile, + TLSKeyFile: cfg.tlsKeyFile, + CoderAccessURL: cfg.coderAccessURL, + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + AllowedPorts: cfg.allowedPorts, + UpstreamProxy: cfg.upstreamProxy, + UpstreamProxyCA: cfg.upstreamProxyCA, + AllowedPrivateCIDRs: cfg.allowedPrivateCIDRs, + NewDumper: cfg.newDumper, + Metrics: cfg.metrics, + RefreshProviders: cfg.refreshProviders, } if cfg.certStore != nil { aibridgeOpts.CertStore = cfg.certStore @@ -306,6 +341,10 @@ func newTestProxy(t *testing.T, opts ...testProxyOption) *aibridgeproxyd.Server require.NoError(t, err) t.Cleanup(func() { _ = srv.Close() }) + // Populate the router before the server starts handling traffic. + // Production performs the first reload during boot via pubsub. + require.NoError(t, srv.Reload(t.Context())) + // Wait for the proxy server to be ready. proxyAddr := srv.Addr() require.NotEmpty(t, proxyAddr) @@ -444,10 +483,9 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.Error(t, err) require.Contains(t, err.Error(), "listen address is required") @@ -460,11 +498,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: "", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.Error(t, err) require.Contains(t, err.Error(), "listen address is required") @@ -477,12 +514,11 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - TLSCertFile: "cert.pem", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: "127.0.0.1:0", + TLSCertFile: "cert.pem", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.Error(t, err) require.Contains(t, err.Error(), "tls cert file and tls key file must both be set") @@ -495,12 +531,11 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - TLSKeyFile: "key.pem", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: "127.0.0.1:0", + TLSKeyFile: "key.pem", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.Error(t, err) require.Contains(t, err.Error(), "tls cert file and tls key file must both be set") @@ -513,14 +548,12 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - TLSCertFile: "/nonexistent/cert.pem", - TLSKeyFile: "/nonexistent/key.pem", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, + ListenAddr: "127.0.0.1:0", + TLSCertFile: "/nonexistent/cert.pem", + TLSKeyFile: "/nonexistent/key.pem", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.Error(t, err) require.Contains(t, err.Error(), "load listener TLS certificate") @@ -533,10 +566,9 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: "127.0.0.1:0", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.Error(t, err) require.Contains(t, err.Error(), "coder access URL is required") @@ -549,11 +581,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: " ", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: " ", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.Error(t, err) require.Contains(t, err.Error(), "coder access URL is required") @@ -566,11 +597,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "://invalid", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "://invalid", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.Error(t, err) require.Contains(t, err.Error(), "invalid coder access URL") @@ -583,12 +613,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic}, - AIBridgeProviderFromHost: testProviderFromHost, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.NoError(t, err) require.Equal(t, "localhost", srv.CoderAccessURL().Hostname()) @@ -602,12 +630,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "https://localhost", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic}, - AIBridgeProviderFromHost: testProviderFromHost, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "https://localhost", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.NoError(t, err) require.Equal(t, "localhost", srv.CoderAccessURL().Hostname()) @@ -621,12 +647,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic}, - AIBridgeProviderFromHost: testProviderFromHost, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.NoError(t, err) require.Equal(t, "localhost", srv.CoderAccessURL().Hostname()) @@ -639,10 +663,9 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CoderAccessURL: "http://localhost:3000", - MITMKeyFile: "key.pem", - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + MITMKeyFile: "key.pem", }) require.Error(t, err) require.Contains(t, err.Error(), "cert file and key file are required") @@ -654,10 +677,9 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: "cert.pem", - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: "cert.pem", }) require.Error(t, err) require.Contains(t, err.Error(), "cert file and key file are required") @@ -669,104 +691,15 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: "/nonexistent/cert.pem", - MITMKeyFile: "/nonexistent/key.pem", - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: "/nonexistent/cert.pem", + MITMKeyFile: "/nonexistent/key.pem", }) require.Error(t, err) require.Contains(t, err.Error(), "failed to load MITM certificate") }) - t.Run("MissingDomainAllowlist", func(t *testing.T) { - t.Parallel() - - mitmCertFile, mitmKeyFile := getSharedTestMITMCert(t) - logger := slogtest.Make(t, nil) - - srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - AIBridgeProviderFromHost: testProviderFromHost, - }) - require.NoError(t, err) - t.Cleanup(func() { _ = srv.Close() }) - }) - - t.Run("EmptyDomainAllowlist", func(t *testing.T) { - t.Parallel() - - mitmCertFile, mitmKeyFile := getSharedTestMITMCert(t) - logger := slogtest.Make(t, nil) - - srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{""}, - AIBridgeProviderFromHost: testProviderFromHost, - }) - require.NoError(t, err) - t.Cleanup(func() { _ = srv.Close() }) - }) - - t.Run("InvalidDomainAllowlist", func(t *testing.T) { - t.Parallel() - - mitmCertFile, mitmKeyFile := getSharedTestMITMCert(t) - logger := slogtest.Make(t, nil) - - _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{"[invalid:domain"}, - }) - require.Error(t, err) - require.Contains(t, err.Error(), "invalid domain") - }) - - t.Run("DomainWithNonAllowedPort", func(t *testing.T) { - t.Parallel() - - mitmCertFile, mitmKeyFile := getSharedTestMITMCert(t) - logger := slogtest.Make(t, nil) - - _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{"api.anthropic.com:8443"}, - }) - require.Error(t, err) - require.Contains(t, err.Error(), "invalid port in domain") - }) - - t.Run("AllowlistWithoutProviderMapping", func(t *testing.T) { - t.Parallel() - - mitmCertFile, mitmKeyFile := getSharedTestMITMCert(t) - logger := slogtest.Make(t, nil) - - _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{"unknown.example.com"}, - AIBridgeProviderFromHost: testProviderFromHost, - }) - require.Error(t, err) - require.Contains(t, err.Error(), `domain "unknown.example.com" is in allowlist but has no provider mapping`) - }) - t.Run("InvalidUpstreamProxy", func(t *testing.T) { t.Parallel() @@ -774,13 +707,11 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - UpstreamProxy: "://invalid-url", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + UpstreamProxy: "://invalid-url", }) require.Error(t, err) require.Contains(t, err.Error(), "invalid upstream proxy URL") @@ -793,14 +724,12 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - UpstreamProxy: "https://proxy.example.com:8080", - UpstreamProxyCA: "/nonexistent/ca.pem", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + UpstreamProxy: "https://proxy.example.com:8080", + UpstreamProxyCA: "/nonexistent/ca.pem", }) require.Error(t, err) require.Contains(t, err.Error(), "failed to read upstream proxy CA certificate") @@ -813,13 +742,11 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - UpstreamProxy: "http://:@proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + UpstreamProxy: "http://:@proxy.example.com:8080", }) require.Error(t, err) require.Contains(t, err.Error(), "invalid credentials: both username and password are empty") @@ -832,13 +759,11 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - AllowedPrivateCIDRs: []string{"not-a-cidr"}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + AllowedPrivateCIDRs: []string{"not-a-cidr"}, }) require.Error(t, err) require.Contains(t, err.Error(), "invalid allowed private CIDR") @@ -851,12 +776,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.NoError(t, err) require.NotNil(t, srv) @@ -870,14 +793,12 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - TLSCertFile: listenerCertFile, - TLSKeyFile: listenerKeyFile, - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, + ListenAddr: "127.0.0.1:0", + TLSCertFile: listenerCertFile, + TLSKeyFile: listenerKeyFile, + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.NoError(t, err) require.NotNil(t, srv) @@ -890,13 +811,11 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - UpstreamProxy: "http://proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + UpstreamProxy: "http://proxy.example.com:8080", }) require.NoError(t, err) require.NotNil(t, srv) @@ -910,14 +829,12 @@ func TestNew(t *testing.T) { // Use the shared MITM certificate as the upstream proxy CA (it's a valid PEM cert) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - UpstreamProxy: "https://proxy.example.com:8080", - UpstreamProxyCA: mitmCertFile, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + UpstreamProxy: "https://proxy.example.com:8080", + UpstreamProxyCA: mitmCertFile, }) require.NoError(t, err) require.NotNil(t, srv) @@ -930,13 +847,11 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - UpstreamProxy: "http://proxyuser:proxypass@proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + UpstreamProxy: "http://proxyuser:proxypass@proxy.example.com:8080", }) require.NoError(t, err) require.NotNil(t, srv) @@ -949,13 +864,11 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - UpstreamProxy: "http://proxyuser:@proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + UpstreamProxy: "http://proxyuser:@proxy.example.com:8080", }) require.NoError(t, err) require.NotNil(t, srv) @@ -969,13 +882,11 @@ func TestNew(t *testing.T) { // Username only (no colon) should also succeed (password is optional) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - UpstreamProxy: "http://proxyuser@proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + UpstreamProxy: "http://proxyuser@proxy.example.com:8080", }) require.NoError(t, err) require.NotNil(t, srv) @@ -988,13 +899,11 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - UpstreamProxy: "http://:proxypass@proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + UpstreamProxy: "http://:proxypass@proxy.example.com:8080", }) require.NoError(t, err) require.NotNil(t, srv) @@ -1011,13 +920,11 @@ func TestNew(t *testing.T) { metrics := aibridgeproxyd.NewMetrics(reg) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - Metrics: metrics, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + Metrics: metrics, }) require.NoError(t, err) require.NotNil(t, srv) @@ -1030,13 +937,11 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - AllowedPrivateCIDRs: []string{"127.0.0.1/32"}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + AllowedPrivateCIDRs: []string{"127.0.0.1/32"}, }) require.NoError(t, err) require.NotNil(t, srv) @@ -1053,12 +958,10 @@ func TestClose(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, }) require.NoError(t, err) @@ -1081,13 +984,11 @@ func TestClose(t *testing.T) { metrics := aibridgeproxyd.NewMetrics(reg) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AIBridgeProviderFromHost: testProviderFromHost, - Metrics: metrics, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + Metrics: metrics, }) require.NoError(t, err) @@ -1110,19 +1011,19 @@ func TestProxy_CertCaching(t *testing.T) { t.Parallel() tests := []struct { - name string - domainAllowlist []string - tunneled bool + name string + providerHosts []string + tunneled bool }{ { - name: "AllowlistedDomainCached", - domainAllowlist: nil, // will use targetURL.Hostname() - tunneled: false, + name: "ProviderHostCached", + providerHosts: nil, // will use targetURL.Hostname() + tunneled: false, }, { - name: "NonAllowlistedDomainNotCached", - domainAllowlist: []string{"other.example.com"}, - tunneled: true, + name: "NonProviderHostNotCached", + providerHosts: []string{"other.example.com"}, + tunneled: true, }, } @@ -1135,7 +1036,7 @@ func TestProxy_CertCaching(t *testing.T) { w.WriteHeader(http.StatusOK) }) - // Create a mock aibridged server for allowlisted (MITM'd) requests. + // Create a mock aibridged server for provider-host (MITM'd) requests. aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) })) @@ -1144,10 +1045,10 @@ func TestProxy_CertCaching(t *testing.T) { // Create a cert cache so we can inspect it after the request. certCache := aibridgeproxyd.NewCertCache() - // Configure domain allowlist. - domainAllowlist := tt.domainAllowlist - if domainAllowlist == nil { - domainAllowlist = []string{targetURL.Hostname()} + // Configure provider hosts. + providerHosts := tt.providerHosts + if providerHosts == nil { + providerHosts = []string{targetURL.Hostname()} } // Start the proxy server with the certificate cache. @@ -1155,7 +1056,7 @@ func TestProxy_CertCaching(t *testing.T) { withCoderAccessURL(aibridgedServer.URL), withAllowedPorts(targetURL.Port()), withCertStore(certCache), - withDomainAllowlist(domainAllowlist...), + withProviderHosts(providerHosts...), ) // Build the cert pool for the client to trust: @@ -1189,7 +1090,7 @@ func TestProxy_CertCaching(t *testing.T) { if tt.tunneled { // Certificate should NOT have been cached since request was tunneled. - require.Equal(t, 1, genCalls, "certificate should NOT have been cached for non-allowlisted domain") + require.Equal(t, 1, genCalls, "certificate should NOT have been cached for non-provider-host") } else { // Certificate should have been cached during MITM. require.Equal(t, 0, genCalls, "certificate should have been cached during request") @@ -1233,7 +1134,7 @@ func TestProxy_PortValidation(t *testing.T) { _, _ = w.Write([]byte("hello from target")) }) - // Create a mock aibridged server for allowlisted (MITM'd) requests. + // Create a mock aibridged server for provider-host (MITM'd) requests. aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("hello from aibridged")) @@ -1244,7 +1145,7 @@ func TestProxy_PortValidation(t *testing.T) { srv := newTestProxy(t, withCoderAccessURL(aibridgedServer.URL), withAllowedPorts(tt.allowedPorts(targetURL)...), - withDomainAllowlist(targetURL.Hostname()), + withProviderHosts(targetURL.Hostname()), ) // Make a request through the proxy to the target server. @@ -1309,7 +1210,7 @@ func TestProxy_Authentication(t *testing.T) { _, _ = w.Write([]byte("hello from target")) }) - // Create a mock aibridged server for allowlisted (MITM'd) requests. + // Create a mock aibridged server for provider-host (MITM'd) requests. aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("hello from aibridged")) @@ -1320,7 +1221,7 @@ func TestProxy_Authentication(t *testing.T) { srv := newTestProxy(t, withCoderAccessURL(aibridgedServer.URL), withAllowedPorts(targetURL.Port()), - withDomainAllowlist(targetURL.Hostname()), + withProviderHosts(targetURL.Hostname()), ) if tt.expectSuccess { @@ -1365,18 +1266,18 @@ func TestProxy_MITM(t *testing.T) { t.Parallel() tests := []struct { - name string - domainAllowlist []string - allowedPorts []string - buildTargetURL func(tunneledURL *url.URL) (string, error) - tunneled bool - expectedPath string - provider string + name string + providerHosts []string + allowedPorts []string + buildTargetURL func(tunneledURL *url.URL) (string, error) + tunneled bool + expectedPath string + provider string }{ { - name: "MitmdAnthropic", - domainAllowlist: []string{aibridgeproxyd.HostAnthropic}, - allowedPorts: []string{"443"}, + name: "MitmdAnthropic", + providerHosts: []string{aibridgeproxyd.HostAnthropic}, + allowedPorts: []string{"443"}, buildTargetURL: func(_ *url.URL) (string, error) { return "https://api.anthropic.com/v1/messages", nil }, @@ -1384,9 +1285,9 @@ func TestProxy_MITM(t *testing.T) { provider: "anthropic", }, { - name: "MitmdAnthropicNonDefaultPort", - domainAllowlist: []string{aibridgeproxyd.HostAnthropic}, - allowedPorts: []string{"8443"}, + name: "MitmdAnthropicNonDefaultPort", + providerHosts: []string{aibridgeproxyd.HostAnthropic}, + allowedPorts: []string{"8443"}, buildTargetURL: func(_ *url.URL) (string, error) { return "https://api.anthropic.com:8443/v1/messages", nil }, @@ -1394,9 +1295,9 @@ func TestProxy_MITM(t *testing.T) { provider: "anthropic", }, { - name: "MitmdOpenAI", - domainAllowlist: []string{aibridgeproxyd.HostOpenAI}, - allowedPorts: []string{"443"}, + name: "MitmdOpenAI", + providerHosts: []string{aibridgeproxyd.HostOpenAI}, + allowedPorts: []string{"443"}, buildTargetURL: func(_ *url.URL) (string, error) { return "https://api.openai.com/v1/chat/completions", nil }, @@ -1404,9 +1305,9 @@ func TestProxy_MITM(t *testing.T) { provider: "openai", }, { - name: "MitmdOpenAINonDefaultPort", - domainAllowlist: []string{aibridgeproxyd.HostOpenAI}, - allowedPorts: []string{"8443"}, + name: "MitmdOpenAINonDefaultPort", + providerHosts: []string{aibridgeproxyd.HostOpenAI}, + allowedPorts: []string{"8443"}, buildTargetURL: func(_ *url.URL) (string, error) { return "https://api.openai.com:8443/v1/chat/completions", nil }, @@ -1414,9 +1315,9 @@ func TestProxy_MITM(t *testing.T) { provider: "openai", }, { - name: "TunneledUnknownHost", - domainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - allowedPorts: nil, // will use tunneledURL.Port() + name: "TunneledUnknownHost", + providerHosts: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + allowedPorts: nil, // will use tunneledURL.Port() buildTargetURL: func(tunneledURL *url.URL) (string, error) { return url.JoinPath(tunneledURL.String(), "/some/path") }, @@ -1458,18 +1359,17 @@ func TestProxy_MITM(t *testing.T) { allowedPorts = []string{tunneledURL.Port()} } - // Configure domain allowlist. - domainAllowlist := tt.domainAllowlist - if domainAllowlist == nil { - domainAllowlist = []string{tunneledURL.Hostname()} + // Configure provider hosts. + providerHosts := tt.providerHosts + if providerHosts == nil { + providerHosts = []string{tunneledURL.Hostname()} } // Start the proxy server pointing to our mock aibridged. srv := newTestProxy(t, withCoderAccessURL(aibridgedServer.URL), withAllowedPorts(allowedPorts...), - withDomainAllowlist(domainAllowlist...), - withAIBridgeProviderFromHost(testProviderFromHost), + withProviderHosts(providerHosts...), withMetrics(metrics), ) @@ -1607,8 +1507,7 @@ func TestProxy_MITM_BYOKInjection(t *testing.T) { srv := newTestProxy(t, withCoderAccessURL(aibridgedServer.URL), - withDomainAllowlist(aibridgeproxyd.HostCopilot), - withAIBridgeProviderFromHost(testProviderFromHost), + withProviderHosts(aibridgeproxyd.HostCopilot), ) certPool := getProxyCertPool(t) @@ -1687,8 +1586,8 @@ func TestListenerTLS(t *testing.T) { withAllowedPorts(targetURL.Port()), ) if tt.tunneled { - // Use a domain allowlist that excludes the target server so requests are tunneled. - proxyOpts = append(proxyOpts, withDomainAllowlist(aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI)) + // Configure provider hosts that exclude the target server so requests are tunneled. + proxyOpts = append(proxyOpts, withProviderHosts(aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI)) } srv := newTestProxy(t, proxyOpts...) @@ -1791,14 +1690,10 @@ func TestServeCACert_CompoundPEM(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: compoundCertFile, - MITMKeyFile: keyFile, - DomainAllowlist: []string{"127.0.0.1", "localhost"}, - AIBridgeProviderFromHost: func(host string) string { - return "test-provider" - }, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: compoundCertFile, + MITMKeyFile: keyFile, }) require.NoError(t, err) t.Cleanup(func() { _ = srv.Close() }) @@ -1849,8 +1744,8 @@ func TestUpstreamProxy(t *testing.T) { name string // tunneled determines whether the request should be tunneled through // the upstream proxy (true) or MITM'd by aiproxy (false). - // When true, the target domain is NOT in the allowlist. - // When false, the target domain IS in the allowlist. + // When true, the target domain has no configured provider. + // When false, the target domain has a configured provider. tunneled bool // upstreamProxyTLS determines whether the upstream proxy uses TLS. // When true, aiproxy must be configured with the upstream proxy's CA. @@ -1865,7 +1760,7 @@ func TestUpstreamProxy(t *testing.T) { upstreamProxyAuth string }{ { - name: "NonAllowlistedDomain_TunneledToHTTPUpstreamProxy", + name: "NonProviderHost_TunneledToHTTPUpstreamProxy", tunneled: true, upstreamProxyTLS: false, buildTargetURL: func(finalDestinationURL *url.URL) string { @@ -1873,7 +1768,7 @@ func TestUpstreamProxy(t *testing.T) { }, }, { - name: "NonAllowlistedDomain_TunneledToHTTPSUpstreamProxy", + name: "NonProviderHost_TunneledToHTTPSUpstreamProxy", tunneled: true, upstreamProxyTLS: true, buildTargetURL: func(finalDestinationURL *url.URL) string { @@ -1881,7 +1776,7 @@ func TestUpstreamProxy(t *testing.T) { }, }, { - name: "NonAllowlistedDomain_TunneledToHTTPUpstreamProxyWithAuth", + name: "NonProviderHost_TunneledToHTTPUpstreamProxyWithAuth", tunneled: true, upstreamProxyTLS: false, upstreamProxyAuth: "proxyuser:proxypass", @@ -1890,7 +1785,7 @@ func TestUpstreamProxy(t *testing.T) { }, }, { - name: "NonAllowlistedDomain_TunneledToHTTPUpstreamProxyWithUsernameOnly", + name: "NonProviderHost_TunneledToHTTPUpstreamProxyWithUsernameOnly", tunneled: true, upstreamProxyTLS: false, upstreamProxyAuth: "proxyuser", @@ -1899,7 +1794,7 @@ func TestUpstreamProxy(t *testing.T) { }, }, { - name: "NonAllowlistedDomain_TunneledToHTTPUpstreamProxyWithUsernameAndColon", + name: "NonProviderHost_TunneledToHTTPUpstreamProxyWithUsernameAndColon", tunneled: true, upstreamProxyTLS: false, upstreamProxyAuth: "proxyuser:", @@ -1908,7 +1803,7 @@ func TestUpstreamProxy(t *testing.T) { }, }, { - name: "NonAllowlistedDomain_TunneledToHTTPUpstreamProxyWithTokenAuth", + name: "NonProviderHost_TunneledToHTTPUpstreamProxyWithTokenAuth", tunneled: true, upstreamProxyTLS: false, upstreamProxyAuth: ":proxypass", @@ -1917,7 +1812,7 @@ func TestUpstreamProxy(t *testing.T) { }, }, { - name: "AllowlistedDomain_MITMByAIProxy", + name: "ProviderHost_MITMByAIProxy", tunneled: false, upstreamProxyTLS: false, buildTargetURL: func(_ *url.URL) string { @@ -2057,10 +1952,10 @@ func TestUpstreamProxy(t *testing.T) { parsedTargetURL, err := url.Parse(targetURL) require.NoError(t, err) - // Configure allowlist based on test case: - // - For tunneled requests, api.anthropic.com is in allowlist, but we target a different host. - // - For MITM, api.anthropic.com must be in the allowlist. - domainAllowlist := []string{aibridgeproxyd.HostAnthropic} + // Configure provider hosts based on test case: + // - For tunneled requests, api.anthropic.com has a configured provider, but we target a different host. + // - For MITM, api.anthropic.com must have a configured provider. + providerHosts := []string{aibridgeproxyd.HostAnthropic} // Build upstream proxy URL with optional auth credentials. upstreamProxyURLStr := upstreamProxy.URL @@ -2073,10 +1968,9 @@ func TestUpstreamProxy(t *testing.T) { // Create aiproxy with upstream proxy configured. proxyOpts := []testProxyOption{ withCoderAccessURL(aibridgeServer.URL), - withDomainAllowlist(domainAllowlist...), + withProviderHosts(providerHosts...), withUpstreamProxy(upstreamProxyURLStr), withAllowedPorts("80", "443", parsedTargetURL.Port()), - withAIBridgeProviderFromHost(testProviderFromHost), } if upstreamProxyCAFile != "" { proxyOpts = append(proxyOpts, withUpstreamProxyCA(upstreamProxyCAFile)) @@ -2114,7 +2008,7 @@ func TestUpstreamProxy(t *testing.T) { // Verify the request flow based on test case. if tt.tunneled { require.True(t, upstreamProxyCONNECTReceived, - "upstream proxy should receive CONNECT for non-allowlisted domain") + "upstream proxy should receive CONNECT for non-provider-host") require.Equal(t, finalDestinationURL.Host, upstreamProxyCONNECTHost, "upstream proxy should receive CONNECT to correct host") require.True(t, finalDestinationReceived, @@ -2124,12 +2018,12 @@ func TestUpstreamProxy(t *testing.T) { require.Equal(t, requestBody, finalDestinationBody, "final destination should receive the exact request body") require.False(t, aibridgeReceived, - "aibridge should NOT receive request for non-allowlisted domain") + "aibridge should NOT receive request for non-provider-host") require.Empty(t, aibridgeAuthz, "tunneled requests should not reach aibridge") } else { require.False(t, upstreamProxyCONNECTReceived, - "upstream proxy should NOT receive CONNECT for allowlisted domain") + "upstream proxy should NOT receive CONNECT for provider host") require.True(t, aibridgeReceived, "aibridge should receive the MITM'd request") require.Equal(t, tt.expectedAIBridgePath, aibridgePath, @@ -2141,7 +2035,7 @@ func TestUpstreamProxy(t *testing.T) { require.Equal(t, requestBody, aibridgeBody, "aibridge should receive the exact request body") require.False(t, finalDestinationReceived, - "final destination should NOT receive request for allowlisted domain") + "final destination should NOT receive request for provider host") } // Verify upstream proxy authentication if configured. @@ -2155,7 +2049,7 @@ func TestUpstreamProxy(t *testing.T) { } // TestProxy_MITM_CustomProvider verifies that a non-builtin provider -// (e.g. OpenRouter) whose domain is added to the allowlist is correctly +// (e.g. OpenRouter) whose domain is registered as a provider host is correctly // MITM'd and routed through the proxy to the bridge endpoint. func TestProxy_MITM_CustomProvider(t *testing.T) { t.Parallel() @@ -2177,16 +2071,16 @@ func TestProxy_MITM_CustomProvider(t *testing.T) { })) t.Cleanup(aibridgedServer.Close) - // Wire the custom domain and provider mapping directly, as the - // real daemon would after calling domainsFromProviders. + // Wire the custom domain and provider mapping directly via + // withProviders, equivalent to the snapshot the daemon's Reload + // builds from classified providers in production. srv := newTestProxy(t, withCoderAccessURL(aibridgedServer.URL), - withDomainAllowlist(openrouterDomain), - withAIBridgeProviderFromHost(func(host string) string { - if host == openrouterDomain { - return openrouterProvider - } - return "" + withProviders(aibridgeproxyd.ReloadedProvider{ + Name: openrouterProvider, + Type: "openai", + Host: openrouterDomain, + Status: aibridgeproxyd.ProviderStatusEnabled, }), ) @@ -2307,10 +2201,10 @@ func TestProxy_PrivateIPBlocking(t *testing.T) { // Build the CONNECT target using the configured hostname. connectTarget := fmt.Sprintf("%s:%s", tt.targetHostname, targetURL.Port()) - // Use a domain allowlist that excludes the target so CONNECT requests + // Configure provider hosts that exclude the target so CONNECT requests // go through the tunnel path rather than being MITM'd. opts := []testProxyOption{ - withDomainAllowlist(aibridgeproxyd.HostAnthropic), + withProviderHosts(aibridgeproxyd.HostAnthropic), withAllowedPorts(targetURL.Port()), } @@ -2395,8 +2289,7 @@ func TestProxy_APIDump(t *testing.T) { srv := newTestProxy(t, withCoderAccessURL(aibridgedServer.URL), withAllowedPorts("443"), - withDomainAllowlist(aibridgeproxyd.HostAnthropic), - withAIBridgeProviderFromHost(testProviderFromHost), + withProviderHosts(aibridgeproxyd.HostAnthropic), withNewDumper(func(provider, requestID string) aibridgeproxyd.RoundTripDumper { dumpedProvider = provider dumpedRequestID = requestID @@ -2443,8 +2336,7 @@ func TestProxy_APIDump_ErrorsDoNotAffectProxy(t *testing.T) { srv := newTestProxy(t, withCoderAccessURL(aibridgedServer.URL), withAllowedPorts("443"), - withDomainAllowlist(aibridgeproxyd.HostAnthropic), - withAIBridgeProviderFromHost(testProviderFromHost), + withProviderHosts(aibridgeproxyd.HostAnthropic), withNewDumper(func(_, _ string) aibridgeproxyd.RoundTripDumper { return &failingDumper{} }), diff --git a/enterprise/aibridgeproxyd/reload.go b/enterprise/aibridgeproxyd/reload.go index d235a9be3b..686bab1130 100644 --- a/enterprise/aibridgeproxyd/reload.go +++ b/enterprise/aibridgeproxyd/reload.go @@ -3,7 +3,6 @@ package aibridgeproxyd import ( "context" "net/http" - "net/url" "slices" "strings" @@ -13,23 +12,72 @@ import ( "cdr.dev/slog/v3" ) +// ProviderStatus describes the lifecycle state of a configured AI +// provider for observability and routing purposes. +type ProviderStatus string + +const ( + // ProviderStatusEnabled means the provider is configured, valid, and + // included in the active routing snapshot. + ProviderStatusEnabled ProviderStatus = "enabled" + // ProviderStatusDisabled means the provider exists in configuration + // but is intentionally turned off by an operator. + ProviderStatusDisabled ProviderStatus = "disabled" + // ProviderStatusError means the provider exists in configuration but + // cannot be routed to because of a validation failure (missing or + // invalid base URL, duplicate host, etc.). + ProviderStatusError ProviderStatus = "error" +) + +// ReloadedProvider is one row from the provider configuration together +// with the outcome of evaluating it for routing. Host is populated only +// when Status == ProviderStatusEnabled; Err is populated only when +// Status == ProviderStatusError. +type ReloadedProvider struct { + Name string + Type string + Host string + Status ProviderStatus + Err error +} + +// ProviderReload is the result of a single refresh pass: every +// configured provider with its classification. +type ProviderReload struct { + Providers []ReloadedProvider +} + +// RefreshProvidersFunc returns the live provider classification used by +// Reload to rebuild the proxy's routing snapshot. +type RefreshProvidersFunc func(ctx context.Context) (ProviderReload, error) + // Reload refreshes proxy routing from the configured provider source. // A refresh failure leaves the previous snapshot in place. func (s *Server) Reload(ctx context.Context) error { if s.refreshProviders == nil { return nil } - providers, err := s.refreshProviders(ctx) + reload, err := s.refreshProviders(ctx) if err != nil { return xerrors.Errorf("refresh ai providers for proxy routing: %w", err) } - router, err := buildProviderRouter(ctx, s.logger, providers, s.allowedPorts) + router, err := buildProviderRouter(reload, s.allowedPorts) if err != nil { - return xerrors.Errorf("build provider router (provider_count=%d): %w", len(providers), err) + return xerrors.Errorf("build provider router (provider_count=%d): %w", len(reload.Providers), err) } s.providerRouter.Store(router) + for _, p := range reload.Providers { + if p.Status == ProviderStatusError { + s.logger.Warn(s.ctx, "provider excluded from routing", + slog.F("provider", p.Name), + slog.Error(p.Err), + ) + } + } s.logger.Debug(s.ctx, "aibridgeproxyd router reloaded", + slog.F("provider_count", len(reload.Providers)), slog.F("mitm_host_count", len(router.mitmHosts)), + slog.F("mitm_hosts", router.mitmHosts), ) return nil } @@ -42,7 +90,7 @@ func (s *Server) loadProviderRouter() *providerRouter { } // mitmHostsCondition returns a goproxy ReqConditionFunc that reads the -// allowlist from the atomic router on every match. Using a closure +// MITM host set from the atomic router on every match. Using a closure // instead of goproxy.ReqHostIs(...) lets Reload affect every later // CONNECT without re-registering handlers. func (s *Server) mitmHostsCondition() goproxy.ReqConditionFunc { @@ -54,35 +102,23 @@ func (s *Server) mitmHostsCondition() goproxy.ReqConditionFunc { } } -// buildProviderRouter constructs a router snapshot from a refreshed -// provider list. First provider wins on duplicate hostnames. -func buildProviderRouter(ctx context.Context, logger slog.Logger, providers []ProviderRoute, allowedPorts []string) (*providerRouter, error) { - nameByHost := make(map[string]string, len(providers)) - var domains []string - for _, p := range providers { - if p.BaseURL == "" { - logger.Warn(ctx, "skipping ai provider without base url", - slog.F("provider_name", p.Name), - ) +// buildProviderRouter constructs a router snapshot from a classified +// provider reload. Only providers with Status == ProviderStatusEnabled +// are included in the active routing tables; the refresh function is +// responsible for classifying disabled and errored rows. First entry +// wins on duplicate hostnames as a defense-in-depth measure even though +// the refresh function should mark duplicates as errors. +func buildProviderRouter(reload ProviderReload, allowedPorts []string) (*providerRouter, error) { + nameByHost := make(map[string]string, len(reload.Providers)) + domains := make([]string, 0, len(reload.Providers)) + for _, p := range reload.Providers { + if p.Status != ProviderStatusEnabled { continue } - u, err := url.Parse(p.BaseURL) - if err != nil { - logger.Warn(ctx, "skipping ai provider with invalid base url", - slog.F("provider_name", p.Name), - slog.F("base_url", p.BaseURL), - slog.Error(err), - ) + host := strings.ToLower(p.Host) + if host == "" { continue } - if u.Hostname() == "" { - logger.Warn(ctx, "skipping ai provider base url without hostname", - slog.F("provider_name", p.Name), - slog.F("base_url", p.BaseURL), - ) - continue - } - host := strings.ToLower(u.Hostname()) if _, exists := nameByHost[host]; exists { continue } @@ -95,30 +131,3 @@ func buildProviderRouter(ctx context.Context, logger slog.Logger, providers []Pr } return &providerRouter{mitmHosts: mitmHosts, nameByHost: nameByHost}, nil } - -// buildBootRouter seeds the providerRouter from the boot-time inputs. -// The lookup function is consulted only for hosts in the allowlist; a -// nil function with an empty allowlist is fine and yields an empty -// router (the proxy fails closed until Reload populates it). -func buildBootRouter(domainAllowlist []string, providerFromHost func(string) string, allowedPorts []string) (*providerRouter, error) { - mitmHosts, err := convertDomainsToHosts(domainAllowlist, allowedPorts) - if err != nil { - return nil, xerrors.Errorf("invalid domain allowlist: %w", err) - } - nameByHost := make(map[string]string, len(domainAllowlist)) - for _, domain := range domainAllowlist { - domain = strings.TrimSpace(strings.ToLower(domain)) - if domain == "" { - continue - } - var name string - if providerFromHost != nil { - name = providerFromHost(domain) - } - if name == "" { - return nil, xerrors.Errorf("domain %q is in allowlist but has no provider mapping", domain) - } - nameByHost[domain] = name - } - return &providerRouter{mitmHosts: mitmHosts, nameByHost: nameByHost}, nil -} diff --git a/enterprise/aibridgeproxyd/reload_internal_test.go b/enterprise/aibridgeproxyd/reload_internal_test.go index 7a8b0f9fae..fb985445f3 100644 --- a/enterprise/aibridgeproxyd/reload_internal_test.go +++ b/enterprise/aibridgeproxyd/reload_internal_test.go @@ -12,17 +12,26 @@ import ( "github.com/coder/coder/v2/testutil" ) +func enabledProvider(name, host string) ReloadedProvider { + return ReloadedProvider{ + Name: name, + Type: "openai", + Host: host, + Status: ProviderStatusEnabled, + } +} + func TestServerReloadSwapsProviderRouter(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) - providers := []ProviderRoute{{Name: "old", BaseURL: "https://old.example.com/"}} + reload := ProviderReload{Providers: []ReloadedProvider{enabledProvider("old", "old.example.com")}} srv := &Server{ ctx: ctx, logger: slogtest.Make(t, nil), allowedPorts: []string{"443"}, - refreshProviders: func(context.Context) ([]ProviderRoute, error) { - return providers, nil + refreshProviders: func(context.Context) (ProviderReload, error) { + return reload, nil }, } srv.providerRouter.Store(emptyProviderRouter) @@ -31,7 +40,7 @@ func TestServerReloadSwapsProviderRouter(t *testing.T) { assert.Equal(t, "old", srv.loadProviderRouter().providerFromHost("old.example.com")) assert.Empty(t, srv.loadProviderRouter().providerFromHost("new.example.com")) - providers = []ProviderRoute{{Name: "new", BaseURL: "https://new.example.com/"}} + reload = ProviderReload{Providers: []ReloadedProvider{enabledProvider("new", "new.example.com")}} require.NoError(t, srv.Reload(ctx)) router := srv.loadProviderRouter() @@ -45,17 +54,17 @@ func TestServerReloadPreservesProviderRouterOnRefreshError(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) refreshErr := xerrors.New("refresh failed") - providers := []ProviderRoute{{Name: "old", BaseURL: "https://old.example.com/"}} + reload := ProviderReload{Providers: []ReloadedProvider{enabledProvider("old", "old.example.com")}} failRefresh := false srv := &Server{ ctx: ctx, logger: slogtest.Make(t, nil), allowedPorts: []string{"443"}, - refreshProviders: func(context.Context) ([]ProviderRoute, error) { + refreshProviders: func(context.Context) (ProviderReload, error) { if failRefresh { - return nil, refreshErr + return ProviderReload{}, refreshErr } - return providers, nil + return reload, nil }, } srv.providerRouter.Store(emptyProviderRouter) @@ -73,75 +82,84 @@ func TestServerReloadPreservesProviderRouterOnRefreshError(t *testing.T) { assert.Equal(t, []string{"old.example.com:443"}, after.mitmHosts) } -// TestBuildProviderRouter covers the host-and-routing derivation that -// Reload feeds into the providerRouter. +// TestBuildProviderRouter covers the host-and-routing derivation from +// the classified provider reload. func TestBuildProviderRouter(t *testing.T) { t.Parallel() - t.Run("ExtractsHostnames", func(t *testing.T) { + t.Run("IncludesEnabledOnly", func(t *testing.T) { t.Parallel() - providers := []ProviderRoute{ - {Name: "openai", BaseURL: "https://api.openai.com/v1/"}, - {Name: "anthropic", BaseURL: "https://api.anthropic.com/"}, - {Name: "custom", BaseURL: "https://custom-llm.example.com:8443/api"}, - } + reload := ProviderReload{Providers: []ReloadedProvider{ + enabledProvider("openai", "api.openai.com"), + enabledProvider("anthropic", "api.anthropic.com"), + enabledProvider("custom", "custom-llm.example.com"), + // Host is populated on the non-enabled rows so the Status + // guard, not the empty-host guard, is what excludes them. + {Name: "off", Type: "openai", Host: "disabled.example.com", Status: ProviderStatusDisabled}, + {Name: "bad", Type: "openai", Host: "errored.example.com", Status: ProviderStatusError, Err: xerrors.New("nope")}, + }} - router, err := buildProviderRouter(testutil.Context(t, testutil.WaitShort), slogtest.Make(t, nil), providers, []string{"443"}) + router, err := buildProviderRouter(reload, []string{"443"}) require.NoError(t, err) assert.Equal(t, "openai", router.providerFromHost("api.openai.com")) assert.Equal(t, "anthropic", router.providerFromHost("api.anthropic.com")) assert.Equal(t, "custom", router.providerFromHost("custom-llm.example.com")) assert.Empty(t, router.providerFromHost("unknown.com")) + assert.Empty(t, router.providerFromHost("disabled.example.com"), + "disabled provider must not be routable even with a populated Host") + assert.Empty(t, router.providerFromHost("errored.example.com"), + "errored provider must not be routable even with a populated Host") assert.Contains(t, router.mitmHosts, "api.openai.com:443") assert.Contains(t, router.mitmHosts, "api.anthropic.com:443") - }) - - t.Run("DeduplicatesSameHost", func(t *testing.T) { - t.Parallel() - - providers := []ProviderRoute{ - {Name: "first", BaseURL: "https://api.example.com/v1"}, - {Name: "second", BaseURL: "https://api.example.com/v2"}, - } - - router, err := buildProviderRouter(testutil.Context(t, testutil.WaitShort), slogtest.Make(t, nil), providers, []string{"443"}) - require.NoError(t, err) - - // First provider wins on duplicate host. - assert.Equal(t, "first", router.providerFromHost("api.example.com")) + assert.Len(t, router.mitmHosts, 3) }) t.Run("CaseInsensitive", func(t *testing.T) { t.Parallel() - providers := []ProviderRoute{ - {Name: "provider", BaseURL: "https://API.Example.COM/v1"}, - } + reload := ProviderReload{Providers: []ReloadedProvider{ + {Name: "provider", Type: "openai", Host: "API.Example.COM", Status: ProviderStatusEnabled}, + }} - router, err := buildProviderRouter(testutil.Context(t, testutil.WaitShort), slogtest.Make(t, nil), providers, []string{"443"}) + router, err := buildProviderRouter(reload, []string{"443"}) require.NoError(t, err) assert.Equal(t, "provider", router.providerFromHost("API.Example.COM")) assert.Equal(t, "provider", router.providerFromHost("api.example.com")) }) - t.Run("SkipsEmptyOrMalformedBaseURL", func(t *testing.T) { + t.Run("DefensiveDeduplicatesSameHost", func(t *testing.T) { t.Parallel() - providers := []ProviderRoute{ - {Name: "no-url"}, - {Name: "scheme-only", BaseURL: "https://"}, - {Name: "good", BaseURL: "https://api.good.example.com/"}, - } + // Refresh function should mark the duplicate as ProviderStatusError; + // buildProviderRouter is defensive and tolerates an enabled duplicate + // by giving the first entry the host (first wins). + reload := ProviderReload{Providers: []ReloadedProvider{ + enabledProvider("first", "api.example.com"), + enabledProvider("second", "api.example.com"), + }} - router, err := buildProviderRouter(testutil.Context(t, testutil.WaitShort), slogtest.Make(t, nil), providers, []string{"443"}) + router, err := buildProviderRouter(reload, []string{"443"}) + require.NoError(t, err) + + assert.Equal(t, "first", router.providerFromHost("api.example.com")) + }) + + t.Run("SkipsRowsWithEmptyHost", func(t *testing.T) { + t.Parallel() + + reload := ProviderReload{Providers: []ReloadedProvider{ + {Name: "no-host", Type: "openai", Status: ProviderStatusEnabled}, + enabledProvider("good", "api.good.example.com"), + }} + + router, err := buildProviderRouter(reload, []string{"443"}) require.NoError(t, err) assert.Equal(t, "good", router.providerFromHost("api.good.example.com")) - assert.Empty(t, router.providerFromHost("scheme-only")) assert.Equal(t, []string{"api.good.example.com:443"}, router.mitmHosts) }) } diff --git a/enterprise/aibridgeproxyd/reload_test.go b/enterprise/aibridgeproxyd/reload_test.go index d51aa5bc98..e55d45a372 100644 --- a/enterprise/aibridgeproxyd/reload_test.go +++ b/enterprise/aibridgeproxyd/reload_test.go @@ -5,6 +5,7 @@ import ( "io" "net/http" "net/http/httptest" + "net/url" "slices" "strings" "sync" @@ -55,16 +56,24 @@ func (r *aibridgedRecorder) reset() { r.path = "" } -// providerStore is a mutable [aibridgeproxyd.RefreshProvidersFunc] -// backing for integration tests. set / setErr mutate the snapshot -// returned by the next Reload, mimicking CRUD against the database. +// rawProvider is a (name, base URL) pair representing what the database +// holds before classification, mirroring the ai_providers row shape +// that the production refresh function classifies. +type rawProvider struct { + name string + baseURL string +} + +// providerStore is a mutable RefreshProvidersFunc backing for +// integration tests. set / setErr mutate the snapshot returned by the +// next Reload, mimicking CRUD against the database. type providerStore struct { mu sync.Mutex - providers []aibridgeproxyd.ProviderRoute + providers []rawProvider err error } -func (s *providerStore) set(providers []aibridgeproxyd.ProviderRoute) { +func (s *providerStore) set(providers []rawProvider) { s.mu.Lock() defer s.mu.Unlock() s.providers = providers @@ -77,20 +86,59 @@ func (s *providerStore) setErr(err error) { s.err = err } -func (s *providerStore) refresh(context.Context) ([]aibridgeproxyd.ProviderRoute, error) { +func (s *providerStore) refresh(context.Context) (aibridgeproxyd.ProviderReload, error) { s.mu.Lock() defer s.mu.Unlock() if s.err != nil { - return nil, s.err + return aibridgeproxyd.ProviderReload{}, s.err } - // Return a copy so callers can't mutate our internal snapshot. - return slices.Clone(s.providers), nil + providers := slices.Clone(s.providers) + reload := aibridgeproxyd.ProviderReload{ + Providers: make([]aibridgeproxyd.ReloadedProvider, 0, len(providers)), + } + seenHost := make(map[string]string, len(providers)) + for _, p := range providers { + reload.Providers = append(reload.Providers, classifyRaw(p, seenHost)) + } + return reload, nil } -// newReloadTestHarness boots a proxy with an empty boot allowlist and a -// store-backed RefreshProviders. Production wiring is identical: the -// daemon constructs the proxy without a static allowlist and lets -// Reload populate the router from the database. +// classifyRaw mirrors the production classifier in enterprise/cli so +// the reload tests exercise the same validation rules end-to-end. +func classifyRaw(p rawProvider, seenHost map[string]string) aibridgeproxyd.ReloadedProvider { + out := aibridgeproxyd.ReloadedProvider{Name: p.name, Type: "openai"} + if strings.TrimSpace(p.baseURL) == "" { + out.Status = aibridgeproxyd.ProviderStatusError + out.Err = xerrors.New("base url is empty") + return out + } + u, err := url.Parse(p.baseURL) + if err != nil { + out.Status = aibridgeproxyd.ProviderStatusError + out.Err = xerrors.Errorf("invalid base url %q: %w", p.baseURL, err) + return out + } + host := strings.ToLower(u.Hostname()) + if host == "" { + out.Status = aibridgeproxyd.ProviderStatusError + out.Err = xerrors.Errorf("base url %q has no hostname", p.baseURL) + return out + } + if claimedBy, taken := seenHost[host]; taken { + out.Status = aibridgeproxyd.ProviderStatusError + out.Err = xerrors.Errorf("hostname %q already claimed by provider %q", host, claimedBy) + return out + } + seenHost[host] = p.name + out.Host = host + out.Status = aibridgeproxyd.ProviderStatusEnabled + return out +} + +// newReloadTestHarness boots a proxy with an empty initial router and +// a store-backed RefreshProviders. Production wiring is identical: the +// daemon constructs the proxy without preconfigured provider hosts and +// lets Reload populate the router from the database. func newReloadTestHarness(t *testing.T) *reloadTestHarness { t.Helper() @@ -106,10 +154,6 @@ func newReloadTestHarness(t *testing.T) *reloadTestHarness { srv := newTestProxy(t, withCoderAccessURL(bridged.URL), withAllowedPorts("443"), - // Empty boot allowlist: the router must be populated by Reload, - // matching the production daemon's behavior. - withDomainAllowlist(), - withAIBridgeProviderFromHost(nil), withRefreshProviders(store.refresh), ) @@ -192,16 +236,128 @@ func (h *reloadTestHarness) expectNotRouted(t *testing.T, targetURL string) { "aibridged must not be reached for non-routed host %s", targetURL) } +// TestProxy_StaleTunnelStopsRoutingAfterProviderChange is the +// regression test for a bug where a long-lived CONNECT tunnel that was +// established while a provider was enabled kept routing decrypted +// requests to aibridged after the provider was disabled or renamed. The +// fix re-validates the CONNECT-time provider against the live router on +// every decrypted request and covers both shapes of stale mapping: +// +// - ProviderDisabled: liveProvider == "" (host no longer MITM'd). +// - ProviderRenamed: liveProvider != reqCtx.Provider (host MITM'd, but +// under a new provider name). +func TestProxy_StaleTunnelStopsRoutingAfterProviderChange(t *testing.T) { + t.Parallel() + + testCases := []struct { + name string + // applyChange mutates the store to simulate the provider change + // after the initial routed request succeeds. + applyChange func(*providerStore) + // changeDescription is appended to the second-request assertion + // message so a failure points at the exercised branch. + changeDescription string + }{ + { + name: "ProviderDisabled", + applyChange: func(s *providerStore) { s.set(nil) }, + changeDescription: "after alpha was disabled", + }, + { + name: "ProviderRenamed", + applyChange: func(s *providerStore) { + // Same host, new provider name: the live router still + // MITMs alpha.invalid, but as "alpha-v2". The stale + // CONNECT-time name "alpha" no longer matches. + s.set([]rawProvider{ + {name: "alpha-v2", baseURL: "https://alpha.invalid/v1"}, + }) + }, + changeDescription: "after alpha was renamed to alpha-v2", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + recorder := &aibridgedRecorder{} + bridged := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + recorder.record(r.URL.Path) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("aibridged")) + })) + t.Cleanup(bridged.Close) + + store := &providerStore{} + store.set([]rawProvider{ + {name: "alpha", baseURL: "https://alpha.invalid/v1"}, + }) + + // newTestProxy seeds the router from the store via the + // initial Reload, so the first CONNECT is MITM'd as alpha. + srv := newTestProxy(t, + withCoderAccessURL(bridged.URL), + withAllowedPorts("443"), + withRefreshProviders(store.refresh), + ) + + certPool := getProxyCertPool(t) + client := newProxyClient(t, srv, makeProxyAuthHeader("coder-token"), certPool, false) + // Keep-alives are required: the regression exists only when a + // subsequent request reuses the original CONNECT tunnel. A fresh + // CONNECT would correctly observe the post-reload router. + transport := client.Transport.(*http.Transport) + transport.DisableKeepAlives = false + transport.MaxConnsPerHost = 1 + transport.MaxIdleConnsPerHost = 1 + + sendThroughTunnel := func(path string) (status int, err error) { + ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitShort) + defer cancel() + req, reqErr := http.NewRequestWithContext(ctx, http.MethodPost, "https://alpha.invalid"+path, strings.NewReader(`{}`)) + require.NoError(t, reqErr) + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + return 0, err + } + defer resp.Body.Close() + _, _ = io.Copy(io.Discard, resp.Body) + return resp.StatusCode, nil + } + + // First request: alpha is enabled, the proxy MITMs and routes to + // aibridged under the alpha namespace. + recorder.reset() + status, err := sendThroughTunnel("/v1/messages") + require.NoError(t, err) + require.Equal(t, http.StatusOK, status) + require.Equal(t, "/api/v2/aibridge/alpha/v1/messages", recorder.load(), + "first request must be routed to aibridged while alpha is enabled") + + // Apply the provider change and reload. The atomic router swap + // takes effect immediately, but the client's connection (and + // the proxy's hijacked tunnel) remain open. + tc.applyChange(store) + require.NoError(t, srv.Reload(t.Context())) + + // Second request on the same tunnel: aibridged must NOT see it. + // The connection is hijacked so the request reaches the proxy's + // handleRequest with the stale CONNECT-time provider; the fix + // re-validates against the live router and passes through to + // the original upstream (alpha.invalid, which fails DNS). + recorder.reset() + _, _ = sendThroughTunnel("/v1/should-not-route") + require.Empty(t, recorder.load(), + "%s, aibridged must not receive the request even on a reused tunnel", tc.changeDescription) + }) + } +} + // TestProxy_HotReloadRoutingCRUD drives the proxy through a CRUD-style // sequence of provider changes and asserts on routing after each -// Reload via real HTTPS requests. Each sub-test mutates the store and -// validates that: -// - newly created providers are MITM'd to aibridged with the right -// /api/v2/aibridge// -// - renamed providers route under the new name -// - providers whose BaseURL host changes route the new host and stop -// MITM'ing the old host -// - deleted providers stop being MITM'd; aibridged sees nothing +// Reload via real HTTPS requests. // // Hostnames are .invalid (RFC 2606) so a request that escapes the MITM // path fails fast via DNS rather than reaching a real upstream. @@ -210,30 +366,30 @@ func TestProxy_HotReloadRoutingCRUD(t *testing.T) { h := newReloadTestHarness(t) - // InitialEmptyRouter: no Reload has been called and the boot - // allowlist is empty, so any host falls through to the tunneled + // InitialEmptyRouter: no Reload has been called and no provider + // hosts are configured, so any host falls through to the tunneled // middleware. h.expectNotRouted(t, "https://alpha.invalid/v1/messages") // CreateProvider. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "alpha", BaseURL: "https://alpha.invalid/v1"}, + h.store.set([]rawProvider{ + {name: "alpha", baseURL: "https://alpha.invalid/v1"}, }) require.NoError(t, h.srv.Reload(t.Context())) h.expectRoutedTo(t, "https://alpha.invalid/v1/messages", "/api/v2/aibridge/alpha/v1/messages") // UpdateProviderName: the same BaseURL with a new name must route // under the new name on the next Reload. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "alpha-v2", BaseURL: "https://alpha.invalid/v1"}, + h.store.set([]rawProvider{ + {name: "alpha-v2", baseURL: "https://alpha.invalid/v1"}, }) require.NoError(t, h.srv.Reload(t.Context())) h.expectRoutedTo(t, "https://alpha.invalid/v1/messages", "/api/v2/aibridge/alpha-v2/v1/messages") // UpdateProviderBaseURLHost: moving the provider to a new host must // start MITM'ing the new host and stop MITM'ing the old one. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "alpha-v2", BaseURL: "https://alpha-new.invalid/v1"}, + h.store.set([]rawProvider{ + {name: "alpha-v2", baseURL: "https://alpha-new.invalid/v1"}, }) require.NoError(t, h.srv.Reload(t.Context())) h.expectRoutedTo(t, "https://alpha-new.invalid/v1/messages", "/api/v2/aibridge/alpha-v2/v1/messages") @@ -241,9 +397,9 @@ func TestProxy_HotReloadRoutingCRUD(t *testing.T) { // AddSecondProvider: a second provider added in the same Reload must // route independently from the first. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "alpha-v2", BaseURL: "https://alpha-new.invalid/v1"}, - {Name: "beta", BaseURL: "https://beta.invalid/v1"}, + h.store.set([]rawProvider{ + {name: "alpha-v2", baseURL: "https://alpha-new.invalid/v1"}, + {name: "beta", baseURL: "https://beta.invalid/v1"}, }) require.NoError(t, h.srv.Reload(t.Context())) h.expectRoutedTo(t, "https://alpha-new.invalid/v1/messages", "/api/v2/aibridge/alpha-v2/v1/messages") @@ -251,8 +407,8 @@ func TestProxy_HotReloadRoutingCRUD(t *testing.T) { // DeleteOneProvider: removing alpha must keep beta routed and stop // routing alpha. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "beta", BaseURL: "https://beta.invalid/v1"}, + h.store.set([]rawProvider{ + {name: "beta", baseURL: "https://beta.invalid/v1"}, }) require.NoError(t, h.srv.Reload(t.Context())) h.expectRoutedTo(t, "https://beta.invalid/v1/chat/completions", "/api/v2/aibridge/beta/v1/chat/completions") @@ -268,8 +424,8 @@ func TestProxy_HotReloadRoutingCRUD(t *testing.T) { // RecreateAfterDelete: reintroducing a previously-deleted provider // must route again without restart, confirming the swap is // symmetric. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "alpha", BaseURL: "https://alpha.invalid/v1"}, + h.store.set([]rawProvider{ + {name: "alpha", baseURL: "https://alpha.invalid/v1"}, }) require.NoError(t, h.srv.Reload(t.Context())) h.expectRoutedTo(t, "https://alpha.invalid/v1/messages", "/api/v2/aibridge/alpha/v1/messages") @@ -288,11 +444,11 @@ func TestProxy_HotReloadRoutingInvalidProviders(t *testing.T) { h := newReloadTestHarness(t) // One valid provider and one with an empty BaseURL. The empty - // entry must be silently dropped; the valid one must still - // route. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "no-url"}, - {Name: "valid", BaseURL: "https://valid.invalid/v1"}, + // entry must be classified as error and excluded from routing; + // the valid one must still route. + h.store.set([]rawProvider{ + {name: "no-url"}, + {name: "valid", baseURL: "https://valid.invalid/v1"}, }) require.NoError(t, h.srv.Reload(t.Context())) @@ -304,12 +460,12 @@ func TestProxy_HotReloadRoutingInvalidProviders(t *testing.T) { h := newReloadTestHarness(t) // A BaseURL that fails url.Parse and one whose Hostname() is - // empty must both be dropped. Mixed with a valid entry, only - // the valid one routes. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "malformed", BaseURL: "://not-a-url"}, - {Name: "no-host", BaseURL: "https://"}, - {Name: "valid", BaseURL: "https://valid.invalid/v1"}, + // empty must both be classified as error. Mixed with a valid + // entry, only the valid one routes. + h.store.set([]rawProvider{ + {name: "malformed", baseURL: "://not-a-url"}, + {name: "no-host", baseURL: "https://"}, + {name: "valid", baseURL: "https://valid.invalid/v1"}, }) require.NoError(t, h.srv.Reload(t.Context())) @@ -320,11 +476,11 @@ func TestProxy_HotReloadRoutingInvalidProviders(t *testing.T) { t.Parallel() h := newReloadTestHarness(t) - // Two providers with the same BaseURL host: the first one wins, - // matching buildProviderRouter's documented contract. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "first", BaseURL: "https://shared.invalid/v1"}, - {Name: "second", BaseURL: "https://shared.invalid/v2"}, + // Two providers with the same BaseURL host: the second is + // classified as error and excluded; the first routes. + h.store.set([]rawProvider{ + {name: "first", baseURL: "https://shared.invalid/v1"}, + {name: "second", baseURL: "https://shared.invalid/v2"}, }) require.NoError(t, h.srv.Reload(t.Context())) @@ -337,10 +493,10 @@ func TestProxy_HotReloadRoutingInvalidProviders(t *testing.T) { h := newReloadTestHarness(t) // When every provider is invalid, the router contains no // entries and the proxy fails closed: no host is MITM'd. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "no-url"}, - {Name: "malformed", BaseURL: "://not-a-url"}, - {Name: "no-host", BaseURL: "https://"}, + h.store.set([]rawProvider{ + {name: "no-url"}, + {name: "malformed", baseURL: "://not-a-url"}, + {name: "no-host", baseURL: "https://"}, }) require.NoError(t, h.srv.Reload(t.Context())) @@ -352,15 +508,15 @@ func TestProxy_HotReloadRoutingInvalidProviders(t *testing.T) { h := newReloadTestHarness(t) // Seed a valid snapshot so we have something to preserve. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "alpha", BaseURL: "https://alpha.invalid/v1"}, + h.store.set([]rawProvider{ + {name: "alpha", baseURL: "https://alpha.invalid/v1"}, }) require.NoError(t, h.srv.Reload(t.Context())) h.expectRoutedTo(t, "https://alpha.invalid/v1/messages", "/api/v2/aibridge/alpha/v1/messages") // A refresh error must NOT clear the router: dropping the - // allowlist on every transient DB hiccup would amplify the - // fault into a denial of service. + // provider host set on every transient DB hiccup would + // amplify the fault into a denial of service. h.store.setErr(xerrors.New("simulated db failure")) err := h.srv.Reload(t.Context()) require.Error(t, err) @@ -369,8 +525,8 @@ func TestProxy_HotReloadRoutingInvalidProviders(t *testing.T) { // Recovery: once the store returns providers again, the next // Reload applies the new snapshot. - h.store.set([]aibridgeproxyd.ProviderRoute{ - {Name: "beta", BaseURL: "https://beta.invalid/v1"}, + h.store.set([]rawProvider{ + {name: "beta", baseURL: "https://beta.invalid/v1"}, }) require.NoError(t, h.srv.Reload(t.Context())) h.expectRoutedTo(t, "https://beta.invalid/v1/messages", "/api/v2/aibridge/beta/v1/messages") diff --git a/enterprise/cli/aibridgeproxyd.go b/enterprise/cli/aibridgeproxyd.go index 0f7ba976a5..00cbefaee6 100644 --- a/enterprise/cli/aibridgeproxyd.go +++ b/enterprise/cli/aibridgeproxyd.go @@ -5,7 +5,9 @@ package cli import ( "context" "io" + "net/url" "path/filepath" + "strings" "github.com/prometheus/client_golang/prometheus" "golang.org/x/xerrors" @@ -86,19 +88,67 @@ func newAIBridgeProxyDaemon(coderAPI *coderd.API) (io.Closer, error) { }, nil } +// refreshProxyProviders classifies every ai_providers row as enabled, +// disabled, or error so the proxy router and any observers see the full +// configured set. Disabled rows are excluded from routing; errored rows +// are excluded from routing and surface their failure reason for +// metrics and logs. func refreshProxyProviders(db database.Store) aibridgeproxyd.RefreshProvidersFunc { - return func(ctx context.Context) ([]aibridgeproxyd.ProviderRoute, error) { + return func(ctx context.Context) (aibridgeproxyd.ProviderReload, error) { //nolint:gocritic // AsAIProviderMetadataReader is the correct subject for routing-only access. rows, err := db.GetAIProviders(dbauthz.AsAIProviderMetadataReader(ctx), database.GetAIProvidersParams{ - IncludeDisabled: false, + IncludeDisabled: true, }) if err != nil { - return nil, xerrors.Errorf("load ai providers: %w", err) + return aibridgeproxyd.ProviderReload{}, xerrors.Errorf("load ai providers: %w", err) } - out := make([]aibridgeproxyd.ProviderRoute, 0, len(rows)) + reload := aibridgeproxyd.ProviderReload{ + Providers: make([]aibridgeproxyd.ReloadedProvider, 0, len(rows)), + } + seenHost := make(map[string]string, len(rows)) for _, row := range rows { - out = append(out, aibridgeproxyd.ProviderRoute{Name: row.Name, BaseURL: row.BaseUrl}) + reload.Providers = append(reload.Providers, classifyProviderRow(row, seenHost)) } - return out, nil + return reload, nil } } + +// classifyProviderRow evaluates a single ai_providers row for routing. +// seenHost is mutated to track the first provider that claimed each +// hostname so later duplicates can be flagged as errors. +func classifyProviderRow(row database.AIProvider, seenHost map[string]string) aibridgeproxyd.ReloadedProvider { + out := aibridgeproxyd.ReloadedProvider{ + Name: row.Name, + Type: string(row.Type), + } + if !row.Enabled { + out.Status = aibridgeproxyd.ProviderStatusDisabled + return out + } + if strings.TrimSpace(row.BaseUrl) == "" { + out.Status = aibridgeproxyd.ProviderStatusError + out.Err = xerrors.New("base url is empty") + return out + } + u, err := url.Parse(row.BaseUrl) + if err != nil { + out.Status = aibridgeproxyd.ProviderStatusError + out.Err = xerrors.Errorf("invalid base url %q: %w", row.BaseUrl, err) + return out + } + host := strings.ToLower(u.Hostname()) + if host == "" { + out.Status = aibridgeproxyd.ProviderStatusError + out.Err = xerrors.Errorf("base url %q has no hostname", row.BaseUrl) + return out + } + if claimedBy, taken := seenHost[host]; taken { + out.Status = aibridgeproxyd.ProviderStatusError + out.Err = xerrors.Errorf("hostname %q already claimed by provider %q", host, claimedBy) + return out + } + seenHost[host] = row.Name + out.Host = host + out.Status = aibridgeproxyd.ProviderStatusEnabled + return out +} diff --git a/enterprise/cli/aibridgeproxyd_internal_test.go b/enterprise/cli/aibridgeproxyd_internal_test.go new file mode 100644 index 0000000000..54c6c25f78 --- /dev/null +++ b/enterprise/cli/aibridgeproxyd_internal_test.go @@ -0,0 +1,105 @@ +//go:build !slim + +package cli + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/enterprise/aibridgeproxyd" +) + +// TestClassifyProviderRow covers every branch of the classifier so the +// disabled, error, and enabled paths are exercised through the +// production code instead of relying on classifyRaw, the test mirror in +// reload_test.go. +func TestClassifyProviderRow(t *testing.T) { + t.Parallel() + + enabledRow := func(name, baseURL string) database.AIProvider { + return database.AIProvider{ + Name: name, + Type: database.AiProviderTypeOpenai, + Enabled: true, + BaseUrl: baseURL, + } + } + + t.Run("Enabled", func(t *testing.T) { + t.Parallel() + + seen := map[string]string{} + got := classifyProviderRow(enabledRow("openai", "https://api.openai.com/v1"), seen) + assert.Equal(t, "openai", got.Name) + assert.Equal(t, string(database.AiProviderTypeOpenai), got.Type) + assert.Equal(t, aibridgeproxyd.ProviderStatusEnabled, got.Status) + assert.Equal(t, "api.openai.com", got.Host) + assert.NoError(t, got.Err) + assert.Equal(t, "openai", seen["api.openai.com"]) + }) + + t.Run("DisabledRow", func(t *testing.T) { + t.Parallel() + + seen := map[string]string{} + row := enabledRow("off", "https://api.off.example.com/v1") + row.Enabled = false + got := classifyProviderRow(row, seen) + assert.Equal(t, aibridgeproxyd.ProviderStatusDisabled, got.Status) + assert.Empty(t, got.Host, "disabled provider must not claim a host") + assert.NoError(t, got.Err) + assert.Empty(t, seen, "disabled provider must not occupy a host slot") + }) + + t.Run("EmptyBaseURL", func(t *testing.T) { + t.Parallel() + + seen := map[string]string{} + got := classifyProviderRow(enabledRow("no-url", " "), seen) + assert.Equal(t, aibridgeproxyd.ProviderStatusError, got.Status) + assert.Empty(t, got.Host) + assert.ErrorContains(t, got.Err, "base url is empty") + }) + + t.Run("MalformedBaseURL", func(t *testing.T) { + t.Parallel() + + seen := map[string]string{} + got := classifyProviderRow(enabledRow("bad", "://not-a-url"), seen) + assert.Equal(t, aibridgeproxyd.ProviderStatusError, got.Status) + assert.ErrorContains(t, got.Err, "invalid base url") + }) + + t.Run("BaseURLWithoutHostname", func(t *testing.T) { + t.Parallel() + + seen := map[string]string{} + got := classifyProviderRow(enabledRow("no-host", "https://"), seen) + assert.Equal(t, aibridgeproxyd.ProviderStatusError, got.Status) + assert.ErrorContains(t, got.Err, "no hostname") + }) + + t.Run("DuplicateHostnameFirstWins", func(t *testing.T) { + t.Parallel() + + seen := map[string]string{} + first := classifyProviderRow(enabledRow("first", "https://shared.example.com/v1"), seen) + assert.Equal(t, aibridgeproxyd.ProviderStatusEnabled, first.Status) + + second := classifyProviderRow(enabledRow("second", "https://shared.example.com/v2"), seen) + assert.Equal(t, aibridgeproxyd.ProviderStatusError, second.Status) + assert.ErrorContains(t, second.Err, "already claimed by provider \"first\"") + assert.Equal(t, "first", seen["shared.example.com"], "first wins must not be overwritten") + }) + + t.Run("HostnameLowercased", func(t *testing.T) { + t.Parallel() + + seen := map[string]string{} + got := classifyProviderRow(enabledRow("mixed", "https://API.Example.COM/v1"), seen) + assert.Equal(t, aibridgeproxyd.ProviderStatusEnabled, got.Status) + assert.Equal(t, "api.example.com", got.Host) + }) +}