129 lines
4.5 KiB
Rust
129 lines
4.5 KiB
Rust
use axum::{
|
|
extract::ws::{Message, WebSocket},
|
|
};
|
|
use futures::{SinkExt, StreamExt};
|
|
use shared::{ControlMsg, PeerInfo, UserId};
|
|
use std::net::SocketAddr;
|
|
use tracing::{info, error};
|
|
|
|
use crate::state::{AppState, Peer, Room};
|
|
|
|
pub async fn handle_socket(mut socket: WebSocket, _addr: SocketAddr, state: AppState) {
|
|
let (mut sender, mut receiver) = socket.split();
|
|
let mut user_id: Option<UserId> = None;
|
|
let mut room_id: Option<String> = None;
|
|
|
|
// Authentication / Join Phase
|
|
while let Some(msg) = receiver.next().await {
|
|
match msg {
|
|
Ok(Message::Text(text)) => {
|
|
if let Ok(ControlMsg::Join { room_code, display_name }) = serde_json::from_str(&text) {
|
|
let uid = rand::random::<u32>();
|
|
|
|
info!("User {} ({}) joining room {}", uid, display_name, room_code);
|
|
|
|
let room = state.rooms.entry(room_code.clone())
|
|
.or_insert_with(|| Room::new(room_code.clone()));
|
|
|
|
// Gather existing peers
|
|
let peers_list: Vec<PeerInfo> = room.peers.iter().map(|p| {
|
|
let k = *p.key();
|
|
let v = p.value();
|
|
PeerInfo {
|
|
user_id: k,
|
|
display_name: v.display_name.clone(),
|
|
}
|
|
}).collect();
|
|
|
|
// Notify others
|
|
let _ = room.tx.send(ControlMsg::PeerJoined {
|
|
user_id: uid,
|
|
display_name: display_name.clone(),
|
|
});
|
|
|
|
// Add self to room
|
|
room.peers.insert(uid, Peer {
|
|
id: uid,
|
|
display_name: display_name.clone(),
|
|
addr: None,
|
|
});
|
|
|
|
user_id = Some(uid);
|
|
room_id = Some(room_code.clone());
|
|
|
|
// Send Joined response
|
|
let resp = ControlMsg::Joined {
|
|
self_id: uid,
|
|
room_code: room_code.clone(),
|
|
peers: peers_list,
|
|
};
|
|
if let Err(e) = sender.send(Message::Text(serde_json::to_string(&resp).unwrap().into())).await {
|
|
error!("Failed to send Joined response: {}", e);
|
|
return;
|
|
}
|
|
|
|
break;
|
|
}
|
|
}
|
|
Ok(Message::Close(_)) => return,
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
if user_id.is_none() {
|
|
return;
|
|
}
|
|
|
|
let uid = user_id.unwrap();
|
|
let rid = room_id.unwrap();
|
|
|
|
// Subscribe to room broadcasts
|
|
let mut rx = {
|
|
let room = state.rooms.get(&rid).unwrap();
|
|
room.tx.subscribe()
|
|
};
|
|
|
|
// Main Loop
|
|
loop {
|
|
tokio::select! {
|
|
msg = receiver.next() => {
|
|
match msg {
|
|
Some(Ok(Message::Text(text))) => {
|
|
if let Ok(control) = serde_json::from_str::<ControlMsg>(&text) {
|
|
match control {
|
|
ControlMsg::UpdateStream { .. } => {
|
|
// Broadcast to room
|
|
if let Some(room) = state.rooms.get(&rid) {
|
|
let _ = room.tx.send(control);
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
Some(Ok(Message::Close(_))) => break,
|
|
Some(Err(_)) => break,
|
|
None => break,
|
|
_ => {}
|
|
}
|
|
}
|
|
Ok(msg) = rx.recv() => {
|
|
// Forward broadcast to client
|
|
let _ = sender.send(Message::Text(serde_json::to_string(&msg).unwrap().into())).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Cleanup
|
|
if let Some(room) = state.rooms.get(&rid) {
|
|
if let Some((_, peer)) = room.peers.remove(&uid) {
|
|
// Clean up address mapping if present
|
|
if let Some(addr) = peer.addr {
|
|
state.peers_by_addr.remove(&addr);
|
|
}
|
|
}
|
|
let _ = room.tx.send(ControlMsg::PeerLeft { user_id: uid });
|
|
}
|
|
|
|
info!("User {} left room {}", uid, rid);
|
|
}
|