use axum::Router;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::routing::get;
use futures_util::{SinkExt, StreamExt};
use tracing::debug;
use crate::state::AppState;
use super::config::WebSocketSurface;
pub fn router(_state: &AppState, surface: &WebSocketSurface) -> Router<AppState> {
let path = surface.path.clone();
Router::new().route(
&path,
get(move |ws: WebSocketUpgrade| async move { ws.on_upgrade(handle_socket) }),
)
}
async fn handle_socket(mut socket: WebSocket) {
while let Some(msg) = socket.next().await {
match msg {
Ok(Message::Text(text)) => {
if SinkExt::send(&mut socket, Message::Text(text)).await.is_err() {
break;
}
}
Ok(Message::Binary(data)) => {
if SinkExt::send(&mut socket, Message::Binary(data)).await.is_err() {
break;
}
}
Ok(Message::Close(_)) => break,
Ok(Message::Ping(data)) => {
if SinkExt::send(&mut socket, Message::Pong(data)).await.is_err() {
break;
}
}
Ok(Message::Pong(_)) => {}
Err(err) => {
debug!(%err, "websocket read ended");
break;
}
}
}
}