From 926c1c6697226379b274968ff3aafd2d2263f79b Mon Sep 17 00:00:00 2001 From: Michel de Graaf Date: Wed, 9 Sep 2026 15:44:06 +0200 Subject: [PATCH] feat(download): resume interrupted transfers --- soulseek-rs-lib/src/client/connection.rs | 23 +- soulseek-rs-lib/src/client/operations.rs | 91 +---- soulseek-rs-lib/src/client/tests.rs | 16 +- soulseek-rs-lib/src/download_store.rs | 80 +++- soulseek-rs-lib/src/peer/download_peer.rs | 465 ++++++++++++++++------ soulseek-rs-lib/src/peer/listen.rs | 50 +-- soulseek-rs-lib/src/peer/upload_peer.rs | 229 +++++++---- soulseek-rs-lib/tests/e2e.rs | 2 +- soulseek-rs-lib/tests/resume_e2e.rs | 408 +++++++++++++++++++ 9 files changed, 1032 insertions(+), 332 deletions(-) create mode 100644 soulseek-rs-lib/tests/resume_e2e.rs diff --git a/soulseek-rs-lib/src/client/connection.rs b/soulseek-rs-lib/src/client/connection.rs index e6d0436..3d4f36f 100644 --- a/soulseek-rs-lib/src/client/connection.rs +++ b/soulseek-rs-lib/src/client/connection.rs @@ -1,8 +1,8 @@ use super::{ Arc, Client, ClientContext, ClientOperation, ConnectionType, DownloadPeer, - DownloadStatus, Listen, Peer, PeerRegistry, Receiver, Result, RwLock, - RwLockExt, Sender, ServerActor, ServerMessage, Shares, SoulseekRs, - TcpStream, error, info, mpsc, thread, trace, warn, + Listen, Peer, PeerRegistry, Receiver, Result, RwLock, RwLockExt, Sender, + ServerActor, ServerMessage, Shares, SoulseekRs, TcpStream, error, info, + mpsc, thread, trace, warn, }; /// Ceiling on the wait for a login verdict. Generous enough for a slow server @@ -282,27 +282,12 @@ impl Client { own_username, ); - match download_peer.download_file( - client_context.clone(), - None, - None, - ) { + match download_peer.download_file(client_context, None, None) { Ok((download, filename)) => { trace!( "[client] downloaded {} bytes {:?} ", filename, download.size ); - let _ = download.sender.send(DownloadStatus::Completed); - match client_context.write_safe() { - Ok(mut ctx) => ctx.update_download_with_status( - download.token, - DownloadStatus::Completed, - ), - Err(e) => error!( - "[client] connect_to_peer F write: {}", - e - ), - } } Err(e) => { trace!("[client] failed to download: {}", e); diff --git a/soulseek-rs-lib/src/client/operations.rs b/soulseek-rs-lib/src/client/operations.rs index bedf5a3..81dcd8e 100644 --- a/soulseek-rs-lib/src/client/operations.rs +++ b/soulseek-rs-lib/src/client/operations.rs @@ -135,52 +135,33 @@ impl Client { } } ClientOperation::DownloadFromPeer(token, peer, allowed) => { - let maybe_download = match client_context.read_safe() { - Ok(ctx) => { - ctx.get_download_by_token(token).cloned() - } + let maybe_download = match client_context.write_safe() { + Ok(mut ctx) => ctx + .downloads + .claim_for_peer(token, &peer.username), Err(e) => { - error!("[client] DownloadFromPeer read: {}", e); + error!( + "[client] DownloadFromPeer write: {}", + e + ); continue; } }; - let own_username = own_username.clone(); - let client_context_clone = client_context.clone(); - trace!( "[client] DownloadFromPeer token: {} peer: {:?}", token, peer ); let Some(download) = maybe_download else { - error!( - "Can't find download with token {:?}", - token + debug!( + "[client] transfer token {} is missing, belongs \ + to another peer, or is already claimed; ignoring", + token, ); continue; }; - if matches!( - download.status, - DownloadStatus::InProgress { .. } - | DownloadStatus::Completed - | DownloadStatus::Cancelled - ) { - 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, - }, - ); - } + + let own_username = own_username.clone(); + let client_context = client_context.clone(); thread::spawn(move || { let download_peer = DownloadPeer::new( @@ -191,35 +172,12 @@ impl Client { allowed, own_username, ); - let Some(filename) = - download.filename.split('\\').next_back() - else { - error!( - "Cant find filename to save download: {:?}", - download.filename - ); - return; - }; match download_peer.download_file( - client_context_clone.clone(), + client_context, Some(download.clone()), None, ) { Ok((download, filename)) => { - let _ = download - .sender - .send(DownloadStatus::Completed); - match client_context_clone.write_safe() { - Ok(mut ctx) => ctx - .update_download_with_status( - download.token, - DownloadStatus::Completed, - ), - Err(e) => error!( - "[client] download complete write: {}", - e - ), - } info!( "Successfully downloaded {} bytes to {}", download.size, filename @@ -227,24 +185,9 @@ impl Client { } Err(DownloadError::Cancelled) => {} Err(e) => { - let reason = Some(e.to_string()); - let _ = download.sender.send( - DownloadStatus::Failed(reason.clone()), - ); - match client_context_clone.write_safe() { - Ok(mut ctx) => ctx - .update_download_with_status( - download.token, - DownloadStatus::Failed(reason), - ), - Err(e) => error!( - "[client] download failed write: {}", - e - ), - } error!( "Failed to download file '{}' from {}:{} (token: {}) - Error: {}", - filename, + download.filename, peer.host, peer.port, download.token, diff --git a/soulseek-rs-lib/src/client/tests.rs b/soulseek-rs-lib/src/client/tests.rs index 1c16b42..fb97d6c 100644 --- a/soulseek-rs-lib/src/client/tests.rs +++ b/soulseek-rs-lib/src/client/tests.rs @@ -567,16 +567,17 @@ 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 download_dir = std::env::temp_dir() + .join(format!("soulseek-replayed-transfer-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&download_dir); + std::fs::create_dir_all(&download_dir).unwrap(); 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 mut queued = + download("peer", "f.mp3", 9, DownloadStatus::Queued, sender); + queued.download_directory = download_dir.display().to_string(); + client.context.write().unwrap().add_download(queued); let (ops_tx, ops_rx) = mpsc::channel(); Client::listen_to_client_operations( @@ -625,6 +626,7 @@ fn a_replayed_transfer_response_does_not_start_a_second_transfer() { connections, 1, "a replayed TransferResponse must not dial the peer again" ); + let _ = std::fs::remove_dir_all(download_dir); } #[test] diff --git a/soulseek-rs-lib/src/download_store.rs b/soulseek-rs-lib/src/download_store.rs index 3cd9699..02f274a 100644 --- a/soulseek-rs-lib/src/download_store.rs +++ b/soulseek-rs-lib/src/download_store.rs @@ -61,13 +61,45 @@ impl DownloadStore { let Some(download) = self.get_by_token_mut(token) else { return false; }; - if matches!(download.status, DownloadStatus::Cancelled) { + if download.status.is_terminal() + || matches!( + (&download.status, &status), + ( + DownloadStatus::Paused { .. }, + DownloadStatus::InProgress { .. } + ) + ) + { return false; } download.status = status; true } + /// Atomically reserve a queued download for the peer that owns it. + /// + /// Every file-connection direction uses this before opening a partial, so + /// two sockets cannot both observe `Queued` and write the same target. + pub fn claim_for_peer( + &mut self, + token: u32, + username: &str, + ) -> Option { + let download = self.get_by_token_mut(token)?; + if download.username != username + || !matches!(download.status, DownloadStatus::Queued) + { + return None; + } + + download.status = DownloadStatus::InProgress { + bytes_downloaded: 0, + total_bytes: download.size, + speed_bytes_per_sec: 0.0, + }; + Some(download.clone()) + } + pub fn cancel_by_file( &mut self, username: &str, @@ -373,6 +405,52 @@ mod tests { )); } + #[test] + fn progress_cannot_implicitly_resume_a_paused_download() { + let mut store = DownloadStore::new(); + store.add(make_download( + 1, + DownloadStatus::Paused { + bytes_downloaded: 25, + total_bytes: 100, + }, + )); + + assert!(!store.update_status( + 1, + DownloadStatus::InProgress { + bytes_downloaded: 50, + total_bytes: 100, + speed_bytes_per_sec: 10.0, + }, + )); + assert!(matches!( + store.get_by_token(1).unwrap().status, + DownloadStatus::Paused { + bytes_downloaded: 25, + total_bytes: 100 + } + )); + } + + #[test] + fn a_download_can_only_be_claimed_once_and_by_its_peer() { + let mut store = DownloadStore::new(); + store.add(make_download(1, DownloadStatus::Queued)); + + assert!(store.claim_for_peer(1, "other").is_none()); + let claimed = store.claim_for_peer(1, "peer").unwrap(); + assert!(matches!( + claimed.status, + DownloadStatus::InProgress { + bytes_downloaded: 0, + total_bytes: 100, + speed_bytes_per_sec: 0.0 + } + )); + assert!(store.claim_for_peer(1, "peer").is_none()); + } + #[test] fn remove_queued_skips_active_downloads() { let mut store = DownloadStore::new(); diff --git a/soulseek-rs-lib/src/peer/download_peer.rs b/soulseek-rs-lib/src/peer/download_peer.rs index b941d32..034e54f 100644 --- a/soulseek-rs-lib/src/peer/download_peer.rs +++ b/soulseek-rs-lib/src/peer/download_peer.rs @@ -34,10 +34,9 @@ pub enum DownloadError { DownloadInfoMissing(u32), FileWriteError(io::Error), PathResolutionError(String), - InvalidTokenBytes, LockPoisoned, Cancelled, - IncompleteDownload { received: usize, expected: usize }, + IncompleteDownload { received: u64, expected: u64 }, } impl std::fmt::Display for DownloadError { @@ -60,9 +59,6 @@ impl std::fmt::Display for DownloadError { Self::PathResolutionError(msg) => { write!(f, "Path resolution error: {msg}") } - Self::InvalidTokenBytes => { - write!(f, "Invalid token bytes received") - } Self::LockPoisoned => write!(f, "Lock poisoned"), Self::Cancelled => write!(f, "Download cancelled"), Self::IncompleteDownload { received, expected } => write!( @@ -147,8 +143,9 @@ impl PartFile { .map_err(DownloadError::FileWriteError)?; } - // A `.part` that already meets or exceeds the expected size is left over - // from a different (or corrupted) transfer, not a resume point. + // A `.part` at or beyond the advertised size is restarted. We do not + // persist its username, full remote path, or prior expected size, so + // equality alone cannot prove that its bytes belong to this transfer. let on_disk = fs::metadata(&path).map_or(0, |m| m.len()); let written = if on_disk < download.size { on_disk } else { 0 }; @@ -170,8 +167,9 @@ impl PartFile { /// Peers can coalesce trailing bytes into the last chunk, so overshoot past /// the expected size is dropped rather than written. fn write(&mut self, data: &[u8]) -> Result<(), DownloadError> { - let remaining = (self.expected - self.written) as usize; - let data = &data[..data.len().min(remaining)]; + let remaining = self.expected - self.written; + let write_len = remaining.min(data.len() as u64) as usize; + let data = &data[..write_len]; self.file .write_all(data) .map_err(DownloadError::FileWriteError)?; @@ -193,8 +191,8 @@ impl PartFile { fn finish(self) -> Result { if !self.is_complete() { return Err(DownloadError::IncompleteDownload { - received: self.written as usize, - expected: self.expected as usize, + received: self.written, + expected: self.expected, }); } drop(self.file); @@ -319,77 +317,73 @@ impl DownloadPeer { Ok((download, part)) } - fn begin_pierced_download( - &self, - data: &[u8], - stream: &mut TcpStream, - client_context: &Arc>, - ) -> Result<(Download, PartFile), DownloadError> { - let token_bytes = - data.get(0..4).ok_or(DownloadError::InvalidTokenBytes)?; - let token_array: [u8; 4] = token_bytes - .try_into() - .map_err(|_| DownloadError::InvalidTokenBytes)?; - let token_u32 = u32::from_le_bytes(token_array); - - trace!( - "[download_peer:{}] got token: {} from data chunk", - self.username, token_u32 - ); - - let client_guard = client_context - .read() - .map_err(|_| DownloadError::LockPoisoned)?; - let download_info = - client_guard.get_download_by_token(token_u32).cloned(); - drop(client_guard); - - let download = - download_info.ok_or(DownloadError::TokenNotFound(token_u32))?; - - Self::start_transfer(download, stream, client_context) - } - fn read_download_stream( &self, stream: &mut TcpStream, client_context: &Arc>, download: Option, + matched_download: &mut Option, ) -> Result<(Download, String), DownloadError> { let mut read_buffer = vec![0u8; READ_BUFFER_SIZE]; let mut bytes_since_last_update = 0usize; let mut last_update_time = Instant::now(); - let mut last_data = Instant::now(); - stream - .set_read_timeout(Some(CANCEL_POLL)) - .map_err(DownloadError::ConnectionFailed)?; trace!( "[download_peer:{}] Starting to read data from peer", self.username ); - let mut transfer = match download { - Some(dl) => Some(Self::start_transfer(dl, stream, client_context)?), - None => None, + let (download, mut part) = match download { + Some(download) => { + Self::start_transfer(download, stream, client_context)? + } + None if !self.no_pierce => { + let mut token = [0u8; 4]; + stream + .read_exact(&mut token) + .map_err(DownloadError::StreamReadError)?; + let token = u32::from_le_bytes(token); + trace!( + "[download_peer:{}] got token: {} from data chunk", + self.username, token + ); + let download = client_context + .write() + .map_err(|_| DownloadError::LockPoisoned)? + .downloads + .claim_for_peer(token, &self.username) + .ok_or(DownloadError::TokenNotFound(token))?; + *matched_download = Some(download.clone()); + Self::start_transfer(download, stream, client_context)? + } + None => { + return Err(DownloadError::DownloadInfoMissing(self.token)); + } }; + stream + .set_read_timeout(Some(CANCEL_POLL)) + .map_err(DownloadError::ConnectionFailed)?; + let mut last_data = Instant::now(); loop { - if let Some((dl, _)) = &transfer { - match Self::wait_while_paused(client_context, dl) { - Ok(true) => last_data = Instant::now(), - Ok(false) => {} - Err(e) => { - if matches!(e, DownloadError::Cancelled) - && let Some((_, part)) = transfer.take() - { - part.discard(); - } - return Err(e); + match Self::wait_while_paused(client_context, &download) { + Ok(true) => last_data = Instant::now(), + Ok(false) => {} + Err(error) => { + if matches!(error, DownloadError::Cancelled) { + part.discard(); } + return Err(error); } } + // A zero-byte transfer needs no payload bytes. There is no offset + // acknowledgement in Soulseek's transfer protocol, so the uploader + // waits for us to close and we must finalize without another read. + if part.is_complete() { + break; + } + match stream.read(&mut read_buffer) { Ok(0) => { trace!( @@ -400,24 +394,7 @@ impl DownloadPeer { } Ok(bytes_read) => { last_data = Instant::now(); - let data = &read_buffer[..bytes_read]; - - if transfer.is_none() && !self.no_pierce { - transfer = Some(self.begin_pierced_download( - data, - stream, - client_context, - )?); - continue; - } - - let Some((dl, part)) = transfer.as_mut() else { - return Err(DownloadError::DownloadInfoMissing( - self.token, - )); - }; - - part.write(data)?; + part.write(&read_buffer[..bytes_read])?; bytes_since_last_update += bytes_read; if bytes_since_last_update >= PROGRESS_INTERVAL_BYTES { @@ -429,7 +406,7 @@ impl DownloadPeer { }; Self::report_progress( client_context, - dl, + &download, part.written, speed, ); @@ -457,11 +434,9 @@ impl DownloadPeer { self.username ); - let Some((download, part)) = transfer else { - return Err(DownloadError::DownloadInfoMissing(self.token)); - }; - - Ok((download, part.finish()?)) + let final_path = + Self::finish_transfer(client_context, &download, part)?; + Ok((download, final_path)) } fn report_progress( @@ -481,7 +456,7 @@ impl DownloadPeer { ); } - fn send_download_status( + pub(crate) fn send_download_status( client_context: &Arc>, download: &Download, status: DownloadStatus, @@ -496,6 +471,56 @@ impl DownloadPeer { } } + /// Promote a complete partial and publish `Completed` as one operation with + /// respect to cancellation. Holding the store write lock across the rename + /// gives both possible races deterministic outcomes: cancellation marks the + /// row first and discards the partial, or completion lands first and a later + /// cancellation is refused. + fn finish_transfer( + client_context: &Arc>, + download: &Download, + part: PartFile, + ) -> Result { + loop { + let mut context = client_context + .write() + .map_err(|_| DownloadError::LockPoisoned)?; + let status = context + .get_download_by_token(download.token) + .map(|download| download.status.clone()) + .ok_or(DownloadError::TokenNotFound(download.token))?; + + if matches!(status, DownloadStatus::Cancelled) { + drop(context); + part.discard(); + return Err(DownloadError::Cancelled); + } + if matches!(status, DownloadStatus::Paused { .. }) { + drop(context); + if let Err(error) = + Self::wait_while_paused(client_context, download) + { + if matches!(error, DownloadError::Cancelled) { + part.discard(); + } + return Err(error); + } + continue; + } + + let final_path = part.finish()?; + if !context + .downloads + .update_status(download.token, DownloadStatus::Completed) + { + return Err(DownloadError::TokenNotFound(download.token)); + } + drop(context); + let _ = download.sender.send(DownloadStatus::Completed); + return Ok(final_path); + } + } + fn wait_while_paused( client_context: &Arc>, download: &Download, @@ -527,36 +552,54 @@ impl DownloadPeer { download: Option, stream: Option, ) -> Result<(Download, String), DownloadError> { - trace!( - "[download_peer:{}] download_file: download is present?: {:?}, stream is present?: {:?}, no_pierce: {}", - self.username, - download.is_some(), - stream.is_some(), - self.no_pierce - ); + let mut matched_download = download.clone(); + let result: Result<_, DownloadError> = (|| { + trace!( + "[download_peer:{}] download_file: download is present?: {:?}, stream is present?: {:?}, no_pierce: {}", + self.username, + download.is_some(), + stream.is_some(), + self.no_pierce + ); - let mut stream = match stream { - Some(s) => { - Self::apply_transfer_deadlines(&s)?; - s - } - None => self.establish_connection()?, - }; + let mut stream = match stream { + Some(s) => { + Self::apply_transfer_deadlines(&s)?; + s + } + None => self.establish_connection()?, + }; - trace!("[download_peer:{}] connected", self.username); + trace!("[download_peer:{}] connected", self.username); - self.perform_handshake(&mut stream)?; - trace!("[download_peer:{}] handshake completed", self.username); + self.perform_handshake(&mut stream)?; + trace!("[download_peer:{}] handshake completed", self.username); - let (download, final_path) = - self.read_download_stream(&mut stream, &client_context, download)?; + let (download, final_path) = self.read_download_stream( + &mut stream, + &client_context, + download, + &mut matched_download, + )?; - trace!( - "[download_peer:{}] download completed successfully: {} bytes, saved to: {}", - self.username, download.size, final_path - ); + trace!( + "[download_peer:{}] download completed successfully: {} bytes, saved to: {}", + self.username, download.size, final_path + ); - Ok((download, final_path)) + Ok((download, final_path)) + })(); + if let Err(error) = &result + && !matches!(error, DownloadError::Cancelled) + && let Some(download) = matched_download + { + Self::send_download_status( + &client_context, + &download, + DownloadStatus::Failed(Some(error.to_string())), + ); + } + result } } @@ -580,7 +623,7 @@ mod tests { fn download_into(dir: &std::path::Path, size: u64) -> Download { Download { username: "peer".to_string(), - filename: "song.mp3".to_string(), + filename: "file.bin".to_string(), token: 1, size, download_directory: dir.display().to_string(), @@ -596,13 +639,13 @@ mod tests { let dir = scratch_dir("discard"); let mut part = PartFile::open(&download_into(&dir, 8)).unwrap(); part.write(b"data").unwrap(); - let path = dir.join("song.mp3.part"); + let path = dir.join("file.bin.part"); assert!(path.exists()); part.discard(); assert!(!path.exists()); - assert!(!dir.join("song.mp3").exists()); + assert!(!dir.join("file.bin").exists()); let _ = std::fs::remove_dir_all(&dir); } @@ -625,7 +668,7 @@ mod tests { .write() .unwrap() .downloads - .resume_by_file("peer", "song.mp3") + .resume_by_file("peer", "file.bin") ); }) }; @@ -681,6 +724,131 @@ mod tests { let _ = std::fs::remove_dir_all(&dir); } + #[test] + fn a_pierced_short_transfer_reports_failure() { + let dir = scratch_dir("pierced-short-transfer"); + let (sender, receiver) = mpsc::channel(); + let mut download = download_into(&dir, 10); + download.sender = sender; + let context = Arc::new(RwLock::new(ClientContext::new())); + context.write().unwrap().add_download(download.clone()); + + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let client_stream = + std::net::TcpStream::connect(listener.local_addr().unwrap()) + .unwrap(); + let (mut uploader, _) = listener.accept().unwrap(); + uploader + .set_read_timeout(Some(std::time::Duration::from_secs(2))) + .unwrap(); + let peer = DownloadPeer::new( + download.username.clone(), + "127.0.0.1".into(), + 0, + 99, + false, + "me".into(), + ); + let worker_context = context.clone(); + let worker = std::thread::spawn(move || { + peer.download_file(worker_context, None, Some(client_stream)) + }); + + // Drain the framed PierceFirewall request before sending the raw token. + let mut frame_len = [0u8; 4]; + std::io::Read::read_exact(&mut uploader, &mut frame_len).unwrap(); + let mut frame = vec![0; u32::from_le_bytes(frame_len) as usize]; + std::io::Read::read_exact(&mut uploader, &mut frame).unwrap(); + + std::io::Write::write_all(&mut uploader, &download.token.to_le_bytes()) + .unwrap(); + + let mut offset = [0u8; 8]; + std::io::Read::read_exact(&mut uploader, &mut offset).unwrap(); + assert_eq!(u64::from_le_bytes(offset), 0); + std::io::Write::write_all(&mut uploader, b"0123").unwrap(); + std::net::TcpStream::shutdown(&uploader, std::net::Shutdown::Write) + .unwrap(); + drop(uploader); + + assert!(matches!( + worker.join().unwrap(), + Err(DownloadError::IncompleteDownload { + received: 4, + expected: 10 + }) + )); + assert!(matches!( + receiver.recv_timeout(std::time::Duration::from_secs(2)), + Ok(DownloadStatus::InProgress { .. }) + )); + assert!(matches!( + receiver.recv_timeout(std::time::Duration::from_secs(2)), + Ok(DownloadStatus::Failed(_)) + )); + assert!(matches!( + context + .read() + .unwrap() + .get_download_by_token(download.token) + .unwrap() + .status, + DownloadStatus::Failed(_) + )); + assert_eq!(std::fs::read(dir.join("file.bin.part")).unwrap(), b"0123"); + let _ = std::fs::remove_dir_all(dir); + } + + #[test] + fn a_pierced_token_cannot_claim_another_peers_download() { + let dir = scratch_dir("token-peer-mismatch"); + let mut download = download_into(&dir, 10); + download.username = "victim".into(); + let context = Arc::new(RwLock::new(ClientContext::new())); + context.write().unwrap().add_download(download.clone()); + + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let client_stream = + std::net::TcpStream::connect(listener.local_addr().unwrap()) + .unwrap(); + let (mut attacker, _) = listener.accept().unwrap(); + let peer = DownloadPeer::new( + "attacker".into(), + "127.0.0.1".into(), + 0, + 100, + false, + "me".into(), + ); + let worker_context = context.clone(); + let worker = std::thread::spawn(move || { + peer.download_file(worker_context, None, Some(client_stream)) + }); + + let mut frame_len = [0u8; 4]; + std::io::Read::read_exact(&mut attacker, &mut frame_len).unwrap(); + let mut frame = vec![0; u32::from_le_bytes(frame_len) as usize]; + std::io::Read::read_exact(&mut attacker, &mut frame).unwrap(); + std::io::Write::write_all(&mut attacker, &download.token.to_le_bytes()) + .unwrap(); + + assert!(matches!( + worker.join().unwrap(), + Err(DownloadError::TokenNotFound(1)) + )); + assert!(!dir.join("file.bin.part").exists()); + assert!(matches!( + context + .read() + .unwrap() + .get_download_by_token(download.token) + .unwrap() + .status, + DownloadStatus::Queued + )); + let _ = std::fs::remove_dir_all(dir); + } + #[test] fn start_transfer_refuses_a_cancelled_download() { let dir = scratch_dir("refuse"); @@ -698,7 +866,7 @@ mod tests { DownloadPeer::start_transfer(download, &mut stream, &context); assert!(matches!(outcome, Err(DownloadError::Cancelled))); - assert!(!dir.join("song.mp3.part").exists(), "no .part is opened"); + assert!(!dir.join("file.bin.part").exists(), "no .part is opened"); let _ = std::fs::remove_dir_all(&dir); } @@ -722,7 +890,7 @@ mod tests { #[test] fn a_partial_file_sets_the_resume_offset() { let dir = scratch_dir("resume"); - std::fs::write(dir.join("song.mp3.part"), b"0123").unwrap(); + std::fs::write(dir.join("file.bin.part"), b"0123").unwrap(); let mut part = PartFile::open(&download_into(&dir, 10)).unwrap(); assert_eq!(part.written, 4, "resume from what is already on disk"); @@ -732,17 +900,76 @@ mod tests { assert_eq!(std::fs::read(&final_path).unwrap(), b"0123456789"); assert!( - !dir.join("song.mp3.part").exists(), + !dir.join("file.bin.part").exists(), "the .part is renamed, not left behind" ); let _ = std::fs::remove_dir_all(dir); } + #[test] + fn an_unidentified_exact_size_partial_restarts_from_zero() { + // With no persisted transfer identity, length alone cannot distinguish + // a final write awaiting rename from unrelated same-sized content. + let dir = scratch_dir("complete"); + std::fs::write(dir.join("file.bin.part"), b"0123456789").unwrap(); + + let mut part = PartFile::open(&download_into(&dir, 10)).unwrap(); + assert_eq!(part.written, 0); + assert!(!part.is_complete()); + + part.write(b"abcdefghij").unwrap(); + let final_path = part.finish().unwrap(); + assert_eq!(std::fs::read(final_path).unwrap(), b"abcdefghij"); + assert!(!dir.join("file.bin.part").exists()); + let _ = std::fs::remove_dir_all(dir); + } + + #[test] + fn cancellation_while_a_complete_transfer_is_paused_discards_its_part() { + let dir = scratch_dir("cancel-paused-finish"); + let mut download = download_into(&dir, 10); + let mut part = PartFile::open(&download).unwrap(); + part.write(b"0123456789").unwrap(); + download.status = DownloadStatus::Paused { + bytes_downloaded: 10, + total_bytes: 10, + }; + let context = Arc::new(RwLock::new(ClientContext::new())); + context.write().unwrap().add_download(download.clone()); + let worker_context = context.clone(); + let worker_download = download.clone(); + + let worker = std::thread::spawn(move || { + DownloadPeer::finish_transfer( + &worker_context, + &worker_download, + part, + ) + }); + + std::thread::sleep(std::time::Duration::from_millis(50)); + assert!( + context + .write() + .unwrap() + .downloads + .cancel_by_file(&download.username, &download.filename) + .is_some() + ); + + assert!(matches!( + worker.join().unwrap(), + Err(DownloadError::Cancelled) + )); + assert!(!dir.join("file.bin.part").exists()); + let _ = std::fs::remove_dir_all(dir); + } + #[test] fn an_oversized_partial_file_restarts_from_zero() { // Left over from a different transfer — not a valid resume point. let dir = scratch_dir("stale"); - std::fs::write(dir.join("song.mp3.part"), vec![9u8; 40]).unwrap(); + std::fs::write(dir.join("file.bin.part"), vec![9u8; 40]).unwrap(); let mut part = PartFile::open(&download_into(&dir, 10)).unwrap(); assert_eq!(part.written, 0); @@ -770,11 +997,11 @@ mod tests { }) )); assert_eq!( - std::fs::read(dir.join("song.mp3.part")).unwrap(), + std::fs::read(dir.join("file.bin.part")).unwrap(), b"01234", "the partial stays on disk for the next attempt" ); - assert!(!dir.join("song.mp3").exists()); + assert!(!dir.join("file.bin").exists()); let _ = std::fs::remove_dir_all(dir); } diff --git a/soulseek-rs-lib/src/peer/listen.rs b/soulseek-rs-lib/src/peer/listen.rs index 20e747d..eff0b7b 100644 --- a/soulseek-rs-lib/src/peer/listen.rs +++ b/soulseek-rs-lib/src/peer/listen.rs @@ -12,7 +12,7 @@ use crate::peer::{ConnectionType, DownloadError, 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::{debug, error, info, trace}; /// How long to wait before accepting again after a failure. /// @@ -130,20 +130,10 @@ fn wait_for_download_registration( 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 - | DownloadStatus::Cancelled - ) - }) - .cloned(), + let download = match client_context.write_safe() { + Ok(mut context) => { + context.downloads.claim_for_peer(token, username) + } Err(e) => { error!("[listener] client context lock: {}", e); return None; @@ -224,8 +214,6 @@ fn handle_file_connection( Instant::now() + TOKEN_REGISTRATION_TIMEOUT, ) }); - let failure_token = download.as_ref().map(|d| d.token); - let download_peer = DownloadPeer::new( peer.username.clone(), peer.host.clone(), @@ -241,16 +229,6 @@ fn handle_file_connection( Some(stream), ) { Ok((download, filename)) => { - let _ = download.sender.send(DownloadStatus::Completed); - match context.client_context.write_safe() { - Ok(mut ctx) => ctx.update_download_with_status( - download.token, - DownloadStatus::Completed, - ), - Err(e) => { - error!("[listener] handle_file_connection write: {}", e); - } - } info!( "Successfully downloaded {} bytes to {}", download.size, filename @@ -262,22 +240,6 @@ fn handle_file_connection( "Failed to download file from {}:{} (token: {}) - Error: {}", peer.host, peer.port, token, e ); - // A failed incoming transfer (e.g. a truncated/incomplete download) - // must not leave the download stuck as Queued/InProgress forever. - if let Some(failure_token) = failure_token { - match context.client_context.write_safe() { - Ok(mut ctx) => ctx.update_download_with_status( - failure_token, - DownloadStatus::Failed(Some(e.to_string())), - ), - Err(e) => { - error!( - "[listener] handle_file_connection fail write: {}", - e - ); - } - } - } } } } @@ -491,7 +453,7 @@ impl Listen { #[cfg(test)] mod tests { use super::*; - use crate::types::DownloadMetadata; + use crate::types::{DownloadMetadata, DownloadStatus}; use std::io::Write; use std::net::TcpListener; use std::sync::mpsc; diff --git a/soulseek-rs-lib/src/peer/upload_peer.rs b/soulseek-rs-lib/src/peer/upload_peer.rs index 57138d4..36e8978 100644 --- a/soulseek-rs-lib/src/peer/upload_peer.rs +++ b/soulseek-rs-lib/src/peer/upload_peer.rs @@ -5,7 +5,7 @@ use std::io::{self, Read, Seek, SeekFrom, Write}; use std::net::{TcpStream, ToSocketAddrs}; use std::path::Path; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::time::Duration; +use std::time::{Duration, Instant}; use crate::message::server::MessageFactory; use crate::peer::ConnectionType; @@ -73,19 +73,21 @@ pub fn serve_file( let mut offset = [0u8; 8]; stream.read_exact(&mut offset)?; let offset = u64::from_le_bytes(offset); - if offset > 0 { - file.seek(SeekFrom::Start(offset))?; - bytes_sent.fetch_add(offset, Ordering::Relaxed); + let file_size = file.metadata()?.len(); + if offset > file_size { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("download offset {offset} exceeds file size {file_size}"), + )); } + file.seek(SeekFrom::Start(offset))?; + bytes_sent.fetch_add(offset, Ordering::Relaxed); let mut buffer = vec![0u8; 64 * 1024]; let mut streamed = 0u64; loop { if cancel.load(Ordering::Relaxed) { - return Err(io::Error::new( - io::ErrorKind::Interrupted, - "upload cancelled", - )); + return Err(io::ErrorKind::Interrupted.into()); } let read = file.read(&mut buffer)?; if read == 0 { @@ -97,11 +99,41 @@ pub fn serve_file( } stream.flush()?; - // Half-close instead of lingering: the FIN tells the downloader it has the - // whole file, and TCP still delivers everything written before it. Sleeping - // here held an upload slot for an extra half-second per transfer, which with - // a slot cap is time no other peer can use. - stream.shutdown(std::net::Shutdown::Write).ok(); + // Soulseek assigns closing the file-transfer connection to the downloader + // after it has read exactly the advertised number of bytes. Poll while + // waiting for that EOF so cancelling a finished-but-unacknowledged upload + // remains prompt rather than waiting for the full socket timeout. + stream.set_read_timeout(Some(Duration::from_millis(200)))?; + let close_deadline = Instant::now() + Duration::from_secs(30); + let mut unexpected = [0u8; 1]; + loop { + if cancel.load(Ordering::Relaxed) { + return Err(io::ErrorKind::Interrupted.into()); + } + match stream.read(&mut unexpected) { + Ok(0) => break, + Ok(_) => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "unexpected data after download offset", + )); + } + Err(error) + if matches!( + error.kind(), + io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut + ) => + { + if Instant::now() >= close_deadline { + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "downloader did not close the file connection", + )); + } + } + Err(error) => return Err(error), + } + } trace!("[upload] served {} to {}:{}", path.display(), host, port); Ok(streamed) } @@ -111,8 +143,9 @@ mod tests { use super::serve_file; use std::io::{Read, Write}; use std::net::{TcpListener, TcpStream}; - use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + use std::sync::{Arc, mpsc}; + use std::time::Duration; /// Play the downloader: accept the F connection, drain the PeerInit frame /// and the raw token, then reply with the 8-byte START_DOWNLOAD offset. @@ -143,77 +176,152 @@ mod tests { path } + fn spawn_uploader( + listener: &TcpListener, + path: std::path::PathBuf, + token: u32, + cancel: Arc, + ) -> ( + Arc, + std::thread::JoinHandle>, + ) { + let port = u32::from(listener.local_addr().unwrap().port()); + let bytes_sent = Arc::new(AtomicU64::new(0)); + let sent_counter = bytes_sent.clone(); + let uploader = std::thread::spawn(move || { + serve_file( + "127.0.0.1", + port, + "me", + token, + &path, + &sent_counter, + &cancel, + ) + }); + (bytes_sent, uploader) + } + #[test] fn serve_file_streams_the_file_over_an_f_connection() { let content: Vec = (0..4096u32).map(|i| (i % 251) as u8).collect(); let path = scratch_file("stream", &content); let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let port = u32::from(listener.local_addr().unwrap().port()); - - let bytes_sent = Arc::new(AtomicU64::new(0)); - let sent_counter = bytes_sent.clone(); let dir = path.parent().unwrap().to_path_buf(); - let uploader = std::thread::spawn(move || { - serve_file( - "127.0.0.1", - port, - "me", - 777, - &path, - &sent_counter, - &AtomicBool::new(false), - ) - }); + let (bytes_sent, uploader) = + spawn_uploader(&listener, path, 777, Arc::default()); let (mut stream, token) = accept_and_send_offset(&listener, 0); assert_eq!(token, 777); - let mut received = Vec::new(); - stream.read_to_end(&mut received).unwrap(); + let mut received = vec![0; content.len()]; + stream.read_exact(&mut received).unwrap(); assert_eq!(received, content); - let _ = uploader.join(); + assert!( + !uploader.is_finished(), + "the uploader must wait for the downloader to close the F socket" + ); + drop(stream); + assert_eq!(uploader.join().unwrap().unwrap(), 4096); assert_eq!(bytes_sent.load(Ordering::Relaxed), 4096); let _ = std::fs::remove_dir_all(dir); } #[test] - fn serve_file_resumes_from_a_non_zero_offset() { + fn serve_file_accepts_valid_resume_offsets() { let content: Vec = (0..4096u32).map(|i| (i % 251) as u8).collect(); - let path = scratch_file("resume", &content); + for offset in [1000, content.len()] { + let path = + scratch_file(&format!("valid-offset-{offset}"), &content); + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let dir = path.parent().unwrap().to_path_buf(); + let (bytes_sent, uploader) = + spawn_uploader(&listener, path, 779, Arc::default()); + + let (mut stream, _) = + accept_and_send_offset(&listener, offset as u64); + let mut received = vec![0; content.len() - offset]; + stream.read_exact(&mut received).unwrap(); + assert_eq!(received, content[offset..]); + drop(stream); + + assert_eq!( + uploader.join().unwrap().unwrap(), + (content.len() - offset) as u64 + ); + assert_eq!( + bytes_sent.load(Ordering::Relaxed), + content.len() as u64 + ); + let _ = std::fs::remove_dir_all(dir); + } + } + + #[test] + fn serve_file_rejects_an_offset_past_end_of_file() { + let content: Vec = (0..4096u32).map(|i| (i % 251) as u8).collect(); + let path = scratch_file("invalid-offset", &content); let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let port = u32::from(listener.local_addr().unwrap().port()); - - let bytes_sent = Arc::new(AtomicU64::new(0)); - let sent_counter = bytes_sent.clone(); let dir = path.parent().unwrap().to_path_buf(); + let (bytes_sent, uploader) = + spawn_uploader(&listener, path, 780, Arc::default()); + + let (mut stream, _) = accept_and_send_offset(&listener, 4097); + let mut received = Vec::new(); + stream.read_to_end(&mut received).unwrap(); + assert!(received.is_empty()); + + let error = uploader.join().unwrap().unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::InvalidData); + assert_eq!( + bytes_sent.load(Ordering::Relaxed), + 0, + "an invalid prefix must not be counted as uploaded" + ); + let _ = std::fs::remove_dir_all(dir); + } + + #[test] + fn serve_file_cancels_while_waiting_for_the_downloader_to_close() { + let content: Vec = (0..4096u32).map(|i| (i % 251) as u8).collect(); + let path = scratch_file("cancel-close-wait", &content); + let dir = path.parent().unwrap().to_path_buf(); + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let port = u32::from(listener.local_addr().unwrap().port()); + let cancel = Arc::new(AtomicBool::new(false)); + let cancel_flag = cancel.clone(); + let (result_tx, result_rx) = mpsc::channel(); + let uploader = std::thread::spawn(move || { - serve_file( + let result = serve_file( "127.0.0.1", port, "me", - 779, + 782, &path, - &sent_counter, - &AtomicBool::new(false), - ) + &AtomicU64::new(0), + &cancel_flag, + ); + let _ = result_tx.send(result); }); - // The downloader already holds the first 1000 bytes. - let (mut stream, _) = accept_and_send_offset(&listener, 1000); + let (mut stream, _) = accept_and_send_offset(&listener, 0); + let mut received = vec![0; content.len()]; + stream.read_exact(&mut received).unwrap(); + assert_eq!(received, content); - let mut received = Vec::new(); - stream.read_to_end(&mut received).unwrap(); - assert_eq!(received, content[1000..], "only the missing tail is sent"); + cancel.store(true, Ordering::Relaxed); + let error = result_rx + .recv_timeout(Duration::from_secs(2)) + .expect("cancellation must wake the close wait") + .unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::Interrupted); + drop(stream); let _ = uploader.join(); - assert_eq!( - bytes_sent.load(Ordering::Relaxed), - 4096, - "progress counts the resumed prefix" - ); let _ = std::fs::remove_dir_all(dir); } @@ -223,23 +331,10 @@ mod tests { let dir = path.parent().unwrap().to_path_buf(); let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let port = u32::from(listener.local_addr().unwrap().port()); - // Cancelled before the copy loop starts: the stream must abort with // Interrupted instead of serving the whole file. let cancel = Arc::new(AtomicBool::new(true)); - let cancel_flag = cancel; - let uploader = std::thread::spawn(move || { - serve_file( - "127.0.0.1", - port, - "me", - 778, - &path, - &AtomicU64::new(0), - &cancel_flag, - ) - }); + let (_, uploader) = spawn_uploader(&listener, path, 778, cancel); let (mut stream, _) = accept_and_send_offset(&listener, 0); diff --git a/soulseek-rs-lib/tests/e2e.rs b/soulseek-rs-lib/tests/e2e.rs index f4e7725..4f6149c 100644 --- a/soulseek-rs-lib/tests/e2e.rs +++ b/soulseek-rs-lib/tests/e2e.rs @@ -1182,7 +1182,7 @@ fn an_interrupted_download_resumes_from_its_partial_file() { client.connect().expect("connect"); assert!(client.login().expect("login")); - let filename = "resumed_song.mp3"; + let filename = "resumed_file.bin"; let content: Vec = (0..2000u32).map(|i| (i % 251) as u8).collect(); let size = content.len() as u64; let already_have = 800usize; diff --git a/soulseek-rs-lib/tests/resume_e2e.rs b/soulseek-rs-lib/tests/resume_e2e.rs new file mode 100644 index 0000000..2f5cf55 --- /dev/null +++ b/soulseek-rs-lib/tests/resume_e2e.rs @@ -0,0 +1,408 @@ +//! Server-independent end-to-end coverage for downloads that resume. +//! +//! The subject is a real `Client` with its real peer listener and actor stack. +//! A loopback socket stands in for a Soulseek peer and speaks the public wire +//! format. The configured server is deliberately idle: peer transfers do not +//! require a login once the control connection is registered. + +use std::fs; +use std::io::{Read, Write}; +use std::net::{Shutdown, TcpListener, TcpStream}; +use std::path::{Path, PathBuf}; +use std::sync::mpsc::Receiver; +use std::time::{Duration, Instant}; + +use soulseek_rs::message::Message; +use soulseek_rs::{ + Client, ClientSettings, ClientVersion, DownloadStatus, PeerAddress, +}; + +const IO_TIMEOUT: Duration = Duration::from_secs(10); + +struct ScratchDir(PathBuf); + +impl ScratchDir { + fn new(label: &str) -> Self { + let path = std::env::temp_dir().join(format!( + "soulseek-resume-e2e-{}-{label}", + std::process::id() + )); + let _ = fs::remove_dir_all(&path); + fs::create_dir_all(&path).expect("create scratch directory"); + Self(path) + } + + fn path(&self) -> &Path { + &self.0 + } +} + +impl Drop for ScratchDir { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } +} + +struct Subject { + client: Client, + peer_addr: String, + // Keeping a server listener alive is enough for ServerActor's TCP connect. + // No login or server messages are needed by the peer-transfer path. + _idle_server: TcpListener, +} + +impl Subject { + fn new(username: &str) -> Self { + let idle_server = + TcpListener::bind("127.0.0.1:0").expect("bind idle server"); + let server_addr = + idle_server.local_addr().expect("idle server address"); + let mut client = Client::with_settings(ClientSettings { + username: username.to_string(), + password: "unused".to_string(), + server_address: PeerAddress::new( + "127.0.0.1".to_string(), + server_addr.port(), + ), + enable_listen: true, + listen_port: 0, + shared_directories: Vec::new(), + version: ClientVersion::default(), + }); + client.connect().expect("connect subject"); + let port = client.listen_port().expect("subject listener port"); + Self { + client, + peer_addr: format!("127.0.0.1:{port}"), + _idle_server: idle_server, + } + } +} + +struct ControlPeer { + username: String, + subject_addr: String, + stream: TcpStream, +} + +impl ControlPeer { + fn register(subject: &Subject, username: &str) -> Self { + let mut stream = TcpStream::connect(&subject.peer_addr) + .expect("connect peer control socket"); + stream + .set_read_timeout(Some(IO_TIMEOUT)) + .expect("control read timeout"); + stream + .set_write_timeout(Some(IO_TIMEOUT)) + .expect("control write timeout"); + + // Code 15/16 is a positive registration fence. Receiving the response + // proves the inbound P connection has reached the peer actor and can + // receive QueueUpload messages from `Client::download`. + stream + .write_all(&peer_init_bytes(username, "P", 0)) + .expect("send P PeerInit"); + stream + .write_all(&Message::new().write_int32(15).get_buffer()) + .expect("send UserInfoRequest registration fence"); + stream.flush().expect("flush registration fence"); + let _ = expect_code(&mut stream, 16); + + Self { + username: username.to_string(), + subject_addr: subject.peer_addr.clone(), + stream, + } + } + + fn accept_queue_and_offer(&mut self, filename: &str, size: u64) -> u32 { + let mut queue = expect_code(&mut self.stream, 43); + queue.set_pointer(8); + assert_eq!(queue.read_string(), filename, "queued remote filename"); + + let token = 700_000; + let mut request = Message::new(); + request + .write_int32(40) + .write_int32(1) // upload from the peer to the subject + .write_int32(token) + .write_string(filename) + .write_int64(size); + self.stream + .write_all(&request.get_buffer()) + .expect("send TransferRequest"); + self.stream.flush().expect("flush TransferRequest"); + + let mut response = expect_code(&mut self.stream, 41); + response.set_pointer(8); + assert_eq!(response.read_int32(), token, "response transfer token"); + assert_eq!(response.read_int8(), 1, "transfer should be allowed"); + token + } + + fn open_file_transfer(&self, token: u32) -> (TcpStream, u64) { + let mut stream = TcpStream::connect(&self.subject_addr) + .expect("connect peer file socket"); + stream + .set_read_timeout(Some(IO_TIMEOUT)) + .expect("file read timeout"); + stream + .set_write_timeout(Some(IO_TIMEOUT)) + .expect("file write timeout"); + + let mut init = peer_init_bytes(&self.username, "F", token); + init.extend_from_slice(&token.to_le_bytes()); + stream.write_all(&init).expect("send F PeerInit and token"); + stream.flush().expect("flush F PeerInit and token"); + + let mut offset = [0u8; 8]; + stream + .read_exact(&mut offset) + .expect("read START_DOWNLOAD offset"); + (stream, u64::from_le_bytes(offset)) + } +} + +fn peer_init_bytes( + username: &str, + connection_type: &str, + token: u32, +) -> Vec { + Message::new() + .write_int8(1) + .write_string(username) + .write_string(connection_type) + .write_int32(token) + .get_buffer() +} + +fn read_framed(stream: &mut TcpStream) -> std::io::Result { + let mut len = [0u8; 4]; + stream.read_exact(&mut len)?; + let mut payload = vec![0u8; u32::from_le_bytes(len) as usize]; + stream.read_exact(&mut payload)?; + let mut bytes = len.to_vec(); + bytes.extend_from_slice(&payload); + Ok(Message::new_with_data(bytes)) +} + +fn expect_code(stream: &mut TcpStream, expected: u32) -> Message { + let deadline = Instant::now() + IO_TIMEOUT; + loop { + let message = read_framed(stream).unwrap_or_else(|error| { + panic!("read peer code {expected}: {error}") + }); + if message.get_message_code() == expected { + return message; + } + assert!( + Instant::now() < deadline, + "timed out waiting for peer code {expected}" + ); + } +} + +fn start_download( + subject: &Subject, + control: &mut ControlPeer, + filename: &str, + content_len: usize, + destination: &Path, +) -> (Receiver, TcpStream, u64) { + let (_download, statuses) = subject + .client + .download( + filename.to_string(), + control.username.clone(), + content_len as u64, + destination.display().to_string(), + ) + .expect("queue download"); + let token = control.accept_queue_and_offer(filename, content_len as u64); + let (file_stream, offset) = control.open_file_transfer(token); + (statuses, file_stream, offset) +} + +fn wait_for_status( + statuses: &Receiver, + description: &str, + wanted: impl Fn(&DownloadStatus) -> bool, +) -> DownloadStatus { + let deadline = Instant::now() + IO_TIMEOUT; + loop { + if let Ok(status) = statuses.recv_timeout(Duration::from_millis(100)) + && wanted(&status) + { + return status; + } + assert!( + Instant::now() < deadline, + "status channel did not report {description}" + ); + } +} + +fn assert_stored_status( + subject: &Subject, + filename: &str, + wanted: impl Fn(&DownloadStatus) -> bool, +) { + let status = subject + .client + .get_all_downloads() + .into_iter() + .find(|download| download.filename == filename) + .unwrap_or_else(|| panic!("download {filename} is missing from store")) + .status; + assert!(wanted(&status), "unexpected stored status: {status:?}"); +} + +fn expect_downloader_close(stream: &mut TcpStream) { + let mut unexpected = [0u8; 1]; + assert_eq!( + stream + .read(&mut unexpected) + .expect("wait for downloader to close the F socket"), + 0, + "the downloader sends no bytes after its offset" + ); +} + +fn binary_payload(len: usize) -> Vec { + (0..len).map(|index| index as u8).collect() +} + +#[test] +fn an_actual_short_transfer_resumes_opaque_bytes() { + let subject = Subject::new("resume_binary_subject"); + let mut control = ControlPeer::register(&subject, "resume_binary_peer"); + let destination = ScratchDir::new("binary-tracer"); + let content = binary_payload(3 * 64 * 1024 + 257); + let filename = "opaque.bin"; + let interrupt_at = 64 * 1024 + 1; + let part_path = destination.path().join(format!("{filename}.part")); + let final_path = destination.path().join(filename); + + assert!(content.contains(&0)); + assert!(content.contains(&u8::MAX)); + assert!(std::str::from_utf8(&content).is_err()); + + let (first_statuses, mut first, first_offset) = start_download( + &subject, + &mut control, + filename, + content.len(), + destination.path(), + ); + assert_eq!(first_offset, 0); + first + .write_all(&content[..interrupt_at]) + .expect("write short first transfer"); + first.flush().expect("flush short first transfer"); + first + .shutdown(Shutdown::Write) + .expect("interrupt first transfer"); + drop(first); + + let _ = wait_for_status(&first_statuses, "Failed", |status| { + matches!(status, DownloadStatus::Failed(_)) + }); + assert_stored_status(&subject, filename, |status| { + matches!(status, DownloadStatus::Failed(_)) + }); + assert!(!final_path.exists()); + assert_eq!( + fs::read(&part_path).expect("read retained partial file"), + content[..interrupt_at] + ); + + let (retry_statuses, mut retry, retry_offset) = start_download( + &subject, + &mut control, + filename, + content.len(), + destination.path(), + ); + assert_eq!(retry_offset, interrupt_at as u64); + retry + .write_all(&content[interrupt_at..]) + .expect("write retry tail"); + retry.flush().expect("flush retry tail"); + expect_downloader_close(&mut retry); + drop(retry); + + let _ = wait_for_status(&retry_statuses, "Completed", |status| { + matches!(status, DownloadStatus::Completed) + }); + assert_stored_status(&subject, filename, |status| { + matches!(status, DownloadStatus::Completed) + }); + assert_eq!(fs::read(&final_path).unwrap(), content); + assert!(!part_path.exists()); +} + +#[test] +fn a_live_download_pauses_and_resumes_on_the_same_connection() { + let subject = Subject::new("resume_pause_subject"); + let mut control = ControlPeer::register(&subject, "resume_pause_peer"); + let destination = ScratchDir::new("pause-resume"); + let filename = "paused.bin"; + let content = binary_payload(5 * 64 * 1024 + 257); + let part_path = destination.path().join(format!("{filename}.part")); + + let (statuses, mut peer, offset) = start_download( + &subject, + &mut control, + filename, + content.len(), + destination.path(), + ); + assert_eq!(offset, 0); + + let first_chunk = 2 * 64 * 1024; + peer.write_all(&content[..first_chunk]) + .expect("send first chunk"); + peer.flush().expect("flush first chunk"); + let _ = wait_for_status(&statuses, "initial progress", |status| { + matches!( + status, + DownloadStatus::InProgress { + bytes_downloaded, + .. + } if *bytes_downloaded >= first_chunk as u64 + ) + }); + + assert!(subject.client.pause_download(&control.username, filename)); + let _ = wait_for_status(&statuses, "Paused", |status| { + matches!(status, DownloadStatus::Paused { .. }) + }); + + // Give a read already in flight one timeout cycle to observe the pause. + std::thread::sleep(Duration::from_millis(1_200)); + let stable_len = fs::metadata(&part_path).expect("paused partial").len(); + let queued_while_paused = 8 * 1024; + peer.write_all(&content[first_chunk..first_chunk + queued_while_paused]) + .expect("queue bytes while paused"); + peer.flush().expect("flush bytes while paused"); + std::thread::sleep(Duration::from_millis(300)); + assert_eq!( + fs::metadata(&part_path).expect("paused partial").len(), + stable_len, + "a paused transfer must stop consuming its socket" + ); + + assert!(subject.client.resume_download(&control.username, filename)); + peer.write_all(&content[first_chunk + queued_while_paused..]) + .expect("send tail after resume"); + peer.flush().expect("flush resumed tail"); + expect_downloader_close(&mut peer); + + let _ = wait_for_status(&statuses, "Completed", |status| { + matches!(status, DownloadStatus::Completed) + }); + assert_eq!( + fs::read(destination.path().join(filename)).unwrap(), + content + ); +}