Initial commit: Rust WebSocket/UDP signaling and media relay server

src/handlers.rs

@@ -0,0 +1,129 @@
+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);
+}