From 74b6d12a8ab3c2fe2a91985f92ff2b0a1575892a Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Tue, 13 Jan 2026 11:30:51 +0000 Subject: [PATCH] feat: implement selective MITM with configurable domain allowlist in aibridgeproxyd (#21473) ## Description Implements selective MITM (Man-in-the-Middle) in `aibridgeproxyd` so that only requests to allowlisted domains are intercepted and decrypted. Requests to all other domains are tunneled directly without decryption. ## Changes * New config option: `CODER_AIBRIDGE_PROXY_DOMAIN_ALLOWLIST` (default: `api.anthropic.com`,`api.openai.com`) * Selective MITM: Uses `goproxy.ReqHostIs()` to only intercept `CONNECT` requests to allowlisted hosts * Certificate caching: Now only generates/caches certificates for allowlisted domains * Validation: Startup fails if domain allowlist is empty or contains invalid entries Closes: https://github.com/coder/internal/issues/1182 --- cli/testdata/server-config.yaml.golden | 7 + coderd/apidoc/docs.go | 6 + coderd/apidoc/swagger.json | 6 + codersdk/deployment.go | 20 +- docs/reference/api/general.md | 3 + docs/reference/api/schemas.md | 25 +- enterprise/aibridgeproxyd/aibridgeproxyd.go | 100 +++- .../aibridgeproxyd/aibridgeproxyd_test.go | 474 ++++++++++++------ enterprise/cli/aibridgeproxyd.go | 9 +- site/src/api/typesGenerated.ts | 1 + 10 files changed, 473 insertions(+), 178 deletions(-) diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden index 66a7f910c7..662a45b348 100644 --- a/cli/testdata/server-config.yaml.golden +++ b/cli/testdata/server-config.yaml.golden @@ -787,6 +787,13 @@ aibridgeproxy: # Path to the CA private key file for AI Bridge Proxy. # (default: , type: string) key_file: "" + # Comma-separated list of domains for which HTTPS traffic will be decrypted and + # routed through AI Bridge. Requests to other domains will be tunneled directly + # without decryption. + # (default: api.anthropic.com,api.openai.com, type: string-array) + domain_allowlist: + - api.anthropic.com + - api.openai.com # Configure data retention policies for various database tables. Retention # policies automatically purge old data to reduce database size and improve # performance. Setting a retention duration to 0 disables automatic purging for diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index c65d62fd0d..120beb94bb 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -12055,6 +12055,12 @@ const docTemplate = `{ "cert_file": { "type": "string" }, + "domain_allowlist": { + "type": "array", + "items": { + "type": "string" + } + }, "enabled": { "type": "boolean" }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index a40981e9a7..00d038a2c7 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -10716,6 +10716,12 @@ "cert_file": { "type": "string" }, + "domain_allowlist": { + "type": "array", + "items": { + "type": "string" + } + }, "enabled": { "type": "boolean" }, diff --git a/codersdk/deployment.go b/codersdk/deployment.go index f75c041fe0..db27afb27d 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -3526,6 +3526,17 @@ Write out the current server config as YAML to stdout.`, Group: &deploymentGroupAIBridgeProxy, YAML: "key_file", }, + { + Name: "AI Bridge Proxy Domain Allowlist", + Description: "Comma-separated list of domains for which HTTPS traffic will be decrypted and routed through AI Bridge. Requests to other domains will be tunneled directly without decryption.", + Flag: "aibridge-proxy-domain-allowlist", + Env: "CODER_AIBRIDGE_PROXY_DOMAIN_ALLOWLIST", + Value: &c.AI.BridgeProxyConfig.DomainAllowlist, + Default: "api.anthropic.com,api.openai.com", + Hidden: true, + Group: &deploymentGroupAIBridgeProxy, + YAML: "domain_allowlist", + }, // Retention settings { @@ -3620,10 +3631,11 @@ type AIBridgeBedrockConfig struct { } type AIBridgeProxyConfig struct { - Enabled serpent.Bool `json:"enabled" typescript:",notnull"` - ListenAddr serpent.String `json:"listen_addr" typescript:",notnull"` - CertFile serpent.String `json:"cert_file" typescript:",notnull"` - KeyFile serpent.String `json:"key_file" typescript:",notnull"` + Enabled serpent.Bool `json:"enabled" typescript:",notnull"` + ListenAddr serpent.String `json:"listen_addr" typescript:",notnull"` + CertFile serpent.String `json:"cert_file" typescript:",notnull"` + KeyFile serpent.String `json:"key_file" typescript:",notnull"` + DomainAllowlist serpent.StringArray `json:"domain_allowlist" typescript:",notnull"` } type AIConfig struct { diff --git a/docs/reference/api/general.md b/docs/reference/api/general.md index fc2ca32f10..90ef394aca 100644 --- a/docs/reference/api/general.md +++ b/docs/reference/api/general.md @@ -164,6 +164,9 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \ "ai": { "aibridge_proxy": { "cert_file": "string", + "domain_allowlist": [ + "string" + ], "enabled": true, "key_file": "string", "listen_addr": "string" diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index ab4507cf3b..ad74e5a056 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -597,6 +597,9 @@ ```json { "cert_file": "string", + "domain_allowlist": [ + "string" + ], "enabled": true, "key_file": "string", "listen_addr": "string" @@ -605,12 +608,13 @@ ### Properties -| Name | Type | Required | Restrictions | Description | -|---------------|---------|----------|--------------|-------------| -| `cert_file` | string | false | | | -| `enabled` | boolean | false | | | -| `key_file` | string | false | | | -| `listen_addr` | string | false | | | +| Name | Type | Required | Restrictions | Description | +|--------------------|-----------------|----------|--------------|-------------| +| `cert_file` | string | false | | | +| `domain_allowlist` | array of string | false | | | +| `enabled` | boolean | false | | | +| `key_file` | string | false | | | +| `listen_addr` | string | false | | | ## codersdk.AIBridgeTokenUsage @@ -712,6 +716,9 @@ { "aibridge_proxy": { "cert_file": "string", + "domain_allowlist": [ + "string" + ], "enabled": true, "key_file": "string", "listen_addr": "string" @@ -2624,6 +2631,9 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "ai": { "aibridge_proxy": { "cert_file": "string", + "domain_allowlist": [ + "string" + ], "enabled": true, "key_file": "string", "listen_addr": "string" @@ -3165,6 +3175,9 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "ai": { "aibridge_proxy": { "cert_file": "string", + "domain_allowlist": [ + "string" + ], "enabled": true, "key_file": "string", "listen_addr": "string" diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index 77ac68d9d0..79d7e3bf48 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -10,6 +10,7 @@ import ( "net" "net/http" "net/url" + "slices" "strings" "sync" "time" @@ -69,6 +70,10 @@ 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 is the list of domains to intercept and route through AI Bridge. + // Only requests to these domains will be MITM'd and forwarded to aibridged. + // Requests to other domains will be tunneled directly without decryption. + DomainAllowlist []string } func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) { @@ -78,10 +83,6 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) return nil, xerrors.New("listen address is required") } - if opts.CertFile == "" || opts.KeyFile == "" { - return nil, xerrors.New("cert file and key file are required") - } - if strings.TrimSpace(opts.CoderAccessURL) == "" { return nil, xerrors.New("coder access URL is required") } @@ -90,6 +91,31 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) return nil, xerrors.Errorf("invalid coder access URL %q: %w", opts.CoderAccessURL, err) } + if opts.CertFile == "" || opts.KeyFile == "" { + return nil, xerrors.New("cert file and key file are required") + } + + allowedPorts := opts.AllowedPorts + if len(allowedPorts) == 0 { + allowedPorts = []string{"80", "443"} + } + + if len(opts.DomainAllowlist) == 0 { + return nil, xerrors.New("domain allow list is required") + } + mitmHosts, err := convertDomainsToHosts(opts.DomainAllowlist, allowedPorts) + if err != nil { + return nil, xerrors.Errorf("invalid domain allowlist: %w", err) + } + if len(mitmHosts) == 0 { + return nil, xerrors.New("domain allowlist is empty, at least one domain is required") + } + + logger.Info(ctx, "configured domain allowlist for MITM", + slog.F("domains", opts.DomainAllowlist), + slog.F("hosts", mitmHosts), + ) + // Load CA certificate for MITM certPEM, err := loadMitmCertificate(opts.CertFile, opts.KeyFile) if err != nil { @@ -100,9 +126,6 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) // Cache generated leaf certificates to avoid expensive RSA key generation // and signing on every request to the same hostname. - // TODO(ssncferreira): Currently certs are cached for all MITM'd hosts, but once - // host filtering is implemented, only AI provider certs will be cached. - // Related to https://github.com/coder/internal/issues/1182 if opts.CertStore != nil { proxy.CertStore = opts.CertStore } else { @@ -118,19 +141,19 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) } // Reject CONNECT requests to non-standard ports. - allowedPorts := opts.AllowedPorts - if len(allowedPorts) == 0 { - allowedPorts = []string{"80", "443"} - } proxy.OnRequest().HandleConnectFunc(srv.portMiddleware(allowedPorts)) - // Extract Coder session token from proxy authentication to forward to aibridged. - proxy.OnRequest().HandleConnectFunc(srv.authMiddleware) + // Apply MITM with authentication only to allowlisted hosts. + proxy.OnRequest( + // Only CONNECT requests to these hosts will be intercepted and decrypted. + // All other requests will be tunneled directly to their destination. + goproxy.ReqHostIs(mitmHosts...), + ).HandleConnectFunc( + // Extract Coder session token from proxy authentication to forward to aibridged. + srv.authMiddleware, + ) // Handle decrypted requests: route to aibridged for known AI providers, or passthrough to original destination. - // TODO(ssncferreira): Currently the proxy always behaves as MITM, but this should only happen for known - // AI providers as all other requests should be tunneled. This will be implemented upstack. - // Related to https://github.com/coder/internal/issues/1182 proxy.OnRequest().DoFunc(srv.handleRequest) // Create listener first so we can get the actual address. @@ -249,6 +272,38 @@ func (s *Server) portMiddleware(allowedPorts []string) func(host string, ctx *go } } +// convertDomainsToHosts converts a list of domain names to host:port combinations. +// Each domain is combined with each allowed port. +// Returns an error if a domain includes a port that's not in the allowed ports list. +// For example, ["api.anthropic.com"] with ports ["443"] becomes ["api.anthropic.com:443"]. +func convertDomainsToHosts(domains []string, allowedPorts []string) ([]string, error) { + var hosts []string + for _, domain := range domains { + domain = strings.TrimSpace(strings.ToLower(domain)) + if domain == "" { + continue + } + + // If domain already includes a port, validate it's in the allowed list. + if strings.Contains(domain, ":") { + host, port, err := net.SplitHostPort(domain) + if err != nil { + return nil, xerrors.Errorf("invalid domain %q: %w", domain, err) + } + if !slices.Contains(allowedPorts, port) { + return nil, xerrors.Errorf("invalid port in domain %q: port %s is not in allowed ports %v", domain, port, allowedPorts) + } + hosts = append(hosts, host+":"+port) + } else { + // Otherwise, combine domain with all allowed ports. + for _, port := range allowedPorts { + hosts = append(hosts, domain+":"+port) + } + } + } + return hosts, nil +} + // authMiddleware is a CONNECT middleware that extracts the Coder session token // from the Proxy-Authorization header and stores it in ctx.UserData for use by // downstream request handlers. @@ -317,10 +372,6 @@ func extractCoderTokenFromProxyAuth(proxyAuth string) string { // - Known AI providers return their provider name, used to route to the // corresponding aibridge endpoint. // - Unknown hosts return empty string and are passed through directly. -// -// TODO(ssncferreira): Provider list configurable via domain allowlists will be implemented upstack. -// -// Related to https://github.com/coder/internal/issues/1182. func providerFromURL(reqURL *url.URL) string { if reqURL == nil { return "" @@ -345,14 +396,15 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. // Check if this request is for a supported AI provider. provider := providerFromURL(req.URL) if provider == "" { - // TODO(ssncferreira): After implementing selective MITM, this case should never - // happen since unknown hosts will be tunneled, not decrypted. - // Related to https://github.com/coder/internal/issues/1182 - s.logger.Debug(s.ctx, "passthrough request to unknown host", + // This can happen if a domain is in the allowlist but doesn't have a + // corresponding provider mapping in providerFromURL(). The request was + // decrypted but we don't know how to route it to aibridged. + s.logger.Warn(s.ctx, "decrypted request has no provider mapping, passing through", slog.F("host", req.Host), slog.F("method", req.Method), slog.F("path", originalPath), ) + // Passthrough to the original destination. return req, nil } diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index f9870501e1..1dce7fccf9 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -96,10 +96,11 @@ func generateSharedTestCA() (certFile, keyFile string, err error) { } type testProxyConfig struct { - listenAddr string - coderAccessURL string - allowedPorts []string - certStore *aibridgeproxyd.CertCache + listenAddr string + coderAccessURL string + allowedPorts []string + certStore *aibridgeproxyd.CertCache + domainAllowlist []string } type testProxyOption func(*testProxyConfig) @@ -122,6 +123,12 @@ func withCertStore(store *aibridgeproxyd.CertCache) testProxyOption { } } +func withDomainAllowlist(domains ...string) testProxyOption { + return func(cfg *testProxyConfig) { + cfg.domainAllowlist = domains + } +} + // newTestProxy creates a new AI Bridge Proxy server for testing. // It uses the shared test CA and registers cleanup automatically. // It waits for the proxy server to be ready before returning. @@ -129,8 +136,9 @@ func newTestProxy(t *testing.T, opts ...testProxyOption) *aibridgeproxyd.Server t.Helper() cfg := &testProxyConfig{ - listenAddr: "127.0.0.1:0", - coderAccessURL: "http://localhost:3000", + listenAddr: "127.0.0.1:0", + coderAccessURL: "http://localhost:3000", + domainAllowlist: []string{"127.0.0.1", "localhost"}, } for _, opt := range opts { opt(cfg) @@ -140,11 +148,12 @@ func newTestProxy(t *testing.T, opts ...testProxyOption) *aibridgeproxyd.Server logger := slogtest.Make(t, nil) aibridgeOpts := aibridgeproxyd.Options{ - ListenAddr: cfg.listenAddr, - CoderAccessURL: cfg.coderAccessURL, - CertFile: certFile, - KeyFile: keyFile, - AllowedPorts: cfg.allowedPorts, + ListenAddr: cfg.listenAddr, + CoderAccessURL: cfg.coderAccessURL, + CertFile: certFile, + KeyFile: keyFile, + AllowedPorts: cfg.allowedPorts, + DomainAllowlist: cfg.domainAllowlist, } if cfg.certStore != nil { aibridgeOpts.CertStore = cfg.certStore @@ -169,9 +178,10 @@ func newTestProxy(t *testing.T, opts ...testProxyOption) *aibridgeproxyd.Server return srv } -// newProxyClient creates an HTTP client configured to use the proxy and trust its CA. -// It adds a Proxy-Authorization header with the provided token for authentication. -func newProxyClient(t *testing.T, srv *aibridgeproxyd.Server, proxyAuth string) *http.Client { +// getProxyCertPool returns a cert pool containing the shared test CA certificate. +// This is used for tests where requests are MITM'd by the proxy, so the client +// needs to trust the proxy's CA to verify the generated certificates. +func getProxyCertPool(t *testing.T) *x509.CertPool { t.Helper() certFile, _ := getSharedTestCA(t) @@ -183,6 +193,16 @@ func newProxyClient(t *testing.T, srv *aibridgeproxyd.Server, proxyAuth string) ok := certPool.AppendCertsFromPEM(certPEM) require.True(t, ok) + return certPool +} + +// newProxyClient creates an HTTP client configured to use the proxy. +// It adds a Proxy-Authorization header with the provided token for authentication. +// The certPool parameter specifies which certificates the client should trust. +// For MITM'd requests, use the proxy's CA. For passthrough, use the target server's cert. +func newProxyClient(t *testing.T, srv *aibridgeproxyd.Server, proxyAuth string, certPool *x509.CertPool) *http.Client { + t.Helper() + // Create an HTTP client configured to use the proxy. proxyURL, err := url.Parse("http://" + srv.Addr()) require.NoError(t, err) @@ -207,8 +227,8 @@ func newProxyClient(t *testing.T, srv *aibridgeproxyd.Server, proxyAuth string) } // newTargetServer creates a mock HTTPS server that will be the target of proxied requests. -// It returns the server's parsed URL. The server is automatically closed when the test ends. -func newTargetServer(t *testing.T, handler http.HandlerFunc) *url.URL { +// It returns the server and its parsed URL. The server is automatically closed when the test ends. +func newTargetServer(t *testing.T, handler http.HandlerFunc) (*httptest.Server, *url.URL) { t.Helper() srv := httptest.NewTLSServer(handler) @@ -217,7 +237,7 @@ func newTargetServer(t *testing.T, handler http.HandlerFunc) *url.URL { srvURL, err := url.Parse(srv.URL) require.NoError(t, err) - return srvURL + return srv, srvURL } // makeProxyAuthHeader creates a Proxy-Authorization header value with the given token. @@ -237,10 +257,27 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "", - CoderAccessURL: "http://localhost:3000", - CertFile: certFile, - KeyFile: keyFile, + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{"127.0.0.1", "localhost"}, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "listen address is required") + }) + + t.Run("EmptyListenAddr", func(t *testing.T) { + t.Parallel() + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{"127.0.0.1", "localhost"}, }) require.Error(t, err) require.Contains(t, err.Error(), "listen address is required") @@ -253,9 +290,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", - CertFile: certFile, - KeyFile: keyFile, + ListenAddr: "127.0.0.1:0", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{"127.0.0.1", "localhost"}, }) require.Error(t, err) require.Contains(t, err.Error(), "coder access URL is required") @@ -268,10 +306,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: " ", - CertFile: certFile, - KeyFile: keyFile, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: " ", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{"127.0.0.1", "localhost"}, }) require.Error(t, err) require.Contains(t, err.Error(), "coder access URL is required") @@ -284,10 +323,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: "://invalid", - CertFile: certFile, - KeyFile: keyFile, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "://invalid", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{"127.0.0.1", "localhost"}, }) require.Error(t, err) require.Contains(t, err.Error(), "invalid coder access URL") @@ -299,9 +339,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CoderAccessURL: "http://localhost:3000", - KeyFile: "key.pem", + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + KeyFile: "key.pem", + DomainAllowlist: []string{"127.0.0.1", "localhost"}, }) require.Error(t, err) require.Contains(t, err.Error(), "cert file and key file are required") @@ -313,9 +354,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CoderAccessURL: "http://localhost:3000", - CertFile: "cert.pem", + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + CertFile: "cert.pem", + DomainAllowlist: []string{"127.0.0.1", "localhost"}, }) require.Error(t, err) require.Contains(t, err.Error(), "cert file and key file are required") @@ -327,15 +369,83 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CoderAccessURL: "http://localhost:3000", - CertFile: "/nonexistent/cert.pem", - KeyFile: "/nonexistent/key.pem", + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + CertFile: "/nonexistent/cert.pem", + KeyFile: "/nonexistent/key.pem", + DomainAllowlist: []string{"127.0.0.1", "localhost"}, }) require.Error(t, err) require.Contains(t, err.Error(), "failed to load MITM certificate") }) + t.Run("MissingDomainAllowlist", func(t *testing.T) { + t.Parallel() + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "domain allow list is required") + }) + + t.Run("EmptyDomainAllowlist", func(t *testing.T) { + t.Parallel() + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{""}, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "domain allowlist is empty, at least one domain is required") + }) + + t.Run("InvalidDomainAllowlist", func(t *testing.T) { + t.Parallel() + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{"[invalid:domain"}, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "invalid domain") + }) + + t.Run("DomainWithNonAllowedPort", func(t *testing.T) { + t.Parallel() + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{"api.anthropic.com:8443"}, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "invalid port in domain") + }) + t.Run("Success", func(t *testing.T) { t.Parallel() @@ -343,10 +453,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", - CertFile: certFile, - KeyFile: keyFile, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{"api.anthropic.com", "api.openai.com"}, }) require.NoError(t, err) require.NotNil(t, srv) @@ -363,10 +474,11 @@ 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", - CertFile: certFile, - KeyFile: keyFile, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{"127.0.0.1", "localhost"}, }) require.NoError(t, err) @@ -381,38 +493,86 @@ func TestClose(t *testing.T) { func TestProxy_CertCaching(t *testing.T) { t.Parallel() - // Create a mock HTTPS server that will be the target of the proxied request. - targetURL := newTargetServer(t, func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusOK) - }) + tests := []struct { + name string + domainAllowlist []string + passthrough bool + }{ + { + name: "AllowlistedDomainCached", + domainAllowlist: nil, // will use targetURL.Hostname() + passthrough: false, + }, + { + name: "NonAllowlistedDomainNotCached", + domainAllowlist: []string{"other.example.com"}, + passthrough: true, + }, + } - // Create a cert cache so we can inspect it after the request. - certCache := aibridgeproxyd.NewCertCache() + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() - // Start the proxy server with the certificate cache. - srv := newTestProxy(t, - withAllowedPorts(targetURL.Port()), - withCertStore(certCache), - ) + // Create a mock HTTPS server that will be the target of the proxied request. + targetServer, targetURL := newTargetServer(t, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) - // Make a request through the proxy to the target server. - // This triggers MITM and caches the generated certificate. - client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token")) - req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetURL.String(), nil) - require.NoError(t, err) - resp, err := client.Do(req) - require.NoError(t, err) - defer resp.Body.Close() + // Create a cert cache so we can inspect it after the request. + certCache := aibridgeproxyd.NewCertCache() - // Fetch with a generator that tracks calls: if the certificate was cached - // during the request above, the generator should not be called. - genCalls := 0 - _, err = certCache.Fetch(targetURL.Hostname(), func() (*tls.Certificate, error) { - genCalls++ - return &tls.Certificate{}, nil - }) - require.NoError(t, err) - require.Equal(t, 0, genCalls, "certificate should have been cached during request") + // Configure domain allowlist. + domainAllowlist := tt.domainAllowlist + if domainAllowlist == nil { + domainAllowlist = []string{targetURL.Hostname()} + } + + // Start the proxy server with the certificate cache. + srv := newTestProxy(t, + withAllowedPorts(targetURL.Port()), + withCertStore(certCache), + withDomainAllowlist(domainAllowlist...), + ) + + // Build the cert pool for the client to trust. + // - For MITM'd requests, the client connects through the proxy which generates + // certificates signed by our test CA, so it needs to trust the proxy's CA. + // - For passthrough requests, the client connects directly to the target server + // through a tunnel, so it needs to trust the target's self-signed certificate. + var certPool *x509.CertPool + if tt.passthrough { + certPool = x509.NewCertPool() + certPool.AddCert(targetServer.Certificate()) + } else { + certPool = getProxyCertPool(t) + } + + // Make a request through the proxy to the target server. + client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token"), certPool) + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetURL.String(), nil) + require.NoError(t, err) + resp, err := client.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + // Fetch with a generator that tracks calls. + genCalls := 0 + _, err = certCache.Fetch(targetURL.Hostname(), func() (*tls.Certificate, error) { + genCalls++ + return &tls.Certificate{}, nil + }) + require.NoError(t, err) + + if tt.passthrough { + // 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") + } else { + // Certificate should have been cached during MITM. + require.Equal(t, 0, genCalls, "certificate should have been cached during request") + } + }) + } } func TestProxy_PortValidation(t *testing.T) { @@ -445,16 +605,19 @@ func TestProxy_PortValidation(t *testing.T) { t.Parallel() // Create a target HTTPS server that will be the destination of our proxied request. - targetURL := newTargetServer(t, func(w http.ResponseWriter, r *http.Request) { + _, targetURL := newTargetServer(t, func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("hello from target")) }) // Start the proxy server on a random port to avoid conflicts when running tests in parallel. - srv := newTestProxy(t, withAllowedPorts(tt.allowedPorts(targetURL)...)) + srv := newTestProxy(t, + withAllowedPorts(tt.allowedPorts(targetURL)...), + withDomainAllowlist(targetURL.Hostname()), + ) // Make a request through the proxy to the target server. - client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token")) + client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token"), getProxyCertPool(t)) req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetURL.String(), nil) require.NoError(t, err) @@ -511,16 +674,19 @@ func TestProxy_Authentication(t *testing.T) { t.Parallel() // Create a mock HTTPS server that will be the target of our proxied request. - targetURL := newTargetServer(t, func(w http.ResponseWriter, _ *http.Request) { + _, targetURL := newTargetServer(t, func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("hello from target")) }) // Start the proxy server on a random port to avoid conflicts when running tests in parallel. - srv := newTestProxy(t, withAllowedPorts(targetURL.Port())) + srv := newTestProxy(t, + withAllowedPorts(targetURL.Port()), + withDomainAllowlist(targetURL.Hostname()), + ) // Make a request through the proxy to the target server. - client := newProxyClient(t, srv, tt.proxyAuth) + client := newProxyClient(t, srv, tt.proxyAuth, getProxyCertPool(t)) req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetURL.String(), nil) require.NoError(t, err) resp, err := client.Do(req) @@ -545,44 +711,72 @@ func TestProxy_MITM(t *testing.T) { t.Parallel() tests := []struct { - name string - targetHost string - targetPort string // optional, if empty uses default HTTPS port (443) - targetPath string - expectedPath string - passthrough bool + name string + domainAllowlist []string + allowedPorts []string + buildTargetURL func(passthroughURL *url.URL) (string, error) + passthrough bool + noAIBridgeRouting bool + expectedPath string }{ { - name: "AnthropicMessages", - targetHost: "api.anthropic.com", - targetPath: "/v1/messages", + name: "MitmdAnthropic", + domainAllowlist: []string{"api.anthropic.com"}, + allowedPorts: []string{"443"}, + buildTargetURL: func(_ *url.URL) (string, error) { + return "https://api.anthropic.com/v1/messages", nil + }, expectedPath: "/api/v2/aibridge/anthropic/v1/messages", }, { - name: "AnthropicNonDefaultPort", - targetHost: "api.anthropic.com", - targetPort: "8443", - targetPath: "/v1/messages", + name: "MitmdAnthropicNonDefaultPort", + domainAllowlist: []string{"api.anthropic.com"}, + allowedPorts: []string{"8443"}, + buildTargetURL: func(_ *url.URL) (string, error) { + return "https://api.anthropic.com:8443/v1/messages", nil + }, expectedPath: "/api/v2/aibridge/anthropic/v1/messages", }, { - name: "OpenAIChatCompletions", - targetHost: "api.openai.com", - targetPath: "/v1/chat/completions", + name: "MitmdOpenAI", + domainAllowlist: []string{"api.openai.com"}, + allowedPorts: []string{"443"}, + buildTargetURL: func(_ *url.URL) (string, error) { + return "https://api.openai.com/v1/chat/completions", nil + }, expectedPath: "/api/v2/aibridge/openai/v1/chat/completions", }, { - name: "OpenAINonDefaultPort", - targetHost: "api.openai.com", - targetPort: "8443", - targetPath: "/v1/chat/completions", + name: "MitmdOpenAINonDefaultPort", + domainAllowlist: []string{"api.openai.com"}, + allowedPorts: []string{"8443"}, + buildTargetURL: func(_ *url.URL) (string, error) { + return "https://api.openai.com:8443/v1/chat/completions", nil + }, expectedPath: "/api/v2/aibridge/openai/v1/chat/completions", }, { - name: "UnknownHostPassthrough", - targetPath: "/some/path", + name: "PassthroughUnknownHost", + domainAllowlist: []string{"other.example.com"}, + allowedPorts: nil, // will use passthroughURL.Port() + buildTargetURL: func(passthroughURL *url.URL) (string, error) { + return url.JoinPath(passthroughURL.String(), "/some/path") + }, passthrough: true, }, + // The host is MITM'd but has no provider mapping. + // The request is decrypted but passed through to the original destination + // instead of being routed to aibridge. + { + name: "MitmdWithoutAIBridgeRouting", + domainAllowlist: nil, // will use passthroughURL.Hostname() + allowedPorts: nil, // will use passthroughURL.Port() + buildTargetURL: func(passthroughURL *url.URL) (string, error) { + return url.JoinPath(passthroughURL.String(), "/some/path") + }, + passthrough: false, + noAIBridgeRouting: true, + }, } for _, tt := range tests { @@ -603,50 +797,49 @@ func TestProxy_MITM(t *testing.T) { t.Cleanup(func() { aibridgedServer.Close() }) // Create a mock target server for passthrough tests. - passthroughURL := newTargetServer(t, func(w http.ResponseWriter, _ *http.Request) { + passthroughServer, passthroughURL := newTargetServer(t, func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("hello from passthrough")) }) - // Configure allowed ports based on test case. - // AI provider tests connect to the specified port, or 443 if not specified. - // Passthrough tests connect directly to the local target server's random port. - var allowedPorts []string - switch { - case tt.passthrough: + // Configure allowed ports. + allowedPorts := tt.allowedPorts + if allowedPorts == nil { allowedPorts = []string{passthroughURL.Port()} - case tt.targetPort != "": - allowedPorts = []string{tt.targetPort} - default: - allowedPorts = []string{"443"} + } + + // Configure domain allowlist. + domainAllowlist := tt.domainAllowlist + if domainAllowlist == nil { + domainAllowlist = []string{passthroughURL.Hostname()} } // Start the proxy server pointing to our mock aibridged. srv := newTestProxy(t, withCoderAccessURL(aibridgedServer.URL), withAllowedPorts(allowedPorts...), + withDomainAllowlist(domainAllowlist...), ) // Build the target URL: - // - For passthrough, target the local mock TLS server. - // - For AI providers, use their real hostnames to trigger routing. - // Non-default ports are included explicitly; default port (443) is omitted. - var targetURL string - var err error - switch { - case tt.passthrough: - targetURL, err = url.JoinPath(passthroughURL.String(), tt.targetPath) - require.NoError(t, err) - case tt.targetPort != "": - targetURL, err = url.JoinPath("https://"+tt.targetHost+":"+tt.targetPort, tt.targetPath) - require.NoError(t, err) - default: - targetURL, err = url.JoinPath("https://"+tt.targetHost, tt.targetPath) - require.NoError(t, err) + targetURL, err := tt.buildTargetURL(passthroughURL) + require.NoError(t, err) + + // Build the cert pool for the client to trust. + // - For MITM'd requests, the client connects through the proxy which generates + // certificates signed by our test CA, so it needs to trust the proxy's CA. + // - For passthrough requests, the client connects directly to the target server + // through a tunnel, so it needs to trust the target's self-signed certificate. + var certPool *x509.CertPool + if tt.passthrough { + certPool = x509.NewCertPool() + certPool.AddCert(passthroughServer.Certificate()) + } else { + certPool = getProxyCertPool(t) } // Make a request through the proxy to the target URL. - client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token")) + client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token"), certPool) req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, targetURL, strings.NewReader(`{}`)) require.NoError(t, err) req.Header.Set("Content-Type", "application/json") @@ -659,16 +852,16 @@ func TestProxy_MITM(t *testing.T) { require.NoError(t, err) require.Equal(t, http.StatusOK, resp.StatusCode) - if tt.passthrough { + if tt.passthrough || tt.noAIBridgeRouting { // Verify request went to target server, not aibridged. require.Equal(t, "hello from passthrough", string(body)) require.Empty(t, receivedPath, "aibridged should not receive passthrough requests") - require.Empty(t, receivedAuth, "aibridged should not receive passthrough requests") + require.Empty(t, receivedAuth, "passthrough requests are not authenticated by the proxy") } else { // Verify the request was routed to aibridged correctly. require.Equal(t, "hello from aibridged", string(body)) require.Equal(t, tt.expectedPath, receivedPath) - require.Equal(t, "Bearer test-session-token", receivedAuth) + require.Equal(t, "Bearer test-session-token", receivedAuth, "MITM'd requests must include authentication") } }) } @@ -743,10 +936,11 @@ 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", - CertFile: compoundCertFile, - KeyFile: keyFile, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: compoundCertFile, + KeyFile: keyFile, + DomainAllowlist: []string{"127.0.0.1", "localhost"}, }) require.NoError(t, err) t.Cleanup(func() { _ = srv.Close() }) diff --git a/enterprise/cli/aibridgeproxyd.go b/enterprise/cli/aibridgeproxyd.go index cc613ad2ba..9ee6725b72 100644 --- a/enterprise/cli/aibridgeproxyd.go +++ b/enterprise/cli/aibridgeproxyd.go @@ -18,10 +18,11 @@ func newAIBridgeProxyDaemon(coderAPI *coderd.API) (*aibridgeproxyd.Server, error logger := coderAPI.Logger.Named("aibridgeproxyd") srv, err := aibridgeproxyd.New(ctx, logger, aibridgeproxyd.Options{ - ListenAddr: coderAPI.DeploymentValues.AI.BridgeProxyConfig.ListenAddr.String(), - CoderAccessURL: coderAPI.AccessURL.String(), - CertFile: coderAPI.DeploymentValues.AI.BridgeProxyConfig.CertFile.String(), - KeyFile: coderAPI.DeploymentValues.AI.BridgeProxyConfig.KeyFile.String(), + ListenAddr: coderAPI.DeploymentValues.AI.BridgeProxyConfig.ListenAddr.String(), + CoderAccessURL: coderAPI.AccessURL.String(), + CertFile: coderAPI.DeploymentValues.AI.BridgeProxyConfig.CertFile.String(), + KeyFile: coderAPI.DeploymentValues.AI.BridgeProxyConfig.KeyFile.String(), + DomainAllowlist: coderAPI.DeploymentValues.AI.BridgeProxyConfig.DomainAllowlist.Value(), }) if err != nil { return nil, xerrors.Errorf("failed to start in-memory aibridgeproxy daemon: %w", err) diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 054b280dbf..c76d7509a6 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -71,6 +71,7 @@ export interface AIBridgeProxyConfig { readonly listen_addr: string; readonly cert_file: string; readonly key_file: string; + readonly domain_allowlist: string; } // From codersdk/aibridge.go