From b97572285a3cdd868197d747b0e2ee217ca0b968 Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Mon, 29 Dec 2025 15:31:51 +0000 Subject: [PATCH] feat: add core AI MITM proxy daemon (#21296) ## Description Adds the core AI Bridge MITM proxy daemon. This proxy intercepts HTTPS traffic, decrypts it using a configured CA certificate, and forwards requests to AIBridge for processing. ## Changes * Added `aibridgeproxyd` package with the core proxy server implementation * Added configuration options: `CODER_AIBRIDGE_PROXY_ENABLED`, `CODER_AIBRIDGE_PROXY_LISTEN_ADDR`, `CODER_AIBRIDGE_PROXY_CERT_FILE`, `CODER_AIBRIDGE_PROXY_KEY_FILE` * Added tests for server initialization and MITM functionality Closes https://github.com/coder/internal/issues/1180 --- cli/testdata/coder_server_--help.golden | 14 + cli/testdata/server-config.yaml.golden | 14 + coderd/apidoc/docs.go | 20 ++ coderd/apidoc/swagger.json | 20 ++ codersdk/deployment.go | 57 +++- docs/reference/api/general.md | 6 + docs/reference/api/schemas.md | 45 ++- docs/reference/cli/server.md | 42 +++ enterprise/aibridgeproxyd/aibridgeproxyd.go | 126 +++++++++ .../aibridgeproxyd/aibridgeproxyd_test.go | 259 ++++++++++++++++++ enterprise/cli/aibridgeproxyd.go | 30 ++ enterprise/cli/server.go | 10 + .../cli/testdata/coder_server_--help.golden | 14 + go.mod | 1 + go.sum | 2 + site/src/api/typesGenerated.ts | 9 + 16 files changed, 665 insertions(+), 4 deletions(-) create mode 100644 enterprise/aibridgeproxyd/aibridgeproxyd.go create mode 100644 enterprise/aibridgeproxyd/aibridgeproxyd_test.go create mode 100644 enterprise/cli/aibridgeproxyd.go diff --git a/cli/testdata/coder_server_--help.golden b/cli/testdata/coder_server_--help.golden index 43ed943601..13900f7687 100644 --- a/cli/testdata/coder_server_--help.golden +++ b/cli/testdata/coder_server_--help.golden @@ -139,6 +139,20 @@ AI BRIDGE OPTIONS: Maximum number of AI Bridge requests per second per replica. Set to 0 to disable (unlimited). +AI BRIDGE PROXY OPTIONS: + --aibridge-proxy-cert-file string, $CODER_AIBRIDGE_PROXY_CERT_FILE + Path to the CA certificate file for AI Bridge Proxy. + + --aibridge-proxy-enabled bool, $CODER_AIBRIDGE_PROXY_ENABLED (default: false) + Enable the AI Bridge MITM Proxy for intercepting and decrypting AI + provider requests. + + --aibridge-proxy-key-file string, $CODER_AIBRIDGE_PROXY_KEY_FILE + Path to the CA private key file for AI Bridge Proxy. + + --aibridge-proxy-listen-addr string, $CODER_AIBRIDGE_PROXY_LISTEN_ADDR (default: :8888) + The address the AI Bridge Proxy will listen on. + CLIENT OPTIONS: These options change the behavior of how clients interact with the Coder. Clients include the Coder CLI, Coder Desktop, IDE extensions, and the web UI. diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden index 14ea079f4c..63a28ff3b9 100644 --- a/cli/testdata/server-config.yaml.golden +++ b/cli/testdata/server-config.yaml.golden @@ -765,6 +765,20 @@ aibridge: # (unlimited). # (default: 0, type: int) rateLimit: 0 +aibridgeproxy: + # Enable the AI Bridge MITM Proxy for intercepting and decrypting AI provider + # requests. + # (default: false, type: bool) + enabled: false + # The address the AI Bridge Proxy will listen on. + # (default: :8888, type: string) + listen_addr: :8888 + # Path to the CA certificate file for AI Bridge Proxy. + # (default: , type: string) + cert_file: "" + # Path to the CA private key file for AI Bridge Proxy. + # (default: , type: string) + key_file: "" # 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 e48f97a138..fcee015a12 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -12009,6 +12009,23 @@ const docTemplate = `{ } } }, + "codersdk.AIBridgeProxyConfig": { + "type": "object", + "properties": { + "cert_file": { + "type": "string" + }, + "enabled": { + "type": "boolean" + }, + "key_file": { + "type": "string" + }, + "listen_addr": { + "type": "string" + } + } + }, "codersdk.AIBridgeTokenUsage": { "type": "object", "properties": { @@ -12108,6 +12125,9 @@ const docTemplate = `{ "codersdk.AIConfig": { "type": "object", "properties": { + "aibridge_proxy": { + "$ref": "#/definitions/codersdk.AIBridgeProxyConfig" + }, "bridge": { "$ref": "#/definitions/codersdk.AIBridgeConfig" } diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 577914a491..b3bc29d7bf 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -10673,6 +10673,23 @@ } } }, + "codersdk.AIBridgeProxyConfig": { + "type": "object", + "properties": { + "cert_file": { + "type": "string" + }, + "enabled": { + "type": "boolean" + }, + "key_file": { + "type": "string" + }, + "listen_addr": { + "type": "string" + } + } + }, "codersdk.AIBridgeTokenUsage": { "type": "object", "properties": { @@ -10772,6 +10789,9 @@ "codersdk.AIConfig": { "type": "object", "properties": { + "aibridge_proxy": { + "$ref": "#/definitions/codersdk.AIBridgeProxyConfig" + }, "bridge": { "$ref": "#/definitions/codersdk.AIBridgeConfig" } diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 58725eea84..218cf94eca 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -1217,6 +1217,10 @@ func (c *DeploymentValues) Options() serpent.OptionSet { Name: "AI Bridge", YAML: "aibridge", } + deploymentGroupAIBridgeProxy = serpent.Group{ + Name: "AI Bridge Proxy", + YAML: "aibridgeproxy", + } deploymentGroupRetention = serpent.Group{ Name: "Retention", Description: "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 that data type.", @@ -3443,6 +3447,49 @@ Write out the current server config as YAML to stdout.`, Group: &deploymentGroupAIBridge, YAML: "rateLimit", }, + + // AI Bridge Proxy Options + { + Name: "AI Bridge Proxy Enabled", + Description: "Enable the AI Bridge MITM Proxy for intercepting and decrypting AI provider requests.", + Flag: "aibridge-proxy-enabled", + Env: "CODER_AIBRIDGE_PROXY_ENABLED", + Value: &c.AI.BridgeProxyConfig.Enabled, + Default: "false", + Group: &deploymentGroupAIBridgeProxy, + YAML: "enabled", + }, + { + Name: "AI Bridge Proxy Listen Address", + Description: "The address the AI Bridge Proxy will listen on.", + Flag: "aibridge-proxy-listen-addr", + Env: "CODER_AIBRIDGE_PROXY_LISTEN_ADDR", + Value: &c.AI.BridgeProxyConfig.ListenAddr, + Default: ":8888", + Group: &deploymentGroupAIBridgeProxy, + YAML: "listen_addr", + }, + { + Name: "AI Bridge Proxy Certificate File", + Description: "Path to the CA certificate file for AI Bridge Proxy.", + Flag: "aibridge-proxy-cert-file", + Env: "CODER_AIBRIDGE_PROXY_CERT_FILE", + Value: &c.AI.BridgeProxyConfig.CertFile, + Default: "", + Group: &deploymentGroupAIBridgeProxy, + YAML: "cert_file", + }, + { + Name: "AI Bridge Proxy Key File", + Description: "Path to the CA private key file for AI Bridge Proxy.", + Flag: "aibridge-proxy-key-file", + Env: "CODER_AIBRIDGE_PROXY_KEY_FILE", + Value: &c.AI.BridgeProxyConfig.KeyFile, + Default: "", + Group: &deploymentGroupAIBridgeProxy, + YAML: "key_file", + }, + // Retention settings { Name: "Audit Logs Retention", @@ -3535,8 +3582,16 @@ type AIBridgeBedrockConfig struct { SmallFastModel serpent.String `json:"small_fast_model" typescript:",notnull"` } +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"` +} + type AIConfig struct { - BridgeConfig AIBridgeConfig `json:"bridge,omitempty"` + BridgeConfig AIBridgeConfig `json:"bridge,omitempty"` + BridgeProxyConfig AIBridgeProxyConfig `json:"aibridge_proxy,omitempty"` } type SupportConfig struct { diff --git a/docs/reference/api/general.md b/docs/reference/api/general.md index e9a12a196f..83dae06688 100644 --- a/docs/reference/api/general.md +++ b/docs/reference/api/general.md @@ -162,6 +162,12 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \ }, "agent_stat_refresh_interval": 0, "ai": { + "aibridge_proxy": { + "cert_file": "string", + "enabled": true, + "key_file": "string", + "listen_addr": "string" + }, "bridge": { "anthropic": { "base_url": "string", diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 6929f4bdb6..494b2d30ac 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -592,6 +592,26 @@ | `base_url` | string | false | | | | `key` | string | false | | | +## codersdk.AIBridgeProxyConfig + +```json +{ + "cert_file": "string", + "enabled": true, + "key_file": "string", + "listen_addr": "string" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|---------------|---------|----------|--------------|-------------| +| `cert_file` | string | false | | | +| `enabled` | boolean | false | | | +| `key_file` | string | false | | | +| `listen_addr` | string | false | | | + ## codersdk.AIBridgeTokenUsage ```json @@ -690,6 +710,12 @@ ```json { + "aibridge_proxy": { + "cert_file": "string", + "enabled": true, + "key_file": "string", + "listen_addr": "string" + }, "bridge": { "anthropic": { "base_url": "string", @@ -717,9 +743,10 @@ ### Properties -| Name | Type | Required | Restrictions | Description | -|----------|----------------------------------------------------|----------|--------------|-------------| -| `bridge` | [codersdk.AIBridgeConfig](#codersdkaibridgeconfig) | false | | | +| Name | Type | Required | Restrictions | Description | +|------------------|--------------------------------------------------------------|----------|--------------|-------------| +| `aibridge_proxy` | [codersdk.AIBridgeProxyConfig](#codersdkaibridgeproxyconfig) | false | | | +| `bridge` | [codersdk.AIBridgeConfig](#codersdkaibridgeconfig) | false | | | ## codersdk.APIAllowListTarget @@ -2595,6 +2622,12 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o }, "agent_stat_refresh_interval": 0, "ai": { + "aibridge_proxy": { + "cert_file": "string", + "enabled": true, + "key_file": "string", + "listen_addr": "string" + }, "bridge": { "anthropic": { "base_url": "string", @@ -3126,6 +3159,12 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o }, "agent_stat_refresh_interval": 0, "ai": { + "aibridge_proxy": { + "cert_file": "string", + "enabled": true, + "key_file": "string", + "listen_addr": "string" + }, "bridge": { "anthropic": { "base_url": "string", diff --git a/docs/reference/cli/server.md b/docs/reference/cli/server.md index 6f7aa705a9..d3ea11879d 100644 --- a/docs/reference/cli/server.md +++ b/docs/reference/cli/server.md @@ -1814,6 +1814,48 @@ Maximum number of concurrent AI Bridge requests per replica. Set to 0 to disable Maximum number of AI Bridge requests per second per replica. Set to 0 to disable (unlimited). +### --aibridge-proxy-enabled + +| | | +|-------------|--------------------------------------------| +| Type | bool | +| Environment | $CODER_AIBRIDGE_PROXY_ENABLED | +| YAML | aibridgeproxy.enabled | +| Default | false | + +Enable the AI Bridge MITM Proxy for intercepting and decrypting AI provider requests. + +### --aibridge-proxy-listen-addr + +| | | +|-------------|------------------------------------------------| +| Type | string | +| Environment | $CODER_AIBRIDGE_PROXY_LISTEN_ADDR | +| YAML | aibridgeproxy.listen_addr | +| Default | :8888 | + +The address the AI Bridge Proxy will listen on. + +### --aibridge-proxy-cert-file + +| | | +|-------------|----------------------------------------------| +| Type | string | +| Environment | $CODER_AIBRIDGE_PROXY_CERT_FILE | +| YAML | aibridgeproxy.cert_file | + +Path to the CA certificate file for AI Bridge Proxy. + +### --aibridge-proxy-key-file + +| | | +|-------------|---------------------------------------------| +| Type | string | +| Environment | $CODER_AIBRIDGE_PROXY_KEY_FILE | +| YAML | aibridgeproxy.key_file | + +Path to the CA private key file for AI Bridge Proxy. + ### --audit-logs-retention | | | diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go new file mode 100644 index 0000000000..adf6848aa6 --- /dev/null +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -0,0 +1,126 @@ +package aibridgeproxyd + +import ( + "context" + "crypto/tls" + "crypto/x509" + "errors" + "net/http" + "sync" + "time" + + "github.com/elazarl/goproxy" + "golang.org/x/xerrors" + + "cdr.dev/slog" +) + +// loadMitmOnce ensures the MITM certificate is loaded exactly once. +// goproxy.GoproxyCa is a package-level global variable shared across all +// goproxy.ProxyHttpServer instances in the process. In tests, multiple proxy +// servers run in parallel, and without this guard they would race on writing +// to GoproxyCa. In production, only one server runs, so this has no impact. +var loadMitmOnce sync.Once + +// Server is the AI MITM (Man-in-the-Middle) proxy server. +// It is responsible for: +// - intercepting HTTPS requests to AI providers +// - decrypting requests using the configured CA certificate +// - forwarding requests to aibridged for processing +type Server struct { + logger slog.Logger + proxy *goproxy.ProxyHttpServer + httpServer *http.Server +} + +// Options configures the AI Bridge Proxy server. +type Options struct { + // ListenAddr is the address the proxy server will listen on. + ListenAddr string + // CertFile is the path to the CA certificate file used for MITM. + CertFile string + // KeyFile is the path to the CA private key file used for MITM. + KeyFile string +} + +func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) { + logger.Info(ctx, "initializing AI Bridge Proxy server") + + if opts.ListenAddr == "" { + return nil, xerrors.New("listen address is required") + } + + if opts.CertFile == "" || opts.KeyFile == "" { + return nil, xerrors.New("cert file and key file are required") + } + + // Load CA certificate for MITM + if err := loadMitmCertificate(opts.CertFile, opts.KeyFile); err != nil { + return nil, xerrors.Errorf("failed to load MITM certificate: %w", err) + } + + proxy := goproxy.NewProxyHttpServer() + + // Decrypt all HTTPS requests via MITM. Requests are forwarded to + // the original destination without modification for now. + // TODO(ssncferreira): Route requests to aibridged will be implemented upstack. + // Related to https://github.com/coder/internal/issues/1181 + proxy.OnRequest().HandleConnect(goproxy.AlwaysMitm) + + srv := &Server{ + logger: logger, + proxy: proxy, + } + + // Start HTTP server in background + srv.httpServer = &http.Server{ + Addr: opts.ListenAddr, + Handler: proxy, + ReadHeaderTimeout: 10 * time.Second, + } + + go func() { + logger.Info(ctx, "starting AI Bridge Proxy", slog.F("addr", opts.ListenAddr)) + if err := srv.httpServer.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + logger.Error(ctx, "aibridgeproxyd server error", slog.Error(err)) + } + }() + + return srv, nil +} + +// Close gracefully shuts down the proxy server. +func (s *Server) Close() error { + if s.httpServer == nil { + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + return s.httpServer.Shutdown(ctx) +} + +// loadMitmCertificate loads the CA certificate and private key for MITM proxying. +// This function is safe to call concurrently - the certificate is only loaded once +// into the global goproxy.GoproxyCa variable. +func loadMitmCertificate(certFile, keyFile string) error { + tlsCert, err := tls.LoadX509KeyPair(certFile, keyFile) + if err != nil { + return xerrors.Errorf("load CA certificate: %w", err) + } + + x509Cert, err := x509.ParseCertificate(tlsCert.Certificate[0]) + if err != nil { + return xerrors.Errorf("parse CA certificate: %w", err) + } + + // Only protect the global assignment with sync.Once + loadMitmOnce.Do(func() { + goproxy.GoproxyCa = tls.Certificate{ + Certificate: tlsCert.Certificate, + PrivateKey: tlsCert.PrivateKey, + Leaf: x509Cert, + } + }) + + return nil +} diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go new file mode 100644 index 0000000000..8b5b05773a --- /dev/null +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -0,0 +1,259 @@ +package aibridgeproxyd_test + +import ( + "crypto/rand" + "crypto/rsa" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "io" + "math/big" + "net" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "cdr.dev/slog/sloggers/slogtest" + + "github.com/coder/coder/v2/enterprise/aibridgeproxyd" + "github.com/coder/coder/v2/testutil" +) + +var ( + // testCAOnce ensures the shared CA is generated exactly once. + // sync.Once guarantees single execution even with parallel tests. + // Note: no retry on failure. + testCAOnce sync.Once + // Shared CA certificate and key paths, and any error from generation. + // These are set once by testCAOnce and read by all tests. + testCACert string + testCAKey string + errTestSharedCA error +) + +// getSharedTestCA returns a shared CA certificate for all tests. +// This avoids race conditions with goproxy.GoproxyCa which is a global variable. +// Using sync.Once ensures the CA is generated exactly once, even when tests run +// in parallel. All tests share the same CA, so goproxy.GoproxyCa is only set once. +func getSharedTestCA(t *testing.T) (certFile, keyFile string) { + t.Helper() + + testCAOnce.Do(func() { + testCACert, testCAKey, errTestSharedCA = generateSharedTestCA() + }) + + require.NoError(t, errTestSharedCA, "failed to generate shared test CA") + return testCACert, testCAKey +} + +// generateSharedTestCA creates a shared CA certificate and key for testing. +func generateSharedTestCA() (certFile, keyFile string, err error) { + caKey, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + return "", "", xerrors.Errorf("generate CA key: %w", err) + } + + caTemplate := x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "Shared Test CA"}, + NotBefore: time.Now(), + NotAfter: time.Now().Add(time.Hour), + KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign, + BasicConstraintsValid: true, + IsCA: true, + } + + caCertDER, err := x509.CreateCertificate(rand.Reader, &caTemplate, &caTemplate, &caKey.PublicKey, caKey) + if err != nil { + return "", "", xerrors.Errorf("create CA certificate: %w", err) + } + + tmpDir := os.TempDir() + certPath := filepath.Join(tmpDir, "aibridgeproxyd_test_ca.crt") + keyPath := filepath.Join(tmpDir, "aibridgeproxyd_test_ca.key") + + certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: caCertDER}) + if err := os.WriteFile(certPath, certPEM, 0o600); err != nil { + return "", "", xerrors.Errorf("write cert file: %w", err) + } + + keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(caKey)}) + if err := os.WriteFile(keyPath, keyPEM, 0o600); err != nil { + return "", "", xerrors.Errorf("write key file: %w", err) + } + + return certPath, keyPath, nil +} + +func TestNew(t *testing.T) { + t.Parallel() + + t.Run("MissingListenAddr", func(t *testing.T) { + t.Parallel() + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "", + CertFile: certFile, + KeyFile: keyFile, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "listen address is required") + }) + + t.Run("MissingCertFile", func(t *testing.T) { + t.Parallel() + + logger := slogtest.Make(t, nil) + + _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: ":0", + KeyFile: "key.pem", + }) + require.Error(t, err) + require.Contains(t, err.Error(), "cert file and key file are required") + }) + + t.Run("MissingKeyFile", func(t *testing.T) { + t.Parallel() + + logger := slogtest.Make(t, nil) + + _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: ":0", + CertFile: "cert.pem", + }) + require.Error(t, err) + require.Contains(t, err.Error(), "cert file and key file are required") + }) + + t.Run("InvalidCertFile", func(t *testing.T) { + t.Parallel() + + logger := slogtest.Make(t, nil) + + _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: ":0", + CertFile: "/nonexistent/cert.pem", + KeyFile: "/nonexistent/key.pem", + }) + require.Error(t, err) + require.Contains(t, err.Error(), "failed to load MITM certificate") + }) + + t.Run("Success", func(t *testing.T) { + t.Parallel() + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "127.0.0.1:0", + CertFile: certFile, + KeyFile: keyFile, + }) + require.NoError(t, err) + require.NotNil(t, srv) + + err = srv.Close() + require.NoError(t, err) + }) +} + +func TestClose(t *testing.T) { + t.Parallel() + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "127.0.0.1:0", + CertFile: certFile, + KeyFile: keyFile, + }) + require.NoError(t, err) + + err = srv.Close() + require.NoError(t, err) + + // Calling Close again should not error + err = srv.Close() + require.NoError(t, err) +} + +func TestProxy_MITM(t *testing.T) { + t.Parallel() + + // Create a mock HTTPS server that will be the target of our proxied request. + targetServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("hello from target")) + })) + defer targetServer.Close() + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + // Start the proxy server. + srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "127.0.0.1:8888", + CertFile: certFile, + KeyFile: keyFile, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Close() }) + + // Wait for the proxy server to be ready. + require.Eventually(t, func() bool { + conn, err := net.Dial("tcp", "127.0.0.1:8888") + if err != nil { + return false + } + _ = conn.Close() + return true + }, testutil.WaitShort, testutil.IntervalFast) + + // Load the CA certificate so the client trusts the proxy's MITM certificate. + certPEM, err := os.ReadFile(certFile) + require.NoError(t, err) + certPool := x509.NewCertPool() + certPool.AppendCertsFromPEM(certPEM) + + // Create an HTTP client configured to use the proxy. + proxyURL, err := url.Parse("http://127.0.0.1:8888") + require.NoError(t, err) + + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + TLSClientConfig: &tls.Config{ + MinVersion: tls.VersionTLS12, + RootCAs: certPool, + }, + }, + } + + // Make a request through the proxy to the target server. + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetServer.URL, nil) + require.NoError(t, err) + resp, err := client.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + // Verify the response was successfully proxied. + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, "hello from target", string(body)) +} diff --git a/enterprise/cli/aibridgeproxyd.go b/enterprise/cli/aibridgeproxyd.go new file mode 100644 index 0000000000..6c2a388a98 --- /dev/null +++ b/enterprise/cli/aibridgeproxyd.go @@ -0,0 +1,30 @@ +//go:build !slim + +package cli + +import ( + "context" + + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/enterprise/aibridgeproxyd" + "github.com/coder/coder/v2/enterprise/coderd" +) + +func newAIBridgeProxyDaemon(coderAPI *coderd.API) (*aibridgeproxyd.Server, error) { + ctx := context.Background() + coderAPI.Logger.Debug(ctx, "starting in-memory aibridgeproxy daemon") + + logger := coderAPI.Logger.Named("aibridgeproxyd") + + srv, err := aibridgeproxyd.New(ctx, logger, aibridgeproxyd.Options{ + ListenAddr: coderAPI.DeploymentValues.AI.BridgeProxyConfig.ListenAddr.String(), + CertFile: coderAPI.DeploymentValues.AI.BridgeProxyConfig.CertFile.String(), + KeyFile: coderAPI.DeploymentValues.AI.BridgeProxyConfig.KeyFile.String(), + }) + if err != nil { + return nil, xerrors.Errorf("failed to start in-memory aibridgeproxy daemon: %w", err) + } + + return srv, nil +} diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go index bc77bc54ba..6950144a75 100644 --- a/enterprise/cli/server.go +++ b/enterprise/cli/server.go @@ -165,6 +165,16 @@ func (r *RootCmd) Server(_ func()) *serpent.Command { closers.Add(aibridgeDaemon) } + // In-memory AI Bridge Proxy daemon + if options.DeploymentValues.AI.BridgeProxyConfig.Enabled.Value() { + aiBridgeProxyServer, err := newAIBridgeProxyDaemon(api) + if err != nil { + _ = closers.Close() + return nil, nil, xerrors.Errorf("create aibridgeproxyd: %w", err) + } + closers.Add(aiBridgeProxyServer) + } + return api.AGPL, closers, nil }) diff --git a/enterprise/cli/testdata/coder_server_--help.golden b/enterprise/cli/testdata/coder_server_--help.golden index 20c3c2609e..7e09efc337 100644 --- a/enterprise/cli/testdata/coder_server_--help.golden +++ b/enterprise/cli/testdata/coder_server_--help.golden @@ -140,6 +140,20 @@ AI BRIDGE OPTIONS: Maximum number of AI Bridge requests per second per replica. Set to 0 to disable (unlimited). +AI BRIDGE PROXY OPTIONS: + --aibridge-proxy-cert-file string, $CODER_AIBRIDGE_PROXY_CERT_FILE + Path to the CA certificate file for AI Bridge Proxy. + + --aibridge-proxy-enabled bool, $CODER_AIBRIDGE_PROXY_ENABLED (default: false) + Enable the AI Bridge MITM Proxy for intercepting and decrypting AI + provider requests. + + --aibridge-proxy-key-file string, $CODER_AIBRIDGE_PROXY_KEY_FILE + Path to the CA private key file for AI Bridge Proxy. + + --aibridge-proxy-listen-addr string, $CODER_AIBRIDGE_PROXY_LISTEN_ADDR (default: :8888) + The address the AI Bridge Proxy will listen on. + CLIENT OPTIONS: These options change the behavior of how clients interact with the Coder. Clients include the Coder CLI, Coder Desktop, IDE extensions, and the web UI. diff --git a/go.mod b/go.mod index f209f86007..81961b06e5 100644 --- a/go.mod +++ b/go.mod @@ -482,6 +482,7 @@ require ( github.com/coder/preview v1.0.4 github.com/danieljoos/wincred v1.2.3 github.com/dgraph-io/ristretto/v2 v2.3.0 + github.com/elazarl/goproxy v1.7.2 github.com/fsnotify/fsnotify v1.9.0 github.com/go-git/go-git/v5 v5.16.2 github.com/icholy/replace v0.6.0 diff --git a/go.sum b/go.sum index e188e17e0d..eab5176f08 100644 --- a/go.sum +++ b/go.sum @@ -1049,6 +1049,8 @@ github.com/elastic/go-sysinfo v1.15.1 h1:zBmTnFEXxIQ3iwcQuk7MzaUotmKRp3OabbbWM8T github.com/elastic/go-sysinfo v1.15.1/go.mod h1:jPSuTgXG+dhhh0GKIyI2Cso+w5lPJ5PvVqKlL8LV/Hk= github.com/elastic/go-windows v1.0.0 h1:qLURgZFkkrYyTTkvYpsZIgf83AUsdIHfvlJaqaZ7aSY= github.com/elastic/go-windows v1.0.0/go.mod h1:TsU0Nrp7/y3+VwE82FoZF8gC/XFg/Elz6CcloAxnPgU= +github.com/elazarl/goproxy v1.7.2 h1:Y2o6urb7Eule09PjlhQRGNsqRfPmYI3KKQLFpCAV3+o= +github.com/elazarl/goproxy v1.7.2/go.mod h1:82vkLNir0ALaW14Rc399OTTjyNREgmdL2cVoIbS6XaE= github.com/emersion/go-sasl v0.0.0-20200509203442-7bfe0ed36a21 h1:OJyUGMJTzHTd1XQp98QTaHernxMYzRaOasRir9hUlFQ= github.com/emersion/go-sasl v0.0.0-20200509203442-7bfe0ed36a21/go.mod h1:iL2twTeMvZnrg54ZoPDNfJaJaqy0xIQFuBdrLsmspwQ= github.com/emersion/go-smtp v0.21.2 h1:OLDgvZKuofk4em9fT5tFG5j4jE1/hXnX75UMvcrL4AA= diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index f9cd42e2c3..f55bff3a3d 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -65,6 +65,14 @@ export interface AIBridgeOpenAIConfig { readonly key: string; } +// From codersdk/deployment.go +export interface AIBridgeProxyConfig { + readonly enabled: boolean; + readonly listen_addr: string; + readonly cert_file: string; + readonly key_file: string; +} + // From codersdk/aibridge.go export interface AIBridgeTokenUsage { readonly id: string; @@ -106,6 +114,7 @@ export interface AIBridgeUserPrompt { // From codersdk/deployment.go export interface AIConfig { readonly bridge?: AIBridgeConfig; + readonly aibridge_proxy?: AIBridgeProxyConfig; } // From codersdk/allowlist.go