chore: refactor listener and default peer

src/client.rs

@@ -118,6 +118,8 @@ impl Client {
 
         let listen_port = self.listen_port;
         self.context.write().unwrap().sender = Some(sender.clone());
+        let context = self.context.clone();
+        let own_username = self.username.clone();
 
         let client_sender = sender.clone();
 
@@ -130,9 +132,14 @@ impl Client {
                         server.get_address().get_port()
                     );
 
-                    thread::spawn(move || {
+                    // thread::spawn(move || {
-                        Listen::start(listen_port, client_sender.clone());
+                    //     Listen::start(
-                    });
+                    //         listen_port,
+                    //         client_sender.clone(),
+                    //         context.clone(),
+                    //         own_username,
+                    //     );
+                    // });
                     let mut unlocked_context = self.context.write().unwrap();
                     unlocked_context.server_sender =
                         Some(server.get_sender().clone());
@@ -229,7 +236,7 @@ impl Client {
         let download_initiated = context
             .peers
             .get(&username)
-            .map(|p| p.transfer_request(download.clone()))
+            .map(|p| p.queue_upload(download.filename))
             .is_some();
 
         drop(context);
@@ -303,7 +310,7 @@ impl Client {
                             token,
                             peer
                         );
-                        match maybe_download {
+                        return match maybe_download {
                             Some(download) => {
                                 thread::spawn(move || {
                                     let download_peer = DownloadPeer::new(
@@ -322,7 +329,8 @@ impl Client {
                                                         Some(filename) => {
                                                             match download_peer.download_file(
                                                                 client_context_clone,
-                                                                Some(download.clone())
+                                                                Some(download.clone()),
+                                                                None
                                                             ) {
                                                                 Ok((download, filename)) => {
                                                                     download.sender.send(DownloadStatus::Completed).unwrap();
@@ -347,7 +355,7 @@ impl Client {
                                     token
                                 );
                             }
-                        }
+                        };
                     }
                     ClientOperation::NewPeer(new_peer) => {
                         if client_context
@@ -507,7 +515,7 @@ impl Client {
                                 DefaultPeer::new(peer_clone, sender_clone);
 
                             let connect_result = match stream {
-                                    Some(s) => default_peer.connect_with_socket(s),
+                                    Some(s) => default_peer.connect_with_socket(s, None),
                                     _ => default_peer.connect()
                                 };
 
@@ -543,6 +551,7 @@ impl Client {
                                     match download_peer
                                     .download_file(
                                         client_context2.clone(),
+                                        None,
                                         None
                                     ) {
                                         Ok((download, filename)) => {

src/message/message_reader.rs

@@ -63,6 +63,10 @@ impl MessageReader {
         Ok(())
     }
 
+    pub fn buffer_len(&self) -> usize {
+        self.buffer.len()
+    }
+
     pub fn extract_message(&mut self) -> io::Result<Option<Message>> {
         let bytes_read = self.buffer.len();
         if bytes_read < 4 {

src/message/peer/peer_init.rs

@@ -3,6 +3,7 @@ use std::sync::mpsc::Sender;
 use crate::{
     message::{Message, MessageHandler},
     peer::PeerOperation,
+    trace,
 };
 
 pub struct PeerInit;
@@ -15,6 +16,15 @@ impl MessageHandler<PeerOperation> for PeerInit {
         message.set_pointer(4);
         let _message_code = message.read_int8();
         let username = message.read_string();
+        let connection_type = message.read_string();
+        let token = message.read_int32();
+        trace!(
+            "PeerInit: username: {}, connection_type: {}, token: {}",
+            username,
+            connection_type,
+            token
+        );
+
         sender.send(PeerOperation::SetUsername(username)).unwrap();
     }
 }

src/peer/default_peer.rs

@@ -74,6 +74,7 @@ impl DefaultPeer {
     pub fn connect_with_socket(
         mut self,
         mut stream: TcpStream,
+        reader: Option<MessageReader>,
     ) -> Result<Self, io::Error> {
         if let Some(token) = self.peer.read().unwrap().token {
             let mut message: Vec<u8> = [0, 5, 0, 0, 0, 0].to_vec();
@@ -89,7 +90,7 @@ impl DefaultPeer {
             "[default_peer:{}] connect_with_socket: direct",
             self.peer.read().unwrap().username
         );
-        self.start_read_write_loops(stream)?;
+        self.start_read_write_loops(stream, reader)?;
 
         Ok(self)
     }
@@ -130,7 +131,7 @@ impl DefaultPeer {
                 .write_all(&MessageFactory::build_watch_user(token).get_data())
                 .unwrap();
         }
-        self.start_read_write_loops(stream)?;
+        self.start_read_write_loops(stream, None)?;
 
         Ok(self)
     }
@@ -138,6 +139,7 @@ impl DefaultPeer {
     fn start_read_write_loops(
         &mut self,
         stream: TcpStream,
+        reader: Option<MessageReader>,
     ) -> Result<(), io::Error> {
         debug!(
             "[default_peer:{}] start_read_write_loops",
@@ -148,17 +150,38 @@ impl DefaultPeer {
             Receiver<PeerOperation>,
         ) = mpsc::channel();
 
-        // Set the peer_channel so transfer_request can send messages
         self.peer_channel = Some(peer_sender.clone());
 
-        let mut read_stream = stream.try_clone()?;
+        let mut write_stream = match stream.try_clone() {
-        let mut write_stream = stream; // Use the original stream for writing
+            Ok(s) => s,
+            Err(e) => {
+                error!(
+                    "[default_peer:{}] Failed to clone stream: {}",
+                    self.peer.read().unwrap().username,
+                    e
+                );
+                return Err(e);
+            }
+        };
+        let mut read_stream = stream;
+
+        debug!(
+            "[default_peer:{}] Streams created, spawning threads",
+            self.peer.read().unwrap().username
+        );
 
         let peer = self.peer.clone();
         let peer_clone = self.peer.clone();
         let client_channel_for_read = self.client_channel.clone();
 
+        let mut buffered_reader = reader.unwrap_or_default();
+        let buffer_size = buffered_reader.buffer_len();
+
         self.read_thread = Some(thread::spawn(move || {
+            debug!(
+                "Read thread received reader with {} bytes in buffer",
+                buffer_size
+            );
             let mut handlers = Handlers::new();
             handlers.register_handler(FileSearchResponse);
             handlers.register_handler(TransferRequest);
@@ -173,34 +196,20 @@ impl DefaultPeer {
                 peer_sender,
                 handlers,
             );
+            trace!(
+                "[default_peer:{}] read thread starting",
+                peer.read().unwrap().username
+            );
 
-            let mut buffered_reader = MessageReader::new();
+            if buffer_size > 0 {
-            loop {
+                trace!(
-                match buffered_reader.read_from_socket(&mut read_stream) {
+                    "[default_peer:{}] starting with {} buffered bytes",
-                    Ok(()) => {}
+                    peer.read().unwrap().username,
-                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
+                    buffer_size
-                        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_lock.host, peer_lock.port
                 );
             }
-                        continue;
-                    }
-                    Err(e) => {
-                        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;
-                    }
-                }
 
-                // Extract all available messages from buffer
+            loop {
                 let mut should_terminate = false;
                 loop {
                     match buffered_reader.extract_message() {
@@ -230,13 +239,45 @@ impl DefaultPeer {
                             should_terminate = true;
                             break;
                         }
-                        Ok(_) => break,
+                        Ok(None) => {
-                    }
-                }
-                if should_terminate {
                             break;
                         }
                     }
+                }
+
+                if should_terminate {
+                    break;
+                }
+
+                match buffered_reader.read_from_socket(&mut read_stream) {
+                    Ok(()) => continue,
+                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
+                        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_lock.host, peer_lock.port
+                            );
+                        }
+                        continue;
+                    }
+                    Err(e) => {
+                        let username = peer.read().unwrap().username.clone();
+                        error!("[default_peer:{}] Error reading from peer: {} (kind: {:?}). Terminating read loop.", username, e, e.kind());
+                        let _ = client_channel_for_read
+                            .send(ClientOperation::PeerDisconnected(username));
+                        break;
+                    }
+                }
+            }
+            trace!(
+                "[default_peer:{}] read thread exiting",
+                peer.read().unwrap().username
+            );
         }));
 
         let client_channel = self.client_channel.clone();
@@ -374,6 +415,13 @@ impl DefaultPeer {
         Ok(())
     }
 
+    pub fn queue_upload(&self, filename: String) -> Result<(), io::Error> {
+        let message = MessageFactory::build_queue_upload_message(&filename);
+        if let Some(sender) = &self.peer_channel {
+            sender.send(PeerOperation::SendMessage(message)).unwrap();
+        }
+        Ok(())
+    }
     pub fn transfer_request(
         &self,
         download: Download,

src/peer/download_peer.rs

@@ -236,8 +236,17 @@ impl DownloadPeer {
         self,
         client_context: Arc<RwLock<ClientContext>>,
         mut download: Option<Download>,
+        stream: Option<TcpStream>,
     ) -> Result<(Download, String), io::Error> {
-        let mut stream = self.establish_connection()?;
+        if stream.is_none() {
+            trace!("[download_peer:{}] stream is none", self.username);
+        }
+        let mut stream = stream.unwrap_or_else(|| {
+            self.establish_connection().unwrap_or_else(|e| {
+                panic!("Failed to establish connection: {:?}", e)
+            })
+        });
+
         trace!("[download_peer:{}] connected", self.username);
 
         self.perform_handshake(&mut stream)?;
@@ -270,11 +279,15 @@ impl DownloadPeer {
                         && processor.handle_pierce_token(data, &mut stream)?
                     {
                         let token = data.get(0..4).unwrap();
-                        let token_u32 =
+                        let token_u32 = u32::from_le_bytes(
-                            u32::from_le_bytes(token.try_into().unwrap());
+                            token
+                                .try_into()
+                                .unwrap_or_else(|_| panic!("[download_peer:{}] slice with incorrect length", self.username)),
+                        );
                         trace!(
-                            "[download_peer:{}] received token: {:?} ",
+                            "[download_peer:{}] received token_u32: {:?}  - token: {:?}",
                             self.username,
+                            token_u32,
                             token
                         );
 
@@ -297,8 +310,12 @@ impl DownloadPeer {
                                 download = Some(d.clone());
                             }
                             None => {
+                                let tokens = read
+                                    .download_tokens
+                                    .keys()
+                                    .collect::<Vec<_>>();
                                 panic!(
-                                    "No download info for token {token_u32}"
+                                    "[download_peer:{}] No download info for token {token_u32}, tokens: {:?}", self.username, tokens
                                 );
                             }
                         }
@@ -427,7 +444,9 @@ mod tests {
             "own_username".to_string(),
         );
         let dummy_context = Arc::new(RwLock::new(ClientContext::new()));
-        let _ = download_peer.download_file(dummy_context).unwrap();
+        let _ = download_peer
+            .download_file(dummy_context, None, None)
+            .unwrap();
 
         // Give the client time to send messages
         thread::sleep(Duration::from_millis(10));
@@ -450,7 +469,9 @@ mod tests {
             "own_username".to_string(),
         );
         let dummy_context = Arc::new(RwLock::new(ClientContext::new()));
-        let _ = download_peer.download_file(dummy_context).unwrap();
+        let _ = download_peer
+            .download_file(dummy_context, None, None)
+            .unwrap();
 
         // Give the client time to send messages
         thread::sleep(Duration::from_millis(10));
@@ -611,43 +632,4 @@ mod tests {
             "file.mp3"
         );
     }
-
-    #[test]
-    fn test_create_download_path_from_filename() {
-        // Test with filename provided
-        let path = FileManager::create_download_path_from_filename(
-            Some("/downloads"),
-            "remote_user",
-            123,
-            Some("song.mp3"),
-        );
-        assert_eq!(path, "/downloads/song.mp3");
-
-        // Test with full path filename
-        let path = FileManager::create_download_path_from_filename(
-            Some("/downloads"),
-            "remote_user",
-            123,
-            Some("/remote/path/to/song.mp3"),
-        );
-        assert_eq!(path, "/downloads/song.mp3");
-
-        // Test without filename (fallback)
-        let path = FileManager::create_download_path_from_filename(
-            Some("/downloads"),
-            "remote_user",
-            123,
-            None,
-        );
-        assert_eq!(path, "/downloads/remote_user_123.mp3");
-
-        // Test with default directory
-        let path = FileManager::create_download_path_from_filename(
-            None,
-            "remote_user",
-            123,
-            Some("song.mp3"),
-        );
-        assert_eq!(path, "/tmp/song.mp3");
-    }
 }

src/peer/listen.rs

@@ -1,42 +1,134 @@
 use std::net::TcpListener;
 use std::sync::mpsc::Sender;
+use std::sync::{Arc, RwLock};
 use std::thread;
 
-use crate::client::ClientOperation;
+use crate::client::{ClientContext, ClientOperation};
 
-use crate::peer::{DefaultPeer, Peer};
+use crate::message::MessageReader;
-use crate::{info, trace};
+use crate::peer::{ConnectionType, DefaultPeer, DownloadPeer, Peer};
+use crate::{debug, error, info, trace, DownloadStatus};
 
 pub struct Listen {}
 
 impl Listen {
-    pub fn start(port: u32, client_sender: Sender<ClientOperation>) {
+    pub fn start(
-        let mut index = 0;
+        port: u32,
-
+        client_sender: Sender<ClientOperation>,
+        client_context: Arc<RwLock<ClientContext>>,
+        own_username: String,
+    ) {
         info!("starting listener on port {port}");
         let listener = TcpListener::bind(format!("0.0.0.0:{port}")).unwrap();
         for stream in listener.incoming() {
-            index += 1;
+            let client_sender = client_sender.clone();
-            let client_sender_clone = client_sender.clone();
+            let client_context = client_context.clone();
+            let own_username = own_username.clone();
             thread::spawn(move || {
-                let stream = stream.unwrap();
+                let mut stream = match stream {
-
+                    Ok(s) => s,
+                    Err(e) => {
+                        trace!("Failed to accept connection: {}", e);
+                        return;
+                    }
+                };
                 let peer_ip = stream.peer_addr().unwrap().ip().to_string();
+                let peer_port = stream.peer_addr().unwrap().port();
+
+                let mut reader = MessageReader::new();
+
+                let mut message = loop {
+                    if let Err(e) = reader.read_from_socket(&mut stream) {
+                        trace!("Failed to read from socket: {}", e);
+                        return;
+                    }
+
+                    if let Ok(Some(msg)) = reader.extract_message() {
+                        break msg;
+                    }
+                };
+
+                message.set_pointer(4);
+                let message_code = message.read_int8();
+
+                if message_code == 1 {
+                    let username = message.read_string();
+                    let connection_type = message.read_string();
+                    let token = message.read_int32();
+                    debug!("[listener:{peer_ip}:{peer_port}] peerInit (0)  username: {username} connection_type: {connection_type} token: {token}");
 
                     let peer = Peer::new(
-                    format!("unknown-{index}"),
+                        format!("{}-direct", username.clone()),
-                    crate::peer::ConnectionType::P,
+                        ConnectionType::P,
                         peer_ip.clone(),
-                    port,
+                        peer_port.into(),
                         None,
                         0,
                         0,
                         0,
                     );
-                let default_peer = DefaultPeer::new(peer, client_sender_clone);
 
-                trace!("[listener{index}] new connection from {peer_ip}");
+                    if connection_type == "P" {
-                default_peer.connect_with_socket(stream).unwrap();
+                        debug!("[listener:{peer_ip}:{peer_port}] connection type is P, reader buffer has {} bytes", reader.buffer_len());
+                        let default_peer =
+                            DefaultPeer::new(peer, client_sender);
+                        let default_peer = default_peer
+                            .connect_with_socket(stream, Some(reader))
+                            .unwrap();
+
+                        drop(default_peer);
+                    } else if connection_type == "F" {
+                        let maybe_download = {
+                            let client_context = client_context.read().unwrap();
+                            client_context.download_tokens.get(&token).cloned()
+                        };
+
+                        let client_context_clone = client_context.clone();
+
+                        trace!(
+                            "[client] DownloadFromPeer token: {} peer: {:?}",
+                            token,
+                            peer
+                        );
+
+                        thread::spawn(move || {
+                            let download_peer = DownloadPeer::new(
+                                format!("{}-direct", username.clone()),
+                                peer.host.clone(),
+                                peer.port,
+                                token,
+                                false,
+                                own_username,
+                            );
+
+                            match download_peer.download_file(
+                                client_context_clone,
+                                maybe_download,
+                                Some(stream),
+                            ) {
+                                Ok((download, filename)) => {
+                                    download
+                                        .sender
+                                        .send(DownloadStatus::Completed)
+                                        .unwrap();
+                                    info!("Successfully downloaded {} bytes to {}", download.size, filename);
+                                }
+                                Err(e) => {
+                                    error!(
+                                        "Failed to download file from {}:{} (token: {}) - Error: {}", 
+                                        peer.host, peer.port, token, e
+                                    );
+                                }
+                            }
+                        });
+                    } else {
+                        debug!(
+                            "[listener:{peer_ip}:{peer_port}] connection type is not P or F"
+                        );
+                    }
+                } else {
+                    debug!("[listener:{peer_ip}:{peer_port}] unknown message with code: {message_code}");
+                }
             });
         }
     }