From 297badbddf6cb019de512813b10720afa391ee6b Mon Sep 17 00:00:00 2001 From: Michel de Graaf Date: Thu, 6 Nov 2025 10:59:57 +0100 Subject: [PATCH] fix: direct download --- src/client.rs | 15 +++--- src/message/message_reader.rs | 4 ++ src/peer/default_peer.rs | 3 +- src/peer/download_peer.rs | 83 +++++++++++++++++++----------- src/peer/listen.rs | 96 +++++++++++++++++++++++------------ 5 files changed, 129 insertions(+), 72 deletions(-) diff --git a/src/client.rs b/src/client.rs index b12f7b1..1dda7ed 100644 --- a/src/client.rs +++ b/src/client.rs @@ -34,7 +34,7 @@ pub enum ClientOperation { SearchResult(FileSearchResult), PeerDisconnected(String), PierceFireWall(Peer), - DownloadFromPeer(u32, Peer), + DownloadFromPeer(u32, Peer, bool), UpdateDownloadTokens(Transfer, String), GetPeerAddressResponse { username: String, @@ -168,18 +168,15 @@ impl Client { } pub fn login(&self) -> Result { - // Attempt to login info!("Logging in as {}", self.username); if let Some(server) = &self.server { - let result = server.login(&self.username, &self.password)?; - if result { + return if (server.login(&self.username, &self.password))? { Ok(true) } else { Err(SoulseekRs::AuthenticationFailed) - } - } else { - Err(SoulseekRs::NotConnected) + }; } + Err(SoulseekRs::NotConnected) } #[allow(dead_code)] @@ -306,7 +303,7 @@ impl Client { own_username.clone(), ); } - ClientOperation::DownloadFromPeer(token, peer) => { + ClientOperation::DownloadFromPeer(token, peer, allowed) => { let maybe_download = { let client_context = client_context.read().unwrap(); client_context.download_tokens.get(&token).cloned() @@ -327,7 +324,7 @@ impl Client { peer.host.clone(), peer.port, token, - false, + allowed, own_username, ); let filename: Option<&str> = download diff --git a/src/message/message_reader.rs b/src/message/message_reader.rs index 8925b11..09947ba 100644 --- a/src/message/message_reader.rs +++ b/src/message/message_reader.rs @@ -67,6 +67,10 @@ impl MessageReader { self.buffer.len() } + pub fn get_buffer(&mut self) -> Vec { + self.buffer.drain(..).collect() + } + pub fn extract_message(&mut self) -> io::Result> { let bytes_read = self.buffer.len(); if bytes_read < 4 { diff --git a/src/peer/default_peer.rs b/src/peer/default_peer.rs index 58ea4b4..3fadcde 100644 --- a/src/peer/default_peer.rs +++ b/src/peer/default_peer.rs @@ -58,7 +58,7 @@ impl DefaultPeer { write_thread: None, } } - pub fn disconnect(mut self) { + pub fn disconnect(&mut self) { let username = self.peer.read().unwrap().username.clone(); debug!("[default_peer:{}] disconnect", username); if let Err(e) = self @@ -385,6 +385,7 @@ impl DefaultPeer { .send(ClientOperation::DownloadFromPeer( token, peer_clone.read().unwrap().clone(), + allowed, )) .unwrap(); } diff --git a/src/peer/download_peer.rs b/src/peer/download_peer.rs index 3ab6522..803215e 100644 --- a/src/peer/download_peer.rs +++ b/src/peer/download_peer.rs @@ -49,7 +49,6 @@ impl FileManager { } 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)?; } @@ -72,7 +71,6 @@ impl FileManager { } fn extract_filename_from_path(full_path: &str) -> String { - // Split on both forward slashes and backslashes to handle Windows and Unix paths full_path .split(['/', '\\']) .next_back() @@ -124,6 +122,7 @@ impl StreamProcessor { self.received = true; return Ok(true); // Skip this data chunk } + self.received = true; Ok(false) } @@ -202,27 +201,27 @@ impl DownloadPeer { self.no_pierce ); if self.no_pierce { - let message = MessageFactory::build_peer_init_message( - &self.own_username, - super::ConnectionType::F, - self.token, - ); - trace!( - "[default_peer:{}] sending peer init, token: {} message: {:?} - ", - self.token, - self.username, - message.get_buffer(), - ); - stream.write_all(&message.get_buffer())?; - stream.flush()?; + // let message = MessageFactory::build_peer_init_message( + // &self.own_username, + // super::ConnectionType::F, + // self.token, + // ); + // trace!( + // "[download_peer:{}] sending peer init, token: {} message: {:?} + // ", + // self.username, + // self.token, + // message.get_buffer(), + // ); + // stream.write_all(&message.get_buffer())?; + // stream.flush()?; } else { let message = MessageFactory::build_pierce_firewall_message(self.token); stream.write_all(&message.get_buffer())?; trace!( - "[default_peer:{}] sending pierce firewall message token: {}: {:?}", + "[download_peer:{}] sending pierce firewall message token: {}: {:?}", self.username, self.token, &message.get_buffer() @@ -238,9 +237,13 @@ impl DownloadPeer { mut download: Option, stream: Option, ) -> Result<(Download, String), io::Error> { - if stream.is_none() { - trace!("[download_peer:{}] stream is none", self.username); - } + 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 = stream.unwrap_or_else(|| { self.establish_connection().unwrap_or_else(|e| { panic!("Failed to establish connection: {:?}", e) @@ -260,9 +263,13 @@ impl DownloadPeer { self.username ); + if download.is_some() { + stream + .write_all(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])?; + } + loop { match stream.read(&mut read_buffer) { - // connection closed Ok(0) => { trace!( "[download_peer:{}] connection closed by peer. bytes read: {}", @@ -274,10 +281,7 @@ impl DownloadPeer { Ok(bytes_read) => { let data = &read_buffer[..bytes_read]; - if !self.no_pierce - && !processor.received - && processor.handle_pierce_token(data, &mut stream)? - { + if !self.no_pierce && !processor.received { let token = data.get(0..4).unwrap(); let token_u32 = u32::from_le_bytes( token @@ -285,11 +289,15 @@ impl DownloadPeer { .unwrap_or_else(|_| panic!("[download_peer:{}] slice with incorrect length", self.username)), ); trace!( - "[download_peer:{}] received token_u32: {:?} - token: {:?}", + "[download_peer:{}] got token: {} from data chunk", self.username, - token_u32, - token + token_u32 ); + processor.received = true; + + stream.write_all(&[ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + ])?; let read = client_context.read().unwrap(); @@ -326,7 +334,7 @@ impl DownloadPeer { processor.process_data_chunk(data); if !processor.should_continue(Some( - download.clone().unwrap().size as usize, + download.as_ref().unwrap().size as usize, )) { break; } @@ -340,6 +348,23 @@ impl DownloadPeer { "[download_peer:{}] finished reading data from peer", self.username ); + + // let final_path = FileManager::create_download_path_from_filename( + // "/tmp/".to_string(), + // "test.mp3".to_string(), + // ); + // + // if let Some(parent) = Path::new(&final_path).parent() { + // fs::create_dir_all(parent)?; + // } + // + // fs::write(&final_path, &processor.buffer).unwrap_or_else(|e| { + // panic!( + // "[download_peer] Failed to write to file {}: {}", + // final_path, e + // ) + // }); + let read = client_context.read().unwrap(); let tokens = read.download_tokens.keys().collect::>(); diff --git a/src/peer/listen.rs b/src/peer/listen.rs index 3a8f6e0..134ce11 100644 --- a/src/peer/listen.rs +++ b/src/peer/listen.rs @@ -7,6 +7,7 @@ use crate::client::{ClientContext, ClientOperation}; use crate::message::MessageReader; use crate::peer::{ConnectionType, DefaultPeer, DownloadPeer, Peer}; +use crate::types::Download; use crate::{debug, error, info, trace, DownloadStatus}; pub struct Listen {} @@ -58,7 +59,7 @@ impl Listen { debug!("[listener:{peer_ip}:{peer_port}] peerInit (0) username: {username} connection_type: {connection_type} token: {token}"); let peer = Peer::new( - format!("{}-direct", username.clone()), + username.clone(), ConnectionType::P, peer_ip.clone(), peer_port.into(), @@ -78,49 +79,78 @@ impl Listen { drop(default_peer); } else if connection_type == "F" { - let maybe_download = { - let client_context = client_context.read().unwrap(); - client_context.download_tokens.get(&token).cloned() - }; - - let client_context_clone = client_context.clone(); - trace!( "[client] DownloadFromPeer token: {} peer: {:?}", token, peer ); - thread::spawn(move || { - let download_peer = DownloadPeer::new( - format!("{}-direct", username.clone()), - peer.host.clone(), - peer.port, - token, - false, - own_username, - ); + let mut download: Option = None; + if reader.buffer_len() > 0 { + let buffer = reader.get_buffer(); + trace!("[listener:{peer_ip}:{peer_port}] reader buffer has {} bytes. {:?}", buffer.len(), buffer); - match download_peer.download_file( - client_context_clone, - maybe_download, - Some(stream), - ) { - Ok((download, filename)) => { - download - .sender - .send(DownloadStatus::Completed) - .unwrap(); - info!("Successfully downloaded {} bytes to {}", download.size, filename); - } - Err(e) => { - error!( + let token = buffer.get(0..4).unwrap(); + let token_u32 = u32::from_le_bytes( + token + .try_into() + .unwrap_or_else(|_| panic!("[listener:{}] slice with incorrect length, can't extract transfer_token", username)), + ); + trace!( + "[listener:{}] got transfer_token: {} from data chunk", + username, + token_u32); + + download = client_context + .read() + .unwrap() + .download_tokens + .get(&token_u32) + .cloned(); + + if download.is_none() { + let download_tokens = client_context + .read() + .unwrap() + .download_tokens + .keys() + .cloned() + .collect::>(); + trace!("[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}", token_u32, download_tokens ); + } + } + + let download_peer = DownloadPeer::new( + format!("{}:direct", username.clone()), + peer.host.clone(), + peer.port, + token, + true, + own_username, + ); + + match download_peer.download_file( + client_context, + download, + Some(stream), + ) { + Ok((download, filename)) => { + download + .sender + .send(DownloadStatus::Completed) + .unwrap(); + info!( + "Successfully downloaded {} bytes to {}", + download.size, filename + ); + } + Err(e) => { + error!( "Failed to download file from {}:{} (token: {}) - Error: {}", peer.host, peer.port, token, e ); - } } - }); + } } else { debug!( "[listener:{peer_ip}:{peer_port}] connection type is not P or F"