From 757d5b9f58ad88ada5594788e94b6feee6be6fcc Mon Sep 17 00:00:00 2001
From: Laszlo Nagy <rizsotto@gmail.com>
Date: Sun, 19 Jan 2025 00:37:49 +1100
Subject: [PATCH] rust: add supervisor function to forward signals to child
 processes

---
 rust/bear/Cargo.toml                 |   9 ++
 rust/bear/src/intercept/mod.rs       |   1 +
 rust/bear/src/intercept/supervise.rs | 118 +++++++++++++++++++++++++++
 3 files changed, 128 insertions(+)
 create mode 100644 rust/bear/src/intercept/supervise.rs

diff --git a/rust/bear/Cargo.toml b/rust/bear/Cargo.toml
index c94dbf2c..3ec78028 100644
--- a/rust/bear/Cargo.toml
+++ b/rust/bear/Cargo.toml
@@ -40,6 +40,15 @@ nom.workspace = true
 regex.workspace = true
 rand.workspace = true
 tempfile.workspace = true
+nix = { version = "0.29", optional = true, features = ["signal", "process"] }
+winapi = { version = "0.3", optional = true, features = ["processthreadsapi", "winnt", "handleapi"] }
+signal-hook = "0.3.17"
+
+[target.'cfg(unix)'.dependencies]
+nix = { version = "0.29", features = ["signal", "process"] }
+
+[target.'cfg(windows)'.dependencies]
+winapi = { version = "0.3", features = ["processthreadsapi", "winnt", "handleapi"] }
 
 [profile.release]
 strip = true
diff --git a/rust/bear/src/intercept/mod.rs b/rust/bear/src/intercept/mod.rs
index b9c1edd3..03d2e888 100644
--- a/rust/bear/src/intercept/mod.rs
+++ b/rust/bear/src/intercept/mod.rs
@@ -19,6 +19,7 @@ use std::sync::Arc;
 use std::{env, fmt, thread};
 
 pub mod persistence;
+pub mod supervise;
 pub mod tcp;
 
 /// Declare the environment variables used by the intercept mode.
diff --git a/rust/bear/src/intercept/supervise.rs b/rust/bear/src/intercept/supervise.rs
new file mode 100644
index 00000000..323c6495
--- /dev/null
+++ b/rust/bear/src/intercept/supervise.rs
@@ -0,0 +1,118 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+use anyhow::Result;
+use nix::libc::c_int;
+#[cfg(unix)]
+use nix::sys::signal::{kill, Signal};
+#[cfg(unix)]
+use nix::unistd::Pid;
+use std::process::{Command, ExitStatus};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::thread;
+use std::time::Duration;
+#[cfg(windows)]
+use winapi::shared::minwindef::FALSE;
+#[cfg(windows)]
+use winapi::um::processthreadsapi::{OpenProcess, TerminateProcess};
+#[cfg(windows)]
+use winapi::um::winnt::{PROCESS_TERMINATE, SYNCHRONIZE};
+
+/// This method supervises the execution of a command.
+///
+/// It starts the command and waits for its completion. It also forwards
+/// signals to the child process. The method returns the exit status of the
+/// child process.
+pub fn supervise(command: &mut Command) -> Result<ExitStatus> {
+    let mut child = command.spawn()?;
+
+    let child_pid = child.id();
+    let running = Arc::new(AtomicBool::new(true));
+    let running_in_thread = running.clone();
+
+    let mut signals = signal_hook::iterator::Signals::new([
+        signal_hook::consts::SIGINT,
+        signal_hook::consts::SIGTERM,
+    ])?;
+
+    #[cfg(unix)]
+    {
+        signals.add_signal(signal_hook::consts::SIGHUP)?;
+        signals.add_signal(signal_hook::consts::SIGQUIT)?;
+        signals.add_signal(signal_hook::consts::SIGALRM)?;
+        signals.add_signal(signal_hook::consts::SIGUSR1)?;
+        signals.add_signal(signal_hook::consts::SIGUSR2)?;
+        signals.add_signal(signal_hook::consts::SIGCONT)?;
+        signals.add_signal(signal_hook::consts::SIGSTOP)?;
+    }
+
+    let handler = thread::spawn(move || {
+        for signal in signals.forever() {
+            log::debug!("Received signal: {:?}", signal);
+            if forward_signal(signal, child_pid) {
+                // If the signal caused termination, we should stop the process.
+                running_in_thread.store(false, Ordering::SeqCst);
+                break;
+            }
+        }
+    });
+
+    while running.load(Ordering::SeqCst) {
+        thread::sleep(Duration::from_millis(100));
+    }
+    handler.join().unwrap();
+
+    let exit_status = child.wait()?;
+
+    Ok(exit_status)
+}
+
+#[cfg(windows)]
+fn forward_signal(_: c_int, child_pid: u32) -> bool {
+    let process_handle = unsafe { OpenProcess(PROCESS_TERMINATE | SYNCHRONIZE, FALSE, child_pid) };
+    if process_handle.is_null() {
+        let err = unsafe { winapi::um::errhandling::GetLastError() };
+        log::error!("Failed to open process: {}", err);
+        // If the process handle is not valid, presume the process is not running anymore.
+        return true;
+    }
+
+    let terminated = unsafe { TerminateProcess(process_handle, 1) };
+    if terminated == FALSE {
+        let err = unsafe { winapi::um::errhandling::GetLastError() };
+        log::error!("Failed to terminate process: {}", err);
+    }
+
+    // Ensure proper handle closure
+    unsafe { winapi::um::handleapi::CloseHandle(process_handle) };
+
+    // Return true if the process was terminated.
+    terminated == TRUE
+}
+
+#[cfg(unix)]
+fn forward_signal(signal: c_int, child_pid: u32) -> bool {
+    // Forward the signal to the child process
+    if let Err(e) = kill(
+        Pid::from_raw(child_pid as i32),
+        Signal::try_from(signal).ok(),
+    ) {
+        log::error!("Error forwarding signal: {}", e);
+    }
+
+    // Return true if the process was terminated.
+    match kill(Pid::from_raw(child_pid as i32), None) {
+        Ok(_) => {
+            log::debug!("Checking if the process is still running... yes");
+            false
+        }
+        Err(nix::Error::ESRCH) => {
+            log::debug!("Checking if the process is still running... no");
+            true
+        }
+        Err(_) => {
+            log::debug!("Checking if the process is still running... presume dead");
+            true
+        }
+    }
+}