From 66b5fe867840c321cc54c1a0ef82b5ae6f1149a6 Mon Sep 17 00:00:00 2001
From: Dongjia Zhang <tokazerkje@outlook.com>
Date: Wed, 29 Jun 2022 09:43:42 +0900
Subject: [PATCH] Extend weighted scheduler (#685)

* extend

* fix

* fmt

* more fix
---
 fuzzers/fuzzbench_weighted/src/lib.rs |  5 ++--
 fuzzers/libfuzzer_libpng/src/lib.rs   |  7 ++++--
 libafl/src/schedulers/weighted.rs     | 34 ++++++++++++++++-----------
 3 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/fuzzers/fuzzbench_weighted/src/lib.rs b/fuzzers/fuzzbench_weighted/src/lib.rs
index deae68abec..4bc9d768f5 100644
--- a/fuzzers/fuzzbench_weighted/src/lib.rs
+++ b/fuzzers/fuzzbench_weighted/src/lib.rs
@@ -315,10 +315,11 @@ fn fuzz(
         5,
     )?;
 
-    let power = StdMutationalStage::new(mutator);
+    let power = StdPowerMutationalStage::new(mutator, &edges_observer);
 
     // A minimization+queue policy to get testcasess from the corpus
-    let scheduler = IndexesLenTimeMinimizerScheduler::new(StdWeightedScheduler::new());
+    let scheduler =
+        IndexesLenTimeMinimizerScheduler::new(StdWeightedScheduler::new(Some(PowerSchedule::FAST)));
 
     // A fuzzer with feedbacks and a corpus scheduler
     let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective);
diff --git a/fuzzers/libfuzzer_libpng/src/lib.rs b/fuzzers/libfuzzer_libpng/src/lib.rs
index ab9bceaf10..234018c54a 100644
--- a/fuzzers/libfuzzer_libpng/src/lib.rs
+++ b/fuzzers/libfuzzer_libpng/src/lib.rs
@@ -27,7 +27,9 @@ use libafl::{
     mutators::scheduled::{havoc_mutations, tokens_mutations, StdScheduledMutator},
     mutators::token_mutations::Tokens,
     observers::{HitcountsMapObserver, StdMapObserver, TimeObserver},
-    schedulers::{IndexesLenTimeMinimizerScheduler, StdWeightedScheduler},
+    schedulers::{
+        powersched::PowerSchedule, IndexesLenTimeMinimizerScheduler, StdWeightedScheduler,
+    },
     stages::{calibrate::CalibrationStage, power::StdPowerMutationalStage},
     state::{HasCorpus, HasMetadata, StdState},
     Error,
@@ -139,7 +141,8 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
     let mut stages = tuple_list!(calibration, power);
 
     // A minimization+queue policy to get testcasess from the corpus
-    let scheduler = IndexesLenTimeMinimizerScheduler::new(StdWeightedScheduler::new());
+    let scheduler =
+        IndexesLenTimeMinimizerScheduler::new(StdWeightedScheduler::new(Some(PowerSchedule::FAST)));
 
     // A fuzzer with feedbacks and a corpus scheduler
     let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective);
diff --git a/libafl/src/schedulers/weighted.rs b/libafl/src/schedulers/weighted.rs
index b98811aeff..76ae1ceb8b 100644
--- a/libafl/src/schedulers/weighted.rs
+++ b/libafl/src/schedulers/weighted.rs
@@ -11,7 +11,7 @@ use crate::{
     corpus::{Corpus, SchedulerTestcaseMetaData},
     inputs::Input,
     schedulers::{
-        powersched::SchedulerMetadata,
+        powersched::{PowerSchedule, SchedulerMetadata},
         testcase_score::{CorpusWeightTestcaseScore, TestcaseScore},
         Scheduler,
     },
@@ -89,20 +89,10 @@ crate::impl_serdeany!(WeightedScheduleMetadata);
 /// A corpus scheduler using power schedules with weighted queue item selection algo.
 #[derive(Clone, Debug)]
 pub struct WeightedScheduler<F, I, S> {
+    strat: Option<PowerSchedule>,
     phantom: PhantomData<(F, I, S)>,
 }
 
-impl<F, I, S> Default for WeightedScheduler<F, I, S>
-where
-    F: TestcaseScore<I, S>,
-    I: Input,
-    S: HasCorpus<I> + HasMetadata + HasRand,
-{
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 impl<F, I, S> WeightedScheduler<F, I, S>
 where
     F: TestcaseScore<I, S>,
@@ -111,8 +101,9 @@ where
 {
     /// Create a new [`WeightedScheduler`]
     #[must_use]
-    pub fn new() -> Self {
+    pub fn new(strat: Option<PowerSchedule>) -> Self {
         Self {
+            strat,
             phantom: PhantomData,
         }
     }
@@ -216,7 +207,7 @@ where
     /// Add an entry to the corpus and return its index
     fn on_add(&self, state: &mut S, idx: usize) -> Result<(), Error> {
         if !state.has_metadata::<SchedulerMetadata>() {
-            state.add_metadata(SchedulerMetadata::new(None));
+            state.add_metadata(SchedulerMetadata::new(self.strat));
         }
 
         if !state.has_metadata::<WeightedScheduleMetadata>() {
@@ -294,6 +285,21 @@ where
                 psmeta.set_queue_cycles(psmeta.queue_cycles() + 1);
             }
             *state.corpus_mut().current_mut() = Some(idx);
+
+            // Update the handicap
+            let mut testcase = state.corpus().get(idx)?.borrow_mut();
+            let tcmeta = testcase
+                .metadata_mut()
+                .get_mut::<SchedulerTestcaseMetaData>()
+                .ok_or_else(|| {
+                    Error::key_not_found("SchedulerTestcaseMetaData not found".to_string())
+                })?;
+
+            if tcmeta.handicap() >= 4 {
+                tcmeta.set_handicap(tcmeta.handicap() - 4);
+            } else if tcmeta.handicap() > 0 {
+                tcmeta.set_handicap(tcmeta.handicap() - 1);
+            }
             Ok(idx)
         }
     }