spacebar/rust-nex
Watch
14
1
Fork
You've already forked rust-nex
 1
Code Issues Pull requests Projects Releases Packages 60 Wiki Activity Actions

remove dyn-ness from prudpv1 connection impl

Browse source
This commit is contained in:
Maple Nebel 2026-06-17 17:33:09 +02:00
commit 8b68ddb6f7
3 changed files with 125 additions and 96 deletions
Download patch file Download diff file Expand all files Collapse all files

3
prudpv1/src/executables/proxy_secure.rs
View file

@ -9,6 +9,7 @@ use rnex_core::reggie::UnitPacketRead;
use rnex_core::reggie::UnitPacketWrite; use rnex_core::reggie::UnitPacketWrite;
use rnex_core::rmc::structures::RmcSerialize; use rnex_core::rmc::structures::RmcSerialize;
use rnex_core::rnex_proxy_common::ConnectionInitData; use rnex_core::rnex_proxy_common::ConnectionInitData;
use std::ops::Deref;
use std::time::Duration; use std::time::Duration;
use tokio::net::TcpStream; use tokio::net::TcpStream;
use tokio::task; use tokio::task;
@ -105,7 +106,7 @@ pub async fn start(param: ProxyStartupParam) {
} }
} }
} }
conn.close_connection().await; conn.deref().close_connection().await;
}); });
} }
} }

112
prudpv1/src/prudp/router.rs
View file

