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 <tokazerkje@outlook.com>

.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
 

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"

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::<LlmpPage>();
 
 /// 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<SP> LlmpBroker<SP>
@@ -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<LlmpMsgHookResult, Error>,
     {
-        #[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)

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<T: 'static + 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<CtrlHandlerHolder> = None;
+
+/// Set `ConsoleCtrlHandler` to catch Ctrl-C
+/// # Safety
+/// Same safety considerations as in `setup_exception_handler`
+pub(crate) unsafe fn setup_ctrl_handler<T: 'static + CtrlHandler>(
+    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(),
+    }
+}