Allow sending custom signals to a vm

Author: coolreader18

vm/src/lib.rs

@@ -67,7 +67,7 @@ mod pyobjectrc;
 pub mod readline;
 pub mod scope;
 mod sequence;
-mod signal;
+pub mod signal;
 pub mod sliceable;
 pub mod stdlib;
 pub mod suggestion;

vm/src/signal.rs

@@ -1,8 +1,10 @@
 #![cfg_attr(target_os = "wasi", allow(dead_code))]
-use crate::{PyObjectRef, PyResult, VirtualMachine};
+use crate::{PyResult, VirtualMachine};
+use std::fmt;
 use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::mpsc;
 
-pub const NSIG: usize = 64;
+pub(crate) const NSIG: usize = 64;
 static ANY_TRIGGERED: AtomicBool = AtomicBool::new(false);
 // hack to get around const array repeat expressions, rust issue #79270
 #[allow(clippy::declare_interior_mutable_const)]
@@ -37,9 +39,56 @@ fn trigger_signals(vm: &VirtualMachine) -> PyResult<()> {
             }
         }
     }
+    if let Some(signal_rx) = &vm.signal_rx {
+        for f in signal_rx.rx.try_iter() {
+            f(vm)?;
+        }
+    }
     Ok(())
 }
 
 pub(crate) fn set_triggered() {
     ANY_TRIGGERED.store(true, Ordering::Release);
 }
+
+pub type UserSignal = Box<dyn FnOnce(&VirtualMachine) -> PyResult<()> + Send>;
+
+#[derive(Clone, Debug)]
+pub struct UserSignalSender {
+    tx: mpsc::Sender<UserSignal>,
+}
+
+#[derive(Debug)]
+pub struct UserSignalReceiver {
+    rx: mpsc::Receiver<UserSignal>,
+}
+
+impl UserSignalSender {
+    pub fn send(&self, sig: UserSignal) -> Result<(), UserSignalSendError> {
+        self.tx
+            .send(sig)
+            .map_err(|mpsc::SendError(sig)| UserSignalSendError(sig))?;
+        set_triggered();
+        Ok(())
+    }
+}
+
+pub struct UserSignalSendError(pub UserSignal);
+
+impl fmt::Debug for UserSignalSendError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_struct("UserSignalSendError")
+            .finish_non_exhaustive()
+    }
+}
+
+impl fmt::Display for UserSignalSendError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str("sending a signal to a exited vm")
+    }
+}
+
+pub fn user_signal_channel() -> (UserSignalSender, UserSignalReceiver) {
+    let (tx, rx) = mpsc::channel();
+    (UserSignalSender { tx }, UserSignalReceiver { rx })
+}

vm/src/stdlib/signal.rs

@@ -1,34 +1,9 @@
 use crate::{PyObjectRef, VirtualMachine};
 
 pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
-    use crate::signal::NSIG;
-    use _signal::{SIG_DFL, SIG_ERR, SIG_IGN};
-
     let module = _signal::make_module(vm);
 
-    let sig_dfl = vm.new_pyobj(SIG_DFL as u8);
-    let sig_ign = vm.new_pyobj(SIG_IGN as u8);
-
-    for signum in 1..NSIG {
-        let handler = unsafe { libc::signal(signum as i32, SIG_IGN) };
-        if handler != SIG_ERR {
-            unsafe { libc::signal(signum as i32, handler) };
-        }
-        let py_handler = if handler == SIG_DFL {
-            Some(sig_dfl.clone())
-        } else if handler == SIG_IGN {
-            Some(sig_ign.clone())
-        } else {
-            None
-        };
-        vm.signal_handlers.as_deref().unwrap().borrow_mut()[signum] = py_handler;
-    }
-
-    let int_handler = module
-        .clone()
-        .get_attr("default_int_handler", vm)
-        .expect("_signal does not have this attr?");
-    _signal::signal(libc::SIGINT, int_handler, vm).expect("Failed to set sigint handler");
+    _signal::init_signal_handlers(&module, vm);
 
     module
 }
@@ -61,35 +36,35 @@ pub(crate) mod _signal {
     use nix::unistd::alarm as sig_alarm;
 
     #[cfg(not(windows))]
-    pub use libc::SIG_ERR;
+    use libc::SIG_ERR;
 
     #[cfg(not(windows))]
     #[pyattr]
-    pub use libc::{SIG_DFL, SIG_IGN};
+    use libc::{SIG_DFL, SIG_IGN};
 
     #[cfg(windows)]
     #[pyattr]
-    pub const SIG_DFL: libc::sighandler_t = 0;
+    const SIG_DFL: libc::sighandler_t = 0;
     #[cfg(windows)]
     #[pyattr]
-    pub const SIG_IGN: libc::sighandler_t = 1;
+    const SIG_IGN: libc::sighandler_t = 1;
     #[cfg(windows)]
-    pub const SIG_ERR: libc::sighandler_t = !0;
+    const SIG_ERR: libc::sighandler_t = !0;
 
     #[cfg(all(unix, not(target_os = "redox")))]
     extern "C" {
         fn siginterrupt(sig: i32, flag: i32) -> i32;
     }
 
     #[pyattr]
-    pub use crate::signal::NSIG;
+    use crate::signal::NSIG;
 
     #[pyattr]
-    pub use libc::{SIGABRT, SIGFPE, SIGILL, SIGINT, SIGSEGV, SIGTERM};
+    use libc::{SIGABRT, SIGFPE, SIGILL, SIGINT, SIGSEGV, SIGTERM};
 
     #[cfg(unix)]
     #[pyattr]
-    pub use libc::{
+    use libc::{
         SIGALRM, SIGBUS, SIGCHLD, SIGCONT, SIGHUP, SIGIO, SIGKILL, SIGPIPE, SIGPROF, SIGQUIT,
         SIGSTOP, SIGSYS, SIGTRAP, SIGTSTP, SIGTTIN, SIGTTOU, SIGURG, SIGUSR1, SIGUSR2, SIGVTALRM,
         SIGWINCH, SIGXCPU, SIGXFSZ,
@@ -103,10 +78,36 @@ pub(crate) mod _signal {
         target_os = "netbsd"
     )))]
     #[pyattr]
