From f5ce219844e0c6e80b1f1dfbbce1f96a6251448d Mon Sep 17 00:00:00 2001
From: kim <89579420+NyaaaWhatsUpDoc@users.noreply.github.com>
Date: Mon, 14 Apr 2025 13:12:09 +0000
Subject: [PATCH] [bugfix] websocket header token not always returned (#4009)

* always include headerToken response if provided, because Chrome *sigh*

* wording

* Update internal/api/client/streaming/stream.go

Co-authored-by: Ilia Pozdnyakov <iliazeus@proton.me>

---------

Co-authored-by: Ilia Pozdnyakov <iliazeus@proton.me>
---
 internal/api/client/streaming/stream.go | 45 ++++++++++++-------------
 1 file changed, 22 insertions(+), 23 deletions(-)

diff --git a/internal/api/client/streaming/stream.go b/internal/api/client/streaming/stream.go
index e6d1b80f7..8fafed7c2 100644
--- a/internal/api/client/streaming/stream.go
+++ b/internal/api/client/streaming/stream.go
@@ -18,6 +18,7 @@
 package streaming
 
 import (
+	"cmp"
 	"context"
 	"net/http"
 	"slices"
@@ -153,25 +154,22 @@ var pingMsg = []byte("ping!")
 //			description: bad request
 func (m *Module) StreamGETHandler(c *gin.Context) {
 	var (
-		token         string
-		tokenInHeader bool
-		account       *gtsmodel.Account
-		errWithCode   gtserror.WithCode
+		account     *gtsmodel.Account
+		errWithCode gtserror.WithCode
 	)
 
-	if t := c.Query(AccessTokenQueryKey); t != "" {
-		// Token was provided as
-		// query param, no problem.
-		token = t
-	} else if t := c.GetHeader(AccessTokenHeader); t != "" {
-		// Token was provided in "Sec-Websocket-Protocol" header.
-		//
-		// This is hacky and not technically correct but some
-		// clients do it since Mastodon allows it, so we must
-		// also allow it to avoid breaking expectations.
-		token = t
-		tokenInHeader = true
-	}
+	// Check both query parameter AND header "Sec-Websocket-Protocol"
+	// value for a token. The latter is hacky and not technically
+	// correct, but some client do it since Mastodon allows it, so
+	// we must allow it for compatibility.
+	//
+	// Chrome also *always* expects the "Sec-Websocket-Protocol"
+	// response value to match input, so it must always be checked.
+	queryToken := c.Query(AccessTokenQueryKey)
+	headerToken := c.GetHeader(AccessTokenHeader)
+
+	// Prefer query token else use header token.
+	token := cmp.Or(queryToken, headerToken)
 
 	if token != "" {
 
@@ -229,8 +227,7 @@ func (m *Module) StreamGETHandler(c *gin.Context) {
 		return
 	}
 
-	l := log.
-		WithContext(c.Request.Context()).
+	l := log.WithContext(c.Request.Context()).
 		WithField("streamID", id.NewULID()).
 		WithField("username", account.Username)
 
@@ -241,12 +238,14 @@ func (m *Module) StreamGETHandler(c *gin.Context) {
 	// If the upgrade fails, then Upgrade replies to the client
 	// with an HTTP error response.
 	var responseHeader http.Header
-	if tokenInHeader {
-		// Return the token in the response,
-		// else Chrome fails to connect.
+	if headerToken != "" {
+
+		// Chrome always requires the input header
+		// token to be provided in response, even
+		// if it also provided via query... *shrugs*
 		//
 		// https://developer.mozilla.org/en-US/docs/Web/HTTP/Protocol_upgrade_mechanism#sec-websocket-protocol
-		responseHeader = http.Header{AccessTokenHeader: {token}}
+		responseHeader = http.Header{AccessTokenHeader: {headerToken}}
 	}
 
 	wsConn, err := m.wsUpgrade.Upgrade(c.Writer, c.Request, responseHeader)