diff --git a/Cargo.lock b/Cargo.lock index 1c30310..fbff00b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -308,6 +308,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "chrono" version = "0.4.39" @@ -389,6 +395,16 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctrlc" +version = "3.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46f93780a459b7d656ef7f071fe699c4d3d2cb201c4b24d085b6ddc505276e73" +dependencies = [ + "nix", + "windows-sys 0.59.0", +] + [[package]] name = "deranged" version = "0.3.11" @@ -1400,6 +1416,18 @@ dependencies = [ "tempfile", ] +[[package]] +name = "nix" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "nom" version = "7.1.3" @@ -2270,6 +2298,7 @@ dependencies = [ "async-trait", "bytemuck", "chrono", + "ctrlc", "dotenv", "futures", "hex", diff --git a/Cargo.toml b/Cargo.toml index 40dec26..6d88cf4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,6 +44,7 @@ typenum = "1.18.0" futures = "0.3.31" reqwest = "0.12.18" json = "0.12.4" +ctrlc = "3.4.7" diff --git a/Dockerfile b/Dockerfile index bfc6c66..74f7efd 100644 --- a/Dockerfile +++ b/Dockerfile @@ -8,6 +8,7 @@ RUN apk add --no-cache protobuf-dev git musl-dev lld openssl-dev openssl-libs-st RUN git submodule update --init --recursive +RUN OPENSSL_LIB_DIR=/usr/lib OPENSSL_INCLUDE_DIR=/usr/include/openssl OPENSSL_STATIC=1 RUSTFLAGS="-C relocation-model=static -C linker=ld.lld" cargo test --target x86_64-unknown-linux-musl RUN OPENSSL_LIB_DIR=/usr/lib OPENSSL_INCLUDE_DIR=/usr/include/openssl OPENSSL_STATIC=1 RUSTFLAGS="-C relocation-model=static -C linker=ld.lld" cargo build --profile prod --target x86_64-unknown-linux-musl FROM scratch AS final diff --git a/src/main.rs b/src/main.rs index 51ff59b..a052541 100644 --- a/src/main.rs +++ b/src/main.rs @@ -39,7 +39,7 @@ use std::marker::PhantomData; use std::net::{Ipv4Addr, SocketAddrV4}; use std::ops::{BitAnd, BitOr}; use std::str::FromStr; -use std::sync::{Arc, Weak}; +use std::sync::{Arc, Once, Weak}; use std::time::Duration; use std::{env, fs}; use std::sync::atomic::AtomicU32; @@ -101,6 +101,8 @@ static SECURE_STATION_URL: Lazy = Lazy::new(|| { ) }); +static FORCE_EXIT: Once = Once::new(); + #[tokio::main] async fn main() { CombinedLogger::init(vec![ @@ -125,6 +127,12 @@ async fn main() { ]) .unwrap(); + ctrlc::set_handler(||{ + FORCE_EXIT.call_once_force(|_|{ + println!("attempting exit"); + }); + }).unwrap(); + dotenv::dotenv().ok(); start_servers().await; @@ -286,7 +294,7 @@ async fn start_auth() -> JoinHandle<()> { // let conn = socket_secure.connect(auth_sockaddr).await.unwrap(); - loop { + while !FORCE_EXIT.is_completed() { let Some(conn) = socket_secure.accept().await else { error!("server crashed"); return; @@ -338,7 +346,7 @@ async fn start_secure() -> JoinHandle<()> { // let conn = socket_secure.connect(auth_sockaddr).await.unwrap(); - loop { + while !FORCE_EXIT.is_completed() { let Some(conn) = socket_secure.accept().await else { error!("server crashed"); return; diff --git a/src/prudp/router.rs b/src/prudp/router.rs index d6f8c4b..85c2ba2 100644 --- a/src/prudp/router.rs +++ b/src/prudp/router.rs @@ -4,13 +4,16 @@ use std::marker::PhantomData; use tokio::net::UdpSocket; use std::net::{SocketAddr, SocketAddrV4}; use std::net::SocketAddr::V4; -use std::sync::Arc; +use std::sync::{Arc, Weak}; use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; use tokio::task::JoinHandle; use once_cell::sync::Lazy; use log::{error, info, trace}; use thiserror::Error; +use tokio::select; use tokio::sync::RwLock; +use tokio::time::sleep; use crate::prudp::socket::{new_socket_pair, AnyInternalSocket, CryptoHandler, ExternalSocket}; use crate::prudp::packet::{PRUDPPacket, VirtualPort}; use crate::prudp::router::Error::VirtualPortTaken; @@ -70,15 +73,23 @@ impl Router { } } - async fn server_thread_send_entry(self: Arc, socket: Arc){ + async fn server_thread_send_entry(this: Weak, socket: Arc){ info!("starting datagram thread"); - while self.running.load(Ordering::Relaxed) { + 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) = socket.recv_from(&mut msg_buffer) - .await.expect("Datagram thread crashed due to unexpected error from recv_from"); + 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"); @@ -88,8 +99,10 @@ impl Router { let current_msg = &msg_buffer[0..len]; - tokio::spawn(self.clone().process_prudp_packets(socket.clone(), addr, current_msg.to_vec())); + tokio::spawn(this.process_prudp_packets(socket.clone(), addr, current_msg.to_vec())); } + + println!("exitting datagram") } pub async fn new(addr: SocketAddrV4) -> io::Result<(Arc, JoinHandle<()>)>{ @@ -109,10 +122,10 @@ impl Router { let task = { let socket = socket.clone(); - let server= arc.clone(); + let server= Arc::downgrade(&arc); tokio::spawn(async { - server.server_thread_send_entry(socket).await; + Self::server_thread_send_entry(server, socket).await; }) };