feat: use coder specific header for aibridge authentication from AI proxy (#21590)

## Description

Introduces a new `X-Coder-Token` header for authenticating requests from
AI Proxy to AI Bridge. Previously, the proxy overwrote the
`Authorization` header with the Coder token, which prevented the
original authentication headers from flowing through to upstream
providers.

With this change, AI Proxy sets the Coder token in a separate header,
preserving the original `Authorization` and `X-Api-Key` headers. AI
Bridge uses this header for authentication and removes it before
forwarding requests to upstream providers. For requests that don't come
through AI Proxy, AI Bridge continues to use `Authorization` and
`X-Api-Key` for authentication.

## Changes

* Add `HeaderCoderAuth` constant and update `ExtractAuthToken` to check
headers in the following order: `X-Coder-Token` > `Authorization` >
`X-Api-Key`
* Update AI Proxy to set `X-Coder-Token` instead of overwriting
`Authorization`
* Remove `X-Coder-Token` in AI Bridge before forwarding to upstream
providers
* Add tests for header handling and token extraction priority

Related to: https://github.com/coder/internal/issues/1235

coderd/aibridge/aibridge.go

@@ -6,11 +6,20 @@ import (
 	"strings"
 )
 
+// HeaderCoderAuth is an internal header used to pass the Coder token
+// from AI Proxy to AI Bridge for authentication. This header is stripped
+// by AI Bridge before forwarding requests to upstream providers.
+const HeaderCoderAuth = "X-Coder-Token"
+
 // ExtractAuthToken extracts an authorization token from HTTP headers.
-// It checks the Authorization header (Bearer token) and X-Api-Key header,
-// which represent the different ways clients authenticate against AI providers.
-// If neither are present, an empty string is returned.
+// It checks X-Coder-Token first (set by AI Proxy), then falls back
+// to Authorization header (Bearer token) and X-Api-Key header, which represent
+// the different ways clients authenticate against AI providers.
+// If none are present, an empty string is returned.
 func ExtractAuthToken(header http.Header) string {
+	if token := strings.TrimSpace(header.Get(HeaderCoderAuth)); token != "" {
+		return token
+	}
 	if auth := strings.TrimSpace(header.Get("Authorization")); auth != "" {
 		fields := strings.Fields(auth)
 		if len(fields) == 2 && strings.EqualFold(fields[0], "Bearer") {

enterprise/aibridged/aibridged_test.go

@@ -173,6 +173,46 @@ func TestServeHTTP_FailureModes(t *testing.T) {
 	}
 }
 
+func TestServeHTTP_CoderTokenRemoved(t *testing.T) {
+	t.Parallel()
+
+	mockHandler := &mockHandler{}
+
+	srv, client, pool := newTestServer(t)
+	conn := &mockDRPCConn{}
+	client.EXPECT().DRPCConn().AnyTimes().Return(conn)
+	client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil)
+	pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(mockHandler, nil)
+
+	httpSrv := httptest.NewServer(srv)
+	t.Cleanup(httpSrv.Close)
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, httpSrv.URL+"/openai/v1/chat/completions", nil)
+	require.NoError(t, err)
+
+	// X-Coder-Token is used for authentication and should be stripped.
+	// Other authorization headers should be preserved.
+	req.Header.Set(agplaibridge.HeaderCoderAuth, "coder-token")
+	req.Header.Set("Authorization", "Bearer some-token")
+	req.Header.Set("X-Api-Key", "some-api-key")
+
+	resp, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+
+	// Verify X-Coder-Token was removed before forwarding to handler.
+	require.NotNil(t, mockHandler.headersReceived)
+	require.Empty(t, mockHandler.headersReceived.Get(agplaibridge.HeaderCoderAuth),
+		"X-Coder-Token should be removed before forwarding to handler")
+
+	// Verify other headers were preserved.
+	require.Equal(t, "Bearer some-token", mockHandler.headersReceived.Get("Authorization"))
+	require.Equal(t, "some-api-key", mockHandler.headersReceived.Get("X-Api-Key"))
+}
+
 func TestExtractAuthToken(t *testing.T) {
 	t.Parallel()
 
@@ -184,6 +224,31 @@ func TestExtractAuthToken(t *testing.T) {
 		{
 			name: "none",
 		},
+		{
+			name:    "x-coder-token/empty",
+			headers: map[string]string{agplaibridge.HeaderCoderAuth: ""},
+		},
+		{
+			name:        "x-coder-token/ok",
+			headers:     map[string]string{agplaibridge.HeaderCoderAuth: "coder-token"},
+			expectedKey: "coder-token",
+		},
+		{
+			name: "x-coder-token/priority over authorization",
+			headers: map[string]string{
+				agplaibridge.HeaderCoderAuth: "coder-token",
+				"Authorization":              "Bearer other-token",
+			},
+			expectedKey: "coder-token",
+		},
+		{
+			name: "x-coder-token/priority over x-api-key",
+			headers: map[string]string{
+				agplaibridge.HeaderCoderAuth: "coder-token",
+				"X-Api-Key":                  "api-key",
+			},
+			expectedKey: "coder-token",
+		},
 		{
 			name:    "authorization/invalid",
 			headers: map[string]string{"authorization": "invalid"},
@@ -202,6 +267,14 @@ func TestExtractAuthToken(t *testing.T) {
 			headers:     map[string]string{"AUTHORIZATION": "BEARer key"},
 			expectedKey: "key",
 		},
+		{
+			name: "authorization/priority over x-api-key",
+			headers: map[string]string{
+				"Authorization": "Bearer auth-token",
+				"X-Api-Key":     "api-key",
+			},
+			expectedKey: "auth-token",
+		},
 		{
 			name:    "x-api-key/empty",
 			headers: map[string]string{"X-Api-Key": ""},
@@ -229,9 +302,12 @@ func TestExtractAuthToken(t *testing.T) {
 
 var _ http.Handler = &mockHandler{}
 
-type mockHandler struct{}
+type mockHandler struct {
+	headersReceived http.Header
+}
 
-func (*mockHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
+func (h *mockHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
+	h.headersReceived = r.Header.Clone()
 	rw.WriteHeader(http.StatusOK)
 	_, _ = rw.Write([]byte(r.URL.Path))
 }

enterprise/aibridged/http.go

@@ -43,6 +43,9 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// Remove the Coder token header so it's not forwarded to upstream providers.
+	r.Header.Del(agplaibridge.HeaderCoderAuth)
+
 	client, err := s.Client()
 	if err != nil {
 		logger.Warn(ctx, "failed to connect to coderd", slog.Error(err))

enterprise/aibridgeproxyd/aibridgeproxyd.go

@@ -22,6 +22,7 @@ import (
 
 	"cdr.dev/slog/v3"
 	"github.com/coder/aibridge"
+	agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
 )
 
 // Known AI provider hosts.
@@ -237,7 +238,7 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error)
 		// 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.
+		// Extract Coder token from proxy authentication to forward to aibridged.
 		srv.authMiddleware,
 	)
 
@@ -392,8 +393,8 @@ func convertDomainsToHosts(domains []string, allowedPorts []string) ([]string, e
 	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
+// authMiddleware is a CONNECT middleware that extracts the Coder token from
+// the Proxy-Authorization header and stores it in ctx.UserData for use by
 // downstream request handlers.
 // Requests without valid credentials are rejected.
 //
@@ -424,7 +425,7 @@ func (s *Server) authMiddleware(host string, ctx *goproxy.ProxyCtx) (*goproxy.Co
 	return goproxy.MitmConnect, host
 }
 
-// extractCoderTokenFromProxyAuth extracts the Coder session token from the
+// extractCoderTokenFromProxyAuth extracts the Coder token from the
 // Proxy-Authorization header. The token is expected to be in the password
 // field of basic auth: "Basic base64(username:token)".
 //
@@ -472,8 +473,8 @@ func defaultAIBridgeProvider(host string) string {
 }
 
 // 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.
+//   - Requests to known AI providers are rewritten to aibridged, with the Coder token
+//     (from ctx.UserData, set during CONNECT) set in the X-Coder-Token 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
@@ -494,7 +495,7 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.
 		return req, nil
 	}
 
-	// Get the Coder session token stored during CONNECT.
+	// Get the Coder token stored during CONNECT.
 	coderToken, _ := ctx.UserData.(string)
 
 	// Reject unauthenticated requests to AI providers.
@@ -533,8 +534,10 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.
 	req.URL = aiBridgeParsedURL
 	req.Host = aiBridgeParsedURL.Host
 
-	// Set Authorization header for aibridged authentication.
-	req.Header.Set("Authorization", "Bearer "+coderToken)
+	// Set X-Coder-Token header for aibridged authentication.
+	// Using a separate header preserves the original request headers,
+	// which are forwarded to upstream providers.
+	req.Header.Set(agplaibridge.HeaderCoderAuth, coderToken)
 
 	s.logger.Debug(s.ctx, "routing request to aibridged",
 		slog.F("provider", provider),

enterprise/aibridgeproxyd/aibridgeproxyd_test.go

@@ -26,6 +26,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog/v3/sloggers/slogtest"
+	agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
 	"github.com/coder/coder/v2/enterprise/aibridgeproxyd"
 	"github.com/coder/coder/v2/testutil"
 )
@@ -673,7 +674,7 @@ func TestProxy_CertCaching(t *testing.T) {
 			}
 
 			// Make a request through the proxy to the target server.
-			client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token"), certPool)
+			client := newProxyClient(t, srv, makeProxyAuthHeader("test-token"), certPool)
 			req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetURL.String(), nil)
 			require.NoError(t, err)
 			resp, err := client.Do(req)
@@ -749,7 +750,7 @@ func TestProxy_PortValidation(t *testing.T) {
 			)
 
 			// Make a request through the proxy to the target server.
-			client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token"), getProxyCertPool(t))
+			client := newProxyClient(t, srv, makeProxyAuthHeader("test-token"), getProxyCertPool(t))
 			req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetURL.String(), nil)
 			require.NoError(t, err)
 
@@ -780,7 +781,7 @@ func TestProxy_Authentication(t *testing.T) {
 	}{
 		{
 			name:        "ValidCredentials",
-			proxyAuth:   makeProxyAuthHeader("test-coder-session-token"),
+			proxyAuth:   makeProxyAuthHeader("test-coder-token"),
 			expectError: false,
 		},
 		{
@@ -909,13 +910,12 @@ func TestProxy_MITM(t *testing.T) {
 			t.Parallel()
 
 			// Track what aibridged receives.
-			var receivedPath string
-			var receivedAuth string
+			var receivedPath, receivedCoderToken string
 
 			// Create a mock aibridged server that captures requests.
 			aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				receivedPath = r.URL.Path
-				receivedAuth = r.Header.Get("Authorization")
+				receivedCoderToken = r.Header.Get(agplaibridge.HeaderCoderAuth)
 				w.WriteHeader(http.StatusOK)
 				_, _ = w.Write([]byte("hello from aibridged"))
 			}))
@@ -966,7 +966,7 @@ func TestProxy_MITM(t *testing.T) {
 			}
 
 			// Make a request through the proxy to the target URL.
-			client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token"), certPool)
+			client := newProxyClient(t, srv, makeProxyAuthHeader("test-token"), certPool)
 			req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, targetURL, strings.NewReader(`{}`))
 			require.NoError(t, err)
 			req.Header.Set("Content-Type", "application/json")
@@ -983,12 +983,12 @@ func TestProxy_MITM(t *testing.T) {
 				// Verify request went to target server, not aibridged.
 				require.Equal(t, "hello from tunneled", string(body))
 				require.Empty(t, receivedPath, "aibridged should not receive tunneled requests")
-				require.Empty(t, receivedAuth, "tunneled requests are not authenticated by the proxy")
+				require.Empty(t, receivedCoderToken, "tunneled 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, "MITM'd requests must include authentication")
+				require.Equal(t, "test-token", receivedCoderToken, "MITM'd requests must include Coder token")
 			}
 		})
 	}
@@ -1173,7 +1173,7 @@ func TestUpstreamProxy(t *testing.T) {
 				finalDestinationBody         string
 				aibridgeReceived             bool
 				aibridgePath                 string
-				aibridgeAuthHeader           string
+				aibridgeCoderToken           string
 				aibridgeBody                 string
 			)
 
@@ -1269,7 +1269,7 @@ func TestUpstreamProxy(t *testing.T) {
 			aibridgeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				aibridgeReceived = true
 				aibridgePath = r.URL.Path
-				aibridgeAuthHeader = r.Header.Get("Authorization")
+				aibridgeCoderToken = r.Header.Get(agplaibridge.HeaderCoderAuth)
 				body, err := io.ReadAll(r.Body)
 				if err != nil {
 					http.Error(w, err.Error(), http.StatusInternalServerError)
@@ -1345,6 +1345,8 @@ func TestUpstreamProxy(t *testing.T) {
 					"final destination should receive the exact request body")
 				require.False(t, aibridgeReceived,
 					"aibridge should NOT receive request for non-allowlisted domain")
+				require.Empty(t, aibridgeCoderToken,
+					"tunneled requests should not have Coder token")
 			} else {
 				require.False(t, upstreamProxyCONNECTReceived,
 					"upstream proxy should NOT receive CONNECT for allowlisted domain")
@@ -1352,8 +1354,8 @@ func TestUpstreamProxy(t *testing.T) {
 					"aibridge should receive the MITM'd request")
 				require.Equal(t, tt.expectedAIBridgePath, aibridgePath,
 					"aibridge should receive rewritten path")
-				require.Equal(t, "Bearer test-coder-token", aibridgeAuthHeader,
-					"aibridge should receive auth header extracted from proxy auth")
+				require.Equal(t, "test-coder-token", aibridgeCoderToken,
+					"aibridge should receive Coder token header")
 				require.Equal(t, requestBody, aibridgeBody,
 					"aibridge should receive the exact request body")
 				require.False(t, finalDestinationReceived,