From 9fe1c9984c90267e028f83ca8b3e624aea1e1dc0 Mon Sep 17 00:00:00 2001 From: Michel de Graaf Date: Thu, 13 Aug 2026 13:46:18 +0200 Subject: [PATCH 1/2] fix(client): bound the threads a peer flood can create One popular search produced a ConnectToPeer for every responder, and each one cost a 2 MiB wrapper thread, an unbounded 256 KiB actor, a listener thread per inbound socket, and on failure a 2 MiB thread that slept for 20 seconds. A container with a pids limit died on one album search (issues #32, #33). Four bounds, no behavior change for a quiet session: - P-type registrations run inline on the ops loop; only F transfers keep a worker thread. The dial already ran on the actor thread. - The broker timeout is a deadline swept by the ops loop instead of a sleeping thread per failed dial. - The registry caps live actors (default 512, Client::set_max_peers). At capacity the oldest peer outside the caller's protected set (live downloads, uploads, pending serves, queued messages, browses in flight) is evicted; if none qualify the registration is refused before any thread is spawned. - Concurrent listener handshakes are bounded by a 128-permit semaphore released once the peer-init arrives, so slow handshakes cannot multiply threads while transfers stay unthrottled. A replayed TransferResponse for an InProgress or Completed token is dropped instead of starting a second transfer thread over the same .part file. Fixes #32 Fixes #33 --- soulseek-rs-lib/src/actor/peer_registry.rs | 178 +++++++++++++++++++-- soulseek-rs-lib/src/client/connection.rs | 15 +- soulseek-rs-lib/src/client/mod.rs | 69 +++++++- soulseek-rs-lib/src/client/operations.rs | 112 ++++++++----- soulseek-rs-lib/src/client/rooms.rs | 3 +- soulseek-rs-lib/src/client/tests.rs | 48 ++++++ soulseek-rs-lib/src/peer/listen.rs | 19 ++- soulseek-rs-lib/src/utils/mod.rs | 1 + soulseek-rs-lib/src/utils/semaphore.rs | 73 +++++++++ 9 files changed, 458 insertions(+), 60 deletions(-) create mode 100644 soulseek-rs-lib/src/utils/semaphore.rs diff --git a/soulseek-rs-lib/src/actor/peer_registry.rs b/soulseek-rs-lib/src/actor/peer_registry.rs index 7eec579..3819830 100644 --- a/soulseek-rs-lib/src/actor/peer_registry.rs +++ b/soulseek-rs-lib/src/actor/peer_registry.rs @@ -6,11 +6,12 @@ use crate::peer::Peer; use crate::utils::lock::MutexExt; use crate::{debug, error}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::net::TcpStream; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::mpsc::Sender; use std::sync::{Arc, Mutex}; +use std::time::Instant; /// Source of unique per-actor ids so terminal-outcome eviction can be made /// identity-aware (a replaced actor must not evict its replacement). @@ -18,13 +19,14 @@ static NEXT_PEER_ID: AtomicU64 = AtomicU64::new(1); /// Registered peers keyed by username, each stored with the unique id of the /// actor currently occupying the slot. -type PeerMap = HashMap)>; +type PeerMap = HashMap, Instant)>; pub struct PeerRegistry { peers: Arc>, actor_system: Arc, client_channel: Sender, own_username: String, + capacity: Arc, } impl PeerRegistry { @@ -33,12 +35,28 @@ impl PeerRegistry { actor_system: Arc, client_channel: Sender, own_username: String, + ) -> Self { + Self::with_capacity( + actor_system, + client_channel, + own_username, + Arc::new(AtomicUsize::new(usize::MAX)), + ) + } + + #[must_use] + pub(crate) fn with_capacity( + actor_system: Arc, + client_channel: Sender, + own_username: String, + capacity: Arc, ) -> Self { Self { peers: Arc::new(Mutex::new(HashMap::new())), actor_system, client_channel, own_username, + capacity, } } @@ -47,6 +65,16 @@ impl PeerRegistry { peer: Peer, stream: Option, reader: Option, + ) -> Result, String> { + self.register_peer_protected(peer, stream, reader, &HashSet::new()) + } + + pub(crate) fn register_peer_protected( + &self, + peer: Peer, + stream: Option, + reader: Option, + protected: &HashSet, ) -> Result, String> { let username = peer.username.clone(); let id = NEXT_PEER_ID.fetch_add(1, Ordering::Relaxed); @@ -71,6 +99,28 @@ impl PeerRegistry { .lock_safe() .map_err(|e| format!("peer registry lock poisoned: {e}"))?; + let capacity = self.capacity.load(Ordering::Relaxed).max(1); + if peers.len() >= capacity && !peers.contains_key(&username) { + let victim = peers + .iter() + .filter(|(name, _)| !protected.contains(name.as_str())) + .min_by_key(|(_, (_, _, registered))| *registered) + .map(|(name, _)| name.clone()); + let Some(victim) = victim else { + return Err(format!( + "peer registry at capacity ({capacity}) and every peer \ + is busy; refusing {username}" + )); + }; + if let Some((_, handle, _)) = peers.remove(&victim) { + let _ = handle.stop(); + debug!( + "[peer_registry] evicted idle peer {} to admit {}", + victim, username + ); + } + } + let handle = self .actor_system .try_spawn_with_handle(actor, |actor, handle| { @@ -81,8 +131,8 @@ impl PeerRegistry { // become an orphan pinning a pool worker forever. Eviction on the // replaced actor's later shutdown is identity-aware (keyed on its id), // so stopping it here cannot evict this new connection. - if let Some((_, old_handle)) = - peers.insert(username.clone(), (id, handle.clone())) + if let Some((_, old_handle, _)) = + peers.insert(username.clone(), (id, handle.clone(), Instant::now())) { let _ = old_handle.stop(); debug!( @@ -97,7 +147,9 @@ impl PeerRegistry { #[must_use] pub fn get_peer(&self, username: &str) -> Option> { match self.peers.lock_safe() { - Ok(peers) => peers.get(username).map(|(_, handle)| handle.clone()), + Ok(peers) => { + peers.get(username).map(|(_, handle, _)| handle.clone()) + } Err(e) => { error!("[peer_registry] get_peer: {}", e); None @@ -123,7 +175,7 @@ impl PeerRegistry { debug!("[peer_registry] Removed peer actor for {}", username); } - removed.map(|(_, handle)| handle) + removed.map(|(_, handle, _)| handle) } /// Remove and return the actor for `username` only if it is still the actor @@ -142,8 +194,11 @@ impl PeerRegistry { return None; } }; - if peers.get(username).is_some_and(|(stored, _)| *stored == id) { - let removed = peers.remove(username).map(|(_, handle)| handle); + if peers + .get(username) + .is_some_and(|(stored, _, _)| *stored == id) + { + let removed = peers.remove(username).map(|(_, handle, _)| handle); debug!( "[peer_registry] Removed peer actor {} for {}", id, username @@ -192,6 +247,7 @@ impl Clone for PeerRegistry { actor_system: self.actor_system.clone(), client_channel: self.client_channel.clone(), own_username: self.own_username.clone(), + capacity: self.capacity.clone(), } } } @@ -201,8 +257,112 @@ mod tests { use super::PeerRegistry; use crate::actor::ActorSystem; use crate::peer::{ConnectionType, Peer}; + use std::collections::HashSet; use std::net::{TcpListener, TcpStream}; use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::time::Duration; + + fn loopback_peer(name: &str) -> (Peer, TcpStream, TcpStream) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let stream = TcpStream::connect(addr).unwrap(); + stream.set_nonblocking(true).unwrap(); + let far_end = listener.accept().unwrap().0; + let peer = Peer::new( + name.to_string(), + ConnectionType::P, + "127.0.0.1".to_string(), + u32::from(addr.port()), + None, + 0, + 0, + 0, + ); + (peer, stream, far_end) + } + + fn capped_registry( + capacity: usize, + ) -> ( + PeerRegistry, + std::sync::mpsc::Receiver, + ) { + let system = Arc::new(ActorSystem::new()); + let (tx, rx) = std::sync::mpsc::channel(); + let registry = PeerRegistry::with_capacity( + system, + tx, + "me".to_string(), + Arc::new(AtomicUsize::new(capacity)), + ); + (registry, rx) + } + + #[test] + fn capacity_evicts_the_oldest_idle_peer() { + let (registry, _rx) = capped_registry(2); + + for name in ["a", "b", "c"] { + let (peer, stream, _far_end) = loopback_peer(name); + registry.register_peer(peer, Some(stream), None).unwrap(); + std::thread::sleep(Duration::from_millis(5)); + } + + assert!(!registry.contains("a"), "the oldest peer must be evicted"); + assert!(registry.contains("b")); + assert!(registry.contains("c")); + } + + #[test] + fn capacity_refuses_when_every_peer_is_protected() { + let (registry, _rx) = capped_registry(1); + + let (busy, busy_stream, _busy_far) = loopback_peer("busy"); + registry + .register_peer(busy, Some(busy_stream), None) + .unwrap(); + + let protected: HashSet = + std::iter::once("busy".to_string()).collect(); + let (newcomer, stream, _far) = loopback_peer("newcomer"); + let result = registry.register_peer_protected( + newcomer, + Some(stream), + None, + &protected, + ); + + assert!(result.is_err(), "a full registry of busy peers must refuse"); + assert!(registry.contains("busy")); + assert!(!registry.contains("newcomer")); + } + + #[test] + fn a_returning_username_is_replaced_not_refused_at_capacity() { + let (registry, _rx) = capped_registry(1); + + let (first, first_stream, _first_far) = loopback_peer("bob"); + registry + .register_peer(first, Some(first_stream), None) + .unwrap(); + + let protected: HashSet = + std::iter::once("bob".to_string()).collect(); + let (again, stream, _far) = loopback_peer("bob"); + let result = registry.register_peer_protected( + again, + Some(stream), + None, + &protected, + ); + + assert!( + result.is_ok(), + "a reconnecting username must replace itself" + ); + assert!(registry.contains("bob")); + } #[test] fn remove_peer_if_respects_actor_identity() { diff --git a/soulseek-rs-lib/src/client/connection.rs b/soulseek-rs-lib/src/client/connection.rs index da2bd69..bc4a6a4 100644 --- a/soulseek-rs-lib/src/client/connection.rs +++ b/soulseek-rs-lib/src/client/connection.rs @@ -50,10 +50,11 @@ impl Client { ) = mpsc::channel(); let mut ctx = self.context.write_safe()?; - let peer_registry = PeerRegistry::new( + let peer_registry = PeerRegistry::with_capacity( ctx.actor_system.clone(), sender.clone(), self.username.clone(), + ctx.max_peers.clone(), ); ctx.peer_registry = Some(peer_registry); @@ -175,6 +176,13 @@ impl Client { Ok(()) } + pub fn set_max_peers(&self, max_peers: usize) { + if let Ok(ctx) = self.context.read_safe() { + ctx.max_peers + .store(max_peers.max(1), std::sync::atomic::Ordering::Relaxed); + } + } + #[allow(dead_code)] pub fn remove_peer(&self, username: &str) { let context = match self.context.read_safe() { @@ -214,7 +222,10 @@ impl Client { } }; if let Some(ref registry) = context.peer_registry { - match registry.register_peer(peer_clone, stream, None) { + let protected = context.protected_peers(); + match registry.register_peer_protected( + peer_clone, stream, None, &protected, + ) { Ok(_) => (), Err(e) => { trace!( diff --git a/soulseek-rs-lib/src/client/mod.rs b/soulseek-rs-lib/src/client/mod.rs index f473408..b935ee9 100644 --- a/soulseek-rs-lib/src/client/mod.rs +++ b/soulseek-rs-lib/src/client/mod.rs @@ -23,7 +23,7 @@ use std::{ net::TcpStream, sync::{ Arc, RwLock, - atomic::{AtomicBool, AtomicU32, Ordering}, + atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}, mpsc::{self, Receiver, Sender}, }, thread::{self, sleep}, @@ -58,6 +58,10 @@ pub const DEFAULT_UPLOAD_SLOTS: usize = 10; /// before giving up and failing the download. Matches the direct-dial timeout. const BROKER_CONNECT_TIMEOUT: Duration = Duration::from_secs(20); +const DEFAULT_MAX_PEERS: usize = 512; + +const BROWSE_PROTECT_WINDOW: Duration = Duration::from_mins(1); + /// Source of non-zero correlation tokens for server-brokered connections. static NEXT_CONNECT_TOKEN: AtomicU32 = AtomicU32::new(1); @@ -295,7 +299,8 @@ pub struct ClientContext { private_messages: Vec, /// Correlation tokens for server-brokered (firewalled) connections, mapping /// a token we sent in a ConnectToPeer to the peer we expect back. - pending_connect_tokens: HashMap, + pending_connect_tokens: HashMap, + max_peers: Arc, /// Files we share with peers (read-only after connect). pub shares: Arc, /// The directories the current share index was built from. @@ -311,6 +316,7 @@ pub struct ClientContext { pending_serves: HashMap>, /// Shared-file listings received from peers we browsed. browse_results: HashMap>, + pending_browses: HashMap, /// Latest snapshot of the public chat-room list (from `RoomList`, code 64). room_list: Vec, /// Chat-room events awaiting consumption by the client/UI. @@ -395,6 +401,7 @@ impl ClientContext { searches: HashMap::new(), private_messages: Vec::new(), pending_connect_tokens: HashMap::new(), + max_peers: Arc::new(AtomicUsize::new(DEFAULT_MAX_PEERS)), shares: Arc::new(Shares::empty()), shared_directories: Vec::new(), peer_addresses: HashMap::new(), @@ -403,6 +410,7 @@ impl ClientContext { active_uploads: HashMap::new(), pending_serves: HashMap::new(), browse_results: HashMap::new(), + pending_browses: HashMap::new(), room_list: Vec::new(), room_events: Vec::new(), room_members: HashMap::new(), @@ -563,9 +571,17 @@ impl ClientContext { username: String, directories: Vec, ) { + self.pending_browses.remove(&username); self.browse_results.insert(username, directories); } + pub(crate) fn mark_browse_pending(&mut self, username: &str) { + let now = Instant::now(); + self.pending_browses.retain(|_, deadline| *deadline > now); + self.pending_browses + .insert(username.to_string(), now + BROWSE_PROTECT_WINDOW); + } + /// Remove and return the shared-file listing browsed from `username`. pub fn take_browse_result( &mut self, @@ -577,12 +593,57 @@ impl ClientContext { /// Remember that a server-brokered connection to `username` is pending under /// `token`; the peer will quote it back in a PierceFirewall. pub fn add_pending_connect(&mut self, token: u32, username: String) { - self.pending_connect_tokens.insert(token, username); + self.pending_connect_tokens + .insert(token, (username, Instant::now() + BROKER_CONNECT_TIMEOUT)); } /// Resolve and consume the peer expected for a brokered connection `token`. pub fn take_pending_connect(&mut self, token: u32) -> Option { - self.pending_connect_tokens.remove(&token) + self.pending_connect_tokens + .remove(&token) + .map(|(username, _)| username) + } + + pub fn take_expired_connects(&mut self, now: Instant) -> Vec { + let mut expired = Vec::new(); + self.pending_connect_tokens + .retain(|_, (username, deadline)| { + if now >= *deadline { + expired.push(std::mem::take(username)); + return false; + } + true + }); + expired + } + + pub(crate) fn protected_peers(&self) -> HashSet { + let mut protected: HashSet = self + .downloads + .list() + .iter() + .filter(|d| { + matches!( + d.status, + DownloadStatus::Queued + | DownloadStatus::InProgress { .. } + | DownloadStatus::Paused { .. } + ) + }) + .map(|d| d.username.clone()) + .collect(); + protected + .extend(self.active_uploads.values().map(|u| u.username.clone())); + protected.extend(self.pending_serves.keys().cloned()); + protected.extend(self.pending_peer_messages.keys().cloned()); + let now = Instant::now(); + protected.extend( + self.pending_browses + .iter() + .filter(|(_, deadline)| **deadline > now) + .map(|(username, _)| username.clone()), + ); + protected } /// Record a private message received from another user. diff --git a/soulseek-rs-lib/src/client/operations.rs b/soulseek-rs-lib/src/client/operations.rs index 7bdc715..02b67cd 100644 --- a/soulseek-rs-lib/src/client/operations.rs +++ b/soulseek-rs-lib/src/client/operations.rs @@ -1,12 +1,14 @@ use super::{ - Arc, BROKER_CONNECT_TIMEOUT, Client, ClientContext, ClientOperation, - ConnectionType, Download, DownloadPeer, DownloadStatus, Peer, PeerMessage, + Arc, Client, ClientContext, ClientOperation, ConnectionType, Download, + DownloadPeer, DownloadStatus, Duration, Instant, Peer, PeerMessage, PeerRegistry, Receiver, RwLock, RwLockExt, ServerMessage, - build_search_response, debug, error, info, next_connect_token, sleep, + build_search_response, debug, error, info, mpsc, next_connect_token, thread, trace, warn, }; use crate::message::server::MessageFactory; +const CONNECT_SWEEP_INTERVAL: Duration = Duration::from_secs(1); + impl Client { pub(crate) fn listen_to_client_operations( reader: Receiver, @@ -14,20 +16,40 @@ impl Client { own_username: String, ) { thread::spawn(move || { - for operation in reader { + let mut last_sweep = Instant::now(); + loop { + let next = reader.recv_timeout(CONNECT_SWEEP_INTERVAL); + if last_sweep.elapsed() >= CONNECT_SWEEP_INTERVAL { + last_sweep = Instant::now(); + Self::sweep_expired_connects(&client_context); + } + let operation = match next { + Ok(operation) => operation, + Err(mpsc::RecvTimeoutError::Timeout) => continue, + Err(mpsc::RecvTimeoutError::Disconnected) => break, + }; match operation { ClientOperation::ConnectToPeer(peer) => { - let client_context_clone = client_context.clone(); - let own_username_clone = own_username.clone(); + if matches!(peer.connection_type, ConnectionType::F) { + let client_context_clone = client_context.clone(); + let own_username_clone = own_username.clone(); - thread::spawn(move || { + thread::spawn(move || { + Self::connect_to_peer( + peer, + client_context_clone, + own_username_clone, + None, + ); + }); + } else { Self::connect_to_peer( peer, - client_context_clone, - own_username_clone, + client_context.clone(), + own_username.clone(), None, ); - }); + } } ClientOperation::SearchResult(search_result) => { trace!("[client] SearchResult {:?}", search_result); @@ -122,6 +144,28 @@ impl Client { ); continue; }; + if matches!( + download.status, + DownloadStatus::InProgress { .. } + | DownloadStatus::Completed + ) { + debug!( + "[client] transfer token {} already claimed; \ + ignoring replayed TransferResponse", + token + ); + continue; + } + if let Ok(mut ctx) = client_context.write_safe() { + ctx.update_download_with_status( + token, + DownloadStatus::InProgress { + bytes_downloaded: 0, + total_bytes: download.size, + speed_bytes_per_sec: 0.0, + }, + ); + } thread::spawn(move || { let download_peer = DownloadPeer::new( @@ -285,17 +329,12 @@ impl Client { u8::try_from(obfuscation_type).unwrap_or(0), obfuscated_port, ); - let client_context_clone = client_context.clone(); - let own_username_clone = own_username.clone(); - - thread::spawn(move || { - Self::connect_to_peer( - peer, - client_context_clone, - own_username_clone, - None, - ); - }); + Self::connect_to_peer( + peer, + client_context.clone(), + own_username.clone(), + None, + ); } } ClientOperation::UpdateDownloadTokens( @@ -751,30 +790,19 @@ impl Client { ConnectionType::P, ); let _ = sender.send(ServerMessage::SendMessage(msg)); - - // Bound the brokered attempt: if no PierceFirewall - // consumes the token, fail the peer's queued - // downloads (so the caller's Receiver unblocks) - // and reclaim the token. A successful pierce - // takes the token first, making this a no-op. - let timeout_ctx = client_context.clone(); - let timeout_user = username.clone(); - thread::spawn(move || { - sleep(BROKER_CONNECT_TIMEOUT); - let still_pending = - timeout_ctx.write_safe().is_ok_and(|mut c| { - c.take_pending_connect(token).is_some() - }); - if still_pending { - Self::fail_queued_downloads( - &timeout_ctx, - &timeout_user, - ); - } - }); } } } }); } + + fn sweep_expired_connects(client_context: &Arc>) { + let expired = match client_context.write_safe() { + Ok(mut ctx) => ctx.take_expired_connects(Instant::now()), + Err(_) => return, + }; + for username in expired { + Self::fail_queued_downloads(client_context, &username); + } + } } diff --git a/soulseek-rs-lib/src/client/rooms.rs b/soulseek-rs-lib/src/client/rooms.rs index 41d4d5b..8afc871 100644 --- a/soulseek-rs-lib/src/client/rooms.rs +++ b/soulseek-rs-lib/src/client/rooms.rs @@ -177,7 +177,8 @@ impl Client { let request = crate::message::server::MessageFactory::build_get_share_file_list(); let (connected, registry) = { - let ctx = self.context.read_safe()?; + let mut ctx = self.context.write_safe()?; + ctx.mark_browse_pending(username); ( ctx.peer_registry .as_ref() diff --git a/soulseek-rs-lib/src/client/tests.rs b/soulseek-rs-lib/src/client/tests.rs index 9179497..9b5efcb 100644 --- a/soulseek-rs-lib/src/client/tests.rs +++ b/soulseek-rs-lib/src/client/tests.rs @@ -395,3 +395,51 @@ fn a_clean_disconnect_keeps_queued_uploads_an_error_drops_them() { thread::sleep(Duration::from_millis(10)); } } + +#[test] +fn a_pending_connect_expires_after_the_broker_timeout() { + let mut context = ClientContext::new(); + context.add_pending_connect(7, "ghost".to_string()); + + assert!(context.take_expired_connects(Instant::now()).is_empty()); + + let after_deadline = + Instant::now() + BROKER_CONNECT_TIMEOUT + Duration::from_secs(1); + assert_eq!( + context.take_expired_connects(after_deadline), + vec!["ghost".to_string()] + ); + assert!(context.take_pending_connect(7).is_none()); +} + +#[test] +fn protected_peers_covers_downloads_uploads_and_pending_serves() { + let mut context = ClientContext::new(); + context.add_download(download( + "downloader", + "song.mp3", + 1, + DownloadStatus::Queued, + mpsc::channel().0, + )); + context.add_download(download( + "done", + "old.mp3", + 2, + DownloadStatus::Completed, + mpsc::channel().0, + )); + context + .pending_serves + .insert("waiting".to_string(), vec![9]); + context.mark_browse_pending("browsed"); + + let protected = context.protected_peers(); + assert!(protected.contains("downloader")); + assert!(protected.contains("waiting")); + assert!(protected.contains("browsed")); + assert!(!protected.contains("done")); + + context.store_browse_result("browsed".to_string(), Vec::new()); + assert!(!context.protected_peers().contains("browsed")); +} diff --git a/soulseek-rs-lib/src/peer/listen.rs b/soulseek-rs-lib/src/peer/listen.rs index ed77a24..739e3a0 100644 --- a/soulseek-rs-lib/src/peer/listen.rs +++ b/soulseek-rs-lib/src/peer/listen.rs @@ -11,6 +11,7 @@ use crate::message::{Message, MessageReader}; use crate::peer::{ConnectionType, DownloadPeer, Peer}; use crate::types::Download; use crate::utils::lock::RwLockExt; +use crate::utils::semaphore::{Permit, Semaphore}; use crate::{DownloadStatus, debug, error, info, trace}; /// How long to wait before accepting again after a failure. @@ -25,6 +26,8 @@ const PEER_INIT_MESSAGE_CODE: u8 = 1; /// How long an accepted peer gets to send its peer-init handshake. const PEER_INIT_TIMEOUT: Duration = Duration::from_secs(10); +const MAX_HANDSHAKES: usize = 128; + #[derive(Clone)] struct ConnectionContext { client_sender: Sender, @@ -140,7 +143,13 @@ fn handle_peer_connection( } }; if let Some(ref registry) = client_context.peer_registry { - match registry.register_peer(peer.clone(), Some(stream), Some(reader)) { + let protected = client_context.protected_peers(); + match registry.register_peer_protected( + peer.clone(), + Some(stream), + Some(reader), + &protected, + ) { Ok(_) => (), Err(e) => { error!( @@ -283,6 +292,7 @@ fn handle_pierce_firewall( fn handle_incoming_connection( mut stream: TcpStream, context: ConnectionContext, + permit: Permit, ) { let Ok(peer_addr) = stream.peer_addr() else { error!("[listener] failed to get peer address"); @@ -303,6 +313,7 @@ fn handle_incoming_connection( return; } }; + drop(permit); // A firewalled peer brokered through the server connects back with a // PierceFirewall (code 0) instead of a PeerInit (code 1). @@ -406,6 +417,7 @@ impl Listen { client_context, own_username, }; + let handshakes = Arc::new(Semaphore::new(MAX_HANDSHAKES)); for stream in listener.incoming() { let Ok(stream) = stream else { @@ -421,11 +433,14 @@ impl Listen { continue; }; + let permit = handshakes.acquire(); let context = context.clone(); // One thread per connection: the peer-init handshake blocks, and a // peer that is slow to send one must not stop us accepting anybody // else — a wedged accept loop makes us unreachable to every peer. - thread::spawn(move || handle_incoming_connection(stream, context)); + thread::spawn(move || { + handle_incoming_connection(stream, context, permit); + }); } } } diff --git a/soulseek-rs-lib/src/utils/mod.rs b/soulseek-rs-lib/src/utils/mod.rs index 6b42e19..5bb6848 100644 --- a/soulseek-rs-lib/src/utils/mod.rs +++ b/soulseek-rs-lib/src/utils/mod.rs @@ -3,6 +3,7 @@ pub mod logger; pub mod lock; pub mod md5; pub mod path; +pub(crate) mod semaphore; pub mod zlib; // Re-export commonly used items diff --git a/soulseek-rs-lib/src/utils/semaphore.rs b/soulseek-rs-lib/src/utils/semaphore.rs new file mode 100644 index 0000000..0508551 --- /dev/null +++ b/soulseek-rs-lib/src/utils/semaphore.rs @@ -0,0 +1,73 @@ +use std::sync::{Arc, Condvar, Mutex, PoisonError}; + +pub struct Semaphore { + permits: Mutex, + available: Condvar, +} + +pub struct Permit(Arc); + +impl Semaphore { + #[must_use] + pub const fn new(permits: usize) -> Self { + Self { + permits: Mutex::new(permits), + available: Condvar::new(), + } + } + + #[must_use] + pub fn acquire(self: &Arc) -> Permit { + let mut permits = + self.permits.lock().unwrap_or_else(PoisonError::into_inner); + while *permits == 0 { + permits = self + .available + .wait(permits) + .unwrap_or_else(PoisonError::into_inner); + } + *permits -= 1; + Permit(Arc::clone(self)) + } +} + +impl Drop for Permit { + fn drop(&mut self) { + let mut permits = self + .0 + .permits + .lock() + .unwrap_or_else(PoisonError::into_inner); + *permits += 1; + self.0.available.notify_one(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + #[test] + fn permits_bound_concurrent_holders() { + let semaphore = Arc::new(Semaphore::new(2)); + let first = semaphore.acquire(); + let _second = semaphore.acquire(); + + let (tx, rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let _third = semaphore.acquire(); + tx.send(()).unwrap(); + }); + + assert!( + rx.recv_timeout(Duration::from_millis(100)).is_err(), + "a third acquire must block while both permits are held" + ); + drop(first); + assert!( + rx.recv_timeout(Duration::from_secs(2)).is_ok(), + "releasing a permit must wake the waiter" + ); + } +} From 492b96176670f30325b713be91b154288482aec6 Mon Sep 17 00:00:00 2001 From: Michel de Graaf Date: Thu, 13 Aug 2026 14:10:20 +0200 Subject: [PATCH 2/2] fix(client): survive the review findings on the thread-flood fix Review of the capacity work surfaced three real defects and two gaps: - A peer dribbling bytes re-armed the 10s handshake read timeout on every read and held its accept permit forever; 128 such sockets wedged the accept loop for good. The handshake now has one deadline across all reads, plus a 4 KiB cap on buffered handshake bytes. - A capacity refusal was logged and swallowed, leaving the refused peer's queued downloads stuck until the caller's own timeout. Both registration paths now fail those downloads so the caller gets a verdict, and the listener only announces PeerConnected when the registration actually happened. - The transfer-claim was undone by the worker thread writing Queued back over it; that write is gone, so a replayed TransferResponse can no longer start a second writer on the same .part file. - Offered uploads and queue waiters joined the eviction-protected set, and the browse protection window matches the minutes a real share listing takes rather than one. Regression tests: the broker-timeout sweep failing a queued download end to end, the replay guard observed over a real socket, and the handshake deadline against a dribbling and an oversized peer. --- soulseek-rs-lib/src/client/connection.rs | 48 +++--- soulseek-rs-lib/src/client/mod.rs | 5 +- soulseek-rs-lib/src/client/tests.rs | 95 +++++++++++ soulseek-rs-lib/src/peer/download_peer.rs | 8 - soulseek-rs-lib/src/peer/listen.rs | 196 ++++++++++++++++------ 5 files changed, 271 insertions(+), 81 deletions(-) diff --git a/soulseek-rs-lib/src/client/connection.rs b/soulseek-rs-lib/src/client/connection.rs index bc4a6a4..c9219f9 100644 --- a/soulseek-rs-lib/src/client/connection.rs +++ b/soulseek-rs-lib/src/client/connection.rs @@ -214,28 +214,36 @@ impl Client { ConnectionType::P => { let username = peer.username; - let context = match client_context.read_safe() { - Ok(c) => c, - Err(e) => { - error!("[client] connect_to_peer read: {}", e); - return; + let refused = { + 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 { + let protected = context.protected_peers(); + match registry.register_peer_protected( + peer_clone, stream, None, &protected, + ) { + Ok(_) => false, + Err(e) => { + warn!( + "[client] refusing peer connection for \ + {:?}: {:?}", + username, e + ); + true + } + } + } else { + trace!("PeerRegistry not initialized"); + false } }; - if let Some(ref registry) = context.peer_registry { - let protected = context.protected_peers(); - match registry.register_peer_protected( - peer_clone, stream, None, &protected, - ) { - Ok(_) => (), - Err(e) => { - trace!( - "Failed to spawn peer actor for {:?}: {:?}", - username, e - ); - } - } - } else { - trace!("PeerRegistry not initialized"); + if refused { + Self::fail_queued_downloads(&client_context, &username); } } diff --git a/soulseek-rs-lib/src/client/mod.rs b/soulseek-rs-lib/src/client/mod.rs index b935ee9..dd9cf09 100644 --- a/soulseek-rs-lib/src/client/mod.rs +++ b/soulseek-rs-lib/src/client/mod.rs @@ -60,7 +60,7 @@ const BROKER_CONNECT_TIMEOUT: Duration = Duration::from_secs(20); const DEFAULT_MAX_PEERS: usize = 512; -const BROWSE_PROTECT_WINDOW: Duration = Duration::from_mins(1); +const BROWSE_PROTECT_WINDOW: Duration = Duration::from_mins(5); /// Source of non-zero correlation tokens for server-brokered connections. static NEXT_CONNECT_TOKEN: AtomicU32 = AtomicU32::new(1); @@ -634,6 +634,9 @@ impl ClientContext { .collect(); protected .extend(self.active_uploads.values().map(|u| u.username.clone())); + protected.extend(self.uploads.values().map(|j| j.downloader.clone())); + protected + .extend(self.upload_queue.iter().map(|q| q.downloader.clone())); protected.extend(self.pending_serves.keys().cloned()); protected.extend(self.pending_peer_messages.keys().cloned()); let now = Instant::now(); diff --git a/soulseek-rs-lib/src/client/tests.rs b/soulseek-rs-lib/src/client/tests.rs index 9b5efcb..23e8ca4 100644 --- a/soulseek-rs-lib/src/client/tests.rs +++ b/soulseek-rs-lib/src/client/tests.rs @@ -443,3 +443,98 @@ fn protected_peers_covers_downloads_uploads_and_pending_serves() { context.store_browse_result("browsed".to_string(), Vec::new()); assert!(!context.protected_peers().contains("browsed")); } + +#[test] +fn an_expired_broker_connect_fails_the_queued_downloads() { + let client = Client::new("u", "p"); + let (sender, receiver) = mpsc::channel(); + { + let mut ctx = client.context.write().unwrap(); + ctx.add_download(download( + "ghost", + "f.mp3", + 7, + DownloadStatus::Queued, + sender, + )); + ctx.pending_connect_tokens + .insert(7, ("ghost".to_string(), Instant::now())); + } + let (_ops_tx, ops_rx) = mpsc::channel(); + Client::listen_to_client_operations( + ops_rx, + client.context, + "u".to_string(), + ); + + let status = receiver.recv_timeout(Duration::from_secs(5)); + assert!( + matches!(status, Ok(DownloadStatus::Failed(_))), + "the sweep must fail the queued download, got {status:?}" + ); +} + +#[test] +fn a_replayed_transfer_response_does_not_start_a_second_transfer() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let port = u32::from(listener.local_addr().unwrap().port()); + + let client = Client::new("u", "p"); + let (sender, _receiver) = mpsc::channel(); + client.context.write().unwrap().add_download(download( + "peer", + "f.mp3", + 9, + DownloadStatus::Queued, + sender, + )); + + let (ops_tx, ops_rx) = mpsc::channel(); + Client::listen_to_client_operations( + ops_rx, + client.context, + "u".to_string(), + ); + + let peer = Peer::new( + "peer".to_string(), + ConnectionType::F, + "127.0.0.1".to_string(), + port, + None, + 0, + 0, + 0, + ); + ops_tx + .send(ClientOperation::DownloadFromPeer(9, peer.clone(), true)) + .unwrap(); + ops_tx + .send(ClientOperation::DownloadFromPeer(9, peer, true)) + .unwrap(); + + let deadline = Instant::now() + Duration::from_secs(5); + let mut connections = 0; + while Instant::now() < deadline { + match listener.accept() { + Ok(_) => { + connections += 1; + if connections == 1 { + thread::sleep(Duration::from_millis(500)); + } + } + Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => { + if connections >= 1 { + break; + } + thread::sleep(Duration::from_millis(20)); + } + Err(_) => break, + } + } + assert_eq!( + connections, 1, + "a replayed TransferResponse must not dial the peer again" + ); +} diff --git a/soulseek-rs-lib/src/peer/download_peer.rs b/soulseek-rs-lib/src/peer/download_peer.rs index 3ea609e..accc2b0 100644 --- a/soulseek-rs-lib/src/peer/download_peer.rs +++ b/soulseek-rs-lib/src/peer/download_peer.rs @@ -484,14 +484,6 @@ impl DownloadPeer { self.no_pierce ); - if let Some(ref dl) = download { - let _ = dl.sender.send(DownloadStatus::Queued); - client_context - .write() - .map_err(|_| DownloadError::LockPoisoned)? - .update_download_with_status(dl.token, DownloadStatus::Queued); - } - let mut stream = match stream { Some(s) => { Self::apply_transfer_deadlines(&s)?; diff --git a/soulseek-rs-lib/src/peer/listen.rs b/soulseek-rs-lib/src/peer/listen.rs index 739e3a0..070df8f 100644 --- a/soulseek-rs-lib/src/peer/listen.rs +++ b/soulseek-rs-lib/src/peer/listen.rs @@ -3,16 +3,16 @@ use std::net::{TcpListener, TcpStream}; use std::sync::mpsc::Sender; use std::sync::{Arc, RwLock}; use std::thread; -use std::time::Duration; +use std::time::{Duration, Instant}; -use crate::client::{ClientContext, ClientOperation}; +use crate::client::{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::utils::semaphore::{Permit, Semaphore}; -use crate::{DownloadStatus, debug, error, info, trace}; +use crate::{DownloadStatus, debug, error, info, trace, warn}; /// How long to wait before accepting again after a failure. /// @@ -28,6 +28,8 @@ const PEER_INIT_TIMEOUT: Duration = Duration::from_secs(10); const MAX_HANDSHAKES: usize = 128; +const MAX_PEER_INIT_BYTES: usize = 4096; + #[derive(Clone)] struct ConnectionContext { client_sender: Sender, @@ -44,16 +46,30 @@ struct PeerInitData { fn read_peer_init_message( stream: &mut TcpStream, reader: &mut MessageReader, + deadline: Instant, ) -> io::Result { // An untrusted peer gets a bounded handshake. Without this a peer that // connects and stays silent parks this read forever, pinning a thread that // owes us a peer init. - stream.set_read_timeout(Some(PEER_INIT_TIMEOUT))?; let message = loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "peer init deadline passed", + )); + } + stream.set_read_timeout(Some(remaining))?; reader.read_from_socket(stream)?; if let Ok(Some(msg)) = reader.extract_message() { break msg; } + if reader.buffer_len() > MAX_PEER_INIT_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "peer init larger than any handshake", + )); + } }; // Whatever this connection turns into reads on its own terms from here. @@ -122,7 +138,7 @@ fn handle_peer_connection( stream: TcpStream, reader: MessageReader, context: &ConnectionContext, -) { +) -> bool { // The peer actor multiplexes socket reads with its mailbox on a single // thread: `tick()` reads the socket, but outgoing messages (e.g. a queued // QueueUpload for a download) are delivered through the mailbox between @@ -131,36 +147,44 @@ fn handle_peer_connection( // outbound path and drive this connection non-blocking. if let Err(e) = stream.set_nonblocking(true) { error!("[listener] failed to set peer stream non-blocking: {}", e); - return; + return false; } stream.set_nodelay(true).ok(); - let client_context = match context.client_context.read_safe() { - Ok(c) => c, - Err(e) => { - error!("[listener] handle_peer_connection lock: {}", e); - return; + let registered = { + let client_context = match context.client_context.read_safe() { + Ok(c) => c, + Err(e) => { + error!("[listener] handle_peer_connection lock: {}", e); + return false; + } + }; + if let Some(ref registry) = client_context.peer_registry { + let protected = client_context.protected_peers(); + match registry.register_peer_protected( + peer.clone(), + Some(stream), + Some(reader), + &protected, + ) { + Ok(_) => true, + Err(e) => { + warn!( + "[listener] refusing peer connection for {:?}: {:?}", + peer.username, e + ); + false + } + } + } else { + error!("PeerRegistry not initialized"); + false } }; - if let Some(ref registry) = client_context.peer_registry { - let protected = client_context.protected_peers(); - match registry.register_peer_protected( - peer.clone(), - Some(stream), - Some(reader), - &protected, - ) { - Ok(_) => (), - Err(e) => { - error!( - "Failed to spawn peer actor for {:?}: {:?}", - peer.username, e - ); - } - } - } else { - error!("PeerRegistry not initialized"); + if !registered { + Client::fail_queued_downloads(&context.client_context, &peer.username); } + registered } fn handle_file_connection( @@ -280,13 +304,13 @@ fn handle_pierce_firewall( 0, 0, ); - handle_peer_connection(peer, stream, reader, context); - - // Inbound peers don't self-announce, so nudge the client to flush any - // downloads queued for this now-connected peer. - let _ = context - .client_sender - .send(ClientOperation::PeerConnected(username)); + if handle_peer_connection(peer, stream, reader, context) { + // Inbound peers don't self-announce, so nudge the client to flush any + // downloads queued for this now-connected peer. + let _ = context + .client_sender + .send(ClientOperation::PeerConnected(username)); + } } fn handle_incoming_connection( @@ -304,15 +328,17 @@ fn handle_incoming_connection( let mut reader = MessageReader::new(); // A peer that dials and then goes away is routine, not an error. - let message = match read_peer_init_message(&mut stream, &mut reader) { - Ok(message) => message, - Err(e) => { - debug!( - "[listener:{peer_ip}:{peer_port}] no peer init message: {e}" - ); - return; - } - }; + let deadline = Instant::now() + PEER_INIT_TIMEOUT; + let message = + match read_peer_init_message(&mut stream, &mut reader, deadline) { + Ok(message) => message, + Err(e) => { + debug!( + "[listener:{peer_ip}:{peer_port}] no peer init message: {e}" + ); + return; + } + }; drop(permit); // A firewalled peer brokered through the server connects back with a @@ -354,13 +380,14 @@ fn handle_incoming_connection( match init_data.connection_type { ConnectionType::P => { - handle_peer_connection(peer, stream, reader, &context); - // Inbound peers don't self-announce, so nudge the client to flush - // anything queued for this now-connected peer — same as the - // PierceFirewall path. - let _ = context - .client_sender - .send(ClientOperation::PeerConnected(init_data.username)); + if handle_peer_connection(peer, stream, reader, &context) { + // Inbound peers don't self-announce, so nudge the client to + // flush anything queued for this now-connected peer — same as + // the PierceFirewall path. + let _ = context + .client_sender + .send(ClientOperation::PeerConnected(init_data.username)); + } } ConnectionType::F => handle_file_connection( @@ -444,3 +471,68 @@ impl Listen { } } } + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + use std::net::TcpListener; + + #[test] + fn a_dribbling_handshake_hits_the_total_deadline() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let writer = std::thread::spawn(move || { + let mut socket = TcpStream::connect(addr).unwrap(); + let _ = socket.write_all(&1000u32.to_le_bytes()); + for _ in 0..20 { + if socket.write_all(&[0u8]).is_err() { + return; + } + std::thread::sleep(Duration::from_millis(50)); + } + }); + let (mut stream, _) = listener.accept().unwrap(); + let mut reader = MessageReader::new(); + + let started = Instant::now(); + let result = read_peer_init_message( + &mut stream, + &mut reader, + started + Duration::from_millis(300), + ); + + assert!(result.is_err(), "a dribbled handshake must not succeed"); + assert!( + started.elapsed() < Duration::from_secs(5), + "the deadline must bound the whole handshake, not each read" + ); + let _ = writer.join(); + } + + #[test] + fn an_oversized_handshake_is_rejected() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let writer = std::thread::spawn(move || { + let mut socket = TcpStream::connect(addr).unwrap(); + let junk = vec![0xFFu8; MAX_PEER_INIT_BYTES * 2]; + let _ = socket.write_all(&junk); + std::thread::sleep(Duration::from_millis(500)); + }); + let (mut stream, _) = listener.accept().unwrap(); + let mut reader = MessageReader::new(); + + let result = read_peer_init_message( + &mut stream, + &mut reader, + Instant::now() + Duration::from_secs(5), + ); + + assert!( + result.is_err(), + "an oversized handshake must be rejected, got {result:?}" + ); + let _ = writer.join(); + } +}