chore: refactor listener
This commit is contained in:
Michel de Graaf
+5
-100
@@ -1,7 +1,3 @@
|
||||
use crate::client::ClientContext;
|
||||
use crate::message::server::MessageFactory;
|
||||
use crate::trace;
|
||||
use crate::types::Download;
|
||||
use std::env;
|
||||
use std::fs::{self, File};
|
||||
use std::io::{self, BufWriter, Read, Write};
|
||||
@@ -11,6 +7,11 @@ use std::path::Path;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::client::ClientContext;
|
||||
use crate::message::server::MessageFactory;
|
||||
use crate::trace;
|
||||
use crate::types::Download;
|
||||
|
||||
#[derive(Debug)]
|
||||
#[allow(dead_code)]
|
||||
struct DownloadPaths {
|
||||
@@ -416,103 +417,7 @@ impl DownloadPeer {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{DownloadPeer, FileManager, StreamProcessor};
|
||||
use crate::client::ClientContext;
|
||||
use std::fs;
|
||||
use std::io::Read;
|
||||
use std::net::TcpListener;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
pub fn build_test_server() -> (u16, Arc<RwLock<Vec<Vec<u8>>>>) {
|
||||
let messages = Arc::new(RwLock::new(Vec::<Vec<u8>>::new()));
|
||||
let messages_clone = Arc::clone(&messages);
|
||||
|
||||
// Find an available port by binding to port 0
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let port = listener.local_addr().unwrap().port();
|
||||
|
||||
thread::spawn(move || {
|
||||
for stream in listener.incoming() {
|
||||
match stream {
|
||||
Ok(mut stream) => {
|
||||
let mut buffer = [0u8; 1024];
|
||||
match stream.read(&mut buffer) {
|
||||
Ok(bytes_read) if bytes_read > 0 => {
|
||||
let data = buffer[..bytes_read].to_vec();
|
||||
messages_clone.write().unwrap().push(data);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Give the server time to start
|
||||
thread::sleep(Duration::from_millis(10));
|
||||
|
||||
(port, messages)
|
||||
}
|
||||
|
||||
#[test]
|
||||
pub fn test_connect_no_pierce() {
|
||||
let (port, messages) = build_test_server();
|
||||
let token = 33;
|
||||
let download_peer = DownloadPeer::new(
|
||||
"test_user".to_string(),
|
||||
"127.0.0.1".to_string(),
|
||||
port as u32,
|
||||
token,
|
||||
false,
|
||||
"own_username".to_string(),
|
||||
);
|
||||
let dummy_context = Arc::new(RwLock::new(ClientContext::new()));
|
||||
let _ = download_peer
|
||||
.download_file(dummy_context, None, None)
|
||||
.unwrap();
|
||||
|
||||
// Give the client time to send messages
|
||||
thread::sleep(Duration::from_millis(10));
|
||||
|
||||
let received_messages = messages.read().unwrap();
|
||||
assert_eq!(received_messages.len(), 1);
|
||||
assert_eq!(received_messages[0], vec![0, 0, 0, 0, 33, 0, 0, 0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
pub fn test_connect_with_pierce() {
|
||||
let (port, messages) = build_test_server();
|
||||
let token = 33;
|
||||
let download_peer = DownloadPeer::new(
|
||||
"test_user".to_string(),
|
||||
"127.0.0.1".to_string(),
|
||||
port as u32,
|
||||
token,
|
||||
true, // no_pierce = true, should send init message
|
||||
"own_username".to_string(),
|
||||
);
|
||||
let dummy_context = Arc::new(RwLock::new(ClientContext::new()));
|
||||
let _ = download_peer
|
||||
.download_file(dummy_context, None, None)
|
||||
.unwrap();
|
||||
|
||||
// Give the client time to send messages
|
||||
thread::sleep(Duration::from_millis(10));
|
||||
|
||||
let received_messages = messages.read().unwrap();
|
||||
assert_eq!(received_messages.len(), 1); // Expect one init message
|
||||
assert_eq!(
|
||||
received_messages[0],
|
||||
vec![
|
||||
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
|
||||
]
|
||||
);
|
||||
|
||||
// No messages expected when no_pierce is false
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_file_manager_create_paths_with_custom_path() {
|
||||
|
||||
+204
-98
@@ -1,144 +1,163 @@
|
||||
use std::net::TcpListener;
|
||||
use std::io;
|
||||
use std::net::{TcpListener, TcpStream};
|
||||
use std::sync::mpsc::Sender;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::thread;
|
||||
|
||||
use crate::client::{ClientContext, ClientOperation};
|
||||
|
||||
use crate::message::MessageReader;
|
||||
use crate::message::{Message, MessageReader};
|
||||
use crate::peer::{ConnectionType, DefaultPeer, DownloadPeer, Peer};
|
||||
use crate::types::Download;
|
||||
use crate::{debug, error, info, trace, DownloadStatus};
|
||||
|
||||
pub struct Listen {}
|
||||
const PEER_INIT_MESSAGE_CODE: u8 = 1;
|
||||
|
||||
impl Listen {
|
||||
pub fn start(
|
||||
port: u32,
|
||||
#[derive(Clone)]
|
||||
struct ConnectionContext {
|
||||
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() {
|
||||
let client_sender = client_sender.clone();
|
||||
let client_context = client_context.clone();
|
||||
let own_username = own_username.clone();
|
||||
thread::spawn(move || {
|
||||
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 mut reader = MessageReader::new();
|
||||
|
||||
let mut message = loop {
|
||||
if let Err(e) = reader.read_from_socket(&mut stream) {
|
||||
trace!("Failed to read from socket: {}", e);
|
||||
return;
|
||||
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() {
|
||||
break msg;
|
||||
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 == 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}");
|
||||
if message_code != PEER_INIT_MESSAGE_CODE {
|
||||
return None;
|
||||
}
|
||||
|
||||
let peer = Peer::new(
|
||||
username.clone(),
|
||||
ConnectionType::P,
|
||||
peer_ip.clone(),
|
||||
peer_port.into(),
|
||||
None,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
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
|
||||
);
|
||||
|
||||
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();
|
||||
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
|
||||
}
|
||||
|
||||
fn handle_peer_connection(
|
||||
peer: Peer,
|
||||
stream: TcpStream,
|
||||
reader: MessageReader,
|
||||
client_sender: Sender<ClientOperation>,
|
||||
peer_ip: &str,
|
||||
peer_port: u16,
|
||||
) {
|
||||
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);
|
||||
if let Ok(default_peer) =
|
||||
default_peer.connect_with_socket(stream, Some(reader))
|
||||
{
|
||||
drop(default_peer);
|
||||
} else if connection_type == "F" {
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_file_connection(
|
||||
peer: Peer,
|
||||
stream: TcpStream,
|
||||
mut reader: MessageReader,
|
||||
token: u32,
|
||||
context: &ConnectionContext,
|
||||
peer_ip: &str,
|
||||
peer_port: u16,
|
||||
) {
|
||||
trace!(
|
||||
"[client] DownloadFromPeer token: {} peer: {:?}",
|
||||
token,
|
||||
peer
|
||||
);
|
||||
|
||||
let mut download: Option<Download> = None;
|
||||
if reader.buffer_len() > 0 {
|
||||
let buffer = reader.get_buffer();
|
||||
trace!("[listener:{peer_ip}:{peer_port}] reader buffer has {} bytes. {:?}", buffer.len(), buffer);
|
||||
|
||||
let token = buffer.get(0..4).unwrap();
|
||||
let token_u32 = u32::from_le_bytes(
|
||||
token
|
||||
.try_into()
|
||||
.unwrap_or_else(|_| panic!("[listener:{}] slice with incorrect length, can't extract transfer_token", username)),
|
||||
let download = extract_download_from_buffer(
|
||||
&mut reader,
|
||||
&context.client_context,
|
||||
&peer.username,
|
||||
peer_ip,
|
||||
peer_port,
|
||||
);
|
||||
trace!(
|
||||
"[listener:{}] got transfer_token: {} from data chunk",
|
||||
username,
|
||||
token_u32);
|
||||
|
||||
download = client_context
|
||||
.read()
|
||||
.unwrap()
|
||||
.download_tokens
|
||||
.get(&token_u32)
|
||||
.cloned();
|
||||
|
||||
if download.is_none() {
|
||||
let download_tokens = client_context
|
||||
.read()
|
||||
.unwrap()
|
||||
.download_tokens
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect::<Vec<u32>>();
|
||||
trace!("[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}", token_u32, download_tokens );
|
||||
}
|
||||
}
|
||||
|
||||
let download_peer = DownloadPeer::new(
|
||||
format!("{}:direct", username.clone()),
|
||||
format!("{}:direct", peer.username),
|
||||
peer.host.clone(),
|
||||
peer.port,
|
||||
token,
|
||||
true,
|
||||
own_username,
|
||||
context.own_username.clone(),
|
||||
);
|
||||
|
||||
match download_peer.download_file(
|
||||
client_context,
|
||||
context.client_context.clone(),
|
||||
download,
|
||||
Some(stream),
|
||||
) {
|
||||
Ok((download, filename)) => {
|
||||
download
|
||||
.sender
|
||||
.send(DownloadStatus::Completed)
|
||||
.unwrap();
|
||||
download.sender.send(DownloadStatus::Completed).unwrap();
|
||||
info!(
|
||||
"Successfully downloaded {} bytes to {}",
|
||||
download.size, filename
|
||||
@@ -151,15 +170,102 @@ impl Listen {
|
||||
);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
debug!(
|
||||
"[listener:{peer_ip}:{peer_port}] connection type is not P or F"
|
||||
}
|
||||
|
||||
fn handle_incoming_connection(stream: TcpStream, context: ConnectionContext) {
|
||||
let Ok(peer_addr) = stream.peer_addr() else {
|
||||
error!("[listener] failed to get peer address");
|
||||
return;
|
||||
};
|
||||
|
||||
let peer_ip = peer_addr.ip().to_string();
|
||||
let peer_port = peer_addr.port();
|
||||
let mut stream = stream;
|
||||
let mut reader = MessageReader::new();
|
||||
|
||||
let Ok(message) = read_peer_init_message(&mut stream, &mut reader) else {
|
||||
error!(
|
||||
"[listener:{peer_ip}:{peer_port}] Failed to read peer init message"
|
||||
);
|
||||
return;
|
||||
};
|
||||
|
||||
let Some(init_data) = parse_peer_init_message(message) else {
|
||||
error!("[listener:{peer_ip}:{peer_port}] Invalid or unknown peer init message");
|
||||
return;
|
||||
};
|
||||
|
||||
debug!(
|
||||
"[listener:{peer_ip}:{peer_port}] peerInit (0) username: {} connection_type: {} token: {}",
|
||||
init_data.username, init_data.connection_type, init_data.token
|
||||
);
|
||||
|
||||
let peer = Peer::new(
|
||||
init_data.username.clone(),
|
||||
ConnectionType::P,
|
||||
peer_ip.clone(),
|
||||
peer_port.into(),
|
||||
None,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
);
|
||||
|
||||
match init_data.connection_type {
|
||||
ConnectionType::P => handle_peer_connection(
|
||||
peer,
|
||||
stream,
|
||||
reader,
|
||||
context.client_sender,
|
||||
&peer_ip,
|
||||
peer_port,
|
||||
),
|
||||
ConnectionType::F => handle_file_connection(
|
||||
peer,
|
||||
stream,
|
||||
reader,
|
||||
init_data.token,
|
||||
&context,
|
||||
&peer_ip,
|
||||
peer_port,
|
||||
),
|
||||
ConnectionType::D => {
|
||||
debug!("[listener:{peer_ip}:{peer_port}] connection type is D, not supported yet, closing connection. ");
|
||||
}
|
||||
} else {
|
||||
debug!("[listener:{peer_ip}:{peer_port}] unknown message with code: {message_code}");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub struct Listen {}
|
||||
|
||||
impl Listen {
|
||||
pub fn start(
|
||||
port: u32,
|
||||
client_sender: Sender<ClientOperation>,
|
||||
client_context: Arc<RwLock<ClientContext>>,
|
||||
own_username: String,
|
||||
) {
|
||||
info!("[listener] starting listener on port {port}");
|
||||
|
||||
let listener = TcpListener::bind(format!("0.0.0.0:{port}"))
|
||||
.expect("Failed to bind listener to port");
|
||||
|
||||
let context = ConnectionContext {
|
||||
client_sender,
|
||||
client_context,
|
||||
own_username,
|
||||
};
|
||||
|
||||
for stream in listener.incoming() {
|
||||
let Ok(stream) = stream else {
|
||||
error!(
|
||||
"[listener] Failed to accept connection: {}",
|
||||
stream.unwrap_err()
|
||||
);
|
||||
continue;
|
||||
};
|
||||
|
||||
let context = context.clone();
|
||||
thread::spawn(move || handle_incoming_connection(stream, context));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user