chore: refactor listener

src/peer/download_peer.rs

@@ -1,7 +1,3 @@
-use crate::client::ClientContext;
-use crate::message::server::MessageFactory;
-use crate::trace;
-use crate::types::Download;
 use std::env;
 use std::fs::{self, File};
 use std::io::{self, BufWriter, Read, Write};
@@ -11,6 +7,11 @@ use std::path::Path;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 
+use crate::client::ClientContext;
+use crate::message::server::MessageFactory;
+use crate::trace;
+use crate::types::Download;
+
 #[derive(Debug)]
 #[allow(dead_code)]
 struct DownloadPaths {
@@ -416,103 +417,7 @@ impl DownloadPeer {
 #[cfg(test)]
 mod tests {
     use super::{DownloadPeer, FileManager, StreamProcessor};
-    use crate::client::ClientContext;
     use std::fs;
-    use std::io::Read;
-    use std::net::TcpListener;
-    use std::sync::{Arc, RwLock};
-    use std::thread;
-    use std::time::Duration;
-
-    pub fn build_test_server() -> (u16, Arc<RwLock<Vec<Vec<u8>>>>) {
-        let messages = Arc::new(RwLock::new(Vec::<Vec<u8>>::new()));
-        let messages_clone = Arc::clone(&messages);
-
-        // Find an available port by binding to port 0
-        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
-        let port = listener.local_addr().unwrap().port();
-
-        thread::spawn(move || {
-            for stream in listener.incoming() {
-                match stream {
-                    Ok(mut stream) => {
-                        let mut buffer = [0u8; 1024];
-                        match stream.read(&mut buffer) {
-                            Ok(bytes_read) if bytes_read > 0 => {
-                                let data = buffer[..bytes_read].to_vec();
-                                messages_clone.write().unwrap().push(data);
-                            }
-                            _ => {}
-                        }
-                    }
-                    Err(_) => break,
-                }
-            }
-        });
-
-        // Give the server time to start
-        thread::sleep(Duration::from_millis(10));
-
-        (port, messages)
-    }
-
-    #[test]
-    pub fn test_connect_no_pierce() {
-        let (port, messages) = build_test_server();
-        let token = 33;
-        let download_peer = DownloadPeer::new(
-            "test_user".to_string(),
-            "127.0.0.1".to_string(),
-            port as u32,
-            token,
-            false,
-            "own_username".to_string(),
-        );
-        let dummy_context = Arc::new(RwLock::new(ClientContext::new()));
-        let _ = download_peer
-            .download_file(dummy_context, None, None)
-            .unwrap();
-
-        // Give the client time to send messages
-        thread::sleep(Duration::from_millis(10));
-
-        let received_messages = messages.read().unwrap();
-        assert_eq!(received_messages.len(), 1);
-        assert_eq!(received_messages[0], vec![0, 0, 0, 0, 33, 0, 0, 0]);
-    }
-
-    #[test]
-    pub fn test_connect_with_pierce() {
-        let (port, messages) = build_test_server();
-        let token = 33;
-        let download_peer = DownloadPeer::new(
-            "test_user".to_string(),
-            "127.0.0.1".to_string(),
-            port as u32,
-            token,
-            true, // no_pierce = true, should send init message
-            "own_username".to_string(),
-        );
-        let dummy_context = Arc::new(RwLock::new(ClientContext::new()));
-        let _ = download_peer
-            .download_file(dummy_context, None, None)
-            .unwrap();
-
-        // Give the client time to send messages
-        thread::sleep(Duration::from_millis(10));
-
-        let received_messages = messages.read().unwrap();
-        assert_eq!(received_messages.len(), 1); // Expect one init message
-        assert_eq!(
-            received_messages[0],
-            vec![
-                1, 12, 0, 0, 0, 111, 119, 110, 95, 117, 115, 101, 114, 110, 97,
-                109, 101, 1, 0, 0, 0, 70, 33, 0, 0, 0
-            ]
-        );
-
-        // No messages expected when no_pierce is false
-    }
 
     #[test]
     fn test_file_manager_create_paths_with_custom_path() {

src/peer/listen.rs

@@ -1,15 +1,240 @@
-use std::net::TcpListener;
+use std::io;
+use std::net::{TcpListener, TcpStream};
 use std::sync::mpsc::Sender;
 use std::sync::{Arc, RwLock};
 use std::thread;
 
 use crate::client::{ClientContext, ClientOperation};
 
-use crate::message::MessageReader;
+use crate::message::{Message, MessageReader};
 use crate::peer::{ConnectionType, DefaultPeer, DownloadPeer, Peer};
 use crate::types::Download;
 use crate::{debug, error, info, trace, DownloadStatus};
 
+const PEER_INIT_MESSAGE_CODE: u8 = 1;
+
+#[derive(Clone)]
+struct ConnectionContext {
+    client_sender: Sender<ClientOperation>,
+    client_context: Arc<RwLock<ClientContext>>,
+    own_username: String,
+}
+
+struct PeerInitData {
+    username: String,
+    connection_type: ConnectionType,
+    token: u32,
+}
+
+fn read_peer_init_message(
+    stream: &mut TcpStream,
+    reader: &mut MessageReader,
+) -> io::Result<Message> {
+    loop {
+        reader.read_from_socket(stream)?;
+
+        if let Ok(Some(msg)) = reader.extract_message() {
+            return Ok(msg);
+        }
+    }
+}
+
+fn parse_peer_init_message(mut message: Message) -> Option<PeerInitData> {
+    message.set_pointer(4);
+    let message_code = message.read_int8();
+
+    if message_code != PEER_INIT_MESSAGE_CODE {
+        return None;
+    }
+
+    Some(PeerInitData {
+        username: message.read_string(),
+        connection_type: message.read_string().parse().unwrap(),
+        token: message.read_int32(),
+    })
+}
+
+fn parse_token_from_buffer(buffer: &[u8], username: &str) -> Option<u32> {
+    let token_bytes = buffer.get(0..4)?;
+    let token = u32::from_le_bytes(
+        token_bytes
+            .try_into()
+            .unwrap_or_else(|_| panic!("[listener:{}] slice with incorrect length, can't extract transfer_token", username)),
+    );
+    Some(token)
+}
+
+fn extract_download_from_buffer(
+    reader: &mut MessageReader,
+    client_context: &Arc<RwLock<ClientContext>>,
+    username: &str,
+    peer_ip: &str,
+    peer_port: u16,
+) -> Option<Download> {
+    if reader.buffer_len() == 0 {
+        return None;
+    }
+
+    let buffer = reader.get_buffer();
+    trace!(
+        "[listener:{peer_ip}:{peer_port}] reader buffer has {} bytes. {:?}",
+        buffer.len(),
+        buffer
+    );
+
+    let token = parse_token_from_buffer(&buffer, username)?;
+    trace!(
+        "[listener:{}] got transfer_token: {} from data chunk",
+        username,
+        token
+    );
+
+    let context = client_context.read().unwrap();
+    let download = context.download_tokens.get(&token).cloned();
+
+    if download.is_none() {
+        let download_tokens: Vec<u32> =
+            context.download_tokens.keys().cloned().collect();
+        trace!("[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}", token, download_tokens);
+    }
+
+    download
+}
+
+fn handle_peer_connection(
+    peer: Peer,
+    stream: TcpStream,
+    reader: MessageReader,
+    client_sender: Sender<ClientOperation>,
+    peer_ip: &str,
+    peer_port: u16,
+) {
+    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);
+    if let Ok(default_peer) =
+        default_peer.connect_with_socket(stream, Some(reader))
+    {
+        drop(default_peer);
+    }
+}
+
+fn handle_file_connection(
+    peer: Peer,
+    stream: TcpStream,
+    mut reader: MessageReader,
+    token: u32,
+    context: &ConnectionContext,
+    peer_ip: &str,
+    peer_port: u16,
+) {
+    trace!(
+        "[client] DownloadFromPeer token: {} peer: {:?}",
+        token,
+        peer
+    );
+
+    let download = extract_download_from_buffer(
+        &mut reader,
+        &context.client_context,
+        &peer.username,
+        peer_ip,
+        peer_port,
+    );
+
+    let download_peer = DownloadPeer::new(
+        format!("{}:direct", peer.username),
+        peer.host.clone(),
+        peer.port,
+        token,
+        true,
+        context.own_username.clone(),
+    );
+
+    match download_peer.download_file(
+        context.client_context.clone(),
+        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
+            );
+        }
+    }
+}
+
+fn handle_incoming_connection(stream: TcpStream, context: ConnectionContext) {
+    let Ok(peer_addr) = stream.peer_addr() else {
+        error!("[listener] failed to get peer address");
+        return;
+    };
+
+    let peer_ip = peer_addr.ip().to_string();
+    let peer_port = peer_addr.port();
+    let mut stream = stream;
+    let mut reader = MessageReader::new();
+
+    let Ok(message) = read_peer_init_message(&mut stream, &mut reader) else {
+        error!(
+            "[listener:{peer_ip}:{peer_port}] Failed to read peer init message"
+        );
+        return;
+    };
+
+    let Some(init_data) = parse_peer_init_message(message) else {
+        error!("[listener:{peer_ip}:{peer_port}] Invalid or unknown peer init message");
+        return;
+    };
+
+    debug!(
+        "[listener:{peer_ip}:{peer_port}] peerInit (0)  username: {} connection_type: {} token: {}",
+        init_data.username, init_data.connection_type, init_data.token
+    );
+
+    let peer = Peer::new(
+        init_data.username.clone(),
+        ConnectionType::P,
+        peer_ip.clone(),
+        peer_port.into(),
+        None,
+        0,
+        0,
+        0,
+    );
+
+    match init_data.connection_type {
+        ConnectionType::P => handle_peer_connection(
+            peer,
+            stream,
+            reader,
+            context.client_sender,
+            &peer_ip,
+            peer_port,
+        ),
+        ConnectionType::F => handle_file_connection(
+            peer,
+            stream,
+            reader,
+            init_data.token,
+            &context,
+            &peer_ip,
+            peer_port,
+        ),
+        ConnectionType::D => {
+            debug!("[listener:{peer_ip}:{peer_port}] connection type is D, not supported yet, closing connection. ");
+        }
+    }
+}
+
 pub struct Listen {}
 
 impl Listen {
@@ -19,147 +244,28 @@ impl Listen {
         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();
+        info!("[listener] starting listener on port {port}");
+
+        let listener = TcpListener::bind(format!("0.0.0.0:{port}"))
+            .expect("Failed to bind listener to port");
+
+        let context = ConnectionContext {
+            client_sender,
+            client_context,
+            own_username,
+        };
+
         for stream in listener.incoming() {
-            let client_sender = client_sender.clone();
-            let client_context = client_context.clone();
-            let own_username = own_username.clone();
-            thread::spawn(move || {
-                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 Ok(stream) = stream else {
+                error!(
+                    "[listener] Failed to accept connection: {}",
+                    stream.unwrap_err()
+                );
+                continue;
+            };
 
-                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(
-                        username.clone(),
-                        ConnectionType::P,
-                        peer_ip.clone(),
-                        peer_port.into(),
-                        None,
-                        0,
-                        0,
-                        0,
-                    );
-
-                    if connection_type == "P" {
-                        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" {
-                        trace!(
-                            "[client] DownloadFromPeer token: {} peer: {:?}",
-                            token,
-                            peer
-                        );
-
-                        let mut download: Option<Download> = None;
-                        if reader.buffer_len() > 0 {
-                            let buffer = reader.get_buffer();
-                            trace!("[listener:{peer_ip}:{peer_port}] reader buffer has {} bytes. {:?}", buffer.len(), buffer);
-
-                            let token = buffer.get(0..4).unwrap();
-                            let token_u32 = u32::from_le_bytes(
-                            token
-                                .try_into()
-                                .unwrap_or_else(|_| panic!("[listener:{}] slice with incorrect length, can't extract transfer_token", username)),
-                        );
-                            trace!(
-                            "[listener:{}] got transfer_token: {} from data chunk",
-                            username,
-                            token_u32);
-
-                            download = client_context
-                                .read()
-                                .unwrap()
-                                .download_tokens
-                                .get(&token_u32)
-                                .cloned();
-
-                            if download.is_none() {
-                                let download_tokens = client_context
-                                    .read()
-                                    .unwrap()
-                                    .download_tokens
-                                    .keys()
-                                    .cloned()
-                                    .collect::<Vec<u32>>();
-                                trace!("[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}", token_u32, download_tokens );
-                            }
-                        }
-
-                        let download_peer = DownloadPeer::new(
-                            format!("{}:direct", username.clone()),
-                            peer.host.clone(),
-                            peer.port,
-                            token,
-                            true,
-                            own_username,
-                        );
-
-                        match download_peer.download_file(
-                            client_context,
-                            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}");
-                }
-            });
+            let context = context.clone();
+            thread::spawn(move || handle_incoming_connection(stream, context));
         }
     }
 }