chore: refactor default peer

This commit is contained in:
Michel de Graaf
2025-11-04 18:48:32 +01:00
parent c346c56297
commit d11c6782a7
8 changed files with 154 additions and 168 deletions
+4 -3
View File
@@ -310,7 +310,7 @@ impl Client {
download.username.clone(),
peer.host.clone(),
peer.port,
download.token,
token,
false,
own_username,
);
@@ -500,7 +500,7 @@ impl Client {
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);
trace!("[client] connecting to {}, with connection_type: {}, and token {:?}", peer.username, peer.connection_type, peer.token);
match peer.connection_type {
ConnectionType::P => {
let default_peer =
@@ -530,7 +530,7 @@ impl Client {
}
ConnectionType::F => {
trace!("[client] downloading from: {}", peer.username);
trace!("[client] downloading from: {}, {:?}", peer.username, peer.token);
let download_peer = DownloadPeer::new(
peer.username,
peer.host,
@@ -547,6 +547,7 @@ impl Client {
) {
Ok((download, filename)) => {
trace!("[client] downloaded {} bytes {:?} ", filename, download.size);
download.sender.send(DownloadStatus::Completed).unwrap();
}
Err(e) => {
trace!("[client] failed to download: {}", e);
+1
View File
@@ -358,6 +358,7 @@ impl Message {
))),
},
MessageType::Peer => match code {
1 => Ok("PeerInit"),
4 => Ok("GetShareFileList"),
5 => Ok("SharedFileListResponse"),
9 => Ok("FileSearchResponse"),
+2
View File
@@ -1,5 +1,6 @@
mod file_search_response;
mod get_share_file_list;
mod peer_init;
mod place_in_queue_response;
mod transfer_request;
mod transfer_response;
@@ -8,6 +9,7 @@ mod upload_failed;
// Re-export handlers
pub use file_search_response::FileSearchResponse;
pub use get_share_file_list::GetShareFileList;
pub use peer_init::PeerInit;
pub use place_in_queue_response::PlaceInQueueResponse;
pub use transfer_request::TransferRequest;
pub use transfer_response::TransferResponse;
+20
View File
@@ -0,0 +1,20 @@
use std::sync::mpsc::Sender;
use crate::{
message::{Message, MessageHandler},
peer::PeerOperation,
};
pub struct PeerInit;
impl MessageHandler<PeerOperation> for PeerInit {
fn get_code(&self) -> u8 {
1
}
fn handle(&self, message: &mut Message, sender: Sender<PeerOperation>) {
message.set_pointer(4);
let _message_code = message.read_int8();
let username = message.read_string();
sender.send(PeerOperation::SetUsername(username)).unwrap();
}
}
-2
View File
@@ -3,12 +3,10 @@ use std::sync::mpsc::Sender;
use crate::{message::Message, server::ServerOperation};
// IMPORTANT: note the generic trait is now `MessageHandler<ServerOperation>`
use crate::message::MessageHandler;
pub struct LoginHandler;
// Implement `MessageHandler` specifically for `ServerOperation`.
impl MessageHandler<ServerOperation> for LoginHandler {
fn get_code(&self) -> u8 {
1
+89 -59
View File
@@ -1,5 +1,5 @@
use crate::dispatcher::MessageDispatcher;
use crate::message::peer::FileSearchResponse;
use crate::message::peer::{FileSearchResponse, PeerInit};
use crate::message::peer::{
GetShareFileList, PlaceInQueueResponse, TransferRequest, TransferResponse,
UploadFailedHandler,
@@ -9,8 +9,8 @@ use crate::message::{Handlers, Message, MessageReader, MessageType};
use crate::types::{Download, FileSearchResult, Transfer};
use core::result::Result::Ok;
use std::sync::mpsc;
use std::sync::mpsc::{Receiver, Sender};
use std::sync::{mpsc, Arc, RwLock};
use std::thread::{self, JoinHandle};
use crate::client::ClientOperation;
@@ -24,7 +24,7 @@ use std::time::Duration;
#[allow(dead_code)]
#[derive(Debug)]
pub struct DefaultPeer {
peer: Peer,
peer: Arc<RwLock<Peer>>,
peer_channel: Option<Sender<PeerOperation>>,
client_channel: Sender<ClientOperation>,
read_thread: Option<JoinHandle<()>>,
@@ -45,11 +45,13 @@ pub enum PeerOperation {
filename: String,
place: u32,
},
SetUsername(String),
}
impl DefaultPeer {
pub fn new(peer: Peer, client_channel: Sender<ClientOperation>) -> Self {
Self {
peer,
peer: Arc::new(RwLock::new(peer)),
peer_channel: None,
client_channel,
read_thread: None,
@@ -57,10 +59,12 @@ impl DefaultPeer {
}
}
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()),
) {
let username = self.peer.read().unwrap().username.clone();
debug!("[default_peer:{}] disconnect", username);
if let Err(e) = self
.client_channel
.send(ClientOperation::PeerDisconnected(username))
{
error!("Failed to send disconnect notification: {}", e);
}
@@ -71,7 +75,7 @@ impl DefaultPeer {
mut self,
mut stream: TcpStream,
) -> Result<Self, io::Error> {
if let Some(token) = self.peer.token {
if let Some(token) = self.peer.read().unwrap().token {
let mut message: Vec<u8> = [0, 5, 0, 0, 0, 0].to_vec();
message.extend_from_slice(&token.to_le_bytes());
stream.write_all(&message).unwrap();
@@ -83,25 +87,35 @@ impl DefaultPeer {
trace!(
"[default_peer:{}] connect_with_socket: direct",
self.peer.username
self.peer.read().unwrap().username
);
self.start_read_write_loops(stream)?;
Ok(self)
}
pub fn connect(mut self) -> Result<Self, io::Error> {
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()
.ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "Invalid address")
})?;
trace!("[default_peer:{}] connect: indirect", self.peer.username);
pub fn connect(mut self) -> Result<Self, io::Error> {
{
let peer = self.peer.read().unwrap();
println!(
"[default_peer] Connecting to {} on port {}",
peer.host, peer.port
);
}
let socket_address = {
let peer = self.peer.read().unwrap();
format!("{}:{}", peer.host, peer.port)
}
.to_socket_addrs()?
.next()
.ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "Invalid address")
})?;
trace!(
"[default_peer:{}] connect: indirect",
self.peer.read().unwrap().username
);
let mut stream = TcpStream::connect_timeout(
&socket_address,
Duration::from_secs(20),
@@ -111,7 +125,7 @@ impl DefaultPeer {
stream.set_write_timeout(Some(Duration::from_secs(5)))?;
stream.set_nodelay(true)?;
if let Some(token) = self.peer.token {
if let Some(token) = self.peer.read().unwrap().token {
stream
.write_all(&MessageFactory::build_watch_user(token).get_data())
.unwrap();
@@ -127,7 +141,7 @@ impl DefaultPeer {
) -> Result<(), io::Error> {
debug!(
"[default_peer:{}] start_read_write_loops",
self.peer.username
self.peer.read().unwrap().username
);
let (peer_sender, peer_reader): (
Sender<PeerOperation>,
@@ -152,6 +166,7 @@ impl DefaultPeer {
handlers.register_handler(GetShareFileList);
handlers.register_handler(UploadFailedHandler);
handlers.register_handler(PlaceInQueueResponse);
handlers.register_handler(PeerInit);
let dispatcher = MessageDispatcher::new(
"default_peer".to_string(),
@@ -167,19 +182,20 @@ impl DefaultPeer {
continue
}
Err(ref e) if e.kind() == io::ErrorKind::TimedOut => {
debug!(
"Read operation timed out in default peer {:}:{:}",
peer.host, peer.port
);
{
let peer_lock = peer.read().unwrap();
debug!(
"Read operation timed out in default peer {:}:{:}",
peer_lock.host, peer_lock.port
);
}
continue;
}
Err(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(),
),
);
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;
}
}
@@ -191,7 +207,7 @@ impl DefaultPeer {
Ok(Some(mut message)) => {
trace!(
"[default_peer:{}] ← {:?}",
peer.username,
peer.read().unwrap().username,
message
.get_message_name(
MessageType::Peer,
@@ -206,10 +222,10 @@ impl DefaultPeer {
"Error extracting message in default peer: {}. Terminating read loop.",
e
);
let username =
peer.read().unwrap().username.clone();
let _ = client_channel_for_read.send(
ClientOperation::PeerDisconnected(
peer.username.clone(),
),
ClientOperation::PeerDisconnected(username),
);
should_terminate = true;
break;
@@ -225,7 +241,7 @@ impl DefaultPeer {
let client_channel = self.client_channel.clone();
let peer_channel = self.peer_channel.clone();
let peer_username = self.peer.username.clone();
let peer_clone_for_write = self.peer.clone();
self.write_thread = Some(thread::spawn(move || loop {
match peer_reader.recv() {
@@ -233,7 +249,7 @@ impl DefaultPeer {
PeerOperation::SendMessage(message) => {
trace!(
"[default_peer:{}] ➡ {:?} - {:?}",
peer_username,
peer_clone_for_write.read().unwrap().username,
message
.get_message_name(
MessageType::Peer,
@@ -253,10 +269,14 @@ impl DefaultPeer {
if let Err(e) =
write_stream.write_all(&message.get_buffer())
{
error!("Error writing message to stream: {} - {}. Terminating write loop.", peer_username, e);
error!("Error writing message to stream: {} - {}. Terminating write loop.", peer_clone_for_write.read().unwrap().username, e);
let _ = client_channel.send(
ClientOperation::PeerDisconnected(
peer_username.clone(),
peer_clone_for_write
.read()
.unwrap()
.username
.clone(),
),
);
break;
@@ -271,12 +291,13 @@ impl DefaultPeer {
PeerOperation::TransferRequest(transfer) => {
debug!(
"[default_peer:{}] TransferRequest for {}",
peer_username, transfer.token
peer_clone_for_write.read().unwrap().username,
transfer.token
);
client_channel
.send(ClientOperation::UpdateDownloadTokens(
transfer.clone(),
peer_clone.username.clone(),
peer_clone.read().unwrap().username.clone(),
))
.unwrap();
@@ -287,7 +308,7 @@ impl DefaultPeer {
trace!(
"[default_peer:{}] TransferResponse for {:?}",
peer_username,
peer_clone_for_write.read().unwrap().username,
transfer_response.get_buffer()
);
if let Some(sender) = peer_channel.clone() {
@@ -304,38 +325,47 @@ impl DefaultPeer {
reason,
} => {
debug!(
"[default_peer:{}] transfer response token: {} allowed: {}",
peer_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_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_username, token);
debug!("[default_peer:{}] Transfer allowed, ready to connect with token {:}",peer_clone_for_write.read().unwrap().username, token);
client_channel
.send(ClientOperation::DownloadFromPeer(
token,
peer_clone.clone(),
peer_clone.read().unwrap().clone(),
))
.unwrap();
}
}
PeerOperation::PlaceInQueueResponse { filename, place } => {
debug!(
"[default_peer:{}] Place in queue response - file: {}, place: {}",
peer_username, filename, place
"[default_peer:{}] Place in queue response - file: {}, place: {}",
peer_clone_for_write.read().unwrap().username, filename, place
);
}
PeerOperation::SetUsername(username) => {
trace!(
"[default_peer:{}] SetUsername: {}",
peer_clone_for_write.read().unwrap().username,
username
);
peer_clone.write().unwrap().username = username;
}
},
Err(_) => {
debug!("[default_peer:{}] Peer channel closed. Terminating write loop.", peer_username);
debug!("[default_peer:{}] Peer channel closed. Terminating write loop.", peer_clone_for_write.read().unwrap().username);
break;
}
}
@@ -362,7 +392,7 @@ impl DefaultPeer {
impl Drop for DefaultPeer {
fn drop(&mut self) {
self.peer_channel = None;
trace!("[default_peer:{}] drop", self.peer.username);
trace!("[default_peer:{}] drop", self.peer.read().unwrap().username);
if let Some(handle) = self.read_thread.take() {
debug!("Joining read thread...");
+13 -10
View File
@@ -217,16 +217,15 @@ impl DownloadPeer {
stream.write_all(&message.get_buffer())?;
stream.flush()?;
} else {
stream.write_all(
&MessageFactory::build_pierce_firewall_message(self.token)
.get_buffer(),
)?;
let message =
MessageFactory::build_pierce_firewall_message(self.token);
stream.write_all(&message.get_buffer())?;
trace!(
"[default_peer:{}] sending pierce firewall message token: {}: {:?}",
self.username,
self.token,
MessageFactory::build_pierce_firewall_message(self.token)
.get_buffer()
&message.get_buffer()
);
stream.flush()?;
}
@@ -257,8 +256,9 @@ impl DownloadPeer {
// connection closed
Ok(0) => {
trace!(
"[download_peer:{}] connection closed by peer",
self.username
"[download_peer:{}] connection closed by peer. bytes read: {}",
self.username,
processor.total_bytes
);
break;
}
@@ -323,10 +323,13 @@ impl DownloadPeer {
"[download_peer:{}] finished reading data from peer",
self.username
);
let read = client_context.read().unwrap();
let tokens = read.download_tokens.keys().collect::<Vec<_>>();
let download = download.clone().unwrap_or_else(|| {
panic!(
"[download_peer] No download info found for token: {}",
self.token
"[download_peer] No download info found for token: {}, tokens: {:?} ",
self.token, tokens
)
});
+25 -94
View File
@@ -1,112 +1,43 @@
use std::io::Write;
use std::net::TcpListener;
use std::sync::mpsc::Sender;
use std::{io, net::TcpListener};
use std::thread;
use crate::client::ClientOperation;
use crate::message::server::MessageFactory;
use crate::message::MessageReader;
use crate::{debug, error, info, trace, warn, FileSearchResult};
use crate::peer::{DefaultPeer, Peer};
use crate::{info, trace};
pub struct Listen {}
impl Listen {
pub fn start(port: u32, client_sender: Sender<ClientOperation>) {
let mut index = 0;
info!("starting listener on port {port}");
let listener = TcpListener::bind(format!("0.0.0.0:{port}")).unwrap();
for stream in listener.incoming() {
let mut stream = stream.unwrap();
index += 1;
let client_sender_clone = client_sender.clone();
thread::spawn(move || {
let stream = stream.unwrap();
let mut buffered_reader = MessageReader::new();
loop {
match buffered_reader.read_from_socket(&mut stream) {
Ok(()) => {}
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;
}
Err(e) => {
error!("Error reading from server: {}", e);
break;
}
}
let peer_ip = stream.peer_addr().unwrap().ip().to_string();
loop {
match buffered_reader.extract_message() {
Ok(Some(mut message)) => {
message.set_pointer(4);
let message_code = message.read_int8();
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);
trace!(
"[listener] Received message with code: {}",
message_code
);
match message_code {
0 => {
let token = message.read_int32();
debug!("[listener] received Pierce Firewall, token: {}", token);
stream
.write_all(
&MessageFactory::build_pierce_firewall_message(
token,
)
.get_data(),
)
.unwrap();
stream.flush().unwrap();
}
1 => {
// handle handover to default_peer with socket in the future
// let tcp_stream =
// read_stream.try_clone().unwrap();
//
// let username = message.read_string();
// let connection_type =
// message.read_string().parse().unwrap();
// let token = message.read_int32();
}
9 => {
trace!("[listener] Handling file search response: 9",);
message.set_pointer(8);
let file_search =
match FileSearchResult::new_from_message(
&mut message,
) {
Ok(result) => result,
Err(e) => {
trace!("[listener] malformed filesearch_result: {:?}, message: {:?}", e, message);
return;
}
};
client_sender
.send(ClientOperation::SearchResult(
file_search,
))
.unwrap();
break;
}
_ => {
debug!(
"[listener] unknown message_code: {}",
message_code
);
}
}
}
Err(e) => {
warn!("[listener] Error extracting message: {}", e);
break;
}
Ok(None) => break,
}
}
}
// info!("[listener] Connection established!");
trace!("[listener{index}] new connection from {peer_ip}");
default_peer.connect_with_socket(stream).unwrap();
});
}
}
}