Alternative scheduled count strategy (#1252)

* early return generalization stage * scheduled count * aaa * compile * fix * implement alternative scheduled count strategy --------- Co-authored-by: toka <[email protected]>
AFLplusplus · May 4, 2023 · a8e64be · a8e64be
1 parent 53659f8
commit a8e64be
Show file tree

Hide file tree

Showing 14 changed files with 63 additions and 108 deletions.
diff --git 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<State = Self::State>;
 
@@ -249,7 +252,7 @@ where
     CS: Scheduler,
     F: Feedback<CS::State>,
     OF: Feedback<CS::State>,
-    CS::State: HasClientPerfMonitor,
+    CS::State: HasClientPerfMonitor + HasCorpus,
 {
     scheduler: CS,
     feedback: F,
@@ -262,7 +265,7 @@ where
     CS: Scheduler,
     F: Feedback<CS::State>,
     OF: Feedback<CS::State>,
-    CS::State: HasClientPerfMonitor,
+    CS::State: HasClientPerfMonitor + HasCorpus,
 {
     type State = CS::State;
 }
@@ -272,7 +275,7 @@ where
     CS: Scheduler,
     F: Feedback<CS::State>,
     OF: Feedback<CS::State>,
-    CS::State: HasClientPerfMonitor,
+    CS::State: HasClientPerfMonitor + HasCorpus,
 {
     type Scheduler = CS;
 
@@ -290,7 +293,7 @@ where
     CS: Scheduler,
     F: Feedback<CS::State>,
     OF: Feedback<CS::State>,
-    CS::State: HasClientPerfMonitor,
+    CS::State: HasClientPerfMonitor + HasCorpus,
 {
     type Feedback = F;
 
@@ -308,7 +311,7 @@ where
     CS: Scheduler,
     F: Feedback<CS::State>,
     OF: Feedback<CS::State>,
-    CS::State: HasClientPerfMonitor,
+    CS::State: HasClientPerfMonitor + HasCorpus,
 {
     type Objective = OF;
 
@@ -327,7 +330,7 @@ where
     F: Feedback<CS::State>,
     OF: Feedback<CS::State>,
     OT: ObserversTuple<CS::State> + 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<EM>(
@@ -549,7 +552,7 @@ where
     EM: ProgressReporter + EventProcessor<E, Self, State = CS::State>,
     F: Feedback<CS::State>,
     OF: Feedback<CS::State>,
-    CS::State: HasClientPerfMonitor + HasExecutions + HasMetadata,
+    CS::State: HasClientPerfMonitor + HasExecutions + HasMetadata + HasCorpus + HasTestcase,
     ST: StagesTuple<E, EM, CS::State, Self>,
 {
     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<CS::State>,
     OF: Feedback<CS::State>,
-    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<CS::State>,
     E: Executor<EM, Self> + HasObservers<State = CS::State>,
     EM: UsesState<State = CS::State>,
-    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
@@ -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<CorpusId>,
-    ) -> 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
@@ -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<CorpusId>,
-    ) -> 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<S> {
 
 impl<S> UsesState for RandScheduler<S>
 where
-    S: UsesInput,
+    S: UsesInput + HasTestcase,
 {
     type State = S;
 }
 
 impl<S> Scheduler for RandScheduler<S>
 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<CorpusId>,
-    ) -> Result<(), Error> {
-        *state.corpus_mut().current_mut() = next_idx;
-        Ok(())
-    }
 }
 
 impl<S> RandScheduler<S> {

diff --git 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::<SchedulerTestcaseMetadata>()?;
 
             if tcmeta.handicap() >= 4 {

diff --git 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,15 +89,15 @@ where
 
 impl<F, S> UsesState for ProbabilitySamplingScheduler<F, S>
 where
-    S: UsesInput,
+    S: UsesInput + HasTestcase,
 {
     type State = S;
 }
 
 impl<F, S> Scheduler for ProbabilitySamplingScheduler<F, S>
 where
     F: TestcaseScore<S>,
-    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<CorpusId>,
-    ) -> Result<(), Error> {
-        *state.corpus_mut().current_mut() = next_idx;
-        Ok(())
-    }
 }
 
 impl<F, S> Default for ProbabilitySamplingScheduler<F, S>

diff --git 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<S> RemovableScheduler for QueueScheduler<S> where S: HasCorpus {}
+impl<S> RemovableScheduler for QueueScheduler<S> where S: HasCorpus + HasTestcase {}
 
 impl<S> Scheduler for QueueScheduler<S>
 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<CorpusId>,
-    ) -> Result<(), Error> {
-        *state.corpus_mut().current_mut() = next_idx;
-        Ok(())
-    }
 }
 
 impl<S> QueueScheduler<S> {

diff --git 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<S> RemovableScheduler for TuneableScheduler<S> where S: HasCorpus + HasMetadata {}
+impl<S> RemovableScheduler for TuneableScheduler<S> where S: HasCorpus + HasMetadata + HasTestcase {}
 
 impl<S> Scheduler for TuneableScheduler<S>
 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<CorpusId>,
-    ) -> Result<(), Error> {
-        *state.corpus_mut().current_mut() = next_idx;
-        Ok(())
-    }
 }
diff --git 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::<SchedulerTestcaseMetadata>()?;
 
             if tcmeta.handicap() >= 4 {

diff --git 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
@@ -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);