[Windows] Setup ASAN death callback (#908)

* step 1

* i forgot to change this

* add handler

* doc

* fmt

* move to libafl_targets

* fix

* windows

* clp

* fix

* clp

* cfg

* fix

* clp

Co-authored-by: Dominik Maier <domenukk@gmail.com>

libafl/src/bolts/os/windows_exceptions.rs

@@ -23,9 +23,14 @@ pub use windows::Win32::{
 
 use crate::Error;
 
-//const EXCEPTION_CONTINUE_EXECUTION: c_long = -1;
+// For VEH
+const EXCEPTION_CONTINUE_EXECUTION: c_long = -1;
+
+// For VEH
 //const EXCEPTION_CONTINUE_SEARCH: c_long = 0;
-const EXCEPTION_EXECUTE_HANDLER: c_long = 1;
+
+// For SEH
+//const EXCEPTION_EXECUTE_HANDLER: c_long = 1;
 
 // From https://github.com/Alexpux/mingw-w64/blob/master/mingw-w64-headers/crt/signal.h
 pub const SIGINT: i32 = 2;
@@ -314,9 +319,9 @@ unsafe fn internal_handle_exception(
         Some(handler_holder) => {
             let handler = &mut **handler_holder.handler.get();
             handler.handle(exception_code, exception_pointers);
-            EXCEPTION_EXECUTE_HANDLER
+            EXCEPTION_CONTINUE_EXECUTION
         }
-        None => EXCEPTION_EXECUTE_HANDLER,
+        None => EXCEPTION_CONTINUE_EXECUTION,
     }
 }
 
@@ -376,7 +381,7 @@ pub unsafe fn setup_exception_handler<T: 'static + Handler>(handler: &mut T) ->
     // SetUnhandledFilter does not work with frida since the stack is changed and exception handler is lost with Stalker enabled.
     // See https://github.com/AFLplusplus/LibAFL/pull/403
     AddVectoredExceptionHandler(
-        1,
+        0,
         Some(core::mem::transmute(handle_exception as *const c_void)),
     );
     Ok(())

libafl/src/executors/inprocess.rs

@@ -188,6 +188,7 @@ where
     {
         let handlers = InProcessHandlers::new::<Self, EM, OF, Z, H>()?;
         #[cfg(windows)]
+        // Some initialization necessary for windows.
         unsafe {
             /*
                 See https://github.com/AFLplusplus/LibAFL/pull/403
@@ -873,6 +874,149 @@ mod unix_signal_handler {
     }
 }
 
+/// Same as `inproc_crash_handler`, but this is called when address sanitizer exits, not from the exception handler
+#[cfg(all(windows, feature = "std"))]
+pub mod windows_asan_handler {
+    use alloc::string::String;
+    use core::{
+        ptr,
+        sync::atomic::{compiler_fence, Ordering},
+    };
+    #[cfg(feature = "std")]
+    use std::io::{stdout, Write};
+
+    use windows::Win32::System::Threading::{
+        EnterCriticalSection, LeaveCriticalSection, RTL_CRITICAL_SECTION,
+    };
+
+    use crate::{
+        corpus::{Corpus, Testcase},
+        events::{Event, EventFirer, EventRestarter},
+        executors::{inprocess::GLOBAL_STATE, Executor, ExitKind, HasObservers},
+        feedbacks::Feedback,
+        fuzzer::HasObjective,
+        inputs::UsesInput,
+        observers::ObserversTuple,
+        state::{HasClientPerfMonitor, HasMetadata, HasSolutions},
+    };
+
+    /// # Safety
+    /// ASAN deatch handler
+    pub unsafe extern "C" fn asan_death_handler<E, EM, OF, Z>()
+    where
+        E: Executor<EM, Z> + HasObservers,
+        EM: EventFirer<State = E::State> + EventRestarter<State = E::State>,
+        OF: Feedback<E::State>,
+        E::State: HasSolutions + HasClientPerfMonitor,
+        Z: HasObjective<Objective = OF, State = E::State>,
+    {
+        let mut data = &mut GLOBAL_STATE;
+        // Have we set a timer_before?
+        if !(data.tp_timer as *mut windows::Win32::System::Threading::TP_TIMER).is_null() {
+            /*
+                We want to prevent the timeout handler being run while the main thread is executing the crash handler
+                Timeout handler runs if it has access to the critical section or data.in_target == 0
+                Writing 0 to the data.in_target makes the timeout handler makes the timeout handler invalid.
+            */
+            compiler_fence(Ordering::SeqCst);
+            EnterCriticalSection(data.critical as *mut RTL_CRITICAL_SECTION);
+            compiler_fence(Ordering::SeqCst);
+            data.in_target = 0;
+            compiler_fence(Ordering::SeqCst);
+            LeaveCriticalSection(data.critical as *mut RTL_CRITICAL_SECTION);
+            compiler_fence(Ordering::SeqCst);
+        }
+
+        #[cfg(feature = "std")]
+        eprintln!("ASAN detected crash!");
+        if data.current_input_ptr.is_null() {
+            #[cfg(feature = "std")]
+            {
+                eprintln!("Double crash\n");
+                eprintln!(
+                "ASAN detected crash but we're not in the target... Bug in the fuzzer? Exiting.",
+                );
+            }
+            #[cfg(feature = "std")]
+            {
+                eprintln!("Type QUIT to restart the child");
+                let mut line = String::new();
+                while line.trim() != "QUIT" {
+                    std::io::stdin().read_line(&mut line).unwrap();
+                }
+            }
+
+            // TODO tell the parent to not restart
+        } else {
+            let executor = data.executor_mut::<E>();
+            // reset timer
+            if !data.tp_timer.is_null() {
+                executor.post_run_reset();
+                data.tp_timer = ptr::null_mut();
+            }
+
+            let state = data.state_mut::<E::State>();
+            let fuzzer = data.fuzzer_mut::<Z>();
+            let event_mgr = data.event_mgr_mut::<EM>();
+            let observers = executor.observers_mut();
+
+            #[cfg(feature = "std")]
+            eprintln!("Child crashed!");
+            #[cfg(feature = "std")]
+            drop(stdout().flush());
+
+            // Make sure we don't crash in the crash handler forever.
+            let input = data.take_current_input::<<E::State as UsesInput>::Input>();
+
+            #[cfg(feature = "std")]
+            eprintln!("Child crashed!");
+            #[cfg(feature = "std")]
+            drop(stdout().flush());
+
+            observers
+                .post_exec_all(state, input, &ExitKind::Crash)
+                .expect("Observers post_exec_all failed");
+
+            let interesting = fuzzer
+                .objective_mut()
+                .is_interesting(state, event_mgr, input, observers, &ExitKind::Crash)
+                .expect("In crash handler objective failure.");
+
+            if interesting {
+                let new_input = input.clone();
+                let mut new_testcase = Testcase::new(new_input);
+                new_testcase.add_metadata(ExitKind::Crash);
+                fuzzer
+                    .objective_mut()
+                    .append_metadata(state, &mut new_testcase)
+                    .expect("Failed adding metadata");
+                state
+                    .solutions_mut()
+                    .add(new_testcase)
+                    .expect("In crash handler solutions failure.");
+                event_mgr
+                    .fire(
+                        state,
+                        Event::Objective {
+                            objective_size: state.solutions().count(),
+                        },
+                    )
+                    .expect("Could not send crashing input");
+            }
+
+            event_mgr.on_restart(state).unwrap();
+
+            #[cfg(feature = "std")]
+            eprintln!("Waiting for broker...");
+            event_mgr.await_restart_safe();
+            #[cfg(feature = "std")]
+            eprintln!("Bye!");
+        }
+        // Don't need to exit, Asan will exit for us
+        // ExitProcess(1);
+    }
+}
+
 #[cfg(all(windows, feature = "std"))]
 mod windows_exception_handler {
     #[cfg(feature = "std")]
@@ -890,7 +1034,9 @@ mod windows_exception_handler {
         panic,
     };
 
-    use windows::Win32::System::Threading::ExitProcess;
+    use windows::Win32::System::Threading::{
+        EnterCriticalSection, ExitProcess, LeaveCriticalSection, RTL_CRITICAL_SECTION,
+    };
 
     use crate::{
         bolts::os::windows_exceptions::{
@@ -904,6 +1050,7 @@ mod windows_exception_handler {
         },
         feedbacks::Feedback,
         fuzzer::HasObjective,
+        inputs::UsesInput,
         observers::ObserversTuple,
         state::{HasClientPerfMonitor, HasMetadata, HasSolutions},
     };
@@ -935,12 +1082,6 @@ mod windows_exception_handler {
         }
     }
 
-    use windows::Win32::System::Threading::{
-        EnterCriticalSection, LeaveCriticalSection, RTL_CRITICAL_SECTION,
-    };
-
-    use crate::inputs::UsesInput;
-
     /// invokes the `post_exec` hook on all observer in case of panic
     #[cfg(feature = "std")]
     pub fn setup_panic_hook<E, EM, OF, Z>()
@@ -1030,6 +1171,7 @@ mod windows_exception_handler {
         }));
     }
 
