feat: seperate into lib and cli packages
This commit is contained in:
Michel de Graaf
@@ -0,0 +1,458 @@
|
||||
use std::env;
|
||||
use std::fs;
|
||||
use std::io::{self, Read, Write};
|
||||
use std::net::TcpStream;
|
||||
use std::net::ToSocketAddrs;
|
||||
use std::path::{Path, PathBuf};
|
||||
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;
|
||||
|
||||
const START_DOWNLOAD: [u8; 8] =
|
||||
[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
|
||||
const READ_BUFFER_SIZE: usize = 8192;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum DownloadError {
|
||||
ConnectionFailed(io::Error),
|
||||
InvalidAddress(String),
|
||||
HandshakeFailed(io::Error),
|
||||
StreamReadError(io::Error),
|
||||
StreamWriteError(io::Error),
|
||||
TokenNotFound(u32),
|
||||
DownloadInfoMissing(u32),
|
||||
FileWriteError(io::Error),
|
||||
PathResolutionError(String),
|
||||
InvalidTokenBytes,
|
||||
LockPoisoned,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for DownloadError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::ConnectionFailed(e) => write!(f, "Connection failed: {}", e),
|
||||
Self::InvalidAddress(addr) => {
|
||||
write!(f, "Invalid address: {}", addr)
|
||||
}
|
||||
Self::HandshakeFailed(e) => write!(f, "Handshake failed: {}", e),
|
||||
Self::StreamReadError(e) => write!(f, "Stream read error: {}", e),
|
||||
Self::StreamWriteError(e) => write!(f, "Stream write error: {}", e),
|
||||
Self::TokenNotFound(token) => {
|
||||
write!(f, "Token not found: {}", token)
|
||||
}
|
||||
Self::DownloadInfoMissing(token) => {
|
||||
write!(f, "Download info missing for token: {}", token)
|
||||
}
|
||||
Self::FileWriteError(e) => write!(f, "File write error: {}", e),
|
||||
Self::PathResolutionError(msg) => {
|
||||
write!(f, "Path resolution error: {}", msg)
|
||||
}
|
||||
Self::InvalidTokenBytes => {
|
||||
write!(f, "Invalid token bytes received")
|
||||
}
|
||||
Self::LockPoisoned => write!(f, "Lock poisoned"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for DownloadError {}
|
||||
|
||||
impl From<io::Error> for DownloadError {
|
||||
fn from(error: io::Error) -> Self {
|
||||
Self::StreamReadError(error)
|
||||
}
|
||||
}
|
||||
|
||||
struct FileManager;
|
||||
|
||||
impl FileManager {
|
||||
fn expand_path(path: &str) -> PathBuf {
|
||||
if let Some(stripped) = path.strip_prefix('~') {
|
||||
if let Ok(home) = env::var("HOME") {
|
||||
PathBuf::from(home).join(stripped.trim_start_matches('/'))
|
||||
} else {
|
||||
PathBuf::from(path)
|
||||
}
|
||||
} else {
|
||||
PathBuf::from(path)
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_filename_from_path(full_path: &str) -> &str {
|
||||
full_path
|
||||
.split(['/', '\\'])
|
||||
.next_back()
|
||||
.unwrap_or(full_path)
|
||||
}
|
||||
|
||||
fn create_download_path_from_filename(
|
||||
output_directory: PathBuf,
|
||||
filename: &str,
|
||||
) -> PathBuf {
|
||||
let filename_only = Self::extract_filename_from_path(filename);
|
||||
output_directory.join(filename_only)
|
||||
}
|
||||
}
|
||||
|
||||
struct StreamProcessor {
|
||||
total_bytes: usize,
|
||||
received: bool,
|
||||
buffer: Vec<u8>,
|
||||
}
|
||||
|
||||
impl StreamProcessor {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
total_bytes: 0,
|
||||
received: false,
|
||||
buffer: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn process_data_chunk(&mut self, data: &[u8]) {
|
||||
self.buffer.extend_from_slice(data);
|
||||
self.total_bytes += data.len();
|
||||
}
|
||||
|
||||
fn should_continue(&self, expected_size: Option<usize>) -> bool {
|
||||
if let Some(size) = expected_size {
|
||||
self.total_bytes < size
|
||||
} else {
|
||||
true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct DownloadPeer {
|
||||
username: String,
|
||||
host: String,
|
||||
port: u32,
|
||||
#[allow(dead_code)]
|
||||
own_username: String,
|
||||
token: u32,
|
||||
no_pierce: bool,
|
||||
}
|
||||
|
||||
impl DownloadPeer {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
username: String,
|
||||
host: String,
|
||||
port: u32,
|
||||
token: u32,
|
||||
no_pierce: bool,
|
||||
own_username: String,
|
||||
) -> Self {
|
||||
Self {
|
||||
username,
|
||||
host,
|
||||
port,
|
||||
own_username,
|
||||
token,
|
||||
no_pierce,
|
||||
}
|
||||
}
|
||||
|
||||
fn establish_connection(&self) -> Result<TcpStream, DownloadError> {
|
||||
let socket_address = format!("{}:{}", self.host, self.port)
|
||||
.to_socket_addrs()
|
||||
.map_err(DownloadError::ConnectionFailed)?
|
||||
.next()
|
||||
.ok_or_else(|| {
|
||||
DownloadError::InvalidAddress(format!(
|
||||
"{}:{}",
|
||||
self.host, self.port
|
||||
))
|
||||
})?;
|
||||
|
||||
let stream = TcpStream::connect_timeout(
|
||||
&socket_address,
|
||||
Duration::from_secs(20),
|
||||
)
|
||||
.map_err(DownloadError::ConnectionFailed)?;
|
||||
|
||||
stream
|
||||
.set_read_timeout(Some(Duration::from_secs(30)))
|
||||
.map_err(DownloadError::ConnectionFailed)?;
|
||||
stream
|
||||
.set_write_timeout(Some(Duration::from_secs(5)))
|
||||
.map_err(DownloadError::ConnectionFailed)?;
|
||||
stream
|
||||
.set_nodelay(true)
|
||||
.map_err(DownloadError::ConnectionFailed)?;
|
||||
|
||||
Ok(stream)
|
||||
}
|
||||
|
||||
fn perform_handshake(
|
||||
&self,
|
||||
stream: &mut TcpStream,
|
||||
) -> Result<(), DownloadError> {
|
||||
trace!(
|
||||
"[download_peer:{}] performing handshake no_pierce: {}",
|
||||
self.username,
|
||||
self.no_pierce
|
||||
);
|
||||
|
||||
if !self.no_pierce {
|
||||
let message =
|
||||
MessageFactory::build_pierce_firewall_message(self.token);
|
||||
stream
|
||||
.write_all(&message.get_buffer())
|
||||
.map_err(DownloadError::HandshakeFailed)?;
|
||||
trace!(
|
||||
"[download_peer:{}] sending pierce firewall message token: {}: {:?}",
|
||||
self.username,
|
||||
self.token,
|
||||
&message.get_buffer()
|
||||
);
|
||||
stream.flush().map_err(DownloadError::HandshakeFailed)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_pierce_firewall_response(
|
||||
&self,
|
||||
data: &[u8],
|
||||
stream: &mut TcpStream,
|
||||
client_context: &Arc<RwLock<ClientContext>>,
|
||||
) -> Result<Download, DownloadError> {
|
||||
let token_bytes =
|
||||
data.get(0..4).ok_or(DownloadError::InvalidTokenBytes)?;
|
||||
let token_array: [u8; 4] = token_bytes
|
||||
.try_into()
|
||||
.map_err(|_| DownloadError::InvalidTokenBytes)?;
|
||||
let token_u32 = u32::from_le_bytes(token_array);
|
||||
|
||||
trace!(
|
||||
"[download_peer:{}] got token: {} from data chunk",
|
||||
self.username,
|
||||
token_u32
|
||||
);
|
||||
|
||||
stream
|
||||
.write_all(&START_DOWNLOAD)
|
||||
.map_err(DownloadError::StreamWriteError)?;
|
||||
|
||||
let client_guard = client_context
|
||||
.read()
|
||||
.map_err(|_| DownloadError::LockPoisoned)?;
|
||||
let download_info =
|
||||
client_guard.download_tokens.get(&token_u32).cloned();
|
||||
drop(client_guard);
|
||||
|
||||
download_info.ok_or(DownloadError::TokenNotFound(token_u32))
|
||||
}
|
||||
|
||||
fn read_download_stream(
|
||||
&self,
|
||||
stream: &mut TcpStream,
|
||||
client_context: &Arc<RwLock<ClientContext>>,
|
||||
mut download: Option<Download>,
|
||||
) -> Result<(Vec<u8>, Download), DownloadError> {
|
||||
let mut processor = StreamProcessor::new();
|
||||
let mut read_buffer = [1u8; READ_BUFFER_SIZE];
|
||||
|
||||
trace!(
|
||||
"[download_peer:{}] Starting to read data from peer",
|
||||
self.username
|
||||
);
|
||||
|
||||
if download.is_some() {
|
||||
stream
|
||||
.write_all(&START_DOWNLOAD)
|
||||
.map_err(DownloadError::StreamWriteError)?;
|
||||
}
|
||||
|
||||
loop {
|
||||
match stream.read(&mut read_buffer) {
|
||||
Ok(0) => {
|
||||
trace!(
|
||||
"[download_peer:{}] connection closed by peer. bytes read: {}",
|
||||
self.username,
|
||||
processor.total_bytes
|
||||
);
|
||||
break;
|
||||
}
|
||||
Ok(bytes_read) => {
|
||||
let data = &read_buffer[..bytes_read];
|
||||
|
||||
if !self.no_pierce && !processor.received {
|
||||
let new_download = self
|
||||
.handle_pierce_firewall_response(
|
||||
data,
|
||||
stream,
|
||||
client_context,
|
||||
)?;
|
||||
trace!(
|
||||
"[download_peer:{}] got download info for token: {} - filename: {}",
|
||||
self.username,
|
||||
self.token,
|
||||
new_download.filename
|
||||
);
|
||||
download = Some(new_download);
|
||||
processor.received = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
processor.process_data_chunk(data);
|
||||
|
||||
let expected_size = download
|
||||
.as_ref()
|
||||
.ok_or(DownloadError::DownloadInfoMissing(self.token))?
|
||||
.size as usize;
|
||||
|
||||
if !processor.should_continue(Some(expected_size)) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(DownloadError::StreamReadError(e));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
trace!(
|
||||
"[download_peer:{}] finished reading data from peer",
|
||||
self.username
|
||||
);
|
||||
|
||||
let download =
|
||||
download.ok_or(DownloadError::DownloadInfoMissing(self.token))?;
|
||||
|
||||
Ok((processor.buffer, download))
|
||||
}
|
||||
|
||||
fn resolve_download_path(
|
||||
&self,
|
||||
download: &Download,
|
||||
) -> Result<String, DownloadError> {
|
||||
let download_directory = &download.download_directory;
|
||||
let mut expanded_path = FileManager::expand_path(download_directory);
|
||||
|
||||
if !expanded_path.is_dir() {
|
||||
expanded_path = expanded_path
|
||||
.parent()
|
||||
.ok_or_else(|| {
|
||||
DownloadError::PathResolutionError(format!(
|
||||
"Cannot resolve parent directory for: {}",
|
||||
expanded_path.display()
|
||||
))
|
||||
})?
|
||||
.to_path_buf();
|
||||
}
|
||||
|
||||
let final_path = FileManager::create_download_path_from_filename(
|
||||
expanded_path,
|
||||
&download.filename,
|
||||
);
|
||||
|
||||
final_path
|
||||
.to_str()
|
||||
.ok_or_else(|| {
|
||||
DownloadError::PathResolutionError(format!(
|
||||
"Path contains invalid UTF-8: {}",
|
||||
final_path.display()
|
||||
))
|
||||
})
|
||||
.map(String::from)
|
||||
}
|
||||
|
||||
fn save_downloaded_file(
|
||||
&self,
|
||||
path: &str,
|
||||
data: &[u8],
|
||||
) -> Result<(), DownloadError> {
|
||||
if let Some(parent) = Path::new(path).parent() {
|
||||
fs::create_dir_all(parent)
|
||||
.map_err(DownloadError::FileWriteError)?;
|
||||
}
|
||||
|
||||
fs::write(path, data).map_err(DownloadError::FileWriteError)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn download_file(
|
||||
self,
|
||||
client_context: Arc<RwLock<ClientContext>>,
|
||||
download: Option<Download>,
|
||||
stream: Option<TcpStream>,
|
||||
) -> Result<(Download, String), DownloadError> {
|
||||
trace!(
|
||||
"[download_peer:{}] download_file: download is present?: {:?}, stream is present?: {:?}, no_pierce: {}",
|
||||
self.username,
|
||||
download.is_some(),
|
||||
stream.is_some(),
|
||||
self.no_pierce
|
||||
);
|
||||
|
||||
let mut stream = match stream {
|
||||
Some(s) => s,
|
||||
None => self.establish_connection()?,
|
||||
};
|
||||
|
||||
trace!("[download_peer:{}] connected", self.username);
|
||||
|
||||
self.perform_handshake(&mut stream)?;
|
||||
trace!("[download_peer:{}] handshake completed", self.username);
|
||||
|
||||
let (buffer, download) =
|
||||
self.read_download_stream(&mut stream, &client_context, download)?;
|
||||
|
||||
let final_path = self.resolve_download_path(&download)?;
|
||||
self.save_downloaded_file(&final_path, &buffer)?;
|
||||
|
||||
trace!(
|
||||
"[download_peer:{}] download completed successfully: {} bytes, saved to: {}",
|
||||
self.username,
|
||||
buffer.len(),
|
||||
final_path
|
||||
);
|
||||
|
||||
Ok((download, final_path))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{DownloadPeer, FileManager};
|
||||
|
||||
#[test]
|
||||
fn test_establish_connection_invalid_address() {
|
||||
let download_peer = DownloadPeer::new(
|
||||
"user".to_string(),
|
||||
"invalid-host".to_string(),
|
||||
9999,
|
||||
123,
|
||||
false,
|
||||
"own_user".to_string(),
|
||||
);
|
||||
let result = download_peer.establish_connection();
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_filename_from_path() {
|
||||
assert_eq!(
|
||||
FileManager::extract_filename_from_path("/path/to/file.mp3"),
|
||||
"file.mp3"
|
||||
);
|
||||
assert_eq!(
|
||||
FileManager::extract_filename_from_path("C:\\path\\to\\file.mp3"),
|
||||
"file.mp3"
|
||||
);
|
||||
assert_eq!(
|
||||
FileManager::extract_filename_from_path("@@bhfrv\\Soulseek Downloads\\complete\\Beatport Top Deep House (2021)\\michel test file.mp3"),
|
||||
"michel test file.mp3"
|
||||
);
|
||||
assert_eq!(
|
||||
FileManager::extract_filename_from_path("file.mp3"),
|
||||
"file.mp3"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,275 @@
|
||||
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::{Message, MessageReader};
|
||||
use crate::peer::{ConnectionType, DownloadPeer, Peer};
|
||||
use crate::types::Download;
|
||||
use crate::{debug, error, info, trace, DownloadStatus};
|
||||
|
||||
const PEER_INIT_MESSAGE_CODE: u8 = 1;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ConnectionContext {
|
||||
#[allow(dead_code)]
|
||||
client_sender: Sender<ClientOperation>,
|
||||
client_context: Arc<RwLock<ClientContext>>,
|
||||
own_username: String,
|
||||
}
|
||||
|
||||
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();
|
||||
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,
|
||||
context: &ConnectionContext,
|
||||
_peer_ip: &str,
|
||||
_peer_port: u16,
|
||||
) {
|
||||
let client_context = context.client_context.read().unwrap();
|
||||
if let Some(ref registry) = client_context.peer_registry {
|
||||
match registry.register_peer(peer.clone(), Some(stream), Some(reader)) {
|
||||
Ok(_) => (),
|
||||
Err(e) => {
|
||||
error!(
|
||||
"Failed to spawn peer actor for {:?}: {:?}",
|
||||
peer.username, e
|
||||
);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
error!("PeerRegistry not initialized");
|
||||
}
|
||||
}
|
||||
|
||||
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 download = extract_download_from_buffer(
|
||||
&mut reader,
|
||||
&context.client_context,
|
||||
&peer.username,
|
||||
peer_ip,
|
||||
peer_port,
|
||||
);
|
||||
|
||||
let download_peer = DownloadPeer::new(
|
||||
format!("{}:direct", peer.username),
|
||||
peer.host.clone(),
|
||||
peer.port,
|
||||
token,
|
||||
true,
|
||||
context.own_username.clone(),
|
||||
);
|
||||
|
||||
match download_peer.download_file(
|
||||
context.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 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(
|
||||
format!("{}:direct", init_data.username),
|
||||
// init_data.username.clone(),
|
||||
init_data.connection_type.clone(),
|
||||
peer_ip.clone(),
|
||||
peer_port.into(),
|
||||
None,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
);
|
||||
|
||||
match init_data.connection_type {
|
||||
ConnectionType::P => handle_peer_connection(
|
||||
peer, stream, reader, &context, &peer_ip, peer_port,
|
||||
),
|
||||
|
||||
ConnectionType::F => {
|
||||
thread::spawn(move || {
|
||||
trace!(
|
||||
"[listener:{peer_ip}:{peer_port}] handling file connection in thread"
|
||||
);
|
||||
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. ");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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();
|
||||
handle_incoming_connection(stream, context);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
mod download_peer;
|
||||
pub mod listen;
|
||||
|
||||
// Export actor types
|
||||
pub use crate::actor::peer_actor::{PeerActor, PeerMessage};
|
||||
pub use crate::actor::peer_registry::PeerRegistry;
|
||||
|
||||
pub use download_peer::DownloadPeer;
|
||||
|
||||
use crate::message::Message;
|
||||
use core::fmt;
|
||||
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 {
|
||||
P,
|
||||
F,
|
||||
D,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ParseConnectionTypeError;
|
||||
|
||||
impl fmt::Display for ParseConnectionTypeError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
write!(f, "invalid connection type")
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ParseConnectionTypeError {}
|
||||
|
||||
impl FromStr for ConnectionType {
|
||||
type Err = ParseConnectionTypeError;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s {
|
||||
"P" => Ok(ConnectionType::P),
|
||||
"F" => Ok(ConnectionType::F),
|
||||
"D" => Ok(ConnectionType::D),
|
||||
_ => Err(ParseConnectionTypeError),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for ConnectionType {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
let s = match self {
|
||||
ConnectionType::P => "P",
|
||||
ConnectionType::F => "F",
|
||||
ConnectionType::D => "D",
|
||||
};
|
||||
write!(f, "{}", s)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
#[allow(dead_code)]
|
||||
pub struct Peer {
|
||||
pub username: String,
|
||||
pub connection_type: ConnectionType,
|
||||
pub host: String,
|
||||
pub port: u32,
|
||||
pub token: Option<u32>,
|
||||
pub privileged: Option<u8>,
|
||||
pub unknown: Option<u8>,
|
||||
pub obfuscated_port: Option<u8>,
|
||||
}
|
||||
impl Peer {
|
||||
#[allow(clippy::too_many_arguments, dead_code)]
|
||||
pub fn new(
|
||||
username: String,
|
||||
connection_type: ConnectionType,
|
||||
host: String,
|
||||
port: u32,
|
||||
token: Option<u32>,
|
||||
privileged: u8,
|
||||
unknown: u8,
|
||||
obfuscated_port: u8,
|
||||
) -> Self {
|
||||
Self {
|
||||
username,
|
||||
connection_type,
|
||||
host,
|
||||
port,
|
||||
token,
|
||||
privileged: Some(privileged),
|
||||
unknown: Some(unknown),
|
||||
obfuscated_port: Some(obfuscated_port),
|
||||
}
|
||||
}
|
||||
#[allow(dead_code)]
|
||||
pub fn new_from_message(message: &mut Message) -> Self {
|
||||
let username = message.read_string();
|
||||
let connection_type = message.read_string().parse().unwrap();
|
||||
|
||||
let mut ip: Vec<i32> = vec![];
|
||||
for _ in 0..4 {
|
||||
ip.push(message.read_int8().into());
|
||||
}
|
||||
let host: String = format!(
|
||||
"{}.{}.{}.{}",
|
||||
ip[3].abs(),
|
||||
ip[2].abs(),
|
||||
ip[1].abs(),
|
||||
ip[0].abs()
|
||||
);
|
||||
|
||||
let (port, token, privileged, unknown, obfuscated_port) = (
|
||||
message.read_int32(),
|
||||
message.read_int32(),
|
||||
message.read_int8(),
|
||||
message.read_int8(),
|
||||
message.read_int8(),
|
||||
);
|
||||
|
||||
Self {
|
||||
username,
|
||||
connection_type,
|
||||
host,
|
||||
port,
|
||||
token: Some(token),
|
||||
privileged: Some(privileged),
|
||||
unknown: Some(unknown),
|
||||
obfuscated_port: Some(obfuscated_port),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn test_new_from_message() {
|
||||
let data: Vec<u8> = [
|
||||
36, 0, 0, 0, 18, 0, 0, 0, 2, 0, 0, 0, 100, 112, 1, 0, 0, 0, 80, 27,
|
||||
231, 37, 45, 186, 8, 0, 0, 178, 78, 25, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
]
|
||||
.to_vec();
|
||||
let mut message = Message::new_with_data(data);
|
||||
message.set_pointer(8);
|
||||
|
||||
let peer = Peer::new_from_message(&mut message);
|
||||
|
||||
assert_eq!(peer.username, "dp");
|
||||
assert!(matches!(peer.connection_type, ConnectionType::P));
|
||||
assert_eq!(peer.host, "45.37.231.27");
|
||||
assert_eq!(peer.port, 2234);
|
||||
assert_eq!(peer.token, Some(1658546));
|
||||
assert_eq!(peer.privileged, Some(0));
|
||||
assert_eq!(peer.unknown, Some(0));
|
||||
assert_eq!(peer.obfuscated_port, Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_from_message2() {
|
||||
let data: Vec<u8> = [
|
||||
42, 0, 0, 0, 18, 0, 0, 0, 8, 0, 0, 0, 103, 114, 97, 110, 100, 112, 97,
|
||||
103, 1, 0, 0, 0, 80, 137, 128, 193, 68, 187, 8, 0, 0, 58, 16, 0, 0, 0,
|
||||
1, 0, 0, 0, 188, 8, 0, 0,
|
||||
]
|
||||
.to_vec();
|
||||
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");
|
||||
assert!(matches!(peer.connection_type, ConnectionType::P));
|
||||
assert_eq!(peer.host, "68.193.128.137");
|
||||
assert_eq!(peer.port, 2235);
|
||||
assert_eq!(peer.token, Some(4154));
|
||||
assert_eq!(peer.privileged, Some(0));
|
||||
assert_eq!(peer.unknown, Some(1));
|
||||
assert_eq!(peer.obfuscated_port, Some(0));
|
||||
}
|
||||
Reference in New Issue
Block a user