chore: refactor default peer

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);

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"),

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;

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<PeerOperation> for PeerInit {
+    fn get_code(&self) -> u8 {
+        1
+    }
+
+    fn handle(&self, message: &mut Message, sender: Sender<PeerOperation>) {
+        message.set_pointer(4);
+        let _message_code = message.read_int8();
+        let username = message.read_string();
+        sender.send(PeerOperation::SetUsername(username)).unwrap();
+    }
+}

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<ServerOperation>`
 use crate::message::MessageHandler;
 
 pub struct LoginHandler;
 
-// Implement `MessageHandler` specifically for `ServerOperation`.
 impl MessageHandler<ServerOperation> for LoginHandler {
     fn get_code(&self) -> u8 {
         1

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<RwLock<Peer>>,
     peer_channel: Option<Sender<PeerOperation>>,
     client_channel: Sender<ClientOperation>,
     read_thread: Option<JoinHandle<()>>,
@@ -45,11 +45,13 @@ pub enum PeerOperation {
         filename: String,
         place: u32,
     },
+    SetUsername(String),
 }
+
 impl DefaultPeer {
     pub fn new(peer: Peer, client_channel: Sender<ClientOperation>) -> 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());
+        let username = self.peer.read().unwrap().username.clone();
-        if let Err(e) = self.client_channel.send(
+        debug!("[default_peer:{}] disconnect", username);
-            ClientOperation::PeerDisconnected(self.peer.username.clone()),
+        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<Self, io::Error> {
-        if let Some(token) = self.peer.token {
+        if let Some(token) = self.peer.read().unwrap().token {
             let mut message: Vec<u8> = [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<Self, io::Error> {
+        {
+            let peer = self.peer.read().unwrap();
             println!(
                 "[default_peer] Connecting to {} on port {}",
-            self.peer.host, self.peer.port
+                peer.host, peer.port
             );
-        let socket_address = format!("{}:{}", self.peer.host, self.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.username);
+        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<PeerOperation>,
@@ -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 => {
+                        {
+                            let peer_lock = peer.read().unwrap();
                             debug!(
                                 "Read operation timed out in default peer {:}:{:}",
-                            peer.host, peer.port
+                                peer_lock.host, peer_lock.port
                             );
+                        }
                         continue;
                     }
                     Err(e) => {
-                        error!("[default_peer:{}] Error reading from peer:  {}. Terminating read loop.",peer.username, e);
+                        let username = peer.read().unwrap().username.clone();
-                        let _ = client_channel_for_read.send(
+                        error!("[default_peer:{}] Error reading from peer:  {}. Terminating read loop.",username, e);
-                            ClientOperation::PeerDisconnected(
+                        let _ = client_channel_for_read
-                                peer.username.clone(),
+                            .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(
+                                ClientOperation::PeerDisconnected(username),
-                                    peer.username.clone(),
-                                ),
                             );
                             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() {
@@ -305,24 +326,24 @@ impl DefaultPeer {
                     } => {
                         debug!(
                                                     "[default_peer:{}] transfer response token: {} allowed: {}",
-                                    peer_username, 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(),
+                                                        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();
                         }
@@ -330,12 +351,21 @@ impl DefaultPeer {
                     PeerOperation::PlaceInQueueResponse { filename, place } => {
                         debug!(
                                             "[default_peer:{}] Place in queue response - file: {}, place: {}",
-                            peer_username, filename, 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...");

src/peer/download_peer.rs

@@ -217,16 +217,15 @@ impl DownloadPeer {
             stream.write_all(&message.get_buffer())?;
             stream.flush()?;
         } else {
-            stream.write_all(
+            let message =
-                &MessageFactory::build_pierce_firewall_message(self.token)
+                MessageFactory::build_pierce_firewall_message(self.token);
-                    .get_buffer(),
+
-            )?;
+            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)
+                &message.get_buffer()
-                    .get_buffer()
             );
             stream.flush()?;
         }
@@ -257,8 +256,9 @@ impl DownloadPeer {
                 // connection closed
                 Ok(0) => {
                     trace!(
-                        "[download_peer:{}] connection closed by peer",
+                        "[download_peer:{}] connection closed by peer. bytes read: {}",
-                        self.username
+                        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::<Vec<_>>();
+
         let download = download.clone().unwrap_or_else(|| {
             panic!(
-                "[download_peer] No download info found for token: {}",
+                "[download_peer] No download info found for token: {}, tokens: {:?} ",
-                self.token
+                self.token, tokens
             )
         });
 

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<ClientOperation>) {
+        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();
+                let peer_ip = stream.peer_addr().unwrap().ip().to_string();
-            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;
-                    }
-                }
 
-                loop {
+                let peer = Peer::new(
-                    match buffered_reader.extract_message() {
+                    format!("unknown-{index}"),
-                        Ok(Some(mut message)) => {
+                    crate::peer::ConnectionType::P,
-                            message.set_pointer(4);
+                    peer_ip.clone(),
-                            let message_code = message.read_int8();
+                    port,
-
+                    None,
-                            trace!(
+                    0,
-                                "[listener] Received message with code: {}",
+                    0,
-                                message_code
+                    0,
                 );
+                let default_peer = DefaultPeer::new(peer, client_sender_clone);
 
-                            match message_code {
+                trace!("[listener{index}] new connection from {peer_ip}");
-                                0 => {
+                default_peer.connect_with_socket(stream).unwrap();
-                                    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!");
         }
     }
 }