diff --git a/libafl/src/events/llmp.rs b/libafl/src/events/llmp.rs index ebd7d08d2a..c5a72f80f3 100644 --- a/libafl/src/events/llmp.rs +++ b/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()?; diff --git a/libafl/src/events/mod.rs b/libafl/src/events/mod.rs index 140bbe14a3..c5c0b64434 100644 --- a/libafl/src/events/mod.rs +++ b/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); } diff --git a/libafl/src/events/simple.rs b/libafl/src/events/simple.rs index e8f54dbddb..5a4bb1fbe7 100644 --- a/libafl/src/events/simple.rs +++ b/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()?; diff --git a/libafl/src/events/tcp.rs b/libafl/src/events/tcp.rs index ef5f66f4eb..8042fd608b 100644 --- a/libafl/src/events/tcp.rs +++ b/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()?;