Merge pull request #37 from michel/issue-32-33-thread-flood

fix(client): bound the threads a peer flood can create

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<String, (u64, ActorHandle<PeerMessage>)>;
+type PeerMap = HashMap<String, (u64, ActorHandle<PeerMessage>, Instant)>;
 
 pub struct PeerRegistry {
     peers: Arc<Mutex<PeerMap>>,
     actor_system: Arc<ActorSystem>,
     client_channel: Sender<ClientOperation>,
     own_username: String,
+    capacity: Arc<AtomicUsize>,
 }
 
 impl PeerRegistry {
@@ -33,12 +35,28 @@ impl PeerRegistry {
         actor_system: Arc<ActorSystem>,
         client_channel: Sender<ClientOperation>,
         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<ActorSystem>,
+        client_channel: Sender<ClientOperation>,
+        own_username: String,
+        capacity: Arc<AtomicUsize>,
     ) -> 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<TcpStream>,
         reader: Option<MessageReader>,
+    ) -> Result<ActorHandle<PeerMessage>, String> {
+        self.register_peer_protected(peer, stream, reader, &HashSet::new())
+    }
+
+    pub(crate) fn register_peer_protected(
+        &self,
+        peer: Peer,
+        stream: Option<TcpStream>,
+        reader: Option<MessageReader>,
+        protected: &HashSet<String>,
     ) -> Result<ActorHandle<PeerMessage>, 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)) =
+        if let Some((_, old_handle, _)) =
-            peers.insert(username.clone(), (id, handle.clone()))
+            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<ActorHandle<PeerMessage>> {
         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) {
+        if peers
-            let removed = peers.remove(username).map(|(_, handle)| handle);
+            .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<crate::client::ClientOperation>,
+    ) {
+        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<String> =
+            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<String> =
+            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() {

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() {
@@ -206,6 +214,7 @@ impl Client {
             ConnectionType::P => {
                 let username = peer.username;
 
+                let refused = {
                     let context = match client_context.read_safe() {
                         Ok(c) => c,
                         Err(e) => {
@@ -214,17 +223,27 @@ impl Client {
                         }
                     };
                     if let Some(ref registry) = context.peer_registry {
-                    match registry.register_peer(peer_clone, stream, None) {
+                        let protected = context.protected_peers();
-                        Ok(_) => (),
+                        match registry.register_peer_protected(
+                            peer_clone, stream, None, &protected,
+                        ) {
+                            Ok(_) => false,
                             Err(e) => {
-                            trace!(
+                                warn!(
-                                "Failed to spawn peer actor for {:?}: {:?}",
+                                    "[client] refusing peer connection for \
+                                     {:?}: {:?}",
                                     username, e
                                 );
+                                true
                             }
                         }
                     } else {
                         trace!("PeerRegistry not initialized");
+                        false
+                    }
+                };
+                if refused {
+                    Self::fail_queued_downloads(&client_context, &username);
                 }
             }
 

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(5);
+
 /// 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<UserMessage>,
     /// 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<u32, String>,
+    pending_connect_tokens: HashMap<u32, (String, Instant)>,
+    max_peers: Arc<AtomicUsize>,
     /// Files we share with peers (read-only after connect).
     pub shares: Arc<Shares>,
     /// The directories the current share index was built from.
@@ -311,6 +316,7 @@ pub struct ClientContext {
     pending_serves: HashMap<String, Vec<u32>>,
     /// Shared-file listings received from peers we browsed.
     browse_results: HashMap<String, Vec<SharedDirectory>>,
+    pending_browses: HashMap<String, Instant>,
     /// Latest snapshot of the public chat-room list (from `RoomList`, code 64).
     room_list: Vec<RoomInfo>,
     /// 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<SharedDirectory>,
     ) {
+        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,60 @@ 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<String> {
-        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<String> {
+        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<String> {
+        let mut protected: HashSet<String> = 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.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();
+        protected.extend(
+            self.pending_browses
+                .iter()
+                .filter(|(_, deadline)| **deadline > now)
+                .map(|(username, _)| username.clone()),
+        );
+        protected
     }
 
     /// Record a private message received from another user.

soulseek-rs-lib/src/client/operations.rs

@@ -1,12 +1,14 @@
 use super::{
-    Arc, BROKER_CONNECT_TIMEOUT, Client, ClientContext, ClientOperation,
+    Arc, Client, ClientContext, ClientOperation, ConnectionType, Download,
-    ConnectionType, Download, DownloadPeer, DownloadStatus, Peer, PeerMessage,
+    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<ClientOperation>,
@@ -14,9 +16,21 @@ 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) => {
+                        if matches!(peer.connection_type, ConnectionType::F) {
                             let client_context_clone = client_context.clone();
                             let own_username_clone = own_username.clone();
 
@@ -28,6 +42,14 @@ impl Client {
                                     None,
                                 );
                             });
+                        } else {
+                            Self::connect_to_peer(
+                                peer,
+                                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,
+                                client_context.clone(),
-                                    own_username_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
+    fn sweep_expired_connects(client_context: &Arc<RwLock<ClientContext>>) {
-                        // consumes the token, fail the peer's queued
+        let expired = match client_context.write_safe() {
-                        // downloads (so the caller's Receiver unblocks)
+            Ok(mut ctx) => ctx.take_expired_connects(Instant::now()),
-                        // and reclaim the token. A successful pierce
+            Err(_) => return,
-                        // takes the token first, making this a no-op.
+        };
-                        let timeout_ctx = client_context.clone();
+        for username in expired {
-                        let timeout_user = username.clone();
+            Self::fail_queued_downloads(client_context, &username);
-                        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,
-                                );
-                            }
-                        });
         }
     }
 }
-        });
-    }
-}

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()

soulseek-rs-lib/src/client/tests.rs

@@ -395,3 +395,146 @@ 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"));
+}
+
+#[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"
+    );
+}

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)?;

