From d11c6782a71fdda9e1f259c1b443015c99a1bbd1 Mon Sep 17 00:00:00 2001 From: Michel de Graaf Date: Tue, 4 Nov 2025 18:48:32 +0100 Subject: [PATCH] chore: refactor default peer --- src/client.rs | 7 +- src/message/mod.rs | 1 + src/message/peer/mod.rs | 2 + src/message/peer/peer_init.rs | 20 +++++ src/message/server/login.rs | 2 - src/peer/default_peer.rs | 148 ++++++++++++++++++++-------------- src/peer/download_peer.rs | 23 +++--- src/peer/listen.rs | 119 ++++++--------------------- 8 files changed, 154 insertions(+), 168 deletions(-) create mode 100644 src/message/peer/peer_init.rs diff --git a/src/client.rs b/src/client.rs index ff1dd1a..6674c53 100644 --- a/src/client.rs +++ b/src/client.rs @@ -310,7 +310,7 @@ impl Client { download.username.clone(), peer.host.clone(), peer.port, - download.token, + token, false, own_username, ); @@ -500,7 +500,7 @@ impl Client { let peer_clone = peer.clone(); let sender_clone = sender; unlocked_context.thread_pool.execute(move || { - trace!("[client] connecting to {}, with connection_type: {}", peer.username, peer.connection_type); + trace!("[client] connecting to {}, with connection_type: {}, and token {:?}", peer.username, peer.connection_type, peer.token); match peer.connection_type { ConnectionType::P => { let default_peer = @@ -530,7 +530,7 @@ impl Client { } ConnectionType::F => { - trace!("[client] downloading from: {}", peer.username); + trace!("[client] downloading from: {}, {:?}", peer.username, peer.token); let download_peer = DownloadPeer::new( peer.username, peer.host, @@ -547,6 +547,7 @@ impl Client { ) { Ok((download, filename)) => { trace!("[client] downloaded {} bytes {:?} ", filename, download.size); + download.sender.send(DownloadStatus::Completed).unwrap(); } Err(e) => { trace!("[client] failed to download: {}", e); diff --git a/src/message/mod.rs b/src/message/mod.rs index 76e2462..82506fa 100644 --- a/src/message/mod.rs +++ b/src/message/mod.rs @@ -358,6 +358,7 @@ impl Message { ))), }, MessageType::Peer => match code { + 1 => Ok("PeerInit"), 4 => Ok("GetShareFileList"), 5 => Ok("SharedFileListResponse"), 9 => Ok("FileSearchResponse"), diff --git a/src/message/peer/mod.rs b/src/message/peer/mod.rs index f122fa9..c876cf9 100644 --- a/src/message/peer/mod.rs +++ b/src/message/peer/mod.rs @@ -1,5 +1,6 @@ mod file_search_response; mod get_share_file_list; +mod peer_init; mod place_in_queue_response; mod transfer_request; mod transfer_response; @@ -8,6 +9,7 @@ mod upload_failed; // Re-export handlers pub use file_search_response::FileSearchResponse; pub use get_share_file_list::GetShareFileList; +pub use peer_init::PeerInit; pub use place_in_queue_response::PlaceInQueueResponse; pub use transfer_request::TransferRequest; pub use transfer_response::TransferResponse; diff --git a/src/message/peer/peer_init.rs b/src/message/peer/peer_init.rs new file mode 100644 index 0000000..882e95c --- /dev/null +++ b/src/message/peer/peer_init.rs @@ -0,0 +1,20 @@ +use std::sync::mpsc::Sender; + +use crate::{ + message::{Message, MessageHandler}, + peer::PeerOperation, +}; + +pub struct PeerInit; +impl MessageHandler for PeerInit { + fn get_code(&self) -> u8 { + 1 + } + + fn handle(&self, message: &mut Message, sender: Sender) { + message.set_pointer(4); + let _message_code = message.read_int8(); + let username = message.read_string(); + sender.send(PeerOperation::SetUsername(username)).unwrap(); + } +} diff --git a/src/message/server/login.rs b/src/message/server/login.rs index 742dd0c..fc0396b 100644 --- a/src/message/server/login.rs +++ b/src/message/server/login.rs @@ -3,12 +3,10 @@ use std::sync::mpsc::Sender; use crate::{message::Message, server::ServerOperation}; -// IMPORTANT: note the generic trait is now `MessageHandler` use crate::message::MessageHandler; pub struct LoginHandler; -// Implement `MessageHandler` specifically for `ServerOperation`. impl MessageHandler for LoginHandler { fn get_code(&self) -> u8 { 1 diff --git a/src/peer/default_peer.rs b/src/peer/default_peer.rs index c95d798..af69191 100644 --- a/src/peer/default_peer.rs +++ b/src/peer/default_peer.rs @@ -1,5 +1,5 @@ use crate::dispatcher::MessageDispatcher; -use crate::message::peer::FileSearchResponse; +use crate::message::peer::{FileSearchResponse, PeerInit}; use crate::message::peer::{ GetShareFileList, PlaceInQueueResponse, TransferRequest, TransferResponse, UploadFailedHandler, @@ -9,8 +9,8 @@ use crate::message::{Handlers, Message, MessageReader, MessageType}; use crate::types::{Download, FileSearchResult, Transfer}; use core::result::Result::Ok; -use std::sync::mpsc; use std::sync::mpsc::{Receiver, Sender}; +use std::sync::{mpsc, Arc, RwLock}; use std::thread::{self, JoinHandle}; use crate::client::ClientOperation; @@ -24,7 +24,7 @@ use std::time::Duration; #[allow(dead_code)] #[derive(Debug)] pub struct DefaultPeer { - peer: Peer, + peer: Arc>, peer_channel: Option>, client_channel: Sender, read_thread: Option>, @@ -45,11 +45,13 @@ pub enum PeerOperation { filename: String, place: u32, }, + SetUsername(String), } + impl DefaultPeer { pub fn new(peer: Peer, client_channel: Sender) -> Self { Self { - peer, + peer: Arc::new(RwLock::new(peer)), peer_channel: None, client_channel, read_thread: None, @@ -57,10 +59,12 @@ impl DefaultPeer { } } pub fn disconnect(mut self) { - debug!("[default_peer:{}] disconnect", self.peer.username.clone()); - if let Err(e) = self.client_channel.send( - ClientOperation::PeerDisconnected(self.peer.username.clone()), - ) { + let username = self.peer.read().unwrap().username.clone(); + debug!("[default_peer:{}] disconnect", username); + if let Err(e) = self + .client_channel + .send(ClientOperation::PeerDisconnected(username)) + { error!("Failed to send disconnect notification: {}", e); } @@ -71,7 +75,7 @@ impl DefaultPeer { mut self, mut stream: TcpStream, ) -> Result { - if let Some(token) = self.peer.token { + if let Some(token) = self.peer.read().unwrap().token { let mut message: Vec = [0, 5, 0, 0, 0, 0].to_vec(); message.extend_from_slice(&token.to_le_bytes()); stream.write_all(&message).unwrap(); @@ -83,25 +87,35 @@ impl DefaultPeer { trace!( "[default_peer:{}] connect_with_socket: direct", - self.peer.username + self.peer.read().unwrap().username ); self.start_read_write_loops(stream)?; Ok(self) } - pub fn connect(mut self) -> Result { - println!( - "[default_peer] Connecting to {} on port {}", - self.peer.host, self.peer.port - ); - let socket_address = format!("{}:{}", self.peer.host, self.peer.port) - .to_socket_addrs()? - .next() - .ok_or_else(|| { - io::Error::new(io::ErrorKind::InvalidInput, "Invalid address") - })?; - trace!("[default_peer:{}] connect: indirect", self.peer.username); + pub fn connect(mut self) -> Result { + { + let peer = self.peer.read().unwrap(); + println!( + "[default_peer] Connecting to {} on port {}", + peer.host, peer.port + ); + } + let socket_address = { + let peer = self.peer.read().unwrap(); + format!("{}:{}", peer.host, peer.port) + } + .to_socket_addrs()? + .next() + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "Invalid address") + })?; + + trace!( + "[default_peer:{}] connect: indirect", + self.peer.read().unwrap().username + ); let mut stream = TcpStream::connect_timeout( &socket_address, Duration::from_secs(20), @@ -111,7 +125,7 @@ impl DefaultPeer { stream.set_write_timeout(Some(Duration::from_secs(5)))?; stream.set_nodelay(true)?; - if let Some(token) = self.peer.token { + if let Some(token) = self.peer.read().unwrap().token { stream .write_all(&MessageFactory::build_watch_user(token).get_data()) .unwrap(); @@ -127,7 +141,7 @@ impl DefaultPeer { ) -> Result<(), io::Error> { debug!( "[default_peer:{}] start_read_write_loops", - self.peer.username + self.peer.read().unwrap().username ); let (peer_sender, peer_reader): ( Sender, @@ -152,6 +166,7 @@ impl DefaultPeer { handlers.register_handler(GetShareFileList); handlers.register_handler(UploadFailedHandler); handlers.register_handler(PlaceInQueueResponse); + handlers.register_handler(PeerInit); let dispatcher = MessageDispatcher::new( "default_peer".to_string(), @@ -167,19 +182,20 @@ impl DefaultPeer { continue } Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { - debug!( - "Read operation timed out in default peer {:}:{:}", - peer.host, peer.port - ); + { + let peer_lock = peer.read().unwrap(); + debug!( + "Read operation timed out in default peer {:}:{:}", + peer_lock.host, peer_lock.port + ); + } continue; } Err(e) => { - error!("[default_peer:{}] Error reading from peer: {}. Terminating read loop.",peer.username, e); - let _ = client_channel_for_read.send( - ClientOperation::PeerDisconnected( - peer.username.clone(), - ), - ); + let username = peer.read().unwrap().username.clone(); + error!("[default_peer:{}] Error reading from peer: {}. Terminating read loop.",username, e); + let _ = client_channel_for_read + .send(ClientOperation::PeerDisconnected(username)); break; } } @@ -191,7 +207,7 @@ impl DefaultPeer { Ok(Some(mut message)) => { trace!( "[default_peer:{}] ← {:?}", - peer.username, + peer.read().unwrap().username, message .get_message_name( MessageType::Peer, @@ -206,10 +222,10 @@ impl DefaultPeer { "Error extracting message in default peer: {}. Terminating read loop.", e ); + let username = + peer.read().unwrap().username.clone(); let _ = client_channel_for_read.send( - ClientOperation::PeerDisconnected( - peer.username.clone(), - ), + ClientOperation::PeerDisconnected(username), ); should_terminate = true; break; @@ -225,7 +241,7 @@ impl DefaultPeer { let client_channel = self.client_channel.clone(); let peer_channel = self.peer_channel.clone(); - let peer_username = self.peer.username.clone(); + let peer_clone_for_write = self.peer.clone(); self.write_thread = Some(thread::spawn(move || loop { match peer_reader.recv() { @@ -233,7 +249,7 @@ impl DefaultPeer { PeerOperation::SendMessage(message) => { trace!( "[default_peer:{}] ➡ {:?} - {:?}", - peer_username, + peer_clone_for_write.read().unwrap().username, message .get_message_name( MessageType::Peer, @@ -253,10 +269,14 @@ impl DefaultPeer { if let Err(e) = write_stream.write_all(&message.get_buffer()) { - error!("Error writing message to stream: {} - {}. Terminating write loop.", peer_username, e); + error!("Error writing message to stream: {} - {}. Terminating write loop.", peer_clone_for_write.read().unwrap().username, e); let _ = client_channel.send( ClientOperation::PeerDisconnected( - peer_username.clone(), + peer_clone_for_write + .read() + .unwrap() + .username + .clone(), ), ); break; @@ -271,12 +291,13 @@ impl DefaultPeer { PeerOperation::TransferRequest(transfer) => { debug!( "[default_peer:{}] TransferRequest for {}", - peer_username, transfer.token + peer_clone_for_write.read().unwrap().username, + transfer.token ); client_channel .send(ClientOperation::UpdateDownloadTokens( transfer.clone(), - peer_clone.username.clone(), + peer_clone.read().unwrap().username.clone(), )) .unwrap(); @@ -287,7 +308,7 @@ impl DefaultPeer { trace!( "[default_peer:{}] TransferResponse for {:?}", - peer_username, + peer_clone_for_write.read().unwrap().username, transfer_response.get_buffer() ); if let Some(sender) = peer_channel.clone() { @@ -304,38 +325,47 @@ impl DefaultPeer { reason, } => { debug!( - "[default_peer:{}] transfer response token: {} allowed: {}", - peer_username, token, allowed - ); + "[default_peer:{}] transfer response token: {} allowed: {}", + peer_clone_for_write.read().unwrap().username, token, allowed + ); if !allowed { if let Some(reason_text) = reason { debug!( - "[default_peer:{}] Transfer rejected: {} - token {}, I will receive TransferRequest soon...", - peer_username.clone(), - reason_text, - token - ); + "[default_peer:{}] Transfer rejected: {} - token {}, I will receive TransferRequest soon...", + peer_clone_for_write.read().unwrap().username.clone(), + reason_text, + token + ); } } else { - debug!("[default_peer:{}] Transfer allowed, ready to connect with token {:}",peer_username, token); + debug!("[default_peer:{}] Transfer allowed, ready to connect with token {:}",peer_clone_for_write.read().unwrap().username, token); client_channel .send(ClientOperation::DownloadFromPeer( token, - peer_clone.clone(), + peer_clone.read().unwrap().clone(), )) .unwrap(); } } PeerOperation::PlaceInQueueResponse { filename, place } => { debug!( - "[default_peer:{}] Place in queue response - file: {}, place: {}", - peer_username, filename, place + "[default_peer:{}] Place in queue response - file: {}, place: {}", + peer_clone_for_write.read().unwrap().username, filename, place + ); + } + PeerOperation::SetUsername(username) => { + trace!( + "[default_peer:{}] SetUsername: {}", + peer_clone_for_write.read().unwrap().username, + username ); + + peer_clone.write().unwrap().username = username; } }, Err(_) => { - debug!("[default_peer:{}] Peer channel closed. Terminating write loop.", peer_username); + debug!("[default_peer:{}] Peer channel closed. Terminating write loop.", peer_clone_for_write.read().unwrap().username); break; } } @@ -362,7 +392,7 @@ impl DefaultPeer { impl Drop for DefaultPeer { fn drop(&mut self) { self.peer_channel = None; - trace!("[default_peer:{}] drop", self.peer.username); + trace!("[default_peer:{}] drop", self.peer.read().unwrap().username); if let Some(handle) = self.read_thread.take() { debug!("Joining read thread..."); diff --git a/src/peer/download_peer.rs b/src/peer/download_peer.rs index d9ea86d..f4b7fca 100644 --- a/src/peer/download_peer.rs +++ b/src/peer/download_peer.rs @@ -217,16 +217,15 @@ impl DownloadPeer { stream.write_all(&message.get_buffer())?; stream.flush()?; } else { - stream.write_all( - &MessageFactory::build_pierce_firewall_message(self.token) - .get_buffer(), - )?; + let message = + MessageFactory::build_pierce_firewall_message(self.token); + + stream.write_all(&message.get_buffer())?; trace!( "[default_peer:{}] sending pierce firewall message token: {}: {:?}", self.username, self.token, - MessageFactory::build_pierce_firewall_message(self.token) - .get_buffer() + &message.get_buffer() ); stream.flush()?; } @@ -257,8 +256,9 @@ impl DownloadPeer { // connection closed Ok(0) => { trace!( - "[download_peer:{}] connection closed by peer", - self.username + "[download_peer:{}] connection closed by peer. bytes read: {}", + self.username, + processor.total_bytes ); break; } @@ -323,10 +323,13 @@ impl DownloadPeer { "[download_peer:{}] finished reading data from peer", self.username ); + let read = client_context.read().unwrap(); + let tokens = read.download_tokens.keys().collect::>(); + let download = download.clone().unwrap_or_else(|| { panic!( - "[download_peer] No download info found for token: {}", - self.token + "[download_peer] No download info found for token: {}, tokens: {:?} ", + self.token, tokens ) }); diff --git a/src/peer/listen.rs b/src/peer/listen.rs index e6aa222..ab21b20 100644 --- a/src/peer/listen.rs +++ b/src/peer/listen.rs @@ -1,112 +1,43 @@ -use std::io::Write; +use std::net::TcpListener; use std::sync::mpsc::Sender; -use std::{io, net::TcpListener}; +use std::thread; use crate::client::ClientOperation; -use crate::message::server::MessageFactory; -use crate::message::MessageReader; -use crate::{debug, error, info, trace, warn, FileSearchResult}; +use crate::peer::{DefaultPeer, Peer}; +use crate::{info, trace}; pub struct Listen {} impl Listen { pub fn start(port: u32, client_sender: Sender) { + let mut index = 0; + info!("starting listener on port {port}"); let listener = TcpListener::bind(format!("0.0.0.0:{port}")).unwrap(); for stream in listener.incoming() { - let mut stream = stream.unwrap(); + index += 1; + let client_sender_clone = client_sender.clone(); + thread::spawn(move || { + let stream = stream.unwrap(); - let mut buffered_reader = MessageReader::new(); - loop { - match buffered_reader.read_from_socket(&mut stream) { - Ok(()) => {} - Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { - continue - } - Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { - debug!("Read operation timed out"); - continue; - } - Err(e) => { - error!("Error reading from server: {}", e); - break; - } - } + let peer_ip = stream.peer_addr().unwrap().ip().to_string(); - loop { - match buffered_reader.extract_message() { - Ok(Some(mut message)) => { - message.set_pointer(4); - let message_code = message.read_int8(); + let peer = Peer::new( + format!("unknown-{index}"), + crate::peer::ConnectionType::P, + peer_ip.clone(), + port, + None, + 0, + 0, + 0, + ); + let default_peer = DefaultPeer::new(peer, client_sender_clone); - trace!( - "[listener] Received message with code: {}", - message_code - ); - - match message_code { - 0 => { - let token = message.read_int32(); - debug!("[listener] received Pierce Firewall, token: {}", token); - stream - .write_all( - &MessageFactory::build_pierce_firewall_message( - token, - ) - .get_data(), - ) - .unwrap(); - stream.flush().unwrap(); - } - 1 => { - // handle handover to default_peer with socket in the future - // let tcp_stream = - // read_stream.try_clone().unwrap(); - // - // let username = message.read_string(); - // let connection_type = - // message.read_string().parse().unwrap(); - // let token = message.read_int32(); - } - 9 => { - trace!("[listener] Handling file search response: 9",); - message.set_pointer(8); - let file_search = - match FileSearchResult::new_from_message( - &mut message, - ) { - Ok(result) => result, - Err(e) => { - trace!("[listener] malformed filesearch_result: {:?}, message: {:?}", e, message); - return; - } - }; - client_sender - .send(ClientOperation::SearchResult( - file_search, - )) - .unwrap(); - break; - } - _ => { - debug!( - "[listener] unknown message_code: {}", - message_code - ); - } - } - } - Err(e) => { - warn!("[listener] Error extracting message: {}", e); - break; - } - Ok(None) => break, - } - } - } - - // info!("[listener] Connection established!"); + trace!("[listener{index}] new connection from {peer_ip}"); + default_peer.connect_with_socket(stream).unwrap(); + }); } } }