diff --git a/soulseek-rs-lib/src/download_store.rs b/soulseek-rs-lib/src/download_store.rs index c07f89b..532a0a3 100644 --- a/soulseek-rs-lib/src/download_store.rs +++ b/soulseek-rs-lib/src/download_store.rs @@ -142,8 +142,8 @@ impl DownloadStore { } } -/// Returns the tokens of downloads matching `username` (and optionally a -/// `filename`) after notifying their senders of `Failed`. +/// Returns the tokens of queued downloads matching `username` (and optionally +/// 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(); diff --git a/soulseek-rs-lib/src/peer/listen.rs b/soulseek-rs-lib/src/peer/listen.rs index 3536be2..18cf6c4 100644 --- a/soulseek-rs-lib/src/peer/listen.rs +++ b/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( - reader: &mut MessageReader, - client_context: &Arc>, - username: &str, - peer_ip: &str, - peer_port: u32, -) -> Option { - if reader.buffer_len() == 0 { - return None; - } +fn take_transfer_token(reader: &mut MessageReader) -> Option { let buffer = reader.get_buffer(); - let token = u32::from_le_bytes(buffer.get(0..4)?.try_into().ok()?); - trace!( - "[listener:{}] got transfer_token: {} from data chunk", - username, token - ); + Some(u32::from_le_bytes(buffer.get(0..4)?.try_into().ok()?)) +} - let context = match client_context.read_safe() { - Ok(c) => c, - Err(e) => { - error!("[listener] client context lock: {}", e); +fn wait_for_download_registration( + client_context: &Arc>, + token: u32, + username: &str, + deadline: Instant, +) -> Option { + 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; + } + }; + if download.is_some() { + return download; + } + if Instant::now() >= deadline { + if let Ok(context) = client_context.read_safe() { + trace!( + "[listener:{}] no claimable download for token {}, known tokens: {:?}", + username, + token, + context.get_download_tokens() + ); + } return None; } - }; - let download = context.get_download_by_token(token).cloned(); - - if download.is_none() { - let download_tokens = context.get_download_tokens(); - trace!( - "[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}", - token, download_tokens - ); + thread::sleep(Duration::from_millis(50)); } - - download } fn handle_peer_connection( @@ -201,13 +214,15 @@ fn handle_file_connection( return; } - let download = extract_download_from_buffer( - &mut reader, - &context.client_context, - &peer.username, - &peer.host, - peer.port, - ); + let download = + take_transfer_token(&mut reader).and_then(|transfer_token| { + wait_for_download_registration( + &context.client_context, + transfer_token, + &peer.username, + Instant::now() + TOKEN_REGISTRATION_TIMEOUT, + ) + }); 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(); diff --git a/soulseek-rs-lib/tests/e2e.rs b/soulseek-rs-lib/tests/e2e.rs index 39c3433..072bc21 100644 --- a/soulseek-rs-lib/tests/e2e.rs +++ b/soulseek-rs-lib/tests/e2e.rs @@ -520,36 +520,43 @@ fn connect_retry(addr: &str, timeout: Duration) -> std::io::Result { } } -fn run_mock_uploader(cfg: &MockUpload) -> std::io::Result { - // 1. P (control) connection: register ourselves with the downloader. +fn open_control_and_await_queue( + cfg: &MockUpload, +) -> std::io::Result { 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 - // becomes the download's expected size, so it must match the content. +/// Build a TransferRequest (peer code 40). The size becomes the download's +/// 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 { + 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 { } } - // 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 { ) } +fn run_mock_uploader_f_first(cfg: &MockUpload) -> std::io::Result { + 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))?; - 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; - } - } + let mut p = open_control_and_await_queue(cfg)?; 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, +) { 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(