Timeout for Inprocess Executor on Windows (#267)

* start working on windows timeout

* salvage Input in timeout handler

* this time inproc_timeout_handler (need clean up later)

* cleaup

* more in inproc_timeout_handler

* fix for linux build

* more fixes for unix, fmt

* revert timeoutexecutor api

* revert baby_fuzzer/src/main.rs

* various fixes

* no unsafe

* remove timer in crash_handler

libafl/build.rs

@@ -7,12 +7,12 @@ fn main() {
     #[cfg(target_os = "windows")]
     #[allow(clippy::ptr_arg, clippy::upper_case_acronyms)]
     windows::build!(
-        Windows::Win32::Foundation::{HANDLE, BOOL, PSTR, CloseHandle, NTSTATUS},
+        Windows::Win32::Foundation::{HANDLE, HWND, BOOL, PSTR, CloseHandle, NTSTATUS},
         Windows::Win32::System::{
             Memory::{CreateFileMappingA, OpenFileMappingA, MapViewOfFile, UnmapViewOfFile, FILE_MAP, PAGE_TYPE},
             Diagnostics::Debug::{SetUnhandledExceptionFilter, EXCEPTION_POINTERS, EXCEPTION_RECORD, LPTOP_LEVEL_EXCEPTION_FILTER},
-            Threading::ExitProcess,
-        }
+            Threading::{CreateTimerQueue, CreateTimerQueueTimer, DeleteTimerQueueEx, DeleteTimerQueueTimer, ExitProcess},
+        },
     );
 
     // Set cfg flags depending on release channel

libafl/src/bolts/shmem.rs

@@ -1063,7 +1063,7 @@ pub mod win32_shmem {
             bindings::{
                 Windows::Win32::Foundation::{CloseHandle, BOOL, HANDLE, PSTR},
                 Windows::Win32::System::Memory::{
-                    CreateFileMappingA, MapViewOfFile, OpenFileMappingA, UnmapViewOfFile, FILE_MAP,
+                    CreateFileMappingA, MapViewOfFile, OpenFileMappingA, UnmapViewOfFile,
                     FILE_MAP_ALL_ACCESS, PAGE_READWRITE,
                 },
             },

libafl/src/executors/inprocess.rs

@@ -35,6 +35,9 @@ use crate::{
     Error,
 };
 
+#[cfg(windows)]
+use core::mem::transmute;
+
 /// The inmem executor simply calls a target function, then returns afterwards.
 #[allow(dead_code)]
 #[derive(Debug)]
@@ -101,7 +104,7 @@ where
                 &self.observers as *const _ as *const c_void,
             );
             data.crash_handler = self.crash_handler;
-            //data.timeout_handler = self.timeout_handler;
+            data.timeout_handler = self.timeout_handler;
             // Direct raw pointers access /aliasing is pretty undefined behavior.
             // Since the state and event may have moved in memory, refresh them right before the signal may happen
             write_volatile(&mut data.state_ptr, _state as *mut _ as *mut c_void);
@@ -211,8 +214,16 @@ where
                     S,
                     Z,
                 > as *const _,
-                // timeout_handler: windows_exception_handler::inproc_timeout_handler::<EM, I, OC, OF, OT, S, Z> as *const _,
-                timeout_handler: ptr::null(),
+                timeout_handler: windows_exception_handler::inproc_timeout_handler::<
+                    EM,
+                    I,
+                    OC,
+                    OF,
+                    OT,
+                    S,
+                    Z,
+                > as *const c_void,
+                // timeout_handler: ptr::null(),
                 phantom: PhantomData,
             })
         }
@@ -247,6 +258,8 @@ pub struct InProcessExecutorHandlerData {
     pub current_input_ptr: *const c_void,
     pub crash_handler: *const c_void,
     pub timeout_handler: *const c_void,
+    #[cfg(windows)]
+    pub timer_queue: *mut c_void,
 }
 
 unsafe impl Send for InProcessExecutorHandlerData {}
