chore: refactor peer connection to actor and non blocking
This commit is contained in:
Michel de Graaf
@@ -1,344 +0,0 @@
|
||||
use crate::actor::{Actor, ActorSystem};
|
||||
use crate::client::ClientContext;
|
||||
use crate::message::{Message, MessageReader};
|
||||
use crate::peer::{ConnectionType, DownloadPeer, Peer};
|
||||
use crate::types::Download;
|
||||
use crate::{debug, error, info, trace, DownloadStatus};
|
||||
|
||||
use std::io;
|
||||
use std::net::{TcpListener, TcpStream};
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
const PEER_INIT_MESSAGE_CODE: u8 = 1;
|
||||
|
||||
struct PeerInitData {
|
||||
username: String,
|
||||
connection_type: ConnectionType,
|
||||
token: u32,
|
||||
}
|
||||
|
||||
fn read_peer_init_message(
|
||||
stream: &mut TcpStream,
|
||||
reader: &mut MessageReader,
|
||||
) -> io::Result<Message> {
|
||||
loop {
|
||||
reader.read_from_socket(stream)?;
|
||||
|
||||
if let Ok(Some(msg)) = reader.extract_message() {
|
||||
return Ok(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_peer_init_message(mut message: Message) -> Option<PeerInitData> {
|
||||
message.set_pointer(4);
|
||||
let message_code = message.read_int8();
|
||||
|
||||
if message_code != PEER_INIT_MESSAGE_CODE {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(PeerInitData {
|
||||
username: message.read_string(),
|
||||
connection_type: message.read_string().parse().unwrap(),
|
||||
token: message.read_int32(),
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_token_from_buffer(buffer: &[u8], username: &str) -> Option<u32> {
|
||||
let token_bytes = buffer.get(0..4)?;
|
||||
let token = u32::from_le_bytes(
|
||||
token_bytes
|
||||
.try_into()
|
||||
.unwrap_or_else(|_| panic!("[listener:{}] slice with incorrect length, can't extract transfer_token", username)),
|
||||
);
|
||||
Some(token)
|
||||
}
|
||||
|
||||
fn extract_download_from_buffer(
|
||||
reader: &mut MessageReader,
|
||||
client_context: &Arc<RwLock<ClientContext>>,
|
||||
username: &str,
|
||||
peer_ip: &str,
|
||||
peer_port: u16,
|
||||
) -> Option<Download> {
|
||||
if reader.buffer_len() == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let buffer = reader.get_buffer();
|
||||
trace!(
|
||||
"[listener:{peer_ip}:{peer_port}] reader buffer has {} bytes. {:?}",
|
||||
buffer.len(),
|
||||
buffer
|
||||
);
|
||||
|
||||
let token = parse_token_from_buffer(&buffer, username)?;
|
||||
trace!(
|
||||
"[listener:{}] got transfer_token: {} from data chunk",
|
||||
username,
|
||||
token
|
||||
);
|
||||
|
||||
let context = client_context.read().unwrap();
|
||||
let download = context.download_tokens.get(&token).cloned();
|
||||
|
||||
if download.is_none() {
|
||||
let download_tokens: Vec<u32> = context.download_tokens.keys().cloned().collect();
|
||||
trace!("[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}", token, download_tokens);
|
||||
}
|
||||
|
||||
download
|
||||
}
|
||||
|
||||
/// Message types for ConnectionHandlerActor
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ConnectionHandlerMessage {
|
||||
/// Process the incoming connection
|
||||
Process,
|
||||
}
|
||||
|
||||
/// Actor that handles a single incoming connection
|
||||
pub struct ConnectionHandlerActor {
|
||||
stream: Option<TcpStream>,
|
||||
client_context: Arc<RwLock<ClientContext>>,
|
||||
own_username: String,
|
||||
peer_ip: String,
|
||||
peer_port: u16,
|
||||
}
|
||||
|
||||
impl ConnectionHandlerActor {
|
||||
pub fn new(
|
||||
stream: TcpStream,
|
||||
client_context: Arc<RwLock<ClientContext>>,
|
||||
own_username: String,
|
||||
) -> Self {
|
||||
let peer_addr = stream.peer_addr().unwrap();
|
||||
let peer_ip = peer_addr.ip().to_string();
|
||||
let peer_port = peer_addr.port();
|
||||
|
||||
Self {
|
||||
stream: Some(stream),
|
||||
client_context,
|
||||
own_username,
|
||||
peer_ip,
|
||||
peer_port,
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_peer_connection(
|
||||
&self,
|
||||
peer: Peer,
|
||||
stream: TcpStream,
|
||||
reader: MessageReader,
|
||||
) {
|
||||
debug!("[listener:{}:{}] connection type is P, reader buffer has {} bytes",
|
||||
self.peer_ip, self.peer_port, reader.buffer_len());
|
||||
|
||||
let client_context = self.client_context.read().unwrap();
|
||||
if let Some(ref registry) = client_context.peer_registry {
|
||||
match registry.register_peer(peer.clone(), stream, Some(reader)) {
|
||||
Ok(_) => trace!("[listener] peer actor spawned for: {}", peer.username),
|
||||
Err(e) => error!("Failed to spawn peer actor for {:?}: {:?}", peer.username, e),
|
||||
}
|
||||
} else {
|
||||
error!("PeerRegistry not initialized");
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_file_connection(
|
||||
&self,
|
||||
peer: Peer,
|
||||
stream: TcpStream,
|
||||
mut reader: MessageReader,
|
||||
token: u32,
|
||||
) {
|
||||
trace!(
|
||||
"[client] DownloadFromPeer token: {} peer: {:?}",
|
||||
token,
|
||||
peer
|
||||
);
|
||||
|
||||
// Set stream to blocking mode - DownloadPeer expects blocking I/O
|
||||
if let Err(e) = stream.set_nonblocking(false) {
|
||||
error!(
|
||||
"Failed to set stream to blocking mode for {}:{} (token: {}) - Error: {}",
|
||||
peer.host, peer.port, token, e
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let download = extract_download_from_buffer(
|
||||
&mut reader,
|
||||
&self.client_context,
|
||||
&peer.username,
|
||||
&self.peer_ip,
|
||||
self.peer_port,
|
||||
);
|
||||
|
||||
let download_peer = DownloadPeer::new(
|
||||
format!("{}:direct", peer.username),
|
||||
peer.host.clone(),
|
||||
peer.port,
|
||||
token,
|
||||
true,
|
||||
self.own_username.clone(),
|
||||
);
|
||||
|
||||
match download_peer.download_file(
|
||||
self.client_context.clone(),
|
||||
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
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn process_connection(&mut self) {
|
||||
let Some(mut stream) = self.stream.take() else {
|
||||
error!("[listener] stream already taken");
|
||||
return;
|
||||
};
|
||||
|
||||
let mut reader = MessageReader::new();
|
||||
|
||||
let Ok(message) = read_peer_init_message(&mut stream, &mut reader) else {
|
||||
error!(
|
||||
"[listener:{}:{}] Failed to read peer init message",
|
||||
self.peer_ip, self.peer_port
|
||||
);
|
||||
return;
|
||||
};
|
||||
|
||||
let Some(init_data) = parse_peer_init_message(message) else {
|
||||
error!("[listener:{}:{}] Invalid or unknown peer init message",
|
||||
self.peer_ip, self.peer_port);
|
||||
return;
|
||||
};
|
||||
|
||||
debug!(
|
||||
"[listener:{}:{}] peerInit (0) username: {} connection_type: {} token: {}",
|
||||
self.peer_ip, self.peer_port, init_data.username, init_data.connection_type, init_data.token
|
||||
);
|
||||
|
||||
let peer = Peer::new(
|
||||
format!("{}:direct", init_data.username),
|
||||
init_data.connection_type.clone(),
|
||||
self.peer_ip.clone(),
|
||||
self.peer_port.into(),
|
||||
None,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
);
|
||||
|
||||
match init_data.connection_type {
|
||||
ConnectionType::P => self.handle_peer_connection(peer, stream, reader),
|
||||
ConnectionType::F => self.handle_file_connection(peer, stream, reader, init_data.token),
|
||||
ConnectionType::D => {
|
||||
debug!("[listener:{}:{}] connection type is D, not supported yet, closing connection.",
|
||||
self.peer_ip, self.peer_port);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Actor for ConnectionHandlerActor {
|
||||
type Message = ConnectionHandlerMessage;
|
||||
|
||||
fn handle(&mut self, msg: Self::Message) {
|
||||
match msg {
|
||||
ConnectionHandlerMessage::Process => self.process_connection(),
|
||||
}
|
||||
}
|
||||
|
||||
fn on_start(&mut self) {
|
||||
trace!("[connection_handler:{}:{}] starting", self.peer_ip, self.peer_port);
|
||||
self.process_connection();
|
||||
}
|
||||
}
|
||||
|
||||
/// Message types for ListenerActor
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ListenerMessage {
|
||||
/// Stop the listener
|
||||
Stop,
|
||||
}
|
||||
|
||||
/// Actor that accepts incoming TCP connections
|
||||
pub struct ListenerActor {
|
||||
listener: Option<TcpListener>,
|
||||
client_context: Arc<RwLock<ClientContext>>,
|
||||
own_username: String,
|
||||
actor_system: Arc<ActorSystem>,
|
||||
}
|
||||
|
||||
impl ListenerActor {
|
||||
pub fn new(
|
||||
port: u32,
|
||||
client_context: Arc<RwLock<ClientContext>>,
|
||||
own_username: String,
|
||||
actor_system: Arc<ActorSystem>,
|
||||
) -> Result<Self, std::io::Error> {
|
||||
let listener = TcpListener::bind(format!("0.0.0.0:{port}"))?;
|
||||
listener.set_nonblocking(true)?;
|
||||
|
||||
info!("[listener] starting listener on port {port}");
|
||||
|
||||
Ok(Self {
|
||||
listener: Some(listener),
|
||||
client_context,
|
||||
own_username,
|
||||
actor_system,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Actor for ListenerActor {
|
||||
type Message = ListenerMessage;
|
||||
|
||||
fn handle(&mut self, msg: Self::Message) {
|
||||
match msg {
|
||||
ListenerMessage::Stop => {
|
||||
info!("[listener] stopping listener");
|
||||
self.listener = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn tick(&mut self) {
|
||||
let Some(ref listener) = self.listener else {
|
||||
return;
|
||||
};
|
||||
|
||||
match listener.accept() {
|
||||
Ok((stream, _addr)) => {
|
||||
trace!("[listener] accepted new connection");
|
||||
let actor = ConnectionHandlerActor::new(
|
||||
stream,
|
||||
self.client_context.clone(),
|
||||
self.own_username.clone(),
|
||||
);
|
||||
self.actor_system.spawn(actor);
|
||||
}
|
||||
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
|
||||
// No incoming connections, this is normal for non-blocking
|
||||
}
|
||||
Err(e) => {
|
||||
error!("[listener] Failed to accept connection: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+196
-62
@@ -15,6 +15,14 @@ use std::io::{self, Error, Write};
|
||||
use std::net::TcpStream;
|
||||
use std::sync::mpsc::{Receiver, Sender};
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
enum ConnectionState {
|
||||
Disconnected,
|
||||
Connecting { since: Instant },
|
||||
Connected,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PeerMessage {
|
||||
@@ -37,46 +45,49 @@ pub enum PeerMessage {
|
||||
ProcessRead,
|
||||
}
|
||||
|
||||
/// Actor that manages a single peer connection
|
||||
pub struct PeerActor {
|
||||
peer: Arc<RwLock<Peer>>,
|
||||
stream: Option<TcpStream>,
|
||||
connection_state: ConnectionState,
|
||||
reader: MessageReader,
|
||||
client_channel: Sender<ClientOperation>,
|
||||
self_handle: Option<ActorHandle<PeerMessage>>,
|
||||
dispatcher: Option<MessageDispatcher<PeerMessage>>,
|
||||
dispatcher_receiver: Option<Receiver<PeerMessage>>,
|
||||
queued_filenames: Vec<String>,
|
||||
queued_messages: Vec<PeerMessage>,
|
||||
}
|
||||
|
||||
impl PeerActor {
|
||||
/// Create a new PeerActor with an existing socket connection
|
||||
pub fn new(
|
||||
peer: Peer,
|
||||
stream: TcpStream,
|
||||
stream: Option<TcpStream>,
|
||||
reader: Option<MessageReader>,
|
||||
client_channel: Sender<ClientOperation>,
|
||||
) -> Self {
|
||||
let connection_state = if stream.is_some() {
|
||||
ConnectionState::Connected
|
||||
} else {
|
||||
ConnectionState::Disconnected
|
||||
};
|
||||
|
||||
Self {
|
||||
peer: Arc::new(RwLock::new(peer)),
|
||||
stream: Some(stream),
|
||||
stream,
|
||||
connection_state,
|
||||
reader: reader.unwrap_or_default(),
|
||||
client_channel,
|
||||
self_handle: None,
|
||||
dispatcher: None,
|
||||
dispatcher_receiver: None,
|
||||
queued_filenames: Vec::new(),
|
||||
queued_messages: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the actor's own handle for self-messaging
|
||||
pub fn set_self_handle(&mut self, handle: ActorHandle<PeerMessage>) {
|
||||
self.self_handle = Some(handle);
|
||||
}
|
||||
|
||||
/// Initialize the message dispatcher with handlers
|
||||
fn initialize_dispatcher(&mut self) {
|
||||
// Create a channel for the dispatcher to send messages to
|
||||
let (dispatcher_sender, dispatcher_receiver) =
|
||||
std::sync::mpsc::channel::<PeerMessage>();
|
||||
|
||||
@@ -98,28 +109,34 @@ impl PeerActor {
|
||||
));
|
||||
}
|
||||
|
||||
/// Process any messages from the dispatcher
|
||||
fn process_dispatcher_messages(&mut self) {
|
||||
// Collect all messages first to avoid borrow issues
|
||||
let messages: Vec<PeerMessage> =
|
||||
if let Some(ref receiver) = self.dispatcher_receiver {
|
||||
let messages: Vec<PeerMessage> = self
|
||||
.dispatcher_receiver
|
||||
.as_ref()
|
||||
.map_or_else(Vec::new, |receiver| {
|
||||
let mut msgs = Vec::new();
|
||||
while let Ok(msg) = receiver.try_recv() {
|
||||
msgs.push(msg);
|
||||
}
|
||||
msgs
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
});
|
||||
|
||||
// Process all collected messages
|
||||
for msg in messages {
|
||||
self.handle_message(msg);
|
||||
}
|
||||
messages
|
||||
.iter()
|
||||
.for_each(|msg| self.handle_message(msg.clone()));
|
||||
}
|
||||
|
||||
/// Internal handler that doesn't call Actor::handle
|
||||
fn handle_message(&mut self, msg: PeerMessage) {
|
||||
if matches!(self.connection_state, ConnectionState::Connecting { .. }) {
|
||||
match &msg {
|
||||
PeerMessage::SetUsername(_) | PeerMessage::ProcessRead => {}
|
||||
_ => {
|
||||
self.queued_messages.push(msg);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match msg {
|
||||
PeerMessage::SendMessage(message) => {
|
||||
self.send_message(message);
|
||||
@@ -202,7 +219,6 @@ impl PeerActor {
|
||||
self.peer.write().unwrap().username = username;
|
||||
}
|
||||
PeerMessage::QueueUpload(filename) => {
|
||||
self.queued_filenames.push(filename.clone());
|
||||
let message =
|
||||
MessageFactory::build_queue_upload_message(&filename);
|
||||
self.send_message(message);
|
||||
@@ -225,21 +241,15 @@ impl PeerActor {
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle reading from the socket and dispatching messages
|
||||
fn process_read(&mut self) {
|
||||
if self.reader.buffer_len() > 0 {
|
||||
self.extract_and_process_messages(); // This will process all messages in the buffer
|
||||
self.extract_and_process_messages();
|
||||
}
|
||||
let username = self.peer.read().unwrap().username.clone();
|
||||
|
||||
// Try to read more data from the socket first
|
||||
{
|
||||
let stream = match self.stream.as_mut() {
|
||||
Some(s) => s,
|
||||
None => {
|
||||
trace!("[peer_actor:{}] process_read: stream is None, returning", username);
|
||||
return;
|
||||
}
|
||||
None => return,
|
||||
};
|
||||
|
||||
match self.reader.read_from_socket(stream) {
|
||||
@@ -263,7 +273,7 @@ impl PeerActor {
|
||||
}
|
||||
}
|
||||
}
|
||||
self.extract_and_process_messages(); // This will process all messages in the buffer
|
||||
self.extract_and_process_messages();
|
||||
}
|
||||
|
||||
fn extract_and_process_messages(&mut self) {
|
||||
@@ -307,11 +317,9 @@ impl PeerActor {
|
||||
}
|
||||
}
|
||||
|
||||
// Process any messages that were dispatched
|
||||
self.process_dispatcher_messages();
|
||||
}
|
||||
|
||||
/// Handle writing a message to the socket
|
||||
fn send_message(&mut self, message: Message) {
|
||||
let stream = match self.stream.as_mut() {
|
||||
Some(s) => s,
|
||||
@@ -365,7 +373,6 @@ impl PeerActor {
|
||||
error!("Failed to send disconnect notification: {}", e);
|
||||
}
|
||||
}
|
||||
/// Disconnect the peer and notify the client
|
||||
fn disconnect(&mut self) {
|
||||
let username = self.peer.read().unwrap().username.clone();
|
||||
debug!("[peer_actor:{}] disconnect", username);
|
||||
@@ -379,6 +386,145 @@ impl PeerActor {
|
||||
error!("Failed to send disconnect notification: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
fn initiate_connection(&mut self) -> bool {
|
||||
let peer = self.peer.read().unwrap();
|
||||
let username = peer.username.clone();
|
||||
let host = peer.host.clone();
|
||||
let port = peer.port;
|
||||
drop(peer);
|
||||
|
||||
let socket_addr =
|
||||
format!("{}:{}", host, port).parse::<std::net::SocketAddr>();
|
||||
|
||||
match socket_addr {
|
||||
Ok(addr) => {
|
||||
if let Ok(stream) = TcpStream::connect(addr) {
|
||||
if let Err(e) = stream.set_nonblocking(true) {
|
||||
error!(
|
||||
"[peer_actor:{}] Failed to set non-blocking: {}",
|
||||
username, e
|
||||
);
|
||||
self.disconnect_with_error(e);
|
||||
return false;
|
||||
}
|
||||
|
||||
stream.set_nodelay(true).ok();
|
||||
|
||||
self.stream = Some(stream);
|
||||
self.connection_state = ConnectionState::Connecting {
|
||||
since: Instant::now(),
|
||||
};
|
||||
true
|
||||
} else {
|
||||
match TcpStream::connect(addr) {
|
||||
Ok(stream) => {
|
||||
if let Err(e) = stream.set_nonblocking(true) {
|
||||
error!("[peer_actor:{}] Failed to set non-blocking: {}", username, e);
|
||||
self.disconnect_with_error(e);
|
||||
return false;
|
||||
}
|
||||
stream.set_nodelay(true).ok();
|
||||
self.stream = Some(stream);
|
||||
self.connection_state =
|
||||
ConnectionState::Connecting {
|
||||
since: Instant::now(),
|
||||
};
|
||||
true
|
||||
}
|
||||
Err(e) => {
|
||||
self.disconnect_with_error(e);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!(
|
||||
"[peer_actor:{}] Invalid socket address {}:{} - {}",
|
||||
username, host, port, e
|
||||
);
|
||||
self.disconnect_with_error(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
e,
|
||||
));
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn check_connection_status(&mut self) {
|
||||
let ConnectionState::Connecting { since } = self.connection_state
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
let username = self.peer.read().unwrap().username.clone();
|
||||
|
||||
if since.elapsed() > Duration::from_secs(20) {
|
||||
error!(
|
||||
"[peer_actor:{}] Connection timeout after 20 seconds",
|
||||
username
|
||||
);
|
||||
self.disconnect_with_error(io::Error::new(
|
||||
io::ErrorKind::TimedOut,
|
||||
"Connection timeout",
|
||||
));
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(ref stream) = self.stream else {
|
||||
return;
|
||||
};
|
||||
|
||||
match stream.peer_addr() {
|
||||
Ok(_) => {
|
||||
self.connection_state = ConnectionState::Connected;
|
||||
self.on_connection_established();
|
||||
}
|
||||
Err(ref e) if e.kind() == io::ErrorKind::NotConnected => {}
|
||||
Err(e) => {
|
||||
error!("[peer_actor:{}] Connection failed: {}", username, e);
|
||||
self.disconnect_with_error(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn on_connection_established(&mut self) {
|
||||
let peer = self.peer.read().unwrap();
|
||||
let username = peer.username.clone();
|
||||
let token = peer.token;
|
||||
drop(peer);
|
||||
|
||||
let Some(ref mut stream) = self.stream else {
|
||||
return;
|
||||
};
|
||||
|
||||
if let Some(token) = token {
|
||||
let handshake_msg = MessageFactory::build_watch_user(token);
|
||||
if let Err(e) = stream.write_all(&handshake_msg.get_data()) {
|
||||
error!(
|
||||
"[peer_actor:{}] Failed to send watch user handshake: {}",
|
||||
username, e
|
||||
);
|
||||
self.disconnect_with_error(e);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
self.initialize_dispatcher();
|
||||
|
||||
let queued = std::mem::take(&mut self.queued_messages);
|
||||
for msg in queued {
|
||||
self.handle_message(msg);
|
||||
}
|
||||
|
||||
if let Some(ref handle) = self.self_handle {
|
||||
handle.send(PeerMessage::ProcessRead).ok();
|
||||
}
|
||||
|
||||
self.process_read();
|
||||
}
|
||||
}
|
||||
|
||||
impl Actor for PeerActor {
|
||||
@@ -389,29 +535,12 @@ impl Actor for PeerActor {
|
||||
}
|
||||
|
||||
fn on_start(&mut self) {
|
||||
let username = self.peer.read().unwrap().username.clone();
|
||||
|
||||
// Initialize the dispatcher
|
||||
self.initialize_dispatcher();
|
||||
|
||||
// Trigger initial read
|
||||
if let Some(ref handle) = self.self_handle {
|
||||
match handle.send(PeerMessage::ProcessRead) {
|
||||
Ok(_) => {}
|
||||
Err(e) => error!(
|
||||
"[peer_actor:{}] FAILED to send ProcessRead: {}",
|
||||
username, e
|
||||
),
|
||||
}
|
||||
if self.stream.is_none() {
|
||||
self.initiate_connection();
|
||||
} else {
|
||||
error!(
|
||||
"[peer_actor:{}] self_handle is None! Cannot send ProcessRead",
|
||||
username
|
||||
);
|
||||
self.connection_state = ConnectionState::Connected;
|
||||
self.on_connection_established();
|
||||
}
|
||||
|
||||
// Also do an immediate read as fallback
|
||||
self.process_read();
|
||||
}
|
||||
|
||||
fn on_stop(&mut self) {
|
||||
@@ -421,15 +550,20 @@ impl Actor for PeerActor {
|
||||
}
|
||||
|
||||
fn tick(&mut self) {
|
||||
// Periodically check for new data on the socket
|
||||
match self.connection_state {
|
||||
ConnectionState::Connecting { .. } => {
|
||||
// Check if connection has been established
|
||||
self.check_connection_status();
|
||||
}
|
||||
ConnectionState::Connected => {
|
||||
// Normal operation - process reads
|
||||
if self.stream.is_some() {
|
||||
self.process_read();
|
||||
} else {
|
||||
let username = self.peer.read().unwrap().username.clone();
|
||||
trace!(
|
||||
"[peer_actor:{}] tick() called but stream is None",
|
||||
username
|
||||
);
|
||||
}
|
||||
}
|
||||
ConnectionState::Disconnected => {
|
||||
// Do nothing
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,16 +1,15 @@
|
||||
use crate::actor::peer_actor::{PeerActor, PeerMessage};
|
||||
use crate::actor::{ActorHandle, ActorSystem};
|
||||
use crate::client::ClientOperation;
|
||||
use crate::debug;
|
||||
use crate::message::MessageReader;
|
||||
use crate::peer::Peer;
|
||||
use crate::{debug, trace};
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::net::TcpStream;
|
||||
use std::sync::mpsc::Sender;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
/// Registry for managing peer actors
|
||||
pub struct PeerRegistry {
|
||||
peers: Arc<Mutex<HashMap<String, ActorHandle<PeerMessage>>>>,
|
||||
actor_system: Arc<ActorSystem>,
|
||||
@@ -29,45 +28,33 @@ impl PeerRegistry {
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn a new peer actor and register it
|
||||
pub fn register_peer(
|
||||
&self,
|
||||
peer: Peer,
|
||||
stream: TcpStream,
|
||||
stream: Option<TcpStream>,
|
||||
reader: Option<MessageReader>,
|
||||
) -> Result<ActorHandle<PeerMessage>, String> {
|
||||
let username = peer.username.clone();
|
||||
|
||||
debug!("[peer_registry] Registering peer actor for {}", username);
|
||||
|
||||
let actor =
|
||||
PeerActor::new(peer, stream, reader, self.client_channel.clone());
|
||||
|
||||
// Spawn the actor with initialization callback to set self_handle
|
||||
let handle =
|
||||
self.actor_system.spawn_with_handle(actor, |actor, handle| {
|
||||
actor.set_self_handle(handle);
|
||||
});
|
||||
|
||||
// Store in registry
|
||||
let mut peers = self.peers.lock().unwrap();
|
||||
peers.insert(username.clone(), handle.clone());
|
||||
|
||||
trace!(
|
||||
"[peer_registry] Peer actor {} registered successfully",
|
||||
username
|
||||
);
|
||||
|
||||
Ok(handle)
|
||||
}
|
||||
|
||||
/// Get a peer actor handle by username
|
||||
pub fn get_peer(&self, username: &str) -> Option<ActorHandle<PeerMessage>> {
|
||||
let peers = self.peers.lock().unwrap();
|
||||
peers.get(username).cloned()
|
||||
}
|
||||
|
||||
/// Remove a peer actor from the registry
|
||||
pub fn remove_peer(
|
||||
&self,
|
||||
username: &str,
|
||||
@@ -82,25 +69,21 @@ impl PeerRegistry {
|
||||
handle
|
||||
}
|
||||
|
||||
/// Get all registered peer usernames
|
||||
pub fn get_all_usernames(&self) -> Vec<String> {
|
||||
let peers = self.peers.lock().unwrap();
|
||||
peers.keys().cloned().collect()
|
||||
}
|
||||
|
||||
/// Get count of registered peers
|
||||
pub fn count(&self) -> usize {
|
||||
let peers = self.peers.lock().unwrap();
|
||||
peers.len()
|
||||
}
|
||||
|
||||
/// Check if a peer is registered
|
||||
pub fn contains(&self, username: &str) -> bool {
|
||||
let peers = self.peers.lock().unwrap();
|
||||
peers.contains_key(username)
|
||||
}
|
||||
|
||||
/// Send a message to a peer by username
|
||||
pub fn send_to_peer(
|
||||
&self,
|
||||
username: &str,
|
||||
|
||||
+41
-113
@@ -318,16 +318,22 @@ impl Client {
|
||||
own_username: String,
|
||||
) {
|
||||
thread::spawn(move || loop {
|
||||
if let Ok(operation) = reader.recv() {
|
||||
match reader.recv() {
|
||||
Ok(operation) => {
|
||||
match operation {
|
||||
ClientOperation::ConnectToPeer(peer) => {
|
||||
trace!("[client] ConnectToPeer {:?}", peer);
|
||||
let client_context_clone = client_context.clone();
|
||||
let own_username_clone = own_username.clone();
|
||||
|
||||
// Spawn in background to avoid blocking the receiver thread
|
||||
thread::spawn(move || {
|
||||
Self::connect_to_peer(
|
||||
peer,
|
||||
client_context.clone(),
|
||||
own_username.clone(),
|
||||
client_context_clone,
|
||||
own_username_clone,
|
||||
None,
|
||||
);
|
||||
});
|
||||
}
|
||||
ClientOperation::SearchResult(file_search) => {
|
||||
trace!("[client] SearchResult {:?}", file_search);
|
||||
@@ -365,10 +371,18 @@ impl Client {
|
||||
own_username.clone(),
|
||||
);
|
||||
}
|
||||
ClientOperation::DownloadFromPeer(token, peer, allowed) => {
|
||||
ClientOperation::DownloadFromPeer(
|
||||
token,
|
||||
peer,
|
||||
allowed,
|
||||
) => {
|
||||
let maybe_download = {
|
||||
let client_context = client_context.read().unwrap();
|
||||
client_context.download_tokens.get(&token).cloned()
|
||||
let client_context =
|
||||
client_context.read().unwrap();
|
||||
client_context
|
||||
.download_tokens
|
||||
.get(&token)
|
||||
.cloned()
|
||||
};
|
||||
let own_username = own_username.clone();
|
||||
let client_context_clone = client_context.clone();
|
||||
@@ -378,7 +392,7 @@ impl Client {
|
||||
token,
|
||||
peer
|
||||
);
|
||||
return match maybe_download {
|
||||
match maybe_download {
|
||||
Some(download) => {
|
||||
thread::spawn(move || {
|
||||
let download_peer = DownloadPeer::new(
|
||||
@@ -518,12 +532,19 @@ impl Client {
|
||||
obfuscation_type.try_into().unwrap(),
|
||||
obfuscated_port.try_into().unwrap(),
|
||||
);
|
||||
let client_context_clone =
|
||||
client_context.clone();
|
||||
let own_username_clone = own_username.clone();
|
||||
|
||||
// Spawn in background to avoid blocking the receiver thread
|
||||
thread::spawn(move || {
|
||||
Self::connect_to_peer(
|
||||
peer,
|
||||
client_context.clone(),
|
||||
own_username.clone(),
|
||||
client_context_clone,
|
||||
own_username_clone,
|
||||
None,
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
ClientOperation::UpdateDownloadTokens(
|
||||
@@ -575,6 +596,11 @@ impl Client {
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("[client] Channel recv error: {:?}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -584,113 +610,18 @@ impl Client {
|
||||
own_username: String,
|
||||
stream: Option<TcpStream>,
|
||||
) {
|
||||
let client_context2 = client_context.clone();
|
||||
let unlocked_context = client_context.read().unwrap();
|
||||
let client_context = client_context.clone();
|
||||
|
||||
trace!("[client] connect_to_peer: {}", peer.username);
|
||||
if let Some(_sender) = unlocked_context.sender.clone() {
|
||||
let peer_clone = peer.clone();
|
||||
trace!("[client] connecting to {}, with connection_type: {}, and token {:?}", peer.username, peer.connection_type, peer.token);
|
||||
match peer.connection_type {
|
||||
ConnectionType::P => {
|
||||
let username = peer.username.clone();
|
||||
|
||||
let mut tcp_stream = match stream {
|
||||
Some(s) => s,
|
||||
None => {
|
||||
let socket_addr = format!(
|
||||
"{}:{}",
|
||||
peer_clone.host, peer_clone.port
|
||||
)
|
||||
.parse::<std::net::SocketAddr>();
|
||||
|
||||
match socket_addr {
|
||||
Ok(addr) => {
|
||||
match TcpStream::connect_timeout(
|
||||
&addr,
|
||||
Duration::from_secs(20),
|
||||
) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
trace!(
|
||||
"Can't connect to {:?} {:?}:{:?} - {:?}",
|
||||
username, peer_clone.host, peer_clone.port, e
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
trace!(
|
||||
"Invalid socket address for {:?}: {:?}",
|
||||
username,
|
||||
e
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Send connection handshake message
|
||||
use std::io::Write;
|
||||
if let Some(token) = peer_clone.token {
|
||||
// For indirect connections, send watch user message
|
||||
let handshake_msg = crate::message::server::MessageFactory::build_watch_user(token);
|
||||
if let Err(e) =
|
||||
tcp_stream.write_all(&handshake_msg.get_data())
|
||||
{
|
||||
trace!("Failed to send watch user handshake to {:?}: {:?}", username, e);
|
||||
return;
|
||||
}
|
||||
trace!(
|
||||
"[client] Sent watch user handshake for token: {}",
|
||||
token
|
||||
);
|
||||
}
|
||||
|
||||
// Configure socket
|
||||
if let Err(e) = tcp_stream
|
||||
.set_read_timeout(Some(Duration::from_secs(5)))
|
||||
{
|
||||
trace!(
|
||||
"Failed to set read timeout for {:?}: {:?}",
|
||||
username,
|
||||
e
|
||||
);
|
||||
return;
|
||||
}
|
||||
if let Err(e) = tcp_stream
|
||||
.set_write_timeout(Some(Duration::from_secs(5)))
|
||||
{
|
||||
trace!(
|
||||
"Failed to set write timeout for {:?}: {:?}",
|
||||
username,
|
||||
e
|
||||
);
|
||||
return;
|
||||
}
|
||||
if let Err(e) = tcp_stream.set_nodelay(true) {
|
||||
trace!(
|
||||
"Failed to set nodelay for {:?}: {:?}",
|
||||
username,
|
||||
e
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// Spawn peer actor
|
||||
let context = client_context2.read().unwrap();
|
||||
let context = client_context.read().unwrap();
|
||||
if let Some(ref registry) = context.peer_registry {
|
||||
match registry
|
||||
.register_peer(peer_clone, tcp_stream, None)
|
||||
{
|
||||
Ok(_) => {
|
||||
trace!(
|
||||
"[client] peer actor spawned for: {}",
|
||||
username
|
||||
);
|
||||
}
|
||||
match registry.register_peer(peer_clone, stream, None) {
|
||||
Ok(_) => (),
|
||||
Err(e) => {
|
||||
trace!(
|
||||
"Failed to spawn peer actor for {:?}: {:?}",
|
||||
@@ -720,7 +651,7 @@ impl Client {
|
||||
);
|
||||
|
||||
match download_peer.download_file(
|
||||
client_context2.clone(),
|
||||
client_context.clone(),
|
||||
None,
|
||||
None,
|
||||
) {
|
||||
@@ -744,9 +675,6 @@ impl Client {
|
||||
error!("ConnectionType::D not implemented")
|
||||
}
|
||||
}
|
||||
} else {
|
||||
debug!("[client] Peer already connected: {}", peer.username);
|
||||
}
|
||||
}
|
||||
fn pierce_firewall(
|
||||
peer: Peer,
|
||||
|
||||
+1
-1
@@ -14,7 +14,7 @@ fn main() {
|
||||
match client.login() {
|
||||
Ok(_) => {
|
||||
match client
|
||||
.search("Eternity Alex Kassian", Duration::from_secs(10))
|
||||
.search("Eternity Alex Kassian", Duration::from_secs(30))
|
||||
{
|
||||
Ok(results) => {
|
||||
if let Some(result) =
|
||||
|
||||
+5
-17
@@ -75,14 +75,7 @@ fn extract_download_from_buffer(
|
||||
if reader.buffer_len() == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let buffer = reader.get_buffer();
|
||||
trace!(
|
||||
"[listener:{peer_ip}:{peer_port}] reader buffer has {} bytes. {:?}",
|
||||
buffer.len(),
|
||||
buffer
|
||||
);
|
||||
|
||||
let token = parse_token_from_buffer(&buffer, username)?;
|
||||
trace!(
|
||||
"[listener:{}] got transfer_token: {} from data chunk",
|
||||
@@ -107,18 +100,13 @@ fn handle_peer_connection(
|
||||
stream: TcpStream,
|
||||
reader: MessageReader,
|
||||
context: &ConnectionContext,
|
||||
peer_ip: &str,
|
||||
peer_port: u16,
|
||||
_peer_ip: &str,
|
||||
_peer_port: u16,
|
||||
) {
|
||||
debug!("[listener:{peer_ip}:{peer_port}] connection type is P, reader buffer has {} bytes", reader.buffer_len());
|
||||
|
||||
// Spawn peer actor directly using the peer registry
|
||||
let client_context = context.client_context.read().unwrap();
|
||||
if let Some(ref registry) = client_context.peer_registry {
|
||||
match registry.register_peer(peer.clone(), stream, Some(reader)) {
|
||||
Ok(_) => {
|
||||
trace!("[listener] peer actor spawned for: {}", peer.username);
|
||||
}
|
||||
match registry.register_peer(peer.clone(), Some(stream), Some(reader)) {
|
||||
Ok(_) => (),
|
||||
Err(e) => {
|
||||
error!(
|
||||
"Failed to spawn peer actor for {:?}: {:?}",
|
||||
@@ -213,7 +201,7 @@ fn handle_incoming_connection(stream: TcpStream, context: ConnectionContext) {
|
||||
);
|
||||
|
||||
let peer = Peer::new(
|
||||
format!("{}:direct", init_data.username), // Add direct suffix to distinguish from regular peers
|
||||
format!("{}:direct", init_data.username),
|
||||
// init_data.username.clone(),
|
||||
init_data.connection_type.clone(),
|
||||
peer_ip.clone(),
|
||||
|
||||
+1
-16
@@ -341,10 +341,6 @@ impl Server {
|
||||
if let Ok(operation) = server_channel.recv() {
|
||||
match operation {
|
||||
ServerOperation::ConnectToPeer(peer) => {
|
||||
trace!(
|
||||
"[server] → ConnectToPeer {}",
|
||||
peer.username
|
||||
);
|
||||
if let Some(op) = match peer.connection_type {
|
||||
ConnectionType::P | ConnectionType::F => {
|
||||
Some(ClientOperation::ConnectToPeer(
|
||||
@@ -353,18 +349,7 @@ impl Server {
|
||||
}
|
||||
ConnectionType::D => None,
|
||||
} {
|
||||
trace!(
|
||||
"[server] → ConnectToPeer {} - op {:?}",
|
||||
peer.username,
|
||||
op
|
||||
);
|
||||
|
||||
match client_channel.send(op) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
error!("[server] Failed to send: {}", e)
|
||||
}
|
||||
}
|
||||
client_channel.send(op).unwrap();
|
||||
}
|
||||
}
|
||||
ServerOperation::LoginStatus(message) => {
|
||||
|
||||
Reference in New Issue
Block a user