diff --git a/cmd/server/main.go b/cmd/server/main.go index 33f948f..70f7c5f 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -113,7 +113,7 @@ func main() { // WebSocket endpoint r.Get("/ws", func(w http.ResponseWriter, r *http.Request) { - gateway.ServeWS(database.DB, hub, logger, w, r) + gateway.ServeWS(database.DB, hub, logger, w, r, cfg.Session.CookieName) }) // API routes diff --git a/internal/gateway/client.go b/internal/gateway/client.go index a595c3f..27392fb 100644 --- a/internal/gateway/client.go +++ b/internal/gateway/client.go @@ -148,7 +148,7 @@ type authMessage struct { // connection first, then waits for an auth message frame containing the // session token. This avoids leaking the token in the URL (which would // appear in server logs, browser history, and Referer headers). -func ServeWS(db *sql.DB, hub *Hub, logger *slog.Logger, w http.ResponseWriter, r *http.Request) { +func ServeWS(db *sql.DB, hub *Hub, logger *slog.Logger, w http.ResponseWriter, r *http.Request, cookieName string) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { logger.Error("websocket upgrade failed", "error", err) @@ -157,27 +157,36 @@ func ServeWS(db *sql.DB, hub *Hub, logger *slog.Logger, w http.ResponseWriter, r // Short deadline for the auth frame; we don't want unauthenticated // connections hanging around consuming resources. - conn.SetReadDeadline(time.Now().Add(10 * time.Second)) - _, raw, err := conn.ReadMessage() - if err != nil { - logger.Warn("ws auth: failed to read auth message", "error", err) - conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"auth timeout"}`)) - conn.Close() - return + // Try cookie-based auth first (browser clients send cookies automatically). + var token string + if cookie, cookieErr := r.Cookie(cookieName); cookieErr == nil && cookie.Value != "" { + token = cookie.Value } - var auth authMessage - if err := json.Unmarshal(raw, &auth); err != nil || auth.Token == "" { - logger.Warn("ws auth: invalid auth message") - conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"missing token"}`)) - conn.Close() - return + // 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 string err = db.QueryRowContext(context.Background(), `SELECT user_id FROM sessions WHERE token = $1 AND expires_at > NOW()`, - auth.Token, + token, ).Scan(&userID) if err != nil { logger.Warn("ws auth: invalid session", "error", err) diff --git a/web/src/components/Layout.tsx b/web/src/components/Layout.tsx index edcc7bf..173dbdb 100644 --- a/web/src/components/Layout.tsx +++ b/web/src/components/Layout.tsx @@ -1,6 +1,7 @@ import { useEffect } from 'react'; import { Link, Outlet, useLocation, useNavigate } from 'react-router-dom'; import { useAuthStore } from '../stores/auth.ts'; +import { useWebSocketStore } from '../stores/ws.ts'; import { ServerBar } from './ServerBar.tsx'; import { ChannelList } from './ChannelList.tsx'; import { MemberList } from './MemberList.tsx'; @@ -12,12 +13,16 @@ export function Layout() { const user = useAuthStore((state) => state.user); const fetchMe = useAuthStore((state) => state.fetchMe); const logout = useAuthStore((state) => state.logout); + const wsConnect = useWebSocketStore((s) => s.connect); + const wsDisconnect = useWebSocketStore((s) => s.disconnect); const navigate = useNavigate(); const location = useLocation(); useEffect(() => { fetchMe(); - }, [fetchMe]); + wsConnect(); + return () => { wsDisconnect(); }; + }, [fetchMe, wsConnect, wsDisconnect]); useEffect(() => { if (!isLoading && !isAuthenticated && location.pathname !== '/login') { diff --git a/web/src/stores/ws.ts b/web/src/stores/ws.ts index adb27b0..ce9c7a3 100644 --- a/web/src/stores/ws.ts +++ b/web/src/stores/ws.ts @@ -19,7 +19,7 @@ export interface WsEvent { interface WebSocketState { socket: WebSocket | null; connected: boolean; - connect: (token: string) => void; + connect: () => void; disconnect: () => void; send: (event: WsEvent) => void; } @@ -64,7 +64,7 @@ export const useWebSocketStore = create((set, get) => ({ socket: null, connected: false, - connect: (token: string) => { + connect: () => { const existing = get().socket; if (existing && (existing.readyState === WebSocket.OPEN || existing.readyState === WebSocket.CONNECTING)) { return; @@ -82,14 +82,15 @@ export const useWebSocketStore = create((set, get) => ({ } reconnectTimeout = window.setTimeout(() => { if (!get().connected) { - get().connect(token); + get().connect(); } }, reconnectDelay); reconnectDelay = Math.min(reconnectDelay * 2, maxReconnectDelay); }; socket.onopen = () => { - socket.send(JSON.stringify({ token })); + // Cookie-based auth: browser sends session cookie automatically. + // No need to send a token frame. }; socket.onmessage = (event) => {