v0.8.14

src/api/websocket.rs

@@ -0,0 +1,40 @@
+#[derive(Debug, Deserialize)]
+struct WsQuery { token: Option<String> }
+async fn websocket(State(state): State<AppState>, Query(query): Query<WsQuery>, ws: WebSocketUpgrade) -> Result<Response, AppError> {
+    let expected = state.config.app_token.trim();
+    if !expected.is_empty() && query.token.as_deref() != Some(expected) { return Err(AppError::Unauthorized); }
+    Ok(ws.on_upgrade(move |socket| websocket_loop(state, socket)))
+}
+
+async fn websocket_loop(state: AppState, mut socket: WebSocket) {
+    let initial = match build_bootstrap(&state).await {
+        Ok(data) => json!({"event":"bootstrap","timestamp":Utc::now(),"data":data}),
+        Err(err) => json!({"event":"error","timestamp":Utc::now(),"data":{"message":err.to_string()}}),
+    };
+    if socket.send(Message::Text(initial.to_string())).await.is_err() { return; }
+    let mut receiver = state.events.subscribe();
+    loop {
+        tokio::select! {
+            event = receiver.recv() => {
+                match event {
+                    Ok(event) => {
+                        if let Ok(text) = serde_json::to_string(&event) {
+                            if socket.send(Message::Text(text)).await.is_err() { break; }
+                        }
+                    }
+                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
+                    Err(_) => break,
+                }
+            }
+            message = socket.next() => {
+                match message {
+                    Some(Ok(Message::Ping(value))) => { if socket.send(Message::Pong(value)).await.is_err() { break; } }
+                    Some(Ok(Message::Text(text))) if text == "ping" => { if socket.send(Message::Text("pong".into())).await.is_err() { break; } }
+                    Some(Ok(Message::Close(_))) | None | Some(Err(_)) => break,
+                    _ => {}
+                }
+            }
+        }
+    }
+}
+