From fe60ea67f80f28ea2f3ca6d3113e066bcfadd0e4 Mon Sep 17 00:00:00 2001 From: Michel de Graaf Date: Thu, 6 Nov 2025 12:18:28 +0100 Subject: [PATCH] chore: refactor listener --- src/peer/download_peer.rs | 105 +---------- src/peer/listen.rs | 388 ++++++++++++++++++++++++-------------- 2 files changed, 252 insertions(+), 241 deletions(-) diff --git a/src/peer/download_peer.rs b/src/peer/download_peer.rs index 803215e..9fcdd0b 100644 --- a/src/peer/download_peer.rs +++ b/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>>>) { - let messages = Arc::new(RwLock::new(Vec::>::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() { diff --git a/src/peer/listen.rs b/src/peer/listen.rs index 134ce11..82a2918 100644 --- a/src/peer/listen.rs +++ b/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, + client_context: Arc>, + 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 { + 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 { + 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 { + 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>, + username: &str, + peer_ip: &str, + peer_port: u16, +) -> Option { + 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 = + 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, + 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>, 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 = 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::>(); - 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)); } } }