From 3517772e92f20b6f31808ec172268b060554a213 Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Mon, 29 Dec 2025 15:56:58 +0000 Subject: [PATCH] feat: add proxy authorization to aibridgeproxyd (#21342) ## Description This PR adds proxy authorization to the AI Bridge Proxy server. Clients provide their Coder session token via the proxy password field on the HTTP Proxy settings (`HTTPS_PROXY=http://ignored:@host:port`), which is then used for forwarding to aibridged to handle authorization. ## Changes * Extract Coder session token from `Proxy-Authorization` header during CONNECT * Reject requests without valid credentials * Store token in `ctx.UserData` for downstream request handlers * Add `Addr()` method to get the actual listening address (useful for tests with port 0) Related to: https://github.com/coder/internal/issues/1181 --- enterprise/aibridgeproxyd/aibridgeproxyd.go | 151 ++++++++- .../aibridgeproxyd/aibridgeproxyd_test.go | 289 ++++++++++++++---- 2 files changed, 373 insertions(+), 67 deletions(-) diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index adf6848aa6..9a5aefc0b4 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -4,8 +4,11 @@ import ( "context" "crypto/tls" "crypto/x509" + "encoding/base64" "errors" + "net" "net/http" + "strings" "sync" "time" @@ -28,9 +31,11 @@ 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 } // Options configures the AI Bridge Proxy server. @@ -41,6 +46,9 @@ type Options struct { CertFile string // KeyFile is the path to the CA private key file used for MITM. KeyFile string + // AllowedPorts is the list of ports allowed for CONNECT requests. + // Defaults to ["80", "443"] if empty. + AllowedPorts []string } func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) { @@ -61,27 +69,43 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) 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{ + ctx: ctx, logger: logger, proxy: proxy, } + // 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. + // 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) + + // 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) + if err != nil { + return nil, xerrors.Errorf("failed to listen on %s: %w", opts.ListenAddr, err) + } + srv.listener = listener + // 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.Info(ctx, "starting AI Bridge Proxy", slog.F("addr", listener.Addr().String())) + if err := srv.httpServer.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) { logger.Error(ctx, "aibridgeproxyd server error", slog.Error(err)) } }() @@ -89,6 +113,15 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) return srv, nil } +// Addr returns the address the server is listening on. +// This is useful when the server was started with port 0. +func (s *Server) Addr() string { + if s.listener == nil { + return "" + } + return s.listener.Addr().String() +} + // Close gracefully shuts down the proxy server. func (s *Server) Close() error { if s.httpServer == nil { @@ -124,3 +157,103 @@ func loadMitmCertificate(certFile, keyFile string) error { return nil } + +// portMiddleware is a CONNECT middleware that rejects requests to non-standard ports. +// This prevents the proxy from being used to tunnel to arbitrary services (SSH, databases, etc.). +func (s *Server) portMiddleware(allowedPorts []string) func(host string, ctx *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) { + allowed := make(map[string]bool, len(allowedPorts)) + for _, p := range allowedPorts { + allowed[p] = true + } + + return func(host string, _ *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) { + _, port, err := net.SplitHostPort(host) + if err != nil { + s.logger.Warn(s.ctx, "rejecting CONNECT with invalid host format", + slog.F("host", host), + slog.Error(err), + ) + return goproxy.RejectConnect, host + } + if port == "" { + s.logger.Warn(s.ctx, "rejecting CONNECT with empty port", + slog.F("host", host), + ) + return goproxy.RejectConnect, host + } + + if !allowed[port] { + s.logger.Warn(s.ctx, "rejecting CONNECT to non-allowed port", + slog.F("host", host), + slog.F("port", port), + ) + return goproxy.RejectConnect, host + } + + return 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. +// Requests without valid credentials are rejected. +// +// Clients provide credentials by setting their HTTP Proxy as: +// +// HTTPS_PROXY=http://ignored:@host:port +// +// The token is extracted from the password field of basic auth. +func (s *Server) authMiddleware(host string, ctx *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) { + proxyAuth := ctx.Req.Header.Get("Proxy-Authorization") + coderToken := extractCoderTokenFromProxyAuth(proxyAuth) + + // Reject requests without valid credentials. + if coderToken == "" { + hasAuth := proxyAuth != "" + s.logger.Warn(s.ctx, "rejecting CONNECT request", + slog.F("host", host), + slog.F("reason", map[bool]string{true: "invalid_credentials", false: "missing_credentials"}[hasAuth]), + ) + return goproxy.RejectConnect, host + } + + // Store the token in UserData for downstream handlers. + // goproxy propagates UserData to subsequent request contexts + // for decrypted requests within this MITM session. + ctx.UserData = coderToken + + return goproxy.MitmConnect, host +} + +// extractCoderTokenFromProxyAuth extracts the Coder session token from the +// Proxy-Authorization header. The token is expected to be in the password +// field of basic auth: "Basic base64(username:token)". +// +// Returns empty string if no valid token is found. +func extractCoderTokenFromProxyAuth(proxyAuth string) string { + if proxyAuth == "" { + return "" + } + + // Expected format: "Basic base64(username:password)" + // Auth scheme is case-insensitive per RFC 7235. + parts := strings.Fields(proxyAuth) + if len(parts) != 2 || !strings.EqualFold(parts[0], "Basic") { + return "" + } + + decoded, err := base64.StdEncoding.DecodeString(parts[1]) + if err != nil { + return "" + } + + // Format: "username:password", password is the Coder token. + // Username is ignored and can be any value. + credentials := strings.SplitN(string(decoded), ":", 2) + if len(credentials) != 2 { + return "" + } + + return credentials[1] +} diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index 8b5b05773a..91369ad073 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -6,6 +6,7 @@ import ( "crypto/tls" "crypto/x509" "crypto/x509/pkix" + "encoding/base64" "encoding/pem" "io" "math/big" @@ -94,6 +95,13 @@ func generateSharedTestCA() (certFile, keyFile string, err error) { return certPath, keyPath, nil } +// makeProxyAuthHeader creates a Proxy-Authorization header value with the given token. +// Format: "Basic base64(username:token)" where username is "ignored". +func makeProxyAuthHeader(token string) string { + credentials := base64.StdEncoding.EncodeToString([]byte("ignored:" + token)) + return "Basic " + credentials +} + func TestNew(t *testing.T) { t.Parallel() @@ -192,68 +200,233 @@ func TestClose(t *testing.T) { require.NoError(t, err) } -func TestProxy_MITM(t *testing.T) { +func TestProxy_PortValidation(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, - }, + tests := []struct { + name string + allowPort bool + expectError bool + }{ + { + name: "AllowedPort", + allowPort: true, + }, + { + name: "RejectedPort", + allowPort: false, + expectError: true, }, } - // 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() + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() - // 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)) + // Create a target HTTPS server that will be the destination 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() + + targetURL, err := url.Parse(targetServer.URL) + require.NoError(t, err) + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + // Configure allowed ports based on test case. + // For allowed case, include the target's random port. + // For rejected case, only allow port 443 which doesn't match the target. + var allowedPorts []string + if tt.allowPort { + allowedPorts = []string{targetURL.Port()} + } else { + allowedPorts = []string{"443"} + } + + // 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, + }) + 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 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://" + 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, + }, + }, + } + + // 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) + if tt.expectError { + require.Error(t, err) + return + } + + 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) + require.Equal(t, "hello from target", string(body)) + }) + } +} + +func TestProxy_Authentication(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + proxyAuth string + expectError bool + }{ + { + name: "ValidCredentials", + proxyAuth: makeProxyAuthHeader("test-coder-session-token"), + expectError: false, + }, + { + name: "MissingCredentials", + proxyAuth: "", + expectError: true, + }, + { + name: "InvalidBase64", + proxyAuth: "Basic not-valid-base64!", + expectError: true, + }, + { + name: "EmptyToken", + proxyAuth: makeProxyAuthHeader(""), + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(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() + + targetURL, err := url.Parse(targetServer.URL) + require.NoError(t, err) + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + // 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()}, + }) + 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 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://" + proxyAddr) + require.NoError(t, err) + + transport := &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + TLSClientConfig: &tls.Config{ + MinVersion: tls.VersionTLS12, + RootCAs: certPool, + }, + } + + if tt.proxyAuth != "" { + transport.ProxyConnectHeader = http.Header{ + "Proxy-Authorization": []string{tt.proxyAuth}, + } + } + + client := &http.Client{Transport: transport} + + // 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) + + if tt.expectError { + require.Error(t, err) + } else { + 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)) + } + }) + } }