From eff40320eb88c7d46505359029f55aeec3b8e155 Mon Sep 17 00:00:00 2001 From: Aarnav Date: Tue, 2 Jul 2024 17:45:20 +0200 Subject: [PATCH] Add Stoppable trait to State which exposes an API to stop the fuzzer (#2325) * add HasStopNext to State which exposes an API to stop the fuzzer. Stops the fuzzer in fuzz_loop or fuzz_loop_for when set to true * fix import * rename HasStopNext to HasShouldStopFuzzing and stop_next to should_stop_fuzzing * added HasShouldStopFuzzing trait constraint for libafl_libfuzzer_runtime fuzzer * rename HasShouldStopFuzzing to Stoppable and add it as a type constraint in libafl_libfuzzer report.rs * rename should_stop_fuzzing -> should_stop * introduce Event::Stop * fix prelude import * Call send_exiting when processing Event::Stop in restartable managers * fix clippy * introduce on_shutdown function in EventProcessor, a function to exit without saving state gracefully. In contrast with on_restart. * call manager.on_shutdown when stopping in fuzz_loop due to state.should_stop * Add missing on_shutdown implementations Check after every stage in Stages::perform_all if should exit and do so. * remove specialization * fix doc * introduce EventProcessor constraint in libafl_libfuzzer_runtime run clippy in libafl_libfuzzer_runtime * fix CentralizedEventManager's on_shutdown not calling inner.on_shutdown * fix bugs in CentralizedLauncher that wouldn't allow children to terminate properly * don't call send_exiting when processing Event::Stop since it will be called when calling on_shutdown anyways * clippy * add set_exit_after so broker does not need to inner_mut to set exit_cleanly_after * return Cow from Event::name_detailed instead of a String * fix missing import in libafl_libfuzzer_runtime * add initate_stop and reset_stop to Stoppable trait to superceed should_stop_mut * clippy * typo * rename initate_stop to request_stop, should_stop to stop_requested and reset_stop to discard_stop_request * fix missing import * windows clippy fix * fix broker typo --- libafl/src/events/broker_hooks/centralized.rs | 2 +- libafl/src/events/broker_hooks/mod.rs | 1 + libafl/src/events/centralized.rs | 14 +++++- libafl/src/events/launcher.rs | 27 +++++------ libafl/src/events/llmp/mgr.rs | 7 +++ libafl/src/events/llmp/mod.rs | 5 +- libafl/src/events/llmp/restarting.rs | 18 +++---- libafl/src/events/mod.rs | 44 ++++++++++------- libafl/src/events/simple.rs | 46 +++++++++++------- libafl/src/events/tcp.rs | 13 ++++- libafl/src/fuzzer/mod.rs | 20 ++++++-- libafl/src/stages/mod.rs | 32 ++++++++++--- libafl/src/state/mod.rs | 47 +++++++++++++++++++ libafl_bolts/examples/llmp_test/main.rs | 6 +-- libafl_bolts/src/llmp.rs | 6 +++ .../libafl_libfuzzer_runtime/src/fuzz.rs | 14 ++++-- .../libafl_libfuzzer_runtime/src/report.rs | 9 ++-- libafl_nyx/src/executor.rs | 19 ++++---- 18 files changed, 235 insertions(+), 95 deletions(-) diff --git a/libafl/src/events/broker_hooks/centralized.rs b/libafl/src/events/broker_hooks/centralized.rs index b95fd42b0f..eaa6c6bb4b 100644 --- a/libafl/src/events/broker_hooks/centralized.rs +++ b/libafl/src/events/broker_hooks/centralized.rs @@ -95,7 +95,7 @@ where event: &Event, ) -> Result { match &event { - Event::NewTestcase { .. } => Ok(BrokerEventResult::Forward), + Event::NewTestcase { .. } | Event::Stop => Ok(BrokerEventResult::Forward), _ => Ok(BrokerEventResult::Handled), } } diff --git a/libafl/src/events/broker_hooks/mod.rs b/libafl/src/events/broker_hooks/mod.rs index 50352502d9..4e3780adb6 100644 --- a/libafl/src/events/broker_hooks/mod.rs +++ b/libafl/src/events/broker_hooks/mod.rs @@ -208,6 +208,7 @@ where Ok(BrokerEventResult::Handled) } Event::CustomBuf { .. } => Ok(BrokerEventResult::Forward), + Event::Stop => Ok(BrokerEventResult::Forward), //_ => Ok(BrokerEventResult::Forward), } } diff --git a/libafl/src/events/centralized.rs b/libafl/src/events/centralized.rs index c4cd759579..ec7bfc271f 100644 --- a/libafl/src/events/centralized.rs +++ b/libafl/src/events/centralized.rs @@ -39,7 +39,7 @@ use crate::{ fuzzer::{EvaluatorObservers, ExecutionProcessor}, inputs::{Input, NopInput, UsesInput}, observers::{ObserversTuple, TimeObserver}, - state::{HasExecutions, HasLastReportTime, NopState, State, UsesState}, + state::{HasExecutions, HasLastReportTime, NopState, State, Stoppable, UsesState}, Error, HasMetadata, }; @@ -279,6 +279,7 @@ where self.inner.should_send() } + #[allow(clippy::match_same_arms)] fn fire( &mut self, state: &mut Self::State, @@ -295,6 +296,7 @@ where true } Event::UpdateExecStats { .. } => true, // send it but this guy won't be handled. the only purpose is to keep this client alive else the broker thinks it is dead and will dc it + Event::Stop => true, _ => false, }; @@ -391,6 +393,11 @@ where self.inner.process(fuzzer, state, executor) } } + + fn on_shutdown(&mut self) -> Result<(), Error> { + self.inner.on_shutdown()?; + self.client.sender_mut().send_exiting() + } } impl EventManager for CentralizedEventManager @@ -476,7 +483,7 @@ impl CentralizedEventManager where EM: UsesState + EventFirer + AdaptiveSerializer + HasEventManagerId, EMH: EventManagerHooksTuple, - S: State, + S: State + Stoppable, SP: ShMemProvider, { #[cfg(feature = "llmp_compression")] @@ -662,6 +669,9 @@ where log::debug!("[{}] {} was discarded...)", process::id(), event_name); } } + Event::Stop => { + state.request_stop(); + } _ => { return Err(Error::unknown(format!( "Received illegal message that message should not have arrived: {:?}.", diff --git a/libafl/src/events/launcher.rs b/libafl/src/events/launcher.rs index 21ce5cc08b..01134f1ddd 100644 --- a/libafl/src/events/launcher.rs +++ b/libafl/src/events/launcher.rs @@ -28,6 +28,8 @@ use std::process::Stdio; #[cfg(all(unix, feature = "std"))] use std::{fs::File, os::unix::io::AsRawFd}; +#[cfg(all(unix, feature = "std", feature = "fork"))] +use libafl_bolts::llmp::Broker; #[cfg(all(unix, feature = "std", feature = "fork"))] use libafl_bolts::llmp::Brokers; #[cfg(all(unix, feature = "std", feature = "fork"))] @@ -716,7 +718,8 @@ where self.time_obs.clone(), )?; - self.main_run_client.take().unwrap()(state, c_mgr, *bind_to) + self.main_run_client.take().unwrap()(state, c_mgr, *bind_to)?; + Err(Error::shutting_down()) } else { // Secondary clients log::debug!("Running secondary client on PID {}", std::process::id()); @@ -733,7 +736,8 @@ where self.time_obs.clone(), )?; - self.secondary_run_client.take().unwrap()(state, c_mgr, *bind_to) + self.secondary_run_client.take().unwrap()(state, c_mgr, *bind_to)?; + Err(Error::shutting_down()) } }?, }; @@ -756,6 +760,7 @@ where }; let mut brokers = Brokers::new(); + let exit_cleanly_after = NonZeroUsize::try_from(self.cores.ids.len()).unwrap(); // Add centralized broker brokers.add(Box::new({ @@ -769,12 +774,14 @@ where let centralized_hooks = tuple_list!(CentralizedLlmpHook::::new()?); // TODO switch to false after solving the bug - LlmpBroker::with_keep_pages_attach_to_tcp( + let mut broker = LlmpBroker::with_keep_pages_attach_to_tcp( self.shmem_provider.clone(), centralized_hooks, self.centralized_broker_port, true, - )? + )?; + broker.set_exit_after(exit_cleanly_after); + broker })); #[cfg(feature = "multi_machine")] @@ -808,19 +815,11 @@ where broker.inner_mut().connect_b2b(remote_broker_addr)?; }; - let exit_cleanly_after = NonZeroUsize::try_from(self.cores.ids.len()).unwrap(); - - broker - .inner_mut() - .set_exit_cleanly_after(exit_cleanly_after); + broker.set_exit_after(exit_cleanly_after); brokers.add(Box::new(broker)); } - - log::debug!( - "Brokers have been initialized on port {}.", - std::process::id() - ); + log::debug!("Broker has been initialized; pid {}.", std::process::id()); // Loop over all the brokers that should be polled brokers.loop_with_timeouts(Duration::from_secs(30), Some(Duration::from_millis(5))); diff --git a/libafl/src/events/llmp/mgr.rs b/libafl/src/events/llmp/mgr.rs index 70fc95ef6c..3980d00a57 100644 --- a/libafl/src/events/llmp/mgr.rs +++ b/libafl/src/events/llmp/mgr.rs @@ -468,6 +468,9 @@ where } } } + Event::Stop => { + state.request_stop(); + } _ => { return Err(Error::unknown(format!( "Received illegal message that message should not have arrived: {:?}.", @@ -626,6 +629,10 @@ where } Ok(count) } + + fn on_shutdown(&mut self) -> Result<(), Error> { + self.send_exiting() + } } impl EventManager for LlmpEventManager diff --git a/libafl/src/events/llmp/mod.rs b/libafl/src/events/llmp/mod.rs index 6ed2f27dc0..a3550588f6 100644 --- a/libafl/src/events/llmp/mod.rs +++ b/libafl/src/events/llmp/mod.rs @@ -20,7 +20,7 @@ use crate::{ executors::{Executor, HasObservers}, fuzzer::{EvaluatorObservers, ExecutionProcessor}, inputs::{Input, InputConverter, NopInput, NopInputConverter, UsesInput}, - state::{HasExecutions, NopState, State, UsesState}, + state::{HasExecutions, NopState, State, Stoppable, UsesState}, Error, HasMetadata, }; @@ -253,7 +253,7 @@ where impl LlmpEventConverter where - S: UsesInput + HasExecutions + HasMetadata, + S: UsesInput + HasExecutions + HasMetadata + Stoppable, SP: ShMemProvider, IC: InputConverter, ICB: InputConverter, @@ -329,6 +329,7 @@ where } Ok(()) } + Event::Stop => Ok(()), _ => Err(Error::unknown(format!( "Received illegal message that message should not have arrived: {:?}.", event.name() diff --git a/libafl/src/events/llmp/restarting.rs b/libafl/src/events/llmp/restarting.rs index c28ddd3c7d..fd52a5c176 100644 --- a/libafl/src/events/llmp/restarting.rs +++ b/libafl/src/events/llmp/restarting.rs @@ -22,15 +22,15 @@ use libafl_bolts::os::startable_self; use libafl_bolts::os::unix_signals::setup_signal_handler; #[cfg(all(feature = "std", feature = "fork", unix))] use libafl_bolts::os::{fork, ForkResult}; -use libafl_bolts::{ - llmp::LlmpBroker, - shmem::ShMemProvider, - tuples::{tuple_list, Handle}, -}; #[cfg(feature = "std")] use libafl_bolts::{ llmp::LlmpConnection, os::CTRL_C_EXIT, shmem::StdShMemProvider, staterestore::StateRestorer, }; +use libafl_bolts::{ + llmp::{Broker, LlmpBroker}, + shmem::ShMemProvider, + tuples::{tuple_list, Handle}, +}; use serde::{Deserialize, Serialize}; #[cfg(feature = "std")] use typed_builder::TypedBuilder; @@ -218,6 +218,10 @@ where self.intermediate_save()?; Ok(res) } + + fn on_shutdown(&mut self) -> Result<(), Error> { + self.send_exiting() + } } #[cfg(feature = "std")] @@ -448,9 +452,7 @@ where }; if let Some(exit_cleanly_after) = self.exit_cleanly_after { - broker - .inner_mut() - .set_exit_cleanly_after(exit_cleanly_after); + broker.set_exit_after(exit_cleanly_after); } broker.loop_with_timeouts(Duration::from_secs(30), Some(Duration::from_millis(5))); diff --git a/libafl/src/events/mod.rs b/libafl/src/events/mod.rs index 4f554ad830..f4de028a57 100644 --- a/libafl/src/events/mod.rs +++ b/libafl/src/events/mod.rs @@ -21,12 +21,7 @@ pub use llmp::*; pub mod tcp; pub mod broker_hooks; -use alloc::{ - borrow::Cow, - boxed::Box, - string::{String, ToString}, - vec::Vec, -}; +use alloc::{borrow::Cow, boxed::Box, string::String, vec::Vec}; use core::{ fmt, hash::{BuildHasher, Hasher}, @@ -353,6 +348,8 @@ where /// Tag of this buffer tag: String, }, + /// Exit gracefully + Stop, /*/// A custom type Custom { // TODO: Allow custom events @@ -364,7 +361,8 @@ impl Event where I: Input, { - fn name(&self) -> &str { + /// Event's corresponding name + pub fn name(&self) -> &str { match self { Event::NewTestcase { .. } => "Testcase", Event::UpdateExecStats { .. } => "Client Heartbeat", @@ -377,21 +375,24 @@ where /*Event::Custom { sender_id: _, /*custom_event} => custom_event.name()*/ } => "todo",*/ + Event::Stop => "Stop", } } - fn name_detailed(&self) -> String { + /// Event's corresponding name with additional info + fn name_detailed(&self) -> Cow<'static, str> { match self { Event::NewTestcase { input, .. } => { - format!("Testcase {}", input.generate_name(None)) + Cow::Owned(format!("Testcase {}", input.generate_name(None))) } - Event::UpdateExecStats { .. } => "Client Heartbeat".to_string(), - Event::UpdateUserStats { .. } => "UserStats".to_string(), + Event::UpdateExecStats { .. } => Cow::Borrowed("Client Heartbeat"), + Event::UpdateUserStats { .. } => Cow::Borrowed("UserStats"), #[cfg(feature = "introspection")] - Event::UpdatePerfMonitor { .. } => "PerfMonitor".to_string(), - Event::Objective { .. } => "Objective".to_string(), - Event::Log { .. } => "Log".to_string(), - Event::CustomBuf { .. } => "CustomBuf".to_string(), + Event::UpdatePerfMonitor { .. } => Cow::Borrowed("PerfMonitor"), + Event::Objective { .. } => Cow::Borrowed("Objective"), + Event::Log { .. } => Cow::Borrowed("Log"), + Event::CustomBuf { .. } => Cow::Borrowed("CustomBuf"), + Event::Stop => Cow::Borrowed("Stop"), /*Event::Custom { sender_id: _, /*custom_event} => custom_event.name()*/ } => "todo",*/ @@ -399,7 +400,7 @@ where } } -/// [`EventFirer`] fire an event. +/// [`EventFirer`] fires an event. pub trait EventFirer: UsesState { /// Send off an [`Event`] to the broker /// @@ -574,6 +575,9 @@ pub trait EventProcessor: UsesState { state: &mut Self::State, executor: &mut E, ) -> Result; + + /// Shutdown gracefully; typically without saving state. + fn on_shutdown(&mut self) -> Result<(), Error>; } /// The id of this [`EventManager`]. /// For multi processed [`EventManager`]s, @@ -662,6 +666,10 @@ where ) -> Result { Ok(0) } + + fn on_shutdown(&mut self) -> Result<(), Error> { + Ok(()) + } } impl EventManager for NopEventManager where @@ -793,6 +801,10 @@ where ) -> Result { self.inner.process(fuzzer, state, executor) } + + fn on_shutdown(&mut self) -> Result<(), Error> { + self.inner.on_shutdown() + } } impl EventManager for MonitorTypedEventManager diff --git a/libafl/src/events/simple.rs b/libafl/src/events/simple.rs index a9b5a2fff2..ce17d038fb 100644 --- a/libafl/src/events/simple.rs +++ b/libafl/src/events/simple.rs @@ -32,7 +32,7 @@ use crate::{ }, inputs::UsesInput, monitors::Monitor, - state::{HasExecutions, HasLastReportTime, State, UsesState}, + state::{HasExecutions, HasLastReportTime, State, Stoppable, UsesState}, Error, HasMetadata, }; #[cfg(feature = "std")] @@ -50,7 +50,7 @@ const _ENV_FUZZER_BROKER_CLIENT_INITIAL: &str = "_AFL_ENV_FUZZER_BROKER_CLIENT"; /// A simple, single-threaded event manager that just logs pub struct SimpleEventManager where - S: UsesInput, + S: UsesInput + Stoppable, { /// The monitor monitor: MT, @@ -64,7 +64,7 @@ where impl Debug for SimpleEventManager where MT: Debug, - S: UsesInput, + S: UsesInput + Stoppable, { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { f.debug_struct("SimpleEventManager") @@ -128,6 +128,10 @@ where } Ok(count) } + + fn on_shutdown(&mut self) -> Result<(), Error> { + self.send_exiting() + } } impl EventManager for SimpleEventManager @@ -163,7 +167,7 @@ where impl HasEventManagerId for SimpleEventManager where MT: Monitor, - S: UsesInput, + S: UsesInput + Stoppable, { fn mgr_id(&self) -> EventManagerId { EventManagerId(0) @@ -173,7 +177,7 @@ where #[cfg(feature = "std")] impl SimpleEventManager where - S: UsesInput, + S: UsesInput + Stoppable, { /// Creates a [`SimpleEventManager`] that just prints to `stdout`. #[must_use] @@ -185,7 +189,7 @@ where impl SimpleEventManager where MT: Monitor, //TODO CE: CustomEvent, - S: UsesInput, + S: UsesInput + Stoppable, { /// Creates a new [`SimpleEventManager`]. pub fn new(monitor: MT) -> Self { @@ -281,22 +285,27 @@ where Ok(BrokerEventResult::Handled) } Event::CustomBuf { .. } => Ok(BrokerEventResult::Forward), - //_ => Ok(BrokerEventResult::Forward), + Event::Stop => Ok(BrokerEventResult::Forward), } } // Handle arriving events in the client #[allow(clippy::needless_pass_by_value, clippy::unused_self)] fn handle_in_client(&mut self, state: &mut S, event: Event) -> Result<(), Error> { - if let Event::CustomBuf { tag, buf } = &event { - for handler in &mut self.custom_buf_handlers { - handler(state, tag, buf)?; + match event { + Event::CustomBuf { buf, tag } => { + for handler in &mut self.custom_buf_handlers { + handler(state, &tag, &buf)?; + } + Ok(()) } - Ok(()) - } else { - Err(Error::unknown(format!( + Event::Stop => { + state.request_stop(); + Ok(()) + } + _ => Err(Error::unknown(format!( "Received illegal message that message should not have arrived: {event:?}." - ))) + ))), } } } @@ -309,7 +318,7 @@ where #[derive(Debug)] pub struct SimpleRestartingEventManager where - S: UsesInput, + S: UsesInput + Stoppable, SP: ShMemProvider, //CE: CustomEvent, { /// The actual simple event mgr @@ -388,6 +397,9 @@ where ) -> Result { self.simple_event_mgr.process(fuzzer, state, executor) } + fn on_shutdown(&mut self) -> Result<(), Error> { + self.send_exiting() + } } #[cfg(feature = "std")] @@ -427,7 +439,7 @@ where impl HasEventManagerId for SimpleRestartingEventManager where MT: Monitor, - S: UsesInput, + S: UsesInput + Stoppable, SP: ShMemProvider, { fn mgr_id(&self) -> EventManagerId { @@ -439,7 +451,7 @@ where #[allow(clippy::type_complexity, clippy::too_many_lines)] impl SimpleRestartingEventManager where - S: UsesInput, + S: UsesInput + Stoppable, SP: ShMemProvider, MT: Monitor, //TODO CE: CustomEvent, { diff --git a/libafl/src/events/tcp.rs b/libafl/src/events/tcp.rs index 2dbde88ca8..fac23b1f07 100644 --- a/libafl/src/events/tcp.rs +++ b/libafl/src/events/tcp.rs @@ -407,7 +407,7 @@ where log::log!((*severity_level).into(), "{message}"); Ok(BrokerEventResult::Handled) } - Event::CustomBuf { .. } => Ok(BrokerEventResult::Forward), + Event::CustomBuf { .. } | Event::Stop => Ok(BrokerEventResult::Forward), //_ => Ok(BrokerEventResult::Forward), } } @@ -657,6 +657,9 @@ where } } } + Event::Stop => { + state.request_stop(); + } _ => { return Err(Error::unknown(format!( "Received illegal message that message should not have arrived: {:?}.", @@ -809,6 +812,10 @@ where Ok(count) } + + fn on_shutdown(&mut self) -> Result<(), Error> { + self.send_exiting() + } } impl EventManager for TcpEventManager @@ -967,6 +974,10 @@ where fn process(&mut self, fuzzer: &mut Z, state: &mut S, executor: &mut E) -> Result { self.tcp_mgr.process(fuzzer, state, executor) } + + fn on_shutdown(&mut self) -> Result<(), Error> { + self.send_exiting() + } } #[cfg(feature = "std")] diff --git a/libafl/src/fuzzer/mod.rs b/libafl/src/fuzzer/mod.rs index ca188a93b6..9ee538ddfd 100644 --- a/libafl/src/fuzzer/mod.rs +++ b/libafl/src/fuzzer/mod.rs @@ -19,7 +19,7 @@ use crate::{ start_timer, state::{ HasCorpus, HasCurrentTestcase, HasExecutions, HasImported, HasLastReportTime, HasSolutions, - UsesState, + Stoppable, UsesState, }, Error, HasMetadata, }; @@ -182,9 +182,9 @@ pub trait Evaluator: UsesState { /// The main fuzzer trait. pub trait Fuzzer: Sized + UsesState where - Self::State: HasMetadata + HasExecutions + HasLastReportTime, + Self::State: HasMetadata + HasExecutions + HasLastReportTime + Stoppable, E: UsesState, - EM: ProgressReporter, + EM: ProgressReporter + EventProcessor, ST: StagesTuple, { /// Fuzz for a single iteration. @@ -216,8 +216,14 @@ where loop { // log::info!("Starting another fuzz_loop"); manager.maybe_report_progress(state, monitor_timeout)?; + if state.stop_requested() { + state.discard_stop_request(); + manager.on_shutdown()?; + break; + } self.fuzz_one(stages, executor, state, manager)?; } + Ok(()) } /// Fuzz for n iterations. @@ -248,6 +254,10 @@ where for _ in 0..iters { manager.maybe_report_progress(state, monitor_timeout)?; + if state.stop_requested() { + state.discard_stop_request(); + break; + } ret = Some(self.fuzz_one(stages, executor, state, manager)?); } @@ -865,7 +875,7 @@ pub mod test { use crate::{ corpus::CorpusId, - events::ProgressReporter, + events::{EventProcessor, ProgressReporter}, stages::{HasCurrentStage, StagesTuple}, state::{HasExecutions, HasLastReportTime, State, UsesState}, Fuzzer, HasMetadata, @@ -901,7 +911,7 @@ pub mod test { impl Fuzzer for NopFuzzer where E: UsesState, - EM: ProgressReporter, + EM: ProgressReporter + EventProcessor, ST: StagesTuple, Self::State: HasMetadata + HasExecutions + HasLastReportTime + HasCurrentStage, { diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 5d290517cb..d4e40c97d1 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -45,13 +45,13 @@ pub use unicode::*; use crate::{ corpus::{CorpusId, HasCurrentCorpusId}, - events::{EventFirer, EventRestarter, HasEventManagerId, ProgressReporter}, + events::{EventFirer, EventProcessor, EventRestarter, HasEventManagerId, ProgressReporter}, executors::{Executor, HasObservers}, inputs::UsesInput, observers::ObserversTuple, schedulers::Scheduler, stages::push::PushStage, - state::{HasCorpus, HasExecutions, HasLastReportTime, HasRand, State, UsesState}, + state::{HasCorpus, HasExecutions, HasLastReportTime, HasRand, State, Stoppable, UsesState}, Error, EvaluatorObservers, ExecutesInput, ExecutionProcessor, HasMetadata, HasNamedMetadata, HasScheduler, }; @@ -135,7 +135,7 @@ where Z: UsesState, S: UsesInput + HasCurrentStage, { - /// Performs all `Stages` in this tuple + /// Performs all `Stages` in this tuple. fn perform_all( &mut self, fuzzer: &mut Z, @@ -174,10 +174,13 @@ where Head: Stage, Tail: StagesTuple + HasConstLen, E: UsesState, - EM: UsesState, + EM: UsesState + EventProcessor, Z: UsesState, Head::State: HasCurrentStage, { + /// Performs all stages in the tuple, + /// Checks after every stage if state wants to stop + /// and returns an [`Error::ShuttingDown`] if so fn perform_all( &mut self, fuzzer: &mut Z, @@ -211,6 +214,12 @@ where } } + if state.stop_requested() { + state.discard_stop_request(); + manager.on_shutdown()?; + return Err(Error::shutting_down()); + } + // Execute the remaining stages self.1.perform_all(fuzzer, executor, state, manager) } @@ -273,10 +282,13 @@ impl StagesTuple for Vec>> where E: UsesState, - EM: UsesState, + EM: UsesState + EventProcessor, Z: UsesState, S: UsesInput + HasCurrentStage + State, { + /// Performs all stages in the `Vec` + /// Checks after every stage if state wants to stop + /// and returns an [`Error::ShuttingDown`] if so fn perform_all( &mut self, fuzzer: &mut Z, @@ -284,8 +296,14 @@ where state: &mut S, manager: &mut EM, ) -> Result<(), Error> { - self.iter_mut() - .try_for_each(|x| x.perform_restartable(fuzzer, executor, state, manager)) + self.iter_mut().try_for_each(|x| { + if state.stop_requested() { + state.discard_stop_request(); + manager.on_shutdown()?; + return Err(Error::shutting_down()); + } + x.perform_restartable(fuzzer, executor, state, manager) + }) } } diff --git a/libafl/src/state/mod.rs b/libafl/src/state/mod.rs index dc9c2d4426..e31efb3efc 100644 --- a/libafl/src/state/mod.rs +++ b/libafl/src/state/mod.rs @@ -55,6 +55,7 @@ pub trait State: + MaybeHasScalabilityMonitor + HasCurrentCorpusId + HasCurrentStage + + Stoppable { } @@ -261,6 +262,9 @@ pub struct StdState { last_report_time: Option, /// The current index of the corpus; used to record for resumable fuzzing. corpus_id: Option, + /// Request the fuzzer to stop at the start of the next stage + /// or at the beginning of the next fuzzing iteration + stop_requested: bool, stage_stack: StageStack, phantom: PhantomData, } @@ -532,6 +536,32 @@ where } } +/// A trait for types that want to expose a stop API +pub trait Stoppable { + /// Check if stop is requested + fn stop_requested(&self) -> bool; + + /// Request to stop + fn request_stop(&mut self); + + /// Discard the stop request + fn discard_stop_request(&mut self); +} + +impl Stoppable for StdState { + fn request_stop(&mut self) { + self.stop_requested = true; + } + + fn discard_stop_request(&mut self) { + self.stop_requested = false; + } + + fn stop_requested(&self) -> bool { + self.stop_requested + } +} + impl HasCurrentStage for StdState { fn set_current_stage_idx(&mut self, idx: StageId) -> Result<(), Error> { self.stage_stack.set_current_stage_idx(idx) @@ -1087,6 +1117,7 @@ where corpus, solutions, max_size: DEFAULT_MAX_SIZE, + stop_requested: false, #[cfg(feature = "introspection")] introspection_monitor: ClientPerfMonitor::new(), #[cfg(feature = "scalability_introspection")] @@ -1135,6 +1166,7 @@ impl HasScalabilityMonitor for StdState { pub struct NopState { metadata: SerdeAnyMap, execution: u64, + stop_requested: bool, rand: StdRand, phantom: PhantomData, } @@ -1147,6 +1179,7 @@ impl NopState { metadata: SerdeAnyMap::new(), execution: 0, rand: StdRand::default(), + stop_requested: false, phantom: PhantomData, } } @@ -1179,6 +1212,20 @@ impl HasExecutions for NopState { } } +impl Stoppable for NopState { + fn request_stop(&mut self) { + self.stop_requested = true; + } + + fn discard_stop_request(&mut self) { + self.stop_requested = false; + } + + fn stop_requested(&self) -> bool { + self.stop_requested + } +} + impl HasLastReportTime for NopState { fn last_report_time(&self) -> &Option { unimplemented!(); diff --git a/libafl_bolts/examples/llmp_test/main.rs b/libafl_bolts/examples/llmp_test/main.rs index eff560778b..d7c166f1a8 100644 --- a/libafl_bolts/examples/llmp_test/main.rs +++ b/libafl_bolts/examples/llmp_test/main.rs @@ -168,6 +168,8 @@ fn main() { fn main() -> Result<(), Box> { /* The main node has a broker, and a few worker threads */ + use libafl_bolts::llmp::Broker; + let mode = std::env::args() .nth(1) .expect("no mode specified, chose 'broker', 'b2b', 'ctr', 'adder', 'large', or 'exiting'"); @@ -193,9 +195,7 @@ fn main() -> Result<(), Box> { )?; broker.inner_mut().launch_tcp_listener_on(port)?; // Exit when we got at least _n_ nodes, and all of them quit. - broker - .inner_mut() - .set_exit_cleanly_after(NonZeroUsize::new(1_usize).unwrap()); + broker.set_exit_after(NonZeroUsize::new(1_usize).unwrap()); broker.loop_with_timeouts(BROKER_TIMEOUT, Some(SLEEP_BETWEEN_FORWARDS)); } "b2b" => { diff --git a/libafl_bolts/src/llmp.rs b/libafl_bolts/src/llmp.rs index 2d3a00f4f3..e3514c83a2 100644 --- a/libafl_bolts/src/llmp.rs +++ b/libafl_bolts/src/llmp.rs @@ -2091,6 +2091,9 @@ pub trait Broker { /// Getter to `exit_after` fn exit_after(&self) -> Option; + /// Setter for `exit_after` + fn set_exit_after(&mut self, n_clients: NonZeroUsize); + /// Getter to `has_clients` fn has_clients(&self) -> bool; @@ -2124,6 +2127,9 @@ where fn exit_after(&self) -> Option { self.inner.exit_cleanly_after } + fn set_exit_after(&mut self, n_clients: NonZeroUsize) { + self.inner.set_exit_cleanly_after(n_clients); + } fn has_clients(&self) -> bool { self.inner.has_clients() diff --git a/libafl_libfuzzer/libafl_libfuzzer_runtime/src/fuzz.rs b/libafl_libfuzzer/libafl_libfuzzer_runtime/src/fuzz.rs index f671bf120a..8639c72401 100644 --- a/libafl_libfuzzer/libafl_libfuzzer_runtime/src/fuzz.rs +++ b/libafl_libfuzzer/libafl_libfuzzer_runtime/src/fuzz.rs @@ -13,7 +13,7 @@ use std::{ use libafl::{ corpus::Corpus, events::{ - launcher::Launcher, EventConfig, ProgressReporter, SimpleEventManager, + launcher::Launcher, EventConfig, EventProcessor, ProgressReporter, SimpleEventManager, SimpleRestartingEventManager, }, executors::ExitKind, @@ -23,7 +23,7 @@ use libafl::{ Monitor, MultiMonitor, }, stages::{HasCurrentStage, StagesTuple}, - state::{HasExecutions, HasLastReportTime, HasSolutions, UsesState}, + state::{HasExecutions, HasLastReportTime, HasSolutions, Stoppable, UsesState}, Error, Fuzzer, HasMetadata, }; use libafl_bolts::{ @@ -66,9 +66,15 @@ fn do_fuzz( ) -> Result<(), Error> where F: Fuzzer, - S: HasMetadata + HasExecutions + UsesInput + HasSolutions + HasLastReportTime + HasCurrentStage, + S: HasMetadata + + HasExecutions + + UsesInput + + HasSolutions + + HasLastReportTime + + HasCurrentStage + + Stoppable, E: UsesState, - EM: ProgressReporter, + EM: ProgressReporter + EventProcessor, ST: StagesTuple, { if let Some(solution) = state.solutions().last() { diff --git a/libafl_libfuzzer/libafl_libfuzzer_runtime/src/report.rs b/libafl_libfuzzer/libafl_libfuzzer_runtime/src/report.rs index 7e4353ee98..557e8d67f6 100644 --- a/libafl_libfuzzer/libafl_libfuzzer_runtime/src/report.rs +++ b/libafl_libfuzzer/libafl_libfuzzer_runtime/src/report.rs @@ -1,13 +1,13 @@ use std::ffi::c_int; use libafl::{ - events::{ProgressReporter, SimpleEventManager}, + events::{EventProcessor, ProgressReporter, SimpleEventManager}, executors::HasObservers, feedbacks::MapFeedbackMetadata, inputs::UsesInput, monitors::SimpleMonitor, stages::{HasCurrentStage, StagesTuple}, - state::{HasExecutions, HasLastReportTime}, + state::{HasExecutions, HasLastReportTime, Stoppable}, Error, Fuzzer, HasMetadata, HasNamedMetadata, }; @@ -29,9 +29,10 @@ where + HasExecutions + UsesInput + HasLastReportTime - + HasCurrentStage, + + HasCurrentStage + + Stoppable, E: HasObservers, - EM: ProgressReporter, + EM: ProgressReporter + EventProcessor, ST: StagesTuple, { let meta = state diff --git a/libafl_nyx/src/executor.rs b/libafl_nyx/src/executor.rs index d9cc30de44..c62ab796de 100644 --- a/libafl_nyx/src/executor.rs +++ b/libafl_nyx/src/executor.rs @@ -121,18 +121,15 @@ where } }; - match self.stdout.as_mut() { - Some(ob) => { - let mut stdout = Vec::new(); - self.helper.nyx_stdout.rewind()?; - self.helper - .nyx_stdout - .read_to_end(&mut stdout) - .map_err(|e| Error::illegal_state(format!("Failed to read Nyx stdout: {e}")))?; + if let Some(ob) = self.stdout.as_mut() { + let mut stdout = Vec::new(); + self.helper.nyx_stdout.rewind()?; + self.helper + .nyx_stdout + .read_to_end(&mut stdout) + .map_err(|e| Error::illegal_state(format!("Failed to read Nyx stdout: {e}")))?; - ob.observe_stdout(&stdout); - } - None => (), + ob.observe_stdout(&stdout); } Ok(exit_kind)