diff --git a/prudpv1/src/executables/proxy_secure.rs b/prudpv1/src/executables/proxy_secure.rs index 3f3b7c6..222c1dd 100644 --- a/prudpv1/src/executables/proxy_secure.rs +++ b/prudpv1/src/executables/proxy_secure.rs @@ -9,6 +9,7 @@ use rnex_core::reggie::UnitPacketRead; use rnex_core::reggie::UnitPacketWrite; use rnex_core::rmc::structures::RmcSerialize; use rnex_core::rnex_proxy_common::ConnectionInitData; +use std::ops::Deref; use std::time::Duration; use tokio::net::TcpStream; use tokio::task; @@ -105,7 +106,7 @@ pub async fn start(param: ProxyStartupParam) { } } } - conn.close_connection().await; + conn.deref().close_connection().await; }); } } diff --git a/prudpv1/src/prudp/router.rs b/prudpv1/src/prudp/router.rs index ca28dd6..3d09123 100644 --- a/prudpv1/src/prudp/router.rs +++ b/prudpv1/src/prudp/router.rs @@ -1,58 +1,68 @@ +use crate::prudp::packet::PRUDPV1Packet; +use crate::prudp::router::Error::VirtualPortTaken; +use crate::prudp::socket::{AnyInternalSocket, CryptoHandler, ExternalSocket, new_socket_pair}; +use log::{error, info}; +use rnex_core::prudp::virtual_port::VirtualPort; use std::io; use std::io::Cursor; use std::marker::PhantomData; -use tokio::net::UdpSocket; -use std::net::{SocketAddr, SocketAddrV4}; use std::net::SocketAddr::V4; +use std::net::{SocketAddr, SocketAddrV4}; use std::sync::{Arc, Weak}; use std::time::Duration; -use tokio::task::JoinHandle; -use log::{error, info}; use thiserror::Error; +use tokio::net::UdpSocket; use tokio::select; use tokio::sync::RwLock; +use tokio::task::JoinHandle; use tokio::time::sleep; -use crate::prudp::socket::{new_socket_pair, AnyInternalSocket, CryptoHandler, ExternalSocket}; -use crate::prudp::packet::{PRUDPV1Packet}; -use rnex_core::prudp::virtual_port::VirtualPort; -use crate::prudp::router::Error::VirtualPortTaken; - - pub struct Router { endpoints: RwLock<[Option>; 16]>, //running: AtomicBool, socket: Arc, - _no_outside_construction: PhantomData<()> + _no_outside_construction: PhantomData<()>, } #[derive(Debug, Error)] -pub enum Error{ +pub enum Error { #[error("tried to register socket to a port which is already taken (port: {0})")] - VirtualPortTaken(u8) + VirtualPortTaken(u8), } - impl Router { - async fn process_prudp_packets<'a>(self: Arc, _socket: Arc, addr: SocketAddrV4, udp_message: Vec){ + async fn process_prudp_packets<'a>( + self: Arc, + _socket: Arc, + addr: SocketAddrV4, + udp_message: Vec, + ) { let mut stream = Cursor::new(&udp_message); while stream.position() as usize != udp_message.len() { - let packet = match PRUDPV1Packet::new(&mut stream){ + let packet = match PRUDPV1Packet::new(&mut stream) { Ok(p) => p, Err(e) => { - error!("Somebody({}) is fucking with the servers or their connection is bad (reason: {})", addr, e); + error!( + "Somebody({}) is fucking with the servers or their connection is bad (reason: {})", + addr, e + ); break; - }, + } }; let connection = packet.source_sockaddr(addr); - let endpoints = self.endpoints.read().await; - let Some(endpoint) = endpoints[packet.header.destination_port.get_port_number() as usize].as_ref() else { - error!("connection to invalid endpoint({}) attempted by {}", packet.header.destination_port.get_port_number(), connection.regular_socket_addr); + let Some(endpoint) = + endpoints[packet.header.destination_port.get_port_number() as usize].as_ref() + else { + error!( + "connection to invalid endpoint({}) attempted by {}", + packet.header.destination_port.get_port_number(), + connection.regular_socket_addr + ); continue; }; @@ -60,38 +70,32 @@ impl Router { // Dont keep the locked structure for too long drop(endpoints); - - tokio::spawn(async move { - endpoint.receive_packet(connection, packet).await - }); + tokio::spawn(async move { endpoint.receive_packet(connection, packet).await }); } } - async fn server_thread_send_entry(this: Weak, socket: Arc){ + async fn server_thread_send_entry(this: Weak, socket: Arc) { info!("starting datagram thread"); while let Some(this) = this.upgrade() { // yes we actually allow the max udp to be read lol let mut msg_buffer = vec![0u8; 65507]; - let (len, addr) = - select! { - r = socket.recv_from(&mut msg_buffer) => { - r.expect("Datagram thread crashed due to unexpected error from recv_from") - } - _ = sleep(Duration::from_secs(5)) => { - continue; - } - }; - + let (len, addr) = select! { + r = socket.recv_from(&mut msg_buffer) => { + r.expect("Datagram thread crashed due to unexpected error from recv_from") + } + _ = sleep(Duration::from_secs(5)) => { + continue; + } + }; let V4(addr) = addr else { error!("somehow got ipv6 packet...? ignoring"); continue; }; - let current_msg = &msg_buffer[0..len]; tokio::spawn(this.process_prudp_packets(socket.clone(), addr, current_msg.to_vec())); @@ -99,8 +103,8 @@ impl Router { println!("exitting datagram") } - - pub async fn new(addr: SocketAddrV4) -> io::Result<(Arc, JoinHandle<()>)>{ + + pub async fn new(addr: SocketAddrV4) -> io::Result<(Arc, JoinHandle<()>)> { // trace!("starting router on {}", addr); let socket = Arc::new(UdpSocket::bind(addr).await?); @@ -109,15 +113,14 @@ impl Router { endpoints: Default::default(), // running: AtomicBool::new(true), socket: socket.clone(), - _no_outside_construction: Default::default() + _no_outside_construction: Default::default(), }; let arc = Arc::new(own_impl); - let task = { let socket = socket.clone(); - let server= Arc::downgrade(&arc); + let server = Arc::downgrade(&arc); tokio::spawn(async { Self::server_thread_send_entry(server, socket).await; @@ -135,29 +138,31 @@ impl Router { }); } - Ok((arc, task)) } - pub fn get_udp_socket(&self) -> Arc{ + pub fn get_udp_socket(&self) -> Arc { self.socket.clone() } // This will remove a socket from the router, this renders all instances of that socket unable // to recieve any more data making the error out on trying to for example recieve connections - pub async fn remove_socket(&self, virtual_port: VirtualPort){ + pub async fn remove_socket(&self, virtual_port: VirtualPort) { self.endpoints.write().await[virtual_port.get_port_number() as usize] = None; } // returns Some(()) i - pub async fn add_socket(&self, virtual_port: VirtualPort, encryption: E) - -> Result{ + pub async fn add_socket( + &self, + virtual_port: VirtualPort, + encryption: E, + ) -> Result, Error> { let mut endpoints = self.endpoints.write().await; let idx = virtual_port.get_port_number() as usize; // dont create the socket if we dont need to - if !endpoints[idx].is_none(){ + if !endpoints[idx].is_none() { return Err(VirtualPortTaken(idx as u8)); } @@ -168,11 +173,14 @@ impl Router { Ok(external) } - pub fn get_own_address(&self) -> SocketAddrV4{ - match self.socket.local_addr().expect("unable to get socket address"){ + pub fn get_own_address(&self) -> SocketAddrV4 { + match self + .socket + .local_addr() + .expect("unable to get socket address") + { SocketAddr::V4(v4) => v4, - _ => unreachable!() + _ => unreachable!(), } } } - diff --git a/prudpv1/src/prudp/socket.rs b/prudpv1/src/prudp/socket.rs index 19bf24d..9fd0274 100644 --- a/prudpv1/src/prudp/socket.rs +++ b/prudpv1/src/prudp/socket.rs @@ -40,7 +40,7 @@ pub struct CommonConnection { struct InternalConnection { common: Arc, - connections: Weak>>>>>, + connections: Weak>>>>, reliable_server_counter: u16, reliable_client_counter: u16, supported_function_version: u32, @@ -53,6 +53,23 @@ struct InternalConnection { unacknowleged_packets: Vec<(Instant, PRUDPV1Packet)>, } +struct InternalConnectionMutex(Mutex>); + +impl AsRef>> + for InternalConnectionMutex +{ + fn as_ref(&self) -> &Mutex> { + &self.0 + } +} + +impl Deref for InternalConnectionMutex { + type Target = Mutex>; + fn deref(&self) -> &Self::Target { + &self.0 + } +} + impl Deref for InternalConnection { type Target = CommonConnection; fn deref(&self) -> &Self::Target { @@ -124,15 +141,26 @@ impl InternalConnection { } } -pub struct ExternalConnection { - sending: SendingConnection, +pub struct ExternalConnection { + sending: SendingConnection, data_receiver: Receiver>, } -#[derive(Clone)] -pub struct SendingConnection { +pub struct SendingConnection { common: Arc, - internal: Weak, + internal: Weak>, +} + +// we couldnt use the implementation the derive would generate here because that +// bakes in the assumption that all type parameters must be `Clone` for the struct +// we are deriving on to be `Clone` as well +impl Clone for SendingConnection { + fn clone(&self) -> Self { + Self { + common: self.common.clone(), + internal: self.internal.clone(), + } + } } pub struct CommonSocket { @@ -146,20 +174,23 @@ pub(super) struct InternalSocket { crypto_handler: T, // perf note: change the code to use RwLock here instead to avoid connections being able to block one another before the data is sent off. internal_connections: Arc< - Mutex>>>>, + Mutex>>>, >, connection_establishment_data_sender: Mutex>>, - connection_sender: Sender, + connection_sender: Sender>, } -pub struct ExternalSocket { +pub struct ExternalSocket { common: Arc, - connection_receiver: Receiver, - internal: Weak, + connection_receiver: Receiver>, + internal: Weak>, } -impl ExternalSocket { - pub async fn connect(&mut self, addr: PRUDPSockAddr) -> Option { +impl ExternalSocket { + pub async fn connect( + &mut self, + addr: PRUDPSockAddr, + ) -> Option> { let socket = self.internal.upgrade()?; socket.connect(addr).await; @@ -167,12 +198,12 @@ impl ExternalSocket { self.connection_receiver.recv().await } - pub async fn accept(&mut self) -> Option { + pub async fn accept(&mut self) -> Option> { self.connection_receiver.recv().await } } -impl Deref for ExternalSocket { +impl Deref for ExternalSocket { type Target = CommonSocket; fn deref(&self) -> &Self::Target { &self.common @@ -194,15 +225,7 @@ pub(super) trait AnyInternalSocket: async fn connect(&self, address: PRUDPSockAddr) -> Option<()>; } -#[async_trait] -pub(super) trait AnyInternalConnection: Send + Sync + 'static { - async fn send_data_packet(&self, data: Vec); - - async fn close_connection(&self); -} - -#[async_trait] -impl AnyInternalConnection for Mutex> { +impl InternalConnectionMutex { async fn send_data_packet(&self, data: Vec) { let pieces = data.chunks(600); let max_piece = pieces.len() - 1; @@ -248,15 +271,9 @@ impl AnyInternalConnection for Mutex InternalSocket { async fn get_connection( &self, addr: PRUDPSockAddr, - ) -> Option>>> { + ) -> Option>> { let connections = self.internal_connections.lock().await; let Some(conn) = connections.get(&addr) else { error!("tried to send data on inactive connection!"); @@ -337,7 +354,7 @@ impl InternalSocket { } async fn connection_thread( - connection: Weak>>, + connection: Weak>, ) { //todo: handle stuff like resending packets if they arent acknowledged in here @@ -417,9 +434,9 @@ impl InternalSocket { supported_function_version, }; - let internal = Arc::new(Mutex::new(internal)); + let internal = Arc::new(InternalConnectionMutex(Mutex::new(internal))); - let dyn_internal: Arc = internal.clone(); + let dyn_internal = internal.clone(); let external = ExternalConnection { sending: SendingConnection { @@ -670,6 +687,7 @@ impl AnyInternalSocket for InternalSocket { if (packet.header.types_and_flags.get_flags() & MULTI_ACK) != 0 { if let Some(conn) = self.get_connection(address).await { + let conn = &**conn; let mut conn = conn.lock().await; if conn.supported_function_version == 1 { @@ -818,7 +836,7 @@ pub(super) fn new_socket_pair( virtual_port: VirtualPort, encryption: T, socket: Arc, -) -> (Arc>, ExternalSocket) { +) -> (Arc>, ExternalSocket) { let common = Arc::new(CommonSocket { virtual_port, _phantom_unconstructible: Default::default(), @@ -835,7 +853,7 @@ pub(super) fn new_socket_pair( socket, }); - let dyn_internal: Arc = internal.clone(); + let dyn_internal = internal.clone(); let external = ExternalSocket { common, @@ -872,32 +890,32 @@ pub trait CryptoHandler: Send + Sync + 'static { fn sign_pre_handshake(&self, packet: &mut PRUDPV1Packet); } -impl Deref for ExternalConnection { - type Target = SendingConnection; +impl Deref for ExternalConnection { + type Target = SendingConnection; fn deref(&self) -> &Self::Target { &self.sending } } -impl Deref for SendingConnection { +impl Deref for SendingConnection { type Target = CommonConnection; fn deref(&self) -> &Self::Target { &self.common } } -impl ExternalConnection { +impl ExternalConnection { pub async fn recv(&mut self) -> Option> { self.data_receiver.recv().await } //todo: make this an actual result instead of an option - pub fn duplicate_sender(&self) -> SendingConnection { + pub fn duplicate_sender(&self) -> SendingConnection { self.sending.clone() } } -impl SendingConnection { +impl SendingConnection { pub async fn send(&self, data: Vec) -> Option<()> { let internal = self.internal.upgrade()?; spawn(async move { @@ -911,6 +929,8 @@ impl SendingConnection { return; }; + let mut internal = internal.lock().await; + internal.close_connection().await; } }