soulseek-rs-lib/src/peer/listen.rs

@@ -3,15 +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::{DownloadStatus, debug, error, info, trace};
+use crate::utils::semaphore::{Permit, Semaphore};
+use crate::{DownloadStatus, debug, error, info, trace, warn};
 
 /// How long to wait before accepting again after a failure.
 ///
@@ -25,6 +26,10 @@ 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;
+
+const MAX_PEER_INIT_BYTES: usize = 4096;
+
 #[derive(Clone)]
 struct ConnectionContext {
     client_sender: Sender<ClientOperation>,
@@ -41,16 +46,30 @@ struct PeerInitData {
 fn read_peer_init_message(
     stream: &mut TcpStream,
     reader: &mut MessageReader,
+    deadline: Instant,
 ) -> io::Result<Message> {
     // 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.
@@ -119,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
@@ -128,30 +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 registered = {
         let client_context = match context.client_context.read_safe() {
             Ok(c) => c,
             Err(e) => {
                 error!("[listener] handle_peer_connection lock: {}", e);
-            return;
+                return false;
             }
         };
         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();
-            Ok(_) => (),
+            match registry.register_peer_protected(
+                peer.clone(),
+                Some(stream),
+                Some(reader),
+                &protected,
+            ) {
+                Ok(_) => true,
                 Err(e) => {
-                error!(
+                    warn!(
-                    "Failed to spawn peer actor for {:?}: {:?}",
+                        "[listener] refusing peer connection for {:?}: {:?}",
                         peer.username, e
                     );
+                    false
                 }
             }
         } else {
             error!("PeerRegistry not initialized");
+            false
         }
+    };
+    if !registered {
+        Client::fail_queued_downloads(&context.client_context, &peer.username);
+    }
+    registered
 }
 
 fn handle_file_connection(
@@ -271,18 +304,19 @@ fn handle_pierce_firewall(
         0,
         0,
     );
-    handle_peer_connection(peer, stream, reader, context);
+    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(
     mut stream: TcpStream,
     context: ConnectionContext,
+    permit: Permit,
 ) {
     let Ok(peer_addr) = stream.peer_addr() else {
         error!("[listener] failed to get peer address");
@@ -294,7 +328,9 @@ 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) {
+    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!(
@@ -303,6 +339,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).
@@ -343,14 +380,15 @@ fn handle_incoming_connection(
 
     match init_data.connection_type {
         ConnectionType::P => {
-            handle_peer_connection(peer, stream, reader, &context);
+            if handle_peer_connection(peer, stream, reader, &context) {
-            // Inbound peers don't self-announce, so nudge the client to flush
+                // Inbound peers don't self-announce, so nudge the client to
-            // anything queued for this now-connected peer — same as the
+                // flush anything queued for this now-connected peer — same as
-            // PierceFirewall path.
+                // the PierceFirewall path.
                 let _ = context
                     .client_sender
                     .send(ClientOperation::PeerConnected(init_data.username));
             }
+        }
 
         ConnectionType::F => handle_file_connection(
             peer,
@@ -406,6 +444,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 +460,79 @@ 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);
+            });
         }
     }
 }
+
+#[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();
+    }
+}

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

soulseek-rs-lib/src/utils/semaphore.rs

@@ -0,0 +1,73 @@
+use std::sync::{Arc, Condvar, Mutex, PoisonError};
+
+pub struct Semaphore {
+    permits: Mutex<usize>,
+    available: Condvar,
+}
+
+pub struct Permit(Arc<Semaphore>);
+
+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<Self>) -> 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"
+        );
+    }
+}