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
+24 -15
View File
@@ -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)