From 0e654eb4bd8444f8fab7aa321831734147d0e317 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Mon, 6 Nov 2023 20:03:40 +0800
Subject: [PATCH] Ctrlc issues (#75)

---
 Cargo.toml       |  2 +-
 src/error.rs     |  4 ++--
 src/tun2proxy.rs | 55 ++++++++++++++++++++++++------------------------
 tests/proxy.rs   |  5 +----
 4 files changed, 32 insertions(+), 34 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 09db757..0b5f4e5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,7 +10,7 @@ crate-type = ["cdylib", "lib"]
 [dependencies]
 base64 = { version = "0.21" }
 clap = { version = "4.4", features = ["derive"] }
-ctrlc = "3.4"
+ctrlc2 = { version = "3.5", features = ["termination"] }
 digest_auth = "0.3"
 dotenvy = "0.15"
 env_logger = "0.10"
diff --git a/src/error.rs b/src/error.rs
index 2e3d393..e2360ed 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -3,8 +3,8 @@ pub enum Error {
     #[error("std::ffi::NulError {0:?}")]
     Nul(#[from] std::ffi::NulError),
 
-    #[error("ctrlc::Error {0:?}")]
-    InterruptHandler(#[from] ctrlc::Error),
+    #[error("ctrlc2::Error {0:?}")]
+    InterruptHandler(#[from] ctrlc2::Error),
 
     #[error(transparent)]
     Io(#[from] std::io::Error),
diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs
index 6792490..994eeab 100644
--- a/src/tun2proxy.rs
+++ b/src/tun2proxy.rs
@@ -1176,53 +1176,50 @@ impl<'a> TunToProxy<'a> {
     }
 
     #[cfg(any(target_os = "windows", target_os = "linux", target_os = "macos"))]
-    fn prepare_exiting_signal_trigger(&mut self) -> Result<()> {
+    fn prepare_exiting_signal_trigger(&mut self) -> Result<std::thread::JoinHandle<()>> {
         let mut exit_trigger = self.exit_trigger.take().ok_or("Already running")?;
-        ctrlc::set_handler(move || {
-            let mut count = 0;
-            loop {
-                match exit_trigger.write(b"EXIT") {
-                    Ok(_) => {
-                        log::trace!("Exit signal triggered successfully");
-                        break;
-                    }
-                    Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
-                        if count > 5 {
-                            log::error!("Send exit signal failed 5 times, exit anyway");
-                            std::process::exit(1);
-                        }
-                        log::trace!("Send exit signal failed, retry in 1 second");
-                        std::thread::sleep(std::time::Duration::from_secs(1));
-                        count += 1;
-                    }
-                    Err(err) => {
-                        log::error!("Failed to send exit signal: \"{}\"", err);
-                        break;
+        let mut count = 0;
+        let handle = ctrlc2::set_handler(move || -> bool {
+            match exit_trigger.write(b"EXIT") {
+                Ok(_) => {
+                    log::trace!("Exit signal triggered successfully");
+                    true
+                }
+                Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
+                    if count > 5 {
+                        log::error!("Send exit signal failed 5 times, exit anyway");
+                        return true; // std::process::exit(1);
                     }
+                    count += 1;
+                    false
+                }
+                Err(err) => {
+                    log::error!("Failed to send exit signal: \"{}\"", err);
+                    true
                 }
             }
         })?;
-        Ok(())
+        Ok(handle)
     }
 
     pub fn run(&mut self) -> Result<(), Error> {
         #[cfg(any(target_os = "windows", target_os = "linux", target_os = "macos"))]
-        self.prepare_exiting_signal_trigger()?;
+        let handle = self.prepare_exiting_signal_trigger()?;
 
         let mut events = Events::with_capacity(1024);
-        loop {
+        let ret = 'exit_point: loop {
             if let Err(err) = self.poll.poll(&mut events, None) {
                 if err.kind() == std::io::ErrorKind::Interrupted {
                     log::debug!("Poll interrupted: \"{err}\", ignored, continue polling");
                     continue;
                 }
-                return Err(err.into());
+                break 'exit_point Err(Error::from(err));
             }
             for event in events.iter() {
                 match event.token() {
                     EXIT_TOKEN => {
                         if self.exiting_event_handler()? {
-                            return Ok(());
+                            break 'exit_point Ok(());
                         }
                     }
                     EXIT_TRIGGER_TOKEN => {
@@ -1236,7 +1233,11 @@ impl<'a> TunToProxy<'a> {
             self.send_to_smoltcp()?;
             self.clearup_expired_connection()?;
             self.clearup_expired_dns_over_tcp()?;
-        }
+        };
+        #[cfg(any(target_os = "windows", target_os = "linux", target_os = "macos"))]
+        handle.join().unwrap();
+        log::trace!("{:?}", ret);
+        ret
     }
 
     fn exiting_event_handler(&mut self) -> Result<bool> {
diff --git a/tests/proxy.rs b/tests/proxy.rs
index a2521ac..4dcc043 100644
--- a/tests/proxy.rs
+++ b/tests/proxy.rs
@@ -3,15 +3,12 @@
 mod tests {
     extern crate reqwest;
 
-    use std::env;
-    use std::net::IpAddr;
-    use std::str::FromStr;
-
     use fork::Fork;
     use nix::sys::signal;
     use nix::unistd::Pid;
     use serial_test::serial;
     use smoltcp::wire::IpCidr;
+    use std::env;
 
     use tun2proxy::setup::{get_default_cidrs, Setup};
     use tun2proxy::util::str_to_cidr;