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
This commit is contained in:
Susana Ferreira
2026-01-13 11:30:51 +00:00
committed by GitHub
parent 64e7a77983
commit 74b6d12a8a
10 changed files with 473 additions and 178 deletions
+7
View File
@@ -787,6 +787,13 @@ aibridgeproxy:
# Path to the CA private key file for AI Bridge Proxy.
# (default: <unset>, 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
+6
View File
@@ -12055,6 +12055,12 @@ const docTemplate = `{
"cert_file": {
"type": "string"
},
"domain_allowlist": {
"type": "array",
"items": {
"type": "string"
}
},
"enabled": {
"type": "boolean"
},
+6
View File
@@ -10716,6 +10716,12 @@
"cert_file": {
"type": "string"
},
"domain_allowlist": {
"type": "array",
"items": {
"type": "string"
}
},
"enabled": {
"type": "boolean"
},
+16 -4
View File
@@ -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 {
+3
View File
@@ -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"
+19 -6
View File
@@ -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"
+76 -24
View File
@@ -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
}
+334 -140
View File
@@ -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() })
+5 -4
View File
@@ -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)
+1
View File
@@ -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