+    /// Timeout handler for windows
     pub unsafe extern "system" fn inproc_timeout_handler<E, EM, OF, Z>(
         _p0: *mut u8,
         global_state: *mut c_void,

libafl_targets/build.rs

@@ -2,6 +2,7 @@
 
 use std::{env, fs::File, io::Write, path::Path};
 
+#[allow(clippy::too_many_lines)]
 fn main() {
     let out_dir = env::var_os("OUT_DIR").unwrap();
     let out_dir = out_dir.to_string_lossy().to_string();
@@ -127,6 +128,15 @@ fn main() {
             .compile("forkserver");
     }
 
+    #[cfg(windows)]
+    {
+        println!("cargo:rerun-if-changed=src/windows_asan.c");
+
+        cc::Build::new()
+            .file(src_dir.join("windows_asan.c"))
+            .compile("windows_asan");
+    }
+
     println!("cargo:rustc-link-search=native={}", &out_dir);
 
     println!("cargo:rerun-if-changed=build.rs");

libafl_targets/src/lib.rs

@@ -101,6 +101,11 @@ pub use cmplog::*;
 #[cfg(feature = "std")]
 pub mod drcov;
 
+#[cfg(all(windows, feature = "std"))]
+pub mod windows_asan;
+#[cfg(all(windows, feature = "std"))]
+pub use windows_asan::*;
+
 #[cfg(any(target_os = "linux", target_os = "freebsd"))]
 pub mod forkserver;
 #[cfg(any(target_os = "linux", target_os = "freebsd"))]

libafl_targets/src/windows_asan.c

@@ -0,0 +1,5 @@
+#include "common.h"
+
+EXT_FUNC_IMPL(__sanitizer_set_death_callback, void, (void), false) {
+  return;
+}

libafl_targets/src/windows_asan.rs

@@ -0,0 +1,35 @@
+//! Setup asan death callbback
+
+use libafl::{
+    events::{EventFirer, EventRestarter},
+    executors::{inprocess::windows_asan_handler::asan_death_handler, Executor, HasObservers},
+    feedbacks::Feedback,
+    state::{HasClientPerfMonitor, HasSolutions},
+    HasObjective,
+};
+
+/// Asan death callback type
+pub type CB = unsafe extern "C" fn() -> ();
+
+extern "C" {
+    fn __sanitizer_set_death_callback(cb: CB);
+}
+
+/// # Safety
+/// Setup asan callback on windows
+// See https://github.com/AFLplusplus/LibAFL/issues/769
+// This is needed to intercept asan error exit
+// When we use AddressSanitizer on windows, the crash handler is not called when ASAN detects an error
+// This is because, on linux, ASAN runtime raises SIGABRT so we can rely on the signal handler
+// but on windows it simply calls TerminateProcess.
+// so we need to call the api by asan to register the callback when asan is about to finish the process.
+pub unsafe fn setup_asan_callback<E, EM, OF, Z>(_executor: &E, _event_mgr: &EM, _fuzzer: &Z)
+where
+    E: Executor<EM, Z> + HasObservers,
+    EM: EventFirer<State = E::State> + EventRestarter<State = E::State>,
+    OF: Feedback<E::State>,
+    E::State: HasSolutions + HasClientPerfMonitor,
+    Z: HasObjective<Objective = OF, State = E::State>,
+{
+    __sanitizer_set_death_callback(asan_death_handler::<E, EM, OF, Z>);
+}