@@ -268,6 +281,8 @@ pub static mut GLOBAL_STATE: InProcessExecutorHandlerData = InProcessExecutorHan
     crash_handler: ptr::null(),
     /// The timeout handler fn
     timeout_handler: ptr::null(),
+    #[cfg(windows)]
+    timer_queue: ptr::null_mut(),
 };
 
 #[cfg(unix)]
@@ -284,7 +299,7 @@ mod unix_signal_handler {
         events::{Event, EventFirer, EventRestarter},
         executors::{
             inprocess::{InProcessExecutorHandlerData, GLOBAL_STATE},
-            timeout::remove_timeout,
+            timeout::unix_remove_timeout,
             ExitKind,
         },
         feedbacks::Feedback,
@@ -432,7 +447,7 @@ mod unix_signal_handler {
         I: Input,
         Z: HasObjective<I, OF, S>,
     {
-        remove_timeout();
+        unix_remove_timeout();
 
         #[cfg(all(target_os = "android", target_arch = "aarch64"))]
         let _context = *(((_context as *mut _ as *mut libc::c_void as usize) + 128)
@@ -557,13 +572,14 @@ mod unix_signal_handler {
 #[cfg(all(windows, feature = "std"))]
 mod windows_exception_handler {
     use alloc::vec::Vec;
+    use core::ffi::c_void;
     use core::{mem::transmute, ptr};
     #[cfg(feature = "std")]
     use std::io::{stdout, Write};
 
     use crate::{
         bolts::{
-            bindings::Windows::Win32::System::Threading::ExitProcess,
+            bindings::Windows::Win32::{Foundation::HANDLE, System::Threading::ExitProcess},
             os::windows_exceptions::{
                 ExceptionCode, Handler, CRASH_EXCEPTIONS, EXCEPTION_POINTERS,
             },
@@ -572,6 +588,7 @@ mod windows_exception_handler {
         events::{Event, EventFirer, EventRestarter},
         executors::{
             inprocess::{InProcessExecutorHandlerData, GLOBAL_STATE},
+            timeout::windows_delete_timer_queue,
             ExitKind,
         },
         feedbacks::Feedback,
@@ -608,6 +625,78 @@ mod windows_exception_handler {
         }
     }
 
+    pub unsafe extern "system" fn inproc_timeout_handler<EM, I, OC, OF, OT, S, Z>(
+        global_state: *mut c_void,
+        _p1: u8,
+    ) where
+        EM: EventFirer<I, S> + EventRestarter<S>,
+        OT: ObserversTuple<I, S>,
+        OC: Corpus<I>,
+        OF: Feedback<I, S>,
+        S: HasSolutions<OC, I> + HasClientPerfStats,
+        I: Input,
+        Z: HasObjective<I, OF, S>,
+    {
+        let data: &mut InProcessExecutorHandlerData =
+            &mut *(global_state as *mut InProcessExecutorHandlerData);
+
+        let state = (data.state_ptr as *mut S).as_mut().unwrap();
+        let event_mgr = (data.event_mgr_ptr as *mut EM).as_mut().unwrap();
+        let fuzzer = (data.fuzzer_ptr as *mut Z).as_mut().unwrap();
+        let observers = (data.observers_ptr as *const OT).as_ref().unwrap();
+
+        if data.current_input_ptr.is_null() {
+            #[cfg(feature = "std")]
+            dbg!("TIMEOUT or SIGUSR2 happened, but currently not fuzzing. Exiting");
+        } else {
+            #[cfg(feature = "std")]
+            println!("Timeout in fuzz run.");
+            #[cfg(feature = "std")]
+            let _res = stdout().flush();
+
+            let input = (data.current_input_ptr as *const I).as_ref().unwrap();
+            data.current_input_ptr = ptr::null();
+
+            let interesting = fuzzer
+                .objective_mut()
+                .is_interesting(state, event_mgr, input, observers, &ExitKind::Timeout)
+                .expect("In timeout handler objective failure.");
+
+            if interesting {
+                let mut new_testcase = Testcase::new(input.clone());
+                fuzzer
+                    .objective_mut()
+                    .append_metadata(state, &mut new_testcase)
+                    .expect("Failed adding metadata");
+                state
+                    .solutions_mut()
+                    .add(new_testcase)
+                    .expect("In timeout handler solutions failure.");
+                event_mgr
+                    .fire(
+                        state,
+                        Event::Objective {
+                            objective_size: state.solutions().count(),
+                        },
+                    )
+                    .expect("Could not send timeouting input");
+            }
+
+            event_mgr.on_restart(state).unwrap();
+
+            #[cfg(feature = "std")]
+            println!("Waiting for broker...");
+            event_mgr.await_restart_safe();
+            #[cfg(feature = "std")]
+            println!("Bye!");
+
+            event_mgr.await_restart_safe();
+
+            ExitProcess(1);
+        }
+        // println!("TIMER INVOKED!");
+    }
+
     pub unsafe fn inproc_crash_handler<EM, I, OC, OF, OT, S, Z>(
         code: ExceptionCode,
         exception_pointers: *mut EXCEPTION_POINTERS,
@@ -621,6 +710,14 @@ mod windows_exception_handler {
         I: Input,
         Z: HasObjective<I, OF, S>,
     {
+        // Have we set a timer_before?
+        match (data.timer_queue as *mut HANDLE).as_mut() {
+            Some(x) => {
+                windows_delete_timer_queue(*x);
+            }
+            None => {}
+        }
+
         #[cfg(feature = "std")]
         println!("Crashed with {}", code);
         if !data.current_input_ptr.is_null() {
@@ -705,6 +802,28 @@ mod windows_exception_handler {
     }
 }
 
+#[cfg(windows)]
+type WAITORTIMERCALLBACK = unsafe extern "system" fn(param0: *mut c_void, param1: u8);
+
+#[cfg(windows)]
+pub trait HasTimeoutHandler {
+    unsafe fn timeout_handler(&self) -> WAITORTIMERCALLBACK;
+}
+
+#[cfg(windows)]
+impl<'a, H, I, OT, S> HasTimeoutHandler for InProcessExecutor<'a, H, I, OT, S>
+where
+    H: FnMut(&I) -> ExitKind,
+    I: Input,
+    OT: ObserversTuple<I, S>,
+{
+    #[inline]
+    unsafe fn timeout_handler(&self) -> WAITORTIMERCALLBACK {
+        let func: WAITORTIMERCALLBACK = transmute(self.timeout_handler);
+        func
+    }
+}
+
 #[cfg(all(feature = "std", unix))]
 pub struct InProcessForkExecutor<'a, H, I, OT, S, SP>
 where

libafl/src/executors/timeout.rs

@@ -10,11 +10,26 @@ use crate::{
     Error,
 };
 
+#[cfg(windows)]
+use crate::executors::inprocess::{HasTimeoutHandler, GLOBAL_STATE};
+
 #[cfg(unix)]
 use core::{mem::zeroed, ptr::null_mut};
 #[cfg(unix)]
 use libc::c_int;
 
+#[cfg(windows)]
+use crate::bolts::bindings::Windows::Win32::{
+    Foundation::HANDLE,
+    System::Threading::{
+        CreateTimerQueue, CreateTimerQueueTimer, DeleteTimerQueueEx, DeleteTimerQueueTimer,
+        WORKER_THREAD_FLAGS,
+    },
+};
+
+#[cfg(windows)]
+use core::{ffi::c_void, ptr::write_volatile};
+
 #[repr(C)]
 #[cfg(unix)]
 struct Timeval {
@@ -38,15 +53,18 @@ extern "C" {
 const ITIMER_REAL: c_int = 0;
 
 /// Reset and remove the timeout
-pub fn remove_timeout() {
-    #[cfg(unix)]
+#[cfg(unix)]
+pub fn unix_remove_timeout() {
     unsafe {
         let mut itimerval_zero: Itimerval = zeroed();
         setitimer(ITIMER_REAL, &mut itimerval_zero, null_mut());
     }
-    #[cfg(windows)]
-    {
-        // TODO
+}
+
+#[cfg(windows)]
+pub fn windows_delete_timer_queue(timer_queue: HANDLE) {
+    unsafe {
+        DeleteTimerQueueEx(timer_queue, HANDLE::NULL);
     }
 }
 
@@ -55,6 +73,12 @@ pub struct TimeoutExecutor<E> {
     executor: E,
     #[cfg(unix)]
     itimerval: Itimerval,
+    #[cfg(windows)]
+    milli_sec: u32,
+    #[cfg(windows)]
+    ph_new_timer: HANDLE,
+    #[cfg(windows)]
+    timer_queue: HANDLE,
 }
 
 impl<E> TimeoutExecutor<E> {
@@ -83,15 +107,73 @@ impl<E> TimeoutExecutor<E> {
 
     #[cfg(windows)]
     pub fn new(executor: E, exec_tmout: Duration) -> Self {
-        Self { executor }
+        let milli_sec = exec_tmout.as_millis() as u32;
+        let timer_queue = unsafe { CreateTimerQueue() };
+        let ph_new_timer = HANDLE::NULL;
+        Self {
+            executor,
+            milli_sec,
+            ph_new_timer,
+            timer_queue,
+        }
     }
 
     /// Retrieve the inner `Executor` that is wrapped by this `TimeoutExecutor`.
     pub fn inner(&mut self) -> &mut E {
         &mut self.executor
     }
+
+    #[cfg(windows)]
+    pub fn windows_reset_timeout(&self) -> Result<(), Error> {
+        unsafe {
+            let code = DeleteTimerQueueTimer(self.timer_queue, self.ph_new_timer, HANDLE::NULL);
+            if !code.as_bool() {
+                return Err(Error::Unknown(format!("DeleteTimerQueueTimer failed.")));
+            }
+        }
+        Ok(())
+    }
 }
 
+#[cfg(windows)]
+impl<E, EM, I, S, Z> Executor<EM, I, S, Z> for TimeoutExecutor<E>
+where
+    E: Executor<EM, I, S, Z> + HasTimeoutHandler,
+    I: Input,
+{
+    fn run_target(
+        &mut self,
+        fuzzer: &mut Z,
+        state: &mut S,
+        mgr: &mut EM,
+        input: &I,
+    ) -> Result<ExitKind, Error> {
+        unsafe {
+            let data = &mut GLOBAL_STATE;
+            write_volatile(
+                &mut data.timer_queue,
+                &mut self.timer_queue as *mut _ as *mut c_void,
+            );
+            let code = CreateTimerQueueTimer(
+                &mut self.ph_new_timer,
+                &self.timer_queue,
+                Some(self.executor.timeout_handler()),
+                &mut GLOBAL_STATE as *mut _ as *mut c_void,
+                self.milli_sec,
+                0,
+                WORKER_THREAD_FLAGS::default(),
+            );
+            if !code.as_bool() {
+                return Err(Error::Unknown("CreateTimerQueue failed.".to_string()));
+            }
+            let ret = self.executor.run_target(fuzzer, state, mgr, input);
+            self.windows_reset_timeout()?;
+            ret
+        }
+    }
+}
+
+#[cfg(unix)]
 impl<E, EM, I, S, Z> Executor<EM, I, S, Z> for TimeoutExecutor<E>
 where
     E: Executor<EM, I, S, Z>,
@@ -107,15 +189,10 @@ where
         #[cfg(unix)]
         unsafe {
             setitimer(ITIMER_REAL, &mut self.itimerval, null_mut());
+            let ret = self.executor.run_target(fuzzer, state, mgr, input);
+            unix_remove_timeout();
+            ret
         }
-        #[cfg(windows)]
-        {
-            // TODO
-        }
-
-        let ret = self.executor.run_target(fuzzer, state, mgr, input);
-        remove_timeout();
-        ret
     }
 }