121 lines
3.6 KiB
Rust
121 lines
3.6 KiB
Rust
//! Real-time WebSocket event streaming.
|
|
|
|
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;
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct WsAuthQuery {
|
|
pub token: Option<String>,
|
|
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> {
|
|
if !authenticate_websocket(&state, &query, &headers).await {
|
|
return Err(ApiError::Unauthenticated(
|
|
"WebSocket authentication required. Supply a valid API token, session, or nx9_session cookie."
|
|
.to_string(),
|
|
));
|
|
}
|
|
|
|
Ok(ws.on_upgrade(move |socket| handle_socket(socket, state)))
|
|
}
|
|
|
|
async fn handle_socket(socket: WebSocket, state: AppState) {
|
|
let (mut sender, mut receiver) = socket.split();
|
|
let mut rx = state.event_tx.subscribe();
|
|
|
|
// 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;
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
// 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 {
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
|
|
// If either side terminates, stop the other task.
|
|
tokio::select! {
|
|
_ = (&mut send_task) => {
|
|
recv_task.abort();
|
|
}
|
|
_ = (&mut recv_task) => {
|
|
send_task.abort();
|
|
}
|
|
}
|
|
}
|