Skip to main content

karbon_framework/http/
ws.rs

1use axum::{
2    extract::ws::{WebSocket, WebSocketUpgrade},
3    response::Response,
4};
5use std::future::Future;
6
7/// Helper to create a WebSocket handler from a simple async function.
8///
9/// ```ignore
10/// use framework::ws::websocket_handler;
11///
12/// async fn handle_socket(mut socket: WebSocket) {
13///     while let Some(Ok(msg)) = socket.recv().await {
14///         if let Message::Text(text) = msg {
15///             socket.send(Message::Text(format!("Echo: {}", text))).await.ok();
16///         }
17///     }
18/// }
19///
20/// // In your router:
21/// Router::new()
22///     .route("/ws", get(|ws: WebSocketUpgrade| websocket_handler(ws, handle_socket)))
23/// ```
24pub fn websocket_handler<F, Fut>(ws: WebSocketUpgrade, handler: F) -> Response
25where
26    F: FnOnce(WebSocket) -> Fut + Send + 'static,
27    Fut: Future<Output = ()> + Send + 'static,
28{
29    ws.on_upgrade(handler)
30}
31
32/// Helper to create a WebSocket handler with shared state.
33///
34/// ```ignore
35/// async fn chat_socket(mut socket: WebSocket, state: AppState) {
36///     // access state.db, state.config, etc.
37/// }
38///
39/// Router::new()
40///     .route("/ws/chat", get(|ws: WebSocketUpgrade, State(state): State<AppState>| {
41///         websocket_handler_with_state(ws, state, chat_socket)
42///     }))
43/// ```
44pub fn websocket_handler_with_state<S, F, Fut>(ws: WebSocketUpgrade, state: S, handler: F) -> Response
45where
46    S: Send + 'static,
47    F: FnOnce(WebSocket, S) -> Fut + Send + 'static,
48    Fut: Future<Output = ()> + Send + 'static,
49{
50    ws.on_upgrade(move |socket| handler(socket, state))
51}