fix(codersdk): use header auth for non-browser websocket dials (#22461) (cherry-pick/v2.31) (#22508)

Cherry-pick of #22461 to `release/2.31`.

Applies the non-browser websocket auth principle from #22226 to
remaining
`codersdk` websocket callsites, replacing cookie-jar session auth with
header-token auth. Fixes `401` failures on deployments with
`--host-prefix-cookie` enabled.

Closes #22461 (cherry-pick)

---------

Co-authored-by: ethan <ethanndickson@gmail.com>

codersdk/provisionerdaemons.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"net/http/cookiejar"
 	"slices"
 	"strings"
 	"time"
@@ -239,20 +238,14 @@ func (c *Client) provisionerJobLogsAfter(ctx context.Context, path string, after
 	if err != nil {
 		return nil, nil, err
 	}
-	jar, err := cookiejar.New(nil)
-	if err != nil {
-		return nil, nil, xerrors.Errorf("create cookie jar: %w", err)
-	}
-	jar.SetCookies(followURL, []*http.Cookie{{
-		Name:  SessionTokenCookie,
-		Value: c.SessionToken(),
-	}})
 	httpClient := &http.Client{
-		Jar:       jar,
 		Transport: c.HTTPClient.Transport,
 	}
 	conn, res, err := websocket.Dial(ctx, followURL.String(), &websocket.DialOptions{
-		HTTPClient:      httpClient,
+		HTTPClient: httpClient,
+		HTTPHeader: http.Header{
+			SessionTokenHeader: []string{c.SessionToken()},
+		},
 		CompressionMode: websocket.CompressionDisabled,
 	})
 	if err != nil {
@@ -325,16 +318,8 @@ func (c *Client) ServeProvisionerDaemon(ctx context.Context, req ServeProvisione
 		headers.Set(ProvisionerDaemonPSK, req.PreSharedKey)
 	}
 	if req.ProvisionerKey == "" && req.PreSharedKey == "" {
-		// use session token if we don't have a PSK or provisioner key.
-		jar, err := cookiejar.New(nil)
-		if err != nil {
-			return nil, xerrors.Errorf("create cookie jar: %w", err)
-		}
-		jar.SetCookies(serverURL, []*http.Cookie{{
-			Name:  SessionTokenCookie,
-			Value: c.SessionToken(),
-		}})
-		httpClient.Jar = jar
+		// Use session token if we don't have a PSK or provisioner key.
+		headers.Set(SessionTokenHeader, c.SessionToken())
 	}
 
 	conn, res, err := websocket.Dial(ctx, serverURL.String(), &websocket.DialOptions{

codersdk/workspaceagents.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"net/http/cookiejar"
 	"strings"
 	"time"
 
@@ -580,24 +579,16 @@ func (c *Client) WatchWorkspaceAgentContainers(ctx context.Context, agentID uuid
 		return nil, nil, err
 	}
 
-	jar, err := cookiejar.New(nil)
-	if err != nil {
-		return nil, nil, xerrors.Errorf("create cookie jar: %w", err)
-	}
-
-	jar.SetCookies(reqURL, []*http.Cookie{{
-		Name:  SessionTokenCookie,
-		Value: c.SessionToken(),
-	}})
-
 	conn, res, err := websocket.Dial(ctx, reqURL.String(), &websocket.DialOptions{
 		// We want `NoContextTakeover` compression to balance improving
 		// bandwidth cost/latency with minimal memory usage overhead.
 		CompressionMode: websocket.CompressionNoContextTakeover,
 		HTTPClient: &http.Client{
-			Jar:       jar,
 			Transport: c.HTTPClient.Transport,
 		},
+		HTTPHeader: http.Header{
+			SessionTokenHeader: []string{c.SessionToken()},
+		},
 	})
 	if err != nil {
 		if res == nil {
@@ -687,20 +678,14 @@ func (c *Client) WorkspaceAgentLogsAfter(ctx context.Context, agentID uuid.UUID,
 		return ch, closeFunc(func() error { return nil }), nil
 	}
 
-	jar, err := cookiejar.New(nil)
-	if err != nil {
-		return nil, nil, xerrors.Errorf("create cookie jar: %w", err)
-	}
-	jar.SetCookies(reqURL, []*http.Cookie{{
-		Name:  SessionTokenCookie,
-		Value: c.SessionToken(),
-	}})
 	httpClient := &http.Client{
-		Jar:       jar,
 		Transport: c.HTTPClient.Transport,
 	}
 	conn, res, err := websocket.Dial(ctx, reqURL.String(), &websocket.DialOptions{
-		HTTPClient:      httpClient,
+		HTTPClient: httpClient,
+		HTTPHeader: http.Header{
+			SessionTokenHeader: []string{c.SessionToken()},
+		},
 		CompressionMode: websocket.CompressionDisabled,
 	})
 	if err != nil {

codersdk/workspacesdk/workspacesdk.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"net"
 	"net/http"
-	"net/http/cookiejar"
 	"net/netip"
 	"os"
 	"strconv"
@@ -363,26 +362,23 @@ func (c *Client) AgentReconnectingPTY(ctx context.Context, opts WorkspaceAgentRe
 	}
 	serverURL.RawQuery = q.Encode()
 
-	// If we're not using a signed token, we need to set the session token as a
-	// cookie.
-	httpClient := c.client.HTTPClient
+	// Shallow-clone the HTTP client so we never inherit a caller-provided
+	// cookie jar. Non-browser websocket auth uses the Coder-Session-Token
+	// header or a signed-token query param — never cookies. A stale jar
+	// cookie would take precedence on the server (cookies are checked
+	// before headers) and cause spurious 401s.
+	wsHTTPClient := *c.client.HTTPClient
+	wsHTTPClient.Jar = nil
+
+	headers := http.Header{}
+	// If we're not using a signed token, set the session token header.
 	if opts.SignedToken == "" {
-		jar, err := cookiejar.New(nil)
-		if err != nil {
-			return nil, xerrors.Errorf("create cookie jar: %w", err)
-		}
-		jar.SetCookies(serverURL, []*http.Cookie{{
-			Name:  codersdk.SessionTokenCookie,
-			Value: c.client.SessionToken(),
-		}})
-		httpClient = &http.Client{
-			Jar:       jar,
-			Transport: c.client.HTTPClient.Transport,
-		}
+		headers.Set(codersdk.SessionTokenHeader, c.client.SessionToken())
 	}
 	//nolint:bodyclose
 	conn, res, err := websocket.Dial(ctx, serverURL.String(), &websocket.DialOptions{
-		HTTPClient: httpClient,
+		HTTPClient: &wsHTTPClient,
+		HTTPHeader: headers,
 	})
 	if err != nil {
 		if res == nil {