@ -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;
use std::io::Cursor; use std::io::Cursor;
use std::marker::PhantomData; use std::marker::PhantomData;
use tokio::net::UdpSocket;
use std::net::{SocketAddr, SocketAddrV4};
use std::net::SocketAddr::V4; use std::net::SocketAddr::V4;
use std::net::{SocketAddr, SocketAddrV4};
use std::sync::{Arc, Weak}; use std::sync::{Arc, Weak};
use std::time::Duration; use std::time::Duration;
use tokio::task::JoinHandle;
use log::{error, info};
use thiserror::Error; use thiserror::Error;
use tokio::net::UdpSocket;
use tokio::select; use tokio::select;
use tokio::sync::RwLock; use tokio::sync::RwLock;
use tokio::task::JoinHandle;
use tokio::time::sleep; 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 { pub struct Router {
endpoints: RwLock<[Option<Arc<dyn AnyInternalSocket>>; 16]>, endpoints: RwLock<[Option<Arc<dyn AnyInternalSocket>>; 16]>,
//running: AtomicBool, //running: AtomicBool,
socket: Arc<UdpSocket>, socket: Arc<UdpSocket>,
_no_outside_construction: PhantomData<()> _no_outside_construction: PhantomData<()>,
} }
#[derive(Debug, Error)] #[derive(Debug, Error)]
pub enum Error{ pub enum Error {
#[error("tried to register socket to a port which is already taken (port: {0})")] #[error("tried to register socket to a port which is already taken (port: {0})")]
VirtualPortTaken(u8) VirtualPortTaken(u8),
} }
impl Router { impl Router {
async fn process_prudp_packets<'a>(self: Arc<Self>, _socket: Arc<UdpSocket>, addr: SocketAddrV4, udp_message: Vec<u8>){ async fn process_prudp_packets<'a>(
self: Arc<Self>,
_socket: Arc<UdpSocket>,
addr: SocketAddrV4,
udp_message: Vec<u8>,
) {
let mut stream = Cursor::new(&udp_message); let mut stream = Cursor::new(&udp_message);
while stream.position() as usize != udp_message.len() { 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, Ok(p) => p,
Err(e) => { 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; break;
}, }
}; };
let connection = packet.source_sockaddr(addr); let connection = packet.source_sockaddr(addr);
let endpoints = self.endpoints.read().await; let endpoints = self.endpoints.read().await;
let Some(endpoint) = endpoints[packet.header.destination_port.get_port_number() as usize].as_ref() else { let Some(endpoint) =
error!("connection to invalid endpoint({}) attempted by {}", packet.header.destination_port.get_port_number(), connection.regular_socket_addr); 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; continue;
}; };
@ -60,38 +70,32 @@ impl Router {
// Dont keep the locked structure for too long // Dont keep the locked structure for too long
drop(endpoints); drop(endpoints);
tokio::spawn(async move { tokio::spawn(async move { endpoint.receive_packet(connection, packet).await });
endpoint.receive_packet(connection, packet).await
});
} }
} }
async fn server_thread_send_entry(this: Weak<Self>, socket: Arc<UdpSocket>){ async fn server_thread_send_entry(this: Weak<Self>, socket: Arc<UdpSocket>) {
info!("starting datagram thread"); info!("starting datagram thread");
while let Some(this) = this.upgrade() { while let Some(this) = this.upgrade() {
// yes we actually allow the max udp to be read lol // yes we actually allow the max udp to be read lol
let mut msg_buffer = vec![0u8; 65507]; let mut msg_buffer = vec![0u8; 65507];
let (len, addr) = let (len, addr) = select! {
select! { r = socket.recv_from(&mut msg_buffer) => {
r = socket.recv_from(&mut msg_buffer) => { r.expect("Datagram thread crashed due to unexpected error from recv_from")
r.expect("Datagram thread crashed due to unexpected error from recv_from") }
} _ = sleep(Duration::from_secs(5)) => {
_ = sleep(Duration::from_secs(5)) => { continue;
continue; }
} };
};
let V4(addr) = addr else { let V4(addr) = addr else {
error!("somehow got ipv6 packet...? ignoring"); error!("somehow got ipv6 packet...? ignoring");
continue; continue;
}; };
let current_msg = &msg_buffer[0..len]; let current_msg = &msg_buffer[0..len];
tokio::spawn(this.process_prudp_packets(socket.clone(), addr, current_msg.to_vec())); tokio::spawn(this.process_prudp_packets(socket.clone(), addr, current_msg.to_vec()));
@ -99,8 +103,8 @@ impl Router {
println!("exitting datagram") println!("exitting datagram")
} }
pub async fn new(addr: SocketAddrV4) -> io::Result<(Arc<Self>, JoinHandle<()>)>{ pub async fn new(addr: SocketAddrV4) -> io::Result<(Arc<Self>, JoinHandle<()>)> {
// trace!("starting router on {}", addr); // trace!("starting router on {}", addr);
let socket = Arc::new(UdpSocket::bind(addr).await?); let socket = Arc::new(UdpSocket::bind(addr).await?);
@ -109,15 +113,14 @@ impl Router {
endpoints: Default::default(), endpoints: Default::default(),
// running: AtomicBool::new(true), // running: AtomicBool::new(true),
socket: socket.clone(), socket: socket.clone(),
_no_outside_construction: Default::default() _no_outside_construction: Default::default(),
}; };
let arc = Arc::new(own_impl); let arc = Arc::new(own_impl);
let task = { let task = {
let socket = socket.clone(); let socket = socket.clone();
let server= Arc::downgrade(&arc); let server = Arc::downgrade(&arc);
tokio::spawn(async { tokio::spawn(async {
Self::server_thread_send_entry(server, socket).await; Self::server_thread_send_entry(server, socket).await;
@ -135,29 +138,31 @@ impl Router {
}); });
} }
Ok((arc, task)) Ok((arc, task))
} }
pub fn get_udp_socket(&self) -> Arc<UdpSocket>{ pub fn get_udp_socket(&self) -> Arc<UdpSocket> {
self.socket.clone() self.socket.clone()
} }
// This will remove a socket from the router, this renders all instances of that socket unable // 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 // 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; self.endpoints.write().await[virtual_port.get_port_number() as usize] = None;
} }
// returns Some(()) i // returns Some(()) i
pub async fn add_socket<E: CryptoHandler>(&self, virtual_port: VirtualPort, encryption: E) pub async fn add_socket<E: CryptoHandler>(
-> Result<ExternalSocket, Error>{ &self,
virtual_port: VirtualPort,
encryption: E,
) -> Result<ExternalSocket<E>, Error> {
let mut endpoints = self.endpoints.write().await; let mut endpoints = self.endpoints.write().await;
let idx = virtual_port.get_port_number() as usize; let idx = virtual_port.get_port_number() as usize;
// dont create the socket if we dont need to // 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)); return Err(VirtualPortTaken(idx as u8));
} }
@ -168,11 +173,14 @@ impl Router {
Ok(external) Ok(external)
} }
pub fn get_own_address(&self) -> SocketAddrV4{ pub fn get_own_address(&self) -> SocketAddrV4 {
match self.socket.local_addr().expect("unable to get socket address"){ match self
.socket
.local_addr()
.expect("unable to get socket address")
{
SocketAddr::V4(v4) => v4, SocketAddr::V4(v4) => v4,
_ => unreachable!() _ => unreachable!(),
} }
} }
} }

