Fix sigint handlers (#1772)

* fix_handler

* cfg

* win

* fix

* toml

* f

* more

* rename and write_volatile

libafl/src/events/llmp.rs

@@ -1275,6 +1275,9 @@ where
                     self.shmem_provider.pre_fork()?;
                     match unsafe { fork() }? {
                         ForkResult::Parent(handle) => {
+                            unsafe {
+                                EVENTMGR_SIGHANDLER_STATE.set_exit_from_main();
+                            }
                             self.shmem_provider.post_fork(false)?;
                             handle.status()
                         }
@@ -1285,6 +1288,11 @@ where
                     }
                 };
 
+                #[cfg(all(unix, not(feature = "fork")))]
+                unsafe {
+                    EVENTMGR_SIGHANDLER_STATE.set_exit_from_main();
+                }
+
                 // On Windows (or in any case without fork), we spawn ourself again
                 #[cfg(any(windows, not(feature = "fork")))]
                 let child_status = startable_self()?.status()?;

libafl/src/events/mod.rs

@@ -56,6 +56,7 @@ use crate::{
 #[cfg(all(unix, feature = "std"))]
 pub static mut EVENTMGR_SIGHANDLER_STATE: ShutdownSignalData = ShutdownSignalData {
     shutting_down: false,
+    exit_from_main: false,
 };
 
 /// A signal handler for releasing `StateRestore` `ShMem`
@@ -64,6 +65,17 @@ pub static mut EVENTMGR_SIGHANDLER_STATE: ShutdownSignalData = ShutdownSignalDat
 #[derive(Debug, Clone)]
 pub struct ShutdownSignalData {
     shutting_down: bool,
+    exit_from_main: bool,
+}
+
+#[cfg(all(unix, feature = "std"))]
+impl ShutdownSignalData {
+    /// Set the flag to true, indicating that this process has allocated shmem
+    pub fn set_exit_from_main(&mut self) {
+        unsafe {
+            core::ptr::write_volatile(core::ptr::addr_of_mut!(self.exit_from_main), true);
+        }
+    }
 }
 
 /// Shutdown handler. `SigTerm`, `SigInterrupt`, `SigQuit` call this
@@ -76,6 +88,25 @@ impl Handler for ShutdownSignalData {
         _info: &mut siginfo_t,
         _context: Option<&mut ucontext_t>,
     ) {
+        /*
+        println!(
+            "in handler! {} {}",
+            self.exit_from_main,
+            std::process::id()
+        );
+        */
+        // if this process has not allocated any shmem. then simply exit()
+        if !self.exit_from_main {
+            unsafe {
+                #[cfg(unix)]
+                libc::_exit(0);
+
+                #[cfg(windows)]
+                windows::Win32::System::Threading::ExitProcess(1);
+            }
+        }
+
+        // else wait till the next is_shutting_down() is called. then the process will exit throught main().
         unsafe {
             core::ptr::write_volatile(core::ptr::addr_of_mut!(self.shutting_down), true);
         }

libafl/src/events/simple.rs

@@ -499,6 +499,11 @@ where
                     shmem_provider.pre_fork()?;
                     match unsafe { fork() }? {
                         ForkResult::Parent(handle) => {
+                            unsafe {
+                                // The parent will later exit through is_shutting down below
+                                // if the process exits gracefully, it cleans up the shmem.
+                                EVENTMGR_SIGHANDLER_STATE.set_exit_from_main();
+                            }
                             shmem_provider.post_fork(false)?;
                             handle.status()
                         }
@@ -509,6 +514,14 @@ where
                     }
                 };
 
+                // Same, as fork version, mark this main thread as the shmem allocator
+                // then it will not call exit or exitprocess in the sigint handler
+                // so that it exits after cleaning up the shmem segments
+                #[cfg(all(unix, not(feature = "fork")))]
+                unsafe {
+                    EVENTMGR_SIGHANDLER_STATE.set_exit_from_main();
+                }
+
                 // On Windows (or in any case without forks), we spawn ourself again
                 #[cfg(any(windows, not(feature = "fork")))]
                 let child_status = startable_self()?.status()?;

libafl/src/events/tcp.rs

@@ -1045,9 +1045,13 @@ where
     #[inline]
     #[allow(clippy::unused_self)]
     fn is_shutting_down() -> bool {
+        #[cfg(unix)]
         unsafe {
             core::ptr::read_volatile(core::ptr::addr_of!(EVENTMGR_SIGHANDLER_STATE.shutting_down))
         }
+
+        #[cfg(windows)]
+        false
     }
 
     /// Launch the restarting manager
@@ -1153,6 +1157,9 @@ where
                     self.shmem_provider.pre_fork()?;
                     match unsafe { fork() }? {
                         ForkResult::Parent(handle) => {
+                            unsafe {
+                                EVENTMGR_SIGHANDLER_STATE.set_exit_from_main();
+                            }
                             self.shmem_provider.post_fork(false)?;
                             handle.status()
                         }
@@ -1163,6 +1170,11 @@ where
                     }
                 };
 
+                #[cfg(all(unix, not(feature = "fork")))]
+                unsafe {
+                    EVENTMGR_SIGHANDLER_STATE.set_exit_from_main();
+                }
+
                 // On Windows (or in any case without fork), we spawn ourself again
                 #[cfg(any(windows, not(feature = "fork")))]
                 let child_status = startable_self()?.status()?;