From 9553537bc6a9011715db9f16e2cc16c391fbd7ea Mon Sep 17 00:00:00 2001 From: Michel de Graaf Date: Thu, 31 Jul 2025 20:06:41 +0200 Subject: [PATCH] feat: implement download peer --- rustfmt.toml | 1 + src/main.rs | 2 +- src/message/mod.rs | 4 + src/message/peer/upload_failed.rs | 2 +- src/message/server/message_factory.rs | 17 +- src/peer/download_peer.rs | 549 ++++++++++++++++++++++++++ src/peer/mod.rs | 18 +- src/server.rs | 1 + src/types.rs | 3 + 9 files changed, 590 insertions(+), 7 deletions(-) create mode 100644 rustfmt.toml create mode 100644 src/peer/download_peer.rs diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..df99c69 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1 @@ +max_width = 80 diff --git a/src/main.rs b/src/main.rs index 1a1aca6..bdbf112 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,7 +9,7 @@ mod types; use client::Client; use server::PeerAddress; -use std::time::{Duration, Instant}; +use std::time::Duration; // PeerAddress::new(String::from("server.slsknet.org"), 2416), fn main() { diff --git a/src/message/mod.rs b/src/message/mod.rs index 00829e1..2e86619 100644 --- a/src/message/mod.rs +++ b/src/message/mod.rs @@ -9,6 +9,7 @@ pub use message_reader::*; use std::str; #[derive(Debug, PartialEq)] +#[allow(dead_code)] pub enum MessageType { Server, Peer, @@ -141,6 +142,7 @@ impl Message { pub fn new_with_data(data: Vec) -> Self { Self { data, pointer: 0 } } + #[allow(dead_code)] pub fn reset_pointer(&mut self) { self.pointer = 0; } @@ -244,6 +246,7 @@ impl Message { self } + #[allow(dead_code)] pub fn write_int64(&mut self, value: u64) -> &mut Self { self.data.extend_from_slice(&value.to_le_bytes()); self @@ -258,6 +261,7 @@ impl Message { self } + #[allow(dead_code)] pub fn write_raw_hex_string(&mut self, val: &str) -> &mut Self { let mut b = Vec::new(); for i in (0..val.len()).step_by(2) { diff --git a/src/message/peer/upload_failed.rs b/src/message/peer/upload_failed.rs index c016013..42de652 100644 --- a/src/message/peer/upload_failed.rs +++ b/src/message/peer/upload_failed.rs @@ -10,7 +10,7 @@ impl MessageHandler for UploadFailedHandler { fn get_code(&self) -> u8 { 46 } - fn handle(&self, message: &mut Message, sender: Sender) { + fn handle(&self, message: &mut Message, _sender: Sender) { let upload_failed = UploadFailed::new_from_message(message); info!("Upload failed for ${}", upload_failed.filename); } diff --git a/src/message/server/message_factory.rs b/src/message/server/message_factory.rs index 5f78e83..b9b888c 100644 --- a/src/message/server/message_factory.rs +++ b/src/message/server/message_factory.rs @@ -1,4 +1,4 @@ -use crate::{message::Message, types::Transfer, utils::md5}; +use crate::{message::Message, peer::ConnectionType, types::Transfer, utils::md5}; pub struct MessageFactory; impl MessageFactory { @@ -91,7 +91,20 @@ impl MessageFactory { } pub fn build_pierce_firewall_message(token: u32) -> Message { Message::new() - .write_int32(0) // PierceFirewall message code + .write_int32(0) // PierceFirewall message code + .write_int32(token) + .clone() + } + + #[allow(dead_code)] + pub fn build_peer_init_message( + own_username: &str, + connection_type: ConnectionType, + token: u32, + ) -> Message { + Message::new() + .write_string(own_username) + .write_string(&connection_type.to_string()) .write_int32(token) .clone() } diff --git a/src/peer/download_peer.rs b/src/peer/download_peer.rs new file mode 100644 index 0000000..3d3a391 --- /dev/null +++ b/src/peer/download_peer.rs @@ -0,0 +1,549 @@ +use crate::message::server::MessageFactory; +use std::fs::{self, File}; +use std::io::{self, BufWriter, Read, Write}; +use std::net::TcpStream; +use std::net::ToSocketAddrs; +use std::path::Path; +use std::thread::sleep; +use std::time::Duration; + +#[derive(Debug)] +struct DownloadPaths { + final_path: String, + incomplete_path: String, +} + +struct FileManager; + +impl FileManager { + fn create_download_paths( + output_path: Option, + username: &str, + token: u32, + ) -> DownloadPaths { + let final_path = match output_path { + Some(path) if !path.is_empty() => path, + _ => format!("/tmp/{}_{}.mp3", username, token), + }; + let incomplete_path = format!("{}.incomplete", final_path); + + DownloadPaths { + final_path, + incomplete_path, + } + } + + fn create_temp_file(path: &str) -> Result, io::Error> { + // Create directory if needed + if let Some(parent) = Path::new(path).parent() { + fs::create_dir_all(parent)?; + } + + let temp_file = File::create(path)?; + Ok(BufWriter::new(temp_file)) + } + + fn finalize_download( + incomplete_path: &str, + final_path: &str, + ) -> Result<(), io::Error> { + fs::rename(incomplete_path, final_path) + } + + fn cleanup_on_error(incomplete_path: Option<&str>) { + if let Some(path) = incomplete_path { + let _ = fs::remove_file(path); + } + } +} + +struct StreamProcessor { + no_pierce: bool, + token: u32, + total_bytes: usize, + received: bool, +} + +impl StreamProcessor { + fn new(no_pierce: bool, token: u32) -> Self { + Self { + no_pierce, + token, + total_bytes: 0, + received: false, + } + } + + fn handle_pierce_token( + &mut self, + data: &[u8], + stream: &mut TcpStream, + ) -> Result { + if !self.no_pierce && !self.received && data.len() >= 4 { + stream + .write_all(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])?; + self.received = true; + return Ok(true); // Skip this data chunk + } + Ok(false) + } + + fn process_data_chunk( + &mut self, + data: &[u8], + writer: &mut Option>, + ) -> Result<(), io::Error> { + if let Some(ref mut w) = writer { + w.write_all(data)?; + + // Flush every 64KB to ensure data reaches disk + if self.total_bytes % (64 * 1024) == 0 { + w.flush()?; + } + } + + self.total_bytes += data.len(); + Ok(()) + } + + fn should_continue(&self, expected_size: Option) -> bool { + if let Some(size) = expected_size { + self.total_bytes < size + } else { + true + } + } +} + +#[allow(dead_code)] +pub struct DownloadPeer { + username: String, + host: String, + port: u32, + own_username: String, + token: u32, + no_pierce: bool, +} + +impl DownloadPeer { + #[allow(dead_code)] + pub fn new( + username: String, + host: String, + port: u32, + token: u32, + no_pierce: bool, + own_username: String, + ) -> Self { + Self { + username, + host, + port, + own_username, + token, + no_pierce, + } + } + + fn establish_connection(&self) -> Result { + let socket_address = format!("{}:{}", self.host, self.port) + .to_socket_addrs()? + .next() + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "Invalid address") + })?; + + let stream = TcpStream::connect_timeout( + &socket_address, + Duration::from_secs(20), + )?; + + stream.set_read_timeout(Some(Duration::from_secs(30)))?; + stream.set_write_timeout(Some(Duration::from_secs(5)))?; + + Ok(stream) + } + + fn perform_handshake( + &self, + stream: &mut TcpStream, + ) -> Result<(), io::Error> { + if self.no_pierce { + let message = MessageFactory::build_peer_init_message( + &self.own_username, + super::ConnectionType::F, + self.token, + ); + stream.write_all(&message.get_data())?; + sleep(Duration::from_millis(1000)); + stream + .write_all(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])?; + } else { + stream.write_all( + &MessageFactory::build_pierce_firewall_message(self.token) + .get_data(), + )?; + } + Ok(()) + } + + pub fn download_file( + self, + expected_size: Option, + output_path: Option, + ) -> Result { + let mut stream = self.establish_connection()?; + + // Setup file management if output path is provided + let (paths, mut writer) = if let Some(output_path) = output_path { + let paths = FileManager::create_download_paths( + Some(output_path), + &self.own_username, + self.token, + ); + let writer = + Some(FileManager::create_temp_file(&paths.incomplete_path)?); + (Some(paths), writer) + } else { + (None, None) + }; + + // Perform handshake + self.perform_handshake(&mut stream)?; + + // Stream data + let mut processor = StreamProcessor::new(self.no_pierce, self.token); + let mut read_buffer = [0u8; 8192]; + + loop { + match stream.read(&mut read_buffer) { + Ok(0) => break, // Connection closed + Ok(bytes_read) => { + let data = &read_buffer[..bytes_read]; + + // Handle pierce token extraction + if processor.handle_pierce_token(data, &mut stream)? { + continue; // Skip this data chunk + } + + // Process file data + processor.process_data_chunk(data, &mut writer)?; + + // Check completion + if !processor.should_continue(expected_size) { + break; + } + } + Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { + continue; + } + Err(e) => { + // Clean up incomplete file on error + drop(writer); + if let Some(ref paths) = paths { + FileManager::cleanup_on_error(Some( + &paths.incomplete_path, + )); + } + return Err(e); + } + } + } + + // Finalize download + if let Some(mut w) = writer { + w.flush()?; + drop(w); + + if let Some(paths) = paths { + FileManager::finalize_download( + &paths.incomplete_path, + &paths.final_path, + )?; + } + } + + Ok(processor.total_bytes) + } +} + +#[cfg(test)] +mod tests { + use super::{DownloadPeer, FileManager, StreamProcessor}; + use std::fs; + use std::io::{Read, Write}; + use std::net::TcpListener; + use std::sync::{Arc, Mutex}; + use std::thread; + use std::time::Duration; + + pub fn build_test_server() -> (u16, Arc>>>) { + let messages = Arc::new(Mutex::new(Vec::>::new())); + let messages_clone = Arc::clone(&messages); + + // Find an available port by binding to port 0 + 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) => { + let mut buffer = [0u8; 1024]; + match stream.read(&mut buffer) { + Ok(bytes_read) if bytes_read > 0 => { + let data = buffer[..bytes_read].to_vec(); + messages_clone.lock().unwrap().push(data); + } + _ => {} + } + } + Err(_) => break, + } + } + }); + + // Give the server time to start + thread::sleep(Duration::from_millis(10)); + + (port, messages) + } + + #[test] + pub fn test_connect_no_pierce() { + let (port, messages) = build_test_server(); + let token = 33; + let download_peer = DownloadPeer::new( + "test_user".to_string(), + "127.0.0.1".to_string(), + port as u32, + token, + false, + "own_username".to_string(), + ); + let _ = download_peer.download_file(None, None).unwrap(); + + // Give the client time to send messages + thread::sleep(Duration::from_millis(10)); + + let received_messages = messages.lock().unwrap(); + assert_eq!(received_messages.len(), 1); + assert_eq!(received_messages[0], vec![0, 0, 0, 0, 33, 0, 0, 0]); + } + + #[test] + pub fn test_connect_with_pierce() { + let (port, messages) = build_test_server(); + let token = 33; + let download_peer = DownloadPeer::new( + "test_user".to_string(), + "127.0.0.1".to_string(), + port as u32, + token, + true, // no_pierce = true, should send init message + "own_username".to_string(), + ); + let _ = download_peer.download_file(None, None).unwrap(); + + // Give the client time to send messages + thread::sleep(Duration::from_millis(10)); + + let received_messages = messages.lock().unwrap(); + assert_eq!(received_messages.len(), 1); // Expect one init message + assert_eq!( + received_messages[0], + vec![ + 12, 0, 0, 0, 111, 119, 110, 95, 117, 115, 101, 114, 110, 97, + 109, 101, 1, 0, 0, 0, 70, 33, 0, 0, 0 + ] + ); + + // 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(Mutex::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 + .lock() + .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(), + port as u32, + 42, + false, + "test_user".to_string(), + ); + + let result = download_peer.download_file( + Some(test_data.len()), + Some("test_download.mp3".to_string()), + ); + + assert!(result.is_ok()); + let bytes_downloaded = result.unwrap(); + assert_eq!(bytes_downloaded, test_data.len()); + + // 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( + Some("custom/path.mp3".to_string()), + "user", + 123, + ); + assert_eq!(paths.final_path, "custom/path.mp3"); + assert_eq!(paths.incomplete_path, "custom/path.mp3.incomplete"); + } + + #[test] + fn test_file_manager_create_paths_with_default() { + let paths = FileManager::create_download_paths(None, "testuser", 456); + assert_eq!(paths.final_path, "/tmp/testuser_456.mp3"); + assert_eq!(paths.incomplete_path, "/tmp/testuser_456.mp3.incomplete"); + } + + #[test] + fn test_file_manager_create_paths_with_empty_string() { + let paths = FileManager::create_download_paths( + Some("".to_string()), + "user", + 789, + ); + assert_eq!(paths.final_path, "/tmp/user_789.mp3"); + assert_eq!(paths.incomplete_path, "/tmp/user_789.mp3.incomplete"); + } + + #[test] + fn test_file_manager_create_temp_file() { + let temp_path = "test_temp.txt.incomplete"; + let result = FileManager::create_temp_file(temp_path); + assert!(result.is_ok()); + + // Verify file was created + assert!(fs::metadata(temp_path).is_ok()); + + // Clean up + let _ = fs::remove_file(temp_path); + } + + #[test] + fn test_file_manager_finalize_download() { + // Create a temporary file + let temp_path = "test_incomplete.txt"; + let final_path = "test_final.txt"; + fs::write(temp_path, "test content").unwrap(); + + let result = FileManager::finalize_download(temp_path, final_path); + assert!(result.is_ok()); + + // Verify rename worked + assert!(fs::metadata(final_path).is_ok()); + assert!(fs::metadata(temp_path).is_err()); // Original should be gone + + // Clean up + let _ = fs::remove_file(final_path); + } + + #[test] + fn test_stream_processor_new() { + let processor = StreamProcessor::new(true, 123); + assert_eq!(processor.no_pierce, true); + assert_eq!(processor.token, 123); + assert_eq!(processor.total_bytes, 0); + assert_eq!(processor.received, false); + } + + #[test] + fn test_stream_processor_should_continue() { + let processor = StreamProcessor::new(false, 123); + + // Without expected size, should always continue + assert!(processor.should_continue(None)); + + // With expected size, should continue if under limit + assert!(processor.should_continue(Some(100))); + } + + #[test] + fn test_stream_processor_should_continue_with_limit() { + let mut processor = StreamProcessor::new(false, 123); + processor.total_bytes = 150; + + // Should not continue if over limit + assert!(!processor.should_continue(Some(100))); + + // Should continue if under limit + assert!(processor.should_continue(Some(200))); + } + + #[test] + fn test_stream_processor_process_data_chunk() { + let mut processor = StreamProcessor::new(false, 123); + let data = b"test data"; + let mut writer = None; + + let result = processor.process_data_chunk(data, &mut writer); + assert!(result.is_ok()); + assert_eq!(processor.total_bytes, data.len()); + } + + #[test] + fn test_establish_connection_invalid_address() { + let download_peer = DownloadPeer::new( + "user".to_string(), + "invalid-host".to_string(), + 9999, + 123, + false, + "own_user".to_string(), + ); + + let result = download_peer.establish_connection(); + assert!(result.is_err()); + } +} diff --git a/src/peer/mod.rs b/src/peer/mod.rs index f92c054..3f25b9e 100644 --- a/src/peer/mod.rs +++ b/src/peer/mod.rs @@ -1,6 +1,7 @@ pub mod default_peer; - pub use default_peer::*; + +pub mod download_peer; pub mod listen; use crate::message::Message; @@ -40,6 +41,17 @@ impl FromStr for ConnectionType { } } +impl fmt::Display for ConnectionType { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let s = match self { + ConnectionType::P => "P", + ConnectionType::F => "F", + ConnectionType::D => "D", + }; + write!(f, "{}", s) + } +} + #[derive(Debug, Clone)] #[allow(dead_code)] pub struct Peer { @@ -128,7 +140,7 @@ fn test_new_from_message() { assert!(matches!(peer.connection_type, ConnectionType::P)); assert_eq!(peer.host, "45.37.25.27"); assert_eq!(peer.port, 2234); - assert_eq!(peer.token, Some("b24e1900".to_string())); + assert_eq!(peer.token, Some(1683702)); assert_eq!(peer.privileged, 0); assert_eq!(peer.unknown, 0); assert_eq!(peer.obfuscated_port, 0); @@ -150,7 +162,7 @@ fn test_new_from_message2() { assert!(matches!(peer.connection_type, ConnectionType::P)); assert_eq!(peer.host, "68.63.128.119"); assert_eq!(peer.port, 2235); - assert_eq!(peer.token, Some("3a100000".to_string())); + assert_eq!(peer.token, Some(4122)); assert_eq!(peer.privileged, 0); assert_eq!(peer.unknown, 1); assert_eq!(peer.obfuscated_port, 0); diff --git a/src/server.rs b/src/server.rs index 592eced..b6f9e3b 100644 --- a/src/server.rs +++ b/src/server.rs @@ -423,6 +423,7 @@ impl Server { self.queue_message(MessageFactory::build_file_search_message(token, query)); } + #[allow(dead_code)] pub fn pierce_firewall(&self, token: u32) { match self.sender.send(ServerOperation::PierceFirewall(token)) { Ok(_) => {} diff --git a/src/types.rs b/src/types.rs index 91a1ee4..24b8b91 100644 --- a/src/types.rs +++ b/src/types.rs @@ -73,6 +73,7 @@ impl FileSearchResult { } #[derive(Debug, Clone)] +#[allow(dead_code)] pub struct Transfer { pub direction: u32, pub token: u32, @@ -81,6 +82,7 @@ pub struct Transfer { } #[derive(Debug, Clone)] +#[allow(dead_code)] pub struct DownloadResult { pub filename: String, pub username: String, @@ -89,6 +91,7 @@ pub struct DownloadResult { } #[derive(Debug, Clone)] +#[allow(dead_code)] pub enum DownloadStatus { Pending, InProgress,