seal/security/
cleanup.rs

1//! Signal handler registration and cleanup coordination
2
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::sync::Arc;
5
6/// Global flag for coordinating cleanup across threads
7static CLEANUP_REQUESTED: AtomicBool = AtomicBool::new(false);
8
9/// Check if cleanup has been requested (e.g., via signal)
10pub fn is_cleanup_requested() -> bool {
11    CLEANUP_REQUESTED.load(Ordering::SeqCst)
12}
13
14/// Request cleanup (sets the global flag)
15pub fn request_cleanup() {
16    CLEANUP_REQUESTED.store(true, Ordering::SeqCst);
17}
18
19/// Register signal handlers for graceful shutdown
20/// Handles SIGINT (Ctrl+C), SIGTERM, and SIGHUP on Unix
21pub fn register_signal_handlers() -> anyhow::Result<()> {
22    // Register SIGINT handler (Ctrl+C) - works on all platforms
23    ctrlc::set_handler(move || {
24        eprintln!("\nReceived interrupt signal, cleaning up...");
25        request_cleanup();
26        std::process::exit(5); // Exit code 5 for signal received
27    })
28    .map_err(|e| anyhow::anyhow!("Failed to set Ctrl+C handler: {}", e))?;
29
30    // Register additional Unix signals (SIGTERM, SIGHUP)
31    #[cfg(unix)]
32    {
33        use signal_hook::consts::signal::*;
34        use signal_hook::flag;
35
36        let cleanup_flag = Arc::new(AtomicBool::new(false));
37
38        // SIGTERM
39        flag::register(SIGTERM, Arc::clone(&cleanup_flag))
40            .map_err(|e| anyhow::anyhow!("Failed to register SIGTERM handler: {}", e))?;
41
42        // SIGHUP
43        flag::register(SIGHUP, Arc::clone(&cleanup_flag))
44            .map_err(|e| anyhow::anyhow!("Failed to register SIGHUP handler: {}", e))?;
45
46        // Spawn thread to monitor signals
47        std::thread::spawn(move || loop {
48            if cleanup_flag.load(Ordering::Relaxed) {
49                eprintln!("\nReceived termination signal, cleaning up...");
50                request_cleanup();
51                std::process::exit(5);
52            }
53            std::thread::sleep(std::time::Duration::from_millis(100));
54        });
55    }
56
57    Ok(())
58}
59
60/// Cleanup guard that ensures cleanup runs on drop
61pub struct CleanupGuard {
62    cleanup_fn: Option<Box<dyn FnOnce() + Send>>,
63}
64
65impl CleanupGuard {
66    /// Create a new cleanup guard with a cleanup function
67    pub fn new<F>(cleanup_fn: F) -> Self
68    where
69        F: FnOnce() + Send + 'static,
70    {
71        Self {
72            cleanup_fn: Some(Box::new(cleanup_fn)),
73        }
74    }
75
76    /// Explicitly run cleanup (also runs automatically on drop)
77    pub fn cleanup(mut self) {
78        if let Some(cleanup_fn) = self.cleanup_fn.take() {
79            cleanup_fn();
80        }
81    }
82}
83
84impl Drop for CleanupGuard {
85    fn drop(&mut self) {
86        if let Some(cleanup_fn) = self.cleanup_fn.take() {
87            cleanup_fn();
88        }
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn test_cleanup_flag() {
98        assert!(!is_cleanup_requested());
99        request_cleanup();
100        assert!(is_cleanup_requested());
101    }
102
103    #[test]
104    fn test_cleanup_guard() {
105        use std::sync::Arc;
106        let cleaned = Arc::new(AtomicBool::new(false));
107        let cleaned_clone = Arc::clone(&cleaned);
108
109        {
110            let _guard = CleanupGuard::new(move || {
111                cleaned_clone.store(true, Ordering::SeqCst);
112            });
113            assert!(!cleaned.load(Ordering::SeqCst));
114        }
115
116        assert!(cleaned.load(Ordering::SeqCst));
117    }
118}
seal/security/cleanup.rs

seal/security/
cleanup.rs
