fix(gateway): resolve dead session cookie infinite rejection loop for WebSocket auth
This commit is contained in:
+45
-33
@@ -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
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user