Skip to main content

nitro_cli/common/
signal_handler.rs

1// Copyright 2020-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3#![deny(missing_docs)]
4#![deny(warnings)]
5
6use log::warn;
7use nix::sys::signal::{SigSet, Signal};
8use nix::sys::signal::{SIGHUP, SIGINT, SIGQUIT, SIGTERM};
9use std::os::unix::io::RawFd;
10use std::thread;
11
12use crate::common::{NitroCliErrorEnum, NitroCliFailure, NitroCliResult};
13use crate::new_nitro_cli_failure;
14
15/// The custom handler of POSIX signals.
16pub struct SignalHandler {
17    sig_set: Option<SigSet>,
18}
19
20impl SignalHandler {
21    /// Create a new `SignalHandler` instance from the given list of signals.
22    pub fn new(signals: &[Signal]) -> Self {
23        let mut sig_set = SigSet::empty();
24        for signal in signals.iter() {
25            sig_set.add(*signal);
26        }
27
28        SignalHandler {
29            sig_set: Some(sig_set),
30        }
31    }
32
33    /// Create a new `SignalHandler` instance from a default list of signals.
34    pub fn new_with_defaults() -> Self {
35        SignalHandler::new(&[SIGINT, SIGQUIT, SIGTERM, SIGHUP])
36    }
37
38    /// Mask (block) all signals covered by the handler.
39    pub fn mask_all(self) -> NitroCliResult<Self> {
40        if let Some(set) = self.sig_set {
41            set.thread_block().map_err(|e| {
42                new_nitro_cli_failure!(
43                    &format!("Masking signals covered by handler failed: {e:?}"),
44                    NitroCliErrorEnum::SignalMaskingError
45                )
46            })?;
47        }
48
49        Ok(self)
50    }
51
52    /// Unmask (unblock) all signals covered by the handler.
53    pub fn unmask_all(self) -> NitroCliResult<Self> {
54        if let Some(set) = self.sig_set {
55            set.thread_unblock().map_err(|e| {
56                new_nitro_cli_failure!(
57                    &format!("Unmasking signals covered by handler failed: {e:?}"),
58                    NitroCliErrorEnum::SignalUnmaskingError
59                )
60            })?;
61        }
62
63        Ok(self)
64    }
65
66    /// Start listening for events on a dedicated thread and handle them using the provided function.
67    pub fn start_handler(&mut self, fd: RawFd, handler: fn(RawFd, Signal) -> bool) {
68        if self.sig_set.is_none() {
69            return;
70        }
71
72        let thread_sig_set = self.sig_set.take().unwrap();
73        thread::spawn(move || {
74            let mut stop = false;
75            while !stop {
76                stop = match thread_sig_set.wait() {
77                    Ok(signal) => handler(fd, signal),
78                    Err(e) => {
79                        warn!("Error listening for signals: {}", e);
80                        true
81                    }
82                };
83            }
84        });
85    }
86}