diff --git a/src/actor/mod.rs b/src/actor/mod.rs new file mode 100644 index 0000000..0b619ad --- /dev/null +++ b/src/actor/mod.rs @@ -0,0 +1,215 @@ +// src/actor/mod.rs +// Lightweight actor system using the existing thread pool + +use std::sync::mpsc::{channel, Receiver, Sender}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use crate::trace; +use crate::utils::thread_pool::ThreadPool; + +pub mod peer_actor; +pub mod peer_registry; + +/// Core actor trait - each actor processes messages +pub trait Actor: Send + 'static { + type Message: Send + Clone + 'static; + + /// Handle a single message + fn handle(&mut self, msg: Self::Message); + + /// Called when actor starts (optional hook) + fn on_start(&mut self) {} + + /// Called when actor stops (optional hook) + fn on_stop(&mut self) {} + + /// Optional periodic tick for background work + fn tick(&mut self) {} +} + +/// Handle to send messages to an actor +#[derive(Clone)] +pub struct ActorHandle { + pub(crate) sender: Sender>, +} + +impl ActorHandle { + /// Send a message to the actor (non-blocking) + pub fn send(&self, msg: M) -> Result<(), String> { + self.sender + .send(ActorMessage::UserMessage(msg)) + .map_err(|e| format!("Failed to send message: {}", e)) + } + + /// Request actor to stop gracefully + pub fn stop(&self) -> Result<(), String> { + self.sender + .send(ActorMessage::Stop) + .map_err(|e| format!("Failed to send stop signal: {}", e)) + } +} + +/// Internal actor message wrapper +pub(crate) enum ActorMessage { + UserMessage(M), + Stop, + #[allow(dead_code)] + Tick, +} + +/// Actor system that manages actor lifecycle +pub struct ActorSystem { + thread_pool: Arc, +} + +impl ActorSystem { + pub fn new(thread_pool: Arc) -> Self { + ActorSystem { thread_pool } + } + + /// Spawn a new actor and return its handle + pub fn spawn(&self, mut actor: A) -> ActorHandle { + let (sender, receiver) = channel::>(); + let handle = ActorHandle { + sender: sender.clone(), + }; + + // Start the actor event loop on the thread pool + self.thread_pool.execute(move || { + actor.on_start(); + Self::run_actor_loop(&mut actor, receiver); + actor.on_stop(); + }); + + handle + } + + /// Spawn a new actor with initialization callback and return its handle + /// The callback receives the actor handle before on_start is called + pub fn spawn_with_handle( + &self, + mut actor: A, + init: F, + ) -> ActorHandle + where + F: FnOnce(&mut A, ActorHandle) + Send + 'static, + { + let (sender, receiver) = channel::>(); + let handle = ActorHandle { + sender: sender.clone(), + }; + let handle_for_init = handle.clone(); + + self.thread_pool.execute(move || { + init(&mut actor, handle_for_init); + actor.on_start(); + Self::run_actor_loop(&mut actor, receiver); + actor.on_stop(); + }); + + handle + } + + fn run_actor_loop( + actor: &mut A, + receiver: Receiver>, + ) { + let tick_interval = Duration::from_millis(100); + let mut last_tick = Instant::now(); + let mut message_count = 0; + let mut tick_count = 0; + + loop { + match receiver.recv_timeout(tick_interval) { + Ok(ActorMessage::UserMessage(msg)) => { + message_count += 1; + actor.handle(msg); + } + Ok(ActorMessage::Stop) => { + trace!( + "[actor_system] Received Stop message, breaking loop" + ); + break; + } + Ok(ActorMessage::Tick) => { + tick_count += 1; + trace!( + "[actor_system] Received explicit Tick message #{}", + tick_count + ); + + actor.tick(); + } + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { + if last_tick.elapsed() >= tick_interval { + tick_count += 1; + // if tick_count % 10 == 0 { + // trace!("[actor_system] Periodic tick #{} (every 1s log)", tick_count); + // } + actor.tick(); + last_tick = Instant::now(); + } + } + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { + trace!( + "[actor_system] Channel disconnected, breaking loop" + ); + break; + } + } + } + trace!("[actor_system] run_actor_loop ENDED - processed {} messages, {} ticks", message_count, tick_count); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct CounterActor { + count: Arc, + } + + impl Actor for CounterActor { + type Message = usize; + + fn handle(&mut self, msg: Self::Message) { + self.count.fetch_add(msg, Ordering::SeqCst); + } + + fn on_start(&mut self) { + println!("Counter actor started"); + } + + fn on_stop(&mut self) { + println!("Counter actor stopped"); + } + } + + #[test] + fn test_actor_system() { + let thread_pool = Arc::new(ThreadPool::new(4)); + let system = ActorSystem::new(thread_pool); + + let count = Arc::new(AtomicUsize::new(0)); + let actor = CounterActor { + count: count.clone(), + }; + + let handle = system.spawn(actor); + + // Send some messages + handle.send(1).unwrap(); + handle.send(2).unwrap(); + handle.send(3).unwrap(); + + // Give actor time to process + std::thread::sleep(Duration::from_millis(100)); + + assert_eq!(count.load(Ordering::SeqCst), 6); + + handle.stop().unwrap(); + } +} diff --git a/src/actor/peer_actor.rs b/src/actor/peer_actor.rs new file mode 100644 index 0000000..36cf6e0 --- /dev/null +++ b/src/actor/peer_actor.rs @@ -0,0 +1,569 @@ +use crate::actor::{Actor, ActorHandle}; +use crate::client::ClientOperation; +use crate::dispatcher::MessageDispatcher; +use crate::message::peer::{ + FileSearchResponse, GetShareFileList, PeerInit, PlaceInQueueResponse, + TransferRequest, TransferResponse, UploadFailedHandler, +}; +use crate::message::server::MessageFactory; +use crate::message::{Handlers, Message, MessageReader, MessageType}; +use crate::peer::Peer; +use crate::types::{Download, FileSearchResult, Transfer}; +use crate::{debug, error, trace, warn}; + +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 { + SendMessage(Message), + FileSearchResult(FileSearchResult), + TransferRequest(Transfer), + UploadFailed(String, String), + TransferResponse { + token: u32, + allowed: bool, + reason: Option, + }, + PlaceInQueueResponse { + filename: String, + place: u32, + }, + SetUsername(String), + QueueUpload(String), + RequestTransfer(Download), + ProcessRead, +} + +pub struct PeerActor { + peer: Arc>, + stream: Option, + connection_state: ConnectionState, + reader: MessageReader, + client_channel: Sender, + self_handle: Option>, + dispatcher: Option>, + dispatcher_receiver: Option>, + queued_messages: Vec, +} + +impl PeerActor { + pub fn new( + peer: Peer, + 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, + connection_state, + reader: reader.unwrap_or_default(), + client_channel, + self_handle: None, + dispatcher: None, + dispatcher_receiver: None, + queued_messages: Vec::new(), + } + } + + pub fn set_self_handle(&mut self, handle: ActorHandle) { + self.self_handle = Some(handle); + } + + fn initialize_dispatcher(&mut self) { + let (dispatcher_sender, dispatcher_receiver) = + std::sync::mpsc::channel::(); + + self.dispatcher_receiver = Some(dispatcher_receiver); + + let mut handlers = Handlers::new(); + handlers.register_handler(FileSearchResponse); + handlers.register_handler(TransferRequest); + handlers.register_handler(TransferResponse); + handlers.register_handler(GetShareFileList); + handlers.register_handler(UploadFailedHandler); + handlers.register_handler(PlaceInQueueResponse); + handlers.register_handler(PeerInit); + + self.dispatcher = Some(MessageDispatcher::new( + "peer_actor".to_string(), + dispatcher_sender, + handlers, + )); + } + + fn process_dispatcher_messages(&mut self) { + 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 + }); + + messages + .iter() + .for_each(|msg| self.handle_message(msg.clone())); + } + + 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); + } + PeerMessage::FileSearchResult(file_search) => { + self.client_channel + .send(ClientOperation::SearchResult(file_search)) + .unwrap(); + } + PeerMessage::TransferRequest(transfer) => { + let username = self.peer.read().unwrap().username.clone(); + debug!( + "[peer_actor:{}] TransferRequest for {}", + username, transfer.token + ); + + self.client_channel + .send(ClientOperation::UpdateDownloadTokens( + transfer.clone(), + username.clone(), + )) + .unwrap(); + + let transfer_response = + MessageFactory::build_transfer_response_message(transfer); + + if let Some(ref handle) = self.self_handle { + if let Err(e) = + handle.send(PeerMessage::SendMessage(transfer_response)) + { + error!("[peer_actor:{}] Failed to send TransferResponse message: {}", username, e); + } + } + } + PeerMessage::TransferResponse { + token, + allowed, + reason, + } => { + let username = self.peer.read().unwrap().username.clone(); + debug!( + "[peer_actor:{}] transfer response token: {} allowed: {}", + username, token, allowed + ); + + if !allowed { + if let Some(reason_text) = reason { + debug!( + "[peer_actor:{}] Transfer rejected: {} - token {}, waiting for TransferRequest...", + username, reason_text, token + ); + } + } else { + debug!( + "[peer_actor:{}] Transfer allowed, ready to connect with token {:}", + username, token + ); + self.client_channel + .send(ClientOperation::DownloadFromPeer( + token, + self.peer.read().unwrap().clone(), + allowed, + )) + .unwrap(); + } + } + PeerMessage::PlaceInQueueResponse { filename, place } => { + let username = self.peer.read().unwrap().username.clone(); + debug!( + "[peer_actor:{}] Place in queue response - file: {}, place: {}", + username, filename, place + ); + } + PeerMessage::SetUsername(username) => { + trace!( + "[peer_actor:{}] SetUsername: {}", + self.peer.read().unwrap().username, + username + ); + self.peer.write().unwrap().username = username; + } + PeerMessage::QueueUpload(filename) => { + let message = + MessageFactory::build_queue_upload_message(&filename); + self.send_message(message); + } + PeerMessage::RequestTransfer(download) => { + let message = MessageFactory::build_transfer_request_message( + &download.filename, + download.token, + ); + self.send_message(message); + } + PeerMessage::ProcessRead => { + self.process_read(); + } + PeerMessage::UploadFailed(username, filename) => { + self.client_channel + .send(ClientOperation::UploadFailed(username, filename)) + .unwrap(); + } + } + } + + fn process_read(&mut self) { + if self.reader.buffer_len() > 0 { + self.extract_and_process_messages(); + } + + { + let stream = match self.stream.as_mut() { + Some(s) => s, + None => return, + }; + + match self.reader.read_from_socket(stream) { + Ok(()) => {} + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {} + Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { + let peer_lock = self.peer.read().unwrap(); + debug!( + "Read operation timed out for peer actor {:}:{:}", + peer_lock.host, peer_lock.port + ); + } + Err(e) => { + let username = self.peer.read().unwrap().username.clone(); + error!( + "[peer_actor:{}] Error reading from peer: {} (kind: {:?}). Disconnecting.", + username, e, e.kind() + ); + self.disconnect_with_error(e); + return; + } + } + } + self.extract_and_process_messages(); + } + + fn extract_and_process_messages(&mut self) { + let username = self.peer.read().unwrap().username.clone(); + let mut extracted_count = 0; + loop { + match self.reader.extract_message() { + Ok(Some(mut message)) => { + extracted_count += 1; + trace!( + "[peer_actor:{}] ← Message #{}: {:?}", + username, + extracted_count, + message + .get_message_name( + MessageType::Peer, + message.get_message_code() as u32 + ) + .map_err(|e| e.to_string()) + ); + if let Some(ref dispatcher) = self.dispatcher { + dispatcher.dispatch(&mut message); + } else { + warn!( + "[peer_actor:{}] No dispatcher available!", + username + ); + } + } + Err(e) => { + warn!( + "[peer_actor:{}] Error extracting message: {}. Disconnecting peer.", + username, e + ); + self.disconnect_with_error(e); + return; + } + Ok(None) => { + break; + } + } + } + + self.process_dispatcher_messages(); + } + + fn send_message(&mut self, message: Message) { + let stream = match self.stream.as_mut() { + Some(s) => s, + None => { + error!("Cannot send message: stream is None"); + return; + } + }; + + let username = self.peer.read().unwrap().username.clone(); + trace!( + "[peer_actor:{}] ➡ {:?}", + username, + message + .get_message_name( + MessageType::Peer, + u32::from_le_bytes( + message.get_slice(0, 4).try_into().unwrap() + ) + ) + .map_err(|e| e.to_string()) + ); + + if let Err(e) = stream.write_all(&message.get_buffer()) { + error!( + "[peer_actor:{}] Error writing message: {}. Disconnecting.", + username, e + ); + self.disconnect_with_error(e); + return; + } + + if let Err(e) = stream.flush() { + error!( + "[peer_actor:{}] Error flushing stream: {}. Disconnecting.", + username, e + ); + self.disconnect_with_error(e); + } + } + + fn disconnect_with_error(&mut self, error: Error) { + let username = self.peer.read().unwrap().username.clone(); + debug!("[peer_actor:{}] disconnect", username); + + self.stream.take(); + + if let Err(e) = self.client_channel.send( + ClientOperation::PeerDisconnected(username, Some(error.into())), + ) { + error!("Failed to send disconnect notification: {}", e); + } + } + fn disconnect(&mut self) { + let username = self.peer.read().unwrap().username.clone(); + debug!("[peer_actor:{}] disconnect", username); + + self.stream.take(); + + if let Err(e) = self + .client_channel + .send(ClientOperation::PeerDisconnected(username, None)) + { + 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 { + type Message = PeerMessage; + + fn handle(&mut self, msg: Self::Message) { + self.handle_message(msg); + } + + fn on_start(&mut self) { + if self.stream.is_none() { + self.initiate_connection(); + } else { + self.connection_state = ConnectionState::Connected; + self.on_connection_established(); + } + } + + fn on_stop(&mut self) { + let username = self.peer.read().unwrap().username.clone(); + trace!("[peer_actor:{}] actor stopping", username); + self.disconnect(); + } + + fn tick(&mut self) { + 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 new file mode 100644 index 0000000..a7dc1b7 --- /dev/null +++ b/src/actor/peer_registry.rs @@ -0,0 +1,116 @@ +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 std::collections::HashMap; +use std::net::TcpStream; +use std::sync::mpsc::Sender; +use std::sync::{Arc, Mutex}; + +pub struct PeerRegistry { + peers: Arc>>>, + actor_system: Arc, + client_channel: Sender, +} + +impl PeerRegistry { + pub fn new( + actor_system: Arc, + client_channel: Sender, + ) -> Self { + Self { + peers: Arc::new(Mutex::new(HashMap::new())), + actor_system, + client_channel, + } + } + + pub fn register_peer( + &self, + peer: Peer, + stream: Option, + reader: Option, + ) -> Result, String> { + let username = peer.username.clone(); + + let actor = + PeerActor::new(peer, stream, reader, self.client_channel.clone()); + + let handle = + self.actor_system.spawn_with_handle(actor, |actor, handle| { + actor.set_self_handle(handle); + }); + + let mut peers = self.peers.lock().unwrap(); + peers.insert(username.clone(), handle.clone()); + + Ok(handle) + } + + pub fn get_peer(&self, username: &str) -> Option> { + let peers = self.peers.lock().unwrap(); + peers.get(username).cloned() + } + + pub fn remove_peer( + &self, + username: &str, + ) -> Option> { + let mut peers = self.peers.lock().unwrap(); + let handle = peers.remove(username); + + if handle.is_some() { + debug!("[peer_registry] Removed peer actor for {}", username); + } + + handle + } + + pub fn get_all_usernames(&self) -> Vec { + let peers = self.peers.lock().unwrap(); + peers.keys().cloned().collect() + } + + pub fn count(&self) -> usize { + let peers = self.peers.lock().unwrap(); + peers.len() + } + + pub fn contains(&self, username: &str) -> bool { + let peers = self.peers.lock().unwrap(); + peers.contains_key(username) + } + + pub fn send_to_peer( + &self, + username: &str, + message: PeerMessage, + ) -> Result<(), String> { + let handle = self.get_peer(username).ok_or_else(|| { + format!("Peer {} not found in registry", username) + })?; + + handle.send(message) + } + + pub fn queue_upload( + &self, + username: &str, + filename: String, + ) -> Result<(), String> { + self.send_to_peer(username, PeerMessage::QueueUpload(filename)) + } +} + +impl Clone for PeerRegistry { + fn clone(&self) -> Self { + Self { + peers: self.peers.clone(), + actor_system: self.actor_system.clone(), + client_channel: self.client_channel.clone(), + } + } +} diff --git a/src/client.rs b/src/client.rs index 1dda7ed..b5aae32 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,10 +1,8 @@ use crate::types::{DownloadResult, DownloadStatus}; use crate::{ + actor::{peer_registry::PeerRegistry, ActorSystem}, error::{Result, SoulseekRs}, - peer::{ - listen::Listen, ConnectionType, DefaultPeer, DownloadPeer, NewPeer, - Peer, - }, + peer::{listen::Listen, ConnectionType, DownloadPeer, NewPeer, Peer}, server::{PeerAddress, Server, ServerOperation}, types::{Download, FileSearchResult}, utils::{md5, thread_pool::ThreadPool}, @@ -25,14 +23,15 @@ use std::{ }; use crate::{debug, error, info, trace}; -const MAX_THREADS: usize = 50; const DEFALT_LISTEN_PORT: u32 = 2234; +#[derive(Debug)] + pub enum ClientOperation { NewPeer(NewPeer), ConnectToPeer(Peer), SearchResult(FileSearchResult), - PeerDisconnected(String), + PeerDisconnected(String, Option), PierceFireWall(Peer), DownloadFromPeer(u32, Peer, bool), UpdateDownloadTokens(Transfer, String), @@ -43,14 +42,15 @@ pub enum ClientOperation { obfuscation_type: u32, obfuscated_port: u16, }, + UploadFailed(String, String), } pub struct ClientContext { - peers: HashMap, + pub peer_registry: Option, sender: Option>, server_sender: Option>, search_results: Vec, pub download_tokens: HashMap, - thread_pool: ThreadPool, + actor_system: Arc, } impl Default for ClientContext { fn default() -> Self { @@ -61,13 +61,20 @@ impl Default for ClientContext { impl ClientContext { #[must_use] pub fn new() -> Self { + let max_threads = thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(8); + + let thread_pool = Arc::new(ThreadPool::new(max_threads)); + let actor_system = Arc::new(ActorSystem::new(thread_pool.clone())); + Self { - peers: HashMap::new(), + peer_registry: None, sender: None, server_sender: None, search_results: Vec::new(), download_tokens: HashMap::new(), - thread_pool: ThreadPool::new(MAX_THREADS), + actor_system, } } } @@ -91,10 +98,6 @@ impl Client { ) -> Self { crate::utils::logger::init(); let context = Arc::new(RwLock::new(ClientContext::new())); - debug!( - "ThreadPool initialized with {} threads", - context.read().unwrap().thread_pool.thread_count() - ); Self { enable_listen, listen_port: listen_port.unwrap_or(DEFALT_LISTEN_PORT), @@ -120,10 +123,18 @@ impl Client { ) = mpsc::channel(); 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(); + // Initialize PeerRegistry with the actor system + { + let mut ctx = self.context.write().unwrap(); + ctx.sender = Some(sender.clone()); + let peer_registry = + PeerRegistry::new(ctx.actor_system.clone(), sender.clone()); + ctx.peer_registry = Some(peer_registry); + } + let client_sender = sender.clone(); self.server = match Server::new( @@ -181,9 +192,11 @@ impl Client { #[allow(dead_code)] pub fn remove_peer(&self, username: &str) { - let mut context = self.context.write().unwrap(); - if let Some(peer) = context.peers.remove(username) { - drop(peer); + let context = self.context.read().unwrap(); + if let Some(ref registry) = context.peer_registry { + if let Some(handle) = registry.remove_peer(username) { + let _ = handle.stop(); + } } } @@ -218,6 +231,7 @@ impl Client { size: u64, download_directory: String, ) -> Result { + info!("[client] Downloading {} from {}", filename, username); let start_time = Instant::now(); let hash = md5::md5(&filename); @@ -239,11 +253,15 @@ impl Client { let mut context = self.context.write().unwrap(); context.download_tokens.insert(token, download.clone()); - let download_initiated = context - .peers - .get(&username) - .map(|p| p.queue_upload(download.filename)) - .is_some(); + + let download_initiated = + if let Some(ref registry) = context.peer_registry { + registry + .queue_upload(&username, download.filename.clone()) + .is_ok() + } else { + false + }; drop(context); @@ -266,235 +284,321 @@ impl Client { }) } + fn process_failed_uploads( + client_context: Arc>, + username: &str, + filename: Option<&str>, + ) { + let context = client_context.read().unwrap(); + let download_tokens = &context.download_tokens; + let failed_downloads: Vec<_> = download_tokens + .iter() + .filter(|(_, download)| { + download.username == username + && (filename.is_none_or(|f| download.filename == *f)) + }) + .map(|(token, download)| { + download.sender.send(DownloadStatus::Failed).unwrap(); + *token + }) + .collect(); + + failed_downloads.iter().for_each(|token| { + client_context + .write() + .unwrap() + .download_tokens + .remove(token); + }); + } + fn listen_to_client_operations( reader: Receiver, client_context: Arc>, own_username: String, ) { thread::spawn(move || loop { - if let Ok(operation) = reader.recv() { - match operation { - ClientOperation::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) => { - let mut context = client_context.write().unwrap(); - if let Some(peer) = context.peers.remove(&username) { - drop(peer); // Explicitly drop to trigger cleanup - } - } - ClientOperation::PierceFireWall(peer) => { - Self::pierce_firewall( - 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(); + match reader.recv() { + Ok(operation) => { + match operation { + ClientOperation::ConnectToPeer(peer) => { + let client_context_clone = client_context.clone(); + let own_username_clone = own_username.clone(); - trace!( + // 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::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( + client_context.clone(), + &username, + None, + ); + } + } + ClientOperation::PierceFireWall(peer) => { + Self::pierce_firewall( + 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(); + + 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 { - Some(filename) => { - match download_peer.download_file( - client_context_clone, - Some(download.clone()), - None - ) { - Ok((download, filename)) => { - download.sender.send(DownloadStatus::Completed).unwrap(); - info!("Successfully downloaded {} bytes to {}", download.size, filename); - } - Err(e) => { - download.sender.send(DownloadStatus::Failed).unwrap(); - error!( - "Failed to download file '{}' from {}:{} (token: {}) - Error: {}", - filename, peer.host, peer.port, download.token, e - ); - } - } - } - None => error!("Cant find filename to save download: {:?}", download.filename), - } - }); - } - None => { - error!( - "Can't find download with token {:?}", - token - ); - } - }; - } - ClientOperation::NewPeer(new_peer) => { - if client_context - .read() - .unwrap() - .peers - .contains_key(&new_peer.username) - { - 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 addr = new_peer.tcp_stream.peer_addr().unwrap(); - let host = addr.ip().to_string(); - let port: u32 = addr.port().into(); - - 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, - }; - - 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 + 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, + Some(download.clone()), + None + ) { + Ok((download, filename)) => { + download.sender.send(DownloadStatus::Completed).unwrap(); + info!("Successfully downloaded {} bytes to {}", download.size, filename); + } + Err(e) => { + download.sender.send(DownloadStatus::Failed).unwrap(); + error!( + "Failed to download file '{}' from {}:{} (token: {}) - Error: {}", + filename, peer.host, peer.port, download.token, e + ); + } + } + } + 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 let Some(_peer) = - client_context.read().unwrap().peers.get(&username) - { - // 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, + 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 addr = new_peer.tcp_stream.peer_addr().unwrap(); + let host = addr.ip().to_string(); + let port: u32 = addr.port().into(); + + 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::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, + 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.remove(&key); + } + } + ClientOperation::UploadFailed(username, filename) => { + Self::process_failed_uploads( + client_context.clone(), + &username, + Some(&filename), ); } } - 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.remove(&key); - } - } + } + Err(e) => { + error!("[client] Channel recv error: {:?}", e); + break; } } }); @@ -506,76 +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(); - let sender_clone = sender; - unlocked_context.thread_pool.execute(move || { - trace!("[client] connecting to {}, with connection_type: {}, and token {:?}", peer.username, peer.connection_type, peer.token); - match peer.connection_type { - ConnectionType::P => { - let default_peer = - DefaultPeer::new(peer_clone, 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 connect_result = match stream { - Some(s) => default_peer.connect_with_socket(s, None), - _ => default_peer.connect() - }; - - match connect_result { - Ok(p) => { - trace!("[client] connected to: {}", peer.username); - client_context2.write().unwrap().peers.insert(peer.username, p); - } - Err(e) => { - trace!( - "Can't connect to {:?} {:?}:{:?} {:?} - {:?}", - peer.username, - peer.host, - peer.port, - peer.connection_type, - e - ); - } - } - } - - 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(); - } - Err(e) => { - trace!("[client] failed to download: {}", e); - } - } - } - ConnectionType::D => { - error!("ConnectionType::D not implemented") + 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!( + "Failed to spawn peer actor for {:?}: {:?}", + username, + e + ); } } - }); - } else { - debug!("[client] Peer already connected: {}", peer.username); + } 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_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/error.rs b/src/error.rs index 74fddab..81adadd 100644 --- a/src/error.rs +++ b/src/error.rs @@ -3,7 +3,6 @@ use std::{error::Error, fmt}; /// Custom error type for the Soulseek download library #[derive(Debug)] pub enum SoulseekRs { - /// Network-related errors (I/O, connection failures) NetworkError(std::io::Error), /// Authentication failed during login AuthenticationFailed, diff --git a/src/lib.rs b/src/lib.rs index 558d920..7bedec5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,5 @@ // Core modules +pub mod actor; pub mod client; pub mod dispatcher; pub mod error; diff --git a/src/main.rs b/src/main.rs index 64c7710..bbc4458 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,14 +6,16 @@ fn main() { PeerAddress::new(String::from("server.slsknet.org"), 2242), String::from("xxxxxx"), String::from("xxxxxx"), - true, + false, None, ); client.connect(); match client.login() { Ok(_) => { - match client.search("michel test file", Duration::from_secs(4)) { + match client + .search("Eternity Alex Kassian", Duration::from_secs(30)) + { Ok(results) => { if let Some(result) = results.iter().find(|r| !r.files.is_empty()) diff --git a/src/message/handlers.rs b/src/message/handlers.rs index c8f4432..0f343c8 100644 --- a/src/message/handlers.rs +++ b/src/message/handlers.rs @@ -3,12 +3,12 @@ use std::collections::HashMap; use crate::message::Message; use std::sync::mpsc::Sender; -pub trait MessageHandler { +pub trait MessageHandler: Send { fn get_code(&self) -> u8; fn handle(&self, message: &mut Message, sender: Sender); } pub struct Handlers { - handlers: HashMap>>, + handlers: HashMap + Send>>, } impl Default for Handlers { @@ -31,7 +31,10 @@ impl Handlers { self.handlers.insert(handler.get_code(), Box::new(handler)); self } - pub fn get_handler(&self, code: u8) -> Option<&dyn MessageHandler> { + pub fn get_handler( + &self, + code: u8, + ) -> Option<&(dyn MessageHandler + Send)> { self.handlers.get(&code).map(|v| &**v) } } diff --git a/src/message/peer/file_search_response.rs b/src/message/peer/file_search_response.rs index 539c171..33f09f1 100644 --- a/src/message/peer/file_search_response.rs +++ b/src/message/peer/file_search_response.rs @@ -1,21 +1,21 @@ use crate::message::{Message, MessageHandler}; -use crate::peer::PeerOperation; +use crate::peer::PeerMessage; use crate::types::FileSearchResult; use std::sync::mpsc::Sender; pub struct FileSearchResponse; -impl MessageHandler for FileSearchResponse { +impl MessageHandler for FileSearchResponse { fn get_code(&self) -> u8 { 9 } - fn handle(&self, message: &mut Message, sender: Sender) { + fn handle(&self, message: &mut Message, sender: Sender) { let file_search = match FileSearchResult::new_from_message(message) { Ok(result) => result, Err(_) => return, // Skip malformed search results }; sender - .send(PeerOperation::FileSearchResult(file_search)) + .send(PeerMessage::FileSearchResult(file_search)) .unwrap(); } } diff --git a/src/message/peer/get_share_file_list.rs b/src/message/peer/get_share_file_list.rs index 60f7d11..683b94f 100644 --- a/src/message/peer/get_share_file_list.rs +++ b/src/message/peer/get_share_file_list.rs @@ -1,17 +1,17 @@ use crate::{ message::{server::MessageFactory, Message, MessageHandler}, - peer::PeerOperation, + peer::PeerMessage, }; use std::sync::mpsc::Sender; pub struct GetShareFileList; -impl MessageHandler for GetShareFileList { +impl MessageHandler for GetShareFileList { fn get_code(&self) -> u8 { 4 } - fn handle(&self, _message: &mut Message, sender: Sender) { + fn handle(&self, _message: &mut Message, sender: Sender) { let message = MessageFactory::build_shared_folders_message(100, 800); - sender.send(PeerOperation::SendMessage(message)).unwrap(); + sender.send(PeerMessage::SendMessage(message)).unwrap(); } } diff --git a/src/message/peer/peer_init.rs b/src/message/peer/peer_init.rs index 5d37f77..ea1ff0a 100644 --- a/src/message/peer/peer_init.rs +++ b/src/message/peer/peer_init.rs @@ -2,17 +2,17 @@ use std::sync::mpsc::Sender; use crate::{ message::{Message, MessageHandler}, - peer::PeerOperation, + peer::PeerMessage, trace, }; pub struct PeerInit; -impl MessageHandler for PeerInit { +impl MessageHandler for PeerInit { fn get_code(&self) -> u8 { 1 } - fn handle(&self, message: &mut Message, sender: Sender) { + fn handle(&self, message: &mut Message, sender: Sender) { message.set_pointer(4); let _message_code = message.read_int8(); let username = message.read_string(); @@ -25,6 +25,6 @@ impl MessageHandler for PeerInit { token ); - sender.send(PeerOperation::SetUsername(username)).unwrap(); + sender.send(PeerMessage::SetUsername(username)).unwrap(); } } diff --git a/src/message/peer/place_in_queue_response.rs b/src/message/peer/place_in_queue_response.rs index 08ac7cf..0d4a38e 100644 --- a/src/message/peer/place_in_queue_response.rs +++ b/src/message/peer/place_in_queue_response.rs @@ -1,22 +1,22 @@ use crate::{ message::{Message, MessageHandler}, - peer::PeerOperation, + peer::PeerMessage, }; use std::sync::mpsc::Sender; pub struct PlaceInQueueResponse; -impl MessageHandler for PlaceInQueueResponse { +impl MessageHandler for PlaceInQueueResponse { fn get_code(&self) -> u8 { 43 } - fn handle(&self, message: &mut Message, sender: Sender) { + fn handle(&self, message: &mut Message, sender: Sender) { let filename = message.read_string(); let place = message.read_int32(); sender - .send(PeerOperation::PlaceInQueueResponse { filename, place }) + .send(PeerMessage::PlaceInQueueResponse { filename, place }) .unwrap(); } } diff --git a/src/message/peer/transfer_request.rs b/src/message/peer/transfer_request.rs index 828b6b0..a4d7499 100644 --- a/src/message/peer/transfer_request.rs +++ b/src/message/peer/transfer_request.rs @@ -1,20 +1,18 @@ use crate::{ message::{Message, MessageHandler}, - peer::PeerOperation, + peer::PeerMessage, types::Transfer, }; use std::sync::mpsc::Sender; pub struct TransferRequest; -impl MessageHandler for TransferRequest { +impl MessageHandler for TransferRequest { fn get_code(&self) -> u8 { 40 } - fn handle(&self, message: &mut Message, sender: Sender) { + fn handle(&self, message: &mut Message, sender: Sender) { let transfer = Transfer::new_from_message(message); - sender - .send(PeerOperation::TransferRequest(transfer)) - .unwrap(); + sender.send(PeerMessage::TransferRequest(transfer)).unwrap(); } } diff --git a/src/message/peer/transfer_response.rs b/src/message/peer/transfer_response.rs index baa19d2..7191b09 100644 --- a/src/message/peer/transfer_response.rs +++ b/src/message/peer/transfer_response.rs @@ -1,23 +1,23 @@ use crate::{ message::{Message, MessageHandler}, - peer::PeerOperation, + peer::PeerMessage, }; use std::sync::mpsc::Sender; pub struct TransferResponse; -impl MessageHandler for TransferResponse { +impl MessageHandler for TransferResponse { fn get_code(&self) -> u8 { 41 } - fn handle(&self, message: &mut Message, sender: Sender) { + fn handle(&self, message: &mut Message, sender: Sender) { let token = message.read_int32(); let allowed = message.read_int8(); let reason = (allowed == 0).then(|| message.read_string()); sender - .send(PeerOperation::TransferResponse { + .send(PeerMessage::TransferResponse { token, allowed: allowed == 1, reason, diff --git a/src/message/peer/upload_failed.rs b/src/message/peer/upload_failed.rs index d03ea53..27f1ec9 100644 --- a/src/message/peer/upload_failed.rs +++ b/src/message/peer/upload_failed.rs @@ -1,17 +1,17 @@ use crate::info; use crate::{ message::{Message, MessageHandler}, - peer::PeerOperation, + peer::PeerMessage, types::UploadFailed, }; use std::sync::mpsc::Sender; pub struct UploadFailedHandler; -impl MessageHandler for UploadFailedHandler { +impl MessageHandler for UploadFailedHandler { fn get_code(&self) -> u8 { 46 } - fn handle(&self, message: &mut Message, _sender: Sender) { + fn handle(&self, message: &mut Message, _sender: Sender) { let upload_failed = UploadFailed::new_from_message(message); info!("Upload failed for ${}", upload_failed.filename); } diff --git a/src/peer/default_peer.rs b/src/peer/default_peer.rs deleted file mode 100644 index 85faac3..0000000 --- a/src/peer/default_peer.rs +++ /dev/null @@ -1,450 +0,0 @@ -use crate::dispatcher::MessageDispatcher; -use crate::message::peer::{FileSearchResponse, PeerInit}; -use crate::message::peer::{ - GetShareFileList, PlaceInQueueResponse, TransferRequest, TransferResponse, - UploadFailedHandler, -}; -use crate::message::server::MessageFactory; -use crate::message::{Handlers, Message, MessageReader, MessageType}; -use crate::types::{Download, FileSearchResult, Transfer}; - -use core::result::Result::Ok; -use std::sync::mpsc::{Receiver, Sender}; -use std::sync::{mpsc, Arc, RwLock}; -use std::thread::{self, JoinHandle}; - -use crate::client::ClientOperation; -use crate::peer::Peer; -use crate::{debug, error, trace, warn}; -use std::io::{self, Write}; -use std::net::TcpStream; -use std::net::ToSocketAddrs; -use std::time::Duration; - -#[allow(dead_code)] -#[derive(Debug)] -pub struct DefaultPeer { - peer: Arc>, - peer_channel: Option>, - client_channel: Sender, - read_thread: Option>, - write_thread: Option>, -} - -#[allow(dead_code)] -pub enum PeerOperation { - SendMessage(Message), - FileSearchResult(FileSearchResult), - TransferRequest(Transfer), - TransferResponse { - token: u32, - allowed: bool, - reason: Option, - }, - PlaceInQueueResponse { - filename: String, - place: u32, - }, - SetUsername(String), -} - -impl DefaultPeer { - pub fn new(peer: Peer, client_channel: Sender) -> Self { - Self { - peer: Arc::new(RwLock::new(peer)), - peer_channel: None, - client_channel, - read_thread: None, - write_thread: None, - } - } - pub fn disconnect(&mut self) { - let username = self.peer.read().unwrap().username.clone(); - debug!("[default_peer:{}] disconnect", username); - if let Err(e) = self - .client_channel - .send(ClientOperation::PeerDisconnected(username)) - { - error!("Failed to send disconnect notification: {}", e); - } - - self.peer_channel.take(); - } - - 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(); - message.extend_from_slice(&token.to_le_bytes()); - stream.write_all(&message).unwrap(); - } - - stream.set_read_timeout(Some(Duration::from_secs(5)))?; - stream.set_write_timeout(Some(Duration::from_secs(5)))?; - stream.set_nodelay(true)?; - - trace!( - "[default_peer:{}] connect_with_socket: direct", - self.peer.read().unwrap().username - ); - self.start_read_write_loops(stream, reader)?; - - Ok(self) - } - - pub fn connect(mut self) -> Result { - { - let peer = self.peer.read().unwrap(); - trace!( - "[default_peer] Connecting to {} on port {}", - peer.host, - peer.port - ); - } - let socket_address = { - let peer = self.peer.read().unwrap(); - format!("{}:{}", peer.host, peer.port) - } - .to_socket_addrs()? - .next() - .ok_or_else(|| { - io::Error::new(io::ErrorKind::InvalidInput, "Invalid address") - })?; - - trace!( - "[default_peer:{}] connect: indirect", - self.peer.read().unwrap().username - ); - let mut stream = TcpStream::connect_timeout( - &socket_address, - Duration::from_secs(20), - )?; - - stream.set_read_timeout(Some(Duration::from_secs(5)))?; - stream.set_write_timeout(Some(Duration::from_secs(5)))?; - stream.set_nodelay(true)?; - - if let Some(token) = self.peer.read().unwrap().token { - stream - .write_all(&MessageFactory::build_watch_user(token).get_data()) - .unwrap(); - } - self.start_read_write_loops(stream, None)?; - - Ok(self) - } - - fn start_read_write_loops( - &mut self, - stream: TcpStream, - reader: Option, - ) -> Result<(), io::Error> { - let (peer_sender, peer_reader): ( - Sender, - Receiver, - ) = mpsc::channel(); - - self.peer_channel = Some(peer_sender.clone()); - - 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; - - 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 || { - let mut handlers = Handlers::new(); - handlers.register_handler(FileSearchResponse); - handlers.register_handler(TransferRequest); - handlers.register_handler(TransferResponse); - handlers.register_handler(GetShareFileList); - handlers.register_handler(UploadFailedHandler); - handlers.register_handler(PlaceInQueueResponse); - handlers.register_handler(PeerInit); - - let dispatcher = MessageDispatcher::new( - "default_peer".to_string(), - 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 - ); - } - - loop { - let mut should_terminate = false; - loop { - match buffered_reader.extract_message() { - Ok(Some(mut message)) => { - trace!( - "[default_peer:{}] ← {:?}", - peer.read().unwrap().username, - message - .get_message_name( - MessageType::Peer, - message.get_message_code() as u32 - ) - .map_err(|e| e.to_string()) - ); - dispatcher.dispatch(&mut message) - } - Err(e) => { - warn!( - "Error extracting message in default peer: {}. Terminating read loop.", - e - ); - let username = - peer.read().unwrap().username.clone(); - let _ = client_channel_for_read.send( - ClientOperation::PeerDisconnected(username), - ); - should_terminate = true; - 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(); - let peer_channel = self.peer_channel.clone(); - let peer_clone_for_write = self.peer.clone(); - - self.write_thread = Some(thread::spawn(move || loop { - match peer_reader.recv() { - Ok(operation) => match operation { - PeerOperation::SendMessage(message) => { - trace!( - "[default_peer:{}] ➡ {:?} - {:?}", - peer_clone_for_write.read().unwrap().username, - message - .get_message_name( - MessageType::Peer, - u32::from_le_bytes( - message - .get_slice(0, 4) - .try_into() - .unwrap() - ) - ) - .map_err(|e| e.to_string()), - u32::from_le_bytes( - message.get_slice(0, 4).try_into().unwrap() - ) - ); - - if let Err(e) = - write_stream.write_all(&message.get_buffer()) - { - error!("Error writing message to stream: {} - {}. Terminating write loop.", peer_clone_for_write.read().unwrap().username, e); - let _ = client_channel.send( - ClientOperation::PeerDisconnected( - peer_clone_for_write - .read() - .unwrap() - .username - .clone(), - ), - ); - break; - } - write_stream.flush().unwrap(); - } - PeerOperation::FileSearchResult(file_search) => { - client_channel - .send(ClientOperation::SearchResult(file_search)) - .unwrap(); - } - PeerOperation::TransferRequest(transfer) => { - debug!( - "[default_peer:{}] TransferRequest for {}", - peer_clone_for_write.read().unwrap().username, - transfer.token - ); - client_channel - .send(ClientOperation::UpdateDownloadTokens( - transfer.clone(), - peer_clone.read().unwrap().username.clone(), - )) - .unwrap(); - - let transfer_response = - MessageFactory::build_transfer_response_message( - transfer.clone(), - ); - - trace!( - "[default_peer:{}] TransferResponse for {:?}", - peer_clone_for_write.read().unwrap().username, - transfer_response.get_buffer() - ); - if let Some(sender) = peer_channel.clone() { - sender - .send(PeerOperation::SendMessage( - transfer_response, - )) - .unwrap(); - } - } - PeerOperation::TransferResponse { - token, - allowed, - reason, - } => { - debug!( - "[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 - ); - } - } else { - debug!("[default_peer:{}] Transfer allowed, ready to connect with token {:}",peer_clone_for_write.read().unwrap().username, token); - client_channel - .send(ClientOperation::DownloadFromPeer( - token, - peer_clone.read().unwrap().clone(), - allowed, - )) - .unwrap(); - } - } - PeerOperation::PlaceInQueueResponse { filename, place } => { - debug!( - "[default_peer:{}] Place in queue response - file: {}, place: {}", - peer_clone_for_write.read().unwrap().username, filename, place - ); - } - PeerOperation::SetUsername(username) => { - trace!( - "[default_peer:{}] SetUsername: {}", - peer_clone_for_write.read().unwrap().username, - username - ); - - peer_clone.write().unwrap().username = username; - } - }, - Err(_) => { - debug!("[default_peer:{}] Peer channel closed. Terminating write loop.", peer_clone_for_write.read().unwrap().username); - break; - } - } - })); - - 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, - ) -> Result<(), io::Error> { - let message = MessageFactory::build_transfer_request_message( - &download.filename, - download.token, - ); - if let Some(sender) = &self.peer_channel { - sender.send(PeerOperation::SendMessage(message)).unwrap(); - } - Ok(()) - } -} - -impl Drop for DefaultPeer { - fn drop(&mut self) { - self.peer_channel = None; - trace!("[default_peer:{}] drop", self.peer.read().unwrap().username); - - if let Some(handle) = self.read_thread.take() { - debug!("Joining read thread..."); - match handle.join() { - Ok(_) => debug!("Read thread joined successfully."), - Err(e) => error!("Read thread panicked: {:?}", e), - } - } - - if let Some(handle) = self.write_thread.take() { - debug!("Joining write thread..."); - match handle.join() { - Ok(_) => debug!("Write thread joined successfully."), - Err(e) => error!("Write thread panicked: {:?}", e), - } - } - } -} diff --git a/src/peer/listen.rs b/src/peer/listen.rs index 82a2918..080fc31 100644 --- a/src/peer/listen.rs +++ b/src/peer/listen.rs @@ -7,7 +7,7 @@ use std::thread; use crate::client::{ClientContext, ClientOperation}; use crate::message::{Message, MessageReader}; -use crate::peer::{ConnectionType, DefaultPeer, DownloadPeer, Peer}; +use crate::peer::{ConnectionType, DownloadPeer, Peer}; use crate::types::Download; use crate::{debug, error, info, trace, DownloadStatus}; @@ -15,6 +15,7 @@ const PEER_INIT_MESSAGE_CODE: u8 = 1; #[derive(Clone)] struct ConnectionContext { + #[allow(dead_code)] client_sender: Sender, client_context: Arc>, own_username: String, @@ -74,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", @@ -105,17 +99,23 @@ fn handle_peer_connection( peer: Peer, stream: TcpStream, reader: MessageReader, - client_sender: Sender, - peer_ip: &str, - peer_port: u16, + context: &ConnectionContext, + _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); + let client_context = context.client_context.read().unwrap(); + if let Some(ref registry) = client_context.peer_registry { + match registry.register_peer(peer.clone(), Some(stream), Some(reader)) { + Ok(_) => (), + Err(e) => { + error!( + "Failed to spawn peer actor for {:?}: {:?}", + peer.username, e + ); + } + } + } else { + error!("PeerRegistry not initialized"); } } @@ -201,8 +201,9 @@ fn handle_incoming_connection(stream: TcpStream, context: ConnectionContext) { ); let peer = Peer::new( - init_data.username.clone(), - ConnectionType::P, + format!("{}:direct", init_data.username), + // init_data.username.clone(), + init_data.connection_type.clone(), peer_ip.clone(), peer_port.into(), None, @@ -213,22 +214,25 @@ fn handle_incoming_connection(stream: TcpStream, context: ConnectionContext) { 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, + peer, stream, reader, &context, &peer_ip, peer_port, ), + + ConnectionType::F => { + thread::spawn(move || { + trace!( + "[listener:{peer_ip}:{peer_port}] handling file connection in thread" + ); + 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. "); } @@ -265,7 +269,7 @@ impl Listen { }; let context = context.clone(); - thread::spawn(move || handle_incoming_connection(stream, context)); + handle_incoming_connection(stream, context); } } } diff --git a/src/peer/mod.rs b/src/peer/mod.rs index 7032047..9386c18 100644 --- a/src/peer/mod.rs +++ b/src/peer/mod.rs @@ -1,8 +1,10 @@ -mod default_peer; mod download_peer; pub mod listen; -pub use default_peer::{DefaultPeer, PeerOperation}; +// Export actor types +pub use crate::actor::peer_actor::{PeerActor, PeerMessage}; +pub use crate::actor::peer_registry::PeerRegistry; + pub use download_peer::DownloadPeer; use crate::message::Message; diff --git a/src/server.rs b/src/server.rs index 4ad5082..795773c 100644 --- a/src/server.rs +++ b/src/server.rs @@ -341,17 +341,11 @@ impl Server { if let Ok(operation) = server_channel.recv() { match operation { ServerOperation::ConnectToPeer(peer) => { - debug!( - "[server] ConnectToPeer {} - ConnectionType {}", - peer.username, peer.connection_type - ); - if let Some(op) = match peer.connection_type { - ConnectionType::P => { - Some(ClientOperation::ConnectToPeer(peer)) - } - ConnectionType::F => { - Some(ClientOperation::ConnectToPeer(peer)) + ConnectionType::P | ConnectionType::F => { + Some(ClientOperation::ConnectToPeer( + peer.clone(), + )) } ConnectionType::D => None, } {