wala-rust

commit 2150cdf50ac1ffd672e2ff43459511122fbcaeb1
parent 46c4e53313bb4b4009da239b86150dad3726fd5f
Author: lash <dev@holbrook.no>
Date:   Tue, 28 Mar 2023 09:08:33 +0100

Add signal handling option

Diffstat:
M
Cargo.toml
 | 
4
++++
M
src/main.rs
 | 
33
++++++++++++++++++++++++++++-----

2 files changed, 32 insertions(+), 5 deletions(-)
diff --git a/Cargo.toml b/Cargo.toml
@@ -64,6 +64,10 @@ optional = true
 version = "^7.0.0"
 optional = true
 
+[dependencies.signal-hook]
+version = "^0.3.15"
+optional = true
+
 [features]
 pgpauth = ["sequoia-openpgp", "base64", "nettle"]
 dev = []
diff --git a/src/main.rs b/src/main.rs
@@ -23,10 +23,15 @@ use std::io::{
     Seek,
     empty,
 };
+use std::time::Duration;
 
-use env_logger;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
 
+use env_logger;
 use ascii::AsciiStr;
+use signal_hook::flag;
+use signal_hook::consts;
 
 use wala::auth::{
     AuthSpec,
@@ -255,16 +260,34 @@ fn main() {
     };
     let srv = Server::new(srv_cfg).unwrap();
 
-    loop {
-        let b = srv.recv();
-        let mut req: Request;
+    let term = Arc::new(AtomicBool::new(false));
+    signal_hook::flag::register(signal_hook::consts::SIGINT, Arc::clone(&term)).unwrap();
+
+    #[cfg(feature = "docker")]
+    signal_hook::flag::register(signal_hook::consts::SIGTERM, Arc::clone(&term)).unwrap();
+    
+    const loop_timeout: Duration = Duration::new(1, 0);
+
+    while !term.load(Ordering::Relaxed) {
+
+        let b = srv.recv_timeout(loop_timeout);
+        let mut hasreq: Option<Request>;
         match b {
-            Ok(v) => req = v,
+            Ok(v) => hasreq = v,
             Err(e) => {
                 error!("{}", e);
                 break;
             }
         };
+        let mut req: Request;
+        match hasreq {
+            Some(v) => {
+                req = v;
+            },
+            None => {
+                continue 
+            }
+        };
 
         let method = req.method().clone();
         match &method {