diff --git a/Cargo.toml b/Cargo.toml index f9678a4..86976ad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,11 @@ harness = false name = "issue_97" path = "tests/main/issue_97.rs" +[[test]] +harness = false +name = "deinit" +path = "tests/main/deinit.rs" + [dev-dependencies] signal-hook = "0.3" diff --git a/src/block_outcome.rs b/src/block_outcome.rs new file mode 100644 index 0000000..bf7dbc5 --- /dev/null +++ b/src/block_outcome.rs @@ -0,0 +1,5 @@ +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum BlockOutcome { + Awaited, + HandlerRemoved, +} diff --git a/src/error.rs b/src/error.rs index f13625a..64c5d14 100644 --- a/src/error.rs +++ b/src/error.rs @@ -10,6 +10,8 @@ pub enum Error { MultipleHandlers, /// Unexpected system error. System(std::io::Error), + /// Handler was removed + HandlerRemoved, } impl Error { @@ -18,6 +20,7 @@ impl Error { Error::NoSuchSignal(_) => "Signal could not be found from the system", Error::MultipleHandlers => "Ctrl-C signal handler already registered", Error::System(_) => "Unexpected system error", + Error::HandlerRemoved => "Handler was removed", } } } diff --git a/src/lib.rs b/src/lib.rs index 7bdbd30..d5673f6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -50,14 +50,16 @@ mod error; mod platform; +use block_outcome::BlockOutcome; pub use platform::Signal; mod signal; pub use signal::*; +mod block_outcome; pub use error::Error; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Mutex; -use std::thread; +use std::thread::{self, JoinHandle}; static INIT: AtomicBool = AtomicBool::new(false); static INIT_LOCK: Mutex<()> = Mutex::new(()); @@ -90,7 +92,7 @@ static INIT_LOCK: Mutex<()> = Mutex::new(()); /// /// # Panics /// Any panic in the handler will not be caught and will cause the signal handler thread to stop. -pub fn set_handler(user_handler: F) -> Result<(), Error> +pub fn set_handler(user_handler: F) -> Result, Error> where F: FnMut() + 'static + Send, { @@ -102,14 +104,14 @@ where /// # Errors /// Will return an error if another handler exists or if a system error occurred while setting the /// handler. -pub fn try_set_handler(user_handler: F) -> Result<(), Error> +pub fn try_set_handler(user_handler: F) -> Result, Error> where F: FnMut() + 'static + Send, { init_and_set_handler(user_handler, false) } -fn init_and_set_handler(user_handler: F, overwrite: bool) -> Result<(), Error> +fn init_and_set_handler(user_handler: F, overwrite: bool) -> Result, Error> where F: FnMut() + 'static + Send, { @@ -117,16 +119,16 @@ where let _guard = INIT_LOCK.lock().unwrap(); if !INIT.load(Ordering::Relaxed) { - set_handler_inner(user_handler, overwrite)?; + let result = set_handler_inner(user_handler, overwrite)?; INIT.store(true, Ordering::Release); - return Ok(()); + return Ok(result); } } Err(Error::MultipleHandlers) } -fn set_handler_inner(mut user_handler: F, overwrite: bool) -> Result<(), Error> +fn set_handler_inner(mut user_handler: F, overwrite: bool) -> Result, Error> where F: FnMut() + 'static + Send, { @@ -138,11 +140,123 @@ where .name("ctrl-c".into()) .spawn(move || loop { unsafe { - platform::block_ctrl_c().expect("Critical system error while waiting for Ctrl-C"); + match platform::block_ctrl_c() { + Ok(BlockOutcome::Awaited) => {}, + Ok(BlockOutcome::HandlerRemoved) => break, + Err(err) => panic!("Critical system error while waiting for Ctrl-C: {err:?}") + }; } user_handler(); }) + .map_err(Error::System) +} + + +/// Same as [`ctrlc::set_handler`], but uses [`std::ops::FnOnce`] as a handler that only handles one interrupt. +/// +/// Register signal handler for Ctrl-C. +/// +/// Starts a new dedicated signal handling thread. Should only be called at the start of the program, or after +/// last `_once`-handler already fired (for example, via `.join()`). +/// +/// # Example +/// ```no_run +/// +/// # use ctrlc::*; +/// +/// let fires = 0; +/// let handle = ctrlc::set_handler_once(move || fires + 1).unwrap(); +/// +/// // interrupt_and_wait(); // platform-dependant +/// +/// // First unwrap for thread join, second one to `Option` because handler can be removed without firing. +/// let fires = handle.join().unwrap().unwrap(); +/// assert_eq!(fires, 1); +/// +/// assert!(ctrlc::remove_all_handlers().is_err()); // This handler should be already removed after firing once. +/// ``` +pub fn set_handler_once(user_handler: F) -> Result>, Error> +where + F: FnOnce() -> T + 'static + Send, + T: 'static + Send +{ + init_and_set_handler_once(user_handler, true) +} + +/// The same as [`ctrlc::try_set_handler`] but uses [`std::ops::FnOnce`] as a handler that only handles one interrupt. +/// The same as [`ctrlc::set_handler_once`] but errors if a handler already exists for the signal(s). +/// +/// # Errors +/// Will return an error if another handler exists or if a system error occurred while setting the +/// handler. +pub fn try_set_handler_once(user_handler: F) -> Result>, Error> +where + F: FnOnce() -> T + 'static + Send, + T: 'static + Send +{ + init_and_set_handler_once(user_handler, false) +} + + +fn init_and_set_handler_once(user_handler: F, overwrite: bool) -> Result>, Error> +where + F: FnOnce() -> T + 'static + Send, + T: 'static + Send +{ + if !INIT.load(Ordering::Acquire) { + let _guard = INIT_LOCK.lock().unwrap(); + + if !INIT.load(Ordering::Relaxed) { + let handle = set_handler_inner_once(user_handler, overwrite)?; + INIT.store(true, Ordering::Release); + return Ok(handle); + } + } + + Err(Error::MultipleHandlers) +} + + +fn set_handler_inner_once(user_handler: F, overwrite: bool) -> Result>, Error> +where + F: FnOnce() -> T + 'static + Send, + T: 'static + Send, +{ + unsafe { + platform::init_os_handler(overwrite)?; + } + + let thread = thread::Builder::new() + .name("ctrl-c".into()) + .spawn(move || { + let outcome = unsafe { + platform::block_ctrl_c().expect("Critical system error while waiting for Ctrl-C") + }; + if outcome == BlockOutcome::HandlerRemoved { + return None; + } + let result = user_handler(); + + match remove_all_handlers() { + Ok(()) | + Err(Error::HandlerRemoved) => {}, + _ => eprintln!("[ctrlc] System error after waiting for Ctrl-C"), + }; + Some(result) + }) .map_err(Error::System)?; + Ok(thread) +} + +/// Removes all previously added handlers +pub fn remove_all_handlers() -> Result<(), Error> { + if !INIT.load(Ordering::Acquire) { + return Err(Error::HandlerRemoved); + } + unsafe { + platform::deinit_os_handler()?; + INIT.store(false, Ordering::Relaxed); + } Ok(()) } diff --git a/src/platform/unix/mod.rs b/src/platform/unix/mod.rs index 23a6015..5ae56d0 100644 --- a/src/platform/unix/mod.rs +++ b/src/platform/unix/mod.rs @@ -7,7 +7,10 @@ // notice may not be copied, modified, or distributed except // according to those terms. +use crate::block_outcome::BlockOutcome; use crate::error::Error as CtrlcError; +use nix::sys::signal::SigAction; +use nix::sys::signal::SigHandler; use nix::unistd; use std::os::fd::BorrowedFd; use std::os::fd::IntoRawFd; @@ -79,6 +82,14 @@ fn pipe2(flags: nix::fcntl::OFlag) -> nix::Result<(RawFd, RawFd)> { Ok((pipe.0.into_raw_fd(), pipe.1.into_raw_fd())) } +unsafe fn close_pipe() { + // Try to close the pipes. close() should not fail, + // but if it does, there isn't much we can do + let _ = unistd::close(PIPE.1); + let _ = unistd::close(PIPE.0); + PIPE = (-1, -1); +} + /// Register os signal handler. /// /// Must be called before calling [`block_ctrl_c()`](fn.block_ctrl_c.html) @@ -91,41 +102,29 @@ fn pipe2(flags: nix::fcntl::OFlag) -> nix::Result<(RawFd, RawFd)> { pub unsafe fn init_os_handler(overwrite: bool) -> Result<(), Error> { use nix::fcntl; use nix::sys::signal; - + PIPE = pipe2(fcntl::OFlag::O_CLOEXEC)?; - let close_pipe = |e: nix::Error| -> Error { - // Try to close the pipes. close() should not fail, - // but if it does, there isn't much we can do - let _ = unistd::close(PIPE.1); - let _ = unistd::close(PIPE.0); - e - }; - // Make sure we never block on write in the os handler. if let Err(e) = fcntl::fcntl(PIPE.1, fcntl::FcntlArg::F_SETFL(fcntl::OFlag::O_NONBLOCK)) { - return Err(close_pipe(e)); + close_pipe(); + return Err(e); } let handler = signal::SigHandler::Handler(os_handler); - #[cfg(not(target_os = "nto"))] - let new_action = signal::SigAction::new( - handler, - signal::SaFlags::SA_RESTART, - signal::SigSet::empty(), - ); - // SA_RESTART is not supported on QNX Neutrino 7.1 and before - #[cfg(target_os = "nto")] - let new_action = - signal::SigAction::new(handler, signal::SaFlags::empty(), signal::SigSet::empty()); + let new_action = sig_handler_to_sig_action(handler); let sigint_old = match signal::sigaction(signal::Signal::SIGINT, &new_action) { Ok(old) => old, - Err(e) => return Err(close_pipe(e)), + Err(e) => { + close_pipe(); + return Err(e) + } }; if !overwrite && sigint_old.handler() != signal::SigHandler::SigDfl { signal::sigaction(signal::Signal::SIGINT, &sigint_old).unwrap(); - return Err(close_pipe(nix::Error::EEXIST)); + close_pipe(); + return Err(nix::Error::EEXIST); } #[cfg(feature = "termination")] @@ -134,33 +133,80 @@ pub unsafe fn init_os_handler(overwrite: bool) -> Result<(), Error> { Ok(old) => old, Err(e) => { signal::sigaction(signal::Signal::SIGINT, &sigint_old).unwrap(); - return Err(close_pipe(e)); + close_pipe(); + return Err(e); } }; if !overwrite && sigterm_old.handler() != signal::SigHandler::SigDfl { signal::sigaction(signal::Signal::SIGINT, &sigint_old).unwrap(); signal::sigaction(signal::Signal::SIGTERM, &sigterm_old).unwrap(); - return Err(close_pipe(nix::Error::EEXIST)); + close_pipe(); + return Err(nix::Error::EEXIST); } let sighup_old = match signal::sigaction(signal::Signal::SIGHUP, &new_action) { Ok(old) => old, Err(e) => { signal::sigaction(signal::Signal::SIGINT, &sigint_old).unwrap(); signal::sigaction(signal::Signal::SIGTERM, &sigterm_old).unwrap(); - return Err(close_pipe(e)); + close_pipe(); + return Err(e); } }; if !overwrite && sighup_old.handler() != signal::SigHandler::SigDfl { signal::sigaction(signal::Signal::SIGINT, &sigint_old).unwrap(); signal::sigaction(signal::Signal::SIGTERM, &sigterm_old).unwrap(); signal::sigaction(signal::Signal::SIGHUP, &sighup_old).unwrap(); - return Err(close_pipe(nix::Error::EEXIST)); + close_pipe(); + return Err(nix::Error::EEXIST); } } Ok(()) } +#[allow(dead_code)] +pub unsafe fn deinit_os_handler() -> Result<(), Error> { + use nix::sys::signal; + if !is_handler_init() { + return Err(nix::Error::ENOENT); + } + + let new_action = sig_handler_to_sig_action(signal::SigHandler::SigDfl); + + let _ = signal::sigaction(signal::Signal::SIGINT, &new_action); + + #[cfg(feature = "termination")] + { + let _ = signal::sigaction(signal::Signal::SIGTERM, &new_action); + let _ = signal::sigaction(signal::Signal::SIGHUP, &new_action); + } + close_pipe(); + + Ok(()) +} + +#[allow(dead_code)] +pub unsafe fn is_handler_init() -> bool { + return PIPE.0 != -1 && PIPE.1 != -1; +} + +unsafe fn sig_handler_to_sig_action(handler: SigHandler) -> SigAction { + use nix::sys::signal; + + #[cfg(not(target_os = "nto"))] + let action = signal::SigAction::new( + handler, + signal::SaFlags::SA_RESTART, + signal::SigSet::empty(), + ); + + // SA_RESTART is not supported on QNX Neutrino 7.1 and before + #[cfg(target_os = "nto")] + let action = signal::SigAction::new(handler, signal::SaFlags::empty(), signal::SigSet::empty()); + + action +} + /// Blocks until a Ctrl-C signal is received. /// /// Must be called after calling [`init_os_handler()`](fn.init_os_handler.html). @@ -169,21 +215,25 @@ pub unsafe fn init_os_handler(overwrite: bool) -> Result<(), Error> { /// Will return an error if a system error occurred. /// #[inline] -pub unsafe fn block_ctrl_c() -> Result<(), CtrlcError> { - use std::io; +pub unsafe fn block_ctrl_c() -> Result { let mut buf = [0u8]; // TODO: Can we safely convert the pipe fd into a std::io::Read // with std::os::unix::io::FromRawFd, this would handle EINTR // and everything for us. loop { - match unistd::read(PIPE.0, &mut buf[..]) { + let pipe = std::ptr::read_volatile(&raw const PIPE); + match unistd::read(pipe.0, &mut buf[..]) { Ok(1) => break, - Ok(_) => return Err(CtrlcError::System(io::ErrorKind::UnexpectedEof.into())), + + Ok(_) | + Err(nix::errno::Errno::EBADF) + => return Ok(BlockOutcome::HandlerRemoved), + Err(nix::errno::Errno::EINTR) => {} Err(e) => return Err(e.into()), } } - Ok(()) + Ok(BlockOutcome::Awaited) } diff --git a/src/platform/windows/mod.rs b/src/platform/windows/mod.rs index 187b156..19df855 100644 --- a/src/platform/windows/mod.rs +++ b/src/platform/windows/mod.rs @@ -8,55 +8,116 @@ // according to those terms. use std::io; +use std::io::ErrorKind; use std::ptr; + use windows_sys::Win32::Foundation::{CloseHandle, BOOL, HANDLE, WAIT_FAILED, WAIT_OBJECT_0}; use windows_sys::Win32::System::Console::SetConsoleCtrlHandler; use windows_sys::Win32::System::Threading::{ CreateSemaphoreA, ReleaseSemaphore, WaitForSingleObject, INFINITE, }; +use crate::block_outcome::BlockOutcome; + /// Platform specific error type pub type Error = io::Error; /// Platform specific signal type pub type Signal = u32; -const MAX_SEM_COUNT: i32 = 255; -static mut SEMAPHORE: HANDLE = 0 as HANDLE; +const MAX_SEM_COUNT: i32 = 65535; const TRUE: BOOL = 1; const FALSE: BOOL = 0; +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct OsHandler { + semaphore: HANDLE, +} + +static mut HANDLER: Option = None; + unsafe extern "system" fn os_handler(_: u32) -> BOOL { - // Assuming this always succeeds. Can't really handle errors in any meaningful way. - ReleaseSemaphore(SEMAPHORE, 1, ptr::null_mut()); - TRUE + if let Some(handler) = HANDLER { + // Assuming this always succeeds. Can't really handle errors in any meaningful way. + ReleaseSemaphore(handler.semaphore, 1, ptr::null_mut()); + TRUE + } else { + // We have no handler set. Not sure how the hell this function was even called then. + // But okay, just mark this as not handled (FALSE). + FALSE + } } -/// Register os signal handler. +/// Register OS signal handler. /// /// Must be called before calling [`block_ctrl_c()`](fn.block_ctrl_c.html) /// and should only be called once. /// /// # Errors /// Will return an error if a system error occurred. -/// #[inline] -pub unsafe fn init_os_handler(_overwrite: bool) -> Result<(), Error> { - SEMAPHORE = CreateSemaphoreA(ptr::null_mut(), 0, MAX_SEM_COUNT, ptr::null()); - if SEMAPHORE.is_null() { +pub unsafe fn init_os_handler(overwrite: bool) -> Result<(), Error> { + if is_handler_init() { + if !overwrite { + return Err(ErrorKind::AlreadyExists.into()) + } else { + deinit_os_handler()?; + } + } + assert!(!is_handler_init()); + + let semaphore = CreateSemaphoreA(ptr::null_mut(), 0, MAX_SEM_COUNT, ptr::null()); + if semaphore.is_null() { return Err(io::Error::last_os_error()); } + + // Remove OUR handlers if those exist + // It does not make sense to have multiple of same(!) OS handlers added. + let mut handlers_removed = 0; + while SetConsoleCtrlHandler(Some(os_handler), FALSE) == TRUE { + handlers_removed += 1; + } + if handlers_removed > 0 { + // This does not interfere with our ability to add handlers, but it is unexpected for there + // to be more than one. + eprintln!( + "[ctrlc] Somehow {handlers_removed} other OS {} of `ctrlc` was added before. Probably a bug.", + if handlers_removed == 1 { "handler" } else { "handlers" } + ); + } + + // Set our custom handler if SetConsoleCtrlHandler(Some(os_handler), TRUE) == FALSE { let e = io::Error::last_os_error(); - CloseHandle(SEMAPHORE); - SEMAPHORE = 0 as HANDLE; + CloseHandle(semaphore); return Err(e); } + HANDLER = Some(OsHandler { semaphore }); + Ok(()) } +/// Unregisters OS signal handler set by [`ctrlc::platform::init_os_handler`]. +#[inline] +pub unsafe fn deinit_os_handler() -> Result<(), Error> { + if let Some(handler) = HANDLER { + HANDLER = None; + CloseHandle(handler.semaphore); + SetConsoleCtrlHandler(Some(os_handler), FALSE); // Remove the handler callback + Ok(()) + } else { + Err(ErrorKind::NotFound.into()) + } +} + +pub unsafe fn is_handler_init() -> bool { + #[allow(static_mut_refs)] + return HANDLER.is_some(); +} + + /// Blocks until a Ctrl-C signal is received. /// /// Must be called after calling [`init_os_handler()`](fn.init_os_handler.html). @@ -65,10 +126,18 @@ pub unsafe fn init_os_handler(_overwrite: bool) -> Result<(), Error> { /// Will return an error if a system error occurred. /// #[inline] -pub unsafe fn block_ctrl_c() -> Result<(), Error> { - match WaitForSingleObject(SEMAPHORE, INFINITE) { - WAIT_OBJECT_0 => Ok(()), - WAIT_FAILED => Err(io::Error::last_os_error()), +pub unsafe fn block_ctrl_c() -> Result { + let handler = HANDLER.ok_or::(ErrorKind::NotFound.into())?; + + match WaitForSingleObject(handler.semaphore, INFINITE) { + WAIT_OBJECT_0 => Ok(BlockOutcome::Awaited), + WAIT_FAILED => { + if Some(handler) != HANDLER { + Ok(BlockOutcome::HandlerRemoved) + } else { + Err(io::Error::last_os_error()) + } + }, ret => Err(io::Error::new( io::ErrorKind::Other, format!( diff --git a/tests/main/deinit.rs b/tests/main/deinit.rs new file mode 100644 index 0000000..e828b7e --- /dev/null +++ b/tests/main/deinit.rs @@ -0,0 +1,93 @@ +// Copyright (c) 2023 CtrlC developers +// Licensed under the Apache License, Version 2.0 +// or the MIT +// license , +// at your option. All files in the project carrying such +// notice may not be copied, modified, or distributed except +// according to those terms. + +#[macro_use] +mod harness; +use harness::{platform, run_harness}; + +use std::sync::{atomic::{AtomicUsize, Ordering}, Arc}; +use std::time::Duration; + +fn interrupt_and_wait() { + unsafe { platform::raise_ctrl_c(); } + std::thread::sleep(Duration::from_millis(10)); +} + +fn test_deinit() { + let total_fires = Arc::new(AtomicUsize::new(0)); + let fires = 0; + + // 1. First handler + let handle = ctrlc::set_handler_once({ + let total_fires = total_fires.clone(); + move || { + total_fires.fetch_add(1, Ordering::Relaxed); + fires + 1 + } + }).unwrap(); + interrupt_and_wait(); + + // First unwrap for thread join, second one to `Option` because handler can be removed without firing. + let fires = handle.join().unwrap().unwrap(); + assert_eq!(fires, 1); + assert_eq!(total_fires.load(Ordering::Relaxed), 1); + + assert!(ctrlc::remove_all_handlers().is_err()); // This handler should be already removed after firing once. + + // 2. Second handler + let handle = ctrlc::set_handler_once(|| 42).unwrap(); + interrupt_and_wait(); + + assert_eq!(handle.join().unwrap().unwrap(), 42); + assert_eq!(total_fires.load(Ordering::Relaxed), 1); + + assert!(ctrlc::remove_all_handlers().is_err()); // This handler should be already removed after firing once. + + // 3. Test with a non-once handler + ctrlc::set_handler({ + let total_fires = total_fires.clone(); + move || { + total_fires.fetch_add(1, Ordering::Relaxed); + } + }).unwrap(); + interrupt_and_wait(); + interrupt_and_wait(); + interrupt_and_wait(); + interrupt_and_wait(); + interrupt_and_wait(); + + assert_eq!(total_fires.load(Ordering::Relaxed), 6); + + ctrlc::remove_all_handlers().unwrap(); + + // 4. First handler again + let handle = ctrlc::set_handler_once({ + let total_fires = total_fires.clone(); + move || { + total_fires.fetch_add(1, Ordering::Relaxed); + fires + 1 + } + }).unwrap(); + interrupt_and_wait(); + + let fires = handle.join().unwrap().unwrap(); + + assert_eq!(fires, 2); + assert_eq!(total_fires.load(Ordering::Relaxed), 7); + + assert!(ctrlc::remove_all_handlers().is_err()); // This handler should be already removed after firing once. +} + +fn tests() { + run_tests!(test_deinit); +} + +fn main() { + run_harness(tests); +}