diff --git a/src/client.rs b/src/client.rs index 1e99540..ff1dd1a 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,3 +1,4 @@ +use crate::types::{DownloadResult, DownloadStatus}; use crate::{ error::{Result, SoulseekRs}, peer::{ @@ -204,20 +205,23 @@ impl Client { size: u64, download_directory: String, ) -> Result { - use crate::types::{DownloadResult, DownloadStatus}; - use std::time::{Duration, Instant}; - let start_time = Instant::now(); let hash = md5::md5(&filename); let token = u32::from_str_radix(&hash[0..5], 16)?; + let (download_sender, download_receivier): ( + Sender, + Receiver, + ) = mpsc::channel(); + let download = Download { username: username.clone(), filename: filename.clone(), token, size, download_directory, + sender: download_sender, }; let mut context = self.context.write().unwrap(); @@ -230,9 +234,6 @@ impl Client { drop(context); - let timeout = Duration::from_secs(150); - let check_interval = Duration::from_millis(100); - if !download_initiated { return Ok(DownloadResult { filename, @@ -242,14 +243,12 @@ impl Client { }); } - while start_time.elapsed() < timeout { - std::thread::sleep(check_interval); - } + let status = download_receivier.recv().unwrap(); Ok(DownloadResult { filename, username, - status: DownloadStatus::TimedOut, + status, elapsed_time: start_time.elapsed(), }) } @@ -299,6 +298,11 @@ impl Client { let own_username = own_username.clone(); let client_context_clone = client_context.clone(); + trace!( + "[client] DownloadFromPeer token: {} peer: {:?}", + token, + peer + ); match maybe_download { Some(download) => { thread::spawn(move || { @@ -318,13 +322,14 @@ impl Client { Some(filename) => { match download_peer.download_file( client_context_clone, - Some(download.size as usize), - Some(download.download_directory) + Some(download.clone()) ) { - Ok((bytes, filename)) => { - info!("Successfully downloaded {} bytes to {}", bytes, filename); + Ok((download, filename)) => { + download.sender.send(DownloadStatus::Completed).unwrap(); + info!("Successfully downloaded {} bytes to {}", download.size, filename); } Err(e) => { + download.sender.send(DownloadStatus::Failed).unwrap(); error!( "Failed to download file '{}' from {}:{} (token: {}) - Error: {}", filename, peer.host, peer.port, download.token, e @@ -470,6 +475,7 @@ impl Client { size: transfer.size, download_directory: download .download_directory, + sender: download.sender.clone(), }, ); context.download_tokens.remove(&key); @@ -537,11 +543,10 @@ impl Client { match download_peer .download_file( client_context2.clone(), - None, None ) { - Ok((bytes, filename)) => { - trace!("[client] downloaded {} bytes {:?} ", filename,bytes); + Ok((download, filename)) => { + trace!("[client] downloaded {} bytes {:?} ", filename, download.size); } Err(e) => { trace!("[client] failed to download: {}", e); diff --git a/src/peer/download_peer.rs b/src/peer/download_peer.rs index a71e6f2..d9ea86d 100644 --- a/src/peer/download_peer.rs +++ b/src/peer/download_peer.rs @@ -1,6 +1,7 @@ use crate::client::ClientContext; use crate::message::server::MessageFactory; use crate::trace; +use crate::types::Download; use std::env; use std::fs::{self, File}; use std::io::{self, BufWriter, Read, Write}; @@ -80,24 +81,14 @@ impl FileManager { } fn create_download_path_from_filename( - output_directory: Option<&str>, - remote_username: &str, - token: u32, - filename: Option<&str>, + output_directory: String, + filename: String, ) -> String { - let base_dir = output_directory.unwrap_or("/tmp"); - - let mut expanded_base_dir = FileManager::expand_path(base_dir); - if expanded_base_dir.ends_with('/') { - expanded_base_dir.pop(); - } - - if let Some(filename) = filename { - let clean_filename = Self::extract_filename_from_path(filename); - format!("{expanded_base_dir}/{clean_filename}") - } else { - format!("{expanded_base_dir}/{remote_username}_{token}") - } + format!( + "{}/{}", + output_directory.trim_end_matches('/'), + Self::extract_filename_from_path(&filename) + ) } } @@ -245,9 +236,8 @@ impl DownloadPeer { pub fn download_file( self, client_context: Arc>, - mut expected_size: Option, - mut download_path: Option, - ) -> Result<(usize, String), io::Error> { + mut download: Option, + ) -> Result<(Download, String), io::Error> { let mut stream = self.establish_connection()?; trace!("[download_peer:{}] connected", self.username); @@ -256,7 +246,6 @@ impl DownloadPeer { let mut processor = StreamProcessor::new(self.no_pierce, self.token); let mut read_buffer = [1u8; 8192]; - let mut download_info: Option = None; trace!( "[download_peer:{}] Starting to read data from peer", @@ -265,12 +254,13 @@ impl DownloadPeer { loop { match stream.read(&mut read_buffer) { + // connection closed Ok(0) => { trace!( "[download_peer:{}] connection closed by peer", self.username ); - break; // Connection closed + break; } Ok(bytes_read) => { let data = &read_buffer[..bytes_read]; @@ -288,30 +278,27 @@ impl DownloadPeer { token ); - download_info = client_context - .read() - .unwrap() - .download_tokens - .get(&token_u32) - .cloned(); + let read = client_context.read().unwrap(); - match &download_info { - Some(download) => { + let download_info = + read.download_tokens.get(&token_u32); + + match download_info { + Some(d) => { trace!( "[download_peer:{}] got download info for token: {} - filename: {}", self.username, token_u32, - download.filename + d.filename ); - expected_size = Some(download.size as usize); - download_path = - Some(download.download_directory.clone()); + trace!("[download_peer:{}] setting download info for token: {}", + self.username, + token_u32); + download = Some(d.clone()); } None => { - trace!( - "[download_peer:{}] no download info for token: {}", - self.username, - self.token + panic!( + "No download info for token {token_u32}" ); } } @@ -321,7 +308,9 @@ impl DownloadPeer { processor.process_data_chunk(data); - if !processor.should_continue(expected_size) { + if !processor.should_continue(Some( + download.clone().unwrap().size as usize, + )) { break; } } @@ -334,45 +323,39 @@ impl DownloadPeer { "[download_peer:{}] finished reading data from peer", self.username ); - - let output_directory = download_path.as_ref().map(|path| { - if Path::new(path).is_dir() { - path.as_str() - } else if let Some(parent) = Path::new(path).parent() { - parent.to_str().unwrap_or("/tmp") - } else { - "/tmp" - } + let download = download.clone().unwrap_or_else(|| { + panic!( + "[download_peer] No download info found for token: {}", + self.token + ) }); - let filename = download_info.as_ref().map(|d| d.filename.as_str()); + let download_directory = download.download_directory.clone(); + let mut expaned_path = FileManager::expand_path(&download_directory); + if !Path::new(&expaned_path).is_dir() { + expaned_path = Path::new(&expaned_path) + .parent() + .unwrap() + .to_str() + .unwrap() + .to_string() + } - let final_path = if let Some(output_path) = &download_path { - if Path::new(output_path).is_dir() || output_path.ends_with('/') { - // output_path is a directory, append filename - FileManager::create_download_path_from_filename( - Some(output_path), - &self.username, - self.token, - filename, - ) - } else { - output_path.clone() - } - } else { - FileManager::create_download_path_from_filename( - output_directory, - &self.username, - self.token, - filename, - ) - }; + let final_path = FileManager::create_download_path_from_filename( + expaned_path, + download.filename.clone(), + ); if let Some(parent) = Path::new(&final_path).parent() { fs::create_dir_all(parent)?; } - fs::write(&final_path, &processor.buffer)?; + fs::write(&final_path, &processor.buffer).unwrap_or_else(|e| { + panic!( + "[download_peer] Failed to write to file {}: {}", + final_path, e + ) + }); trace!( "[download_peer:{}] download completed successfully: {} bytes, saved to: {}", @@ -381,7 +364,7 @@ impl DownloadPeer { final_path ); - Ok((processor.total_bytes, final_path)) + Ok((download.clone(), final_path)) } } @@ -390,7 +373,7 @@ mod tests { use super::{DownloadPeer, FileManager, StreamProcessor}; use crate::client::ClientContext; use std::fs; - use std::io::{Read, Write}; + use std::io::Read; use std::net::TcpListener; use std::sync::{Arc, RwLock}; use std::thread; @@ -441,9 +424,7 @@ mod tests { "own_username".to_string(), ); let dummy_context = Arc::new(RwLock::new(ClientContext::new())); - let _ = download_peer - .download_file(dummy_context, None, None) - .unwrap(); + let _ = download_peer.download_file(dummy_context).unwrap(); // Give the client time to send messages thread::sleep(Duration::from_millis(10)); @@ -466,9 +447,7 @@ mod tests { "own_username".to_string(), ); let dummy_context = Arc::new(RwLock::new(ClientContext::new())); - let _ = download_peer - .download_file(dummy_context, None, None) - .unwrap(); + let _ = download_peer.download_file(dummy_context).unwrap(); // Give the client time to send messages thread::sleep(Duration::from_millis(10)); @@ -486,77 +465,6 @@ mod tests { // No messages expected when no_pierce is false } - #[test] - pub fn test_download_file() { - let test_data = b"test file content"; - let messages = Arc::new(RwLock::new(Vec::>::new())); - let messages_clone = Arc::clone(&messages); - - // Create a test server that sends test data - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let port = listener.local_addr().unwrap().port(); - - thread::spawn(move || { - for stream in listener.incoming() { - match stream { - Ok(mut stream) => { - // Set non-blocking to avoid hanging - stream.set_nonblocking(true).unwrap(); - - // Read handshake data - let mut buffer = [0u8; 1024]; - std::thread::sleep(Duration::from_millis(10)); - if let Ok(bytes_read) = stream.read(&mut buffer) { - if bytes_read > 0 { - messages_clone - .write() - .unwrap() - .push(buffer[..bytes_read].to_vec()); - } - } - - // Send pierce token (4 bytes) then test data - stream.set_nonblocking(false).unwrap(); - let _ = stream.write_all(&[42, 0, 0, 0]); - std::thread::sleep(Duration::from_millis(10)); - let _ = stream.write_all(test_data); - } - Err(_) => break, - } - } - }); - - thread::sleep(Duration::from_millis(10)); - - let download_peer = DownloadPeer::new( - "remote_user".to_string(), - "127.0.0.1".to_string(), - u32::from(port), - 42, - false, - "test_user".to_string(), - ); - - let dummy_context = Arc::new(RwLock::new(ClientContext::new())); - let result = download_peer.download_file( - dummy_context, - Some(test_data.len()), - Some("test_download.mp3".to_string()), - ); - - assert!(result.is_ok()); - let (bytes_downloaded, file_path) = result.unwrap(); - assert_eq!(bytes_downloaded, test_data.len()); - assert_eq!(file_path, "test_download.mp3"); - - // Verify the file was written correctly - let downloaded_data = fs::read("test_download.mp3").unwrap(); - assert_eq!(downloaded_data, test_data); - - // Clean up test file - let _ = fs::remove_file("test_download.mp3"); - } - #[test] fn test_file_manager_create_paths_with_custom_path() { let paths = FileManager::create_download_paths( diff --git a/src/types.rs b/src/types.rs index faf6416..85cbbe8 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::mpsc::Sender}; use crate::{error::Result, message::Message, utils::zlib::deflate}; @@ -88,6 +88,7 @@ pub struct Download { pub token: u32, pub size: u64, pub download_directory: String, + pub sender: Sender, } #[derive(Debug, Clone)]