-    pub use libc::{SIGPWR, SIGSTKFLT};
+    use libc::{SIGPWR, SIGSTKFLT};
+
+    pub(super) fn init_signal_handlers(module: &PyObjectRef, vm: &VirtualMachine) {
+        let sig_dfl = vm.new_pyobj(SIG_DFL as u8);
+        let sig_ign = vm.new_pyobj(SIG_IGN as u8);
+
+        for signum in 1..NSIG {
+            let handler = unsafe { libc::signal(signum as i32, SIG_IGN) };
+            if handler != SIG_ERR {
+                unsafe { libc::signal(signum as i32, handler) };
+            }
+            let py_handler = if handler == SIG_DFL {
+                Some(sig_dfl.clone())
+            } else if handler == SIG_IGN {
+                Some(sig_ign.clone())
+            } else {
+                None
+            };
+            vm.signal_handlers.as_deref().unwrap().borrow_mut()[signum] = py_handler;
+        }
+
+        let int_handler = module
+            .clone()
+            .get_attr("default_int_handler", vm)
+            .expect("_signal does not have this attr?");
+        signal(libc::SIGINT, int_handler, vm).expect("Failed to set sigint handler");
+    }
 
     #[pyfunction]
-    pub(super) fn signal(
+    pub fn signal(
         signalnum: i32,
         handler: PyObjectRef,
         vm: &VirtualMachine,

vm/src/vm.rs

@@ -24,8 +24,7 @@ use crate::{
     import,
     protocol::{PyIterIter, PyIterReturn, PyMapping},
     scope::Scope,
-    signal::NSIG,
-    stdlib,
+    signal, stdlib,
     types::PyComparisonOp,
     IdProtocol, ItemProtocol, PyArithmeticValue, PyContext, PyGenericObject, PyLease, PyMethod,
     PyObject, PyObjectRef, PyObjectWrap, PyRef, PyRefExact, PyResult, PyValue, TryFromObject,
@@ -58,7 +57,8 @@ pub struct VirtualMachine {
     pub trace_func: RefCell<PyObjectRef>,
     pub use_tracing: Cell<bool>,
     pub recursion_limit: Cell<usize>,
-    pub signal_handlers: Option<Box<RefCell<[Option<PyObjectRef>; NSIG]>>>,
+    pub(crate) signal_handlers: Option<Box<RefCell<[Option<PyObjectRef>; signal::NSIG]>>>,
+    pub(crate) signal_rx: Option<signal::UserSignalReceiver>,
     pub repr_guards: RefCell<HashSet<usize>>,
     pub state: PyRc<PyGlobalState>,
     pub initialized: bool,
@@ -262,7 +262,11 @@ impl VirtualMachine {
         let trace_func = RefCell::new(ctx.none());
         // hack to get around const array repeat expressions, rust issue #79270
         const NONE: Option<PyObjectRef> = None;
-        let signal_handlers = RefCell::new([NONE; NSIG]);
+        // putting it in a const optimizes better, prevents linear initialization of the array
+        #[allow(clippy::declare_interior_mutable_const)]
+        const SIGNAL_HANDLERS: RefCell<[Option<PyObjectRef>; signal::NSIG]> =
+            RefCell::new([NONE; signal::NSIG]);
+        let signal_handlers = Some(Box::new(SIGNAL_HANDLERS));
 
         let module_inits = stdlib::get_module_inits();
 
@@ -285,7 +289,8 @@ impl VirtualMachine {
             trace_func,
             use_tracing: Cell::new(false),
             recursion_limit: Cell::new(if cfg!(debug_assertions) { 256 } else { 1000 }),
-            signal_handlers: Some(Box::new(signal_handlers)),
+            signal_handlers,
+            signal_rx: None,
             repr_guards: RefCell::default(),
             state: PyRc::new(PyGlobalState {
                 settings,
@@ -407,6 +412,11 @@ impl VirtualMachine {
         state.frozen.extend(frozen);
     }
 
+    /// Set the custom signal channel for the interpreter
+    pub fn set_user_signal_channel(&mut self, signal_rx: signal::UserSignalReceiver) {
+        self.signal_rx = Some(signal_rx);
+    }
+
     /// Start a new thread with access to the same interpreter.
     ///
     /// # Note
@@ -463,6 +473,7 @@ impl VirtualMachine {
             use_tracing: Cell::new(false),
             recursion_limit: self.recursion_limit.clone(),
             signal_handlers: None,
+            signal_rx: None,
             repr_guards: RefCell::default(),
             state: self.state.clone(),
             initialized: self.initialized,