diff --git a/libafl/src/schedulers/powersched.rs b/libafl/src/schedulers/powersched.rs index 0601e06104..ff0afed54c 100644 --- a/libafl/src/schedulers/powersched.rs +++ b/libafl/src/schedulers/powersched.rs @@ -65,12 +65,17 @@ impl SchedulerMetadata { } } - /// The powerschedule strategy + /// The `PowerSchedule` #[must_use] pub fn strat(&self) -> Option { self.strat } + /// Set the `PowerSchedule` + pub fn set_strat(&mut self, strat: Option) { + self.strat = strat; + } + /// The measured exec time during calibration #[must_use] pub fn exec_time(&self) -> Duration { diff --git a/libafl/src/schedulers/weighted.rs b/libafl/src/schedulers/weighted.rs index 7cbeff9ab1..369bcf7177 100644 --- a/libafl/src/schedulers/weighted.rs +++ b/libafl/src/schedulers/weighted.rs @@ -101,6 +101,8 @@ pub struct WeightedScheduler { map_observer_handle: Handle, last_hash: usize, phantom: PhantomData<(F, O, S)>, + /// Cycle `PowerSchedule` on completion of every queue cycle. + cycle_schedules: bool, } impl WeightedScheduler @@ -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 { @@ -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 { + 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 UsesState for WeightedScheduler @@ -343,6 +374,9 @@ where if runs_in_current_cycle >= corpus_counts { let psmeta = state.metadata_mut::()?; psmeta.set_queue_cycles(psmeta.queue_cycles() + 1); + if self.cycle_schedules { + self.cycle_schedule(psmeta)?; + } } self.set_current_scheduled(state, Some(idx))?;