diff --git a/soulseek-rs-lib/src/actor/peer_actor.rs b/soulseek-rs-lib/src/actor/peer_actor.rs index b41032f..f81fec2 100644 --- a/soulseek-rs-lib/src/actor/peer_actor.rs +++ b/soulseek-rs-lib/src/actor/peer_actor.rs @@ -9,6 +9,7 @@ use crate::message::server::MessageFactory; use crate::message::{Handlers, Message, MessageReader, MessageType}; use crate::peer::Peer; use crate::types::{Download, SearchResult, Transfer}; +use crate::utils::lock::RwLockExt; use crate::{debug, error, trace, warn}; use std::io::{self, Error, Write}; @@ -80,6 +81,26 @@ impl PeerActor { self.self_handle = Some(handle); } + fn peer_username(&self) -> String { + match self.peer.read_safe() { + Ok(p) => p.username.clone(), + Err(e) => { + error!("[peer_actor] peer lock poisoned: {}", e); + "".to_string() + } + } + } + + fn peer_snapshot(&self) -> Option { + match self.peer.read_safe() { + Ok(p) => Some(p.clone()), + Err(e) => { + error!("[peer_actor] peer lock poisoned: {}", e); + None + } + } + } + fn initialize_dispatcher(&mut self) { let (dispatcher_sender, dispatcher_receiver) = std::sync::mpsc::channel::(); @@ -135,23 +156,35 @@ impl PeerActor { self.send_message(message); } PeerMessage::FileSearchResult(file_search) => { - self.client_channel + if let Err(e) = self + .client_channel .send(ClientOperation::SearchResult(file_search)) - .unwrap(); + { + error!( + "[peer:{}] failed to forward search result: {}", + self.peer_username(), + e + ); + } } PeerMessage::TransferRequest(transfer) => { - let username = self.peer.read().unwrap().username.clone(); + let username = self.peer_username(); debug!( "[peer:{}] TransferRequest for {}", username, transfer.token ); - self.client_channel - .send(ClientOperation::UpdateDownloadTokens( + if let Err(e) = self.client_channel.send( + ClientOperation::UpdateDownloadTokens( transfer.clone(), username.clone(), - )) - .unwrap(); + ), + ) { + error!( + "[peer:{}] failed to send UpdateDownloadTokens: {}", + username, e + ); + } let transfer_response = MessageFactory::build_transfer_response_message(transfer); @@ -171,7 +204,7 @@ impl PeerActor { allowed, reason, } => { - let username = self.peer.read().unwrap().username.clone(); + let username = self.peer_username(); debug!( "[peer:{}] transfer response token: {} allowed: {}", username, token, allowed @@ -189,17 +222,25 @@ impl PeerActor { "[peer:{}] Transfer allowed, ready to connect with token {:}", username, token ); - self.client_channel - .send(ClientOperation::DownloadFromPeer( + let Some(peer_snapshot) = self.peer_snapshot() else { + return; + }; + if let Err(e) = self.client_channel.send( + ClientOperation::DownloadFromPeer( token, - self.peer.read().unwrap().clone(), + peer_snapshot, allowed, - )) - .unwrap(); + ), + ) { + error!( + "[peer:{}] failed to send DownloadFromPeer: {}", + username, e + ); + } } } PeerMessage::PlaceInQueueResponse { filename, place } => { - let username = self.peer.read().unwrap().username.clone(); + let username = self.peer_username(); debug!( "[peer:{}] Place in queue response - file: {}, place: {}", username, filename, place @@ -208,10 +249,15 @@ impl PeerActor { PeerMessage::SetUsername(username) => { trace!( "[peer:{}] SetUsername: {}", - self.peer.read().unwrap().username, + self.peer_username(), username ); - self.peer.write().unwrap().username = username; + match self.peer.write_safe() { + Ok(mut p) => p.username = username, + Err(e) => { + error!("[peer_actor] SetUsername write: {}", e) + } + } } PeerMessage::QueueUpload(filename) => { let message = @@ -229,9 +275,15 @@ impl PeerActor { self.process_read(); } PeerMessage::UploadFailed(username, filename) => { - self.client_channel + if let Err(e) = self + .client_channel .send(ClientOperation::UploadFailed(username, filename)) - .unwrap(); + { + error!( + "[peer_actor] failed to forward UploadFailed: {}", + e + ); + } } } } @@ -251,14 +303,15 @@ impl PeerActor { 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 - ); + if let Ok(peer_lock) = self.peer.read_safe() { + debug!( + "Read operation timed out for peer actor {:}:{:}", + peer_lock.host, peer_lock.port + ); + } } Err(e) => { - let username = self.peer.read().unwrap().username.clone(); + let username = self.peer_username(); error!( "[peer:{}] Error reading from peer: {} (kind: {:?}). Disconnecting.", username, @@ -274,7 +327,7 @@ impl PeerActor { } fn extract_and_process_messages(&mut self) { - let username = self.peer.read().unwrap().username.clone(); + let username = self.peer_username(); let mut extracted_count = 0; loop { match self.reader.extract_message() { @@ -315,6 +368,7 @@ impl PeerActor { } fn send_message(&mut self, message: Message) { + let username = self.peer_username(); let stream = match self.stream.as_mut() { Some(s) => s, None => { @@ -323,7 +377,6 @@ impl PeerActor { } }; - let username = self.peer.read().unwrap().username.clone(); trace!( "[peer:{}] ➡ {:?}", username, @@ -356,7 +409,7 @@ impl PeerActor { } fn disconnect_with_error(&mut self, error: Error) { - let username = self.peer.read().unwrap().username.clone(); + let username = self.peer_username(); debug!("[peer:{}] disconnect", username); self.stream.take(); @@ -368,7 +421,7 @@ impl PeerActor { } } fn disconnect(&mut self) { - let username = self.peer.read().unwrap().username.clone(); + let username = self.peer_username(); debug!("[peer:{}] disconnect", username); self.stream.take(); @@ -382,11 +435,13 @@ impl PeerActor { } 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 (username, host, port) = match self.peer.read_safe() { + Ok(peer) => (peer.username.clone(), peer.host.clone(), peer.port), + Err(e) => { + error!("[peer_actor] initiate_connection peer lock: {}", e); + return false; + } + }; let socket_addr = format!("{}:{}", host, port).parse::(); @@ -438,7 +493,7 @@ impl PeerActor { return; }; - let username = self.peer.read().unwrap().username.clone(); + let username = self.peer_username(); if since.elapsed() > Duration::from_secs(20) { error!("[peer:{}] Connection timeout after 20 seconds", username); @@ -467,10 +522,16 @@ impl PeerActor { } fn on_connection_established(&mut self) { - let peer = self.peer.read().unwrap(); - let username = peer.username.clone(); - let token = peer.token; - drop(peer); + let (username, token) = match self.peer.read_safe() { + Ok(peer) => (peer.username.clone(), peer.token), + Err(e) => { + error!( + "[peer_actor] on_connection_established peer lock: {}", + e + ); + return; + } + }; let Some(ref mut stream) = self.stream else { return; @@ -520,7 +581,7 @@ impl Actor for PeerActor { } fn on_stop(&mut self) { - let username = self.peer.read().unwrap().username.clone(); + let username = self.peer_username(); trace!("[peer:{}] actor stopping", username); self.disconnect(); } diff --git a/soulseek-rs-lib/src/actor/peer_registry.rs b/soulseek-rs-lib/src/actor/peer_registry.rs index a7dc1b7..6dcc1e2 100644 --- a/soulseek-rs-lib/src/actor/peer_registry.rs +++ b/soulseek-rs-lib/src/actor/peer_registry.rs @@ -1,9 +1,10 @@ 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::utils::lock::MutexExt; +use crate::{debug, error}; use std::collections::HashMap; use std::net::TcpStream; @@ -44,22 +45,36 @@ impl PeerRegistry { actor.set_self_handle(handle); }); - let mut peers = self.peers.lock().unwrap(); + let mut peers = self + .peers + .lock_safe() + .map_err(|e| format!("peer registry lock poisoned: {}", e))?; 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() + match self.peers.lock_safe() { + Ok(peers) => peers.get(username).cloned(), + Err(e) => { + error!("[peer_registry] get_peer: {}", e); + None + } + } } pub fn remove_peer( &self, username: &str, ) -> Option> { - let mut peers = self.peers.lock().unwrap(); + let mut peers = match self.peers.lock_safe() { + Ok(p) => p, + Err(e) => { + error!("[peer_registry] remove_peer: {}", e); + return None; + } + }; let handle = peers.remove(username); if handle.is_some() { @@ -70,18 +85,33 @@ impl PeerRegistry { } pub fn get_all_usernames(&self) -> Vec { - let peers = self.peers.lock().unwrap(); - peers.keys().cloned().collect() + match self.peers.lock_safe() { + Ok(peers) => peers.keys().cloned().collect(), + Err(e) => { + error!("[peer_registry] get_all_usernames: {}", e); + Vec::new() + } + } } pub fn count(&self) -> usize { - let peers = self.peers.lock().unwrap(); - peers.len() + match self.peers.lock_safe() { + Ok(peers) => peers.len(), + Err(e) => { + error!("[peer_registry] count: {}", e); + 0 + } + } } pub fn contains(&self, username: &str) -> bool { - let peers = self.peers.lock().unwrap(); - peers.contains_key(username) + match self.peers.lock_safe() { + Ok(peers) => peers.contains_key(username), + Err(e) => { + error!("[peer_registry] contains: {}", e); + false + } + } } pub fn send_to_peer( diff --git a/soulseek-rs-lib/src/actor/server_actor.rs b/soulseek-rs-lib/src/actor/server_actor.rs index b6dfd2b..a98c51a 100644 --- a/soulseek-rs-lib/src/actor/server_actor.rs +++ b/soulseek-rs-lib/src/actor/server_actor.rs @@ -17,6 +17,7 @@ use crate::message::{Handlers, MessageType}; use crate::message::{Message, MessageReader}; use crate::peer::ConnectionType; use crate::peer::Peer; +use crate::utils::lock::RwLockExt; use std::io::{self, Error, Write}; use std::net::{TcpStream, ToSocketAddrs}; @@ -338,9 +339,12 @@ impl ServerActor { self.dispatcher_receiver = Some(dispatcher_receiver); self.dispatcher_sender = Some(dispatcher_sender.clone()); - self.client_channel + if let Err(e) = self + .client_channel .send(ClientOperation::SetServerSender(dispatcher_sender.clone())) - .unwrap(); + { + error!("[server] failed to send SetServerSender: {}", e); + } let mut handlers = Handlers::new(); @@ -401,7 +405,7 @@ impl ServerActor { } { - logged_in = context.read().unwrap().logged_in + logged_in = context.read_safe()?.logged_in } if logged_in.is_some() { @@ -409,7 +413,8 @@ impl ServerActor { } } - if logged_in.unwrap() { + let logged_in = logged_in.ok_or(SoulseekRs::Timeout)?; + if logged_in { info!("Logged in as {}", username); self.queue_message(MessageFactory::build_shared_folders_message( 1, 499, @@ -425,7 +430,7 @@ impl ServerActor { } } - Ok(logged_in.unwrap()) + Ok(logged_in) } pub fn file_search(&mut self, token: u32, query: &str) { @@ -455,12 +460,18 @@ impl ServerActor { Some(ClientOperation::ConnectToPeer(peer.clone())) } ConnectionType::D => None, - } { - self.client_channel.send(op).unwrap(); + } && let Err(e) = self.client_channel.send(op) + { + error!("[server] failed to send ConnectToPeer: {}", e); } } ServerMessage::LoginStatus(message) => { - self.context.write().unwrap().logged_in = Some(message); + match self.context.write_safe() { + Ok(mut ctx) => ctx.logged_in = Some(message), + Err(e) => { + error!("[server] LoginStatus write: {}", e) + } + } } ServerMessage::PierceFirewall(token) => { self.send_message( @@ -525,9 +536,14 @@ impl ServerActor { break; } - if let Some(logged_in) = - context.read().unwrap().logged_in - { + let logged_in = match context.read_safe() { + Ok(ctx) => ctx.logged_in, + Err(e) => { + let _ = response.send(Err(e)); + break; + } + }; + if let Some(logged_in) = logged_in { let result = if logged_in { Ok(true) } else { diff --git a/soulseek-rs-lib/src/client.rs b/soulseek-rs-lib/src/client.rs index c289102..b4c1a93 100644 --- a/soulseek-rs-lib/src/client.rs +++ b/soulseek-rs-lib/src/client.rs @@ -8,7 +8,7 @@ use crate::{ error::{Result, SoulseekRs}, peer::{ConnectionType, DownloadPeer, NewPeer, Peer, listen::Listen}, types::{Download, Search, SearchResult}, - utils::{md5, thread_pool::ThreadPool}, + utils::{lock::RwLockExt, md5, thread_pool::ThreadPool}, }; use std::{ collections::HashMap, @@ -336,13 +336,13 @@ impl Client { } } - pub fn connect(&mut self) { + pub fn connect(&mut self) -> Result<()> { let (sender, message_reader): ( Sender, Receiver, ) = mpsc::channel(); - let mut ctx = self.context.write().unwrap(); + let mut ctx = self.context.write_safe()?; ctx.sender = Some(sender.clone()); let peer_registry = PeerRegistry::new(ctx.actor_system.clone(), sender.clone()); @@ -385,6 +385,8 @@ impl Client { self.context.clone(), self.username.clone(), ); + + Ok(()) } pub fn login(&self) -> Result { @@ -408,7 +410,13 @@ impl Client { #[allow(dead_code)] pub fn remove_peer(&self, username: &str) { - let context = self.context.read().unwrap(); + let context = match self.context.read_safe() { + Ok(c) => c, + Err(e) => { + error!("[client] remove_peer: {}", e); + return; + } + }; if let Some(ref registry) = context.peer_registry && let Some(handle) = registry.remove_peer(username) { @@ -437,7 +445,7 @@ impl Client { let token = u32::from_str_radix(&hash[0..5], 16)?; // Initialize new search with query string as key - self.context.write().unwrap().searches.insert( + self.context.write_safe()?.searches.insert( query.to_string(), Search { token, @@ -476,21 +484,21 @@ impl Client { pub fn get_search_results_count(&self, search_key: &str) -> usize { self.context - .read() - .unwrap() - .searches - .get(search_key) - .map(|s| s.results.len()) + .read_safe() + .ok() + .and_then(|ctx| { + ctx.searches.get(search_key).map(|s| s.results.len()) + }) .unwrap_or(0) } pub fn get_search_results(&self, search_key: &str) -> Vec { self.context - .read() - .unwrap() - .searches - .get(search_key) - .map(|s| s.results.clone()) + .read_safe() + .ok() + .and_then(|ctx| { + ctx.searches.get(search_key).map(|s| s.results.clone()) + }) .unwrap_or_default() } @@ -505,15 +513,27 @@ impl Client { } pub fn get_all_searches(&self) -> HashMap { - self.context.read().unwrap().searches.clone() + self.context + .read_safe() + .map(|ctx| ctx.searches.clone()) + .unwrap_or_default() } pub fn get_all_downloads(&self) -> Vec { - self.context.read().unwrap().get_downloads().clone() + self.context + .read_safe() + .map(|ctx| ctx.get_downloads().clone()) + .unwrap_or_default() } pub fn pause_download(&self, username: &str, filename: &str) -> bool { - let mut context = self.context.write().unwrap(); + let mut context = match self.context.write_safe() { + Ok(c) => c, + Err(e) => { + error!("[client] pause_download: {}", e); + return false; + } + }; let Some(download) = context.get_download_by_file_mut(username, filename) else { @@ -539,7 +559,13 @@ impl Client { } pub fn resume_download(&self, username: &str, filename: &str) -> bool { - let mut context = self.context.write().unwrap(); + let mut context = match self.context.write_safe() { + Ok(c) => c, + Err(e) => { + error!("[client] resume_download: {}", e); + return false; + } + }; let Some(download) = context.get_download_by_file_mut(username, filename) else { @@ -569,10 +595,15 @@ impl Client { username: &str, filename: &str, ) -> bool { - self.context - .write() - .unwrap() - .remove_queued_download_by_file(username, filename) + match self.context.write_safe() { + Ok(mut ctx) => { + ctx.remove_queued_download_by_file(username, filename) + } + Err(e) => { + error!("[client] remove_queued_download: {}", e); + false + } + } } pub fn download( @@ -602,7 +633,7 @@ impl Client { sender: download_sender, }; - let mut context = self.context.write().unwrap(); + let mut context = self.context.write_safe()?; context.add_download(download.clone()); let download_initiated = @@ -619,8 +650,7 @@ impl Client { if !download_initiated { let _ = download.sender.send(DownloadStatus::Failed); self.context - .write() - .unwrap() + .write_safe()? .update_download_with_status(token, DownloadStatus::Failed); } @@ -632,25 +662,38 @@ impl Client { username: &str, filename: Option<&str>, ) { - let context = client_context.read().unwrap(); - let failed_downloads: Vec<_> = context - .get_downloads() - .iter() - .filter(|download| { - download.username == username - && (filename.is_none_or(|f| download.filename == *f)) - }) - .map(|download| { - let _ = download.sender.send(DownloadStatus::Failed); - download.token - }) - .collect(); - drop(context); + let failed_downloads: Vec<_> = match client_context.read_safe() { + Ok(context) => context + .get_downloads() + .iter() + .filter(|download| { + download.username == username + && (filename.is_none_or(|f| download.filename == *f)) + }) + .map(|download| { + let _ = download.sender.send(DownloadStatus::Failed); + download.token + }) + .collect(), + Err(e) => { + error!("[client] process_failed_uploads read: {}", e); + return; + } + }; failed_downloads.iter().for_each(|token| { - let mut context = client_context.write().unwrap(); - context.update_download_with_status(*token, DownloadStatus::Failed); - context.remove_download(*token); + match client_context.write_safe() { + Ok(mut context) => { + context.update_download_with_status( + *token, + DownloadStatus::Failed, + ); + context.remove_download(*token); + } + Err(e) => { + error!("[client] process_failed_uploads write: {}", e); + } + } }); } @@ -683,8 +726,18 @@ impl Client { "[client] SearchResult {:?}", search_result ); - let mut context = - client_context.write().unwrap(); + let mut context = match client_context + .write_safe() + { + Ok(c) => c, + Err(e) => { + error!( + "[client] SearchResult write: {}", + e + ); + continue; + } + }; let result_token = search_result.token; // Find the search with matching token @@ -699,7 +752,16 @@ impl Client { username, error, ) => { - let context = client_context.read().unwrap(); + let context = match client_context.read_safe() { + Ok(c) => c, + Err(e) => { + error!( + "[client] PeerDisconnected read: {}", + e + ); + continue; + } + }; if let Some(ref registry) = context.peer_registry && let Some(handle) = @@ -731,12 +793,19 @@ impl Client { peer, allowed, ) => { - let maybe_download = { - let client_context = - client_context.read().unwrap(); - client_context + let maybe_download = match client_context + .read_safe() + { + Ok(ctx) => ctx .get_download_by_token(token) - .cloned() + .cloned(), + Err(e) => { + error!( + "[client] DownloadFromPeer read: {}", + e + ); + continue; + } }; let own_username = own_username.clone(); let client_context_clone = @@ -777,7 +846,10 @@ impl Client { filename, )) => { let _ = download.sender.send(DownloadStatus::Completed); - client_context_clone.write().unwrap().update_download_with_status(download.token, DownloadStatus::Completed); + match client_context_clone.write_safe() { + Ok(mut ctx) => ctx.update_download_with_status(download.token, DownloadStatus::Completed), + Err(e) => error!("[client] download complete write: {}", e), + } info!( "Successfully downloaded {} bytes to {}", download.size, @@ -786,7 +858,10 @@ impl Client { } Err(e) => { let _ = download.sender.send(DownloadStatus::Failed); - client_context_clone.write().unwrap().update_download_with_status(download.token, DownloadStatus::Failed); + match client_context_clone.write_safe() { + Ok(mut ctx) => ctx.update_download_with_status(download.token, DownloadStatus::Failed), + Err(e) => error!("[client] download failed write: {}", e), + } error!( "Failed to download file '{}' from {}:{} (token: {}) - Error: {}", filename, @@ -814,34 +889,61 @@ impl Client { }; } 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 peer_exists = match client_context + .read_safe() + { + Ok(ctx) => ctx + .peer_registry + .as_ref() + .map(|r| r.contains(&new_peer.username)) + .unwrap_or(false), + Err(e) => { + error!("[client] NewPeer read: {}", e); + continue; + } + }; 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(ServerMessage::GetPeerAddress( - new_peer.username.clone(), - )) - .unwrap(); + } else { + let send_result = client_context + .read_safe() + .ok() + .and_then(|ctx| { + ctx.server_sender.as_ref().map( + |s| { + s.send( + ServerMessage::GetPeerAddress( + new_peer + .username + .clone(), + ), + ) + }, + ) + }); + if let Some(Err(e)) = send_result { + error!( + "[client] NewPeer send GetPeerAddress: {}", + e + ); + } } - let addr = - new_peer.tcp_stream.peer_addr().unwrap(); + let addr = match new_peer.tcp_stream.peer_addr() + { + Ok(a) => a, + Err(e) => { + error!( + "[client] NewPeer peer_addr: {}", + e + ); + continue; + } + }; let host = addr.ip().to_string(); let port: u32 = addr.port().into(); @@ -879,13 +981,22 @@ impl Client { obfuscated_port ); - let peer_exists = client_context - .read() - .unwrap() - .peer_registry - .as_ref() - .map(|r| r.contains(&username)) - .unwrap_or(false); + let peer_exists = match client_context + .read_safe() + { + Ok(ctx) => ctx + .peer_registry + .as_ref() + .map(|r| r.contains(&username)) + .unwrap_or(false), + Err(e) => { + error!( + "[client] GetPeerAddressResponse read: {}", + e + ); + continue; + } + }; if peer_exists { // don't know if i should update? and or reconnect the peer @@ -933,8 +1044,18 @@ impl Client { transfer, username, ) => { - let mut context = - client_context.write().unwrap(); + let mut context = match client_context + .write_safe() + { + Ok(c) => c, + Err(e) => { + error!( + "[client] UpdateDownloadTokens write: {}", + e + ); + continue; + } + }; let download_to_update = context .get_downloads() @@ -981,9 +1102,18 @@ impl Client { ); } ClientOperation::SetServerSender(sender) => { - client_context.write().unwrap().server_sender = - Some(sender); - debug!("[client] Server sender initialized"); + match client_context.write_safe() { + Ok(mut ctx) => { + ctx.server_sender = Some(sender); + debug!( + "[client] Server sender initialized" + ); + } + Err(e) => error!( + "[client] SetServerSender write: {}", + e + ), + } } } } @@ -1013,7 +1143,13 @@ impl Client { ConnectionType::P => { let username = peer.username.clone(); - let context = client_context.read().unwrap(); + let context = match client_context.read_safe() { + Ok(c) => c, + Err(e) => { + error!("[client] connect_to_peer read: {}", e); + return; + } + }; if let Some(ref registry) = context.peer_registry { match registry.register_peer(peer_clone, stream, None) { Ok(_) => (), @@ -1054,13 +1190,16 @@ impl Client { filename, download.size ); let _ = download.sender.send(DownloadStatus::Completed); - client_context - .write() - .unwrap() - .update_download_with_status( + match client_context.write_safe() { + Ok(mut ctx) => ctx.update_download_with_status( download.token, DownloadStatus::Completed, - ); + ), + Err(e) => error!( + "[client] connect_to_peer F write: {}", + e + ), + } } Err(e) => { trace!("[client] failed to download: {}", e); @@ -1079,7 +1218,13 @@ impl Client { ) { debug!("Piercing firewall for peer: {:?}", peer); - let context = client_context.read().unwrap(); + let context = match client_context.read_safe() { + Ok(c) => c, + Err(e) => { + error!("[client] pierce_firewall read: {}", e); + return; + } + }; if let Some(server_sender) = &context.server_sender { if let Some(token) = peer.token { match server_sender.send(ServerMessage::PierceFirewall(token)) { diff --git a/soulseek-rs-lib/src/error.rs b/soulseek-rs-lib/src/error.rs index 81adadd..1410dbb 100644 --- a/soulseek-rs-lib/src/error.rs +++ b/soulseek-rs-lib/src/error.rs @@ -18,6 +18,8 @@ pub enum SoulseekRs { NotConnected, /// Compression/decompression error CompressionError(String), + /// A lock was poisoned by a panic in another thread + LockPoisoned, } impl fmt::Display for SoulseekRs { @@ -39,6 +41,9 @@ impl fmt::Display for SoulseekRs { SoulseekRs::CompressionError(msg) => { write!(f, "Compression error: {}", msg) } + SoulseekRs::LockPoisoned => { + write!(f, "Lock poisoned by panicking thread") + } } } } diff --git a/soulseek-rs-lib/src/peer/listen.rs b/soulseek-rs-lib/src/peer/listen.rs index 692a260..28b23f7 100644 --- a/soulseek-rs-lib/src/peer/listen.rs +++ b/soulseek-rs-lib/src/peer/listen.rs @@ -9,6 +9,7 @@ use crate::client::{ClientContext, ClientOperation}; use crate::message::{Message, MessageReader}; use crate::peer::{ConnectionType, DownloadPeer, Peer}; use crate::types::Download; +use crate::utils::lock::RwLockExt; use crate::{DownloadStatus, debug, error, info, trace}; const PEER_INIT_MESSAGE_CODE: u8 = 1; @@ -82,7 +83,13 @@ fn extract_download_from_buffer( username, token ); - let context = client_context.read().unwrap(); + let context = match client_context.read_safe() { + Ok(c) => c, + Err(e) => { + error!("[listener] client context lock: {}", e); + return None; + } + }; let download = context.get_download_by_token(token).cloned(); if download.is_none() { @@ -104,7 +111,13 @@ fn handle_peer_connection( _peer_ip: &str, _peer_port: u16, ) { - let client_context = context.client_context.read().unwrap(); + let client_context = match context.client_context.read_safe() { + Ok(c) => c, + Err(e) => { + error!("[listener] handle_peer_connection lock: {}", e); + return; + } + }; if let Some(ref registry) = client_context.peer_registry { match registry.register_peer(peer.clone(), Some(stream), Some(reader)) { Ok(_) => (), @@ -158,14 +171,15 @@ fn handle_file_connection( ) { Ok((download, filename)) => { let _ = download.sender.send(DownloadStatus::Completed); - context - .client_context - .write() - .unwrap() - .update_download_with_status( + match context.client_context.write_safe() { + Ok(mut ctx) => ctx.update_download_with_status( download.token, DownloadStatus::Completed, - ); + ), + Err(e) => { + error!("[listener] handle_file_connection write: {}", e) + } + } info!( "Successfully downloaded {} bytes to {}", download.size, filename diff --git a/soulseek-rs-lib/src/utils/lock.rs b/soulseek-rs-lib/src/utils/lock.rs new file mode 100644 index 0000000..17588b6 --- /dev/null +++ b/soulseek-rs-lib/src/utils/lock.rs @@ -0,0 +1,28 @@ +use std::sync::{Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard}; + +use crate::error::{Result, SoulseekRs}; + +pub trait RwLockExt { + fn read_safe(&self) -> Result>; + fn write_safe(&self) -> Result>; +} + +impl RwLockExt for RwLock { + fn read_safe(&self) -> Result> { + self.read().map_err(|_| SoulseekRs::LockPoisoned) + } + + fn write_safe(&self) -> Result> { + self.write().map_err(|_| SoulseekRs::LockPoisoned) + } +} + +pub trait MutexExt { + fn lock_safe(&self) -> Result>; +} + +impl MutexExt for Mutex { + fn lock_safe(&self) -> Result> { + self.lock().map_err(|_| SoulseekRs::LockPoisoned) + } +} diff --git a/soulseek-rs-lib/src/utils/mod.rs b/soulseek-rs-lib/src/utils/mod.rs index 0fde4e4..bd215a2 100644 --- a/soulseek-rs-lib/src/utils/mod.rs +++ b/soulseek-rs-lib/src/utils/mod.rs @@ -1,5 +1,6 @@ #[macro_use] pub mod logger; +pub mod lock; pub mod md5; pub mod thread_pool; pub mod zlib; diff --git a/soulseek-rs-lib/src/utils/thread_pool.rs b/soulseek-rs-lib/src/utils/thread_pool.rs index 1515898..41e5bda 100644 --- a/soulseek-rs-lib/src/utils/thread_pool.rs +++ b/soulseek-rs-lib/src/utils/thread_pool.rs @@ -2,6 +2,8 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, mpsc}; use std::thread; +use crate::utils::lock::MutexExt; + type Job = Box; enum Message { @@ -91,7 +93,10 @@ impl Worker { ) -> Worker { let thread = thread::spawn(move || { loop { - let message = receiver.lock().unwrap().recv(); + let message = match receiver.lock_safe() { + Ok(rx) => rx.recv(), + Err(_) => break, + }; match message { Ok(Message::NewJob(job)) => { // let active = diff --git a/soulseek-rs/src/main.rs b/soulseek-rs/src/main.rs index cfb2f66..4aea41d 100644 --- a/soulseek-rs/src/main.rs +++ b/soulseek-rs/src/main.rs @@ -92,7 +92,9 @@ fn main() -> Result<()> { }; let mut client = Client::with_settings(settings); - client.connect(); + client.connect().map_err(|e| { + color_eyre::eyre::eyre!("Failed to connect: {}", e) + })?; client.login().map_err(|e| { color_eyre::eyre::eyre!("Failed to login: {}", e) })?; @@ -144,7 +146,9 @@ fn search_and_download(config: SearchConfig) -> Result<()> { }; let mut client = Client::with_settings(settings); - client.connect(); + client + .connect() + .map_err(|e| color_eyre::eyre::eyre!("Failed to connect: {}", e))?; client .login() .map_err(|e| color_eyre::eyre::eyre!("Failed to login: {}", e))?;