From 1c048a17ccd39686eec6fba3a71540b0222291f1 Mon Sep 17 00:00:00 2001
From: lokka30 <59464084+lokka30@users.noreply.github.com>
Date: Sun, 27 Oct 2024 23:29:59 +0800
Subject: [PATCH] Implement cross platform signal handling with `tokio` (#183)

* ctrlc: also handle `sigterm` and `sighup` via `termination` feature

This commit enables the `termination` feature in `ctrlc` so that
`sigterm` and `sighup` are also handled on UNIX systems. This in turn
improves Docker compatibility since containers now listen to the
`SIGTERM` signal it sends to containers when trying to stop the
container via `Ctrl-C` whilst attached.

* Replace `ctrlc` signal handling with `tokio`

This removes the `ctrlc` crate from the project, replacing (and upgrading) its functionality to also handle `sigterm` and `sighup`.

Sigterm handling is working on UNIX with this commit, I haven't tested Windows compatibility with `Ctrl-C` nor have I tested `sighup` or `sigint`.

Of course, compiles and runs on UNIX. Haven't tested Windows.

* fix: remove unused import; collapse if statement

* main: fix per-os signal handling

Will only compile Windows-specific/Unix-specific code when on their respective environments.

Handles signals in a separate tokio thread so it doesn't suspend the main thread.

Tested to work on UNIX (macOS) but not tested on Windows.

* main: cargo fmt

* main: fix setup_sighandler fn signature on windows

* main: attempt 2 at fixing windows sighandler

* main: enable ctrlc for all non-unix systems

Previously, ctrlc usage in tokio was only enabled on Windows systems
(since UNIX systems have their own signal handling in Pumpkin). Instead
of limiting ctrlc to just Windows, we can make it run on any non-UNIX
system since it should be platform independent anyhow.
---
 Cargo.toml          |  1 +
 pumpkin/Cargo.toml  |  2 --
 pumpkin/src/main.rs | 59 +++++++++++++++++++++++++++++++++++++--------
 3 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 9e1f2106..fa2adc6c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -30,6 +30,7 @@ tokio = { version = "1.41", features = [
     "rt-multi-thread",
     "sync",
     "io-std",
+    "signal",
 ] }
 
 thiserror = "1.0"
diff --git a/pumpkin/Cargo.toml b/pumpkin/Cargo.toml
index 5e83ea91..52dbafc0 100644
--- a/pumpkin/Cargo.toml
+++ b/pumpkin/Cargo.toml
@@ -37,8 +37,6 @@ rand = "0.8.5"
 
 num-bigint = "0.4"
 
-ctrlc = "3.4"
-
 # encryption
 rsa = "0.9.6"
 rsa-der = "0.3.0"
diff --git a/pumpkin/src/main.rs b/pumpkin/src/main.rs
index 9823afda..038ceb59 100644
--- a/pumpkin/src/main.rs
+++ b/pumpkin/src/main.rs
@@ -25,6 +25,10 @@ use client::Client;
 use server::{ticker::Ticker, Server};
 use std::io::{self};
 use tokio::io::{AsyncBufReadExt, BufReader};
+#[cfg(not(unix))]
+use tokio::signal::ctrl_c;
+#[cfg(unix)]
+use tokio::signal::unix::{signal, SignalKind};
 
 use std::sync::Arc;
 
@@ -94,16 +98,13 @@ async fn main() -> io::Result<()> {
     //     .enable_all()
     //     .build()
     //     .unwrap();
-    ctrlc::set_handler(|| {
-        log::warn!(
-            "{}",
-            TextComponent::text("Stopping Server")
-                .color_named(NamedColor::Red)
-                .to_pretty_console()
-        );
-        std::process::exit(0);
-    })
-    .unwrap();
+
+    tokio::spawn(async {
+        setup_sighandler()
+            .await
+            .expect("Unable to setup signal handlers");
+    });
+
     // ensure rayon is built outside of tokio scope
     rayon::ThreadPoolBuilder::new().build_global().unwrap();
     let default_panic = std::panic::take_hook();
@@ -203,6 +204,44 @@ async fn main() -> io::Result<()> {
     }
 }
 
+fn handle_interrupt() {
+    log::warn!(
+        "{}",
+        TextComponent::text("Received interrupt signal; stopping server...")
+            .color_named(NamedColor::Red)
+            .to_pretty_console()
+    );
+    std::process::exit(0);
+}
+
+// Non-UNIX Ctrl-C handling
+#[cfg(not(unix))]
+async fn setup_sighandler() -> io::Result<()> {
+    if ctrl_c().await.is_ok() {
+        handle_interrupt();
+    }
+
+    Ok(())
+}
+
+// Unix signal handling
+#[cfg(unix)]
+async fn setup_sighandler() -> io::Result<()> {
+    if signal(SignalKind::interrupt())?.recv().await.is_some() {
+        handle_interrupt();
+    }
+
+    if signal(SignalKind::hangup())?.recv().await.is_some() {
+        handle_interrupt();
+    }
+
+    if signal(SignalKind::terminate())?.recv().await.is_some() {
+        handle_interrupt();
+    }
+
+    Ok(())
+}
+
 fn setup_console(server: Arc<Server>) {
     tokio::spawn(async move {
         let stdin = tokio::io::stdin();