From 13f4783f65c35ce58f3f291e9a65a526b06b829b Mon Sep 17 00:00:00 2001 From: Michel de Graaf Date: Mon, 10 Nov 2025 21:12:39 +0100 Subject: [PATCH] chore: refactor peer connection to actor and non blocking --- src/actor/listener_actor.rs | 344 ------------------- src/actor/peer_actor.rs | 262 +++++++++++---- src/actor/peer_registry.rs | 21 +- src/client.rs | 638 ++++++++++++++++-------------------- src/main.rs | 2 +- src/peer/listen.rs | 22 +- src/server.rs | 17 +- 7 files changed, 490 insertions(+), 816 deletions(-) delete mode 100644 src/actor/listener_actor.rs diff --git a/src/actor/listener_actor.rs b/src/actor/listener_actor.rs deleted file mode 100644 index 19d4313..0000000 --- a/src/actor/listener_actor.rs +++ /dev/null @@ -1,344 +0,0 @@ -use crate::actor::{Actor, ActorSystem}; -use crate::client::ClientContext; -use crate::message::{Message, MessageReader}; -use crate::peer::{ConnectionType, DownloadPeer, Peer}; -use crate::types::Download; -use crate::{debug, error, info, trace, DownloadStatus}; - -use std::io; -use std::net::{TcpListener, TcpStream}; -use std::sync::{Arc, RwLock}; - -const PEER_INIT_MESSAGE_CODE: u8 = 1; - -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 -} - -/// Message types for ConnectionHandlerActor -#[derive(Debug, Clone)] -pub enum ConnectionHandlerMessage { - /// Process the incoming connection - Process, -} - -/// Actor that handles a single incoming connection -pub struct ConnectionHandlerActor { - stream: Option, - client_context: Arc>, - own_username: String, - peer_ip: String, - peer_port: u16, -} - -impl ConnectionHandlerActor { - pub fn new( - stream: TcpStream, - client_context: Arc>, - own_username: String, - ) -> Self { - let peer_addr = stream.peer_addr().unwrap(); - let peer_ip = peer_addr.ip().to_string(); - let peer_port = peer_addr.port(); - - Self { - stream: Some(stream), - client_context, - own_username, - peer_ip, - peer_port, - } - } - - fn handle_peer_connection( - &self, - peer: Peer, - stream: TcpStream, - reader: MessageReader, - ) { - debug!("[listener:{}:{}] connection type is P, reader buffer has {} bytes", - self.peer_ip, self.peer_port, reader.buffer_len()); - - let client_context = self.client_context.read().unwrap(); - if let Some(ref registry) = client_context.peer_registry { - match registry.register_peer(peer.clone(), stream, Some(reader)) { - Ok(_) => trace!("[listener] peer actor spawned for: {}", peer.username), - Err(e) => error!("Failed to spawn peer actor for {:?}: {:?}", peer.username, e), - } - } else { - error!("PeerRegistry not initialized"); - } - } - - fn handle_file_connection( - &self, - peer: Peer, - stream: TcpStream, - mut reader: MessageReader, - token: u32, - ) { - trace!( - "[client] DownloadFromPeer token: {} peer: {:?}", - token, - peer - ); - - // Set stream to blocking mode - DownloadPeer expects blocking I/O - if let Err(e) = stream.set_nonblocking(false) { - error!( - "Failed to set stream to blocking mode for {}:{} (token: {}) - Error: {}", - peer.host, peer.port, token, e - ); - return; - } - - let download = extract_download_from_buffer( - &mut reader, - &self.client_context, - &peer.username, - &self.peer_ip, - self.peer_port, - ); - - let download_peer = DownloadPeer::new( - format!("{}:direct", peer.username), - peer.host.clone(), - peer.port, - token, - true, - self.own_username.clone(), - ); - - match download_peer.download_file( - self.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 process_connection(&mut self) { - let Some(mut stream) = self.stream.take() else { - error!("[listener] stream already taken"); - return; - }; - - let mut reader = MessageReader::new(); - - let Ok(message) = read_peer_init_message(&mut stream, &mut reader) else { - error!( - "[listener:{}:{}] Failed to read peer init message", - self.peer_ip, self.peer_port - ); - return; - }; - - let Some(init_data) = parse_peer_init_message(message) else { - error!("[listener:{}:{}] Invalid or unknown peer init message", - self.peer_ip, self.peer_port); - return; - }; - - debug!( - "[listener:{}:{}] peerInit (0) username: {} connection_type: {} token: {}", - self.peer_ip, self.peer_port, init_data.username, init_data.connection_type, init_data.token - ); - - let peer = Peer::new( - format!("{}:direct", init_data.username), - init_data.connection_type.clone(), - self.peer_ip.clone(), - self.peer_port.into(), - None, - 0, - 0, - 0, - ); - - match init_data.connection_type { - ConnectionType::P => self.handle_peer_connection(peer, stream, reader), - ConnectionType::F => self.handle_file_connection(peer, stream, reader, init_data.token), - ConnectionType::D => { - debug!("[listener:{}:{}] connection type is D, not supported yet, closing connection.", - self.peer_ip, self.peer_port); - } - } - } -} - -impl Actor for ConnectionHandlerActor { - type Message = ConnectionHandlerMessage; - - fn handle(&mut self, msg: Self::Message) { - match msg { - ConnectionHandlerMessage::Process => self.process_connection(), - } - } - - fn on_start(&mut self) { - trace!("[connection_handler:{}:{}] starting", self.peer_ip, self.peer_port); - self.process_connection(); - } -} - -/// Message types for ListenerActor -#[derive(Debug, Clone)] -pub enum ListenerMessage { - /// Stop the listener - Stop, -} - -/// Actor that accepts incoming TCP connections -pub struct ListenerActor { - listener: Option, - client_context: Arc>, - own_username: String, - actor_system: Arc, -} - -impl ListenerActor { - pub fn new( - port: u32, - client_context: Arc>, - own_username: String, - actor_system: Arc, - ) -> Result { - let listener = TcpListener::bind(format!("0.0.0.0:{port}"))?; - listener.set_nonblocking(true)?; - - info!("[listener] starting listener on port {port}"); - - Ok(Self { - listener: Some(listener), - client_context, - own_username, - actor_system, - }) - } -} - -impl Actor for ListenerActor { - type Message = ListenerMessage; - - fn handle(&mut self, msg: Self::Message) { - match msg { - ListenerMessage::Stop => { - info!("[listener] stopping listener"); - self.listener = None; - } - } - } - - fn tick(&mut self) { - let Some(ref listener) = self.listener else { - return; - }; - - match listener.accept() { - Ok((stream, _addr)) => { - trace!("[listener] accepted new connection"); - let actor = ConnectionHandlerActor::new( - stream, - self.client_context.clone(), - self.own_username.clone(), - ); - self.actor_system.spawn(actor); - } - Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { - // No incoming connections, this is normal for non-blocking - } - Err(e) => { - error!("[listener] Failed to accept connection: {}", e); - } - } - } -} diff --git a/src/actor/peer_actor.rs b/src/actor/peer_actor.rs index 4be1294..36cf6e0 100644 --- a/src/actor/peer_actor.rs +++ b/src/actor/peer_actor.rs @@ -15,6 +15,14 @@ use std::io::{self, Error, Write}; use std::net::TcpStream; use std::sync::mpsc::{Receiver, Sender}; use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant}; + +#[derive(Debug, Clone)] +enum ConnectionState { + Disconnected, + Connecting { since: Instant }, + Connected, +} #[derive(Debug, Clone)] pub enum PeerMessage { @@ -37,46 +45,49 @@ pub enum PeerMessage { ProcessRead, } -/// Actor that manages a single peer connection pub struct PeerActor { peer: Arc>, stream: Option, + connection_state: ConnectionState, reader: MessageReader, client_channel: Sender, self_handle: Option>, dispatcher: Option>, dispatcher_receiver: Option>, - queued_filenames: Vec, + queued_messages: Vec, } impl PeerActor { - /// Create a new PeerActor with an existing socket connection pub fn new( peer: Peer, - stream: TcpStream, + stream: Option, reader: Option, client_channel: Sender, ) -> Self { + let connection_state = if stream.is_some() { + ConnectionState::Connected + } else { + ConnectionState::Disconnected + }; + Self { peer: Arc::new(RwLock::new(peer)), - stream: Some(stream), + stream, + connection_state, reader: reader.unwrap_or_default(), client_channel, self_handle: None, dispatcher: None, dispatcher_receiver: None, - queued_filenames: Vec::new(), + queued_messages: Vec::new(), } } - /// Set the actor's own handle for self-messaging pub fn set_self_handle(&mut self, handle: ActorHandle) { self.self_handle = Some(handle); } - /// Initialize the message dispatcher with handlers fn initialize_dispatcher(&mut self) { - // Create a channel for the dispatcher to send messages to let (dispatcher_sender, dispatcher_receiver) = std::sync::mpsc::channel::(); @@ -98,28 +109,34 @@ impl PeerActor { )); } - /// Process any messages from the dispatcher fn process_dispatcher_messages(&mut self) { - // Collect all messages first to avoid borrow issues - let messages: Vec = - if let Some(ref receiver) = self.dispatcher_receiver { + let messages: Vec = self + .dispatcher_receiver + .as_ref() + .map_or_else(Vec::new, |receiver| { let mut msgs = Vec::new(); while let Ok(msg) = receiver.try_recv() { msgs.push(msg); } msgs - } else { - Vec::new() - }; + }); - // Process all collected messages - for msg in messages { - self.handle_message(msg); - } + messages + .iter() + .for_each(|msg| self.handle_message(msg.clone())); } - /// Internal handler that doesn't call Actor::handle fn handle_message(&mut self, msg: PeerMessage) { + if matches!(self.connection_state, ConnectionState::Connecting { .. }) { + match &msg { + PeerMessage::SetUsername(_) | PeerMessage::ProcessRead => {} + _ => { + self.queued_messages.push(msg); + return; + } + } + } + match msg { PeerMessage::SendMessage(message) => { self.send_message(message); @@ -202,7 +219,6 @@ impl PeerActor { self.peer.write().unwrap().username = username; } PeerMessage::QueueUpload(filename) => { - self.queued_filenames.push(filename.clone()); let message = MessageFactory::build_queue_upload_message(&filename); self.send_message(message); @@ -225,21 +241,15 @@ impl PeerActor { } } - /// Handle reading from the socket and dispatching messages fn process_read(&mut self) { if self.reader.buffer_len() > 0 { - self.extract_and_process_messages(); // This will process all messages in the buffer + self.extract_and_process_messages(); } - let username = self.peer.read().unwrap().username.clone(); - // Try to read more data from the socket first { let stream = match self.stream.as_mut() { Some(s) => s, - None => { - trace!("[peer_actor:{}] process_read: stream is None, returning", username); - return; - } + None => return, }; match self.reader.read_from_socket(stream) { @@ -263,7 +273,7 @@ impl PeerActor { } } } - self.extract_and_process_messages(); // This will process all messages in the buffer + self.extract_and_process_messages(); } fn extract_and_process_messages(&mut self) { @@ -307,11 +317,9 @@ impl PeerActor { } } - // Process any messages that were dispatched self.process_dispatcher_messages(); } - /// Handle writing a message to the socket fn send_message(&mut self, message: Message) { let stream = match self.stream.as_mut() { Some(s) => s, @@ -365,7 +373,6 @@ impl PeerActor { error!("Failed to send disconnect notification: {}", e); } } - /// Disconnect the peer and notify the client fn disconnect(&mut self) { let username = self.peer.read().unwrap().username.clone(); debug!("[peer_actor:{}] disconnect", username); @@ -379,6 +386,145 @@ impl PeerActor { error!("Failed to send disconnect notification: {}", e); } } + + fn initiate_connection(&mut self) -> bool { + let peer = self.peer.read().unwrap(); + let username = peer.username.clone(); + let host = peer.host.clone(); + let port = peer.port; + drop(peer); + + let socket_addr = + format!("{}:{}", host, port).parse::(); + + match socket_addr { + Ok(addr) => { + if let Ok(stream) = TcpStream::connect(addr) { + if let Err(e) = stream.set_nonblocking(true) { + error!( + "[peer_actor:{}] Failed to set non-blocking: {}", + username, e + ); + self.disconnect_with_error(e); + return false; + } + + stream.set_nodelay(true).ok(); + + self.stream = Some(stream); + self.connection_state = ConnectionState::Connecting { + since: Instant::now(), + }; + true + } else { + match TcpStream::connect(addr) { + Ok(stream) => { + if let Err(e) = stream.set_nonblocking(true) { + error!("[peer_actor:{}] Failed to set non-blocking: {}", username, e); + self.disconnect_with_error(e); + return false; + } + stream.set_nodelay(true).ok(); + self.stream = Some(stream); + self.connection_state = + ConnectionState::Connecting { + since: Instant::now(), + }; + true + } + Err(e) => { + self.disconnect_with_error(e); + false + } + } + } + } + Err(e) => { + error!( + "[peer_actor:{}] Invalid socket address {}:{} - {}", + username, host, port, e + ); + self.disconnect_with_error(io::Error::new( + io::ErrorKind::InvalidInput, + e, + )); + false + } + } + } + + fn check_connection_status(&mut self) { + let ConnectionState::Connecting { since } = self.connection_state + else { + return; + }; + + let username = self.peer.read().unwrap().username.clone(); + + if since.elapsed() > Duration::from_secs(20) { + error!( + "[peer_actor:{}] Connection timeout after 20 seconds", + username + ); + self.disconnect_with_error(io::Error::new( + io::ErrorKind::TimedOut, + "Connection timeout", + )); + return; + } + + let Some(ref stream) = self.stream else { + return; + }; + + match stream.peer_addr() { + Ok(_) => { + self.connection_state = ConnectionState::Connected; + self.on_connection_established(); + } + Err(ref e) if e.kind() == io::ErrorKind::NotConnected => {} + Err(e) => { + error!("[peer_actor:{}] Connection failed: {}", username, e); + self.disconnect_with_error(e); + } + } + } + + fn on_connection_established(&mut self) { + let peer = self.peer.read().unwrap(); + let username = peer.username.clone(); + let token = peer.token; + drop(peer); + + let Some(ref mut stream) = self.stream else { + return; + }; + + if let Some(token) = token { + let handshake_msg = MessageFactory::build_watch_user(token); + if let Err(e) = stream.write_all(&handshake_msg.get_data()) { + error!( + "[peer_actor:{}] Failed to send watch user handshake: {}", + username, e + ); + self.disconnect_with_error(e); + return; + } + } + + self.initialize_dispatcher(); + + let queued = std::mem::take(&mut self.queued_messages); + for msg in queued { + self.handle_message(msg); + } + + if let Some(ref handle) = self.self_handle { + handle.send(PeerMessage::ProcessRead).ok(); + } + + self.process_read(); + } } impl Actor for PeerActor { @@ -389,29 +535,12 @@ impl Actor for PeerActor { } fn on_start(&mut self) { - let username = self.peer.read().unwrap().username.clone(); - - // Initialize the dispatcher - self.initialize_dispatcher(); - - // Trigger initial read - if let Some(ref handle) = self.self_handle { - match handle.send(PeerMessage::ProcessRead) { - Ok(_) => {} - Err(e) => error!( - "[peer_actor:{}] FAILED to send ProcessRead: {}", - username, e - ), - } + if self.stream.is_none() { + self.initiate_connection(); } else { - error!( - "[peer_actor:{}] self_handle is None! Cannot send ProcessRead", - username - ); + self.connection_state = ConnectionState::Connected; + self.on_connection_established(); } - - // Also do an immediate read as fallback - self.process_read(); } fn on_stop(&mut self) { @@ -421,15 +550,20 @@ impl Actor for PeerActor { } fn tick(&mut self) { - // Periodically check for new data on the socket - if self.stream.is_some() { - self.process_read(); - } else { - let username = self.peer.read().unwrap().username.clone(); - trace!( - "[peer_actor:{}] tick() called but stream is None", - username - ); + match self.connection_state { + ConnectionState::Connecting { .. } => { + // Check if connection has been established + self.check_connection_status(); + } + ConnectionState::Connected => { + // Normal operation - process reads + if self.stream.is_some() { + self.process_read(); + } + } + ConnectionState::Disconnected => { + // Do nothing + } } } } diff --git a/src/actor/peer_registry.rs b/src/actor/peer_registry.rs index 992396e..a7dc1b7 100644 --- a/src/actor/peer_registry.rs +++ b/src/actor/peer_registry.rs @@ -1,16 +1,15 @@ use crate::actor::peer_actor::{PeerActor, PeerMessage}; use crate::actor::{ActorHandle, ActorSystem}; use crate::client::ClientOperation; +use crate::debug; use crate::message::MessageReader; use crate::peer::Peer; -use crate::{debug, trace}; use std::collections::HashMap; use std::net::TcpStream; use std::sync::mpsc::Sender; use std::sync::{Arc, Mutex}; -/// Registry for managing peer actors pub struct PeerRegistry { peers: Arc>>>, actor_system: Arc, @@ -29,45 +28,33 @@ impl PeerRegistry { } } - /// Spawn a new peer actor and register it pub fn register_peer( &self, peer: Peer, - stream: TcpStream, + stream: Option, reader: Option, ) -> Result, String> { let username = peer.username.clone(); - debug!("[peer_registry] Registering peer actor for {}", username); - let actor = PeerActor::new(peer, stream, reader, self.client_channel.clone()); - // Spawn the actor with initialization callback to set self_handle let handle = self.actor_system.spawn_with_handle(actor, |actor, handle| { actor.set_self_handle(handle); }); - // Store in registry let mut peers = self.peers.lock().unwrap(); peers.insert(username.clone(), handle.clone()); - trace!( - "[peer_registry] Peer actor {} registered successfully", - username - ); - Ok(handle) } - /// Get a peer actor handle by username pub fn get_peer(&self, username: &str) -> Option> { let peers = self.peers.lock().unwrap(); peers.get(username).cloned() } - /// Remove a peer actor from the registry pub fn remove_peer( &self, username: &str, @@ -82,25 +69,21 @@ impl PeerRegistry { handle } - /// Get all registered peer usernames pub fn get_all_usernames(&self) -> Vec { let peers = self.peers.lock().unwrap(); peers.keys().cloned().collect() } - /// Get count of registered peers pub fn count(&self) -> usize { let peers = self.peers.lock().unwrap(); peers.len() } - /// Check if a peer is registered pub fn contains(&self, username: &str) -> bool { let peers = self.peers.lock().unwrap(); peers.contains_key(username) } - /// Send a message to a peer by username pub fn send_to_peer( &self, username: &str, diff --git a/src/client.rs b/src/client.rs index e5adc61..b5aae32 100644 --- a/src/client.rs +++ b/src/client.rs @@ -318,82 +318,96 @@ impl Client { own_username: String, ) { thread::spawn(move || loop { - if let Ok(operation) = reader.recv() { - match operation { - ClientOperation::ConnectToPeer(peer) => { - trace!("[client] ConnectToPeer {:?}", peer); - Self::connect_to_peer( - peer, - client_context.clone(), - own_username.clone(), - None, - ); - } - ClientOperation::SearchResult(file_search) => { - trace!("[client] SearchResult {:?}", file_search); - client_context - .write() - .unwrap() - .search_results - .push(file_search.clone()); - } - ClientOperation::PeerDisconnected(username, error) => { - let context = client_context.read().unwrap(); - if let Some(ref registry) = context.peer_registry { - if let Some(handle) = - registry.remove_peer(&username) - { - let _ = handle.stop(); - } + match reader.recv() { + Ok(operation) => { + match operation { + ClientOperation::ConnectToPeer(peer) => { + let client_context_clone = client_context.clone(); + let own_username_clone = own_username.clone(); + + // Spawn in background to avoid blocking the receiver thread + thread::spawn(move || { + Self::connect_to_peer( + peer, + client_context_clone, + own_username_clone, + None, + ); + }); } - if let Some(error) = error { - error!( + ClientOperation::SearchResult(file_search) => { + trace!("[client] SearchResult {:?}", file_search); + client_context + .write() + .unwrap() + .search_results + .push(file_search.clone()); + } + ClientOperation::PeerDisconnected(username, error) => { + let context = client_context.read().unwrap(); + if let Some(ref registry) = context.peer_registry { + if let Some(handle) = + registry.remove_peer(&username) + { + let _ = handle.stop(); + } + } + if let Some(error) = error { + error!( "[client] Peer {} disconnected with error: {:?}", username, error ); - Self::process_failed_uploads( + Self::process_failed_uploads( + client_context.clone(), + &username, + None, + ); + } + } + ClientOperation::PierceFireWall(peer) => { + Self::pierce_firewall( + peer, client_context.clone(), - &username, - None, + own_username.clone(), ); } - } - ClientOperation::PierceFireWall(peer) => { - Self::pierce_firewall( + ClientOperation::DownloadFromPeer( + token, peer, - client_context.clone(), - own_username.clone(), - ); - } - ClientOperation::DownloadFromPeer(token, peer, allowed) => { - let maybe_download = { - let client_context = client_context.read().unwrap(); - client_context.download_tokens.get(&token).cloned() - }; - let own_username = own_username.clone(); - let client_context_clone = client_context.clone(); + allowed, + ) => { + let maybe_download = { + let client_context = + client_context.read().unwrap(); + client_context + .download_tokens + .get(&token) + .cloned() + }; + let own_username = own_username.clone(); + let client_context_clone = client_context.clone(); - trace!( + trace!( "[client] DownloadFromPeer token: {} peer: {:?}", token, peer ); - return match maybe_download { - Some(download) => { - thread::spawn(move || { - let download_peer = DownloadPeer::new( - download.username.clone(), - peer.host.clone(), - peer.port, - token, - allowed, - own_username, - ); - let filename: Option<&str> = download - .filename - .split('\\') - .next_back(); - match filename { + match maybe_download { + Some(download) => { + thread::spawn(move || { + let download_peer = DownloadPeer::new( + download.username.clone(), + peer.host.clone(), + peer.port, + token, + allowed, + own_username, + ); + let filename: Option<&str> = download + .filename + .split('\\') + .next_back(); + match filename { Some(filename) => { match download_peer.download_file( client_context_clone, @@ -415,164 +429,176 @@ impl Client { } None => error!("Cant find filename to save download: {:?}", download.filename), } - }); - } - None => { - error!( - "Can't find download with token {:?}", - token - ); - } - }; - } - ClientOperation::NewPeer(new_peer) => { - let peer_exists = client_context - .read() - .unwrap() - .peer_registry - .as_ref() - .map(|r| r.contains(&new_peer.username)) - .unwrap_or(false); - - if peer_exists { - debug!( - "Already connected to {}", - new_peer.username - ); - } else if let Some(server_sender) = - &client_context.read().unwrap().server_sender - { - server_sender - .send(ServerOperation::GetPeerAddress( - new_peer.username.clone(), - )) - .unwrap(); + }); + } + None => { + error!( + "Can't find download with token {:?}", + token + ); + } + }; } + ClientOperation::NewPeer(new_peer) => { + let peer_exists = client_context + .read() + .unwrap() + .peer_registry + .as_ref() + .map(|r| r.contains(&new_peer.username)) + .unwrap_or(false); - let addr = new_peer.tcp_stream.peer_addr().unwrap(); - let host = addr.ip().to_string(); - let port: u32 = addr.port().into(); + if peer_exists { + debug!( + "Already connected to {}", + new_peer.username + ); + } else if let Some(server_sender) = + &client_context.read().unwrap().server_sender + { + server_sender + .send(ServerOperation::GetPeerAddress( + new_peer.username.clone(), + )) + .unwrap(); + } - let peer = Peer { - username: new_peer.username.clone(), - connection_type: new_peer.connection_type, - host, - port, - token: Some(new_peer.token), - privileged: None, - obfuscated_port: None, - unknown: None, - }; + let addr = new_peer.tcp_stream.peer_addr().unwrap(); + let host = addr.ip().to_string(); + let port: u32 = addr.port().into(); - Self::connect_to_peer( - peer, - client_context.clone(), - own_username.clone(), - Some(new_peer.tcp_stream), - ); - } - ClientOperation::GetPeerAddressResponse { - username, - host, - port, - obfuscation_type, - obfuscated_port, - } => { - debug!( - "Received peer address for {}: {}:{} (obf_type: {}, obf_port: {})", - username, host, port, obfuscation_type, obfuscated_port - ); - - let peer_exists = client_context - .read() - .unwrap() - .peer_registry - .as_ref() - .map(|r| r.contains(&username)) - .unwrap_or(false); - - if peer_exists { - // don't know if i should update? and or reconnect the peer - // debug!( - // "existing peer: {:?}, new peer details: - // username: {}, - // host: {}, - // port: {} - // obfuscation_type: {} - // obfuscated_port: {}", - // peer, - // username, - // host, - // port, - // obfuscation_type, - // obfuscated_port, - // ); - } else { - let peer = Peer::new( - username, - ConnectionType::P, + let peer = Peer { + username: new_peer.username.clone(), + connection_type: new_peer.connection_type, host, port, - None, - 0, - obfuscation_type.try_into().unwrap(), - obfuscated_port.try_into().unwrap(), - ); + token: Some(new_peer.token), + privileged: None, + obfuscated_port: None, + unknown: None, + }; + Self::connect_to_peer( peer, client_context.clone(), own_username.clone(), - None, + Some(new_peer.tcp_stream), ); } - } - ClientOperation::UpdateDownloadTokens( - transfer, - username, - ) => { - let mut context = client_context.write().unwrap(); + ClientOperation::GetPeerAddressResponse { + username, + host, + port, + obfuscation_type, + obfuscated_port, + } => { + debug!( + "Received peer address for {}: {}:{} (obf_type: {}, obf_port: {})", + username, host, port, obfuscation_type, obfuscated_port + ); - let key_to_remove = context - .download_tokens - .iter() - .find_map(|(key, d)| { - if d.username == username - && d.filename == transfer.filename - { - Some((*key, d.clone())) - } else { - None - } - }); + let peer_exists = client_context + .read() + .unwrap() + .peer_registry + .as_ref() + .map(|r| r.contains(&username)) + .unwrap_or(false); - if let Some((key, download)) = key_to_remove { - trace!( + if peer_exists { + // don't know if i should update? and or reconnect the peer + // debug!( + // "existing peer: {:?}, new peer details: + // username: {}, + // host: {}, + // port: {} + // obfuscation_type: {} + // obfuscated_port: {}", + // peer, + // username, + // host, + // port, + // obfuscation_type, + // obfuscated_port, + // ); + } else { + let peer = Peer::new( + username, + ConnectionType::P, + host, + port, + None, + 0, + obfuscation_type.try_into().unwrap(), + obfuscated_port.try_into().unwrap(), + ); + let client_context_clone = + client_context.clone(); + let own_username_clone = own_username.clone(); + + // Spawn in background to avoid blocking the receiver thread + thread::spawn(move || { + Self::connect_to_peer( + peer, + client_context_clone, + own_username_clone, + None, + ); + }); + } + } + ClientOperation::UpdateDownloadTokens( + transfer, + username, + ) => { + let mut context = client_context.write().unwrap(); + + let key_to_remove = context + .download_tokens + .iter() + .find_map(|(key, d)| { + if d.username == username + && d.filename == transfer.filename + { + Some((*key, d.clone())) + } else { + None + } + }); + + if let Some((key, download)) = key_to_remove { + trace!( "[client] UpdateDownloadTokens found {key}, transfer: {:?}", transfer ); - context.download_tokens.insert( - transfer.token, - Download { - username: username.clone(), - filename: transfer.filename, - token: transfer.token, - size: transfer.size, - download_directory: download - .download_directory, - sender: download.sender.clone(), - }, + context.download_tokens.insert( + transfer.token, + Download { + username: username.clone(), + filename: transfer.filename, + token: transfer.token, + size: transfer.size, + download_directory: download + .download_directory, + sender: download.sender.clone(), + }, + ); + context.download_tokens.remove(&key); + } + } + ClientOperation::UploadFailed(username, filename) => { + Self::process_failed_uploads( + client_context.clone(), + &username, + Some(&filename), ); - context.download_tokens.remove(&key); } } - ClientOperation::UploadFailed(username, filename) => { - Self::process_failed_uploads( - client_context.clone(), - &username, - Some(&filename), - ); - } + } + Err(e) => { + error!("[client] Channel recv error: {:?}", e); + break; } } }); @@ -584,168 +610,70 @@ impl Client { own_username: String, stream: Option, ) { - let client_context2 = client_context.clone(); - let unlocked_context = client_context.read().unwrap(); + let client_context = client_context.clone(); - trace!("[client] connect_to_peer: {}", peer.username); - if let Some(_sender) = unlocked_context.sender.clone() { - let peer_clone = peer.clone(); - trace!("[client] connecting to {}, with connection_type: {}, and token {:?}", peer.username, peer.connection_type, peer.token); - match peer.connection_type { - ConnectionType::P => { - let username = peer.username.clone(); + let peer_clone = peer.clone(); + trace!("[client] connecting to {}, with connection_type: {}, and token {:?}", peer.username, peer.connection_type, peer.token); + match peer.connection_type { + ConnectionType::P => { + let username = peer.username.clone(); - let mut tcp_stream = match stream { - Some(s) => s, - None => { - let socket_addr = format!( - "{}:{}", - peer_clone.host, peer_clone.port - ) - .parse::(); - - match socket_addr { - Ok(addr) => { - match TcpStream::connect_timeout( - &addr, - Duration::from_secs(20), - ) { - Ok(s) => s, - Err(e) => { - trace!( - "Can't connect to {:?} {:?}:{:?} - {:?}", - username, peer_clone.host, peer_clone.port, e - ); - return; - } - } - } - Err(e) => { - trace!( - "Invalid socket address for {:?}: {:?}", - username, - e - ); - return; - } - } - } - }; - - // Send connection handshake message - use std::io::Write; - if let Some(token) = peer_clone.token { - // For indirect connections, send watch user message - let handshake_msg = crate::message::server::MessageFactory::build_watch_user(token); - if let Err(e) = - tcp_stream.write_all(&handshake_msg.get_data()) - { - trace!("Failed to send watch user handshake to {:?}: {:?}", username, e); - return; - } - trace!( - "[client] Sent watch user handshake for token: {}", - token - ); - } - - // Configure socket - if let Err(e) = tcp_stream - .set_read_timeout(Some(Duration::from_secs(5))) - { - trace!( - "Failed to set read timeout for {:?}: {:?}", - username, - e - ); - return; - } - if let Err(e) = tcp_stream - .set_write_timeout(Some(Duration::from_secs(5))) - { - trace!( - "Failed to set write timeout for {:?}: {:?}", - username, - e - ); - return; - } - if let Err(e) = tcp_stream.set_nodelay(true) { - trace!( - "Failed to set nodelay for {:?}: {:?}", - username, - e - ); - return; - } - - // Spawn peer actor - let context = client_context2.read().unwrap(); - if let Some(ref registry) = context.peer_registry { - match registry - .register_peer(peer_clone, tcp_stream, None) - { - Ok(_) => { - trace!( - "[client] peer actor spawned for: {}", - username - ); - } - Err(e) => { - trace!( - "Failed to spawn peer actor for {:?}: {:?}", - username, - e - ); - } - } - } else { - trace!("PeerRegistry not initialized"); - } - } - - ConnectionType::F => { - trace!( - "[client] downloading from: {}, {:?}", - peer.username, - peer.token - ); - let download_peer = DownloadPeer::new( - peer.username, - peer.host, - peer.port, - peer.token.unwrap(), - false, - own_username.clone(), - ); - - match download_peer.download_file( - client_context2.clone(), - None, - None, - ) { - Ok((download, filename)) => { - trace!( - "[client] downloaded {} bytes {:?} ", - filename, - download.size - ); - download - .sender - .send(DownloadStatus::Completed) - .unwrap(); - } + let context = client_context.read().unwrap(); + if let Some(ref registry) = context.peer_registry { + match registry.register_peer(peer_clone, stream, None) { + Ok(_) => (), Err(e) => { - trace!("[client] failed to download: {}", e); + trace!( + "Failed to spawn peer actor for {:?}: {:?}", + username, + e + ); } } - } - ConnectionType::D => { - error!("ConnectionType::D not implemented") + } else { + trace!("PeerRegistry not initialized"); } } - } else { - debug!("[client] Peer already connected: {}", peer.username); + + ConnectionType::F => { + trace!( + "[client] downloading from: {}, {:?}", + peer.username, + peer.token + ); + let download_peer = DownloadPeer::new( + peer.username, + peer.host, + peer.port, + peer.token.unwrap(), + false, + own_username.clone(), + ); + + match download_peer.download_file( + client_context.clone(), + None, + None, + ) { + Ok((download, filename)) => { + trace!( + "[client] downloaded {} bytes {:?} ", + filename, + download.size + ); + download + .sender + .send(DownloadStatus::Completed) + .unwrap(); + } + Err(e) => { + trace!("[client] failed to download: {}", e); + } + } + } + ConnectionType::D => { + error!("ConnectionType::D not implemented") + } } } fn pierce_firewall( diff --git a/src/main.rs b/src/main.rs index 55d15bd..bbc4458 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,7 +14,7 @@ fn main() { match client.login() { Ok(_) => { match client - .search("Eternity Alex Kassian", Duration::from_secs(10)) + .search("Eternity Alex Kassian", Duration::from_secs(30)) { Ok(results) => { if let Some(result) = diff --git a/src/peer/listen.rs b/src/peer/listen.rs index 5f90be5..080fc31 100644 --- a/src/peer/listen.rs +++ b/src/peer/listen.rs @@ -75,14 +75,7 @@ fn extract_download_from_buffer( 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", @@ -107,18 +100,13 @@ fn handle_peer_connection( stream: TcpStream, reader: MessageReader, context: &ConnectionContext, - peer_ip: &str, - peer_port: u16, + _peer_ip: &str, + _peer_port: u16, ) { - debug!("[listener:{peer_ip}:{peer_port}] connection type is P, reader buffer has {} bytes", reader.buffer_len()); - - // Spawn peer actor directly using the peer registry let client_context = context.client_context.read().unwrap(); if let Some(ref registry) = client_context.peer_registry { - match registry.register_peer(peer.clone(), stream, Some(reader)) { - Ok(_) => { - trace!("[listener] peer actor spawned for: {}", peer.username); - } + match registry.register_peer(peer.clone(), Some(stream), Some(reader)) { + Ok(_) => (), Err(e) => { error!( "Failed to spawn peer actor for {:?}: {:?}", @@ -213,7 +201,7 @@ fn handle_incoming_connection(stream: TcpStream, context: ConnectionContext) { ); let peer = Peer::new( - format!("{}:direct", init_data.username), // Add direct suffix to distinguish from regular peers + format!("{}:direct", init_data.username), // init_data.username.clone(), init_data.connection_type.clone(), peer_ip.clone(), diff --git a/src/server.rs b/src/server.rs index d8cc7d9..795773c 100644 --- a/src/server.rs +++ b/src/server.rs @@ -341,10 +341,6 @@ impl Server { if let Ok(operation) = server_channel.recv() { match operation { ServerOperation::ConnectToPeer(peer) => { - trace!( - "[server] → ConnectToPeer {}", - peer.username - ); if let Some(op) = match peer.connection_type { ConnectionType::P | ConnectionType::F => { Some(ClientOperation::ConnectToPeer( @@ -353,18 +349,7 @@ impl Server { } ConnectionType::D => None, } { - trace!( - "[server] → ConnectToPeer {} - op {:?}", - peer.username, - op - ); - - match client_channel.send(op) { - Ok(_) => {} - Err(e) => { - error!("[server] Failed to send: {}", e) - } - } + client_channel.send(op).unwrap(); } } ServerOperation::LoginStatus(message) => {