kitten/rust-nex
Watch
1
0
Fork
You've already forked rust-nex
 0
forked from spacebar/rust-nex
Code Pull requests Activity

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::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;
});
}
}

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::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<Arc<dyn AnyInternalSocket>>; 16]>,
//running: AtomicBool,
socket: Arc<UdpSocket>,
_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<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);
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<Self>, socket: Arc<UdpSocket>){
async fn server_thread_send_entry(this: Weak<Self>, socket: Arc<UdpSocket>) {
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<Self>, JoinHandle<()>)>{
pub async fn new(addr: SocketAddrV4) -> io::Result<(Arc<Self>, 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<UdpSocket>{
pub fn get_udp_socket(&self) -> Arc<UdpSocket> {
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<E: CryptoHandler>(&self, virtual_port: VirtualPort, encryption: E)
-> Result<ExternalSocket, Error>{
pub async fn add_socket<E: CryptoHandler>(
&self,
virtual_port: VirtualPort,
encryption: E,
) -> Result<ExternalSocket<E>, 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!(),
}
}
}

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

@ -40,7 +40,7 @@ pub struct CommonConnection {
struct InternalConnection<E: CryptoHandlerConnectionInstance> {
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_client_counter: u16,
supported_function_version: u32,
@ -53,6 +53,23 @@ struct InternalConnection<E: CryptoHandlerConnectionInstance> {
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> {
type Target = CommonConnection;
fn deref(&self) -> &Self::Target {
@ -124,15 +141,26 @@ impl<E: CryptoHandlerConnectionInstance> InternalConnection<E> {
}
}
pub struct ExternalConnection {
sending: SendingConnection,
pub struct ExternalConnection<T: CryptoHandlerConnectionInstance> {
sending: SendingConnection<T>,
data_receiver: Receiver<Vec<u8>>,
}
#[derive(Clone)]
pub struct SendingConnection {
pub struct SendingConnection<T: CryptoHandlerConnectionInstance> {
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 {
@ -146,20 +174,23 @@ pub(super) struct InternalSocket<T: CryptoHandler> {
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<BTreeMap<PRUDPSockAddr, Arc<Mutex<InternalConnection<T::CryptoConnectionInstance>>>>>,
Mutex<BTreeMap<PRUDPSockAddr, Arc<InternalConnectionMutex<T::CryptoConnectionInstance>>>>,
>,
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>,
connection_receiver: Receiver<ExternalConnection>,
internal: Weak<dyn AnyInternalSocket>,
connection_receiver: Receiver<ExternalConnection<T::CryptoConnectionInstance>>,
internal: Weak<InternalSocket<T>>,
}
impl ExternalSocket {
pub async fn connect(&mut self, addr: PRUDPSockAddr) -> Option<ExternalConnection> {
impl<T: CryptoHandler> ExternalSocket<T> {
pub async fn connect(
&mut self,
addr: PRUDPSockAddr,
) -> Option<ExternalConnection<T::CryptoConnectionInstance>> {
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<ExternalConnection> {
pub async fn accept(&mut self) -> Option<ExternalConnection<T::CryptoConnectionInstance>> {
self.connection_receiver.recv().await
}
}
impl Deref for ExternalSocket {
impl<T: CryptoHandler> Deref for ExternalSocket<T> {
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<u8>);
async fn close_connection(&self);
}
#[async_trait]
impl<T: CryptoHandlerConnectionInstance> AnyInternalConnection for Mutex<InternalConnection<T>> {
impl<E: CryptoHandlerConnectionInstance> InternalConnectionMutex<E> {
async fn send_data_packet(&self, data: Vec<u8>) {
let pieces = data.chunks(600);
let max_piece = pieces.len() - 1;
@ -248,15 +271,9 @@ impl<T: CryptoHandlerConnectionInstance> AnyInternalConnection for Mutex<Interna
locked.unacknowleged_packets.push((Instant::now(), packet));
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(
@ -280,7 +297,7 @@ impl<T: CryptoHandler> InternalSocket<T> {
async fn get_connection(
&self,
addr: PRUDPSockAddr,
) -> Option<Arc<Mutex<InternalConnection<T::CryptoConnectionInstance>>>> {
) -> Option<Arc<InternalConnectionMutex<T::CryptoConnectionInstance>>> {
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<T: CryptoHandler> InternalSocket<T> {
}
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
@ -417,9 +434,9 @@ impl<T: CryptoHandler> InternalSocket<T> {
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 {
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 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<T: CryptoHandler>(
virtual_port: VirtualPort,
encryption: T,
socket: Arc<UdpSocket>,
) -> (Arc<InternalSocket<T>>, ExternalSocket) {
) -> (Arc<InternalSocket<T>>, ExternalSocket<T>) {
let common = Arc::new(CommonSocket {
virtual_port,
_phantom_unconstructible: Default::default(),
@ -835,7 +853,7 @@ pub(super) fn new_socket_pair<T: CryptoHandler>(
socket,
});
let dyn_internal: Arc<dyn AnyInternalSocket> = 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<T: CryptoHandlerConnectionInstance> Deref for ExternalConnection<T> {
type Target = SendingConnection<T>;
fn deref(&self) -> &Self::Target {
&self.sending
}
}
impl Deref for SendingConnection {
impl<T: CryptoHandlerConnectionInstance> Deref for SendingConnection<T> {
type Target = CommonConnection;
fn deref(&self) -> &Self::Target {
&self.common
}
}
impl ExternalConnection {
impl<E: CryptoHandlerConnectionInstance> ExternalConnection<E> {
pub async fn recv(&mut self) -> Option<Vec<u8>> {
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<E> {
self.sending.clone()
}
}
impl SendingConnection {
impl<E: CryptoHandlerConnectionInstance> SendingConnection<E> {
pub async fn send(&self, data: Vec<u8>) -> 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;
}
}