use std::sync::{Arc, Mutex};
use std::time::Instant;
use axum::extract::ws::{Message, WebSocket};
use axum::extract::{State, WebSocketUpgrade};
use axum::response::IntoResponse;
use futures::{SinkExt, StreamExt};
use tokio::sync::mpsc;
use tracing::{info, warn};
use uuid::Uuid;
use crate::message::{ClientMessage, ServerMessage};
use crate::room::ConnectedClient;
use crate::state::AppState;
pub async fn ws_handler(
ws: WebSocketUpgrade,
State(state): State<AppState>,
) -> impl IntoResponse {
ws.on_upgrade(move |socket| handle_socket(socket, state))
}
async fn handle_socket(socket: WebSocket, state: AppState) {
let (mut ws_sender, mut ws_receiver) = socket.split();
let (client_tx, mut client_rx) = mpsc::channel::<String>(256);
let client_id = Uuid::new_v4();
let config = state.config.clone();
let last_pong = Arc::new(Mutex::new(Instant::now()));
let write_task = {
let last_pong = Arc::clone(&last_pong);
tokio::spawn(async move {
let ping_interval =
tokio::time::Duration::from_secs(config.heartbeat.ping_interval_secs);
let pong_timeout =
tokio::time::Duration::from_secs(config.heartbeat.pong_timeout_secs);
let mut ping_timer = tokio::time::interval(ping_interval);
ping_timer.tick().await;
loop {
tokio::select! {
Some(text) = client_rx.recv() => {
if ws_sender.send(Message::Text(text)).await.is_err() {
break;
}
}
_ = ping_timer.tick() => {
if last_pong.lock().unwrap().elapsed() > pong_timeout {
warn!(client_id = %client_id, "pong timeout, disconnecting");
let _ = ws_sender.close().await;
break;
}
let ping_msg = serde_json::to_string(&ServerMessage::Ping).unwrap();
if ws_sender.send(Message::Text(ping_msg)).await.is_err() {
break;
}
}
else => break,
}
}
})
};
let max_message_size = config.limits.max_message_size;
let max_username_len = config.limits.max_username_len;
let mut current_room: Option<String> = None;
let mut current_username: String = String::new();
while let Some(msg_result) = ws_receiver.next().await {
let msg = match msg_result {
Ok(Message::Text(text)) => text,
Ok(Message::Close(_)) => break,
_ => continue,
};
if msg.len() > max_message_size {
let err = ServerMessage::Error {
code: "message_too_large".into(),
message: format!("max message size is {} bytes", max_message_size),
};
let payload = serde_json::to_string(&err).unwrap();
let _ = client_tx.try_send(payload);
continue;
}
let client_msg: ClientMessage = match serde_json::from_str(&msg) {
Ok(m) => m,
Err(e) => {
warn!(%e, "invalid message from client");
let err = ServerMessage::Error {
code: "invalid_message".into(),
message: format!("invalid JSON: {}", e),
};
let payload = serde_json::to_string(&err).unwrap();
let _ = client_tx.try_send(payload);
continue;
}
};
match client_msg {
ClientMessage::Join { room, username } => {
if username.len() > max_username_len {
let err = ServerMessage::Error {
code: "username_too_long".into(),
message: format!("max username length is {}", max_username_len),
};
let payload = serde_json::to_string(&err).unwrap();
let _ = client_tx.try_send(payload);
continue;
}
let client = ConnectedClient {
id: client_id,
username: username.clone(),
sender: client_tx.clone(),
};
match state.join_room(&room, client).await {
Ok(()) => {
info!(%client_id, %username, %room, "client joined room");
current_room = Some(room.clone());
current_username = username.clone();
let joined = ServerMessage::Joined {
room: room.clone(),
username: username.clone(),
};
let payload = serde_json::to_string(&joined).unwrap();
let _ = client_tx.try_send(payload);
let user_joined = ServerMessage::UserJoined {
room: room.clone(),
username: username.clone(),
};
state.broadcast(&room, &user_joined, Some(client_id)).await;
}
Err(e) => {
let err = ServerMessage::Error {
code: "join_failed".into(),
message: e,
};
let payload = serde_json::to_string(&err).unwrap();
let _ = client_tx.try_send(payload);
}
}
}
ClientMessage::Send { room, text } => {
if current_username.is_empty() {
let err = ServerMessage::Error {
code: "not_joined".into(),
message: "you must join a room first".into(),
};
let payload = serde_json::to_string(&err).unwrap();
let _ = client_tx.try_send(payload);
continue;
}
let at = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
let msg = ServerMessage::Message {
room: room.clone(),
from: current_username.clone(),
text,
at,
};
state.broadcast(&room, &msg, None).await;
}
ClientMessage::Private { to, text } => {
if current_username.is_empty() {
let err = ServerMessage::Error {
code: "not_joined".into(),
message: "you must join a room first".into(),
};
let payload = serde_json::to_string(&err).unwrap();
let _ = client_tx.try_send(payload);
continue;
}
match state.send_private(¤t_username, &to, &text).await {
Ok(()) => {
info!(from = %current_username, to = %to, "private message sent");
}
Err(e) => {
let err = ServerMessage::Error {
code: "private_failed".into(),
message: e,
};
let payload = serde_json::to_string(&err).unwrap();
let _ = client_tx.try_send(payload);
}
}
}
ClientMessage::Ping => {
*last_pong.lock().unwrap() = Instant::now();
let pong = ServerMessage::Pong;
let payload = serde_json::to_string(&pong).unwrap();
let _ = client_tx.try_send(payload);
}
}
}
if let Some(room) = current_room {
if let Some(username) = state.leave_room(&room, client_id).await {
info!(%client_id, %username, %room, "client left room");
let user_left = ServerMessage::UserLeft {
room: room.clone(),
username,
};
state.broadcast(&room, &user_left, None).await;
}
}
write_task.abort();
info!(%client_id, "client disconnected");
}