From a996f13deb02c74f28829761a57efab0f4c07230 Mon Sep 17 00:00:00 2001 From: Hatter Jiang Date: Sun, 5 Nov 2023 16:18:42 +0800 Subject: [PATCH] feat: optimize code --- src/main.rs | 256 ++++++++++++++++++++++++++-------------------------- 1 file changed, 129 insertions(+), 127 deletions(-) diff --git a/src/main.rs b/src/main.rs index 3fd0c53..34a24b4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -128,11 +128,6 @@ async fn inner_handle_connection( Ok(()) } -fn client_not_in_room(tx: &Tx, addr: SocketAddr) { - information!("Client is not in room: {}", addr); - RoomMessageDown::create_error_reply(&None, "Client is not in room").send(tx); -} - fn client_exit(handle_context: &HandleContext, room_id: &Option, client_id: &Option) { let room_map = &handle_context.room_map; if let (Some(room_id), Some(client_id)) = (room_id, client_id) { @@ -184,159 +179,166 @@ fn handle_text_message(handle_context: &mut HandleContext, tx: &Tx, addr: Socket RoomMessageDown::create_error_reply(&room_message.message_id, "Client is already in room").send(tx); return Ok(()); } - if let (Some(msg_room_id), Some(msg_client_id)) = (room_message.room_id, room_message.client_id) { - let mut room_map = handle_context.room_map.lock().unwrap(); - match room_map.get_mut(&msg_room_id) { - Some(client_map) => { - match client_map.get(&msg_client_id) { - Some(peer_addr) => { - if peer_addr == &addr { - information!("Duplicate enter to room: {:?} - {:?}", msg_room_id, msg_client_id); - RoomMessageDown::create_success_reply( - &room_message.message_id, "Duplicate enter to room").send(tx); - } else { - information!( - "Replace client: {:?} - {:?}, from {:?} -> {:?}", - msg_room_id, msg_client_id, peer_addr, addr - ); - let client_replaced_message = format!("Client replaced {:?} -> {:?}", peer_addr, addr); - if let Some(tx) = handle_context.peer_map.lock().unwrap().remove(peer_addr) { - tx.send_close(); - } - client_map.insert(msg_client_id.clone(), addr); - RoomMessageDown::create_success_reply(&room_message.message_id, client_replaced_message).send(tx); - } - } - None => { - information!("Enter room: {:?} - {:?}, addr: {}", msg_room_id, msg_client_id, addr); - client_map.insert(msg_client_id.clone(), addr); - - let m = RoomMessageDown::create_peer_enter(&room_message.message_id, msg_client_id.clone()); - let mm = serde_json::to_string(&m)?; - for (client_id, client_addr) in client_map { - if client_id != &msg_client_id { - if let Some(client_tx) = handle_context.peer_map.lock().unwrap().get(client_addr) { - client_tx.send_text(mm.clone()); - } - } - } - + let (msg_room_id, msg_client_id) = match (room_message.room_id, room_message.client_id) { + (Some(msg_room_id), Some(msg_client_id)) => (msg_room_id, msg_client_id), + _ => { + RoomMessageDown::create_error_reply(&room_message.message_id, "Room id and client id must both assigned").send(tx); + return Ok(()); + } + }; + let mut room_map = handle_context.room_map.lock().unwrap(); + match room_map.get_mut(&msg_room_id) { + Some(client_map) => { + match client_map.get(&msg_client_id) { + Some(peer_addr) => { + if peer_addr == &addr { + information!("Duplicate enter to room: {:?} - {:?}", msg_room_id, msg_client_id); RoomMessageDown::create_success_reply( - &room_message.message_id, format!("Client entered: {:?}", addr)).send(tx); + &room_message.message_id, "Duplicate enter to room").send(tx); + } else { + information!("Replace client: {:?} - {:?}, from {:?} -> {:?}", msg_room_id, msg_client_id, peer_addr, addr); + let client_replaced_message = format!("Client replaced {:?} -> {:?}", peer_addr, addr); + if let Some(tx) = handle_context.peer_map.lock().unwrap().remove(peer_addr) { + tx.send_close(); + } + client_map.insert(msg_client_id.clone(), addr); + RoomMessageDown::create_success_reply(&room_message.message_id, client_replaced_message).send(tx); } } - } - None => { - information!("Create and enter room: {:?} - {:?}, addr: {}", msg_room_id, msg_client_id, addr); - let mut client_map = BTreeMap::new(); - client_map.insert(msg_client_id.clone(), addr); - room_map.insert(msg_room_id.clone(), client_map); - RoomMessageDown::create_success_reply( - &room_message.message_id, format!("Client create room: {:?}", addr)).send(tx); + None => { + information!("Enter room: {:?} - {:?}, addr: {}", msg_room_id, msg_client_id, addr); + client_map.insert(msg_client_id.clone(), addr); + + let m = RoomMessageDown::create_peer_enter(&room_message.message_id, msg_client_id.clone()); + let mm = serde_json::to_string(&m)?; + for (client_id, client_addr) in client_map { + if client_id != &msg_client_id { + if let Some(client_tx) = handle_context.peer_map.lock().unwrap().get(client_addr) { + client_tx.send_text(mm.clone()); + } + } + } + + RoomMessageDown::create_success_reply( + &room_message.message_id, format!("Client entered: {:?}", addr)).send(tx); + } } } - handle_context.room_id = Some(msg_room_id); - handle_context.client_id = Some(msg_client_id); - } else { - RoomMessageDown::create_error_reply(&room_message.message_id, "Room id and client id must both assigned").send(tx); + None => { + information!("Create and enter room: {:?} - {:?}, addr: {}", msg_room_id, msg_client_id, addr); + let mut client_map = BTreeMap::new(); + client_map.insert(msg_client_id.clone(), addr); + room_map.insert(msg_room_id.clone(), client_map); + RoomMessageDown::create_success_reply( + &room_message.message_id, format!("Client create room: {:?}", addr)).send(tx); + } } + // Room created + handle_context.room_id = Some(msg_room_id); + handle_context.client_id = Some(msg_client_id); } RoomMessageType::Exit => { - if let (Some(room_id), Some(client_id)) = (&handle_context.room_id, &handle_context.client_id) { - let mut room_map = handle_context.room_map.lock().unwrap(); - if let Some(client_map) = room_map.get_mut(room_id) { - client_map.remove(client_id); - handle_context.peer_map.lock().unwrap().remove(&addr); - tx.send_close(); - } else { - warning!("Not in room: {:?} - {:?}", room_id, client_id); - } + let (room_id, client_id) = match (&handle_context.room_id, &handle_context.client_id) { + (Some(room_id), Some(client_id)) => (room_id, client_id), + _ => return client_not_in_room(tx, addr), + }; + let mut room_map = handle_context.room_map.lock().unwrap(); + if let Some(client_map) = room_map.get_mut(room_id) { + client_map.remove(client_id); + handle_context.peer_map.lock().unwrap().remove(&addr); + tx.send_close(); } else { - client_not_in_room(tx, addr); + warning!("Not in room: {:?} - {:?}", room_id, client_id); } } RoomMessageType::Destroy => { - if let (Some(room_id), Some(client_id)) = (&handle_context.room_id, &handle_context.client_id) { - information!("Destroy room: {:?} - {:?}", room_id, client_id); - let mut room_map = handle_context.room_map.lock().unwrap(); - let client_map = room_map.remove(room_id); - if let Some(client_map) = client_map { - for (_client_id, client_addr) in client_map { - if let Some(client_tx) = handle_context.peer_map.lock().unwrap().remove(&client_addr) { - client_tx.send_close(); - } + let (room_id, client_id) = match (&handle_context.room_id, &handle_context.client_id) { + (Some(room_id), Some(client_id)) => (room_id, client_id), + _ => return client_not_in_room(tx, addr), + }; + information!("Destroy room: {:?} - {:?}", room_id, client_id); + let mut room_map = handle_context.room_map.lock().unwrap(); + let client_map = room_map.remove(room_id); + if let Some(client_map) = client_map { + for (_client_id, client_addr) in client_map { + if let Some(client_tx) = handle_context.peer_map.lock().unwrap().remove(&client_addr) { + client_tx.send_close(); } } - } else { - client_not_in_room(tx, addr); } } RoomMessageType::ListPeers => { - if let (Some(room_id), Some(client_id)) = (&handle_context.room_id, &handle_context.client_id) { - information!("List room peers: {:?} - {:?}", room_id, client_id); - let room_map = handle_context.room_map.lock().unwrap(); - let client_map = room_map.get(room_id); - let client_ids = client_map - .map(|m| m.keys().cloned().collect::>()) - .unwrap_or_else(Vec::new); - let client_ids_data = serde_json::to_string(&client_ids)?; - RoomMessageDown::create_peerlist_message(&room_message.message_id, Some(client_ids_data)).send(tx); - } else { - client_not_in_room(tx, addr); - } + let (room_id, client_id) = match (&handle_context.room_id, &handle_context.client_id) { + (Some(room_id), Some(client_id)) => (room_id, client_id), + _ => return client_not_in_room(tx, addr), + }; + information!("List room peers: {:?} - {:?}", room_id, client_id); + let room_map = handle_context.room_map.lock().unwrap(); + let client_map = room_map.get(room_id); + let client_ids = client_map + .map(|m| m.keys().cloned().collect::>()) + .unwrap_or_else(Vec::new); + let client_ids_data = serde_json::to_string(&client_ids)?; + RoomMessageDown::create_peerlist_message(&room_message.message_id, Some(client_ids_data)).send(tx); } RoomMessageType::Broadcast => { - if let (Some(room_id), Some(client_id)) = (&handle_context.room_id, &handle_context.client_id) { - information!("Broadcast room message: {:?} - {:?}", room_id, client_id); - let room_map = handle_context.room_map.lock().unwrap(); - let client_map = room_map.get(room_id); + let (room_id, client_id) = match (&handle_context.room_id, &handle_context.client_id) { + (Some(room_id), Some(client_id)) => (room_id, client_id), + _ => return client_not_in_room(tx, addr), + }; + information!("Broadcast room message: {:?} - {:?}", room_id, client_id); + let room_map = handle_context.room_map.lock().unwrap(); + let client_map = room_map.get(room_id); - let m = RoomMessageDown::create_broadcast_message( - &room_message.message_id, client_id, room_message.data); - let mm = serde_json::to_string(&m)?; - if let Some(client_map) = client_map { - let mut sent_messages = 0; - for (peer_client_id, peer_client_addr) in client_map { - if peer_client_id != client_id { - if let Some(peer_tx) = handle_context.peer_map.lock().unwrap().get(peer_client_addr) { - sent_messages += 1; - peer_tx.send_text(mm.clone()); - } + let m = RoomMessageDown::create_broadcast_message( + &room_message.message_id, client_id, room_message.data); + let mm = serde_json::to_string(&m)?; + if let Some(client_map) = client_map { + let mut sent_messages = 0; + for (peer_client_id, peer_client_addr) in client_map { + if peer_client_id != client_id { + if let Some(peer_tx) = handle_context.peer_map.lock().unwrap().get(peer_client_addr) { + sent_messages += 1; + peer_tx.send_text(mm.clone()); } } - RoomMessageDown::create_success_reply( - &room_message.message_id, format!("Send message to {} peers", sent_messages)).send(tx); } - } else { - client_not_in_room(tx, addr); + RoomMessageDown::create_success_reply( + &room_message.message_id, format!("Send message to {} peers", sent_messages)).send(tx); } } RoomMessageType::Peer => { - if let (Some(room_id), Some(client_id)) = (&handle_context.room_id, &handle_context.client_id) { - information!("Send peer message: {:?} - {:?} -> {:?}", room_id, client_id, room_message.peer_id); - if let Some(peer_id) = room_message.peer_id.as_ref() { - let room_map = handle_context.room_map.lock().unwrap(); - let client_map = room_map.get(room_id); + let (room_id, client_id) = match (&handle_context.room_id, &handle_context.client_id) { + (Some(room_id), Some(client_id)) => (room_id, client_id), + _ => return client_not_in_room(tx, addr), + }; + information!("Send peer message: {:?} - {:?} -> {:?}", room_id, client_id, room_message.peer_id); + if let Some(peer_id) = room_message.peer_id.as_ref() { + let room_map = handle_context.room_map.lock().unwrap(); + let client_map = room_map.get(room_id); - let m = RoomMessageDown::create_peer_message(&room_message.message_id, peer_id, room_message.data); - let mm = serde_json::to_string(&m)?; - if let Some(client_map) = client_map { - if let Some(peer_client_addr) = client_map.get(peer_id) { - if let Some(peer_tx) = handle_context.peer_map.lock().unwrap().get(peer_client_addr) { - peer_tx.send_text(mm); - } - RoomMessageDown::create_success_reply( - &room_message.message_id, format!("Message sent to: {}", peer_id)).send(tx); - } else { - RoomMessageDown::create_fail_replay( - &room_message.message_id, 404, format!("Peer not found: {}", peer_id)).send(tx); + let m = RoomMessageDown::create_peer_message(&room_message.message_id, peer_id, room_message.data); + let mm = serde_json::to_string(&m)?; + if let Some(client_map) = client_map { + if let Some(peer_client_addr) = client_map.get(peer_id) { + if let Some(peer_tx) = handle_context.peer_map.lock().unwrap().get(peer_client_addr) { + peer_tx.send_text(mm); } + RoomMessageDown::create_success_reply( + &room_message.message_id, format!("Message sent to: {}", peer_id)).send(tx); + } else { + RoomMessageDown::create_fail_replay( + &room_message.message_id, 404, format!("Peer not found: {}", peer_id)).send(tx); } } - } else { - client_not_in_room(tx, addr); } } } Ok(()) -} \ No newline at end of file +} + +fn client_not_in_room(tx: &Tx, addr: SocketAddr) -> XResult<()> { + information!("Client is not in room: {}", addr); + RoomMessageDown::create_error_reply(&None, "Client is not in room").send(tx); + Ok(()) +}