fix: direct download

This commit is contained in:
Michel de Graaf
2025-11-06 10:59:57 +01:00
parent c623d27e65
commit 297badbddf
5 changed files with 129 additions and 72 deletions
+5 -8
View File
@@ -34,7 +34,7 @@ pub enum ClientOperation {
SearchResult(FileSearchResult), SearchResult(FileSearchResult),
PeerDisconnected(String), PeerDisconnected(String),
PierceFireWall(Peer), PierceFireWall(Peer),
DownloadFromPeer(u32, Peer), DownloadFromPeer(u32, Peer, bool),
UpdateDownloadTokens(Transfer, String), UpdateDownloadTokens(Transfer, String),
GetPeerAddressResponse { GetPeerAddressResponse {
username: String, username: String,
@@ -168,19 +168,16 @@ impl Client {
} }
pub fn login(&self) -> Result<bool> { pub fn login(&self) -> Result<bool> {
// Attempt to login
info!("Logging in as {}", self.username); info!("Logging in as {}", self.username);
if let Some(server) = &self.server { if let Some(server) = &self.server {
let result = server.login(&self.username, &self.password)?; return if (server.login(&self.username, &self.password))? {
if result {
Ok(true) Ok(true)
} else { } else {
Err(SoulseekRs::AuthenticationFailed) Err(SoulseekRs::AuthenticationFailed)
};
} }
} else {
Err(SoulseekRs::NotConnected) Err(SoulseekRs::NotConnected)
} }
}
#[allow(dead_code)] #[allow(dead_code)]
pub fn remove_peer(&self, username: &str) { pub fn remove_peer(&self, username: &str) {
@@ -306,7 +303,7 @@ impl Client {
own_username.clone(), own_username.clone(),
); );
} }
ClientOperation::DownloadFromPeer(token, peer) => { ClientOperation::DownloadFromPeer(token, peer, allowed) => {
let maybe_download = { let maybe_download = {
let client_context = client_context.read().unwrap(); let client_context = client_context.read().unwrap();
client_context.download_tokens.get(&token).cloned() client_context.download_tokens.get(&token).cloned()
@@ -327,7 +324,7 @@ impl Client {
peer.host.clone(), peer.host.clone(),
peer.port, peer.port,
token, token,
false, allowed,
own_username, own_username,
); );
let filename: Option<&str> = download let filename: Option<&str> = download
+4
View File
@@ -67,6 +67,10 @@ impl MessageReader {
self.buffer.len() self.buffer.len()
} }
pub fn get_buffer(&mut self) -> Vec<u8> {
self.buffer.drain(..).collect()
}
pub fn extract_message(&mut self) -> io::Result<Option<Message>> { pub fn extract_message(&mut self) -> io::Result<Option<Message>> {
let bytes_read = self.buffer.len(); let bytes_read = self.buffer.len();
if bytes_read < 4 { if bytes_read < 4 {
+2 -1
View File
@@ -58,7 +58,7 @@ impl DefaultPeer {
write_thread: None, write_thread: None,
} }
} }
pub fn disconnect(mut self) { pub fn disconnect(&mut self) {
let username = self.peer.read().unwrap().username.clone(); let username = self.peer.read().unwrap().username.clone();
debug!("[default_peer:{}] disconnect", username); debug!("[default_peer:{}] disconnect", username);
if let Err(e) = self if let Err(e) = self
@@ -385,6 +385,7 @@ impl DefaultPeer {
.send(ClientOperation::DownloadFromPeer( .send(ClientOperation::DownloadFromPeer(
token, token,
peer_clone.read().unwrap().clone(), peer_clone.read().unwrap().clone(),
allowed,
)) ))
.unwrap(); .unwrap();
} }
+54 -29
View File
@@ -49,7 +49,6 @@ impl FileManager {
} }
fn create_temp_file(path: &str) -> Result<BufWriter<File>, io::Error> { fn create_temp_file(path: &str) -> Result<BufWriter<File>, io::Error> {
// Create directory if needed
if let Some(parent) = Path::new(path).parent() { if let Some(parent) = Path::new(path).parent() {
fs::create_dir_all(parent)?; fs::create_dir_all(parent)?;
} }
@@ -72,7 +71,6 @@ impl FileManager {
} }
fn extract_filename_from_path(full_path: &str) -> String { fn extract_filename_from_path(full_path: &str) -> String {
// Split on both forward slashes and backslashes to handle Windows and Unix paths
full_path full_path
.split(['/', '\\']) .split(['/', '\\'])
.next_back() .next_back()
@@ -124,6 +122,7 @@ impl StreamProcessor {
self.received = true; self.received = true;
return Ok(true); // Skip this data chunk return Ok(true); // Skip this data chunk
} }
self.received = true;
Ok(false) Ok(false)
} }
@@ -202,27 +201,27 @@ impl DownloadPeer {
self.no_pierce self.no_pierce
); );
if self.no_pierce { if self.no_pierce {
let message = MessageFactory::build_peer_init_message( // let message = MessageFactory::build_peer_init_message(
&self.own_username, // &self.own_username,
super::ConnectionType::F, // super::ConnectionType::F,
self.token, // self.token,
); // );
trace!( // trace!(
"[default_peer:{}] sending peer init, token: {} message: {:?} // "[download_peer:{}] sending peer init, token: {} message: {:?}
", // ",
self.token, // self.username,
self.username, // self.token,
message.get_buffer(), // message.get_buffer(),
); // );
stream.write_all(&message.get_buffer())?; // stream.write_all(&message.get_buffer())?;
stream.flush()?; // stream.flush()?;
} else { } else {
let message = let message =
MessageFactory::build_pierce_firewall_message(self.token); MessageFactory::build_pierce_firewall_message(self.token);
stream.write_all(&message.get_buffer())?; stream.write_all(&message.get_buffer())?;
trace!( trace!(
"[default_peer:{}] sending pierce firewall message token: {}: {:?}", "[download_peer:{}] sending pierce firewall message token: {}: {:?}",
self.username, self.username,
self.token, self.token,
&message.get_buffer() &message.get_buffer()
@@ -238,9 +237,13 @@ impl DownloadPeer {
mut download: Option<Download>, mut download: Option<Download>,
stream: Option<TcpStream>, stream: Option<TcpStream>,
) -> Result<(Download, String), io::Error> { ) -> Result<(Download, String), io::Error> {
if stream.is_none() { trace!(
trace!("[download_peer:{}] stream is none", self.username); "[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(|| { let mut stream = stream.unwrap_or_else(|| {
self.establish_connection().unwrap_or_else(|e| { self.establish_connection().unwrap_or_else(|e| {
panic!("Failed to establish connection: {:?}", e) panic!("Failed to establish connection: {:?}", e)
@@ -260,9 +263,13 @@ impl DownloadPeer {
self.username self.username
); );
if download.is_some() {
stream
.write_all(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])?;
}
loop { loop {
match stream.read(&mut read_buffer) { match stream.read(&mut read_buffer) {
// connection closed
Ok(0) => { Ok(0) => {
trace!( trace!(
"[download_peer:{}] connection closed by peer. bytes read: {}", "[download_peer:{}] connection closed by peer. bytes read: {}",
@@ -274,10 +281,7 @@ impl DownloadPeer {
Ok(bytes_read) => { Ok(bytes_read) => {
let data = &read_buffer[..bytes_read]; let data = &read_buffer[..bytes_read];
if !self.no_pierce if !self.no_pierce && !processor.received {
&& !processor.received
&& processor.handle_pierce_token(data, &mut stream)?
{
let token = data.get(0..4).unwrap(); let token = data.get(0..4).unwrap();
let token_u32 = u32::from_le_bytes( let token_u32 = u32::from_le_bytes(
token token
@@ -285,11 +289,15 @@ impl DownloadPeer {
.unwrap_or_else(|_| panic!("[download_peer:{}] slice with incorrect length", self.username)), .unwrap_or_else(|_| panic!("[download_peer:{}] slice with incorrect length", self.username)),
); );
trace!( trace!(
"[download_peer:{}] received token_u32: {:?} - token: {:?}", "[download_peer:{}] got token: {} from data chunk",
self.username, self.username,
token_u32, token_u32
token
); );
processor.received = true;
stream.write_all(&[
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
])?;
let read = client_context.read().unwrap(); let read = client_context.read().unwrap();
@@ -326,7 +334,7 @@ impl DownloadPeer {
processor.process_data_chunk(data); processor.process_data_chunk(data);
if !processor.should_continue(Some( if !processor.should_continue(Some(
download.clone().unwrap().size as usize, download.as_ref().unwrap().size as usize,
)) { )) {
break; break;
} }
@@ -340,6 +348,23 @@ impl DownloadPeer {
"[download_peer:{}] finished reading data from peer", "[download_peer:{}] finished reading data from peer",
self.username 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 read = client_context.read().unwrap();
let tokens = read.download_tokens.keys().collect::<Vec<_>>(); let tokens = read.download_tokens.keys().collect::<Vec<_>>();
+45 -15
View File
@@ -7,6 +7,7 @@ use crate::client::{ClientContext, ClientOperation};
use crate::message::MessageReader; use crate::message::MessageReader;
use crate::peer::{ConnectionType, DefaultPeer, DownloadPeer, Peer}; use crate::peer::{ConnectionType, DefaultPeer, DownloadPeer, Peer};
use crate::types::Download;
use crate::{debug, error, info, trace, DownloadStatus}; use crate::{debug, error, info, trace, DownloadStatus};
pub struct Listen {} 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}"); debug!("[listener:{peer_ip}:{peer_port}] peerInit (0) username: {username} connection_type: {connection_type} token: {token}");
let peer = Peer::new( let peer = Peer::new(
format!("{}-direct", username.clone()), username.clone(),
ConnectionType::P, ConnectionType::P,
peer_ip.clone(), peer_ip.clone(),
peer_port.into(), peer_port.into(),
@@ -78,32 +79,59 @@ impl Listen {
drop(default_peer); drop(default_peer);
} else if connection_type == "F" { } 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!( trace!(
"[client] DownloadFromPeer token: {} peer: {:?}", "[client] DownloadFromPeer token: {} peer: {:?}",
token, token,
peer peer
); );
thread::spawn(move || { let mut download: Option<Download> = None;
if reader.buffer_len() > 0 {
let buffer = reader.get_buffer();
trace!("[listener:{peer_ip}:{peer_port}] reader buffer has {} bytes. {:?}", buffer.len(), buffer);
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::<Vec<u32>>();
trace!("[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}", token_u32, download_tokens );
}
}
let download_peer = DownloadPeer::new( let download_peer = DownloadPeer::new(
format!("{}-direct", username.clone()), format!("{}:direct", username.clone()),
peer.host.clone(), peer.host.clone(),
peer.port, peer.port,
token, token,
false, true,
own_username, own_username,
); );
match download_peer.download_file( match download_peer.download_file(
client_context_clone, client_context,
maybe_download, download,
Some(stream), Some(stream),
) { ) {
Ok((download, filename)) => { Ok((download, filename)) => {
@@ -111,7 +139,10 @@ impl Listen {
.sender .sender
.send(DownloadStatus::Completed) .send(DownloadStatus::Completed)
.unwrap(); .unwrap();
info!("Successfully downloaded {} bytes to {}", download.size, filename); info!(
"Successfully downloaded {} bytes to {}",
download.size, filename
);
} }
Err(e) => { Err(e) => {
error!( error!(
@@ -120,7 +151,6 @@ impl Listen {
); );
} }
} }
});
} else { } else {
debug!( debug!(
"[listener:{peer_ip}:{peer_port}] connection type is not P or F" "[listener:{peer_ip}:{peer_port}] connection type is not P or F"