103 lines
3.4 KiB
Rust
103 lines
3.4 KiB
Rust
|
//! WebSocket interface — real-time bidirectional communication.
|
||
|
|
//!
|
||
|
|
//! Enables web clients and other WS-capable consumers to connect
|
||
|
|
//! and participate in sessions. Built on Axum's WebSocket support.
|
||
|
|
|
||
|
|
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
|
||
|
|
use axum::response::IntoResponse;
|
||
|
|
use axum::routing::get;
|
||
|
|
use axum::Router;
|
||
|
|
use futures_util::stream::StreamExt;
|
||
|
|
use futures_util::SinkExt;
|
||
|
|
use std::sync::Arc;
|
||
|
|
use tracing::info;
|
||
|
|
|
||
|
|
/// Shared application state for the WS server.
|
||
|
|
pub struct WsState {
|
||
|
|
pub store_base_dir: std::path::PathBuf,
|
||
|
|
pub session_id: Option<String>,
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Build the WebSocket router.
|
||
|
|
pub fn build_router(state: Arc<WsState>) -> Router {
|
||
|
|
Router::new()
|
||
|
|
.route("/ws", get(ws_handler))
|
||
|
|
.with_state(state)
|
||
|
|
}
|
||
|
|
|
||
|
|
/// WebSocket upgrade handler.
|
||
|
|
async fn ws_handler(
|
||
|
|
ws: WebSocketUpgrade,
|
||
|
|
axum::extract::State(state): axum::extract::State<Arc<WsState>>,
|
||
|
|
) -> impl IntoResponse {
|
||
|
|
ws.on_upgrade(move |socket| handle_socket(socket, state))
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Handle an established WebSocket connection.
|
||
|
|
async fn handle_socket(mut socket: WebSocket, state: Arc<WsState>) {
|
||
|
|
// Channel for sending text messages to the WebSocket send task.
|
||
|
|
// The receiver side runs in a spawned task that forwards each
|
||
|
|
// string as a `Message::Text` to the client.
|
||
|
|
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
||
|
|
|
||
|
|
info!("WebSocket client connected");
|
||
|
|
|
||
|
|
// Send a welcome message
|
||
|
|
let welcome = serde_json::json!({
|
||
|
|
"type": "connected",
|
||
|
|
"session": state.session_id,
|
||
|
|
"message": "Connected to Zesdex WebSocket server"
|
||
|
|
});
|
||
|
|
// axum 0.8 Message::Text wraps Utf8Bytes; convert via .into()
|
||
|
|
let _ = socket.send(Message::Text(welcome.to_string().into())).await;
|
||
|
|
|
||
|
|
// Split the socket into sender and receiver halves
|
||
|
|
let (mut sender, mut receiver) = socket.split();
|
||
|
|
|
||
|
|
// Spawn task to forward messages from channel to WebSocket sender
|
||
|
|
let send_task = tokio::spawn(async move {
|
||
|
|
while let Some(msg) = rx.recv().await {
|
||
|
|
if sender.send(Message::Text(msg.into())).await.is_err() {
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
// Receive messages from the client
|
||
|
|
// receiver is SplitStream<WebSocket> — use StreamExt::next()
|
||
|
|
while let Some(Ok(msg)) = receiver.next().await {
|
||
|
|
match msg {
|
||
|
|
Message::Text(text) => {
|
||
|
|
// Convert Utf8Bytes -> String for JSON serialisation
|
||
|
|
let text_str = text.to_string();
|
||
|
|
info!("Received WS message: {text_str}");
|
||
|
|
// Echo back for now
|
||
|
|
let response = serde_json::json!({
|
||
|
|
"type": "echo",
|
||
|
|
"data": text_str
|
||
|
|
});
|
||
|
|
let _ = tx.send(response.to_string());
|
||
|
|
}
|
||
|
|
Message::Close(_) => break,
|
||
|
|
_ => {}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
send_task.abort();
|
||
|
|
info!("WebSocket client disconnected");
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Run the WebSocket server standalone.
|
||
|
|
pub async fn run_server(port: u16) -> anyhow::Result<()> {
|
||
|
|
let state = Arc::new(WsState {
|
||
|
|
store_base_dir: std::path::PathBuf::from("."),
|
||
|
|
session_id: None,
|
||
|
|
});
|
||
|
|
let app = build_router(state);
|
||
|
|
let addr = std::net::SocketAddr::from(([0, 0, 0, 0], port));
|
||
|
|
info!("WebSocket server listening on ws://{addr}");
|
||
|
|
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||
|
|
axum::serve(listener, app).await?;
|
||
|
|
Ok(())
|
||
|
|
}
|