From 410b7a4d6b579c43b3505f21e4936ff2570929d5 Mon Sep 17 00:00:00 2001
From: hobokenchicken <dustin@dustin.coffee>
Date: Mon, 27 Jul 2026 12:01:48 -0400
Subject: [PATCH] fix(gateway): resolve dead session cookie infinite rejection
 loop for WebSocket auth

---
 internal/gateway/client.go | 78 ++++++++++++++++++++++----------------
 web/src/stores/ws.ts       | 11 +++++-
 2 files changed, 54 insertions(+), 35 deletions(-)

diff --git a/internal/gateway/client.go b/internal/gateway/client.go
index 375cafd..e742363 100644
--- a/internal/gateway/client.go
+++ b/internal/gateway/client.go
@@ -263,40 +263,52 @@ func ServeWS(db *sql.DB, hub *Hub, logger *slog.Logger, w http.ResponseWriter, r
 		return
 	}
 
-	// Try cookie-based auth first (browser clients).
-	var token string
-	if cookie, cookieErr := r.Cookie(cookieName); cookieErr == nil && cookie.Value != "" {
-		token = cookie.Value
-	}
-
-	// Fall back to message-frame auth (TUI, bots).
-	if token == "" {
-		conn.SetReadDeadline(time.Now().Add(10 * time.Second))
-		_, raw, readErr := conn.ReadMessage()
-		if readErr != nil {
-			logger.Warn("ws auth: failed to read auth message", "error", readErr)
-			conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"auth timeout"}`))
-			conn.Close()
-			return
-		}
-		var auth authMessage
-		if jsonErr := json.Unmarshal(raw, &auth); jsonErr != nil || auth.Token == "" {
-			logger.Warn("ws auth: invalid auth message")
-			conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"missing token"}`))
-			conn.Close()
-			return
-		}
-		token = auth.Token
-	}
-
 	var userID, username string
-	err = db.QueryRowContext(context.Background(),
-		`SELECT u.id, u.username FROM sessions s JOIN users u ON u.id = s.user_id WHERE s.token = $1 AND s.expires_at > NOW()`,
-		token,
-	).Scan(&userID, &username)
-	if err != nil {
-		logger.Warn("ws auth: invalid session", "error", err)
-		conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"invalid session"}`))
+
+	// Collect candidate tokens from Query param, Authorization header, and Cookies.
+	var candidateTokens []string
+	if qToken := r.URL.Query().Get("token"); qToken != "" {
+		candidateTokens = append(candidateTokens, qToken)
+	}
+	authHeader := r.Header.Get("Authorization")
+	if len(authHeader) > 7 && authHeader[:7] == "Bearer " {
+		candidateTokens = append(candidateTokens, authHeader[7:])
+	}
+	for _, c := range r.Cookies() {
+		if c.Name == cookieName && c.Value != "" {
+			candidateTokens = append(candidateTokens, c.Value)
+		}
+	}
+
+	// Test candidates against Postgres sessions
+	for _, token := range candidateTokens {
+		err := db.QueryRowContext(context.Background(),
+			`SELECT u.id, u.username FROM sessions s JOIN users u ON u.id = s.user_id WHERE s.token = $1 AND s.expires_at > NOW()`,
+			token,
+		).Scan(&userID, &username)
+		if err == nil && userID != "" {
+			break
+		}
+	}
+
+	// Fall back to message-frame auth if no candidate token authenticated (wait up to 5s)
+	if userID == "" {
+		conn.SetReadDeadline(time.Now().Add(5 * time.Second))
+		_, raw, readErr := conn.ReadMessage()
+		if readErr == nil {
+			var auth authMessage
+			if jsonErr := json.Unmarshal(raw, &auth); jsonErr == nil && auth.Token != "" {
+				_ = db.QueryRowContext(context.Background(),
+					`SELECT u.id, u.username FROM sessions s JOIN users u ON u.id = s.user_id WHERE s.token = $1 AND s.expires_at > NOW()`,
+					auth.Token,
+				).Scan(&userID, &username)
+			}
+		}
+	}
+
+	if userID == "" {
+		logger.Warn("ws auth: no valid session found")
+		conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"error","error":"invalid_session"}`))
 		conn.Close()
 		return
 	}
diff --git a/web/src/stores/ws.ts b/web/src/stores/ws.ts
index 9202e72..64cce98 100644
--- a/web/src/stores/ws.ts
+++ b/web/src/stores/ws.ts
@@ -95,8 +95,14 @@ export const useWebSocketStore = create<WebSocketState>((set, get) => ({
     };
 
     socket.onopen = () => {
-      // Cookie-based auth: browser sends session cookie automatically.
-      // No need to send a token frame.
+      const token = localStorage.getItem('dumpster_session_token');
+      if (token) {
+        try {
+          socket.send(JSON.stringify({ token }));
+        } catch {
+          // ignore
+        }
+      }
     };
 
     socket.onmessage = (event) => {
@@ -119,6 +125,7 @@ export const useWebSocketStore = create<WebSocketState>((set, get) => ({
 
       if (data.type === 'error') {
         console.error('ws error:', data);
+        useAuthStore.getState().fetchMe();
         socket.close();
         return;
       }