From a8e64be1695f2321d7f101b71cb2073c21ac7f2c Mon Sep 17 00:00:00 2001 From: Addison Crump Date: Thu, 4 May 2023 13:15:28 +0200 Subject: [PATCH] Alternative scheduled count strategy (#1252) * early return generalization stage * scheduled count * aaa * compile * fix * implement alternative scheduled count strategy --------- Co-authored-by: toka --- libafl/src/fuzzer/mod.rs | 33 ++++++++++++------- libafl/src/schedulers/ecofuzz.rs | 20 ----------- libafl/src/schedulers/mod.rs | 31 +++++++++-------- libafl/src/schedulers/powersched.rs | 5 --- .../src/schedulers/probabilistic_sampling.rs | 16 ++------- libafl/src/schedulers/queue.rs | 16 ++------- libafl/src/schedulers/tuneable.rs | 16 ++------- libafl/src/schedulers/weighted.rs | 5 --- libafl/src/stages/calibrate.rs | 2 -- libafl/src/stages/generalization.rs | 4 +++ libafl/src/stages/mod.rs | 4 +-- libafl/src/stages/push/mod.rs | 12 +++---- libafl/src/stages/push/mutational.rs | 2 +- libafl/src/stages/tmin.rs | 5 ++- 14 files changed, 63 insertions(+), 108 deletions(-) diff --git a/libafl/src/fuzzer/mod.rs b/libafl/src/fuzzer/mod.rs index 0e8c231984..6ea268dbb6 100644 --- a/libafl/src/fuzzer/mod.rs +++ b/libafl/src/fuzzer/mod.rs @@ -13,7 +13,7 @@ use crate::monitors::PerfFeature; use crate::state::NopState; use crate::{ bolts::current_time, - corpus::{Corpus, CorpusId, Testcase}, + corpus::{Corpus, CorpusId, HasTestcase, Testcase}, events::{Event, EventConfig, EventFirer, EventProcessor, ProgressReporter}, executors::{Executor, ExitKind, HasObservers}, feedbacks::Feedback, @@ -31,7 +31,10 @@ use crate::{ const STATS_TIMEOUT_DEFAULT: Duration = Duration::from_secs(15); /// Holds a scheduler -pub trait HasScheduler: UsesState { +pub trait HasScheduler: UsesState +where + Self::State: HasCorpus, +{ /// The [`Scheduler`] for this fuzzer type Scheduler: Scheduler; @@ -249,7 +252,7 @@ where CS: Scheduler, F: Feedback, OF: Feedback, - CS::State: HasClientPerfMonitor, + CS::State: HasClientPerfMonitor + HasCorpus, { scheduler: CS, feedback: F, @@ -262,7 +265,7 @@ where CS: Scheduler, F: Feedback, OF: Feedback, - CS::State: HasClientPerfMonitor, + CS::State: HasClientPerfMonitor + HasCorpus, { type State = CS::State; } @@ -272,7 +275,7 @@ where CS: Scheduler, F: Feedback, OF: Feedback, - CS::State: HasClientPerfMonitor, + CS::State: HasClientPerfMonitor + HasCorpus, { type Scheduler = CS; @@ -290,7 +293,7 @@ where CS: Scheduler, F: Feedback, OF: Feedback, - CS::State: HasClientPerfMonitor, + CS::State: HasClientPerfMonitor + HasCorpus, { type Feedback = F; @@ -308,7 +311,7 @@ where CS: Scheduler, F: Feedback, OF: Feedback, - CS::State: HasClientPerfMonitor, + CS::State: HasClientPerfMonitor + HasCorpus, { type Objective = OF; @@ -327,7 +330,7 @@ where F: Feedback, OF: Feedback, OT: ObserversTuple + Serialize + DeserializeOwned, - CS::State: HasCorpus + HasSolutions + HasClientPerfMonitor + HasExecutions, + CS::State: HasCorpus + HasSolutions + HasClientPerfMonitor + HasExecutions + HasCorpus, { /// Evaluate if a set of observation channels has an interesting state fn process_execution( @@ -549,7 +552,7 @@ where EM: ProgressReporter + EventProcessor, F: Feedback, OF: Feedback, - CS::State: HasClientPerfMonitor + HasExecutions + HasMetadata, + CS::State: HasClientPerfMonitor + HasExecutions + HasMetadata + HasCorpus + HasTestcase, ST: StagesTuple, { fn fuzz_one( @@ -588,6 +591,14 @@ where #[cfg(feature = "introspection")] state.introspection_monitor_mut().mark_manager_time(); + { + let mut testcase = state.testcase_mut(idx)?; + let scheduled_count = testcase.scheduled_count(); + + // increase scheduled count, this was fuzz_level in afl + testcase.set_scheduled_count(scheduled_count + 1); + } + Ok(idx) } } @@ -597,7 +608,7 @@ where CS: Scheduler, F: Feedback, OF: Feedback, - CS::State: UsesInput + HasExecutions + HasClientPerfMonitor, + CS::State: UsesInput + HasExecutions + HasClientPerfMonitor + HasCorpus, { /// Create a new `StdFuzzer` with standard behavior. pub fn new(scheduler: CS, feedback: F, objective: OF) -> Self { @@ -665,7 +676,7 @@ where OF: Feedback, E: Executor + HasObservers, EM: UsesState, - CS::State: UsesInput + HasExecutions + HasClientPerfMonitor, + CS::State: UsesInput + HasExecutions + HasClientPerfMonitor + HasCorpus, { /// Runs the input and triggers observers and feedback fn execute_input( diff --git a/libafl/src/schedulers/ecofuzz.rs b/libafl/src/schedulers/ecofuzz.rs index 1a05475894..aac6084592 100644 --- a/libafl/src/schedulers/ecofuzz.rs +++ b/libafl/src/schedulers/ecofuzz.rs @@ -334,26 +334,6 @@ where Ok(id) } - - /// Set current fuzzed corpus id and `scheduled_count` - fn set_current_scheduled( - &mut self, - state: &mut Self::State, - next_idx: Option, - ) -> Result<(), Error> { - let current_idx = *state.corpus().current(); - - if let Some(idx) = current_idx { - let mut testcase = state.testcase_mut(idx)?; - let scheduled_count = testcase.scheduled_count(); - - // increase scheduled count, this was fuzz_level in afl - testcase.set_scheduled_count(scheduled_count + 1); - } - - *state.corpus_mut().current_mut() = next_idx; - Ok(()) - } } /// The weight for each corpus entry diff --git a/libafl/src/schedulers/mod.rs b/libafl/src/schedulers/mod.rs index ba86dd6c29..63c6733f93 100644 --- a/libafl/src/schedulers/mod.rs +++ b/libafl/src/schedulers/mod.rs @@ -34,7 +34,7 @@ pub use tuneable::*; use crate::{ bolts::rands::Rand, - corpus::{Corpus, CorpusId, Testcase}, + corpus::{Corpus, CorpusId, HasTestcase, Testcase}, inputs::UsesInput, observers::ObserversTuple, random_corpus_id, @@ -43,7 +43,10 @@ use crate::{ }; /// The scheduler also implemnts `on_remove` and `on_replace` if it implements this stage. -pub trait RemovableScheduler: Scheduler { +pub trait RemovableScheduler: Scheduler +where + Self::State: HasCorpus, +{ /// Removed the given entry from the corpus at the given index fn on_remove( &mut self, @@ -67,7 +70,10 @@ pub trait RemovableScheduler: Scheduler { /// The scheduler define how the fuzzer requests a testcase from the corpus. /// It has hooks to corpus add/replace/remove to allow complex scheduling algorithms to collect data. -pub trait Scheduler: UsesState { +pub trait Scheduler: UsesState +where + Self::State: HasCorpus, +{ /// Added an entry to the corpus at the given index fn on_add(&mut self, _state: &mut Self::State, _idx: CorpusId) -> Result<(), Error>; // Add parent_id here if it has no inner @@ -94,7 +100,10 @@ pub trait Scheduler: UsesState { &mut self, state: &mut Self::State, next_idx: Option, - ) -> Result<(), Error>; + ) -> Result<(), Error> { + *state.corpus_mut().current_mut() = next_idx; + Ok(()) + } } /// Feed the fuzzer simply with a random testcase on request @@ -105,14 +114,14 @@ pub struct RandScheduler { impl UsesState for RandScheduler where - S: UsesInput, + S: UsesInput + HasTestcase, { type State = S; } impl Scheduler for RandScheduler where - S: HasCorpus + HasRand, + S: HasCorpus + HasRand + HasTestcase, { fn on_add(&mut self, state: &mut Self::State, idx: CorpusId) -> Result<(), Error> { // Set parent id @@ -136,16 +145,6 @@ where Ok(id) } } - - /// Set current fuzzed corpus id and `scheduled_count`. You should call this from `next` - fn set_current_scheduled( - &mut self, - state: &mut Self::State, - next_idx: Option, - ) -> Result<(), Error> { - *state.corpus_mut().current_mut() = next_idx; - Ok(()) - } } impl RandScheduler { diff --git a/libafl/src/schedulers/powersched.rs b/libafl/src/schedulers/powersched.rs index 2e6be4aa01..0dad9b3230 100644 --- a/libafl/src/schedulers/powersched.rs +++ b/libafl/src/schedulers/powersched.rs @@ -336,11 +336,6 @@ where if let Some(idx) = current_idx { let mut testcase = state.testcase_mut(idx)?; - let scheduled_count = testcase.scheduled_count(); - - // increase scheduled count, this was fuzz_level in afl - testcase.set_scheduled_count(scheduled_count + 1); - let tcmeta = testcase.metadata_mut::()?; if tcmeta.handicap() >= 4 { diff --git a/libafl/src/schedulers/probabilistic_sampling.rs b/libafl/src/schedulers/probabilistic_sampling.rs index 63030c94fe..66adc6447c 100644 --- a/libafl/src/schedulers/probabilistic_sampling.rs +++ b/libafl/src/schedulers/probabilistic_sampling.rs @@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize}; use crate::{ bolts::rands::Rand, - corpus::{Corpus, CorpusId}, + corpus::{Corpus, CorpusId, HasTestcase}, inputs::UsesInput, schedulers::{Scheduler, TestcaseScore}, state::{HasCorpus, HasMetadata, HasRand, UsesState}, @@ -89,7 +89,7 @@ where impl UsesState for ProbabilitySamplingScheduler where - S: UsesInput, + S: UsesInput + HasTestcase, { type State = S; } @@ -97,7 +97,7 @@ where impl Scheduler for ProbabilitySamplingScheduler where F: TestcaseScore, - S: HasCorpus + HasMetadata + HasRand, + S: HasCorpus + HasMetadata + HasRand + HasTestcase, { fn on_add(&mut self, state: &mut Self::State, idx: CorpusId) -> Result<(), Error> { let current_idx = *state.corpus().current(); @@ -135,16 +135,6 @@ where Ok(ret) } } - - /// Set current fuzzed corpus id and `scheduled_count` - fn set_current_scheduled( - &mut self, - state: &mut Self::State, - next_idx: Option, - ) -> Result<(), Error> { - *state.corpus_mut().current_mut() = next_idx; - Ok(()) - } } impl Default for ProbabilitySamplingScheduler diff --git a/libafl/src/schedulers/queue.rs b/libafl/src/schedulers/queue.rs index 9e2ebaee0e..62ae1e48f9 100644 --- a/libafl/src/schedulers/queue.rs +++ b/libafl/src/schedulers/queue.rs @@ -4,7 +4,7 @@ use alloc::borrow::ToOwned; use core::marker::PhantomData; use crate::{ - corpus::{Corpus, CorpusId}, + corpus::{Corpus, CorpusId, HasTestcase}, inputs::UsesInput, schedulers::{RemovableScheduler, Scheduler}, state::{HasCorpus, UsesState}, @@ -24,11 +24,11 @@ where type State = S; } -impl RemovableScheduler for QueueScheduler where S: HasCorpus {} +impl RemovableScheduler for QueueScheduler where S: HasCorpus + HasTestcase {} impl Scheduler for QueueScheduler where - S: HasCorpus, + S: HasCorpus + HasTestcase, { fn on_add(&mut self, state: &mut Self::State, idx: CorpusId) -> Result<(), Error> { // Set parent id @@ -57,16 +57,6 @@ where Ok(id) } } - - /// Set current fuzzed corpus id and `scheduled_count` - fn set_current_scheduled( - &mut self, - state: &mut Self::State, - next_idx: Option, - ) -> Result<(), Error> { - *state.corpus_mut().current_mut() = next_idx; - Ok(()) - } } impl QueueScheduler { diff --git a/libafl/src/schedulers/tuneable.rs b/libafl/src/schedulers/tuneable.rs index 37e75985ac..004b9c7b65 100644 --- a/libafl/src/schedulers/tuneable.rs +++ b/libafl/src/schedulers/tuneable.rs @@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize}; use super::RemovableScheduler; use crate::{ - corpus::{Corpus, CorpusId}, + corpus::{Corpus, CorpusId, HasTestcase}, impl_serdeany, inputs::UsesInput, schedulers::Scheduler, @@ -92,11 +92,11 @@ where type State = S; } -impl RemovableScheduler for TuneableScheduler where S: HasCorpus + HasMetadata {} +impl RemovableScheduler for TuneableScheduler where S: HasCorpus + HasMetadata + HasTestcase {} impl Scheduler for TuneableScheduler where - S: HasCorpus + HasMetadata, + S: HasCorpus + HasMetadata + HasTestcase, { fn on_add(&mut self, state: &mut Self::State, idx: CorpusId) -> Result<(), Error> { // Set parent id @@ -126,14 +126,4 @@ where self.set_current_scheduled(state, Some(id))?; Ok(id) } - - /// Set current fuzzed corpus id and `scheduled_count` - fn set_current_scheduled( - &mut self, - state: &mut Self::State, - next_idx: Option, - ) -> Result<(), Error> { - *state.corpus_mut().current_mut() = next_idx; - Ok(()) - } } diff --git a/libafl/src/schedulers/weighted.rs b/libafl/src/schedulers/weighted.rs index ed242016ea..d80aecbf75 100644 --- a/libafl/src/schedulers/weighted.rs +++ b/libafl/src/schedulers/weighted.rs @@ -408,11 +408,6 @@ where if let Some(idx) = current_idx { let mut testcase = state.testcase_mut(idx)?; - let scheduled_count = testcase.scheduled_count(); - - // increase scheduled count, this was fuzz_level in afl - testcase.set_scheduled_count(scheduled_count + 1); - let tcmeta = testcase.metadata_mut::()?; if tcmeta.handicap() >= 4 { diff --git a/libafl/src/stages/calibrate.rs b/libafl/src/stages/calibrate.rs index dbd3932aa4..84b863fbba 100644 --- a/libafl/src/stages/calibrate.rs +++ b/libafl/src/stages/calibrate.rs @@ -257,10 +257,8 @@ where psmeta.set_bitmap_entries(psmeta.bitmap_entries() + 1); let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut(); - let scheduled_count = testcase.scheduled_count(); testcase.set_exec_time(total_time / (iter as u32)); - testcase.set_scheduled_count(scheduled_count + 1); // log::trace!("time: {:#?}", testcase.exec_time()); // If the testcase doesn't have its own `SchedulerTestcaseMetadata`, create it. diff --git a/libafl/src/stages/generalization.rs b/libafl/src/stages/generalization.rs index 3332d2b2eb..8ab9d53aea 100644 --- a/libafl/src/stages/generalization.rs +++ b/libafl/src/stages/generalization.rs @@ -82,6 +82,10 @@ where { let corpus = state.corpus(); let mut testcase = corpus.get(corpus_idx)?.borrow_mut(); + if testcase.scheduled_count() > 0 { + return Ok(()); + } + corpus.load_input_into(&mut testcase)?; } mark_feature_time!(state, PerfFeature::GetInputFromCorpus); diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index bec7eb5c07..5e82155b5a 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -66,7 +66,7 @@ use crate::{ inputs::UsesInput, observers::ObserversTuple, schedulers::Scheduler, - state::{HasClientPerfMonitor, HasExecutions, HasMetadata, HasRand, UsesState}, + state::{HasClientPerfMonitor, HasCorpus, HasExecutions, HasMetadata, HasRand, UsesState}, Error, EvaluatorObservers, ExecutesInput, ExecutionProcessor, HasScheduler, }; @@ -248,7 +248,7 @@ where impl Stage for PushStageAdapter where CS: Scheduler, - CS::State: HasClientPerfMonitor + HasExecutions + HasMetadata + HasRand, + CS::State: HasClientPerfMonitor + HasExecutions + HasMetadata + HasRand + HasCorpus, E: Executor + HasObservers, EM: EventFirer + EventRestarter diff --git a/libafl/src/stages/push/mod.rs b/libafl/src/stages/push/mod.rs index bb0ce56030..87838681fb 100644 --- a/libafl/src/stages/push/mod.rs +++ b/libafl/src/stages/push/mod.rs @@ -23,7 +23,7 @@ use crate::{ inputs::UsesInput, observers::ObserversTuple, schedulers::Scheduler, - state::{HasClientPerfMonitor, HasExecutions, HasMetadata, HasRand}, + state::{HasClientPerfMonitor, HasCorpus, HasExecutions, HasMetadata, HasRand}, Error, EvaluatorObservers, ExecutionProcessor, HasScheduler, }; @@ -38,7 +38,7 @@ where CS: Scheduler, EM: EventFirer + EventRestarter + HasEventManagerId, OT: ObserversTuple, - CS::State: HasClientPerfMonitor + HasRand, + CS::State: HasClientPerfMonitor + HasRand + HasCorpus, Z: ExecutionProcessor + EvaluatorObservers + HasScheduler, @@ -59,7 +59,7 @@ where CS: Scheduler, EM: EventFirer + EventRestarter + HasEventManagerId, OT: ObserversTuple, - CS::State: HasClientPerfMonitor + HasRand, + CS::State: HasClientPerfMonitor + HasRand + HasCorpus, Z: ExecutionProcessor + EvaluatorObservers + HasScheduler, @@ -84,7 +84,7 @@ where CS: Scheduler, EM: EventFirer + EventRestarter + HasEventManagerId, OT: ObserversTuple, - CS::State: HasClientPerfMonitor + HasRand, + CS::State: HasClientPerfMonitor + HasRand + HasCorpus, Z: ExecutionProcessor + EvaluatorObservers + HasScheduler, @@ -116,7 +116,7 @@ where CS: Scheduler, EM: EventFirer + EventRestarter + HasEventManagerId, OT: ObserversTuple, - CS::State: HasClientPerfMonitor + HasRand, + CS::State: HasClientPerfMonitor + HasRand + HasCorpus, Z: ExecutionProcessor + EvaluatorObservers + HasScheduler, @@ -184,7 +184,7 @@ where pub trait PushStage: Iterator where CS: Scheduler, - CS::State: HasClientPerfMonitor + HasRand + HasExecutions + HasMetadata, + CS::State: HasClientPerfMonitor + HasRand + HasExecutions + HasMetadata + HasCorpus, EM: EventFirer + EventRestarter + HasEventManagerId + ProgressReporter, OT: ObserversTuple, Z: ExecutionProcessor diff --git a/libafl/src/stages/push/mutational.rs b/libafl/src/stages/push/mutational.rs index 592432f703..c759260f9c 100644 --- a/libafl/src/stages/push/mutational.rs +++ b/libafl/src/stages/push/mutational.rs @@ -43,7 +43,7 @@ where EM: EventFirer + EventRestarter + HasEventManagerId, M: Mutator, OT: ObserversTuple, - CS::State: HasClientPerfMonitor + HasRand + Clone + Debug, + CS::State: HasClientPerfMonitor + HasRand + HasCorpus + Clone + Debug, Z: ExecutionProcessor + EvaluatorObservers + HasScheduler, diff --git a/libafl/src/stages/tmin.rs b/libafl/src/stages/tmin.rs index d2c25f4931..9a2ce48940 100644 --- a/libafl/src/stages/tmin.rs +++ b/libafl/src/stages/tmin.rs @@ -192,6 +192,7 @@ where CS: Scheduler, M: Mutator, Z: ExecutionProcessor, + CS::State: HasCorpus, { type State = CS::State; } @@ -200,7 +201,8 @@ impl Stage for StdTMinMutationalStage where CS: Scheduler + RemovableScheduler, - CS::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize + HasClientPerfMonitor, + CS::State: + HasCorpus + HasSolutions + HasExecutions + HasMaxSize + HasClientPerfMonitor + HasCorpus, ::Input: HasLen + Hash, E: Executor + HasObservers, EM: EventFirer, @@ -285,6 +287,7 @@ where CS: Scheduler, M: Mutator, Z: ExecutionProcessor, + CS::State: HasCorpus, { /// Creates a new minimising mutational stage that will minimize provided corpus entries pub fn new(mutator: M, factory: FF, runs: usize) -> Self {