fix(chatd): extract session token from cookie for relay header (#22649)

## Problem When a browser connects to the chat stream via WebSocket, it authenticates using cookies only — the native WebSocket API cannot set custom headers like `Coder-Session-Token`. The relay between replicas copies the original request's `Cookie` header but did **not** set the `Coder-Session-Token` header as a fallback. This causes a **401 on the worker replica** when `EnableHostPrefix` is enabled, because the `HTTPCookies.Middleware` strips bare `coder_session_token` cookies (expecting the `__Host-` prefix). Without a `Coder-Session-Token` header fallback, `apiKeyMiddleware` finds no valid credentials. ### Root Cause The data flow: 1. Browser → subscriber replica: `Cookie: __Host-coder_session_token=xxx` (browser sends prefixed cookie) 2. Subscriber's `HTTPCookies.Middleware` normalizes: `Cookie: coder_session_token=xxx` (strips prefix) 3. `relayHeaders()` copies `Cookie: coder_session_token=xxx` to relay request 4. Worker replica's `HTTPCookies.Middleware` sees bare `coder_session_token` → **strips it** (expects `__Host-` prefix) 5. `apiKeyMiddleware` → `APITokenFromRequest`: no cookie, no header → **401** ## Fix Modified `relayHeaders()` to extract the session token value from the `Cookie` header and set it as the `Coder-Session-Token` header when no explicit session token header is already present. The header is never stripped by middleware, so the worker replica can always authenticate. ## Testing - **`TestRelayHeaders`**: Unit tests for the updated `relayHeaders()` function covering all scenarios (cookie-only, header+cookie, no auth, nil source) - **`TestExtractSessionTokenFromCookieHeader`**: Unit tests for the helper function - **`TestChatStreamRelay/RelayCookieOnlyAuth`**: Integration test with plain HTTP, cookie-only WebSocket auth - **`TestChatStreamRelay/RelayCookieOnlyAuthWithHostPrefix`**: Integration test with `EnableHostPrefix=true`, confirming the 401 is fixed - **`cookieOnlySessionTokenProvider`**: Test helper that simulates browser WebSocket behavior (sets Cookie header only on WebSocket dials, no custom headers) ## Files Changed - `enterprise/coderd/chatd/chatd.go` — `relayHeaders()` fix + `extractSessionTokenFromCookieHeader()` helper - `enterprise/coderd/chatd/relay_headers_internal_test.go` — unit tests (new file) - `enterprise/coderd/chats_test.go` — integration tests + test helper type
2026-09-24 15:04:27 +08:00 · 2026-03-05 05:11:07 +00:00
parent f609de860f
commit 94a2e440a8
2 changed files with 417 additions and 27 deletions
@@ -5,6 +5,7 @@ import (
 	"math"
 	"net/http"
 	"net/url"
+	"strings"
 	"time"

 	"github.com/google/uuid"
@@ -450,8 +451,9 @@ func dialRelay(
 	relayCtx, relayCancel := context.WithCancel(ctx)
 	sdkClient := codersdk.New(baseURL)
 	sdkClient.HTTPClient = cfg.ReplicaHTTPClient
-	sdkClient.SessionTokenProvider = relayHeaderTokenProvider{
-		header: relayHeaders(requestHeader, replicaID),
+	sdkClient.SessionTokenProvider = relayTokenProvider{
+		token:     extractSessionToken(requestHeader),
+		replicaID: replicaID,
 	}
 	sourceEvents, sourceStream, err := sdkClient.StreamChat(relayCtx, chatID, &codersdk.StreamChatOptions{
 		AfterID: ptr.Ref(int64(math.MaxInt64)),
@@ -532,44 +534,57 @@ drainInitial:
 	return snapshot, events, cancelFn, nil
 }

-type relayHeaderTokenProvider struct {
-	header http.Header
+// relayTokenProvider authenticates relay requests to the worker
+// replica using the session token extracted from the original
+// browser request. It also stamps each request with the relay
+// source header so the worker can identify it as an inter-replica
+// call.
+type relayTokenProvider struct {
+	token     string
+	replicaID uuid.UUID
 }

-func (p relayHeaderTokenProvider) AsRequestOption() codersdk.RequestOption {
+func (p relayTokenProvider) AsRequestOption() codersdk.RequestOption {
 	return func(req *http.Request) {
-		for key, values := range p.header {
-			for _, value := range values {
-				req.Header.Add(key, value)
-			}
-		}
+		req.Header.Set(codersdk.SessionTokenHeader, p.token)
+		req.Header.Set(RelaySourceHeader, p.replicaID.String())
 	}
 }

-func (p relayHeaderTokenProvider) SetDialOption(opts *websocket.DialOptions) {
+func (p relayTokenProvider) SetDialOption(opts *websocket.DialOptions) {
 	if opts.HTTPHeader == nil {
 		opts.HTTPHeader = make(http.Header)
 	}
-	for key, values := range p.header {
-		for _, value := range values {
-			opts.HTTPHeader.Add(key, value)
-		}
-	}
+	opts.HTTPHeader.Set(codersdk.SessionTokenHeader, p.token)
+	opts.HTTPHeader.Set(RelaySourceHeader, p.replicaID.String())
 }

-func (p relayHeaderTokenProvider) GetSessionToken() string {
-	return p.header.Get(codersdk.SessionTokenHeader)
+func (p relayTokenProvider) GetSessionToken() string {
+	return p.token
 }

-func relayHeaders(source http.Header, replicaID uuid.UUID) http.Header {
-	header := make(http.Header)
-	if source != nil {
-		for _, key := range []string{codersdk.SessionTokenHeader, authorizationHeader, cookieHeader} {
-			for _, value := range source.Values(key) {
-				header.Add(key, value)
-			}
+// extractSessionToken returns the session token carried by the
+// given request headers. It mirrors the priority order used by
+// apiKeyMiddleware: cookie, then Coder-Session-Token header, then
+// Authorization: Bearer header.
+func extractSessionToken(header http.Header) string {
+	if header == nil {
+		return ""
+	}
+	// Cookie (browser WebSocket upgrade — most common relay case).
+	if raw := header.Get(cookieHeader); raw != "" {
+		r := &http.Request{Header: http.Header{cookieHeader: {raw}}}
+		if c, err := r.Cookie(codersdk.SessionTokenCookie); err == nil && c.Value != "" {
+			return c.Value
 		}
 	}
-	header.Set(RelaySourceHeader, replicaID.String())
-	return header
+	// Coder-Session-Token header (SDK / CLI callers).
+	if v := header.Get(codersdk.SessionTokenHeader); v != "" {
+		return v
+	}
+	// Authorization: Bearer <token>.
+	if v := header.Get(authorizationHeader); len(v) > 7 && strings.EqualFold(v[:7], "bearer ") {
+		return strings.TrimSpace(v[7:])
+	}
+	return ""
 }
@@ -19,6 +19,7 @@ import (
 	"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
 	"github.com/coder/coder/v2/enterprise/coderd/license"
 	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/websocket"
 )

 func TestChatStreamRelay(t *testing.T) {
@@ -367,6 +368,343 @@ func TestChatStreamRelay(t *testing.T) {
 		close(streamingChunks)
 	})

+	// This test verifies that the relay works when the subscriber
+	// replica's incoming request authenticates via cookies only,
+	// exactly as a browser WebSocket upgrade does. Browsers cannot
+	// set custom headers (like Coder-Session-Token) on WebSocket
+	// connections, so the relay must forward the Cookie header and
+	// the worker replica must accept it.
+	//
+	// Previous tests used SetSessionToken() which sets the
+	// Coder-Session-Token header, masking this code path.
+	t.Run("RelayCookieOnlyAuth", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		db, pubsub := dbtestutil.NewDB(t)
+		firstClient, _ := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				Database: db,
+				Pubsub:   pubsub,
+			},
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureHighAvailability: 1,
+				},
+			},
+		})
+
+		secondClient, _ := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				Database: db,
+				Pubsub:   pubsub,
+			},
+			DontAddLicense:   true,
+			DontAddFirstUser: true,
+		})
+
+		//nolint:gocritic // Test uses owner client session token for cookie-based relay auth.
+		sessionToken := firstClient.SessionToken()
+
+		// Configure the second client to authenticate via cookies		// only for WebSocket dials, matching browser behavior.
+		// For regular HTTP API calls we still need the header.
+		secondClient.SetSessionToken(sessionToken)
+		secondClient.SessionTokenProvider = cookieOnlySessionTokenProvider{
+			token:     sessionToken,
+			targetURL: secondClient.URL,
+		}
+
+		replicas, err := secondClient.Replicas(ctx)
+		require.NoError(t, err)
+		require.Len(t, replicas, 2)
+		firstReplicaID := replicaIDForClientURL(t, firstClient.URL, replicas)
+		secondReplicaID := replicaIDForClientURL(t, secondClient.URL, replicas)
+
+		streamingChunks := make(chan chattest.OpenAIChunk, 8)
+		chatStreamStarted := make(chan struct{}, 1)
+		openai := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
+			if req.Stream {
+				select {
+				case chatStreamStarted <- struct{}{}:
+				default:
+				}
+				return chattest.OpenAIResponse{StreamingChunks: streamingChunks}
+			}
+			return chattest.OpenAINonStreamingResponse("ok")
+		})
+
+		//nolint:gocritic // Test uses owner client to configure providers.
+		provider, err := firstClient.CreateChatProvider(ctx, codersdk.CreateChatProviderConfigRequest{
+			Provider:    "openai",
+			DisplayName: "OpenAI",
+			APIKey:      "test",
+			BaseURL:     openai,
+		})
+		require.NoError(t, err)
+
+		model, err := firstClient.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{
+			Provider:             provider.Provider,
+			Model:                "gpt-4",
+			DisplayName:          "GPT-4",
+			ContextLimit:         &[]int64{1000}[0],
+			CompressionThreshold: &[]int32{70}[0],
+		})
+		require.NoError(t, err)
+
+		chat, err := firstClient.CreateChat(ctx, codersdk.CreateChatRequest{
+			Content: []codersdk.ChatInputPart{{
+				Type: codersdk.ChatInputPartTypeText,
+				Text: "Test cookie-only relay",
+			}},
+			ModelConfigID: &model.ID,
+		})
+		require.NoError(t, err)
+		require.Equal(t, codersdk.ChatStatusPending, chat.Status)
+
+		var runningChat database.Chat
+		require.Eventually(t, func() bool {
+			current, getErr := db.GetChatByID(ctx, chat.ID)
+			if getErr != nil {
+				return false
+			}
+			if current.Status != database.ChatStatusRunning || !current.WorkerID.Valid {
+				return false
+			}
+			runningChat = current
+			return true
+		}, testutil.WaitLong, testutil.IntervalFast)
+
+		var localClient *codersdk.Client
+		var relayClient *codersdk.Client
+		switch runningChat.WorkerID.UUID {
+		case firstReplicaID:
+			localClient = firstClient
+			relayClient = secondClient
+		case secondReplicaID:
+			localClient = secondClient
+			relayClient = firstClient
+		default:
+			require.FailNowf(
+				t,
+				"worker replica was not recognized",
+				"worker %s was not one of %s or %s",
+				runningChat.WorkerID.UUID,
+				firstReplicaID,
+				secondReplicaID,
+			)
+		}
+
+		firstEvents, firstStream, err := localClient.StreamChat(ctx, chat.ID, nil)
+		require.NoError(t, err)
+		defer firstStream.Close()
+
+		select {
+		case <-chatStreamStarted:
+		case <-ctx.Done():
+			require.FailNowf(
+				t,
+				"timed out waiting for OpenAI stream request",
+				"chat stream did not start: %v",
+				ctx.Err(),
+			)
+		}
+
+		firstChunkText := "cookie-relay-part-one"
+		streamingChunks <- chattest.OpenAITextChunks(firstChunkText)[0]
+		firstEvent := waitForStreamTextPart(ctx, t, firstEvents, firstChunkText)
+		require.Equal(t, "assistant", firstEvent.MessagePart.Role)
+
+		// Subscribe from the non-worker replica with cookie-only
+		// auth. This triggers the relay dial. If the relay doesn't
+		// correctly forward cookies, this fails with 401.
+		secondEvents, secondStream, err := relayClient.StreamChat(ctx, chat.ID, nil)
+		require.NoError(t, err)
+		defer secondStream.Close()
+
+		secondSnapshotEvent := waitForStreamTextPart(ctx, t, secondEvents, firstChunkText)
+		require.Equal(t, "assistant", secondSnapshotEvent.MessagePart.Role)
+
+		secondChunkText := "cookie-relay-part-two"
+		streamingChunks <- chattest.OpenAITextChunks(secondChunkText)[0]
+		waitForStreamTextPart(ctx, t, firstEvents, secondChunkText)
+		waitForStreamTextPart(ctx, t, secondEvents, secondChunkText)
+
+		close(streamingChunks)
+	})
+
+	// This test verifies that cookie-only relay auth works when
+	// EnableHostPrefix is true. When the subscriber replica's
+	// HTTPCookies.Middleware normalizes __Host-coder_session_token
+	// to coder_session_token, the relay forwards the bare cookie.
+	// On the worker replica, the same middleware must not strip it.
+	//
+	// The fix ensures relayHeaders also extracts the token value
+	// and sets the Coder-Session-Token header so the worker
+	// replica can authenticate regardless of cookie prefix config.
+	t.Run("RelayCookieOnlyAuthWithHostPrefix", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		db, pubsub := dbtestutil.NewDB(t)
+		hostPrefixValues := coderdtest.DeploymentValues(t, func(dv *codersdk.DeploymentValues) {
+			dv.HTTPCookies.EnableHostPrefix = true
+			dv.HTTPCookies.Secure = true
+		})
+		firstClient, _ := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				Database:         db,
+				Pubsub:           pubsub,
+				DeploymentValues: hostPrefixValues,
+			},
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureHighAvailability: 1,
+				},
+			},
+		})
+
+		secondClient, _ := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				Database:         db,
+				Pubsub:           pubsub,
+				DeploymentValues: hostPrefixValues,
+			},
+			DontAddLicense:   true,
+			DontAddFirstUser: true,
+		})
+
+		//nolint:gocritic // Test uses owner client session token for cookie-based relay auth.
+		sessionToken := firstClient.SessionToken()
+
+		// Use cookie-only auth for WebSocket, as browsers do.		// With EnableHostPrefix, the browser would have
+		// __Host-coder_session_token but the middleware
+		// normalizes it. The relay copies the normalized cookie.
+		secondClient.SetSessionToken(sessionToken)
+		secondClient.SessionTokenProvider = cookieOnlySessionTokenProvider{
+			token:      sessionToken,
+			targetURL:  secondClient.URL,
+			hostPrefix: true,
+		}
+
+		replicas, err := secondClient.Replicas(ctx)
+		require.NoError(t, err)
+		require.Len(t, replicas, 2)
+		firstReplicaID := replicaIDForClientURL(t, firstClient.URL, replicas)
+		secondReplicaID := replicaIDForClientURL(t, secondClient.URL, replicas)
+
+		streamingChunks := make(chan chattest.OpenAIChunk, 8)
+		chatStreamStarted := make(chan struct{}, 1)
+		openai := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
+			if req.Stream {
+				select {
+				case chatStreamStarted <- struct{}{}:
+				default:
+				}
+				return chattest.OpenAIResponse{StreamingChunks: streamingChunks}
+			}
+			return chattest.OpenAINonStreamingResponse("ok")
+		})
+
+		//nolint:gocritic // Test uses owner client to configure providers.
+		provider, err := firstClient.CreateChatProvider(ctx, codersdk.CreateChatProviderConfigRequest{
+			Provider:    "openai",
+			DisplayName: "OpenAI",
+			APIKey:      "test",
+			BaseURL:     openai,
+		})
+		require.NoError(t, err)
+
+		model, err := firstClient.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{
+			Provider:             provider.Provider,
+			Model:                "gpt-4",
+			DisplayName:          "GPT-4",
+			ContextLimit:         &[]int64{1000}[0],
+			CompressionThreshold: &[]int32{70}[0],
+		})
+		require.NoError(t, err)
+
+		chat, err := firstClient.CreateChat(ctx, codersdk.CreateChatRequest{
+			Content: []codersdk.ChatInputPart{{
+				Type: codersdk.ChatInputPartTypeText,
+				Text: "Test host-prefix relay",
+			}},
+			ModelConfigID: &model.ID,
+		})
+		require.NoError(t, err)
+		require.Equal(t, codersdk.ChatStatusPending, chat.Status)
+
+		var runningChat database.Chat
+		require.Eventually(t, func() bool {
+			current, getErr := db.GetChatByID(ctx, chat.ID)
+			if getErr != nil {
+				return false
+			}
+			if current.Status != database.ChatStatusRunning || !current.WorkerID.Valid {
+				return false
+			}
+			runningChat = current
+			return true
+		}, testutil.WaitLong, testutil.IntervalFast)
+
+		var localClient *codersdk.Client
+		var relayClient *codersdk.Client
+		switch runningChat.WorkerID.UUID {
+		case firstReplicaID:
+			localClient = firstClient
+			relayClient = secondClient
+		case secondReplicaID:
+			localClient = secondClient
+			relayClient = firstClient
+		default:
+			require.FailNowf(
+				t,
+				"worker replica was not recognized",
+				"worker %s was not one of %s or %s",
+				runningChat.WorkerID.UUID,
+				firstReplicaID,
+				secondReplicaID,
+			)
+		}
+
+		firstEvents, firstStream, err := localClient.StreamChat(ctx, chat.ID, nil)
+		require.NoError(t, err)
+		defer firstStream.Close()
+
+		select {
+		case <-chatStreamStarted:
+		case <-ctx.Done():
+			require.FailNowf(
+				t,
+				"timed out waiting for OpenAI stream request",
+				"chat stream did not start: %v",
+				ctx.Err(),
+			)
+		}
+
+		firstChunkText := "hostprefix-relay-part-one"
+		streamingChunks <- chattest.OpenAITextChunks(firstChunkText)[0]
+		firstEvent := waitForStreamTextPart(ctx, t, firstEvents, firstChunkText)
+		require.Equal(t, "assistant", firstEvent.MessagePart.Role)
+
+		// This subscribe triggers the relay. With the bug, the
+		// worker replica's HTTPCookies.Middleware strips the bare
+		// coder_session_token cookie and there's no fallback
+		// Coder-Session-Token header, causing a 401.
+		secondEvents, secondStream, err := relayClient.StreamChat(ctx, chat.ID, nil)
+		require.NoError(t, err)
+		defer secondStream.Close()
+
+		secondSnapshotEvent := waitForStreamTextPart(ctx, t, secondEvents, firstChunkText)
+		require.Equal(t, "assistant", secondSnapshotEvent.MessagePart.Role)
+
+		secondChunkText := "hostprefix-relay-part-two"
+		streamingChunks <- chattest.OpenAITextChunks(secondChunkText)[0]
+		waitForStreamTextPart(ctx, t, firstEvents, secondChunkText)
+		waitForStreamTextPart(ctx, t, secondEvents, secondChunkText)
+
+		close(streamingChunks)
+	})
+
 	t.Run("RelaySnapshotIncludesBufferedParts", func(t *testing.T) {
 		t.Parallel()
 		ctx := testutil.Context(t, testutil.WaitLong)
@@ -716,3 +1054,40 @@ func findChatModelConfigByID(
 	require.FailNowf(t, "missing model config", "model config %s not found", id)
 	return codersdk.ChatModelConfig{}
 }
+
+// cookieOnlySessionTokenProvider authenticates HTTP requests via the
+// Coder-Session-Token header (for regular API calls) but
+// authenticates WebSocket dials via Cookie only, matching how
+// browsers behave (the native WebSocket constructor cannot set
+// custom headers).
+type cookieOnlySessionTokenProvider struct {
+	token     string
+	targetURL *url.URL
+	// hostPrefix, when true, sends the cookie with the
+	// __Host- prefix as browsers do with secure cookies.
+	hostPrefix bool
+}
+
+func (p cookieOnlySessionTokenProvider) AsRequestOption() codersdk.RequestOption {
+	return func(req *http.Request) {
+		req.Header.Set(codersdk.SessionTokenHeader, p.token)
+	}
+}
+
+func (p cookieOnlySessionTokenProvider) GetSessionToken() string {
+	return p.token
+}
+
+func (p cookieOnlySessionTokenProvider) SetDialOption(opts *websocket.DialOptions) {
+	// Browsers send cookies automatically on WebSocket upgrades
+	// but cannot send custom headers. Simulate this by setting
+	// only the Cookie header.
+	if opts.HTTPHeader == nil {
+		opts.HTTPHeader = make(http.Header)
+	}
+	cookieName := codersdk.SessionTokenCookie
+	if p.hostPrefix {
+		cookieName = "__Host-" + cookieName
+	}
+	opts.HTTPHeader.Set("Cookie", cookieName+"="+p.token)
+}