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:<coder-session-token>@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
2026-09-24 15:04:27 +08:00 · 2025-12-29 15:56:58 +00:00
parent b97572285a
commit 3517772e92
2 changed files with 373 additions and 67 deletions
@@ -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:<coder-token>@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]
+}
@@ -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))
+			}
+		})
+	}
 }