From 01ae72c6cc24ca3de0c5542adea6c2c545f06e1b Mon Sep 17 00:00:00 2001 From: David Westgate Date: Wed, 29 Nov 2023 16:36:50 -0800 Subject: [PATCH] handle quitting finally --- src/client.rs | 14 ++++++++++---- src/server.rs | 48 +++++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 53 insertions(+), 9 deletions(-) diff --git a/src/client.rs b/src/client.rs index c7fc7fe..be23806 100644 --- a/src/client.rs +++ b/src/client.rs @@ -59,6 +59,7 @@ fn read_messages(mut stream: TcpStream, nick: &str) { } fn process_message(msg_bytes: &[u8], nick: &str) { + println!(); match msg_bytes[0] { codes::ERROR => { @@ -70,11 +71,9 @@ fn process_message(msg_bytes: &[u8], nick: &str) { println!( "Nickname already in use on server. Connect again with a different one" ); - disconnect(); } codes::error::SERVER_FULL => { println!("Server is full. Try again later"); - disconnect(); } _ => { println!("Error code: {:x?}", msg_bytes[1]); @@ -113,6 +112,10 @@ fn process_message(msg_bytes: &[u8], nick: &str) { let message = String::from_utf8(msg_bytes[1..msg_bytes.len()].to_vec()).unwrap(); println!("{}", message); } + codes::QUIT => { + println!("Server has closed the connection. Stopping client"); + std::process::exit(0); + } _ => { #[cfg(debug_assertions)] println!("BAD RESPONSE = {:x?} ", msg_bytes[0]); @@ -120,7 +123,10 @@ fn process_message(msg_bytes: &[u8], nick: &str) { } } -fn disconnect() {} +fn disconnect(stream: &mut TcpStream) { + stream.write_all(&[codes::QUIT]).unwrap(); + stream.shutdown(std::net::Shutdown::Both).unwrap(); +} fn help() { clear(); @@ -179,7 +185,7 @@ pub fn start() { Some(cmd) => { let param: Option<&str> = args.next(); match cmd { - "/quit" => disconnect(), + "/quit" => {disconnect(&mut stream); break}, "/rooms" => no_param_op(codes::client::LIST_ROOMS, &mut stream), "/users" => no_param_op(codes::client::LIST_USERS, &mut stream), "/list" => match param { diff --git a/src/server.rs b/src/server.rs index b76350d..cb8250f 100644 --- a/src/server.rs +++ b/src/server.rs @@ -30,7 +30,6 @@ impl Server { fn message(room: &str, msg: &str, sender: &str, server: &Arc>) { - println!("message fn {} {}", room, msg); let size = room.len() + msg.len() + sender.len() + 3; let mut out_buf: Vec = vec![0; size]; @@ -59,7 +58,7 @@ fn message(room: &str, msg: &str, sender: &str, server: &Arc>) { byte += 1; } - let mut guard = server.lock().unwrap(); + let mut guard: std::sync::MutexGuard<'_, Server> = server.lock().unwrap(); let server: &mut Server = guard.deref_mut(); let room_users: Option<&Vec> = server.rooms.get(room); @@ -95,12 +94,22 @@ fn broadcast(op: u8, server: &Arc>, message: &str) { } let mut unlocked_server: std::sync::MutexGuard<'_, Server> = server.lock().unwrap(); - let streams = unlocked_server.users.values_mut(); + let streams: std::collections::hash_map::ValuesMut<'_, String, TcpStream> = unlocked_server.users.values_mut(); for stream in streams { stream.write_all(&out_buf).unwrap(); } } +fn disconnect_all(server: &Arc>,) { + let mut guard: std::sync::MutexGuard<'_, Server> = server.lock().unwrap(); + let users: std::collections::hash_map::ValuesMut<'_, String, TcpStream> = guard.users.values_mut(); + users.for_each(|user: &mut TcpStream| { + user.write(&[codes::QUIT]).unwrap(); + user.shutdown(std::net::Shutdown::Both); + }) + +} + fn handle_client( server: &Arc>, stream: &mut TcpStream, @@ -193,6 +202,9 @@ fn handle_client( } } } + codes::QUIT => { + remove_user(server, nickname,stream); + } _ => { #[cfg(debug_assertions)] println!("Unspecified client Op, {:x?}", cmd_bytes); @@ -202,6 +214,18 @@ fn handle_client( // } } +fn remove_user(server: &Arc>, nickname: &str, stream: &mut TcpStream) { + let mut guard: std::sync::MutexGuard<'_, Server> = server.lock().unwrap(); + let server: &mut Server = guard.deref_mut(); + let mut rooms: &mut HashMap> = &mut server.rooms; + rooms.values_mut().for_each(|room: &mut Vec| { + room.retain(|u| !u.eq(nickname)); + }); + let users: &mut HashMap = &mut server.users; + users.remove(nickname); + +} + fn register_nick(server: &Arc>, nickname: &str, stream: &mut TcpStream) { // Check for nickname collision let mut unlocked_server: std::sync::MutexGuard<'_, Server> = server.lock().unwrap(); @@ -214,10 +238,13 @@ fn register_nick(server: &Arc>, nickname: &str, stream: &mut TcpSt } else { // Add the user to the user list let clone = stream.try_clone().expect("fail to clone"); - unlocked_server.users.insert(nickname.to_string(), clone); + let addr = clone.peer_addr().unwrap().to_string(); + unlocked_server.users.insert(nickname.to_string(), clone); // Send response ok stream.write_all(&[codes::RESPONSE_OK]).unwrap(); + + println!("{} has registered nickname {}", addr, nickname); } } @@ -285,7 +312,11 @@ pub fn start() { thread::spawn(move || { let nickname: String; + println!("IP {} has connected", stream.peer_addr().unwrap().to_string()); match stream.read(&mut buf_in) { + Ok(0) => { + println!("IP {} has closed the connection", stream.peer_addr().unwrap().to_string()); + } Ok(size) => { let cmd_bytes: &[u8] = &buf_in[0..1]; let param_bytes: &[u8] = &buf_in[1..size]; @@ -294,6 +325,10 @@ pub fn start() { register_nick(&server_inner, &nickname, &mut stream); loop { match stream.read(&mut buf_in) { + Ok(0) => { + println!("IP {} wit nickname {} has closed the connection", stream.peer_addr().unwrap().to_string(), nickname); + break; + } Ok(size) => { let cmd_bytes: &[u8] = &buf_in[0..1]; let param_bytes: &[u8] = &buf_in[1..size]; @@ -345,7 +380,10 @@ pub fn start() { match inp.parse::() { Ok(num) => match num { 0 => { - println!("Goodbye"); + println!("Stopping Server"); + disconnect_all(&server); + break; + } 1 => println!("Users: {:?}", server.lock().unwrap().users), 2 => println!("Rooms: {:?}", server.lock().unwrap().rooms),