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