llm_manager/backend/
ws_server.rs

1use std::sync::Arc;
2
3use anyhow::{Context, Result, anyhow};
4use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
5use axum::http::StatusCode;
6use axum::response::IntoResponse;
7use axum::{Router, response::Html, routing::get};
8use futures_util::{SinkExt, StreamExt};
9use tokio::sync::broadcast;
10use tokio::task::JoinHandle;
11use tower_http::trace::TraceLayer;
12use tracing::{error, info, warn};
13
14use crate::models::WsMetrics;
15
16#[derive(Clone)]
17pub struct WsAppState {
18    pub metrics_rx: Arc<broadcast::Receiver<WsMetrics>>,
19    pub auth_key: Option<String>,
20}
21
22pub async fn start_ws_server(
23    port: u16,
24    metrics_rx: Arc<broadcast::Receiver<WsMetrics>>,
25    auth_key: Option<String>,
26    tls_config: Option<axum_server::tls_rustls::RustlsConfig>,
27    host: String,
28    mut shutdown_rx: tokio::sync::watch::Receiver<bool>,
29) -> Result<JoinHandle<()>> {
30    let state = WsAppState {
31        metrics_rx,
32        auth_key,
33    };
34
35    let app = Router::new()
36        .route("/dashboard", get(serve_dashboard))
37        .route("/ws", get(ws_handler))
38        .route("/health", get(|| async { "OK" }))
39        .layer(TraceLayer::new_for_http())
40        .with_state(state);
41
42    let addr = format!("{host}:{port}");
43
44    match tls_config {
45        Some(tls_cfg) => {
46            let socket_addr: std::net::SocketAddr = addr
47                .parse()
48                .map_err(|e| anyhow!("Invalid bind address {addr} for TLS: {e}"))?;
49            let tls_listener = axum_server::bind_rustls(socket_addr, tls_cfg);
50            let shutdown_fut = async move {
51                let _ = shutdown_rx.wait_for(|v| *v).await;
52            };
53            let handle = tokio::spawn(async move {
54                tokio::select! {
55                    result = tls_listener.serve(app.into_make_service()) => {
56                        if let Err(e) = result {
57                            error!("WebSocket server error: {e}");
58                        }
59                    }
60                    _ = shutdown_fut => {},
61                }
62            });
63            info!("WebSocket server listening on https://{addr}");
64            Ok(handle)
65        }
66        None => {
67            let listener = tokio::net::TcpListener::bind(&addr)
68                .await
69                .with_context(|| format!("Failed to bind WebSocket server to {addr}"))?;
70            let handle = tokio::spawn(async move {
71                let _ = axum::serve(listener, app)
72                    .with_graceful_shutdown(async move {
73                        let _ = shutdown_rx.wait_for(|v| *v).await;
74                    })
75                    .await;
76            });
77            info!("WebSocket server listening on http://{addr}");
78            Ok(handle)
79        }
80    }
81}
82
83pub fn stop_ws_server(handle: JoinHandle<()>) {
84    handle.abort();
85}
86
87async fn serve_dashboard(
88    axum::extract::State(state): axum::extract::State<WsAppState>,
89) -> Html<String> {
90    let auth_json = serde_json::to_string(&state.auth_key).unwrap_or_else(|_| "null".to_string());
91    let escaped = html_escape_attr(&auth_json);
92    let meta_tag = format!(r#"<meta name="ws-auth" content="{}">"#, escaped);
93    let html = include_str!("../dashboard.html");
94    Html(html.replacen("<body>", &format!("<body>{}", meta_tag), 1))
95}
96
97/// Escape a string for safe placement inside an HTML attribute value.
98fn html_escape_attr(s: &str) -> String {
99    let mut out = String::with_capacity(s.len());
100    for c in s.chars() {
101        match c {
102            '&' => out.push_str("&amp;"),
103            '"' => out.push_str("&quot;"),
104            '<' => out.push_str("&lt;"),
105            '>' => out.push_str("&gt;"),
106            c => out.push(c),
107        }
108    }
109    out
110}
111
112fn constant_time_not_eq(a: &[u8], b: &[u8]) -> bool {
113    if a.len() != b.len() {
114        return true;
115    }
116    let mut result: u8 = 0;
117    for (x, y) in a.iter().zip(b.iter()) {
118        result |= x ^ y;
119    }
120    result != 0
121}
122
123async fn ws_handler(
124    ws: WebSocketUpgrade,
125    axum::extract::State(state): axum::extract::State<WsAppState>,
126    axum::http::request::Parts { headers, .. }: axum::http::request::Parts,
127) -> impl IntoResponse {
128    if let Some(ref expected) = state.auth_key {
129        let provided = headers
130            .get(axum::http::header::SEC_WEBSOCKET_PROTOCOL)
131            .and_then(|v| v.to_str().ok())
132            .map(|s| s.to_string());
133        if let Some(ref provided_str) = provided {
134            if constant_time_not_eq(provided_str.as_bytes(), expected.as_bytes()) {
135                return StatusCode::UNAUTHORIZED.into_response();
136            }
137        } else {
138            return StatusCode::UNAUTHORIZED.into_response();
139        }
140    }
141    ws.on_upgrade(move |socket| handle_socket(socket, state))
142}
143
144async fn handle_socket(socket: WebSocket, state: WsAppState) {
145    let mut rx = state.metrics_rx.resubscribe();
146    info!("WebSocket client connected");
147
148    let (mut sender, mut receiver) = socket.split();
149
150    loop {
151        tokio::select! {
152            biased;
153            _ = receiver.next() => {
154                info!("WebSocket client disconnected");
155                break;
156            }
157            metrics = rx.recv() => match metrics {
158                Ok(m) => {
159                    let json = match serde_json::to_string(&m) {
160                        Ok(j) => j,
161                        Err(e) => {
162                            error!("Failed to serialize metrics: {e}");
163                            continue;
164                        }
165                    };
166                    if sender.send(Message::Text(json.into())).await.is_err() {
167                        info!("WebSocket client disconnected");
168                        break;
169                    }
170                }
171                Err(broadcast::error::RecvError::Lagged(n)) => {
172                    warn!("WebSocket client lagged behind, skipped {n} metrics");
173                }
174                Err(broadcast::error::RecvError::Closed) => {
175                    break;
176                }
177            },
178        }
179    }
180}
llm_manager/backend/ws_server.rs

llm_manager/backend/
ws_server.rs
