Chore: add signal-handling to run-blocking

Author: nichtsfrei

rust/Cargo.lock

@@ -20,6 +20,8 @@ tokio-rustls = { workspace = true }
 tracing = { workspace = true }
 uuid = { workspace = true }
 futures = { workspace = true }
+signal-hook-tokio = { version = "0.4.0", features = ["futures-v0_3"] }
+signal-hook = "0.4.3"
 
 [dev-dependencies]
 insta = { workspace = true }

rust/crates/greenbone-scanner-framework/Cargo.toml

@@ -5,7 +5,7 @@ use std::{
     pin::Pin,
     sync::{
         Arc, RwLock,
-        atomic::{AtomicUsize, Ordering},
+        atomic::{AtomicI32, AtomicUsize, Ordering},
     },
 };
 
@@ -51,6 +51,10 @@ mod post_scans_id;
 mod tls;
 use post_scans::PostScansHandler;
 use post_scans_id::{PostScansId, PostScansIdHandler};
+
+use futures::stream::StreamExt;
+use signal_hook::consts::signal::*;
+use signal_hook_tokio::Signals;
 use tokio::net::TcpListener;
 
 pub trait ExternalError: core::error::Error + Send + Sync + 'static {}
@@ -375,8 +379,29 @@ impl RuntimeBuilder<runtime_builder_states::DeleteScanIDSet> {
     }
 }
 
+static RUN: AtomicI32 = AtomicI32::new(0);
+
+async fn handle_signals(mut signals: Signals) {
+    while let Some(signal) = signals.next().await {
+        match signal {
+            SIGHUP => {
+                tracing::info!("Ignoring SIGHUP signal.");
+            }
+            SIGTERM | SIGINT | SIGQUIT => {
+                tracing::info!(signal, "Exit based on signal.");
+                RUN.store(128 + signal, Ordering::Relaxed);
+            }
+            _ => unreachable!(),
+        }
+    }
+}
+
 impl RuntimeBuilder<runtime_builder_states::End> {
-    pub async fn run_blocking(self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    pub async fn run_blocking(self) -> Result<i32, Box<dyn std::error::Error + Send + Sync>> {
+        let signals = Signals::new([SIGHUP, SIGTERM, SIGINT, SIGQUIT])?;
+        let _handle_guard = signals.handle();
+        tokio::task::spawn(handle_signals(signals));
+
         let scanner = Arc::new(self.build_scanner());
         let tls_config = match &self.tls {
             Some(x) => Some(tls::tls_config(x)?),
@@ -395,7 +420,7 @@ impl RuntimeBuilder<runtime_builder_states::End> {
             let config = Arc::new(tls_config.config);
             let tls_acceptor = tokio_rustls::TlsAcceptor::from(config);
 
-            loop {
+            while RUN.load(Ordering::Relaxed) == 0 {
                 let (tcp_stream, _remote_addr) = incoming.accept().await?;
                 let tls_acceptor = tls_acceptor.clone();
                 let identifier = tls_config.client_identifier.clone();
@@ -433,7 +458,7 @@ impl RuntimeBuilder<runtime_builder_states::End> {
         } else {
             use hyper::server::conn::http1::Builder;
             tracing::info!("listening on http://{}", self.listener_address);
-            loop {
+            while RUN.load(Ordering::Relaxed) == 0 {
                 let (tcp_stream, _remote_addr) = incoming.accept().await?;
                 let ctx = scanner.clone();
                 let handlers = handlers.clone();
@@ -459,6 +484,7 @@ impl RuntimeBuilder<runtime_builder_states::End> {
                 });
             }
         }
+        Ok(RUN.load(Ordering::Relaxed))
     }
 }
 

rust/crates/greenbone-scanner-framework/src/lib.rs

@@ -18,6 +18,7 @@ mod vts;
 use sqlx::migrate::Migrator;
 use std::{
     marker::{Send, Sync},
+    process::ExitCode,
     sync::Arc,
 };
 
@@ -59,8 +60,7 @@ async fn setup_sqlite(config: &Config) -> Result<SqlitePool> {
     Ok(result)
 }
 
-#[tokio::main]
-async fn main() -> Result<()> {
+async fn _main() -> Result<i32> {
     let config = Config::load();
     let _guard = config.logging.init();
 
@@ -111,3 +111,14 @@ async fn main() -> Result<()> {
         .run_blocking()
         .await
 }
+
+#[tokio::main]
+async fn main() -> ExitCode {
+    match _main().await {
+        Ok(x) => ExitCode::from(x as u8),
+        Err(error) => {
+            tracing::error!(%error, "Unexpected error result");
+            ExitCode::from(1)
+        }
+    }
+}