From 47b3846bca0b3ca23a1cd2652e780320db71151d Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Wed, 21 Jan 2026 19:06:19 +0000 Subject: [PATCH] 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 | 15 +++- enterprise/aibridged/aibridged_test.go | 80 ++++++++++++++++++- enterprise/aibridged/http.go | 3 + enterprise/aibridgeproxyd/aibridgeproxyd.go | 21 ++--- .../aibridgeproxyd/aibridgeproxyd_test.go | 28 ++++--- 5 files changed, 120 insertions(+), 27 deletions(-) diff --git a/coderd/aibridge/aibridge.go b/coderd/aibridge/aibridge.go index cb656b5ec5..4a9adee62e 100644 --- a/coderd/aibridge/aibridge.go +++ b/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") { diff --git a/enterprise/aibridged/aibridged_test.go b/enterprise/aibridged/aibridged_test.go index 32469c6dfe..6e0bebb7ad 100644 --- a/enterprise/aibridged/aibridged_test.go +++ b/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)) } diff --git a/enterprise/aibridged/http.go b/enterprise/aibridged/http.go index d2a6b1a49f..087702be03 100644 --- a/enterprise/aibridged/http.go +++ b/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)) diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index 62da187200..8886d9b0a3 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/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), diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index 6430420de0..3c0103a303 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/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,