fix: WebSocket cookie auth + connect on layout mount

This commit is contained in:
2026-06-29 14:30:29 -04:00
parent 90289fede1
commit 369737565e
4 changed files with 36 additions and 21 deletions
+1 -1
View File
@@ -113,7 +113,7 @@ func main() {
// WebSocket endpoint // WebSocket endpoint
r.Get("/ws", func(w http.ResponseWriter, r *http.Request) { 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 // API routes
+16 -7
View File
@@ -148,7 +148,7 @@ type authMessage struct {
// connection first, then waits for an auth message frame containing the // connection first, then waits for an auth message frame containing the
// session token. This avoids leaking the token in the URL (which would // session token. This avoids leaking the token in the URL (which would
// appear in server logs, browser history, and Referer headers). // 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) conn, err := upgrader.Upgrade(w, r, nil)
if err != nil { if err != nil {
logger.Error("websocket upgrade failed", "error", err) 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 // Short deadline for the auth frame; we don't want unauthenticated
// connections hanging around consuming resources. // connections hanging around consuming resources.
// 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
}
// Fall back to message-frame auth (TUI, bots).
if token == "" {
conn.SetReadDeadline(time.Now().Add(10 * time.Second)) conn.SetReadDeadline(time.Now().Add(10 * time.Second))
_, raw, err := conn.ReadMessage() _, raw, readErr := conn.ReadMessage()
if err != nil { if readErr != nil {
logger.Warn("ws auth: failed to read auth message", "error", err) logger.Warn("ws auth: failed to read auth message", "error", readErr)
conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"auth timeout"}`)) conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"auth timeout"}`))
conn.Close() conn.Close()
return return
} }
var auth authMessage var auth authMessage
if err := json.Unmarshal(raw, &auth); err != nil || auth.Token == "" { if jsonErr := json.Unmarshal(raw, &auth); jsonErr != nil || auth.Token == "" {
logger.Warn("ws auth: invalid auth message") logger.Warn("ws auth: invalid auth message")
conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"missing token"}`)) conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"missing token"}`))
conn.Close() conn.Close()
return return
} }
token = auth.Token
}
var userID string var userID string
err = db.QueryRowContext(context.Background(), err = db.QueryRowContext(context.Background(),
`SELECT user_id FROM sessions WHERE token = $1 AND expires_at > NOW()`, `SELECT user_id FROM sessions WHERE token = $1 AND expires_at > NOW()`,
auth.Token, token,
).Scan(&userID) ).Scan(&userID)
if err != nil { if err != nil {
logger.Warn("ws auth: invalid session", "error", err) logger.Warn("ws auth: invalid session", "error", err)
+6 -1
View File
@@ -1,6 +1,7 @@
import { useEffect } from 'react'; import { useEffect } from 'react';
import { Link, Outlet, useLocation, useNavigate } from 'react-router-dom'; import { Link, Outlet, useLocation, useNavigate } from 'react-router-dom';
import { useAuthStore } from '../stores/auth.ts'; import { useAuthStore } from '../stores/auth.ts';
import { useWebSocketStore } from '../stores/ws.ts';
import { ServerBar } from './ServerBar.tsx'; import { ServerBar } from './ServerBar.tsx';
import { ChannelList } from './ChannelList.tsx'; import { ChannelList } from './ChannelList.tsx';
import { MemberList } from './MemberList.tsx'; import { MemberList } from './MemberList.tsx';
@@ -12,12 +13,16 @@ export function Layout() {
const user = useAuthStore((state) => state.user); const user = useAuthStore((state) => state.user);
const fetchMe = useAuthStore((state) => state.fetchMe); const fetchMe = useAuthStore((state) => state.fetchMe);
const logout = useAuthStore((state) => state.logout); const logout = useAuthStore((state) => state.logout);
const wsConnect = useWebSocketStore((s) => s.connect);
const wsDisconnect = useWebSocketStore((s) => s.disconnect);
const navigate = useNavigate(); const navigate = useNavigate();
const location = useLocation(); const location = useLocation();
useEffect(() => { useEffect(() => {
fetchMe(); fetchMe();
}, [fetchMe]); wsConnect();
return () => { wsDisconnect(); };
}, [fetchMe, wsConnect, wsDisconnect]);
useEffect(() => { useEffect(() => {
if (!isLoading && !isAuthenticated && location.pathname !== '/login') { if (!isLoading && !isAuthenticated && location.pathname !== '/login') {
+5 -4
View File
@@ -19,7 +19,7 @@ export interface WsEvent {
interface WebSocketState { interface WebSocketState {
socket: WebSocket | null; socket: WebSocket | null;
connected: boolean; connected: boolean;
connect: (token: string) => void; connect: () => void;
disconnect: () => void; disconnect: () => void;
send: (event: WsEvent) => void; send: (event: WsEvent) => void;
} }
@@ -64,7 +64,7 @@ export const useWebSocketStore = create<WebSocketState>((set, get) => ({
socket: null, socket: null,
connected: false, connected: false,
connect: (token: string) => { connect: () => {
const existing = get().socket; const existing = get().socket;
if (existing && (existing.readyState === WebSocket.OPEN || existing.readyState === WebSocket.CONNECTING)) { if (existing && (existing.readyState === WebSocket.OPEN || existing.readyState === WebSocket.CONNECTING)) {
return; return;
@@ -82,14 +82,15 @@ export const useWebSocketStore = create<WebSocketState>((set, get) => ({
} }
reconnectTimeout = window.setTimeout(() => { reconnectTimeout = window.setTimeout(() => {
if (!get().connected) { if (!get().connected) {
get().connect(token); get().connect();
} }
}, reconnectDelay); }, reconnectDelay);
reconnectDelay = Math.min(reconnectDelay * 2, maxReconnectDelay); reconnectDelay = Math.min(reconnectDelay * 2, maxReconnectDelay);
}; };
socket.onopen = () => { 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) => { socket.onmessage = (event) => {