diff --git a/soulseek-rs-lib/src/client/mod.rs b/soulseek-rs-lib/src/client/mod.rs index cbfb62b..dd71552 100644 --- a/soulseek-rs-lib/src/client/mod.rs +++ b/soulseek-rs-lib/src/client/mod.rs @@ -59,6 +59,18 @@ fn next_upload_token() -> u32 { struct UploadJob { downloader: String, real_path: std::path::PathBuf, + virtual_path: String, + size: u64, +} + +/// Live bookkeeping for an upload being served (or recently finished). +struct ActiveUpload { + username: String, + filename: String, + size: u64, + bytes_sent: Arc, + cancel: Arc, + status: crate::types::UploadStatus, } /// Build a `FileSearchResponse` for `query` against `shares`, or `None` if @@ -210,6 +222,7 @@ pub struct ClientContext { pending_peer_messages: HashMap>, /// Uploads we have offered, keyed by our transfer token. uploads: HashMap, + active_uploads: HashMap, /// Upload tokens waiting for the downloader's address to be resolved. pending_serves: HashMap>, /// Shared-file listings received from peers we browsed. @@ -507,6 +520,7 @@ impl ClientContext { peer_addresses: HashMap::new(), pending_peer_messages: HashMap::new(), uploads: HashMap::new(), + active_uploads: HashMap::new(), pending_serves: HashMap::new(), browse_results: HashMap::new(), room_list: Vec::new(), @@ -657,6 +671,54 @@ impl Client { .unwrap_or_default() } + /// Snapshot of the uploads served this session (active and finished), + /// most recent last. + #[must_use] + pub fn uploads(&self) -> Vec { + self.context.read_safe().map_or_else( + |_| Vec::new(), + |ctx| { + let mut tokens: Vec<&u32> = ctx.active_uploads.keys().collect(); + tokens.sort_unstable(); + tokens + .into_iter() + .map(|token| { + let upload = &ctx.active_uploads[token]; + crate::types::UploadInfo { + username: upload.username.clone(), + filename: upload.filename.clone(), + size: upload.size, + bytes_sent: upload + .bytes_sent + .load(std::sync::atomic::Ordering::Relaxed), + status: upload.status.clone(), + } + }) + .collect() + }, + ) + } + + /// Ask an in-progress upload to `username` of `filename` to stop. + /// Returns whether a matching in-progress upload was found. + pub fn cancel_upload(&self, username: &str, filename: &str) -> bool { + self.context.read_safe().is_ok_and(|ctx| { + let mut found = false; + for upload in ctx.active_uploads.values() { + if upload.username == username + && upload.filename == filename + && upload.status == crate::types::UploadStatus::InProgress + { + upload + .cancel + .store(true, std::sync::atomic::Ordering::Relaxed); + found = true; + } + } + found + }) + } + /// Replace the shared directories at runtime: rescan into a fresh /// index (served to peers from then on) and re-announce the new /// folder/file counts to the server. diff --git a/soulseek-rs-lib/src/client/operations.rs b/soulseek-rs-lib/src/client/operations.rs index 6919e8c..38d820e 100644 --- a/soulseek-rs-lib/src/client/operations.rs +++ b/soulseek-rs-lib/src/client/operations.rs @@ -652,6 +652,8 @@ impl Client { UploadJob { downloader: downloader.clone(), real_path, + virtual_path: filename.clone(), + size, }, ); (ctx.peer_registry.clone(), size) diff --git a/soulseek-rs-lib/src/client/uploads.rs b/soulseek-rs-lib/src/client/uploads.rs index 5aa65f2..e395318 100644 --- a/soulseek-rs-lib/src/client/uploads.rs +++ b/soulseek-rs-lib/src/client/uploads.rs @@ -1,7 +1,9 @@ use super::{ - Arc, Client, ClientContext, DownloadStatus, RwLock, RwLockExt, - collect_failed_tokens, error, thread, + ActiveUpload, Arc, Client, ClientContext, DownloadStatus, RwLock, + RwLockExt, collect_failed_tokens, error, thread, }; +use crate::types::UploadStatus; +use std::sync::atomic::{AtomicBool, AtomicU64}; impl Client { /// Consume the upload job for `token` and stream the file to `host:port` @@ -19,14 +21,47 @@ impl Client { let Some(job) = ctx.uploads.remove(&token) else { return; }; + let bytes_sent = Arc::new(AtomicU64::new(0)); + let cancel = Arc::new(AtomicBool::new(false)); + ctx.active_uploads.insert( + token, + ActiveUpload { + username: job.downloader.clone(), + filename: job.virtual_path.clone(), + size: job.size, + bytes_sent: bytes_sent.clone(), + cancel: cancel.clone(), + status: UploadStatus::InProgress, + }, + ); drop(ctx); let own = own_username.to_string(); let real_path = job.real_path; + let context = client_context.clone(); thread::spawn(move || { - if let Err(e) = crate::peer::upload_peer::serve_file( - &host, port, &own, token, &real_path, - ) { - error!("[client] serve {}: {}", real_path.display(), e); + let result = crate::peer::upload_peer::serve_file( + &host, + port, + &own, + token, + &real_path, + &bytes_sent, + &cancel, + ); + let status = match &result { + Ok(()) => UploadStatus::Completed, + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => { + UploadStatus::Cancelled + } + Err(e) => { + error!("[client] serve {}: {}", real_path.display(), e); + UploadStatus::Failed(e.to_string()) + } + }; + if let Ok(mut ctx) = context.write_safe() + && let Some(upload) = ctx.active_uploads.get_mut(&token) + { + upload.status = status; } }); } diff --git a/soulseek-rs-lib/src/peer/upload_peer.rs b/soulseek-rs-lib/src/peer/upload_peer.rs index b0aecd3..30786a0 100644 --- a/soulseek-rs-lib/src/peer/upload_peer.rs +++ b/soulseek-rs-lib/src/peer/upload_peer.rs @@ -4,6 +4,7 @@ use std::fs::File; use std::io::{self, Read, Write}; use std::net::{TcpStream, ToSocketAddrs}; use std::path::Path; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::Duration; use crate::message::server::MessageFactory; @@ -17,6 +18,9 @@ use crate::trace; /// download is matched by token), then the downloader sends an 8-byte /// START_DOWNLOAD offset before we stream the file. /// +/// `bytes_sent` is updated as the transfer progresses, and setting `cancel` +/// aborts the stream with an [`io::ErrorKind::Interrupted`] error. +/// /// # Errors /// Returns any I/O error opening the file or talking to the peer. pub fn serve_file( @@ -25,6 +29,8 @@ pub fn serve_file( own_username: &str, token: u32, path: &Path, + bytes_sent: &AtomicU64, + cancel: &AtomicBool, ) -> io::Result<()> { let mut file = File::open(path)?; @@ -53,7 +59,21 @@ pub fn serve_file( let mut offset = [0u8; 8]; stream.read_exact(&mut offset)?; - io::copy(&mut file, &mut stream)?; + let mut buffer = [0u8; 64 * 1024]; + loop { + if cancel.load(Ordering::Relaxed) { + return Err(io::Error::new( + io::ErrorKind::Interrupted, + "upload cancelled", + )); + } + let read = file.read(&mut buffer)?; + if read == 0 { + break; + } + stream.write_all(&buffer[..read])?; + bytes_sent.fetch_add(read as u64, Ordering::Relaxed); + } stream.flush()?; // Linger so the downloader drains everything before the socket closes. @@ -67,6 +87,8 @@ mod tests { use super::serve_file; use std::io::{Read, Write}; use std::net::TcpListener; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; #[test] fn serve_file_streams_the_file_over_an_f_connection() { @@ -80,8 +102,18 @@ mod tests { 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 uploader = std::thread::spawn(move || { - serve_file("127.0.0.1", port, "me", 777, &path) + serve_file( + "127.0.0.1", + port, + "me", + 777, + &path, + &sent_counter, + &AtomicBool::new(false), + ) }); // Downloader side: accept, read PeerInit(F) frame + raw token, send the @@ -104,6 +136,53 @@ mod tests { assert_eq!(received, content); let _ = uploader.join(); + assert_eq!(bytes_sent.load(Ordering::Relaxed), 4096); + let _ = std::fs::remove_dir_all(dir); + } + + #[test] + fn serve_file_stops_when_cancelled() { + let dir = std::env::temp_dir() + .join(format!("soulseek-upload-cancel-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("big.bin"); + std::fs::write(&path, vec![7u8; 1024 * 1024]).unwrap(); + + 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.clone(); + let uploader = std::thread::spawn(move || { + serve_file( + "127.0.0.1", + port, + "me", + 778, + &path, + &AtomicU64::new(0), + &cancel_flag, + ) + }); + + let (mut stream, _) = listener.accept().unwrap(); + let mut len_buf = [0u8; 4]; + stream.read_exact(&mut len_buf).unwrap(); + let mut payload = vec![0u8; u32::from_le_bytes(len_buf) as usize]; + stream.read_exact(&mut payload).unwrap(); + let mut token = [0u8; 4]; + stream.read_exact(&mut token).unwrap(); + stream.write_all(&[0u8; 8]).unwrap(); + stream.flush().unwrap(); + + let mut received = Vec::new(); + let _ = stream.read_to_end(&mut received); + assert!(received.is_empty(), "no file bytes after cancellation"); + + let err = uploader.join().unwrap().unwrap_err(); + assert_eq!(err.kind(), std::io::ErrorKind::Interrupted); let _ = std::fs::remove_dir_all(dir); } } diff --git a/soulseek-rs-lib/src/types.rs b/soulseek-rs-lib/src/types.rs index a22486f..6c50854 100644 --- a/soulseek-rs-lib/src/types.rs +++ b/soulseek-rs-lib/src/types.rs @@ -251,3 +251,24 @@ mod tests { assert_eq!(transfer.size, 0); } } + +/// Lifecycle of a file we are serving to a peer. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum UploadStatus { + InProgress, + Completed, + Cancelled, + Failed(String), +} + +/// A snapshot of one upload for display purposes. +#[derive(Debug, Clone)] +pub struct UploadInfo { + /// The peer receiving the file. + pub username: String, + /// The peer-facing virtual path being served. + pub filename: String, + pub size: u64, + pub bytes_sent: u64, + pub status: UploadStatus, +} diff --git a/soulseek-rs-lib/tests/e2e.rs b/soulseek-rs-lib/tests/e2e.rs index 10bea24..a1ccbc5 100644 --- a/soulseek-rs-lib/tests/e2e.rs +++ b/soulseek-rs-lib/tests/e2e.rs @@ -1245,6 +1245,24 @@ fn two_real_clients_search_and_download() { .expect("downloaded file should exist"); assert_eq!(written, content, "downloaded bytes should match the source"); + // The uploader side tracked the transfer and saw it complete. + let deadline = Instant::now() + Duration::from_secs(5); + let mut uploads = sharer.uploads(); + while Instant::now() < deadline + && !uploads + .iter() + .any(|u| u.status == soulseek_rs::types::UploadStatus::Completed) + { + std::thread::sleep(Duration::from_millis(200)); + uploads = sharer.uploads(); + } + let upload = uploads + .iter() + .find(|u| u.status == soulseek_rs::types::UploadStatus::Completed) + .expect("uploader should record a completed upload"); + assert_eq!(upload.username, "e2e_leecher"); + assert_eq!(upload.bytes_sent, size); + let _ = std::fs::remove_dir_all(share_dir); let _ = std::fs::remove_dir_all(download_dir); }