Add support for cycling PowerSchedule on completion of a queue cycle …

…in WeightedScheduler (AFLplusplus#2300)

* add support for cycling PowerSchedule on completion of a queue cycle in WeightedScheduler

* improve doc

* make fn cycle_schedule private

* rename cycle_schedules to cycling_scheduler

libafl/src/schedulers/powersched.rs

@@ -65,12 +65,17 @@ impl SchedulerMetadata {
         }
     }
 
-    /// The powerschedule strategy
+    /// The `PowerSchedule`
     #[must_use]
     pub fn strat(&self) -> Option<PowerSchedule> {
         self.strat
     }
 
+    /// Set the `PowerSchedule`
+    pub fn set_strat(&mut self, strat: Option<PowerSchedule>) {
+        self.strat = strat;
+    }
+
     /// The measured exec time during calibration
     #[must_use]
     pub fn exec_time(&self) -> Duration {

libafl/src/schedulers/weighted.rs

@@ -101,6 +101,8 @@ pub struct WeightedScheduler<C, F, O, S> {
     map_observer_handle: Handle<C>,
     last_hash: usize,
     phantom: PhantomData<(F, O, S)>,
+    /// Cycle `PowerSchedule` on completion of every queue cycle.
+    cycle_schedules: bool,
 }
 
 impl<C, F, O, S> WeightedScheduler<C, F, O, S>
@@ -127,10 +129,18 @@ where
             map_observer_handle: map_observer.handle(),
             last_hash: 0,
             table_invalidated: true,
+            cycle_schedules: false,
             phantom: PhantomData,
         }
     }
 
+    /// Cycle the `PowerSchedule` on completion of a queue cycle
+    #[must_use]
+    pub fn cycling_scheduler(mut self) -> Self {
+        self.cycle_schedules = true;
+        self
+    }
+
     #[must_use]
     /// Getter for `strat`
     pub fn strat(&self) -> &Option<PowerSchedule> {
@@ -220,6 +230,27 @@ where
         wsmeta.set_alias_table(alias_table);
         Ok(())
     }
+
+    /// Cycles the strategy of the scheduler; tries to mimic AFL++'s cycling formula
+    fn cycle_schedule(
+        &mut self,
+        metadata: &mut SchedulerMetadata,
+    ) -> Result<PowerSchedule, Error> {
+        let next_strat = match metadata.strat().ok_or(Error::illegal_argument(
+            "No strategy specified when initializing scheduler; cannot cycle!",
+        ))? {
+            PowerSchedule::EXPLORE => PowerSchedule::EXPLOIT,
+            PowerSchedule::COE => PowerSchedule::LIN,
+            PowerSchedule::LIN => PowerSchedule::QUAD,
+            PowerSchedule::FAST => PowerSchedule::COE,
+            PowerSchedule::QUAD => PowerSchedule::FAST,
+            PowerSchedule::EXPLOIT => PowerSchedule::EXPLORE,
+        };
+        metadata.set_strat(Some(next_strat));
+        // We need to recalculate the scores of testcases.
+        self.table_invalidated = true;
+        Ok(next_strat)
+    }
 }
 
 impl<C, F, O, S> UsesState for WeightedScheduler<C, F, O, S>
@@ -343,6 +374,9 @@ where
             if runs_in_current_cycle >= corpus_counts {
                 let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
                 psmeta.set_queue_cycles(psmeta.queue_cycles() + 1);
+                if self.cycle_schedules {
+                    self.cycle_schedule(psmeta)?;
+                }
             }
 
             self.set_current_scheduled(state, Some(idx))?;