From b336411516d70dfb2d74cafda556ad0d58fe0c0d Mon Sep 17 00:00:00 2001 From: mkravchik Date: Tue, 5 Dec 2023 22:03:00 +0200 Subject: [PATCH] Adding support for shutdown upon Ctrl+C on Windows for LLMP (#1704) * Adding support for shutdown upon Ctrl+C on Windows for LLMP * PR comments and clippy suggestions addressed * Enable CI for PR branches and manually triggered CI * Removed an empty line that broke compilation on some platforms * Trying to fix nostd compilation * Trying to fix nostd compilation for nightly toolchain * Removing use that is unused on some platforms * Trying to fix build on the nightly toolchain * Trying to fix build on the nightly toolchain, take 2 * Unifying LlmpShutdownSignalHandler * Fmt fix * Making the handler pub(crate) * Nightly toolchain fmt fixes --------- Co-authored-by: Dongjia "toka" Zhang --- .github/workflows/build_and_test.yml | 4 +- libafl_bolts/Cargo.toml | 2 +- libafl_bolts/src/llmp.rs | 60 ++++++++++++++----- libafl_bolts/src/os/windows_exceptions.rs | 72 ++++++++++++++++++++++- 4 files changed, 119 insertions(+), 19 deletions(-) diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 4c510f90ab..db6b576a37 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -2,10 +2,10 @@ name: build and test on: push: - branches: [ main ] + branches: [ main, 'pr/**' ] pull_request: branches: [ main ] - + workflow_dispatch: env: CARGO_TERM_COLOR: always diff --git a/libafl_bolts/Cargo.toml b/libafl_bolts/Cargo.toml index da641da1a9..31143b794a 100644 --- a/libafl_bolts/Cargo.toml +++ b/libafl_bolts/Cargo.toml @@ -121,7 +121,7 @@ libc = "0.2" # For (*nix) libc uds = { version = "0.4", optional = true, default-features = false } [target.'cfg(windows)'.dependencies] -windows = { version = "0.51.1", features = ["Win32_Foundation", "Win32_System_Threading", "Win32_System_Diagnostics_Debug", "Win32_System_Kernel", "Win32_System_Memory", "Win32_Security", "Win32_System_SystemInformation"] } +windows = { version = "0.51.1", features = ["Win32_Foundation", "Win32_System_Threading", "Win32_System_Diagnostics_Debug", "Win32_System_Kernel", "Win32_System_Memory", "Win32_Security", "Win32_System_SystemInformation", "Win32_System_Console"] } [target.'cfg(windows)'.build-dependencies] windows = "0.51.1" diff --git a/libafl_bolts/src/llmp.rs b/libafl_bolts/src/llmp.rs index 1524e61dc4..a6ad46cd16 100644 --- a/libafl_bolts/src/llmp.rs +++ b/libafl_bolts/src/llmp.rs @@ -100,6 +100,8 @@ use crate::current_time; use crate::os::unix_signals::setup_signal_handler; #[cfg(unix)] use crate::os::unix_signals::{siginfo_t, ucontext_t, Handler, Signal}; +#[cfg(all(windows, feature = "std"))] +use crate::os::windows_exceptions::{setup_ctrl_handler, CtrlHandler}; use crate::{ shmem::{ShMem, ShMemDescription, ShMemId, ShMemProvider}, ClientId, Error, @@ -175,7 +177,7 @@ const EOP_MSG_SIZE: usize = const LLMP_PAGE_HEADER_LEN: usize = size_of::(); /// The llmp broker registers a signal handler for cleanups on `SIGINT`. -#[cfg(unix)] +#[cfg(any(unix, all(windows, feature = "std")))] static mut LLMP_SIGHANDLER_STATE: LlmpShutdownSignalHandler = LlmpShutdownSignalHandler { shutting_down: false, }; @@ -1962,7 +1964,9 @@ where } /// A signal handler for the [`LlmpBroker`]. -#[cfg(unix)] +/// On unix, it handles signals +/// On Windows - control signals (e.g., CTRL+C) +#[cfg(any(unix, all(windows, feature = "std")))] #[derive(Debug, Clone)] pub struct LlmpShutdownSignalHandler { shutting_down: bool, @@ -1986,6 +1990,18 @@ impl Handler for LlmpShutdownSignalHandler { } } +#[cfg(all(windows, feature = "std"))] +impl CtrlHandler for LlmpShutdownSignalHandler { + #[allow(clippy::not_unsafe_ptr_arg_deref)] + fn handle(&mut self, ctrl_type: u32) -> bool { + log::info!("LLMP: Received shutdown signal, ctrl_type {:?}", ctrl_type); + unsafe { + ptr::write_volatile(&mut self.shutting_down, true); + } + true + } +} + /// The broker forwards all messages to its own bus-like broadcast map. /// It may intercept messages passing through. impl LlmpBroker @@ -2242,15 +2258,37 @@ where /// Internal function, returns true when shuttdown is requested by a `SIGINT` signal #[inline] - #[cfg(unix)] + #[cfg(any(unix, all(windows, feature = "std")))] #[allow(clippy::unused_self)] fn is_shutting_down(&self) -> bool { unsafe { ptr::read_volatile(ptr::addr_of!(LLMP_SIGHANDLER_STATE.shutting_down)) } } + #[cfg(any(all(unix, not(miri)), all(windows, feature = "std")))] + fn setup_handlers() { + #[cfg(all(unix, not(miri)))] + if let Err(e) = unsafe { setup_signal_handler(&mut LLMP_SIGHANDLER_STATE) } { + // We can live without a proper ctrl+c signal handler - Ignore. + log::info!("Failed to setup signal handlers: {e}"); + } else { + log::info!("Successfully setup signal handlers"); + } + + #[cfg(all(windows, feature = "std"))] + if let Err(e) = unsafe { setup_ctrl_handler(&mut LLMP_SIGHANDLER_STATE) } { + // We can live without a proper ctrl+c signal handler - Ignore. + log::info!("Failed to setup control handlers: {e}"); + } else { + log::info!( + "{}: Broker successfully setup control handlers", + std::process::id().to_string() + ); + } + } + /// Always returns true on platforms, where no shutdown signal handlers are supported #[inline] - #[cfg(not(unix))] + #[cfg(not(any(unix, all(windows, feature = "std"))))] #[allow(clippy::unused_self)] fn is_shutting_down(&self) -> bool { false @@ -2280,11 +2318,8 @@ where { use super::current_milliseconds; - #[cfg(all(unix, not(miri)))] - if let Err(_e) = unsafe { setup_signal_handler(&mut LLMP_SIGHANDLER_STATE) } { - // We can live without a proper ctrl+c signal handler. Print and ignore. - log::info!("Failed to setup signal handlers: {_e}"); - } + #[cfg(any(all(unix, not(miri)), all(windows, feature = "std")))] + Self::setup_handlers(); let timeout = timeout.as_millis() as u64; let mut end_time = current_milliseconds() + timeout; @@ -2344,11 +2379,8 @@ where where F: FnMut(ClientId, Tag, Flags, &[u8]) -> Result, { - #[cfg(all(unix, not(miri)))] - if let Err(_e) = unsafe { setup_signal_handler(&mut LLMP_SIGHANDLER_STATE) } { - // We can live without a proper ctrl+c signal handler. Print and ignore. - log::info!("Failed to setup signal handlers: {_e}"); - } + #[cfg(any(all(unix, not(miri)), all(windows, feature = "std")))] + Self::setup_handlers(); while !self.is_shutting_down() { self.once(on_new_msg) diff --git a/libafl_bolts/src/os/windows_exceptions.rs b/libafl_bolts/src/os/windows_exceptions.rs index 900d5b9a07..ffc092de38 100644 --- a/libafl_bolts/src/os/windows_exceptions.rs +++ b/libafl_bolts/src/os/windows_exceptions.rs @@ -11,10 +11,12 @@ use core::{ }; use std::os::raw::{c_long, c_void}; +use log::info; use num_enum::TryFromPrimitive; pub use windows::Win32::{ - Foundation::NTSTATUS, + Foundation::{BOOL, NTSTATUS}, System::{ + Console::{SetConsoleCtrlHandler, CTRL_BREAK_EVENT, CTRL_C_EVENT, PHANDLER_ROUTINE}, Diagnostics::Debug::{ AddVectoredExceptionHandler, UnhandledExceptionFilter, EXCEPTION_POINTERS, }, @@ -321,11 +323,21 @@ unsafe fn internal_handle_exception( .unwrap(); match &EXCEPTION_HANDLERS[index] { Some(handler_holder) => { + info!( + "{:?}: Handling exception {}", + std::process::id(), + exception_code + ); let handler = &mut **handler_holder.handler.get(); handler.handle(exception_code, exception_pointers); EXCEPTION_CONTINUE_EXECUTION } None => { + info!( + "{:?}: No handler for exception {}", + std::process::id(), + exception_code + ); // Go to Default one let handler_holder = &EXCEPTION_HANDLERS[EXCEPTION_HANDLERS_SIZE - 1] .as_ref() @@ -351,7 +363,7 @@ pub unsafe extern "system" fn handle_exception( .unwrap() .ExceptionCode; let exception_code = ExceptionCode::try_from(code.0).unwrap(); - // log::info!("Received exception; code: {}", exception_code); + log::info!("Received exception; code: {}", exception_code); internal_handle_exception(exception_code, exception_pointers) } @@ -406,3 +418,59 @@ pub unsafe fn setup_exception_handler(handler: &mut T) -> ); Ok(()) } + +#[cfg(feature = "alloc")] +pub(crate) trait CtrlHandler { + /// Handle an exception + fn handle(&mut self, ctrl_type: u32) -> bool; +} + +struct CtrlHandlerHolder { + handler: UnsafeCell<*mut dyn CtrlHandler>, +} + +/// Keep track of which handler is registered for which exception +static mut CTRL_HANDLER: Option = None; + +/// Set `ConsoleCtrlHandler` to catch Ctrl-C +/// # Safety +/// Same safety considerations as in `setup_exception_handler` +pub(crate) unsafe fn setup_ctrl_handler( + handler: &mut T, +) -> Result<(), Error> { + write_volatile( + &mut CTRL_HANDLER, + Some(CtrlHandlerHolder { + handler: UnsafeCell::new(handler as *mut dyn CtrlHandler), + }), + ); + compiler_fence(Ordering::SeqCst); + + // Log the result of SetConsoleCtrlHandler + let result = SetConsoleCtrlHandler(Some(ctrl_handler), true); + match result { + Ok(()) => { + info!("SetConsoleCtrlHandler succeeded"); + Ok(()) + } + Err(err) => { + info!("SetConsoleCtrlHandler failed"); + Err(Error::from(err)) + } + } +} + +unsafe extern "system" fn ctrl_handler(ctrl_type: u32) -> BOOL { + match &CTRL_HANDLER { + Some(handler_holder) => { + info!("{:?}: Handling ctrl {}", std::process::id(), ctrl_type); + let handler = &mut *handler_holder.handler.get(); + if let Some(ctrl_handler) = handler.as_mut() { + (*ctrl_handler).handle(ctrl_type).into() + } else { + false.into() + } + } + None => false.into(), + } +}