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

chore: refactor listener and default peer

Browse Source
This commit is contained in:
Michel de Graaf
2025-11-04 20:48:23 +01:00
parent d11c6782a7
commit 9de47892ba
6 changed files with 265 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/client.rs
+17 -8
View File
@@ -118,6 +118,8 @@ impl Client {
let listen_port = self.listen_port;
self.context.write().unwrap().sender = Some(sender.clone());
let context = self.context.clone();
let own_username = self.username.clone();
let client_sender = sender.clone();
@@ -130,9 +132,14 @@ impl Client {
server.get_address().get_port()
);
thread::spawn(move || {
Listen::start(listen_port, client_sender.clone());
});
// thread::spawn(move || {
// Listen::start(
// listen_port,
// client_sender.clone(),
// context.clone(),
// own_username,
// );
// });
let mut unlocked_context = self.context.write().unwrap();
unlocked_context.server_sender =
Some(server.get_sender().clone());
@@ -229,7 +236,7 @@ impl Client {
let download_initiated = context
.peers
.get(&username)
.map(|p| p.transfer_request(download.clone()))
.map(|p| p.queue_upload(download.filename))
.is_some();
drop(context);
@@ -303,7 +310,7 @@ impl Client {
token,
peer
);
match maybe_download {
return match maybe_download {
Some(download) => {
thread::spawn(move || {
let download_peer = DownloadPeer::new(
@@ -322,7 +329,8 @@ impl Client {
Some(filename) => {
match download_peer.download_file(
client_context_clone,
Some(download.clone())
Some(download.clone()),
None
) {
Ok((download, filename)) => {
download.sender.send(DownloadStatus::Completed).unwrap();
@@ -347,7 +355,7 @@ impl Client {
token
);
}
}
};
}
ClientOperation::NewPeer(new_peer) => {
if client_context
@@ -507,7 +515,7 @@ impl Client {
DefaultPeer::new(peer_clone, sender_clone);
let connect_result = match stream {
Some(s) => default_peer.connect_with_socket(s),
Some(s) => default_peer.connect_with_socket(s, None),
_ => default_peer.connect()
};
@@ -543,6 +551,7 @@ impl Client {
match download_peer
.download_file(
client_context2.clone(),
None,
None
) {
Ok((download, filename)) => {
src/message/message_reader.rs
+4
View File
@@ -63,6 +63,10 @@ impl MessageReader {
Ok(())
}
pub fn buffer_len(&self) -> usize {
self.buffer.len()
}
pub fn extract_message(&mut self) -> io::Result<Option<Message>> {
let bytes_read = self.buffer.len();
if bytes_read < 4 {
src/message/peer/peer_init.rs
+10
View File
@@ -3,6 +3,7 @@ use std::sync::mpsc::Sender;
use crate::{
message::{Message, MessageHandler},
peer::PeerOperation,
trace,
};
pub struct PeerInit;
@@ -15,6 +16,15 @@ impl MessageHandler<PeerOperation> for PeerInit {
message.set_pointer(4);
let _message_code = message.read_int8();
let username = message.read_string();
let connection_type = message.read_string();
let token = message.read_int32();
trace!(
"PeerInit: username: {}, connection_type: {}, token: {}",
username,
connection_type,
token
);
sender.send(PeerOperation::SetUsername(username)).unwrap();
}
}
src/peer/default_peer.rs
+91 -43
View File
@@ -74,6 +74,7 @@ impl DefaultPeer {
pub fn connect_with_socket(
mut self,
mut stream: TcpStream,
reader: Option<MessageReader>,
) -> Result<Self, io::Error> {
if let Some(token) = self.peer.read().unwrap().token {
let mut message: Vec<u8> = [0, 5, 0, 0, 0, 0].to_vec();
@@ -89,7 +90,7 @@ impl DefaultPeer {
"[default_peer:{}] connect_with_socket: direct",
self.peer.read().unwrap().username
);
self.start_read_write_loops(stream)?;
self.start_read_write_loops(stream, reader)?;
Ok(self)
}
@@ -130,7 +131,7 @@ impl DefaultPeer {
.write_all(&MessageFactory::build_watch_user(token).get_data())
.unwrap();
}
self.start_read_write_loops(stream)?;
self.start_read_write_loops(stream, None)?;
Ok(self)
}
@@ -138,6 +139,7 @@ impl DefaultPeer {
fn start_read_write_loops(
&mut self,
stream: TcpStream,
reader: Option<MessageReader>,
) -> Result<(), io::Error> {
debug!(
"[default_peer:{}] start_read_write_loops",
@@ -148,17 +150,38 @@ impl DefaultPeer {
Receiver<PeerOperation>,
) = mpsc::channel();
// Set the peer_channel so transfer_request can send messages
self.peer_channel = Some(peer_sender.clone());
let mut read_stream = stream.try_clone()?;
let mut write_stream = stream; // Use the original stream for writing
let mut write_stream = match stream.try_clone() {
Ok(s) => s,
Err(e) => {
error!(
"[default_peer:{}] Failed to clone stream: {}",
self.peer.read().unwrap().username,
e
);
return Err(e);
}
};
let mut read_stream = stream;
debug!(
"[default_peer:{}] Streams created, spawning threads",
self.peer.read().unwrap().username
);
let peer = self.peer.clone();
let peer_clone = self.peer.clone();
let client_channel_for_read = self.client_channel.clone();
let mut buffered_reader = reader.unwrap_or_default();
let buffer_size = buffered_reader.buffer_len();
self.read_thread = Some(thread::spawn(move || {
debug!(
"Read thread received reader with {} bytes in buffer",
buffer_size
);
let mut handlers = Handlers::new();
handlers.register_handler(FileSearchResponse);
handlers.register_handler(TransferRequest);
@@ -173,34 +196,20 @@ impl DefaultPeer {
peer_sender,
handlers,
);
trace!(
"[default_peer:{}] read thread starting",
peer.read().unwrap().username
);
if buffer_size > 0 {
trace!(
"[default_peer:{}] starting with {} buffered bytes",
peer.read().unwrap().username,
buffer_size
);
}
let mut buffered_reader = MessageReader::new();
loop {
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::TimedOut => {
{
let peer_lock = peer.read().unwrap();
debug!(
"Read operation timed out in default peer {:}:{:}",
peer_lock.host, peer_lock.port
);
}
continue;
}
Err(e) => {
let username = peer.read().unwrap().username.clone();
error!("[default_peer:{}] Error reading from peer: {}. Terminating read loop.",username, e);
let _ = client_channel_for_read
.send(ClientOperation::PeerDisconnected(username));
break;
}
}
// Extract all available messages from buffer
let mut should_terminate = false;
loop {
match buffered_reader.extract_message() {
@@ -230,13 +239,45 @@ impl DefaultPeer {
should_terminate = true;
break;
}
Ok(_) => break,
Ok(None) => {
break;
}
}
}
if should_terminate {
break;
}
match buffered_reader.read_from_socket(&mut read_stream) {
Ok(()) => continue,
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
continue
}
Err(ref e) if e.kind() == io::ErrorKind::TimedOut => {
{
let peer_lock = peer.read().unwrap();
debug!(
"Read operation timed out in default peer {:}:{:}",
peer_lock.host, peer_lock.port
);
}
continue;
}
Err(e) => {
let username = peer.read().unwrap().username.clone();
error!("[default_peer:{}] Error reading from peer: {} (kind: {:?}). Terminating read loop.", username, e, e.kind());
let _ = client_channel_for_read
.send(ClientOperation::PeerDisconnected(username));
break;
}
}
}
trace!(
"[default_peer:{}] read thread exiting",
peer.read().unwrap().username
);
}));
let client_channel = self.client_channel.clone();
@@ -325,18 +366,18 @@ impl DefaultPeer {
reason,
} => {
debug!(
"[default_peer:{}] transfer response token: {} allowed: {}",
peer_clone_for_write.read().unwrap().username, token, allowed
);
"[default_peer:{}] transfer response token: {} allowed: {}",
peer_clone_for_write.read().unwrap().username, token, allowed
);
if !allowed {
if let Some(reason_text) = reason {
debug!(
"[default_peer:{}] Transfer rejected: {} - token {}, I will receive TransferRequest soon...",
peer_clone_for_write.read().unwrap().username.clone(),
reason_text,
token
);
"[default_peer:{}] Transfer rejected: {} - token {}, I will receive TransferRequest soon...",
peer_clone_for_write.read().unwrap().username.clone(),
reason_text,
token
);
}
} else {
debug!("[default_peer:{}] Transfer allowed, ready to connect with token {:}",peer_clone_for_write.read().unwrap().username, token);
@@ -350,9 +391,9 @@ impl DefaultPeer {
}
PeerOperation::PlaceInQueueResponse { filename, place } => {
debug!(
"[default_peer:{}] Place in queue response - file: {}, place: {}",
peer_clone_for_write.read().unwrap().username, filename, place
);
"[default_peer:{}] Place in queue response - file: {}, place: {}",
peer_clone_for_write.read().unwrap().username, filename, place
);
}
PeerOperation::SetUsername(username) => {
trace!(
@@ -374,6 +415,13 @@ impl DefaultPeer {
Ok(())
}
pub fn queue_upload(&self, filename: String) -> Result<(), io::Error> {
let message = MessageFactory::build_queue_upload_message(&filename);
if let Some(sender) = &self.peer_channel {
sender.send(PeerOperation::SendMessage(message)).unwrap();
}
Ok(())
}
pub fn transfer_request(
&self,
download: Download,
src/peer/download_peer.rs
+28 -46
View File
@@ -236,8 +236,17 @@ impl DownloadPeer {
self,
client_context: Arc<RwLock<ClientContext>>,
mut download: Option<Download>,
stream: Option<TcpStream>,
) -> Result<(Download, String), io::Error> {
let mut stream = self.establish_connection()?;
if stream.is_none() {
trace!("[download_peer:{}] stream is none", self.username);
}
let mut stream = stream.unwrap_or_else(|| {
self.establish_connection().unwrap_or_else(|e| {
panic!("Failed to establish connection: {:?}", e)
})
});
trace!("[download_peer:{}] connected", self.username);
self.perform_handshake(&mut stream)?;
@@ -270,11 +279,15 @@ impl DownloadPeer {
&& 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());
let token_u32 = u32::from_le_bytes(
token
.try_into()
.unwrap_or_else(|_| panic!("[download_peer:{}] slice with incorrect length", self.username)),
);
trace!(
"[download_peer:{}] received token: {:?} ",
"[download_peer:{}] received token_u32: {:?} - token: {:?}",
self.username,
token_u32,
token
);
@@ -297,8 +310,12 @@ impl DownloadPeer {
download = Some(d.clone());
}
None => {
let tokens = read
.download_tokens
.keys()
.collect::<Vec<_>>();
panic!(
"No download info for token {token_u32}"
"[download_peer:{}] No download info for token {token_u32}, tokens: {:?}", self.username, tokens
);
}
}
@@ -427,7 +444,9 @@ mod tests {
"own_username".to_string(),
);
let dummy_context = Arc::new(RwLock::new(ClientContext::new()));
let _ = download_peer.download_file(dummy_context).unwrap();
let _ = download_peer
.download_file(dummy_context, None, None)
.unwrap();
// Give the client time to send messages
thread::sleep(Duration::from_millis(10));
@@ -450,7 +469,9 @@ mod tests {
"own_username".to_string(),
);
let dummy_context = Arc::new(RwLock::new(ClientContext::new()));
let _ = download_peer.download_file(dummy_context).unwrap();
let _ = download_peer
.download_file(dummy_context, None, None)
.unwrap();
// Give the client time to send messages
thread::sleep(Duration::from_millis(10));
@@ -611,43 +632,4 @@ mod tests {
"file.mp3"
);
}
#[test]
fn test_create_download_path_from_filename() {
// Test with filename provided
let path = FileManager::create_download_path_from_filename(
Some("/downloads"),
"remote_user",
123,
Some("song.mp3"),
);
assert_eq!(path, "/downloads/song.mp3");
// Test with full path filename
let path = FileManager::create_download_path_from_filename(
Some("/downloads"),
"remote_user",
123,
Some("/remote/path/to/song.mp3"),
);
assert_eq!(path, "/downloads/song.mp3");
// Test without filename (fallback)
let path = FileManager::create_download_path_from_filename(
Some("/downloads"),
"remote_user",
123,
None,
);
assert_eq!(path, "/downloads/remote_user_123.mp3");
// Test with default directory
let path = FileManager::create_download_path_from_filename(
None,
"remote_user",
123,
Some("song.mp3"),
);
assert_eq!(path, "/tmp/song.mp3");
}
}
src/peer/listen.rs
+115 -23
View File
@@ -1,42 +1,134 @@
use std::net::TcpListener;
use std::sync::mpsc::Sender;
use std::sync::{Arc, RwLock};
use std::thread;
use crate::client::ClientOperation;
use crate::client::{ClientContext, ClientOperation};
use crate::peer::{DefaultPeer, Peer};
use crate::{info, trace};
use crate::message::MessageReader;
use crate::peer::{ConnectionType, DefaultPeer, DownloadPeer, Peer};
use crate::{debug, error, info, trace, DownloadStatus};
pub struct Listen {}
impl Listen {
pub fn start(port: u32, client_sender: Sender<ClientOperation>) {
let mut index = 0;
pub fn start(
port: u32,
client_sender: Sender<ClientOperation>,
client_context: Arc<RwLock<ClientContext>>,
own_username: String,
) {
info!("starting listener on port {port}");
let listener = TcpListener::bind(format!("0.0.0.0:{port}")).unwrap();
for stream in listener.incoming() {
index += 1;
let client_sender_clone = client_sender.clone();
let client_sender = client_sender.clone();
let client_context = client_context.clone();
let own_username = own_username.clone();
thread::spawn(move || {
let stream = stream.unwrap();
let mut stream = match stream {
Ok(s) => s,
Err(e) => {
trace!("Failed to accept connection: {}", e);
return;
}
};
let peer_ip = stream.peer_addr().unwrap().ip().to_string();
let peer_port = stream.peer_addr().unwrap().port();
let peer = Peer::new(
format!("unknown-{index}"),
crate::peer::ConnectionType::P,
peer_ip.clone(),
port,
None,
0,
0,
0,
);
let default_peer = DefaultPeer::new(peer, client_sender_clone);
let mut reader = MessageReader::new();
trace!("[listener{index}] new connection from {peer_ip}");
default_peer.connect_with_socket(stream).unwrap();
let mut message = loop {
if let Err(e) = reader.read_from_socket(&mut stream) {
trace!("Failed to read from socket: {}", e);
return;
}
if let Ok(Some(msg)) = reader.extract_message() {
break msg;
}
};
message.set_pointer(4);
let message_code = message.read_int8();
if message_code == 1 {
let username = message.read_string();
let connection_type = message.read_string();
let token = message.read_int32();
debug!("[listener:{peer_ip}:{peer_port}] peerInit (0) username: {username} connection_type: {connection_type} token: {token}");
let peer = Peer::new(
format!("{}-direct", username.clone()),
ConnectionType::P,
peer_ip.clone(),
peer_port.into(),
None,
0,
0,
0,
);
if connection_type == "P" {
debug!("[listener:{peer_ip}:{peer_port}] connection type is P, reader buffer has {} bytes", reader.buffer_len());
let default_peer =
DefaultPeer::new(peer, client_sender);
let default_peer = default_peer
.connect_with_socket(stream, Some(reader))
.unwrap();
drop(default_peer);
} else if connection_type == "F" {
let maybe_download = {
let client_context = client_context.read().unwrap();
client_context.download_tokens.get(&token).cloned()
};
let client_context_clone = client_context.clone();
trace!(
"[client] DownloadFromPeer token: {} peer: {:?}",
token,
peer
);
thread::spawn(move || {
let download_peer = DownloadPeer::new(
format!("{}-direct", username.clone()),
peer.host.clone(),
peer.port,
token,
false,
own_username,
);
match download_peer.download_file(
client_context_clone,
maybe_download,
Some(stream),
) {
Ok((download, filename)) => {
download
.sender
.send(DownloadStatus::Completed)
.unwrap();
info!("Successfully downloaded {} bytes to {}", download.size, filename);
}
Err(e) => {
error!(
"Failed to download file from {}:{} (token: {}) - Error: {}",
peer.host, peer.port, token, e
);
}
}
});
} else {
debug!(
"[listener:{peer_ip}:{peer_port}] connection type is not P or F"
);
}
} else {
debug!("[listener:{peer_ip}:{peer_port}] unknown message with code: {message_code}");
}
});
}
}