use futures_util::{SinkExt, StreamExt};
use poem::{
web::websocket::{Message, WebSocket},
web::FromRequest,
Endpoint,
IntoResponse,
Request,
Response,
Result,
};
pub async fn ws_middleware<E>(next: E, req: Request) -> Result<Response>
where
E: Endpoint,
{
let path = req.uri().path();
let path = path.to_string();
println!("request: {}", path);
let resp: Response = if path.ends_with("ws") {
let (req, mut body) = req.split();
let ws = WebSocket::from_request(&req, &mut body).await.unwrap();
let res = hello_ws(ws, path).await;
println!("is ws");
res.into_response()
} else {
let res = next.call(req).await.unwrap();
println!("not ws");
res.into_response()
};
Ok(resp)
}
pub async fn hello_ws(ws: WebSocket, path: String) -> impl IntoResponse {
ws.on_upgrade(|socket| async move {
let (mut tx, rx) = socket.split();
let msg = Message::Text(format!("Hello, WebSocket {}", &path).into());
let _ = tx.send(msg).await;
let _ = rx
.for_each(|msg| async move {
println!("Received message: {:?}", msg);
})
.await;
})
}