diff --git a/src/client.rs b/src/client.rs index 6674c53..cc85f0c 100644 --- a/src/client.rs +++ b/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 || { - Listen::start(listen_port, client_sender.clone()); - }); + // thread::spawn(move || { + // 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)) => { diff --git a/src/message/message_reader.rs b/src/message/message_reader.rs index 95c77d6..8925b11 100644 --- a/src/message/message_reader.rs +++ b/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> { let bytes_read = self.buffer.len(); if bytes_read < 4 { diff --git a/src/message/peer/peer_init.rs b/src/message/peer/peer_init.rs index 882e95c..5d37f77 100644 --- a/src/message/peer/peer_init.rs +++ b/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 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(); } } diff --git a/src/peer/default_peer.rs b/src/peer/default_peer.rs index af69191..58ea4b4 100644 --- a/src/peer/default_peer.rs +++ b/src/peer/default_peer.rs @@ -74,6 +74,7 @@ impl DefaultPeer { pub fn connect_with_socket( mut self, mut stream: TcpStream, + reader: Option, ) -> Result { if let Some(token) = self.peer.read().unwrap().token { let mut message: Vec = [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, ) -> Result<(), io::Error> { debug!( "[default_peer:{}] start_read_write_loops", @@ -148,17 +150,38 @@ impl DefaultPeer { Receiver, ) = 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 = stream; // Use the original stream for writing + let mut write_stream = match stream.try_clone() { + 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 + ); + + if buffer_size > 0 { + trace!( + "[default_peer:{}] starting with {} buffered bytes", + peer.read().unwrap().username, + buffer_size + ); + } - let mut buffered_reader = MessageReader::new(); loop { - match buffered_reader.read_from_socket(&mut read_stream) { - Ok(()) => {} - 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: {}. Terminating read loop.",username, e); - let _ = client_channel_for_read - .send(ClientOperation::PeerDisconnected(username)); - break; - } - } - - // Extract all available messages from buffer 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) => { + 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(); @@ -325,18 +366,18 @@ impl DefaultPeer { reason, } => { debug!( - "[default_peer:{}] transfer response token: {} allowed: {}", - peer_clone_for_write.read().unwrap().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_clone_for_write.read().unwrap().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_clone_for_write.read().unwrap().username, token); @@ -350,9 +391,9 @@ impl DefaultPeer { } PeerOperation::PlaceInQueueResponse { filename, place } => { debug!( - "[default_peer:{}] Place in queue response - file: {}, place: {}", - peer_clone_for_write.read().unwrap().username, filename, place - ); + "[default_peer:{}] Place in queue response - file: {}, place: {}", + peer_clone_for_write.read().unwrap().username, filename, place + ); } PeerOperation::SetUsername(username) => { trace!( @@ -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, diff --git a/src/peer/download_peer.rs b/src/peer/download_peer.rs index f4b7fca..3ab6522 100644 --- a/src/peer/download_peer.rs +++ b/src/peer/download_peer.rs @@ -236,8 +236,17 @@ impl DownloadPeer { self, client_context: Arc>, mut download: Option, + stream: Option, ) -> 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 = - u32::from_le_bytes(token.try_into().unwrap()); + let token_u32 = u32::from_le_bytes( + 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::>(); 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"); - } } diff --git a/src/peer/listen.rs b/src/peer/listen.rs index ab21b20..3a8f6e0 100644 --- a/src/peer/listen.rs +++ b/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::{info, trace}; +use crate::message::MessageReader; +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) { - let mut index = 0; - + pub fn start( + port: u32, + client_sender: Sender, + client_context: Arc>, + 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_clone = client_sender.clone(); + let client_sender = 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 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); + let mut reader = MessageReader::new(); - trace!("[listener{index}] new connection from {peer_ip}"); - default_peer.connect_with_socket(stream).unwrap(); + 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!("{}-direct", 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" { + 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}"); + } }); } }