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}