From cd3dca40c5a6e12e57d7ec74715cd1e041c82e9b Mon Sep 17 00:00:00 2001 From: Michel de Graaf Date: Thu, 31 Jul 2025 22:24:34 +0200 Subject: [PATCH] feat: rename module to soulseek-rs --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/client.rs | 391 ++++++++++++++++---- src/dispatcher.rs | 17 +- src/error.rs | 75 ++++ src/lib.rs | 2 + src/main.rs | 36 +- src/message/mod.rs | 34 +- src/message/peer/file_search_response.rs | 7 +- src/message/peer/mod.rs | 4 + src/message/peer/place_in_queue_response.rs | 22 ++ src/message/peer/transfer_request.rs | 2 + src/message/peer/transfer_response.rs | 27 ++ src/message/server/connect_to_peer.rs | 11 +- src/message/server/get_peer_address.rs | 37 ++ src/message/server/message_factory.rs | 44 ++- src/message/server/mod.rs | 3 +- src/peer/default_peer.rs | 244 ++++++++---- src/peer/download_peer.rs | 2 +- src/peer/listen.rs | 66 ++-- src/peer/mod.rs | 68 ++-- src/server.rs | 200 +++++----- src/types.rs | 24 +- src/utils/thread_pool.rs | 49 ++- src/utils/zlib.rs | 43 ++- test.md | 11 + todo.md | 8 + 27 files changed, 1072 insertions(+), 359 deletions(-) create mode 100644 src/error.rs create mode 100644 src/message/peer/place_in_queue_response.rs create mode 100644 src/message/peer/transfer_response.rs create mode 100644 src/message/server/get_peer_address.rs create mode 100644 test.md create mode 100644 todo.md diff --git a/Cargo.lock b/Cargo.lock index 4aaadf8..4bc772f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3,5 +3,5 @@ version = 4 [[package]] -name = "soulseek_download" +name = "soulseek-rs" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index d678d04..fbba1f9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "soulseek_download" +name = "soulseek-rs" version = "0.1.0" edition = "2021" diff --git a/src/client.rs b/src/client.rs index 31c530b..a26d757 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,37 +1,50 @@ use crate::{ - peer::{ConnectionType, DefaultPeer, Peer}, + error::{Result, SoulseekRs}, + peer::{ + listen::Listen, ConnectionType, DefaultPeer, DownloadPeer, NewPeer, + Peer, + }, server::{PeerAddress, Server, ServerOperation}, - types::{FileSearchResult, Transfer}, + types::{Download, FileSearchResult}, utils::{md5, thread_pool::ThreadPool}, }; use std::{ collections::HashMap, + net::TcpStream, sync::{ mpsc::{Receiver, Sender}, Mutex, }, - thread, + thread::{self, sleep}, }; use std::{ sync::{mpsc, Arc}, time::{Duration, Instant}, }; -use crate::{debug, error, info, trace, warn}; -const MAX_THREADS: usize = 100; +use crate::{debug, error, info, trace}; +const MAX_THREADS: usize = 50; pub enum ClientOperation { + NewPeer(NewPeer), ConnectToPeer(Peer), SearchResult(FileSearchResult), PeerDisconnected(String), - TransferRequest(Transfer), PierceFireWall(Peer), + DownloadFromPeer(u32, Peer), + GetPeerAddressResponse { + username: String, + host: String, + port: u32, + obfuscation_type: u32, + obfuscated_port: u16, + }, } struct ClientContext { peers: HashMap, sender: Option>, server_sender: Option>, search_results: Vec, - downloads: HashMap, + downloads: HashMap, thread_pool: ThreadPool, } impl ClientContext { @@ -61,12 +74,17 @@ impl Client { password: String, ) -> Self { crate::utils::logger::init(); + let context = Arc::new(Mutex::new(ClientContext::new())); + debug!( + "ThreadPool initialized with {} threads", + context.lock().unwrap().thread_pool.thread_count() + ); Self { address, username, password, server: None, - context: Arc::new(Mutex::new(ClientContext::new())), + context, } } @@ -78,7 +96,8 @@ impl Client { self.context.lock().unwrap().sender = Some(sender.clone()); - // self.read_form_channel(message_reader); + let client_sender = sender.clone(); + self.server = match Server::new(self.address.clone(), sender) { Ok(server) => { info!( @@ -87,13 +106,17 @@ impl Client { server.get_address().get_port() ); - // Store the server sender in the client context - self.context.lock().unwrap().server_sender = + // thread::spawn(move || { + // Listen::start(2234, client_sender.clone()); + // }); + let mut unlocked_context = self.context.lock().unwrap(); + unlocked_context.server_sender = Some(server.get_sender().clone()); Self::listen_to_client_operations( message_reader, self.context.clone(), + self.username.clone(), ); Some(server) } @@ -104,24 +127,18 @@ impl Client { }; } - pub fn login(&self) -> Result { + 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.unwrap() { + let result = server.login(&self.username, &self.password)?; + if result { Ok(true) } else { - Err(std::io::Error::new( - std::io::ErrorKind::Other, - "Error logging in", - )) + Err(SoulseekRs::AuthenticationFailed) } } else { - Err(std::io::Error::new( - std::io::ErrorKind::Other, - "Not connected to server", - )) + Err(SoulseekRs::NotConnected) } } @@ -137,78 +154,85 @@ impl Client { &self, query: &str, timeout: Duration, - ) -> Vec { + ) -> Result> { info!("Searching for {}", query); if let Some(server) = &self.server { let hash = md5::md5(query); - let token = u32::from_str_radix(&hash[0..5], 16).unwrap(); + let token = u32::from_str_radix(&hash[0..5], 16)?; server.file_search(token, query); } else { - warn!("Not connected to server"); + return Err(SoulseekRs::NotConnected); } let start = Instant::now(); loop { + sleep(Duration::from_millis(500)); if start.elapsed() >= timeout { break; } } - return self.context.lock().unwrap().search_results.clone(); + Ok(self.context.lock().unwrap().search_results.clone()) } pub fn download( &self, filename: String, username: String, - ) -> crate::types::DownloadResult { + size: u64, + ) -> Result { use crate::types::{DownloadResult, DownloadStatus}; use std::time::{Duration, Instant}; - println!("Downloading {} from {}", filename, username); let start_time = Instant::now(); - // Start the download request - let context = self.context.lock().unwrap(); + let hash = md5::md5(&filename); + let token = u32::from_str_radix(&hash[0..5], 16)?; + + let download = Download { + username: username.clone(), + filename: filename.clone(), + token, + size, + }; + + let mut context = self.context.lock().unwrap(); + context.downloads.insert(token, download.clone()); let download_initiated = context .peers .get(&username) - .map(|p| p.transfer_request(filename.clone())) + .map(|p| p.transfer_request(download.clone())) .is_some(); drop(context); - let timeout = Duration::from_secs(50); + let timeout = Duration::from_secs(150); let check_interval = Duration::from_millis(100); if !download_initiated { - return DownloadResult { + return Ok(DownloadResult { filename, username, status: DownloadStatus::Failed, elapsed_time: start_time.elapsed(), - }; + }); } - // Non-blocking wait loop while start_time.elapsed() < timeout { - // Check download status (for now just wait, actual download logic will be implemented later) std::thread::sleep(check_interval); - - // TODO: Check actual download progress here - // For now, we'll just wait for the timeout } - DownloadResult { + Ok(DownloadResult { filename, username, status: DownloadStatus::TimedOut, elapsed_time: start_time.elapsed(), - } + }) } fn listen_to_client_operations( reader: Receiver, client_context: Arc>, + own_username: String, ) { thread::spawn(move || loop { if let Ok(operation) = reader.recv() { @@ -219,10 +243,17 @@ impl Client { ConnectionType::F => { debug!("Peer with F {:?}", peer) } - ConnectionType::D => (), + ConnectionType::D => { + error!("ConnectionType::D not implemented for peer: {}", peer.username); + } }; - Self::connect_to_peer(peer, client_context.clone()) + Self::connect_to_peer( + peer, + client_context.clone(), + own_username.clone(), + None, + ) } ClientOperation::SearchResult(file_search) => { client_context @@ -237,55 +268,272 @@ impl Client { drop(peer); // Explicitly drop to trigger cleanup } } - ClientOperation::TransferRequest(transfer) => { - client_context + ClientOperation::PierceFireWall(peer) => { + Self::pierce_firewall( + peer, + client_context.clone(), + own_username.clone(), + ); + } + 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() + }; + let own_username = own_username.clone(); + + match maybe_download { + Some(download) => { + thread::spawn(move || { + let download_peer = DownloadPeer::new( + download.username.clone(), + peer.host.clone(), + peer.port, + download.token, + false, + own_username, + ); + 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), + } + }); + } + None => { + error!( + "Can't find download with token {:?}", + token + ); + } + }; + } + ClientOperation::NewPeer(new_peer) => { + if client_context .lock() .unwrap() - .downloads - .insert(transfer.token, transfer); + .peers + .contains_key(&new_peer.username) + { + debug!( + "Already connected to {}", + new_peer.username + ); + } else if let Some(server_sender) = + &client_context.lock().unwrap().server_sender + { + server_sender + .send(ServerOperation::GetPeerAddress( + new_peer.username.clone(), + )) + .unwrap(); + } + + let addr = new_peer.tcp_stream.peer_addr().unwrap(); + let host = addr.ip().to_string(); + let port: u32 = addr.port().into(); + + let peer = Peer { + username: new_peer.username.clone(), + connection_type: new_peer.connection_type, + host, + port, + token: Some(new_peer.token), + privileged: None, + obfuscated_port: None, + unknown: None, + }; + + Self::connect_to_peer( + peer, + client_context.clone(), + own_username.clone(), + Some(new_peer.tcp_stream), + ); } - ClientOperation::PierceFireWall(peer) => { - Self::pierce_firewall(peer, client_context.clone()); + ClientOperation::GetPeerAddressResponse { + username, + host, + port, + obfuscation_type, + obfuscated_port, + } => { + debug!( + "Received peer address for {}: {}:{} (obf_type: {}, obf_port: {})", + username, host, port, obfuscation_type, obfuscated_port + ); + + match client_context + .lock() + .unwrap() + .peers + .get(&username) + { + Some(_peer) => { + // don't know if i should update? and or reconnect the peer + // debug!( + // "existing peer: {:?}, new peer details: + // username: {}, + // host: {}, + // port: {} + // obfuscation_type: {} + // obfuscated_port: {}", + // peer, + // username, + // host, + // port, + // obfuscation_type, + // obfuscated_port, + // ); + } + None => { + let peer = Peer::new( + username, + ConnectionType::P, + host, + port, + None, + 0, + obfuscation_type.try_into().unwrap(), + obfuscated_port.try_into().unwrap(), + ); + Self::connect_to_peer( + peer, + client_context.clone(), + own_username.clone(), + None, + ); + } + } } } } }); } - fn connect_to_peer(peer: Peer, client_context: Arc>) { - let context = client_context.clone(); - let unlocked_context = context.lock().unwrap(); + fn connect_to_peer( + peer: Peer, + client_context: Arc>, + own_username: String, + stream: Option, + ) { + let client_context2 = client_context.clone(); + let unlocked_context = client_context.lock().unwrap(); - if let Some(sender) = &unlocked_context.sender { + 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.clone(); + let sender_clone = sender; unlocked_context.thread_pool.execute(move || { - let default_peer = - DefaultPeer::new(peer_clone, sender_clone); - match default_peer.connect() { - Ok(p) => { - let mut context = client_context.lock().unwrap(); - context.peers.insert(peer.username, p); + trace!("[client] connecting to {}, with connection_type: {}", peer.username, peer.connection_type); + match peer.connection_type { + ConnectionType::P => { + let default_peer = + DefaultPeer::new(peer_clone, sender_clone); + + let connect_result = match stream { + Some(s) => default_peer.connect_with_socket(s), + None => default_peer.connect() + }; + + match connect_result { + Ok(p) => { + trace!("[client] connected to: {}", peer.username); + client_context2.lock().unwrap().peers.insert(peer.username, p); + } + Err(e) => { + trace!( + "Can't connect to {:?} {:?}:{:?} {:?} - {:?}", + peer.username, + peer.host, + peer.port, + peer.connection_type, + e + ); + } + } } - Err(e) => { - trace!( - "Can't connect to {} {}:{} {:?} - {}", - peer.username, - peer.host, - peer.port, - peer.connection_type, - e - ); + + ConnectionType::F => { + let context = client_context2.lock().unwrap(); + let download = + context.downloads.get(&peer.token.unwrap()); + + match download { + Some(download) => { + let download_peer = DownloadPeer::new( + peer.username, + peer.host, + peer.port, + peer.token.unwrap(), + true, + 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!(), + } + } + None => todo!(), + } + } + ConnectionType::D => { + error!("ConnectionType::D not implemented") } } }); + } else { + // debug!("[client] Peer already connected: {}", peer.username); } } else { error!("No sender found"); } } - fn pierce_firewall(peer: Peer, client_context: Arc>) { + fn pierce_firewall( + peer: Peer, + client_context: Arc>, + own_username: String, + ) { debug!("Piercing firewall for peer: {:?}", peer); let context = client_context.lock().unwrap(); @@ -309,7 +557,6 @@ impl Client { } drop(context); - // Also try to connect to the peer directly - Self::connect_to_peer(peer, client_context); + Self::connect_to_peer(peer, client_context, own_username, None); } } diff --git a/src/dispatcher.rs b/src/dispatcher.rs index f780235..a362bbb 100644 --- a/src/dispatcher.rs +++ b/src/dispatcher.rs @@ -3,15 +3,23 @@ use std::sync::mpsc::Sender; use crate::warn; -/// Add `` to make it generic over any operation type. pub struct MessageDispatcher { + owner_name: String, sender: Sender, handlers: Handlers, } impl MessageDispatcher { - pub fn new(sender: Sender, handlers: Handlers) -> Self { - MessageDispatcher { sender, handlers } + pub fn new( + owner_name: String, + sender: Sender, + handlers: Handlers, + ) -> Self { + MessageDispatcher { + owner_name, + sender, + handlers, + } } pub fn dispatch(&self, message: &mut Message) { @@ -22,7 +30,8 @@ impl MessageDispatcher { handler.handle(message, self.sender.clone()); } else { warn!( - "No handler found for message code: {:?}", + "[{}:dispatcher] No handler found for message code: {}", + self.owner_name, message.get_message_code() ); } diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..74fddab --- /dev/null +++ b/src/error.rs @@ -0,0 +1,75 @@ +use std::{error::Error, fmt}; + +/// Custom error type for the Soulseek download library +#[derive(Debug)] +pub enum SoulseekRs { + /// Network-related errors (I/O, connection failures) + NetworkError(std::io::Error), + /// Authentication failed during login + AuthenticationFailed, + /// Error parsing messages or data + ParseError(String), + /// Operation timed out + Timeout, + /// Connection was closed unexpectedly + ConnectionClosed, + /// Invalid message format or content + InvalidMessage(String), + /// Server not connected + NotConnected, + /// Compression/decompression error + CompressionError(String), +} + +impl fmt::Display for SoulseekRs { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + SoulseekRs::NetworkError(err) => { + write!(f, "Network error: {}", err) + } + SoulseekRs::AuthenticationFailed => { + write!(f, "Authentication failed") + } + SoulseekRs::ParseError(msg) => write!(f, "Parse error: {}", msg), + SoulseekRs::Timeout => write!(f, "Operation timed out"), + SoulseekRs::ConnectionClosed => write!(f, "Connection closed"), + SoulseekRs::InvalidMessage(msg) => { + write!(f, "Invalid message: {}", msg) + } + SoulseekRs::NotConnected => write!(f, "Not connected to server"), + SoulseekRs::CompressionError(msg) => { + write!(f, "Compression error: {}", msg) + } + } + } +} + +impl Error for SoulseekRs { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + SoulseekRs::NetworkError(err) => Some(err), + _ => None, + } + } +} + +impl From for SoulseekRs { + fn from(err: std::io::Error) -> Self { + SoulseekRs::NetworkError(err) + } +} + +impl From for SoulseekRs { + fn from(err: std::num::ParseIntError) -> Self { + SoulseekRs::ParseError(format!("Integer parse error: {}", err)) + } +} + +impl From for SoulseekRs { + fn from(err: String) -> Self { + SoulseekRs::CompressionError(err) + } +} + +/// Result type alias for the Soulseek library +pub type Result = std::result::Result; diff --git a/src/lib.rs b/src/lib.rs index e5434c6..558d920 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,7 @@ // Core modules pub mod client; pub mod dispatcher; +pub mod error; pub mod message; pub mod peer; pub mod server; @@ -18,6 +19,7 @@ pub mod prelude { // Re-export commonly used types pub use client::Client; +pub use error::{Result, SoulseekRs}; pub use server::PeerAddress; pub use types::{ DownloadResult, DownloadStatus, File, FileSearchResult, Transfer, diff --git a/src/main.rs b/src/main.rs index a535044..fec9b5a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,24 +1,40 @@ fn main() { - use soulseek_download::{Client, PeerAddress}; + use soulseek_rs::{Client, PeerAddress}; use std::time::Duration; let mut client = Client::new( - PeerAddress::new(String::from("server.slsknet.org"), 2242), + // PeerAddress::new(String::from("server.slsknet.org"), 2242), + PeerAddress::new(String::from("127.0.0.1"), 2242), String::from("insane_in_the_brain3"), String::from("13375137"), ); client.connect(); match client.login() { - Ok(_) => { - let results = client.search("Fantazia", Duration::from_secs(10)); - println!("Search results: {} - {:?} ", results.len(), results); - if !results.is_empty() && !results[0].files.is_empty() { - let file = results[0].files[0].clone(); - let download_result = client.download(file.name, file.username); - println!("Download result: {:?}", download_result); + 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); + } + } + } else { + eprint!("No results") + } } - } + Err(e) => { + eprintln!("Failed to search: {}", e); + } + }, Err(e) => { eprintln!("Failed to login: {}", e); } diff --git a/src/message/mod.rs b/src/message/mod.rs index 48680d9..bc29443 100644 --- a/src/message/mod.rs +++ b/src/message/mod.rs @@ -5,11 +5,10 @@ mod message_reader; pub mod peer; pub mod server; -// Re-export commonly used items pub use handlers::{Handlers, MessageHandler}; pub use message_reader::MessageReader; -use std::str; +use std::{str, usize}; #[derive(Debug, PartialEq)] #[allow(dead_code)] @@ -148,6 +147,10 @@ impl Message { self.data[4] } + pub fn get_message_code_send(&self) -> u8 { + self.data[0] + } + pub fn new_with_data(data: Vec) -> Self { Self { data, pointer: 0 } } @@ -222,6 +225,21 @@ impl Message { val } + pub fn read_raw_byte(&mut self) -> Vec { + if self.pointer + 4 > self.data.len() { + return vec![]; + } + + let val = vec![ + self.data[self.pointer], + self.data[self.pointer + 1], + self.data[self.pointer + 2], + self.data[self.pointer + 3], + ]; + self.pointer += 4; + val + } + pub fn read_int32(&mut self) -> u32 { if self.pointer + 4 > self.data.len() { return 0; @@ -250,6 +268,11 @@ impl Message { self } + pub fn write_int8(&mut self, value: u8) -> &mut Self { + self.data.extend_from_slice(&value.to_le_bytes()); + self + } + pub fn write_int32(&mut self, value: u32) -> &mut Self { self.data.extend_from_slice(&value.to_le_bytes()); self @@ -311,9 +334,16 @@ impl Message { 41 => Ok("Relogged"), 42 => Ok("UserSearch"), 64 => Ok("RoomList"), + 69 => Ok("PrivilegedUsers"), + 71 => Ok("HaveNoParent"), + 83 => Ok("ParentMinSpeed"), + 84 => Ok("ParentSpeedRatio"), 92 => Ok("CheckPrivileges"), + 93 => Ok("EmbeddedMessage"), 100 => Ok("AcceptChildren"), 102 => Ok("PossibleParents"), + 104 => Ok("WishlistInterval"), + 160 => Ok("ExcludedSearchPhrases"), 1001 => Ok("CantConnectToPeer"), _ => { Err(Error(format!("Unknown server message code: {}", code))) diff --git a/src/message/peer/file_search_response.rs b/src/message/peer/file_search_response.rs index d4582f1..539c171 100644 --- a/src/message/peer/file_search_response.rs +++ b/src/message/peer/file_search_response.rs @@ -9,7 +9,10 @@ impl MessageHandler for FileSearchResponse { 9 } fn handle(&self, message: &mut Message, sender: Sender) { - let file_search = FileSearchResult::new_from_message(message); + let file_search = match FileSearchResult::new_from_message(message) { + Ok(result) => result, + Err(_) => return, // Skip malformed search results + }; sender .send(PeerOperation::FileSearchResult(file_search)) @@ -40,7 +43,7 @@ fn test_new_from_message() { let mut message = Message::new_with_data(data); message.set_pointer(8); - let file_search = FileSearchResult::new_from_message(&mut message); + let file_search = FileSearchResult::new_from_message(&mut message).unwrap(); assert_eq!(file_search.token, 882125677); assert_eq!(file_search.files.len(), 2); let file = &file_search.files[0]; diff --git a/src/message/peer/mod.rs b/src/message/peer/mod.rs index 51b67fa..f122fa9 100644 --- a/src/message/peer/mod.rs +++ b/src/message/peer/mod.rs @@ -1,10 +1,14 @@ mod file_search_response; mod get_share_file_list; +mod place_in_queue_response; mod transfer_request; +mod transfer_response; mod upload_failed; // Re-export handlers pub use file_search_response::FileSearchResponse; pub use get_share_file_list::GetShareFileList; +pub use place_in_queue_response::PlaceInQueueResponse; pub use transfer_request::TransferRequest; +pub use transfer_response::TransferResponse; pub use upload_failed::UploadFailedHandler; diff --git a/src/message/peer/place_in_queue_response.rs b/src/message/peer/place_in_queue_response.rs new file mode 100644 index 0000000..ed8cb8b --- /dev/null +++ b/src/message/peer/place_in_queue_response.rs @@ -0,0 +1,22 @@ +use crate::{ + message::{Message, MessageHandler}, + peer::PeerOperation, +}; +use std::sync::mpsc::Sender; + +pub struct PlaceInQueueResponse; + +impl MessageHandler for PlaceInQueueResponse { + fn get_code(&self) -> u8 { + 43 + } + + fn handle(&self, message: &mut Message, sender: Sender) { + let filename = message.read_string(); + let place = message.read_int32() as u32; + + sender + .send(PeerOperation::PlaceInQueueResponse { filename, place }) + .unwrap(); + } +} diff --git a/src/message/peer/transfer_request.rs b/src/message/peer/transfer_request.rs index 828b6b0..693205b 100644 --- a/src/message/peer/transfer_request.rs +++ b/src/message/peer/transfer_request.rs @@ -1,6 +1,7 @@ use crate::{ message::{Message, MessageHandler}, peer::PeerOperation, + trace, types::Transfer, }; use std::sync::mpsc::Sender; @@ -11,6 +12,7 @@ impl MessageHandler for TransferRequest { 40 } fn handle(&self, message: &mut Message, sender: Sender) { + trace!("transfer_response: {:?}", message.get_buffer()); let transfer = Transfer::new_from_message(message); sender diff --git a/src/message/peer/transfer_response.rs b/src/message/peer/transfer_response.rs new file mode 100644 index 0000000..baa19d2 --- /dev/null +++ b/src/message/peer/transfer_response.rs @@ -0,0 +1,27 @@ +use crate::{ + message::{Message, MessageHandler}, + peer::PeerOperation, +}; +use std::sync::mpsc::Sender; + +pub struct TransferResponse; + +impl MessageHandler for TransferResponse { + fn get_code(&self) -> u8 { + 41 + } + + fn handle(&self, message: &mut Message, sender: Sender) { + let token = message.read_int32(); + let allowed = message.read_int8(); + let reason = (allowed == 0).then(|| message.read_string()); + + sender + .send(PeerOperation::TransferResponse { + token, + allowed: allowed == 1, + reason, + }) + .unwrap(); + } +} diff --git a/src/message/server/connect_to_peer.rs b/src/message/server/connect_to_peer.rs index 1b2c5a1..141fee2 100644 --- a/src/message/server/connect_to_peer.rs +++ b/src/message/server/connect_to_peer.rs @@ -1,7 +1,6 @@ use crate::message::{Message, MessageHandler}; -use crate::peer::{ConnectionType, Peer}; +use crate::peer::Peer; use crate::server::ServerOperation; -use crate::trace; use std::sync::mpsc::Sender; pub struct ConnectToPeerHandler; impl MessageHandler for ConnectToPeerHandler { @@ -10,12 +9,8 @@ impl MessageHandler for ConnectToPeerHandler { } fn handle(&self, message: &mut Message, sender: Sender) { let peer = Peer::new_from_message(message); - match peer.connection_type { - ConnectionType::P => (), - ConnectionType::F => trace!("[server] ConnectToPeer {:?}", peer), - ConnectionType::D => (), - } - + println!("ConnectToPeerHandler: {:?}", peer); + println!("ConnectToPeerHandler: {:?}", peer.username); sender.send(ServerOperation::ConnectToPeer(peer)).unwrap(); } } diff --git a/src/message/server/get_peer_address.rs b/src/message/server/get_peer_address.rs new file mode 100644 index 0000000..625c7d9 --- /dev/null +++ b/src/message/server/get_peer_address.rs @@ -0,0 +1,37 @@ +use crate::message::{Message, MessageHandler}; +use crate::server::ServerOperation; +use std::sync::mpsc::Sender; + +pub struct GetPeerAddressHandler; + +impl MessageHandler for GetPeerAddressHandler { + fn get_code(&self) -> u8 { + 3 + } + + fn handle(&self, message: &mut Message, sender: Sender) { + let username = message.read_string(); + + // Read IP address as 4 bytes + let mut ip: Vec = vec![]; + for _ in 0..4 { + ip.push(message.read_int8()); + } + let host = format!("{}.{}.{}.{}", ip[3], ip[2], ip[1], ip[0]); + + let port = message.read_int32(); + let obfuscation_type = message.read_int32(); + let obfuscated_port = message.read_int32() as u16; + println!("GetPeerAddressHandler: {:?}", username); // Debug print + + sender + .send(ServerOperation::GetPeerAddressResponse { + username, + host, + port, + obfuscation_type, + obfuscated_port, + }) + .unwrap(); + } +} diff --git a/src/message/server/message_factory.rs b/src/message/server/message_factory.rs index 27c5299..e1d06af 100644 --- a/src/message/server/message_factory.rs +++ b/src/message/server/message_factory.rs @@ -4,6 +4,13 @@ use crate::{ pub struct MessageFactory; impl MessageFactory { + pub fn build_get_peer_address(username: &str) -> Message { + let mut message = Message::new(); + + message.write_int32(3); + message.write_string(username); + message + } pub fn build_login_message(username: &str, password: &str) -> Message { // Message::new_with_data( // [ @@ -31,9 +38,8 @@ impl MessageFactory { .write_string(password) .write_int32(157) // version .write_string(&hash) - .write_int32(100); //minor version - - message + .write_int32(100) + .clone() } pub fn build_shared_folders_message( @@ -47,14 +53,6 @@ impl MessageFactory { .clone() } pub fn build_file_search_message(token: u32, query: &str) -> Message { - // Message::new_with_data( - // [ - // 26, 0, 0, 0, 219, 178, 47, 28, 11, 0, 0, 0, 116, 104, 101, 32, 119, 101, 101, 107, 101, - // 110, 100, 0, 0, - // ] - // .to_vec(), - // ) - // .clone() Message::new() .write_int32(26) .write_int32(token) @@ -68,10 +66,15 @@ impl MessageFactory { .clone() } pub fn build_no_parent_message() -> Message { - Message::new().write_int32(71).write_int32(1).clone() + Message::new().write_int32(71).write_bool(true).clone() } - pub fn build_set_wait_port_message() -> Message { - Message::new().write_int32(2).write_int32(2234).clone() + pub fn build_set_wait_port_message(port: u32) -> Message { + Message::new() + .write_int32(2) + .write_int32(port) + // .write_int32(0) + // .write_int32(port) // should be different port + .clone() } pub fn build_watch_user(token: u32) -> Message { Message::new() @@ -87,11 +90,23 @@ impl MessageFactory { .clone() } + pub fn build_transfer_request_message( + filename: &str, + token: u32, + ) -> Message { + Message::new() + .write_int32(40) // code + .write_int32(0) // direction + .write_int32(token) + .write_string(filename) + .clone() + } pub fn build_transfer_response_message(transfer: Transfer) -> Message { Message::new() .write_int32(41) .write_int32(transfer.token) .write_bool(true) + .write_int64(transfer.size) .clone() } pub fn build_pierce_firewall_message(token: u32) -> Message { @@ -108,6 +123,7 @@ impl MessageFactory { token: u32, ) -> Message { Message::new() + .write_int8(1) .write_string(own_username) .write_string(&connection_type.to_string()) .write_int32(token) diff --git a/src/message/server/mod.rs b/src/message/server/mod.rs index 8a0db56..07aaa79 100644 --- a/src/message/server/mod.rs +++ b/src/message/server/mod.rs @@ -1,6 +1,7 @@ mod connect_to_peer; mod excluded_search_phrases; mod file_search; +mod get_peer_address; mod login; mod message_factory; mod message_user; @@ -10,10 +11,10 @@ mod privileged_users; mod room_list; mod wish_list_interval; -// Re-export handlers pub use connect_to_peer::ConnectToPeerHandler; pub use excluded_search_phrases::ExcludedSearchPhrasesHandler; pub use file_search::FileSearchHandler; +pub use get_peer_address::GetPeerAddressHandler; pub use login::LoginHandler; pub use message_factory::MessageFactory; pub use message_user::MessageUser; diff --git a/src/peer/default_peer.rs b/src/peer/default_peer.rs index e9c2673..c3e5fa1 100644 --- a/src/peer/default_peer.rs +++ b/src/peer/default_peer.rs @@ -1,11 +1,12 @@ use crate::dispatcher::MessageDispatcher; use crate::message::peer::FileSearchResponse; use crate::message::peer::{ - GetShareFileList, TransferRequest, UploadFailedHandler, + GetShareFileList, PlaceInQueueResponse, TransferRequest, TransferResponse, + UploadFailedHandler, }; use crate::message::server::MessageFactory; use crate::message::{Handlers, Message, MessageReader, MessageType}; -use crate::types::{FileSearchResult, Transfer}; +use crate::types::{Download, FileSearchResult, Transfer}; use std::sync::mpsc; use std::sync::mpsc::{Receiver, Sender}; @@ -20,6 +21,7 @@ use std::net::ToSocketAddrs; use std::time::Duration; #[allow(dead_code)] +#[derive(Debug)] pub struct DefaultPeer { peer: Peer, peer_channel: Option>, @@ -33,6 +35,15 @@ pub enum PeerOperation { SendMessage(Message), FileSearchResult(FileSearchResult), TransferRequest(Transfer), + TransferResponse { + token: u32, + allowed: bool, + reason: Option, + }, + PlaceInQueueResponse { + filename: String, + place: u32, + }, } impl DefaultPeer { pub fn new(peer: Peer, client_channel: Sender) -> Self { @@ -44,7 +55,43 @@ impl DefaultPeer { write_thread: None, } } + pub fn disconnect(mut self) { + debug!("[default_peer:{}] disconnect", self.peer.username.clone()); + if let Err(e) = self.client_channel.send( + ClientOperation::PeerDisconnected(self.peer.username.clone()), + ) { + error!("Failed to send disconnect notification: {}", e); + } + + self.peer_channel.take(); + } + + pub fn connect_with_socket( + mut self, + mut stream: TcpStream, + ) -> Result { + if let Some(token) = self.peer.token { + let mut message: Vec = [0, 5, 0, 0, 0, 0].to_vec(); + message.extend_from_slice(&token.to_le_bytes()); + stream.write_all(&message).unwrap(); + } + + stream.set_read_timeout(Some(Duration::from_secs(5)))?; + stream.set_write_timeout(Some(Duration::from_secs(5)))?; + + trace!( + "[default_peer:{}] connect_with_socket: direct", + self.peer.username + ); + self.start_read_write_loops(stream)?; + + Ok(self) + } pub fn connect(mut self) -> Result { + println!( + "[default_peer] Connecting to {} on port {}", + self.peer.host, self.peer.port + ); let socket_address = format!("{}:{}", self.peer.host, self.peer.port) .to_socket_addrs()? .next() @@ -52,6 +99,7 @@ impl DefaultPeer { io::Error::new(io::ErrorKind::InvalidInput, "Invalid address") })?; + trace!("[default_peer:{}] connect: indirect", self.peer.username); let mut stream = TcpStream::connect_timeout( &socket_address, Duration::from_secs(20), @@ -74,6 +122,10 @@ impl DefaultPeer { &mut self, stream: TcpStream, ) -> Result<(), io::Error> { + debug!( + "[default_peer:{}] start_read_write_loops", + self.peer.username + ); let (peer_sender, peer_reader): ( Sender, Receiver, @@ -86,17 +138,23 @@ impl DefaultPeer { let mut write_stream = stream; // Use the original stream for writing let peer = self.peer.clone(); + let peer_clone = self.peer.clone(); let client_channel_for_read = self.client_channel.clone(); - // Spawn the reader thread self.read_thread = Some(thread::spawn(move || { let mut handlers = Handlers::new(); handlers.register_handler(FileSearchResponse); handlers.register_handler(TransferRequest); + handlers.register_handler(TransferResponse); handlers.register_handler(GetShareFileList); handlers.register_handler(UploadFailedHandler); + handlers.register_handler(PlaceInQueueResponse); - let dispatcher = MessageDispatcher::new(peer_sender, handlers); + let dispatcher = MessageDispatcher::new( + "default_peer".to_string(), + peer_sender, + handlers, + ); let mut buffered_reader = MessageReader::new(); loop { @@ -113,7 +171,7 @@ impl DefaultPeer { continue; } Err(e) => { - error!("Error reading from peer: {}. Terminating read loop.", e); + error!("[default_peer:{}] Error reading from peer: {}. Terminating read loop.",peer.username, e); let _ = client_channel_for_read.send( ClientOperation::PeerDisconnected( peer.username.clone(), @@ -147,6 +205,7 @@ impl DefaultPeer { peer.username.clone(), ), ); + break; } Ok(None) => continue, @@ -158,77 +217,109 @@ impl DefaultPeer { let peer_channel = self.peer_channel.clone(); let peer_username = self.peer.username.clone(); - self.write_thread = Some(thread::spawn(move || { - loop { - match peer_reader.recv() { - Ok(operation) => { - match operation { - PeerOperation::SendMessage(message) => { - trace!( - "[default_peer:{:?}] ➡ {:?} - {:?}", - peer_username, - message - .get_message_name( - MessageType::Peer, - u32::from_le_bytes( - message - .get_slice(0, 4) - .try_into() - .unwrap() - ) - ) - .map_err(|e| e.to_string()), - message - ); + self.write_thread = Some(thread::spawn(move || loop { + match peer_reader.recv() { + Ok(operation) => match operation { + PeerOperation::SendMessage(message) => { + trace!( + "[default_peer:{:?}] ➡ {:?} - {:?}", + peer_username, + message + .get_message_name( + MessageType::Peer, + u32::from_le_bytes( + message + .get_slice(0, 4) + .try_into() + .unwrap() + ) + ) + .map_err(|e| e.to_string()), + message + ); - if let Err(e) = write_stream - .write_all(&message.get_buffer()) - { - error!("Error writing message to stream: {} - {}. Terminating write loop.", peer_username, e); - let _ = client_channel.send( - ClientOperation::PeerDisconnected( - peer_username.clone(), - ), - ); - break; - } - } - PeerOperation::FileSearchResult(file_search) => { - client_channel - .send(ClientOperation::SearchResult( - file_search, - )) - .unwrap(); - } - PeerOperation::TransferRequest(transfer) => { - debug!( - "[default_peer:{}] transfer request: {:?}", - peer_username, transfer - ); - client_channel - .send(ClientOperation::TransferRequest( - transfer.clone(), - )) - .unwrap(); - - let transfer_response = - MessageFactory::build_transfer_response_message(transfer.clone()); - - if let Some(sender) = peer_channel.clone() { - sender - .send(PeerOperation::SendMessage( - transfer_response, - )) - .unwrap(); - } - } + if let Err(e) = + write_stream.write_all(&message.get_buffer()) + { + error!("Error writing message to stream: {} - {}. Terminating write loop.", peer_username, e); + let _ = client_channel.send( + ClientOperation::PeerDisconnected( + peer_username.clone(), + ), + ); + break; } } - Err(_) => { - // The sender has been dropped, the peer is shutting down. - debug!("Peer channel closed. Terminating write loop."); + PeerOperation::FileSearchResult(file_search) => { + client_channel + .send(ClientOperation::SearchResult(file_search)) + .unwrap(); + } + PeerOperation::TransferRequest(transfer) => { + debug!( + "[default_peer:{:}] TransferRequest for {}", + peer_username, transfer.token + ); + + let transfer_response = + MessageFactory::build_transfer_response_message( + transfer.clone(), + ); + + if let Some(sender) = peer_channel.clone() { + sender + .send(PeerOperation::SendMessage( + transfer_response, + )) + .unwrap(); + } + + debug!( + "[default_peer:{:}] TransferResponse for {}", + peer_username, transfer.token + ); + break; } + PeerOperation::TransferResponse { + token, + allowed, + reason, + } => { + debug!( + "[default_peer:{}] transfer response token: {} allowed: {}", + peer_username, token, allowed + ); + + if !allowed { + if let Some(reason_text) = reason { + debug!( + "[default_peer:{}] Transfer rejected: {} - token {}, I will receive TransferRequest soon...", + peer_username.clone(), + reason_text, + token + ); + } + } else { + debug!("[default_peer:{:}] Transfer allowed, ready to connect with token {:}",peer_username, token); + client_channel + .send(ClientOperation::DownloadFromPeer( + token, + peer_clone.clone(), + )) + .unwrap(); + } + } + PeerOperation::PlaceInQueueResponse { filename, place } => { + debug!( + "[default_peer:{}] Place in queue response - file: {}, place: {}", + peer_username, filename, place + ); + } + }, + Err(_) => { + debug!("[default_peer:{}] Peer channel closed. Terminating write loop.", peer_username); + break; } } })); @@ -236,8 +327,14 @@ impl DefaultPeer { Ok(()) } - pub fn transfer_request(&self, filename: String) -> Result<(), io::Error> { - let message = MessageFactory::build_queue_upload_message(&filename); + pub fn transfer_request( + &self, + download: Download, + ) -> Result<(), io::Error> { + let message = MessageFactory::build_transfer_request_message( + &download.filename, + download.token, + ); if let Some(sender) = &self.peer_channel { sender.send(PeerOperation::SendMessage(message)).unwrap(); } @@ -248,6 +345,7 @@ impl DefaultPeer { impl Drop for DefaultPeer { fn drop(&mut self) { self.peer_channel = None; + trace!("[default_peer:{}] drop", self.peer.username); if let Some(handle) = self.read_thread.take() { debug!("Joining read thread..."); diff --git a/src/peer/download_peer.rs b/src/peer/download_peer.rs index dfa660b..1c92408 100644 --- a/src/peer/download_peer.rs +++ b/src/peer/download_peer.rs @@ -359,7 +359,7 @@ mod tests { assert_eq!( received_messages[0], vec![ - 12, 0, 0, 0, 111, 119, 110, 95, 117, 115, 101, 114, 110, 97, + 1, 12, 0, 0, 0, 111, 119, 110, 95, 117, 115, 101, 114, 110, 97, 109, 101, 1, 0, 0, 0, 70, 33, 0, 0, 0 ] ); diff --git a/src/peer/listen.rs b/src/peer/listen.rs index c88e01b..7d5a7af 100644 --- a/src/peer/listen.rs +++ b/src/peer/listen.rs @@ -1,30 +1,29 @@ use std::sync::mpsc::Sender; use std::{io, net::TcpListener}; +use crate::client::ClientOperation; use crate::message::MessageReader; -use crate::server::ServerOperation; +use crate::peer::NewPeer; use crate::{debug, error, info, trace, warn}; pub struct Listen {} impl Listen { - pub fn start(port: u32, _server_channel: Sender) { + pub fn start(port: u32, client_sender: Sender) { info!("starting listener on port {port}"); let listener = TcpListener::bind(format!("0.0.0.0:{port}")).unwrap(); for stream in listener.incoming() { let mut read_stream = stream.unwrap(); + let mut stop = false; let mut buffered_reader = MessageReader::new(); - loop { + while !stop { match buffered_reader.read_from_socket(&mut read_stream) { Ok(_) => {} 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; @@ -37,39 +36,50 @@ impl Listen { match buffered_reader.extract_message() { Ok(Some(mut message)) => { - debug!( - "Received message: {:?}", - message.get_message_code() + message.set_pointer(4); + let message_code = message.read_int8(); + + trace!( + "[listener] Received message with code: {}", + message_code ); - trace!("{:?}", message.get_data()); - message.print_hex(); - if message.get_message_code() == 1 { - let size = message.read_int32(); - message.set_pointer(8); - let typex = message.read_string(); + 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, + ); - // let user = message.read_string(); - // // let connection_type: ConnectionType = - // // message.read_string().parse().unwrap(); - // // let token = message.read_int32(); - // - debug!("type: {:?}", typex); - debug!("user: {:?}", size); - - // server_channel - // .send(ServerOperation::ConnectToPeer(peer)) - // .unwrap(); + 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!("Error extracting message: {}", e) + warn!("[listener] Error extracting message: {}", e) } Ok(None) => continue, } } - info!("Connection established!"); + // info!("[listener] Connection established!"); } } } diff --git a/src/peer/mod.rs b/src/peer/mod.rs index 666cf89..7032047 100644 --- a/src/peer/mod.rs +++ b/src/peer/mod.rs @@ -2,12 +2,38 @@ mod default_peer; mod download_peer; pub mod listen; -// Re-export commonly used items pub use default_peer::{DefaultPeer, PeerOperation}; +pub use download_peer::DownloadPeer; use crate::message::Message; use core::fmt; -use std::str::FromStr; +use std::{net::TcpStream, str::FromStr}; + +#[derive(Debug)] +#[allow(dead_code)] +pub struct NewPeer { + pub username: String, + pub connection_type: ConnectionType, + pub token: u32, + pub tcp_stream: TcpStream, +} +impl NewPeer { + pub fn new_from_message( + message: &mut Message, + tcp_stream: TcpStream, + ) -> Self { + let username = message.read_string(); + let connection_type = message.read_string().parse().unwrap(); + let token = message.read_int32(); + + Self { + username, + connection_type, + token, + tcp_stream, + } + } +} #[derive(Debug, Clone)] pub enum ConnectionType { @@ -16,7 +42,6 @@ pub enum ConnectionType { D, } -// Define the error type #[derive(Debug, Clone)] pub struct ParseConnectionTypeError; @@ -28,7 +53,6 @@ impl fmt::Display for ParseConnectionTypeError { impl std::error::Error for ParseConnectionTypeError {} -// Implement FromStr for ConnectionType impl FromStr for ConnectionType { type Err = ParseConnectionTypeError; @@ -61,9 +85,9 @@ pub struct Peer { pub host: String, pub port: u32, pub token: Option, - pub privileged: u8, - pub unknown: u8, - pub obfuscated_port: u8, + pub privileged: Option, + pub unknown: Option, + pub obfuscated_port: Option, } impl Peer { #[allow(clippy::too_many_arguments, dead_code)] @@ -72,7 +96,7 @@ impl Peer { connection_type: ConnectionType, host: String, port: u32, - token: u32, + token: Option, privileged: u8, unknown: u8, obfuscated_port: u8, @@ -82,10 +106,10 @@ impl Peer { connection_type, host, port, - token: Some(token), - privileged, - unknown, - obfuscated_port, + token, + privileged: Some(privileged), + unknown: Some(unknown), + obfuscated_port: Some(obfuscated_port), } } #[allow(dead_code)] @@ -119,9 +143,9 @@ impl Peer { host, port, token: Some(token), - privileged, - unknown, - obfuscated_port, + privileged: Some(privileged), + unknown: Some(unknown), + obfuscated_port: Some(obfuscated_port), } } } @@ -142,9 +166,9 @@ fn test_new_from_message() { assert_eq!(peer.host, "45.37.231.27"); assert_eq!(peer.port, 2234); assert_eq!(peer.token, Some(1658546)); - assert_eq!(peer.privileged, 0); - assert_eq!(peer.unknown, 0); - assert_eq!(peer.obfuscated_port, 0); + assert_eq!(peer.privileged, Some(0)); + assert_eq!(peer.unknown, Some(0)); + assert_eq!(peer.obfuscated_port, Some(0)); } #[test] @@ -158,6 +182,8 @@ fn test_new_from_message2() { let mut message = Message::new_with_data(data); message.set_pointer(8); + println!("code: {}", message.get_message_code_u32()); + let peer = Peer::new_from_message(&mut message); assert_eq!(peer.username, "grandpag"); @@ -165,7 +191,7 @@ fn test_new_from_message2() { assert_eq!(peer.host, "68.193.128.137"); assert_eq!(peer.port, 2235); assert_eq!(peer.token, Some(4154)); - assert_eq!(peer.privileged, 0); - assert_eq!(peer.unknown, 1); - assert_eq!(peer.obfuscated_port, 0); + assert_eq!(peer.privileged, Some(0)); + assert_eq!(peer.unknown, Some(1)); + assert_eq!(peer.obfuscated_port, Some(0)); } diff --git a/src/server.rs b/src/server.rs index a462627..4ffea9c 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,8 +1,10 @@ use crate::client::ClientOperation; use crate::dispatcher::MessageDispatcher; +use crate::error::{Result, SoulseekRs}; use crate::message::server::ConnectToPeerHandler; use crate::message::server::ExcludedSearchPhrasesHandler; use crate::message::server::FileSearchHandler; +use crate::message::server::GetPeerAddressHandler; use crate::message::server::LoginHandler; use crate::message::server::MessageFactory; use crate::message::server::MessageUser; @@ -11,10 +13,9 @@ use crate::message::server::ParentSpeedRatioHandler; use crate::message::server::PrivilegedUsersHandler; use crate::message::server::RoomListHandler; use crate::message::server::WishListIntervalHandler; -use crate::message::Handlers; use crate::message::Message; use crate::message::MessageReader; -use crate::peer::listen::Listen; +use crate::message::{Handlers, MessageType}; use crate::peer::ConnectionType; use crate::peer::Peer; @@ -26,7 +27,7 @@ use std::sync::{Arc, Barrier, Mutex}; use std::thread::{self}; use std::time::{Duration, Instant}; -use crate::{debug, error, info, warn}; +use crate::{debug, error, info, trace, warn}; #[derive(Debug, Clone)] pub struct PeerAddress { @@ -69,14 +70,6 @@ impl Context { logged_in: Option::None, } } - - // pub fn get_messages_for_user(&self, username: String) -> Option<&UserMessage> { - // self.user_messages.get(&username) - // } - // - // pub fn get_rooms(&mut self) -> &mut Rooms { - // &mut self.rooms - // } } #[derive(Debug, Clone)] pub struct UserMessage { @@ -186,6 +179,14 @@ pub enum ServerOperation { #[allow(dead_code)] ConnectToPeer(Peer), PierceFirewall(u32), + GetPeerAddress(String), + GetPeerAddressResponse { + username: String, + host: String, + port: u32, + obfuscation_type: u32, + obfuscated_port: u16, + }, } #[derive(Debug)] @@ -195,11 +196,10 @@ pub struct Server { context: Arc>, } impl Server { - /// Create a new instance of Server, returning a Result pub fn new( address: PeerAddress, client_channel: Sender, - ) -> Result { + ) -> std::result::Result { let (sender, server_channel): ( Sender, Receiver, @@ -230,7 +230,7 @@ impl Server { &mut self, server_channel: Receiver, client_channel: Sender, - ) -> Result<(), io::Error> { + ) -> std::result::Result<(), io::Error> { let socket_address = format!("{}:{}", self.address.host, self.address.port) .to_socket_addrs()? @@ -279,9 +279,11 @@ impl Server { handlers.register_handler(ParentSpeedRatioHandler); handlers.register_handler(PrivilegedUsersHandler); handlers.register_handler(FileSearchHandler); + handlers.register_handler(GetPeerAddressHandler); handlers.register_handler(ConnectToPeerHandler); - let dispatcher = MessageDispatcher::new(sender, handlers); + let dispatcher = + MessageDispatcher::new("server".to_string(), sender, handlers); let mut buffered_reader = MessageReader::new(); loop { @@ -305,16 +307,16 @@ impl Server { match buffered_reader.extract_message() { Ok(Some(mut message)) => { - // trace!( - // "[server] ← {:?}", - // message - // .get_message_name( - // MessageType::Server, - // message.get_message_code() as u32 - // ) - // .map_err(|e| e.to_string()) - // ); - + 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) => { @@ -325,10 +327,6 @@ impl Server { } }); - // let monitor_handle = thread::spawn(move || loop { - // thread::sleep(Duration::from_secs(1)); - // println!("Stream state: connected={}", stream.peer_addr().is_ok()); - // }); let context = self.context.clone(); thread::spawn(move || { write_barrier.wait(); @@ -336,72 +334,98 @@ impl Server { if let Ok(operation) = server_channel.recv() { match operation { ServerOperation::ConnectToPeer(peer) => { - match peer.connection_type { + debug!( + "[server] ConnectToPeer {} - ConnectionType {}", + peer.username, peer.connection_type + ); + + if let Some(op) = match peer.connection_type { ConnectionType::P => { - match client_channel.send( - ClientOperation::ConnectToPeer(peer), - ) { - Ok(_) => {} - Err(_e) => {} - } + Some(ClientOperation::ConnectToPeer(peer)) } ConnectionType::F => { - match client_channel.send( - ClientOperation::PierceFireWall(peer), - ) { - Ok(_) => { - debug!("Sent PierceFireWall operation for F-type connection"); - } - Err(_e) => { - error!("Failed to send PierceFireWall operation"); - } - } + Some(ClientOperation::PierceFireWall(peer)) } - ConnectionType::D => {} + ConnectionType::D => None, + } { + client_channel.send(op).unwrap(); } } ServerOperation::LoginStatus(message) => { context.lock().unwrap().logged_in = Some(message); } ServerOperation::PierceFirewall(token) => { - let pierce_message = - MessageFactory::build_pierce_firewall_message( + if let Err(e) = write_stream.write_all( + &MessageFactory::build_pierce_firewall_message( token, + ) + .get_buffer(), + ) { + error!( + "Error writing PierceFirewall message: {}", + e ); - match write_stream - .write_all(&pierce_message.get_buffer()) - { - Ok(_) => { - debug!("Sent PierceFirewall message with token: {}", token); - } - Err(e) => { - error!("Error writing PierceFirewall message: {}", e); - break; - } + break; } + debug!( + "Sent PierceFirewall message with token: {}", + token + ); } ServerOperation::SendMessage(message) => { - // trace!( - // "[server] ➡ {:?}", - // message - // .get_message_name( - // MessageType::Server, - // u32::from_le_bytes( - // message.get_slice(0, 4).try_into().unwrap() - // ) - // ) - // .map_err(|e| e.to_string()), - // ); match write_stream.write_all(&message.get_buffer()) { - Ok(_) => {} - Err(e) => { - error!( - "Error writing message to stream : {}", - e - ); - break; - } + Ok(_) => trace!( + "[server] → {:?} - {:}", + message + .get_message_name( + MessageType::Server, + message.get_message_code_send() + as u32 + ) + .map_err(|e| e.to_string()), + message.get_message_code_send() as u32 + ), + Err(e) => error!( + "Error writing message to server: {}", + e + ), + } + } + ServerOperation::GetPeerAddress(username) => { + if let Err(e) = write_stream.write_all( + &MessageFactory::build_get_peer_address( + &username, + ) + .get_buffer(), + ) { + error!("Error writing get_peeer_address message: {}", e); + break; + } + } + ServerOperation::GetPeerAddressResponse { + username, + host, + port, + obfuscation_type, + obfuscated_port, + } => { + debug!( + "Received GetPeerAddress response for {}: {}:{} (obf_type: {}, obf_port: {})", + username, host, port, obfuscation_type, obfuscated_port + ); + + // Forward the response to the client channel + if let Err(e) = client_channel.send( + ClientOperation::GetPeerAddressResponse { + username, + host, + port, + obfuscation_type, + obfuscated_port, + }, + ) { + error!("Error forwarding GetPeerAddress response to client: {}", e); } } } @@ -418,16 +442,7 @@ impl Server { Err(e) => error!("Failed to send: {}", e), } } - fn start_listener(&self, server_channel: Sender) { - thread::spawn(move || Listen::start(2234, server_channel)); - } - - pub fn login( - &self, - username: &str, - password: &str, - ) -> Result { - self.start_listener(self.sender.clone()); + pub fn login(&self, username: &str, password: &str) -> Result { // Send the login message self.queue_message(MessageFactory::build_login_message( username, password, @@ -441,10 +456,7 @@ impl Server { loop { if start.elapsed() > timeout { warn!("Timeout waiting for login response"); - return Err(std::io::Error::new( - std::io::ErrorKind::TimedOut, - "Timeout waiting for login response", - )); + return Err(SoulseekRs::Timeout); } { @@ -458,12 +470,14 @@ impl Server { if logged_in.unwrap() { info!("Logged in as {}", username); - self.queue_message(MessageFactory::build_set_wait_port_message()); self.queue_message(MessageFactory::build_shared_folders_message( 1, 499, )); self.queue_message(MessageFactory::build_no_parent_message()); self.queue_message(MessageFactory::build_set_status_message(2)); + self.queue_message(MessageFactory::build_set_wait_port_message( + 2235, + )); } Ok(logged_in.unwrap()) diff --git a/src/types.rs b/src/types.rs index 24b8b91..8af715f 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1,13 +1,13 @@ use std::collections::HashMap; -use crate::{message::Message, utils::zlib::deflate}; +use crate::{error::Result, message::Message, utils::zlib::deflate}; #[derive(Debug, Clone)] #[allow(dead_code)] pub struct File { pub username: String, pub name: String, - pub size: u32, + pub size: u64, pub attribs: HashMap, } pub struct UploadFailed { @@ -27,14 +27,15 @@ pub struct FileSearchResult { pub files: Vec, pub slots: u8, pub speed: u32, + pub username: String, } impl FileSearchResult { - pub fn new_from_message(message: &mut Message) -> Self { + pub fn new_from_message(message: &mut Message) -> Result { let pointer = message.get_pointer(); let size = message.get_size(); let data: Vec = message.get_slice(pointer, size); - let deflated = deflate(&data).unwrap(); + let deflated = deflate(&data)?; let mut message = Message::new_with_data(deflated); let username = message.read_string(); @@ -44,8 +45,7 @@ impl FileSearchResult { for _ in 0..n_files { message.read_int8(); let name = message.read_string(); - let size = message.read_int32(); - message.read_int32(); + let size = message.read_int64(); message.read_string(); let n_attribs = message.read_int32(); let mut attribs: HashMap = HashMap::new(); @@ -63,12 +63,13 @@ impl FileSearchResult { let slots = message.read_int8(); let speed = message.read_int32(); - Self { + Ok(Self { token, files, slots, speed, - } + username, + }) } } @@ -80,6 +81,13 @@ pub struct Transfer { pub filename: String, pub size: u64, } +#[derive(Debug, Clone)] +pub struct Download { + pub username: String, + pub filename: String, + pub token: u32, + pub size: u64, +} #[derive(Debug, Clone)] #[allow(dead_code)] diff --git a/src/utils/thread_pool.rs b/src/utils/thread_pool.rs index 55a55b3..32c4b92 100644 --- a/src/utils/thread_pool.rs +++ b/src/utils/thread_pool.rs @@ -1,3 +1,4 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{mpsc, Arc, Mutex}; use std::thread; @@ -6,6 +7,7 @@ type Job = Box; pub struct ThreadPool { workers: Vec, sender: mpsc::Sender, + active_threads: Arc, } impl ThreadPool { @@ -14,13 +16,23 @@ impl ThreadPool { let (sender, receiver) = mpsc::channel(); let receiver = Arc::new(Mutex::new(receiver)); + let active_threads = Arc::new(AtomicUsize::new(0)); let mut workers = Vec::with_capacity(size); for id in 0..size { - workers.push(Worker::new(id, Arc::clone(&receiver))); + workers.push(Worker::new( + id, + Arc::clone(&receiver), + Arc::clone(&active_threads), + size, + )); } - ThreadPool { workers, sender } + ThreadPool { + workers, + sender, + active_threads, + } } pub fn execute(&self, f: F) @@ -30,6 +42,14 @@ impl ThreadPool { let job = Box::new(f); self.sender.send(job).unwrap(); } + + pub fn thread_count(&self) -> usize { + self.workers.len() + } + + pub fn active_thread_count(&self) -> usize { + self.active_threads.load(Ordering::SeqCst) + } } impl Drop for ThreadPool { @@ -48,12 +68,35 @@ struct Worker { } impl Worker { - fn new(id: usize, receiver: Arc>>) -> Worker { + fn new( + id: usize, + receiver: Arc>>, + active_threads: Arc, + total_threads: usize, + ) -> Worker { let thread = thread::spawn(move || loop { let job = receiver.lock().unwrap().recv(); match job { Ok(job) => { + let active = + active_threads.fetch_add(1, Ordering::SeqCst) + 1; + trace!( + "Thread {} started job (active: {}/{})", + id, + active, + total_threads + ); + job(); + + let active = + active_threads.fetch_sub(1, Ordering::SeqCst) - 1; + trace!( + "Thread {} finished job (active: {}/{})", + id, + active, + total_threads + ); } Err(_) => { break; diff --git a/src/utils/zlib.rs b/src/utils/zlib.rs index e039f15..8b52421 100644 --- a/src/utils/zlib.rs +++ b/src/utils/zlib.rs @@ -17,7 +17,7 @@ impl BitReader { } } - fn read_byte(&mut self) -> Result { + fn read_byte(&mut self) -> std::result::Result { self.numbits = 0; // discard unread bits if self.pos >= self.mem.len() { return Err("End of data".to_string()); @@ -27,7 +27,7 @@ impl BitReader { Ok(b) } - fn read_bit(&mut self) -> Result { + fn read_bit(&mut self) -> std::result::Result { if self.numbits <= 0 { self.b = self.read_byte()?; self.numbits = 8; @@ -39,7 +39,7 @@ impl BitReader { Ok(bit) } - fn read_bits(&mut self, n: usize) -> Result { + fn read_bits(&mut self, n: usize) -> std::result::Result { let mut o = 0u32; for i in 0..n { o |= (self.read_bit()? as u32) << i; @@ -47,7 +47,7 @@ impl BitReader { Ok(o) } - fn read_bytes(&mut self, n: usize) -> Result { + fn read_bytes(&mut self, n: usize) -> std::result::Result { // read bytes as an integer in little-endian let mut o = 0u32; for i in 0..n { @@ -57,32 +57,38 @@ impl BitReader { } } -pub fn deflate(input: &[u8]) -> Result, String> { +use crate::error::{Result, SoulseekRs}; + +pub fn deflate(input: &[u8]) -> Result> { let mut r = BitReader::new(input.to_vec()); let cmf = r.read_byte()?; let cm = cmf & 15; // Compression method if cm != 8 { // only CM=8 is supported - return Err("invalid CM".to_string()); + return Err(SoulseekRs::CompressionError("invalid CM".to_string())); } let cinfo = (cmf >> 4) & 15; // Compression info if cinfo > 7 { - return Err("invalid CINFO".to_string()); + return Err(SoulseekRs::CompressionError("invalid CINFO".to_string())); } let flg = r.read_byte()?; if ((cmf as u32) * 256 + (flg as u32)) % 31 != 0 { - return Err("CMF+FLG checksum failed".to_string()); + return Err(SoulseekRs::CompressionError( + "CMF+FLG checksum failed".to_string(), + )); } let fdict = (flg >> 5) & 1; // preset dictionary? if fdict != 0 { - return Err("preset dictionary not supported".to_string()); + return Err(SoulseekRs::CompressionError( + "preset dictionary not supported".to_string(), + )); } - let out = inflate(&mut r)?; // decompress DEFLATE data + let out = inflate(&mut r).map_err(SoulseekRs::CompressionError)?; // decompress DEFLATE data let _adler32 = r.read_bytes(4)?; // Adler-32 checksum (for this exercise, we ignore it) Ok(out) } -fn inflate(r: &mut BitReader) -> Result, String> { +fn inflate(r: &mut BitReader) -> std::result::Result, String> { let mut bfinal = 0; let mut out = Vec::new(); while bfinal == 0 { @@ -101,7 +107,7 @@ fn inflate(r: &mut BitReader) -> Result, String> { fn inflate_block_no_compression( r: &mut BitReader, o: &mut Vec, -) -> Result<(), String> { +) -> std::result::Result<(), String> { let len = r.read_bytes(2)?; let _nlen = r.read_bytes(2)?; for _ in 0..len { @@ -157,7 +163,10 @@ impl HuffmanTree { } } -fn decode_symbol(r: &mut BitReader, t: &HuffmanTree) -> Result { +fn decode_symbol( + r: &mut BitReader, + t: &HuffmanTree, +) -> std::result::Result { let mut node = &t.root; while node.left.is_some() || node.right.is_some() { let b = r.read_bit()?; @@ -192,7 +201,7 @@ fn inflate_block_data( literal_length_tree: &HuffmanTree, distance_tree: &HuffmanTree, out: &mut Vec, -) -> Result<(), String> { +) -> std::result::Result<(), String> { loop { let sym = decode_symbol(r, literal_length_tree)?; if sym <= 255 { @@ -257,7 +266,7 @@ const CODE_LENGTH_CODES_ORDER: [usize; 19] = [ fn decode_trees( r: &mut BitReader, -) -> Result<(HuffmanTree, HuffmanTree), String> { +) -> std::result::Result<(HuffmanTree, HuffmanTree), String> { // The number of literal/length codes let hlit = r.read_bits(5)? + 257; @@ -325,7 +334,7 @@ fn decode_trees( fn inflate_block_dynamic( r: &mut BitReader, o: &mut Vec, -) -> Result<(), String> { +) -> std::result::Result<(), String> { let (literal_length_tree, distance_tree) = decode_trees(r)?; inflate_block_data(r, &literal_length_tree, &distance_tree, o) } @@ -333,7 +342,7 @@ fn inflate_block_dynamic( fn inflate_block_fixed( r: &mut BitReader, o: &mut Vec, -) -> Result<(), String> { +) -> std::result::Result<(), String> { let mut bl = Vec::new(); bl.extend(vec![8; 144]); // 0-143: 8 bits bl.extend(vec![9; 112]); // 144-255: 9 bits diff --git a/test.md b/test.md new file mode 100644 index 0000000..f7c3031 --- /dev/null +++ b/test.md @@ -0,0 +1,11 @@ +[2025-10-01 13:30:27] [Conn] Sending message of type to user insane_in_the_brain2 on new connection +[2025-10-01 13:30:27] [Msg] OUT: {'user': 'insane_in_the_brain2', 'ip_address': None, 'port': None, 'obfuscation_type': None, 'obfuscated_port': None} +[2025-10-01 13:30:27] [Conn] Requesting address for user insane_in_the_brain2 +[2025-10-01 13:30:27] [Conn] Attempting direct connection of type P to user insane_in_the_brain2, address ('127.0.0.1', 53507) +[2025-10-01 13:30:27] [Msg] OUT: {'token': 567916, 'user': 'insane_in_the_brain2', 'conn_type': 'P', 'ip_address': None, 'port': None, 'privileged': None, 'obfuscation_type': None, 'obfuscated_port': None} +[2025-10-01 13:30:27] [Conn] Requesting indirect connection to user insane_in_the_brain2 with token 567916 +[2025-10-01 13:30:27] [Msg] IN: {'user': 'insane_in_the_brain2', 'ip_address': '127.0.0.1', 'port': 53507, 'obfuscation_type': 1, 'obfuscated_port': 53508} +[2025-10-01 13:30:27] [Conn] Established outgoing connection of type P with user insane_in_the_brain2. List of outgoing messages: [] +[2025-10-01 13:30:27] [Conn] Sending peer init message of type P to user insane_in_the_brain2 +[2025-10-01 13:30:27] [Msg] OUT: {'sock': , 'init_user': 'insane_in_the_brain6', 'target_user': 'insane_in_the_brain2', 'conn_type': 'P', 'outgoing_msgs': [], 'token': 0} +[2025-10-01 13:30:27] [Msg] OUT:

{'search_username': 'insane_in_the_brain6', 'token': 1, 'freeulslots': True, 'ulspeed': 0, 'inqueue': 0, 'unknown': 0} diff --git a/todo.md b/todo.md new file mode 100644 index 0000000..eece091 --- /dev/null +++ b/todo.md @@ -0,0 +1,8 @@ +# default peer + +- lifetime: disconnect after it sends a few messages + +# peer download + +- make queue download queue with thread pool +- download progress