diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs index 7cd21de..36fc4cc 100644 --- a/src/gateway/mod.rs +++ b/src/gateway/mod.rs @@ -608,6 +608,7 @@ fn build_router(state: Arc) -> Router { routing::get(http::download_attachment), ) .route("/ws", routing::get(ws::ws_handler)) + .route("/ws/logs", routing::get(ws::ws_logs_handler)) .route_layer(middleware::from_fn_with_state( state.auth.clone(), auth::require_auth, diff --git a/src/gateway/ws.rs b/src/gateway/ws.rs index 85eb6ea..4c0208f 100644 --- a/src/gateway/ws.rs +++ b/src/gateway/ws.rs @@ -151,6 +151,102 @@ impl Drop for ConnectionGuard { } } +#[derive(Debug, Default, Deserialize)] +pub struct WsLogsQuery { + level: Option, + search: Option, +} + +pub async fn ws_logs_handler( + ws: WebSocketUpgrade, + Query(query): Query, + Extension(_identity): Extension, +) -> Response { + ws.on_upgrade(|socket| async move { + handle_logs_socket(socket, query).await; + }) +} + +async fn handle_logs_socket(ws: WebSocket, query: WsLogsQuery) { + let Some(tx) = crate::logging::log_sender() else { + return; + }; + let mut rx = tx.subscribe(); + let (mut ws_sender, mut ws_receiver) = ws.split(); + + let min_level = query + .level + .as_deref() + .map(parse_min_level) + .unwrap_or(0); + let search = query + .search + .filter(|s| !s.is_empty()) + .map(|s| s.to_ascii_lowercase()); + + loop { + tokio::select! { + result = rx.recv() => { + match result { + Ok(event) => { + if level_rank(&event.level) < min_level { + continue; + } + if let Some(needle) = &search + && !event.message.to_ascii_lowercase().contains(needle) + && !event.target.to_ascii_lowercase().contains(needle) + { + continue; + } + let json = serde_json::json!({ + "ts": event.ts, + "level": event.level, + "target": event.target, + "message": event.message, + }); + if ws_sender + .send(WsMessage::Text(json.to_string().into())) + .await + .is_err() + { + break; + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue, + Err(tokio::sync::broadcast::error::RecvError::Closed) => break, + } + } + msg = ws_receiver.next() => { + match msg { + Some(Ok(WsMessage::Close(_))) | None => break, + _ => {} + } + } + } + } +} + +fn parse_min_level(level: &str) -> u8 { + match level.to_ascii_uppercase().as_str() { + "DEBUG" => 1, + "INFO" => 2, + "WARN" => 3, + "ERROR" => 4, + _ => 0, + } +} + +fn level_rank(level: &str) -> u8 { + match level { + "TRACE" => 0, + "DEBUG" => 1, + "INFO" => 2, + "WARN" => 3, + "ERROR" => 4, + _ => 2, + } +} + #[cfg(test)] mod tests { use super::*;