diff --git a/src/main.rs b/src/main.rs index 7260c2e..bd2fa26 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,12 +5,11 @@ use tokio_util::sync::CancellationToken; mod config; mod server; use crate::config::Config; -use crate::server::spawn_socks5_server; +use crate::server::server_executor; use flexi_logger::{AdaptiveFormat, Age, Cleanup, Criterion, Duplicate, FileSpec, Logger, Naming}; -#[tokio::main] -async fn main() { +fn main() { Logger::try_with_str("info") .unwrap() .log_to_file(FileSpec::default()) @@ -25,21 +24,15 @@ async fn main() { .start_with_specfile("logspec.toml") .unwrap(); - let cfg = tokio::task::spawn_blocking(Config::get) - .await - .expect("get config"); + let cfg = Config::get(); log::info!("cfg: {:#?}", cfg); let token = CancellationToken::new(); let child_token = token.child_token(); + let handle = std::thread::spawn(move || server_executor(cfg, child_token)); - let (r, _) = tokio::join!( - spawn_socks5_server(cfg, child_token), - tokio::spawn(async move { - tokio::time::sleep(std::time::Duration::from_secs(10)).await; - token.cancel(); - }) - ); + std::thread::sleep(std::time::Duration::from_secs(10)); + token.cancel(); - r.unwrap(); + handle.join().unwrap(); } diff --git a/src/server.rs b/src/server.rs index 41a17fc..ef57612 100644 --- a/src/server.rs +++ b/src/server.rs @@ -12,6 +12,16 @@ use tokio_util::sync::CancellationToken; use crate::config::Config; use crate::config::PasswordAuth; +pub fn server_executor(cfg: Config, token: CancellationToken) { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + spawn_socks5_server(cfg, token).await.unwrap(); + }) +} + pub async fn spawn_socks5_server(cfg: Config, token: CancellationToken) -> Result<()> { let mut server_config = fast_socks5::server::Config::default(); server_config.set_request_timeout(cfg.request_timeout); @@ -34,8 +44,7 @@ pub async fn spawn_socks5_server(cfg: Config, token: CancellationToken) -> Resul log::info!("Listen for socks connections @ {}", &cfg.listen_addr); - // Standard TCP loop - while let Some(socket_res) = or_chancel(incoming.next(), token.child_token()).await { + while let Some(socket_res) = check_cancelled(incoming.next(), token.child_token()).await { match socket_res { Ok(socket) => { let child_token = token.child_token(); @@ -50,13 +59,13 @@ pub async fn spawn_socks5_server(cfg: Config, token: CancellationToken) -> Resul Ok(()) } -async fn or_chancel(future: F, token: CancellationToken) -> Option +async fn check_cancelled(future: F, token: CancellationToken) -> Option where F: Future>, { select! { _ = token.cancelled() => { - log::error!("canceled"); + log::error!("accept canceled"); None } res = future => { @@ -71,7 +80,6 @@ where T: AsyncRead + AsyncWrite + Unpin, { tokio::spawn(async move { - // Wait for either cancellation or a very long time let result = select! { _ = token.cancelled() => { Err("Client connection canceled".to_string())