From fa69b9eff9b204f6723a7c4308ac1835f3151373 Mon Sep 17 00:00:00 2001
From: Dongjia Zhang <tokazerkje@outlook.com>
Date: Thu, 7 Apr 2022 21:00:59 +0900
Subject: [PATCH] Powerschedule::RAND (#596)

---
 fuzzers/fuzzbench_weighted/src/lib.rs   |  2 +-
 libafl/src/schedulers/powersched.rs     |  1 +
 libafl/src/schedulers/testcase_score.rs |  2 +-
 libafl/src/stages/power.rs              | 16 +++++++++++++---
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/fuzzers/fuzzbench_weighted/src/lib.rs b/fuzzers/fuzzbench_weighted/src/lib.rs
index 833dcf4a1c..06ee56f664 100644
--- a/fuzzers/fuzzbench_weighted/src/lib.rs
+++ b/fuzzers/fuzzbench_weighted/src/lib.rs
@@ -314,7 +314,7 @@ fn fuzz(
     )?;
 
     let power =
-        StdPowerMutationalStage::new(&mut state, mutator, &edges_observer, PowerSchedule::FAST);
+        StdPowerMutationalStage::new(&mut state, mutator, &edges_observer, PowerSchedule::RAND);
 
     // A minimization+queue policy to get testcasess from the corpus
     let scheduler = IndexesLenTimeMinimizerScheduler::new(StdWeightedScheduler::new());
diff --git a/libafl/src/schedulers/powersched.rs b/libafl/src/schedulers/powersched.rs
index bf6f274670..2e5ab68faf 100644
--- a/libafl/src/schedulers/powersched.rs
+++ b/libafl/src/schedulers/powersched.rs
@@ -132,6 +132,7 @@ impl PowerScheduleMetadata {
 #[allow(missing_docs)]
 #[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq)]
 pub enum PowerSchedule {
+    RAND,
     EXPLORE,
     EXPLOIT,
     FAST,
diff --git a/libafl/src/schedulers/testcase_score.rs b/libafl/src/schedulers/testcase_score.rs
index 90f1606e4e..412c397042 100644
--- a/libafl/src/schedulers/testcase_score.rs
+++ b/libafl/src/schedulers/testcase_score.rs
@@ -192,7 +192,7 @@ where
         // COE and Fast schedule are fairly different from what are described in the original thesis,
         // This implementation follows the changes made in this pull request https://github.com/AFLplusplus/AFLplusplus/pull/568
         match psmeta.strat() {
-            PowerSchedule::EXPLORE => {
+            PowerSchedule::EXPLORE | PowerSchedule::RAND => {
                 // Nothing happens in EXPLORE
             }
             PowerSchedule::EXPLOIT => {
diff --git a/libafl/src/stages/power.rs b/libafl/src/stages/power.rs
index 88bddf82b1..a97e116771 100644
--- a/libafl/src/stages/power.rs
+++ b/libafl/src/stages/power.rs
@@ -4,6 +4,7 @@ use alloc::string::{String, ToString};
 use core::{fmt::Debug, marker::PhantomData};
 
 use crate::{
+    bolts::rands::Rand,
     corpus::{Corpus, PowerScheduleTestcaseMetaData},
     executors::{Executor, HasObservers},
     fuzzer::Evaluator,
@@ -16,7 +17,7 @@ use crate::{
         TestcaseScore,
     },
     stages::{MutationalStage, Stage},
-    state::{HasClientPerfMonitor, HasCorpus, HasMetadata},
+    state::{HasClientPerfMonitor, HasCorpus, HasMetadata, HasRand},
     Error,
 };
 /// The mutational stage using power schedules
@@ -47,7 +48,7 @@ where
     M: Mutator<I, S>,
     O: MapObserver,
     OT: ObserversTuple<I, S>,
-    S: HasClientPerfMonitor + HasCorpus<I> + HasMetadata,
+    S: HasClientPerfMonitor + HasCorpus<I> + HasMetadata + HasRand,
     Z: Evaluator<E, EM, I, S>,
 {
     /// The mutator, added to this stage
@@ -66,6 +67,15 @@ where
     #[allow(clippy::cast_sign_loss)]
     fn iterations(&self, state: &mut S, corpus_idx: usize) -> Result<usize, Error> {
         // Update handicap
+        let use_random = state
+            .metadata_mut()
+            .get_mut::<PowerScheduleMetadata>()
+            .ok_or_else(|| Error::KeyNotFound("PowerScheduleMetadata not found".to_string()))?
+            .strat() == PowerSchedule::RAND;
+        if use_random {
+            return Ok(1 + state.rand_mut().below(128) as usize)
+        }
+
         let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut();
         let score = F::compute(&mut *testcase, state)? as usize;
         let tcmeta = testcase
@@ -151,7 +161,7 @@ where
     M: Mutator<I, S>,
     O: MapObserver,
     OT: ObserversTuple<I, S>,
-    S: HasClientPerfMonitor + HasCorpus<I> + HasMetadata,
+    S: HasClientPerfMonitor + HasCorpus<I> + HasMetadata + HasRand,
     Z: Evaluator<E, EM, I, S>,
 {
     #[inline]