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
+6 -9
View File
@@ -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<bool> {
// 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
+4
View File
@@ -67,6 +67,10 @@ impl MessageReader {
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>> {
let bytes_read = self.buffer.len();
if bytes_read < 4 {
+2 -1
View File
@@ -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();
}
+54 -29
View File
@@ -49,7 +49,6 @@ impl FileManager {
}
fn create_temp_file(path: &str) -> Result<BufWriter<File>, 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<Download>,
stream: Option<TcpStream>,
) -> 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::<Vec<_>>();
+63 -33
View File
@@ -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<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);
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::<Vec<u32>>();
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"