Finalize nx9-wg production release

crates/nx9-wg-api/src/routes/ws.rs

@@ -4,6 +4,8 @@ use crate::error::{ApiError, ApiResult};
 use crate::state::AppState;
 use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
 use axum::extract::{Query, State};
+use axum::http::HeaderMap;
+use axum::http::header::COOKIE;
 use axum::response::IntoResponse;
 use futures_util::{SinkExt, StreamExt};
 use serde::Deserialize;
@@ -14,24 +16,66 @@ pub struct WsAuthQuery {
     pub session: Option<String>,
 }
 
+/// Extract the NX9 session identifier from the browser session cookie.
+///
+/// The WebUI authenticates through the HttpOnly `nx9_session` cookie.
+/// WebSocket upgrades do not pass through the normal REST authentication
+/// middleware, so the cookie must be authenticated explicitly here.
+fn extract_session_cookie(headers: &HeaderMap) -> Option<&str> {
+    headers
+        .get(COOKIE)
+        .and_then(|value| value.to_str().ok())
+        .and_then(|cookies| {
+            cookies
+                .split(';')
+                .map(str::trim)
+                .find_map(|cookie| cookie.strip_prefix("nx9_session="))
+        })
+        .map(str::trim)
+        .filter(|session_id| !session_id.is_empty())
+}
+
+/// Authenticate a WebSocket request.
+///
+/// Authentication precedence:
+///
+/// 1. Explicit API token: `?token=...`
+/// 2. Explicit session: `?session=...`
+/// 3. Browser session cookie: `nx9_session=...`
+///
+/// The browser WebUI uses the HttpOnly session cookie, so no credential
+/// needs to be exposed in the WebSocket URL.
+async fn authenticate_websocket(
+    state: &AppState,
+    query: &WsAuthQuery,
+    headers: &HeaderMap,
+) -> bool {
+    if let Some(raw_token) = query.token.as_deref() {
+        return state.auth.authenticate_token(raw_token).await.is_ok();
+    }
+
+    if let Some(session_id) = query.session.as_deref() {
+        return state.auth.authenticate_session(session_id).await.is_ok();
+    }
+
+    if let Some(session_id) = extract_session_cookie(headers) {
+        return state.auth.authenticate_session(session_id).await.is_ok();
+    }
+
+    false
+}
+
 /// GET /api/v1/ws
 pub async fn ws_handler(
     ws: WebSocketUpgrade,
     State(state): State<AppState>,
     Query(query): Query<WsAuthQuery>,
+    headers: HeaderMap,
 ) -> ApiResult<impl IntoResponse> {
-    // Authenticate WebSocket connection via query parameters
-    let authenticated = if let Some(ref raw_token) = query.token {
-        state.auth.authenticate_token(raw_token).await.is_ok()
-    } else if let Some(ref session_id) = query.session {
-        state.auth.authenticate_session(session_id).await.is_ok()
-    } else {
-        false
-    };
-
-    if !authenticated {
+    if !authenticate_websocket(&state, &query, &headers).await {
         return Err(ApiError::Unauthenticated(
-            "WebSocket authentication required. Supply ?token=... or ?session=...".to_string(),
+            "WebSocket authentication required. Supply a valid API token, session, or nx9_session cookie."
+                .to_string(),
         ));
     }
 
@@ -42,11 +86,12 @@ async fn handle_socket(socket: WebSocket, state: AppState) {
     let (mut sender, mut receiver) = socket.split();
     let mut rx = state.event_tx.subscribe();
 
-    // Spawn background task to stream broadcast events to client
+    // Stream broadcast events to the connected WebSocket client.
     let mut send_task = tokio::spawn(async move {
         while let Ok(event) = rx.recv().await {
             if let Ok(json) = serde_json::to_string(&event) {
                 let msg = Message::Text(json.into());
+
                 if sender.send(msg).await.is_err() {
                     break;
                 }
@@ -54,7 +99,7 @@ async fn handle_socket(socket: WebSocket, state: AppState) {
         }
     });
 
-    // Client receive loop to handle close/ping/pong
+    // Receive loop handles client close frames and keeps the connection alive.
     let mut recv_task = tokio::spawn(async move {
         while let Some(Ok(msg)) = receiver.next().await {
             if let Message::Close(_) = msg {
@@ -63,9 +108,13 @@ async fn handle_socket(socket: WebSocket, state: AppState) {
         }
     });
 
-    // If either task exits, abort the other
+    // If either side terminates, stop the other task.
     tokio::select! {
-        _ = (&mut send_task) => recv_task.abort(),
-        _ = (&mut recv_task) => send_task.abort(),
+        _ = (&mut send_task) => {
+            recv_task.abort();
+        }
+        _ = (&mut recv_task) => {
+            send_task.abort();
+        }
     }
 }