dumbgoy/soulseek-rs-modded
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix: set nodelay on socket

Browse Source
This commit is contained in:
Michel de Graaf
2025-10-13 10:12:04 +02:00
parent cd3dca40c5
commit 9ccc4fa1e9
9 changed files with 343 additions and 249 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/client.rs
+65 -88
View File
@@ -1,12 +1,10 @@
use crate::{ use crate::{
error::{Result, SoulseekRs}, error::{Result, SoulseekRs},
peer::{ peer::{ConnectionType, DefaultPeer, DownloadPeer, NewPeer, Peer},
listen::Listen, ConnectionType, DefaultPeer, DownloadPeer, NewPeer,
Peer,
},
server::{PeerAddress, Server, ServerOperation}, server::{PeerAddress, Server, ServerOperation},
types::{Download, FileSearchResult}, types::{Download, FileSearchResult},
utils::{md5, thread_pool::ThreadPool}, utils::{md5, thread_pool::ThreadPool},
Transfer,
}; };
use std::{ use std::{
collections::HashMap, collections::HashMap,
@@ -31,6 +29,7 @@ pub enum ClientOperation {
PeerDisconnected(String), PeerDisconnected(String),
PierceFireWall(Peer), PierceFireWall(Peer),
DownloadFromPeer(u32, Peer), DownloadFromPeer(u32, Peer),
UpdateDownloadTokens(Transfer, String),
GetPeerAddressResponse { GetPeerAddressResponse {
username: String, username: String,
host: String, host: String,
@@ -39,12 +38,12 @@ pub enum ClientOperation {
obfuscated_port: u16, obfuscated_port: u16,
}, },
} }
struct ClientContext { pub struct ClientContext {
peers: HashMap<String, DefaultPeer>, peers: HashMap<String, DefaultPeer>,
sender: Option<Sender<ClientOperation>>, sender: Option<Sender<ClientOperation>>,
server_sender: Option<Sender<crate::server::ServerOperation>>, server_sender: Option<Sender<crate::server::ServerOperation>>,
search_results: Vec<FileSearchResult>, search_results: Vec<FileSearchResult>,
downloads: HashMap<u32, Download>, pub download_tokens: HashMap<u32, Download>,
thread_pool: ThreadPool, thread_pool: ThreadPool,
} }
impl ClientContext { impl ClientContext {
@@ -54,7 +53,7 @@ impl ClientContext {
sender: None, sender: None,
server_sender: None, server_sender: None,
search_results: Vec::new(), search_results: Vec::new(),
downloads: HashMap::new(), download_tokens: HashMap::new(),
thread_pool: ThreadPool::new(MAX_THREADS), thread_pool: ThreadPool::new(MAX_THREADS),
} }
} }
@@ -196,7 +195,7 @@ impl Client {
}; };
let mut context = self.context.lock().unwrap(); let mut context = self.context.lock().unwrap();
context.downloads.insert(token, download.clone()); context.download_tokens.insert(token, download.clone());
let download_initiated = context let download_initiated = context
.peers .peers
.get(&username) .get(&username)
@@ -238,16 +237,6 @@ impl Client {
if let Ok(operation) = reader.recv() { if let Ok(operation) = reader.recv() {
match operation { match operation {
ClientOperation::ConnectToPeer(peer) => { ClientOperation::ConnectToPeer(peer) => {
match peer.connection_type {
ConnectionType::P => (),
ConnectionType::F => {
debug!("Peer with F {:?}", peer)
}
ConnectionType::D => {
error!("ConnectionType::D not implemented for peer: {}", peer.username);
}
};
Self::connect_to_peer( Self::connect_to_peer(
peer, peer,
client_context.clone(), client_context.clone(),
@@ -276,20 +265,12 @@ impl Client {
); );
} }
ClientOperation::DownloadFromPeer(token, peer) => { ClientOperation::DownloadFromPeer(token, peer) => {
// {
// let mut context = client_context.lock().unwrap();
// if let Some(default_peer) =
// context.peers.remove(&peer.username)
// {
// default_peer.disconnect();
// }
// }
let maybe_download = { let maybe_download = {
let client_context = client_context.lock().unwrap(); let client_context = client_context.lock().unwrap();
client_context.downloads.get(&token).cloned() client_context.download_tokens.get(&token).cloned()
}; };
let own_username = own_username.clone(); let own_username = own_username.clone();
let client_context_clone = client_context.clone();
match maybe_download { match maybe_download {
Some(download) => { Some(download) => {
@@ -305,24 +286,25 @@ impl Client {
let filename: Option<&str> = let filename: Option<&str> =
download.filename.split('\\').last(); download.filename.split('\\').last();
match filename { match filename {
Some(filename) => { Some(filename) => {
match download_peer.download_file( match download_peer.download_file(
Some(download.size as usize), client_context_clone,
Some(format!("/tmp/{}", filename)), Some(download.size as usize),
) { Some(format!("/tmp/{}", filename)),
Ok(bytes) => { ) {
info!("Successfully downloaded {} bytes to /tmp/{}", bytes, filename); Ok((bytes, filename)) => {
} info!("Successfully downloaded {} bytes to /tmp/{}", bytes, filename);
Err(e) => { }
error!( Err(e) => {
"Failed to download file '{}' from {}:{} (token: {}) - Error: {}", error!(
filename, peer.host, peer.port, download.token, e "Failed to download file '{}' from {}:{} (token: {}) - Error: {}",
); filename, peer.host, peer.port, download.token, e
} );
} }
} }
None => error!("Cant find filename to save download: {:?}", download.filename), }
} None => error!("Cant find filename to save download: {:?}", download.filename),
}
}); });
} }
None => { None => {
@@ -384,9 +366,9 @@ impl Client {
obfuscated_port, obfuscated_port,
} => { } => {
debug!( debug!(
"Received peer address for {}: {}:{} (obf_type: {}, obf_port: {})", "Received peer address for {}: {}:{} (obf_type: {}, obf_port: {})",
username, host, port, obfuscation_type, obfuscated_port username, host, port, obfuscation_type, obfuscated_port
); );
match client_context match client_context
.lock() .lock()
@@ -431,6 +413,20 @@ impl Client {
} }
} }
} }
ClientOperation::UpdateDownloadTokens(
transfer,
username,
) => {
client_context.lock().unwrap().download_tokens.insert(
transfer.token,
Download {
username,
filename: transfer.filename,
token: transfer.token,
size: transfer.size,
},
);
}
} }
} }
}); });
@@ -445,11 +441,11 @@ impl Client {
let client_context2 = client_context.clone(); let client_context2 = client_context.clone();
let unlocked_context = client_context.lock().unwrap(); let unlocked_context = client_context.lock().unwrap();
trace!("[client] connect_to_peer: {}", peer.username);
if let Some(sender) = unlocked_context.sender.clone() { if let Some(sender) = unlocked_context.sender.clone() {
if !unlocked_context.peers.contains_key(&peer.username) { let peer_clone = peer.clone();
let peer_clone = peer.clone(); let sender_clone = sender;
let sender_clone = sender; unlocked_context.thread_pool.execute(move || {
unlocked_context.thread_pool.execute(move || {
trace!("[client] connecting to {}, with connection_type: {}", peer.username, peer.connection_type); trace!("[client] connecting to {}, with connection_type: {}", peer.username, peer.connection_type);
match peer.connection_type { match peer.connection_type {
ConnectionType::P => { ConnectionType::P => {
@@ -480,53 +476,37 @@ impl Client {
} }
ConnectionType::F => { ConnectionType::F => {
let context = client_context2.lock().unwrap(); trace!("[client] downloading from: {}", peer.username);
let download =
context.downloads.get(&peer.token.unwrap());
match download {
Some(download) => {
let download_peer = DownloadPeer::new( let download_peer = DownloadPeer::new(
peer.username, peer.username,
peer.host, peer.host,
peer.port, peer.port,
peer.token.unwrap(), peer.token.unwrap(),
true, false,
own_username.clone(), own_username.clone(),
); );
let filename: Option<&str> = match download_peer
download.filename.split('\\').last(); .download_file(
match filename { client_context2.clone(),
Some(filename) => { None,
download_peer None
.download_file( ) {
Some( Ok((bytes, filename)) => {
download.size as usize, trace!("[client] downloaded {} bytes {:?} ", filename,bytes);
), }
Some(format!( Err(e) => {
"/tmp/{}", trace!("[client] failed to download: {}", e);
filename }
)),
)
.unwrap();
}
None => todo!(),
} }
} }
None => todo!(),
}
}
ConnectionType::D => { ConnectionType::D => {
error!("ConnectionType::D not implemented") error!("ConnectionType::D not implemented")
} }
} }
}); });
} else {
// debug!("[client] Peer already connected: {}", peer.username);
}
} else { } else {
error!("No sender found"); debug!("[client] Peer already connected: {}", peer.username);
} }
} }
fn pierce_firewall( fn pierce_firewall(
@@ -541,10 +521,7 @@ impl Client {
if let Some(token) = peer.token { if let Some(token) = peer.token {
match server_sender.send(ServerOperation::PierceFirewall(token)) match server_sender.send(ServerOperation::PierceFirewall(token))
{ {
Ok(_) => debug!( Ok(_) => (),
"Sent PierceFirewall message with token: {}",
token
),
Err(e) => { Err(e) => {
error!("Failed to send PierceFirewall message: {}", e) error!("Failed to send PierceFirewall message: {}", e)
} }
src/main.rs
+27 -20
View File
@@ -11,30 +11,37 @@ fn main() {
client.connect(); client.connect();
match client.login() { match client.login() {
Ok(_) => match client Ok(_) => {
.search("michel test file", Duration::from_secs(20)) match client.search("michel test file", Duration::from_secs(1)) {
{ Ok(results) => {
Ok(results) => { if let Some(result) =
if let Some(result) = results.iter().find(|r| !r.files.is_empty())
results.iter().find(|r| !r.files.is_empty()) {
{ let file = result.files[0].clone();
let file = result.files[0].clone(); match client.download(
match client.download(file.name, file.username, file.size) { file.name,
Ok(download_result) => { file.username,
println!("Download result: {:?}", download_result); file.size,
} ) {
Err(e) => { Ok(download_result) => {
eprintln!("Failed to download: {}", e); println!(
"Download result: {:?}",
download_result
);
}
Err(e) => {
eprintln!("Failed to download: {}", e);
}
} }
} else {
eprint!("No results")
} }
} else { }
eprint!("No results") Err(e) => {
eprintln!("Failed to search: {}", e);
} }
} }
Err(e) => { }
eprintln!("Failed to search: {}", e);
}
},
Err(e) => { Err(e) => {
eprintln!("Failed to login: {}", e); eprintln!("Failed to login: {}", e);
} }
src/message/peer/transfer_request.rs
-2
View File
@@ -1,7 +1,6 @@
use crate::{ use crate::{
message::{Message, MessageHandler}, message::{Message, MessageHandler},
peer::PeerOperation, peer::PeerOperation,
trace,
types::Transfer, types::Transfer,
}; };
use std::sync::mpsc::Sender; use std::sync::mpsc::Sender;
@@ -12,7 +11,6 @@ impl MessageHandler<PeerOperation> for TransferRequest {
40 40
} }
fn handle(&self, message: &mut Message, sender: Sender<PeerOperation>) { fn handle(&self, message: &mut Message, sender: Sender<PeerOperation>) {
trace!("transfer_response: {:?}", message.get_buffer());
let transfer = Transfer::new_from_message(message); let transfer = Transfer::new_from_message(message);
sender sender
src/message/server/connect_to_peer.rs
+1 -2
View File
@@ -3,14 +3,13 @@ use crate::peer::Peer;
use crate::server::ServerOperation; use crate::server::ServerOperation;
use std::sync::mpsc::Sender; use std::sync::mpsc::Sender;
pub struct ConnectToPeerHandler; pub struct ConnectToPeerHandler;
impl MessageHandler<ServerOperation> for ConnectToPeerHandler { impl MessageHandler<ServerOperation> for ConnectToPeerHandler {
fn get_code(&self) -> u8 { fn get_code(&self) -> u8 {
18 18
} }
fn handle(&self, message: &mut Message, sender: Sender<ServerOperation>) { fn handle(&self, message: &mut Message, sender: Sender<ServerOperation>) {
let peer = Peer::new_from_message(message); let peer = Peer::new_from_message(message);
println!("ConnectToPeerHandler: {:?}", peer);
println!("ConnectToPeerHandler: {:?}", peer.username);
sender.send(ServerOperation::ConnectToPeer(peer)).unwrap(); sender.send(ServerOperation::ConnectToPeer(peer)).unwrap();
} }
} }
src/message/server/message_factory.rs
+2 -2
View File
@@ -106,12 +106,12 @@ impl MessageFactory {
.write_int32(41) .write_int32(41)
.write_int32(transfer.token) .write_int32(transfer.token)
.write_bool(true) .write_bool(true)
.write_int64(transfer.size) // .write_int64(transfer.size)
.clone() .clone()
} }
pub fn build_pierce_firewall_message(token: u32) -> Message { pub fn build_pierce_firewall_message(token: u32) -> Message {
Message::new() Message::new()
.write_int32(0) // PierceFirewall message code .write_int8(0) // PierceFirewall message code
.write_int32(token) .write_int32(token)
.clone() .clone()
} }
src/peer/default_peer.rs
+55 -39
View File
@@ -12,7 +12,7 @@ use std::sync::mpsc;
use std::sync::mpsc::{Receiver, Sender}; use std::sync::mpsc::{Receiver, Sender};
use std::thread::{self, JoinHandle}; use std::thread::{self, JoinHandle};
use crate::client::ClientOperation; use crate::client::{self, ClientOperation};
use crate::peer::Peer; use crate::peer::Peer;
use crate::{debug, error, trace, warn}; use crate::{debug, error, trace, warn};
use std::io::{self, Write}; use std::io::{self, Write};
@@ -78,6 +78,7 @@ impl DefaultPeer {
stream.set_read_timeout(Some(Duration::from_secs(5)))?; stream.set_read_timeout(Some(Duration::from_secs(5)))?;
stream.set_write_timeout(Some(Duration::from_secs(5)))?; stream.set_write_timeout(Some(Duration::from_secs(5)))?;
stream.set_nodelay(true)?;
trace!( trace!(
"[default_peer:{}] connect_with_socket: direct", "[default_peer:{}] connect_with_socket: direct",
@@ -107,6 +108,7 @@ impl DefaultPeer {
stream.set_read_timeout(Some(Duration::from_secs(5)))?; stream.set_read_timeout(Some(Duration::from_secs(5)))?;
stream.set_write_timeout(Some(Duration::from_secs(5)))?; stream.set_write_timeout(Some(Duration::from_secs(5)))?;
stream.set_nodelay(true)?;
if let Some(token) = self.peer.token { if let Some(token) = self.peer.token {
stream stream
@@ -181,34 +183,41 @@ impl DefaultPeer {
} }
} }
match buffered_reader.extract_message() { // Extract all available messages from buffer
Ok(Some(mut message)) => { let mut should_terminate = false;
trace!( loop {
"[default_peer:{:?}] ← {:?}", match buffered_reader.extract_message() {
peer.username, Ok(Some(mut message)) => {
message trace!(
.get_message_name( "[default_peer:{}] ← {:?}",
MessageType::Peer, peer.username,
message.get_message_code() as u32 message
) .get_message_name(
.map_err(|e| e.to_string()) MessageType::Peer,
); message.get_message_code() as u32
dispatcher.dispatch(&mut message) )
.map_err(|e| e.to_string())
);
dispatcher.dispatch(&mut message)
}
Err(e) => {
warn!(
"Error extracting message in default peer: {}. Terminating read loop.",
e
);
let _ = client_channel_for_read.send(
ClientOperation::PeerDisconnected(
peer.username.clone(),
),
);
should_terminate = true;
break;
}
Ok(None) => break,
} }
Err(e) => { }
warn!( if should_terminate {
"Error extracting message in default peer: {}. Terminating read loop.", break;
e
);
let _ = client_channel_for_read.send(
ClientOperation::PeerDisconnected(
peer.username.clone(),
),
);
break;
}
Ok(None) => continue,
} }
} }
})); }));
@@ -222,7 +231,7 @@ impl DefaultPeer {
Ok(operation) => match operation { Ok(operation) => match operation {
PeerOperation::SendMessage(message) => { PeerOperation::SendMessage(message) => {
trace!( trace!(
"[default_peer:{:?}] ➡ {:?} - {:?}", "[default_peer:{}] ➡ {:?} - {:?}",
peer_username, peer_username,
message message
.get_message_name( .get_message_name(
@@ -235,7 +244,9 @@ impl DefaultPeer {
) )
) )
.map_err(|e| e.to_string()), .map_err(|e| e.to_string()),
message u32::from_le_bytes(
message.get_slice(0, 4).try_into().unwrap()
)
); );
if let Err(e) = if let Err(e) =
@@ -249,6 +260,7 @@ impl DefaultPeer {
); );
break; break;
} }
write_stream.flush().unwrap();
} }
PeerOperation::FileSearchResult(file_search) => { PeerOperation::FileSearchResult(file_search) => {
client_channel client_channel
@@ -257,15 +269,26 @@ impl DefaultPeer {
} }
PeerOperation::TransferRequest(transfer) => { PeerOperation::TransferRequest(transfer) => {
debug!( debug!(
"[default_peer:{:}] TransferRequest for {}", "[default_peer:{}] TransferRequest for {}",
peer_username, transfer.token peer_username, transfer.token
); );
client_channel
.send(ClientOperation::UpdateDownloadTokens(
transfer.clone(),
peer_clone.username.clone(),
))
.unwrap();
let transfer_response = let transfer_response =
MessageFactory::build_transfer_response_message( MessageFactory::build_transfer_response_message(
transfer.clone(), transfer.clone(),
); );
trace!(
"[default_peer:{}] TransferResponse for {:?}",
peer_username,
transfer_response.get_buffer()
);
if let Some(sender) = peer_channel.clone() { if let Some(sender) = peer_channel.clone() {
sender sender
.send(PeerOperation::SendMessage( .send(PeerOperation::SendMessage(
@@ -273,13 +296,6 @@ impl DefaultPeer {
)) ))
.unwrap(); .unwrap();
} }
debug!(
"[default_peer:{:}] TransferResponse for {}",
peer_username, transfer.token
);
break;
} }
PeerOperation::TransferResponse { PeerOperation::TransferResponse {
token, token,
@@ -301,7 +317,7 @@ impl DefaultPeer {
); );
} }
} else { } else {
debug!("[default_peer:{:}] Transfer allowed, ready to connect with token {:}",peer_username, token); debug!("[default_peer:{}] Transfer allowed, ready to connect with token {:}",peer_username, token);
client_channel client_channel
.send(ClientOperation::DownloadFromPeer( .send(ClientOperation::DownloadFromPeer(
token, token,
src/peer/download_peer.rs
+131 -35
View File
@@ -1,10 +1,12 @@
use crate::client::ClientContext;
use crate::message::server::MessageFactory; use crate::message::server::MessageFactory;
use crate::trace;
use std::fs::{self, File}; use std::fs::{self, File};
use std::io::{self, BufWriter, Read, Write}; use std::io::{self, BufWriter, Read, Write};
use std::net::TcpStream; use std::net::TcpStream;
use std::net::ToSocketAddrs; use std::net::ToSocketAddrs;
use std::path::Path; use std::path::Path;
use std::thread::sleep; use std::sync::{Arc, Mutex};
use std::time::Duration; use std::time::Duration;
#[derive(Debug)] #[derive(Debug)]
@@ -101,7 +103,6 @@ impl StreamProcessor {
if let Some(ref mut w) = writer { if let Some(ref mut w) = writer {
w.write_all(data)?; w.write_all(data)?;
// Flush every 64KB to ensure data reaches disk
if self.total_bytes % (64 * 1024) == 0 { if self.total_bytes % (64 * 1024) == 0 {
w.flush()?; w.flush()?;
} }
@@ -166,6 +167,7 @@ impl DownloadPeer {
stream.set_read_timeout(Some(Duration::from_secs(30)))?; stream.set_read_timeout(Some(Duration::from_secs(30)))?;
stream.set_write_timeout(Some(Duration::from_secs(5)))?; stream.set_write_timeout(Some(Duration::from_secs(5)))?;
stream.set_nodelay(true)?;
Ok(stream) Ok(stream)
} }
@@ -174,31 +176,55 @@ impl DownloadPeer {
&self, &self,
stream: &mut TcpStream, stream: &mut TcpStream,
) -> Result<(), io::Error> { ) -> Result<(), io::Error> {
trace!(
"[download_peer:{}] performing handshake no_pierce: {}",
self.username,
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,
); );
stream.write_all(&message.get_data())?; trace!(
sleep(Duration::from_millis(1000)); "[default_peer:{}] sending peer init, token: {} message: {:?}
stream ",
.write_all(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])?; self.token,
self.username,
message.get_buffer(),
);
stream.write_all(&message.get_buffer())?;
stream.flush()?;
// stream
// .write_all(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])?;
// stream.flush()?;
//
} else { } else {
stream.write_all( stream.write_all(
&MessageFactory::build_pierce_firewall_message(self.token) &MessageFactory::build_pierce_firewall_message(self.token)
.get_data(), .get_buffer(),
)?; )?;
trace!(
"[default_peer:{}] sending pierce firewall message token: {}: {:?}",
self.username,
self.token,
MessageFactory::build_pierce_firewall_message(self.token)
.get_buffer()
);
stream.flush()?;
} }
Ok(()) Ok(())
} }
pub fn download_file( pub fn download_file(
self, self,
expected_size: Option<usize>, client_context: Arc<Mutex<ClientContext>>,
mut expected_size: Option<usize>,
output_path: Option<String>, output_path: Option<String>,
) -> Result<usize, io::Error> { ) -> Result<(usize, String), io::Error> {
let mut stream = self.establish_connection()?; let mut stream = self.establish_connection()?;
trace!("[download_peer:{}] connected", self.username);
// Setup file management if output path is provided // Setup file management if output path is provided
let (paths, mut writer) = if let Some(output_path) = output_path { let (paths, mut writer) = if let Some(output_path) = output_path {
@@ -209,33 +235,92 @@ impl DownloadPeer {
); );
let writer = let writer =
Some(FileManager::create_temp_file(&paths.incomplete_path)?); Some(FileManager::create_temp_file(&paths.incomplete_path)?);
(Some(paths), writer) (paths, writer)
} else { } else {
(None, None) let paths = FileManager::create_download_paths(
None,
&self.own_username,
self.token,
);
let writer =
Some(FileManager::create_temp_file(&paths.incomplete_path)?);
trace!(
"[download_peer:{}] No output path provided, downloading to temp file: {}",
self.username,
paths.incomplete_path
);
(paths, writer)
}; };
// Perform handshake
self.perform_handshake(&mut stream)?; self.perform_handshake(&mut stream)?;
trace!("[download_peer:{}] handshake completed", self.username);
// Stream data
let mut processor = StreamProcessor::new(self.no_pierce, self.token); let mut processor = StreamProcessor::new(self.no_pierce, self.token);
let mut read_buffer = [0u8; 8192]; let mut read_buffer = [1u8; 8192];
trace!(
"[download_peer:{}] Starting to read data from peer",
self.username
);
loop { loop {
match stream.read(&mut read_buffer) { match stream.read(&mut read_buffer) {
Ok(0) => break, // Connection closed Ok(0) => {
trace!(
"[download_peer:{}] connection closed by peer",
self.username
);
break; // Connection closed
}
Ok(bytes_read) => { Ok(bytes_read) => {
let data = &read_buffer[..bytes_read]; let data = &read_buffer[..bytes_read];
// trace!(
// "[download_peer:{}] received {} bytes",
// self.username,
// bytes_read
// );
if !self.no_pierce
&& !processor.received
&& processor.handle_pierce_token(data, &mut stream)?
{
let token = data.get(0..4).unwrap();
let token_u32 =
u32::from_le_bytes(token.try_into().unwrap());
trace!(
"[download_peer:{}] received token: {:?} ",
self.username,
token
);
match client_context
.lock()
.unwrap()
.download_tokens
.get(&token_u32)
.cloned()
{
Some(download) => {
trace!(
"[download_peer:{}] got download info for token: {}",
self.username,
token_u32
);
expected_size = Some(download.size as usize);
}
None => {
trace!(
"[download_peer:{}] no download info for token: {}",
self.username,
self.token
);
}
}
// Handle pierce token extraction
if processor.handle_pierce_token(data, &mut stream)? {
continue; // Skip this data chunk continue; // Skip this data chunk
} }
// Process file data
processor.process_data_chunk(data, &mut writer)?; processor.process_data_chunk(data, &mut writer)?;
// Check completion
if !processor.should_continue(expected_size) { if !processor.should_continue(expected_size) {
break; break;
} }
@@ -244,38 +329,41 @@ impl DownloadPeer {
continue; continue;
} }
Err(e) => { Err(e) => {
// Clean up incomplete file on error
drop(writer); drop(writer);
if let Some(ref paths) = paths { FileManager::cleanup_on_error(Some(&paths.incomplete_path));
FileManager::cleanup_on_error(Some(
&paths.incomplete_path,
));
}
return Err(e); return Err(e);
} }
} }
} }
trace!(
"[download_peer:{}] finished reading data from peer",
self.username
);
// Finalize download
if let Some(mut w) = writer { if let Some(mut w) = writer {
w.flush()?; w.flush()?;
drop(w); drop(w);
if let Some(paths) = paths { trace!("[download_peer:{}] finalizing download", self.username);
FileManager::finalize_download( FileManager::finalize_download(
&paths.incomplete_path, &paths.incomplete_path,
&paths.final_path, &paths.final_path,
)?; )?;
}
} }
trace!(
"[download_peer:{}] download completed successfully: {} bytes",
self.username,
processor.total_bytes
);
Ok(processor.total_bytes) Ok((processor.total_bytes, paths.final_path.clone()))
} }
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::{DownloadPeer, FileManager, StreamProcessor}; use super::{DownloadPeer, FileManager, StreamProcessor};
use crate::client::ClientContext;
use std::fs; use std::fs;
use std::io::{Read, Write}; use std::io::{Read, Write};
use std::net::TcpListener; use std::net::TcpListener;
@@ -327,7 +415,10 @@ mod tests {
false, false,
"own_username".to_string(), "own_username".to_string(),
); );
let _ = download_peer.download_file(None, None).unwrap(); let dummy_context = Arc::new(Mutex::new(ClientContext::new()));
let _ = download_peer
.download_file(dummy_context, None, None)
.unwrap();
// Give the client time to send messages // Give the client time to send messages
thread::sleep(Duration::from_millis(10)); thread::sleep(Duration::from_millis(10));
@@ -349,7 +440,10 @@ mod tests {
true, // no_pierce = true, should send init message true, // no_pierce = true, should send init message
"own_username".to_string(), "own_username".to_string(),
); );
let _ = download_peer.download_file(None, None).unwrap(); let dummy_context = Arc::new(Mutex::new(ClientContext::new()));
let _ = download_peer
.download_file(dummy_context, None, None)
.unwrap();
// Give the client time to send messages // Give the client time to send messages
thread::sleep(Duration::from_millis(10)); thread::sleep(Duration::from_millis(10));
@@ -418,7 +512,9 @@ mod tests {
"test_user".to_string(), "test_user".to_string(),
); );
let dummy_context = Arc::new(Mutex::new(ClientContext::new()));
let result = download_peer.download_file( let result = download_peer.download_file(
dummy_context,
Some(test_data.len()), Some(test_data.len()),
Some("test_download.mp3".to_string()), Some("test_download.mp3".to_string()),
); );
src/peer/listen.rs
+40 -36
View File
@@ -34,48 +34,52 @@ impl Listen {
} }
} }
match buffered_reader.extract_message() { // Extract all available messages from buffer
Ok(Some(mut message)) => { loop {
message.set_pointer(4); match buffered_reader.extract_message() {
let message_code = message.read_int8(); Ok(Some(mut message)) => {
message.set_pointer(4);
let message_code = message.read_int8();
trace!( trace!(
"[listener] Received message with code: {}", "[listener] Received message with code: {}",
message_code message_code
); );
match message_code { match message_code {
0 => { 0 => {
let token = message.read_string(); let token = message.read_string();
debug!("[listener] received Pierce Firewall, token: {}", token); debug!("[listener] received Pierce Firewall, token: {}", token);
} }
1 => { 1 => {
let tcp_stream = let tcp_stream =
read_stream.try_clone().unwrap(); read_stream.try_clone().unwrap();
let new_peer = NewPeer::new_from_message( let new_peer = NewPeer::new_from_message(
&mut message, &mut message,
tcp_stream, tcp_stream,
); );
trace!("[listener] Handling peer init message: {:?}",new_peer); trace!("[listener] Handling peer init message: {:?}",new_peer);
client_sender client_sender
.send(ClientOperation::NewPeer(new_peer)) .send(ClientOperation::NewPeer(new_peer))
.unwrap(); .unwrap();
stop = true; stop = true;
debug!("[listener] stop"); debug!("[listener] stop");
} }
_ => { _ => {
debug!( debug!(
"[listener] unknown message_code: {}", "[listener] unknown message_code: {}",
message_code message_code
) )
}
} }
} }
Err(e) => {
warn!("[listener] Error extracting message: {}", e);
break;
}
Ok(None) => break,
} }
Err(e) => {
warn!("[listener] Error extracting message: {}", e)
}
Ok(None) => continue,
} }
} }
src/server.rs
+22 -25
View File
@@ -254,6 +254,7 @@ impl Server {
stream.set_read_timeout(Some(Duration::from_secs(5)))?; stream.set_read_timeout(Some(Duration::from_secs(5)))?;
stream.set_write_timeout(Some(Duration::from_secs(5)))?; stream.set_write_timeout(Some(Duration::from_secs(5)))?;
stream.set_nodelay(true)?;
let mut read_stream = stream.try_clone()?; let mut read_stream = stream.try_clone()?;
let mut write_stream = stream.try_clone()?; let mut write_stream = stream.try_clone()?;
@@ -292,9 +293,6 @@ impl Server {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
continue continue
} }
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
continue
}
Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { Err(ref e) if e.kind() == io::ErrorKind::TimedOut => {
debug!("Read operation timed out"); debug!("Read operation timed out");
continue; continue;
@@ -305,24 +303,27 @@ impl Server {
} }
} }
match buffered_reader.extract_message() { loop {
Ok(Some(mut message)) => { match buffered_reader.extract_message() {
trace!( Ok(Some(mut message)) => {
"[server] ← {:?} - {}", trace!(
message "[server] ← {:?} - {}",
.get_message_name( message
MessageType::Server, .get_message_name(
message.get_message_code_u32() MessageType::Server,
) message.get_message_code_u32()
.map_err(|e| e.to_string()), )
message.get_message_code_u32() .map_err(|e| e.to_string()),
); message.get_message_code_u32()
dispatcher.dispatch(&mut message) );
dispatcher.dispatch(&mut message)
}
Err(e) => {
warn!("Error extracting message: {}", e);
break;
}
Ok(None) => break,
} }
Err(e) => {
warn!("Error extracting message: {}", e)
}
Ok(None) => continue,
} }
} }
}); });
@@ -344,7 +345,7 @@ impl Server {
Some(ClientOperation::ConnectToPeer(peer)) Some(ClientOperation::ConnectToPeer(peer))
} }
ConnectionType::F => { ConnectionType::F => {
Some(ClientOperation::PierceFireWall(peer)) Some(ClientOperation::ConnectToPeer(peer))
} }
ConnectionType::D => None, ConnectionType::D => None,
} { } {
@@ -367,10 +368,6 @@ impl Server {
); );
break; break;
} }
debug!(
"Sent PierceFirewall message with token: {}",
token
);
} }
ServerOperation::SendMessage(message) => { ServerOperation::SendMessage(message) => {
match write_stream.write_all(&message.get_buffer()) match write_stream.write_all(&message.get_buffer())