106
prudpv1/src/prudp/socket.rs
View file

@ -40,7 +40,7 @@ pub struct CommonConnection {
struct InternalConnection<E: CryptoHandlerConnectionInstance> { struct InternalConnection<E: CryptoHandlerConnectionInstance> {
common: Arc<CommonConnection>, common: Arc<CommonConnection>,
connections: Weak<Mutex<BTreeMap<PRUDPSockAddr, Arc<Mutex<InternalConnection<E>>>>>>, connections: Weak<Mutex<BTreeMap<PRUDPSockAddr, Arc<InternalConnectionMutex<E>>>>>,
reliable_server_counter: u16, reliable_server_counter: u16,
reliable_client_counter: u16, reliable_client_counter: u16,
supported_function_version: u32, supported_function_version: u32,
@ -53,6 +53,23 @@ struct InternalConnection<E: CryptoHandlerConnectionInstance> {
unacknowleged_packets: Vec<(Instant, PRUDPV1Packet)>, unacknowleged_packets: Vec<(Instant, PRUDPV1Packet)>,
} }
struct InternalConnectionMutex<E: CryptoHandlerConnectionInstance>(Mutex<InternalConnection<E>>);
impl<E: CryptoHandlerConnectionInstance> AsRef<Mutex<InternalConnection<E>>>
for InternalConnectionMutex<E>
{
fn as_ref(&self) -> &Mutex<InternalConnection<E>> {
&self.0
}
}
impl<E: CryptoHandlerConnectionInstance> Deref for InternalConnectionMutex<E> {
type Target = Mutex<InternalConnection<E>>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl<E: CryptoHandlerConnectionInstance> Deref for InternalConnection<E> { impl<E: CryptoHandlerConnectionInstance> Deref for InternalConnection<E> {
type Target = CommonConnection; type Target = CommonConnection;
fn deref(&self) -> &Self::Target { fn deref(&self) -> &Self::Target {
@ -124,15 +141,26 @@ impl<E: CryptoHandlerConnectionInstance> InternalConnection<E> {
} }
} }
pub struct ExternalConnection { pub struct ExternalConnection<T: CryptoHandlerConnectionInstance> {
sending: SendingConnection, sending: SendingConnection<T>,
data_receiver: Receiver<Vec<u8>>, data_receiver: Receiver<Vec<u8>>,
} }
#[derive(Clone)] pub struct SendingConnection<T: CryptoHandlerConnectionInstance> {
pub struct SendingConnection {
common: Arc<CommonConnection>, common: Arc<CommonConnection>,
internal: Weak<dyn AnyInternalConnection>, internal: Weak<InternalConnectionMutex<T>>,
}
// 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<T: CryptoHandlerConnectionInstance> Clone for SendingConnection<T> {
fn clone(&self) -> Self {
Self {
common: self.common.clone(),
internal: self.internal.clone(),
}
}
} }
pub struct CommonSocket { pub struct CommonSocket {
@ -146,20 +174,23 @@ pub(super) struct InternalSocket<T: CryptoHandler> {
crypto_handler: T, 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. // 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< internal_connections: Arc<
Mutex<BTreeMap<PRUDPSockAddr, Arc<Mutex<InternalConnection<T::CryptoConnectionInstance>>>>>, Mutex<BTreeMap<PRUDPSockAddr, Arc<InternalConnectionMutex<T::CryptoConnectionInstance>>>>,
>, >,
connection_establishment_data_sender: Mutex<Option<Sender<PRUDPV1Packet>>>, connection_establishment_data_sender: Mutex<Option<Sender<PRUDPV1Packet>>>,
connection_sender: Sender<ExternalConnection>, connection_sender: Sender<ExternalConnection<T::CryptoConnectionInstance>>,
} }
pub struct ExternalSocket { pub struct ExternalSocket<T: CryptoHandler> {
common: Arc<CommonSocket>, common: Arc<CommonSocket>,
connection_receiver: Receiver<ExternalConnection>, connection_receiver: Receiver<ExternalConnection<T::CryptoConnectionInstance>>,
internal: Weak<dyn AnyInternalSocket>, internal: Weak<InternalSocket<T>>,
} }
impl ExternalSocket { impl<T: CryptoHandler> ExternalSocket<T> {
pub async fn connect(&mut self, addr: PRUDPSockAddr) -> Option<ExternalConnection> { pub async fn connect(
&mut self,
addr: PRUDPSockAddr,
) -> Option<ExternalConnection<T::CryptoConnectionInstance>> {
let socket = self.internal.upgrade()?; let socket = self.internal.upgrade()?;
socket.connect(addr).await; socket.connect(addr).await;
@ -167,12 +198,12 @@ impl ExternalSocket {
self.connection_receiver.recv().await self.connection_receiver.recv().await
} }
pub async fn accept(&mut self) -> Option<ExternalConnection> { pub async fn accept(&mut self) -> Option<ExternalConnection<T::CryptoConnectionInstance>> {
self.connection_receiver.recv().await self.connection_receiver.recv().await
} }
} }
impl Deref for ExternalSocket { impl<T: CryptoHandler> Deref for ExternalSocket<T> {
type Target = CommonSocket; type Target = CommonSocket;
fn deref(&self) -> &Self::Target { fn deref(&self) -> &Self::Target {
&self.common &self.common
@ -194,15 +225,7 @@ pub(super) trait AnyInternalSocket:
async fn connect(&self, address: PRUDPSockAddr) -> Option<()>; async fn connect(&self, address: PRUDPSockAddr) -> Option<()>;
} }
#[async_trait] impl<E: CryptoHandlerConnectionInstance> InternalConnectionMutex<E> {
pub(super) trait AnyInternalConnection: Send + Sync + 'static {
async fn send_data_packet(&self, data: Vec<u8>);
async fn close_connection(&self);
}
#[async_trait]
impl<T: CryptoHandlerConnectionInstance> AnyInternalConnection for Mutex<InternalConnection<T>> {
async fn send_data_packet(&self, data: Vec<u8>) { async fn send_data_packet(&self, data: Vec<u8>) {
let pieces = data.chunks(600); let pieces = data.chunks(600);
let max_piece = pieces.len() - 1; let max_piece = pieces.len() - 1;
@ -248,15 +271,9 @@ impl<T: CryptoHandlerConnectionInstance> AnyInternalConnection for Mutex<Interna
locked.unacknowleged_packets.push((Instant::now(), packet)); locked.unacknowleged_packets.push((Instant::now(), packet));
drop(locked); drop(locked);
sleep(Duration::from_secs(16)).await; sleep(Duration::from_millis(16)).await;
} }
} }
async fn close_connection(&self) {
let mut locked = self.lock().await;
locked.close_connection().await;
}
} }
async fn send_raw_prudp_to_sockaddr( async fn send_raw_prudp_to_sockaddr(
@ -280,7 +297,7 @@ impl<T: CryptoHandler> InternalSocket<T> {
async fn get_connection( async fn get_connection(
&self, &self,
addr: PRUDPSockAddr, addr: PRUDPSockAddr,
) -> Option<Arc<Mutex<InternalConnection<T::CryptoConnectionInstance>>>> { ) -> Option<Arc<InternalConnectionMutex<T::CryptoConnectionInstance>>> {
let connections = self.internal_connections.lock().await; let connections = self.internal_connections.lock().await;
let Some(conn) = connections.get(&addr) else { let Some(conn) = connections.get(&addr) else {
error!("tried to send data on inactive connection!"); error!("tried to send data on inactive connection!");
@ -337,7 +354,7 @@ impl<T: CryptoHandler> InternalSocket<T> {
} }
async fn connection_thread( async fn connection_thread(
connection: Weak<Mutex<InternalConnection<T::CryptoConnectionInstance>>>, connection: Weak<InternalConnectionMutex<T::CryptoConnectionInstance>>,
) { ) {
//todo: handle stuff like resending packets if they arent acknowledged in here //todo: handle stuff like resending packets if they arent acknowledged in here
@ -417,9 +434,9 @@ impl<T: CryptoHandler> InternalSocket<T> {
supported_function_version, supported_function_version,
}; };
let internal = Arc::new(Mutex::new(internal)); let internal = Arc::new(InternalConnectionMutex(Mutex::new(internal)));
let dyn_internal: Arc<dyn AnyInternalConnection> = internal.clone(); let dyn_internal = internal.clone();
let external = ExternalConnection { let external = ExternalConnection {
sending: SendingConnection { sending: SendingConnection {
@ -670,6 +687,7 @@ impl<T: CryptoHandler> AnyInternalSocket for InternalSocket<T> {
if (packet.header.types_and_flags.get_flags() & MULTI_ACK) != 0 { if (packet.header.types_and_flags.get_flags() & MULTI_ACK) != 0 {
if let Some(conn) = self.get_connection(address).await { if let Some(conn) = self.get_connection(address).await {
let conn = &**conn;
let mut conn = conn.lock().await; let mut conn = conn.lock().await;
if conn.supported_function_version == 1 { if conn.supported_function_version == 1 {
@ -818,7 +836,7 @@ pub(super) fn new_socket_pair<T: CryptoHandler>(
virtual_port: VirtualPort, virtual_port: VirtualPort,
encryption: T, encryption: T,
socket: Arc<UdpSocket>, socket: Arc<UdpSocket>,
) -> (Arc<InternalSocket<T>>, ExternalSocket) { ) -> (Arc<InternalSocket<T>>, ExternalSocket<T>) {
let common = Arc::new(CommonSocket { let common = Arc::new(CommonSocket {
virtual_port, virtual_port,
_phantom_unconstructible: Default::default(), _phantom_unconstructible: Default::default(),
@ -835,7 +853,7 @@ pub(super) fn new_socket_pair<T: CryptoHandler>(
socket, socket,
}); });
let dyn_internal: Arc<dyn AnyInternalSocket> = internal.clone(); let dyn_internal = internal.clone();
let external = ExternalSocket { let external = ExternalSocket {
common, common,
@ -872,32 +890,32 @@ pub trait CryptoHandler: Send + Sync + 'static {
fn sign_pre_handshake(&self, packet: &mut PRUDPV1Packet); fn sign_pre_handshake(&self, packet: &mut PRUDPV1Packet);
} }
impl Deref for ExternalConnection { impl<T: CryptoHandlerConnectionInstance> Deref for ExternalConnection<T> {
type Target = SendingConnection; type Target = SendingConnection<T>;
fn deref(&self) -> &Self::Target { fn deref(&self) -> &Self::Target {
&self.sending &self.sending
} }
} }
impl Deref for SendingConnection { impl<T: CryptoHandlerConnectionInstance> Deref for SendingConnection<T> {
type Target = CommonConnection; type Target = CommonConnection;
fn deref(&self) -> &Self::Target { fn deref(&self) -> &Self::Target {
&self.common &self.common
} }
} }
impl ExternalConnection { impl<E: CryptoHandlerConnectionInstance> ExternalConnection<E> {
pub async fn recv(&mut self) -> Option<Vec<u8>> { pub async fn recv(&mut self) -> Option<Vec<u8>> {
self.data_receiver.recv().await self.data_receiver.recv().await
} }
//todo: make this an actual result instead of an option //todo: make this an actual result instead of an option
pub fn duplicate_sender(&self) -> SendingConnection { pub fn duplicate_sender(&self) -> SendingConnection<E> {
self.sending.clone() self.sending.clone()
} }
} }
impl SendingConnection { impl<E: CryptoHandlerConnectionInstance> SendingConnection<E> {
pub async fn send(&self, data: Vec<u8>) -> Option<()> { pub async fn send(&self, data: Vec<u8>) -> Option<()> {
let internal = self.internal.upgrade()?; let internal = self.internal.upgrade()?;
spawn(async move { spawn(async move {
@ -911,6 +929,8 @@ impl SendingConnection {
return; return;
}; };
let mut internal = internal.lock().await;
internal.close_connection().await; internal.close_connection().await;
} }
} }