From f735bd5b816526f71af05397dd9e70ef75f28900 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Wed, 9 Oct 2024 18:25:43 +0800
Subject: [PATCH] Ctrl-C handler

---
 src/bin/main.rs | 13 ++++++++++---
 src/lib.rs      | 14 +++++++++++++-
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/bin/main.rs b/src/bin/main.rs
index 03df48a..b67b1d4 100644
--- a/src/bin/main.rs
+++ b/src/bin/main.rs
@@ -8,8 +8,11 @@ fn main() -> Result<(), BoxError> {
             return Err("C API is not supported for server".into());
         }
 
-        let join = ctrlc2::set_handler(|| {
+        let ctrlc_fired = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
+        let ctrlc_fired_clone = ctrlc_fired.clone();
+        let ctrl_handle = ctrlc2::set_handler(move || {
             log::info!("Ctrl-C received, exiting...");
+            ctrlc_fired_clone.store(true, std::sync::atomic::Ordering::SeqCst);
             unsafe { overtls::over_tls_client_stop() };
             true
         })?;
@@ -31,7 +34,9 @@ fn main() -> Result<(), BoxError> {
 
             unsafe { overtls::over_tls_client_run(config_path, opt.verbosity, Some(port_cb), std::ptr::null_mut()) };
 
-            join.join().expect("Couldn't join on the associated thread");
+            if ctrlc_fired.load(std::sync::atomic::Ordering::SeqCst) {
+                ctrl_handle.join().map_err(|e| format!("{:?}", e))?;
+            }
         } else if let Some(url) = opt.url_of_node.as_ref() {
             let url_str = std::ffi::CString::new(url.as_str())?;
             let url_ptr = url_str.as_ptr();
@@ -42,7 +47,9 @@ fn main() -> Result<(), BoxError> {
 
             unsafe { overtls::over_tls_client_run_with_ssr_url(url_ptr, listen_addr, opt.verbosity, Some(port_cb), std::ptr::null_mut()) };
 
-            join.join().expect("Couldn't join on the associated thread");
+            if ctrlc_fired.load(std::sync::atomic::Ordering::SeqCst) {
+                ctrl_handle.join().map_err(|e| format!("{:?}", e))?;
+            }
         } else {
             return Err("Config file or node URL is required".into());
         }
diff --git a/src/lib.rs b/src/lib.rs
index a5c45f1..69c22cb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -67,13 +67,18 @@ pub(crate) fn b64str_to_address(s: &str, url_safe: bool) -> Result<Address> {
 
 #[doc(hidden)]
 pub async fn async_main(config: Config, allow_shutdown: bool, shutdown_token: CancellationToken) -> Result<()> {
+    let ctrlc_fired = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
+    let mut ctrlc_handle = None;
     if allow_shutdown {
         let shutdown_token_clone = shutdown_token.clone();
-        ctrlc2::set_async_handler(async move {
+        let ctrlc_fired_clone = ctrlc_fired.clone();
+        let handle = ctrlc2::set_async_handler(async move {
             log::info!("Ctrl-C received, exiting...");
+            ctrlc_fired_clone.store(true, std::sync::atomic::Ordering::SeqCst);
             shutdown_token_clone.cancel();
         })
         .await;
+        ctrlc_handle = Some(handle);
     }
 
     let main_body = async {
@@ -92,6 +97,13 @@ pub async fn async_main(config: Config, allow_shutdown: bool, shutdown_token: Ca
             return Err("Config is not a client config".into());
         }
 
+        if ctrlc_fired.load(std::sync::atomic::Ordering::SeqCst) {
+            let Some(handle) = ctrlc_handle else {
+                return Ok(());
+            };
+            log::info!("Waiting for Ctrl-C handler to finish...");
+            handle.await.map_err(|e| e.to_string())?;
+        }
         Ok(())
     };