fix(peer): survive download token races on the listener path

A P-connection drop failed every download for its user, including ones
actively transferring on their own healthy F connection:
collect_failed_tokens now only fails queued downloads. And an F
connection could look its token up before the client-operations thread
had registered it: the listener now waits for the registration, and
claims it only for the matching peer and only when not already claimed.

Fixes #42

soulseek-rs-lib/src/download_store.rs

@@ -142,8 +142,8 @@ impl DownloadStore {
     }
 }
 
-/// Returns the tokens of downloads matching `username` (and optionally a
+/// Returns the tokens of queued downloads matching `username` (and optionally
-/// `filename`) after notifying their senders of `Failed`.
+/// a `filename`) after notifying their senders of `Failed`.
 ///
 /// Caller is responsible for then calling `update_status` and `remove` for
 /// each token, typically under a write lock.
@@ -159,6 +159,7 @@ pub fn collect_failed_tokens(
         .filter(|d| {
             d.username == username && filename.is_none_or(|f| d.filename == *f)
         })
+        .filter(|d| matches!(d.status, DownloadStatus::Queued))
         .map(|d| {
             let _ = d.sender.send(DownloadStatus::Failed(Some(
                 "The upload failed on the other side".to_string(),
@@ -305,6 +306,61 @@ mod tests {
         assert!(!store.remove_by_file("peer", "song.mp3"), "idempotent");
     }
 
+    #[test]
+    fn collect_failed_tokens_collects_only_queued_downloads() {
+        let mut store = DownloadStore::new();
+        let (tx_active, rx_active) = mpsc::channel();
+        let (tx_paused, rx_paused) = mpsc::channel();
+        let (tx_completed, rx_completed) = mpsc::channel();
+        let (tx_queued, _rx_queued) = mpsc::channel();
+
+        let mut active = make_download(
+            1,
+            DownloadStatus::InProgress {
+                bytes_downloaded: 25,
+                total_bytes: 100,
+                speed_bytes_per_sec: 10.0,
+            },
+        );
+        active.sender = tx_active;
+
+        let mut paused = make_download(
+            2,
+            DownloadStatus::Paused {
+                bytes_downloaded: 25,
+                total_bytes: 100,
+            },
+        );
+        paused.sender = tx_paused;
+
+        let mut completed = make_download(3, DownloadStatus::Completed);
+        completed.sender = tx_completed;
+
+        let mut queued = make_download(4, DownloadStatus::Queued);
+        queued.sender = tx_queued;
+
+        store.add(active);
+        store.add(paused);
+        store.add(completed);
+        store.add(queued);
+
+        let tokens = collect_failed_tokens(&store, "peer", None);
+
+        assert_eq!(tokens, vec![4], "only the queued download may fail");
+        assert!(
+            rx_active.try_recv().is_err(),
+            "an in-progress transfer must not be notified of failure"
+        );
+        assert!(
+            rx_paused.try_recv().is_err(),
+            "a paused transfer must not be notified of failure"
+        );
+        assert!(
+            rx_completed.try_recv().is_err(),
+            "a completed download must not be notified of failure"
+        );
+    }
+
     #[test]
     fn collect_failed_tokens_notifies_and_lists_matching() {
         let mut store = DownloadStore::new();

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

@@ -26,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 TOKEN_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(10);
+
 const MAX_HANDSHAKES: usize = 128;
 
 const MAX_PEER_INIT_BYTES: usize = 4096;
@@ -116,41 +118,52 @@ fn wait_for_transfer_token(
     Ok(())
 }
 
-fn extract_download_from_buffer(
+fn take_transfer_token(reader: &mut MessageReader) -> Option<u32> {
-    reader: &mut MessageReader,
-    client_context: &Arc<RwLock<ClientContext>>,
-    username: &str,
-    peer_ip: &str,
-    peer_port: u32,
-) -> Option<Download> {
-    if reader.buffer_len() == 0 {
-        return None;
-    }
     let buffer = reader.get_buffer();
-    let token = u32::from_le_bytes(buffer.get(0..4)?.try_into().ok()?);
+    Some(u32::from_le_bytes(buffer.get(0..4)?.try_into().ok()?))
-    trace!(
+}
-        "[listener:{}] got transfer_token: {} from data chunk",
-        username, token
-    );
 
-    let context = match client_context.read_safe() {
+fn wait_for_download_registration(
-        Ok(c) => c,
+    client_context: &Arc<RwLock<ClientContext>>,
+    token: u32,
+    username: &str,
+    deadline: Instant,
+) -> Option<Download> {
+    loop {
+        let download = match client_context.read_safe() {
+            Ok(context) => context
+                .get_download_by_token(token)
+                .filter(|d| {
+                    d.username == username
+                        && !matches!(
+                            d.status,
+                            DownloadStatus::InProgress { .. }
+                                | DownloadStatus::Paused { .. }
+                                | DownloadStatus::Completed
+                        )
+                })
+                .cloned(),
             Err(e) => {
                 error!("[listener] client context lock: {}", e);
                 return None;
             }
         };
-    let download = context.get_download_by_token(token).cloned();
+        if download.is_some() {
-
+            return download;
-    if download.is_none() {
+        }
-        let download_tokens = context.get_download_tokens();
+        if Instant::now() >= deadline {
+            if let Ok(context) = client_context.read_safe() {
                 trace!(
-            "[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}",
+                    "[listener:{}] no claimable download for token {}, known tokens: {:?}",
-            token, download_tokens
+                    username,
+                    token,
+                    context.get_download_tokens()
                 );
             }
-
+            return None;
-    download
+        }
+        thread::sleep(Duration::from_millis(50));
+    }
 }
 
 fn handle_peer_connection(
@@ -201,13 +214,15 @@ fn handle_file_connection(
         return;
     }
 
-    let download = extract_download_from_buffer(
+    let download =
-        &mut reader,
+        take_transfer_token(&mut reader).and_then(|transfer_token| {
+            wait_for_download_registration(
                 &context.client_context,
+                transfer_token,
                 &peer.username,
-        &peer.host,
+                Instant::now() + TOKEN_REGISTRATION_TIMEOUT,
-        peer.port,
+            )
-    );
+        });
     let failure_token = download.as_ref().map(|d| d.token);
 
     let download_peer = DownloadPeer::new(
@@ -474,8 +489,116 @@ impl Listen {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::types::DownloadMetadata;
     use std::io::Write;
     use std::net::TcpListener;
+    use std::sync::mpsc;
+
+    fn registered_download(token: u32) -> Download {
+        Download {
+            username: "peer".to_string(),
+            filename: "song.mp3".to_string(),
+            token,
+            size: 100,
+            download_directory: std::env::temp_dir().display().to_string(),
+            status: DownloadStatus::Queued,
+            sender: mpsc::channel().0,
+            queue_position: None,
+            metadata: DownloadMetadata::default(),
+        }
+    }
+
+    #[test]
+    fn a_download_registered_during_the_wait_is_found() {
+        let client_context = Arc::new(RwLock::new(ClientContext::new()));
+        let registering = {
+            let client_context = client_context.clone();
+            std::thread::spawn(move || {
+                std::thread::sleep(Duration::from_millis(100));
+                client_context
+                    .write()
+                    .unwrap()
+                    .add_download(registered_download(7));
+            })
+        };
+
+        let download = wait_for_download_registration(
+            &client_context,
+            7,
+            "peer",
+            Instant::now() + Duration::from_secs(5),
+        );
+
+        assert!(
+            download.is_some(),
+            "a token registered during the wait must be found"
+        );
+        let _ = registering.join();
+    }
+
+    #[test]
+    fn an_unregistered_token_gives_up_at_the_deadline() {
+        let client_context = Arc::new(RwLock::new(ClientContext::new()));
+
+        let started = Instant::now();
+        let download = wait_for_download_registration(
+            &client_context,
+            7,
+            "peer",
+            started + Duration::from_millis(200),
+        );
+
+        assert!(download.is_none(), "an unknown token must not be invented");
+        assert!(
+            started.elapsed() < Duration::from_secs(2),
+            "the deadline must bound the wait"
+        );
+    }
+
+    #[test]
+    fn a_download_for_another_peer_is_not_claimed() {
+        let client_context = Arc::new(RwLock::new(ClientContext::new()));
+        client_context
+            .write()
+            .unwrap()
+            .add_download(registered_download(7));
+
+        let download = wait_for_download_registration(
+            &client_context,
+            7,
+            "impostor",
+            Instant::now() + Duration::from_millis(200),
+        );
+
+        assert!(
+            download.is_none(),
+            "a token must only match the peer it was registered for"
+        );
+    }
+
+    #[test]
+    fn an_already_claimed_download_is_not_reclaimed() {
+        let client_context = Arc::new(RwLock::new(ClientContext::new()));
+        let mut claimed = registered_download(7);
+        claimed.status = DownloadStatus::InProgress {
+            bytes_downloaded: 25,
+            total_bytes: 100,
+            speed_bytes_per_sec: 10.0,
+        };
+        client_context.write().unwrap().add_download(claimed);
+
+        let download = wait_for_download_registration(
+            &client_context,
+            7,
+            "peer",
+            Instant::now() + Duration::from_millis(200),
+        );
+
+        assert!(
+            download.is_none(),
+            "a transfer already in progress must not be claimed again"
+        );
+    }
 
     #[test]
     fn a_dribbling_handshake_hits_the_total_deadline() {
@@ -557,7 +680,7 @@ mod tests {
 
         assert!(result.is_err(), "a silent peer must not park the thread");
         assert!(
-            started.elapsed() < Duration::from_secs(5),
+            started.elapsed() < Duration::from_secs(2),
             "the deadline must bound the wait"
         );
         let _ = writer.join();

soulseek-rs-lib/tests/e2e.rs

@@ -520,36 +520,43 @@ fn connect_retry(addr: &str, timeout: Duration) -> std::io::Result<TcpStream> {
     }
 }
 
-fn run_mock_uploader(cfg: &MockUpload) -> std::io::Result<u64> {
+fn open_control_and_await_queue(
-    // 1. P (control) connection: register ourselves with the downloader.
+    cfg: &MockUpload,
+) -> std::io::Result<TcpStream> {
     let mut p = connect_retry(&cfg.listen_addr, Duration::from_secs(5))?;
     p.set_read_timeout(Some(Duration::from_secs(10)))?;
     p.write_all(&peer_init_bytes(&cfg.peer_username, "P", 0))?;
     p.flush()?;
     let _ = cfg.ready.send(());
 
-    // 2. Wait for the downloader's QueueUpload (peer code 43).
     loop {
-        let mut msg = read_framed(&mut p)?;
+        let msg = read_framed(&mut p)?;
         if msg.get_message_code() == 43 {
-            msg.set_pointer(8);
-            let _requested = msg.read_string();
             break;
         }
     }
+    Ok(p)
+}
 
-    // 3. Offer the transfer with a TransferRequest (peer code 40). The size here
+/// Build a TransferRequest (peer code 40). The size becomes the download's
-    //    becomes the download's expected size, so it must match the content.
+/// expected size, so it must match the content.
+fn transfer_request(cfg: &MockUpload) -> Message {
     let mut tr = Message::new();
     tr.write_int32(40)
         .write_int32(1) // direction: upload
         .write_int32(cfg.token)
         .write_string(&cfg.filename)
         .write_int64(cfg.content.len() as u64);
-    p.write_all(&tr.get_buffer())?;
+    tr
+}
+
+fn run_mock_uploader(cfg: &MockUpload) -> std::io::Result<u64> {
+    let mut p = open_control_and_await_queue(cfg)?;
+
+    p.write_all(&transfer_request(cfg).get_buffer())?;
     p.flush()?;
 
-    // 4. Wait for the downloader to allow it (TransferResponse, peer code 41).
+    // Wait for the downloader to allow it (TransferResponse, peer code 41).
     loop {
         let msg = read_framed(&mut p)?;
         if msg.get_message_code() == 41 {
@@ -557,7 +564,6 @@ fn run_mock_uploader(cfg: &MockUpload) -> std::io::Result<u64> {
         }
     }
 
-    // 5. Open the F (file) connection to the downloader's listener and stream.
     serve_file_over_f(
         &cfg.listen_addr,
         &cfg.peer_username,
@@ -567,22 +573,38 @@ fn run_mock_uploader(cfg: &MockUpload) -> std::io::Result<u64> {
     )
 }
 
+fn run_mock_uploader_f_first(cfg: &MockUpload) -> std::io::Result<u64> {
+    let mut p = open_control_and_await_queue(cfg)?;
+
+    let listen_addr = cfg.listen_addr.clone();
+    let peer_username = cfg.peer_username.clone();
+    let token = cfg.token;
+    let content = cfg.content.clone();
+    let token_delay = cfg.token_delay;
+    let f = std::thread::spawn(move || {
+        serve_file_over_f(
+            &listen_addr,
+            &peer_username,
+            token,
+            &content,
+            token_delay,
+        )
+    });
+
+    std::thread::sleep(Duration::from_millis(300));
+
+    p.write_all(&transfer_request(cfg).get_buffer())?;
+    p.flush()?;
+
+    f.join()
+        .map_err(|_| std::io::Error::other("F connection thread panicked"))?
+}
+
 fn run_refusing_uploader(
     cfg: &MockUpload,
     refusal: &Message,
 ) -> std::io::Result<()> {
-    let mut p = connect_retry(&cfg.listen_addr, Duration::from_secs(5))?;
+    let mut p = open_control_and_await_queue(cfg)?;
-    p.set_read_timeout(Some(Duration::from_secs(10)))?;
-    p.write_all(&peer_init_bytes(&cfg.peer_username, "P", 0))?;
-    p.flush()?;
-    let _ = cfg.ready.send(());
-
-    loop {
-        let msg = read_framed(&mut p)?;
-        if msg.get_message_code() == 43 {
-            break;
-        }
-    }
 
     p.write_all(&refusal.get_buffer())?;
     p.flush()?;
@@ -724,7 +746,11 @@ fn unique_download_dir() -> PathBuf {
     dir
 }
 
-fn file_downloads_over_p_and_f(peer_username: &str, token_delay: Duration) {
+fn file_downloads_over_p_and_f(
+    peer_username: &str,
+    token_delay: Duration,
+    run_uploader: fn(&MockUpload) -> std::io::Result<u64>,
+) {
     let server = server_or_skip!();
 
     // Downloader: connected to the server with its peer listener enabled.
@@ -755,7 +781,7 @@ fn file_downloads_over_p_and_f(peer_username: &str, token_delay: Duration) {
         token_delay,
     };
     let uploader = std::thread::spawn(move || {
-        if let Err(e) = run_mock_uploader(&cfg) {
+        if let Err(e) = run_uploader(&cfg) {
             eprintln!("[mock uploader] {e}");
         }
     });
@@ -792,12 +818,29 @@ fn file_downloads_over_p_and_f(peer_username: &str, token_delay: Duration) {
 
 #[test]
 fn a_file_downloads_from_a_peer_over_p_and_f_connections() {
-    file_downloads_over_p_and_f("e2e_mockpeer", Duration::ZERO);
+    file_downloads_over_p_and_f(
+        "e2e_mockpeer",
+        Duration::ZERO,
+        run_mock_uploader,
+    );
 }
 
 #[test]
 fn a_file_downloads_when_the_transfer_token_arrives_late() {
-    file_downloads_over_p_and_f("e2e_latepeer", Duration::from_millis(300));
+    file_downloads_over_p_and_f(
+        "e2e_latepeer",
+        Duration::from_millis(300),
+        run_mock_uploader,
+    );
+}
+
+#[test]
+fn a_download_completes_when_the_f_connection_beats_token_registration() {
+    file_downloads_over_p_and_f(
+        "e2e_earlypeer",
+        Duration::ZERO,
+        run_mock_uploader_f_first,
+    );
 }
 
 fn refusal_fails_the_download_quickly(