update examples

Author: ssrlive

Cargo.toml

@@ -33,14 +33,15 @@ tokio = { version = "1", default-features = false, features = [
 
 [dev-dependencies]
 clap = { version = "4", features = ["derive"] }
-ctrlc2 = { version = "3", features = ["tokio", "termination"] }
+ctrlc2 = { version = "3", features = ["async", "termination"] }
 dotenvy = "0.15"
 env_logger = "0.11"
 hickory-proto = "0.25"
 log = "0.4"
 moka = { version = "0.12", features = ["future"] }
 rand = "0.9"
 tokio = { version = "1", features = ["rt-multi-thread"] }
+tokio-util = { version = "0.7", features = [] }
 
 [[example]]
 name = "demo-client"

examples/echo-server.rs

@@ -79,15 +79,13 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
         Ok::<(), std::io::Error>(())
     });
 
-    let (tx, mut rx) = tokio::sync::mpsc::channel::<()>(1);
-
-    ctrlc2::set_async_handler(async move {
-        tx.send(()).await.unwrap();
+    let ctrlc = ctrlc2::AsyncCtrlC::new(|| {
         log::info!("Ctrl-C received, shutting down...");
+        true
     })
-    .await;
+    .expect("cannot create Ctrl+C handler");
 
-    rx.recv().await.unwrap();
+    ctrlc.await;
     log::info!("Exiting...");
 
     Ok(())

examples/s5-server.rs

@@ -5,13 +5,14 @@ use socks5_impl::{
 };
 use std::{
     net::{SocketAddr, ToSocketAddrs},
-    sync::{Arc, atomic::AtomicBool},
+    sync::Arc,
 };
 use tokio::{
     io,
     net::{TcpStream, UdpSocket},
     sync::Mutex,
 };
+use tokio_util::sync::CancellationToken;
 
 /// Simple socks5 proxy server.
 #[derive(clap::Parser, Debug, Clone, PartialEq, Eq)]
@@ -55,52 +56,48 @@ async fn main() -> Result<()> {
     let default = format!("{}={:?}", module_path!(), opt.verbosity);
     env_logger::Builder::from_env(env_logger::Env::default().default_filter_or(default)).init();
 
-    let exiting_flag = Arc::new(AtomicBool::new(false));
-    let exiting_flag_clone = exiting_flag.clone();
+    let token = CancellationToken::new();
+    let cloned_token = token.clone();
 
-    let local_addr = opt.listen_addr;
-
-    ctrlc2::set_async_handler(async move {
-        exiting_flag_clone.store(true, std::sync::atomic::Ordering::Relaxed);
-
-        let addr = if local_addr.is_ipv6() {
-            SocketAddr::from((std::net::Ipv6Addr::LOCALHOST, local_addr.port()))
-        } else {
-            SocketAddr::from((std::net::Ipv4Addr::LOCALHOST, local_addr.port()))
-        };
-        let _ = std::net::TcpStream::connect(addr);
+    let ctrlc = ctrlc2::AsyncCtrlC::new(move || {
         log::info!("");
         log::info!("Ctrl-C received, shutting down...");
+        cloned_token.cancel();
+        true
     })
-    .await;
+    .expect("Failed to set up Ctrl-C handler");
 
     match (opt.username, opt.password) {
         (Some(username), password) => {
             let password = password.unwrap_or_default();
             let auth = Arc::new(auth::UserKeyAuth::new(&username, &password));
-            main_loop(auth, opt.listen_addr, Some(exiting_flag)).await?;
+            main_loop(auth, opt.listen_addr, token).await?;
         }
         _ => {
             let auth = Arc::new(auth::NoAuth);
-            main_loop(auth, opt.listen_addr, Some(exiting_flag)).await?;
+            main_loop(auth, opt.listen_addr, token).await?;
         }
     }
 
+    ctrlc.await;
+
     Ok(())
 }
 
-async fn main_loop<S>(auth: auth::AuthAdaptor<S>, listen_addr: SocketAddr, exiting_flag: Option<Arc<AtomicBool>>) -> Result<()>
+async fn main_loop<S>(auth: auth::AuthAdaptor<S>, listen_addr: SocketAddr, token: CancellationToken) -> Result<()>
 where
     S: Send + Sync + 'static,
 {
     let server = Server::bind(listen_addr, auth).await?;
 
-    while let Ok((conn, _)) = server.accept().await {
-        if let Some(exiting_flag) = &exiting_flag {
-            if exiting_flag.load(std::sync::atomic::Ordering::Relaxed) {
+    loop {
+        let (conn, _) = tokio::select! {
+            _ = token.cancelled() => {
+                log::info!("CancellationToken fired, session will be closed");
                 break;
             }
-        }
+            conn = server.accept() => conn?,
+        };
         tokio::spawn(async move {
             if let Err(err) = handle(conn).await {
                 log::error!("{err}");