diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index 9a5aefc0b4..f57a353e64 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -8,6 +8,7 @@ import ( "errors" "net" "net/http" + "net/url" "strings" "sync" "time" @@ -16,6 +17,14 @@ import ( "golang.org/x/xerrors" "cdr.dev/slog" + + "github.com/coder/aibridge" +) + +// Known AI provider hosts. +const ( + HostAnthropic = "api.anthropic.com" + HostOpenAI = "api.openai.com" ) // loadMitmOnce ensures the MITM certificate is loaded exactly once. @@ -31,17 +40,21 @@ var loadMitmOnce sync.Once // - decrypting requests using the configured CA certificate // - forwarding requests to aibridged for processing type Server struct { - ctx context.Context - logger slog.Logger - proxy *goproxy.ProxyHttpServer - httpServer *http.Server - listener net.Listener + ctx context.Context + logger slog.Logger + proxy *goproxy.ProxyHttpServer + httpServer *http.Server + listener net.Listener + coderAccessURL *url.URL } // Options configures the AI Bridge Proxy server. type Options struct { // ListenAddr is the address the proxy server will listen on. ListenAddr string + // CoderAccessURL is the URL of the Coder deployment where aibridged is running. + // Requests to supported AI providers are forwarded here. + CoderAccessURL 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. @@ -62,6 +75,14 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) 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") + } + coderAccessURL, err := url.Parse(opts.CoderAccessURL) + if err != nil { + return nil, xerrors.Errorf("invalid coder access URL %q: %w", opts.CoderAccessURL, err) + } + // 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) @@ -70,9 +91,10 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) proxy := goproxy.NewProxyHttpServer() srv := &Server{ - ctx: ctx, - logger: logger, - proxy: proxy, + ctx: ctx, + logger: logger, + proxy: proxy, + coderAccessURL: coderAccessURL, } // Reject CONNECT requests to non-standard ports. @@ -83,12 +105,14 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) proxy.OnRequest().HandleConnectFunc(srv.portMiddleware(allowedPorts)) // Extract Coder session token from proxy authentication to forward to aibridged. - // 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().HandleConnectFunc(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. // This is useful in tests where port 0 is used to avoid conflicts. listener, err := net.Listen("tcp", opts.ListenAddr) @@ -257,3 +281,97 @@ func extractCoderTokenFromProxyAuth(proxyAuth string) string { return credentials[1] } + +// providerFromHost maps the request host to the aibridge provider name. +// - 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 "" + } + switch strings.ToLower(reqURL.Hostname()) { + case HostAnthropic: + return aibridge.ProviderAnthropic + case HostOpenAI: + return aibridge.ProviderOpenAI + default: + return "" + } +} + +// handleRequest intercepts HTTP requests after MITM decryption. +// - Requests to known AI providers are rewritten to aibridged, with the Coder session token +// (from ctx.UserData, set during CONNECT) injected in the Authorization header. +// - Unknown hosts are passed through to the original upstream. +func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) { + originalPath := req.URL.Path + + // 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", + slog.F("host", req.Host), + slog.F("method", req.Method), + slog.F("path", originalPath), + ) + return req, nil + } + + // Get the Coder session token stored during CONNECT. + coderToken, _ := ctx.UserData.(string) + + // Reject unauthenticated requests to AI providers. + if coderToken == "" { + s.logger.Warn(s.ctx, "rejecting unauthenticated request to AI provider", + slog.F("host", req.Host), + slog.F("provider", provider), + ) + resp := goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusProxyAuthRequired, "Proxy authentication required") + // Describe to the client how to authenticate with the proxy. + resp.Header.Set("Proxy-Authenticate", `Basic realm="Coder AI Bridge Proxy"`) + return req, resp + } + + // Rewrite the request to point to aibridged. + if s.coderAccessURL == nil || s.coderAccessURL.String() == "" { + s.logger.Error(s.ctx, "coderAccessURL is not configured") + return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Proxy misconfigured") + } + + aiBridgeURL, err := url.JoinPath(s.coderAccessURL.String(), "api/v2/aibridge", provider, originalPath) + if err != nil { + s.logger.Error(s.ctx, "failed to build aibridged URL", slog.Error(err)) + return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Failed to build AI Bridge URL") + } + + aiBridgeParsedURL, err := url.Parse(aiBridgeURL) + if err != nil { + s.logger.Error(s.ctx, "failed to parse aibridged URL", slog.Error(err)) + return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Failed to parse AI Bridge URL") + } + + // Preserve query parameters from the original request. + // Both URL and Host must be set for the request to be properly routed. + aiBridgeParsedURL.RawQuery = req.URL.RawQuery + req.URL = aiBridgeParsedURL + req.Host = aiBridgeParsedURL.Host + + // Set Authorization header for aibridged authentication. + req.Header.Set("Authorization", "Bearer "+coderToken) + + s.logger.Debug(s.ctx, "routing request to aibridged", + slog.F("provider", provider), + slog.F("original_path", originalPath), + slog.F("aibridged_url", aiBridgeParsedURL.String()), + ) + + return req, nil +} diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index 91369ad073..c4ca3011e5 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -16,6 +16,7 @@ import ( "net/url" "os" "path/filepath" + "strings" "sync" "testing" "time" @@ -112,12 +113,60 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "", + ListenAddr: "", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "listen address is required") + }) + + t.Run("MissingCoderAccessURL", 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", CertFile: certFile, KeyFile: keyFile, }) require.Error(t, err) - require.Contains(t, err.Error(), "listen address is required") + require.Contains(t, err.Error(), "coder access URL is required") + }) + + t.Run("EmptyCoderAccessURL", 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: " ", + CertFile: certFile, + KeyFile: keyFile, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "coder access URL is required") + }) + + t.Run("InvalidCoderAccessURL", 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: "://invalid", + CertFile: certFile, + KeyFile: keyFile, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "invalid coder access URL") }) t.Run("MissingCertFile", func(t *testing.T) { @@ -126,8 +175,9 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - KeyFile: "key.pem", + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + KeyFile: "key.pem", }) require.Error(t, err) require.Contains(t, err.Error(), "cert file and key file are required") @@ -139,8 +189,9 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CertFile: "cert.pem", + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + CertFile: "cert.pem", }) require.Error(t, err) require.Contains(t, err.Error(), "cert file and key file are required") @@ -152,9 +203,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CertFile: "/nonexistent/cert.pem", - KeyFile: "/nonexistent/key.pem", + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + CertFile: "/nonexistent/cert.pem", + KeyFile: "/nonexistent/key.pem", }) require.Error(t, err) require.Contains(t, err.Error(), "failed to load MITM certificate") @@ -167,9 +219,10 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CertFile: certFile, - KeyFile: keyFile, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, }) require.NoError(t, err) require.NotNil(t, srv) @@ -186,9 +239,10 @@ func TestClose(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CertFile: certFile, - KeyFile: keyFile, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, }) require.NoError(t, err) @@ -228,7 +282,7 @@ func TestProxy_PortValidation(t *testing.T) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("hello from target")) })) - defer targetServer.Close() + t.Cleanup(func() { targetServer.Close() }) targetURL, err := url.Parse(targetServer.URL) require.NoError(t, err) @@ -248,10 +302,11 @@ func TestProxy_PortValidation(t *testing.T) { // Start the proxy server on a random port to avoid conflicts when running tests in parallel. srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CertFile: certFile, - KeyFile: keyFile, - AllowedPorts: allowedPorts, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + AllowedPorts: allowedPorts, }) require.NoError(t, err) t.Cleanup(func() { _ = srv.Close() }) @@ -352,7 +407,7 @@ func TestProxy_Authentication(t *testing.T) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("hello from target")) })) - defer targetServer.Close() + t.Cleanup(func() { targetServer.Close() }) targetURL, err := url.Parse(targetServer.URL) require.NoError(t, err) @@ -363,10 +418,11 @@ func TestProxy_Authentication(t *testing.T) { // Start the proxy server on a random port to avoid conflicts when running tests in parallel. // The actual port is accessible via srv.Addr(). srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CertFile: certFile, - KeyFile: keyFile, - AllowedPorts: []string{targetURL.Port()}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + AllowedPorts: []string{targetURL.Port()}, }) require.NoError(t, err) t.Cleanup(func() { _ = srv.Close() }) @@ -430,3 +486,181 @@ func TestProxy_Authentication(t *testing.T) { }) } } + +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: "AnthropicMessages", + targetHost: "api.anthropic.com", + targetPath: "/v1/messages", + expectedPath: "/api/v2/aibridge/anthropic/v1/messages", + }, + { + name: "AnthropicNonDefaultPort", + targetHost: "api.anthropic.com", + targetPort: "8443", + targetPath: "/v1/messages", + expectedPath: "/api/v2/aibridge/anthropic/v1/messages", + }, + { + name: "OpenAIChatCompletions", + targetHost: "api.openai.com", + targetPath: "/v1/chat/completions", + expectedPath: "/api/v2/aibridge/openai/v1/chat/completions", + }, + { + name: "OpenAINonDefaultPort", + targetHost: "api.openai.com", + targetPort: "8443", + targetPath: "/v1/chat/completions", + expectedPath: "/api/v2/aibridge/openai/v1/chat/completions", + }, + { + name: "UnknownHostPassthrough", + targetPath: "/some/path", + passthrough: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Track what aibridged receives. + var receivedPath string + var receivedAuth string + + // Create a mock aibridged server. + aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedPath = r.URL.Path + receivedAuth = r.Header.Get("Authorization") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("hello from aibridged")) + })) + t.Cleanup(func() { aibridgedServer.Close() }) + + // Create a mock target server for passthrough tests. + targetServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("hello from passthrough")) + })) + t.Cleanup(func() { targetServer.Close() }) + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + // 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: + parsedTargetURL, err := url.Parse(targetServer.URL) + require.NoError(t, err) + allowedPorts = []string{parsedTargetURL.Port()} + case tt.targetPort != "": + allowedPorts = []string{tt.targetPort} + default: + allowedPorts = []string{"443"} + } + + // Start the proxy server pointing to our mock aibridged. + srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "127.0.0.1:0", + CoderAccessURL: aibridgedServer.URL, + CertFile: certFile, + KeyFile: keyFile, + AllowedPorts: allowedPorts, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Close() }) + + proxyAddr := srv.Addr() + require.NotEmpty(t, proxyAddr) + + // Wait for the proxy server to be ready. + require.Eventually(t, func() bool { + conn, err := net.Dial("tcp", proxyAddr) + if err != nil { + return false + } + _ = conn.Close() + return true + }, testutil.WaitShort, testutil.IntervalFast) + + // Load the CA 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://" + proxyAddr) + require.NoError(t, err) + + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + ProxyConnectHeader: http.Header{ + "Proxy-Authorization": []string{makeProxyAuthHeader("test-session-token")}, + }, + TLSClientConfig: &tls.Config{ + MinVersion: tls.VersionTLS12, + RootCAs: certPool, + }, + }, + } + + // 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 + switch { + case tt.passthrough: + targetURL, err = url.JoinPath(targetServer.URL, 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) + } + + // Make a request through the proxy to the target URL. + req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, targetURL, strings.NewReader(`{}`)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + + resp, err := client.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + + if tt.passthrough { + // 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") + } 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) + } + }) + } +} diff --git a/enterprise/cli/aibridgeproxyd.go b/enterprise/cli/aibridgeproxyd.go index 6c2a388a98..cc613ad2ba 100644 --- a/enterprise/cli/aibridgeproxyd.go +++ b/enterprise/cli/aibridgeproxyd.go @@ -18,9 +18,10 @@ 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(), - 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(), }) if err != nil { return nil, xerrors.Errorf("failed to start in-memory aibridgeproxy daemon: %w", err)