From 283ceaac9b2b7fabe2d3ec9acdab3271252a3db3 Mon Sep 17 00:00:00 2001 From: Dongjia Zhang Date: Sun, 8 May 2022 23:43:02 +0900 Subject: [PATCH] Make weigthed scheduler independent of powersheduler stage (#599) * rename & add metadata in scheduler, not stage * Update testcase_score * rename * fix * update handicap in scheduler * fmt * update fuzzers * doc * fmt * fix * fmt * more * fix * fix * fix * fmt --- fuzzers/fuzzbench/src/lib.rs | 6 +- fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs | 6 +- fuzzers/fuzzbench_qemu/src/fuzzer.rs | 6 +- fuzzers/fuzzbench_text/src/lib.rs | 12 +- fuzzers/fuzzbench_weighted/src/lib.rs | 3 +- fuzzers/libfuzzer_libpng/src/lib.rs | 3 +- .../fuzzer/src/main.rs | 4 +- fuzzers/tutorial/src/lib.rs | 5 +- libafl/src/corpus/mod.rs | 2 +- libafl/src/corpus/testcase.rs | 8 +- libafl/src/schedulers/powersched.rs | 63 +++-- libafl/src/schedulers/testcase_score.rs | 254 ++++++++++-------- libafl/src/schedulers/weighted.rs | 26 +- libafl/src/stages/calibrate.rs | 19 +- libafl/src/stages/power.rs | 43 +-- 15 files changed, 235 insertions(+), 225 deletions(-) diff --git a/fuzzers/fuzzbench/src/lib.rs b/fuzzers/fuzzbench/src/lib.rs index 42049f52f1..4db8b93aca 100644 --- a/fuzzers/fuzzbench/src/lib.rs +++ b/fuzzers/fuzzbench/src/lib.rs @@ -313,11 +313,11 @@ fn fuzz( 5, )?; - let power = - StdPowerMutationalStage::new(&mut state, mutator, &edges_observer, PowerSchedule::FAST); + let power = StdPowerMutationalStage::new(mutator, &edges_observer); // A minimization+queue policy to get testcasess from the corpus - let scheduler = IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new()); + let scheduler = + IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new(PowerSchedule::FAST)); // A fuzzer with feedbacks and a corpus scheduler let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); diff --git a/fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs b/fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs index 5c1a5593f9..67eef8306f 100644 --- a/fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs +++ b/fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs @@ -284,11 +284,11 @@ fn fuzz( 5, )?; - let power = - StdPowerMutationalStage::new(&mut state, mutator, &edges_observer, PowerSchedule::FAST); + let power = StdPowerMutationalStage::new(mutator, &edges_observer); // A minimization+queue policy to get testcasess from the corpus - let scheduler = IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new()); + let scheduler = + IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new(PowerSchedule::FAST)); // A fuzzer with feedbacks and a corpus scheduler let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); diff --git a/fuzzers/fuzzbench_qemu/src/fuzzer.rs b/fuzzers/fuzzbench_qemu/src/fuzzer.rs index 6a4985397b..617b6f05d0 100644 --- a/fuzzers/fuzzbench_qemu/src/fuzzer.rs +++ b/fuzzers/fuzzbench_qemu/src/fuzzer.rs @@ -297,11 +297,11 @@ fn fuzz( 5, )?; - let power = - StdPowerMutationalStage::new(&mut state, mutator, &edges_observer, PowerSchedule::FAST); + let power = StdPowerMutationalStage::new(mutator, &edges_observer); // A minimization+queue policy to get testcasess from the corpus - let scheduler = IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new()); + let scheduler = + IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new(PowerSchedule::FAST)); // A fuzzer with feedbacks and a corpus scheduler let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); diff --git a/fuzzers/fuzzbench_text/src/lib.rs b/fuzzers/fuzzbench_text/src/lib.rs index 44ada277e7..0ed16796ea 100644 --- a/fuzzers/fuzzbench_text/src/lib.rs +++ b/fuzzers/fuzzbench_text/src/lib.rs @@ -374,11 +374,11 @@ fn fuzz_binary( 5, )?; - let power = - StdPowerMutationalStage::new(&mut state, mutator, &edges_observer, PowerSchedule::FAST); + let power = StdPowerMutationalStage::new(mutator, &edges_observer); // A minimization+queue policy to get testcasess from the corpus - let scheduler = IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new()); + let scheduler = + IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new(PowerSchedule::FAST)); // A fuzzer with feedbacks and a corpus scheduler let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); @@ -584,8 +584,7 @@ fn fuzz_text( 5, )?; - let power = - StdPowerMutationalStage::new(&mut state, mutator, &edges_observer, PowerSchedule::FAST); + let power = StdPowerMutationalStage::new(mutator, &edges_observer); let grimoire_mutator = StdScheduledMutator::with_max_stack_pow( tuple_list!( @@ -601,7 +600,8 @@ fn fuzz_text( let grimoire = StdMutationalStage::new(grimoire_mutator); // A minimization+queue policy to get testcasess from the corpus - let scheduler = IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new()); + let scheduler = + IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new(PowerSchedule::FAST)); // A fuzzer with feedbacks and a corpus scheduler let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); diff --git a/fuzzers/fuzzbench_weighted/src/lib.rs b/fuzzers/fuzzbench_weighted/src/lib.rs index 06ee56f664..c6ad9d0cb1 100644 --- a/fuzzers/fuzzbench_weighted/src/lib.rs +++ b/fuzzers/fuzzbench_weighted/src/lib.rs @@ -313,8 +313,7 @@ fn fuzz( 5, )?; - let power = - StdPowerMutationalStage::new(&mut state, mutator, &edges_observer, PowerSchedule::RAND); + let power = StdMutationalStage::new(mutator); // A minimization+queue policy to get testcasess from the corpus let scheduler = IndexesLenTimeMinimizerScheduler::new(StdWeightedScheduler::new()); diff --git a/fuzzers/libfuzzer_libpng/src/lib.rs b/fuzzers/libfuzzer_libpng/src/lib.rs index 6449697093..211b3999e9 100644 --- a/fuzzers/libfuzzer_libpng/src/lib.rs +++ b/fuzzers/libfuzzer_libpng/src/lib.rs @@ -132,8 +132,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations())); let calibration = CalibrationStage::new(&edges_observer); - let power = - StdPowerMutationalStage::new(&mut state, mutator, &edges_observer, PowerSchedule::COE); + let power = StdPowerMutationalStage::new(mutator, &edges_observer); let mut stages = tuple_list!(calibration, power); diff --git a/fuzzers/libfuzzer_stb_image_concolic/fuzzer/src/main.rs b/fuzzers/libfuzzer_stb_image_concolic/fuzzer/src/main.rs index 5bee33261e..b07962154e 100644 --- a/fuzzers/libfuzzer_stb_image_concolic/fuzzer/src/main.rs +++ b/fuzzers/libfuzzer_stb_image_concolic/fuzzer/src/main.rs @@ -40,7 +40,7 @@ use libafl::{ }, StdMapObserver, TimeObserver, }, - schedulers::{IndexesLenTimeMinimizerScheduler, PowerQueueScheduler}, + schedulers::{IndexesLenTimeMinimizerScheduler, QueueScheduler}, stages::{ ConcolicTracingStage, ShadowTracingStage, SimpleConcolicMutationalStage, StdMutationalStage, TracingStage, @@ -150,7 +150,7 @@ fn fuzz( println!("We're a client, let's fuzz :)"); // A minimization+queue policy to get testcasess from the corpus - let scheduler = IndexesLenTimeMinimizerScheduler::new(PowerQueueScheduler::new()); + let scheduler = IndexesLenTimeMinimizerScheduler::new(QueueScheduler::new()); // A fuzzer with feedbacks and a corpus scheduler let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); diff --git a/fuzzers/tutorial/src/lib.rs b/fuzzers/tutorial/src/lib.rs index 2c4a5e5fa6..ba9e049aef 100644 --- a/fuzzers/tutorial/src/lib.rs +++ b/fuzzers/tutorial/src/lib.rs @@ -126,13 +126,12 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re let mutator = LainMutator::new(); let calibration = CalibrationStage::new(&edges_observer); - let power = - StdPowerMutationalStage::new(&mut state, mutator, &edges_observer, PowerSchedule::FAST); + let power = StdPowerMutationalStage::new(mutator, &edges_observer); let mut stages = tuple_list!(calibration, power); // A minimization+queue policy to get testcasess from the corpus - let scheduler = PacketLenMinimizerScheduler::new(PowerQueueScheduler::new()); + let scheduler = PacketLenMinimizerScheduler::new(PowerQueueScheduler::new(PowerSchedule::FAST)); // A fuzzer with feedbacks and a corpus scheduler let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); diff --git a/libafl/src/corpus/mod.rs b/libafl/src/corpus/mod.rs index 4f13353122..1c7f94576d 100644 --- a/libafl/src/corpus/mod.rs +++ b/libafl/src/corpus/mod.rs @@ -1,7 +1,7 @@ //! Corpuses contain the testcases, either in memory, on disk, or somewhere else. pub mod testcase; -pub use testcase::{PowerScheduleTestcaseMetaData, Testcase}; +pub use testcase::{SchedulerTestcaseMetaData, Testcase}; pub mod inmemory; pub use inmemory::InMemoryCorpus; diff --git a/libafl/src/corpus/testcase.rs b/libafl/src/corpus/testcase.rs index 07ec0c3d50..2d7a82feb5 100644 --- a/libafl/src/corpus/testcase.rs +++ b/libafl/src/corpus/testcase.rs @@ -276,7 +276,7 @@ where /// The Metadata for each testcase used in power schedules. #[derive(Serialize, Deserialize, Clone, Debug)] -pub struct PowerScheduleTestcaseMetaData { +pub struct SchedulerTestcaseMetaData { /// Number of bits set in bitmap, updated in calibrate_case bitmap_size: u64, /// Number of queue cycles behind @@ -287,8 +287,8 @@ pub struct PowerScheduleTestcaseMetaData { n_fuzz_entry: usize, } -impl PowerScheduleTestcaseMetaData { - /// Create new [`struct@PowerScheduleTestcaseMetaData`] +impl SchedulerTestcaseMetaData { + /// Create new [`struct@SchedulerTestcaseMetaData`] #[must_use] pub fn new(depth: u64) -> Self { Self { @@ -344,4 +344,4 @@ impl PowerScheduleTestcaseMetaData { } } -crate::impl_serdeany!(PowerScheduleTestcaseMetaData); +crate::impl_serdeany!(SchedulerTestcaseMetaData); diff --git a/libafl/src/schedulers/powersched.rs b/libafl/src/schedulers/powersched.rs index a984ee9ac7..b39a9d2b34 100644 --- a/libafl/src/schedulers/powersched.rs +++ b/libafl/src/schedulers/powersched.rs @@ -6,7 +6,7 @@ use alloc::{ }; use crate::{ - corpus::{Corpus, PowerScheduleTestcaseMetaData}, + corpus::{Corpus, SchedulerTestcaseMetaData}, inputs::Input, schedulers::Scheduler, state::{HasCorpus, HasMetadata}, @@ -17,13 +17,13 @@ use serde::{Deserialize, Serialize}; /// The n fuzz size pub const N_FUZZ_SIZE: usize = 1 << 21; -crate::impl_serdeany!(PowerScheduleMetadata); +crate::impl_serdeany!(SchedulerMetadata); /// The metadata used for power schedules #[derive(Serialize, Deserialize, Clone, Debug)] -pub struct PowerScheduleMetadata { +pub struct SchedulerMetadata { /// Powerschedule strategy - strat: PowerSchedule, + strat: Option, /// Measured exec time during calibration exec_time: Duration, /// Calibration cycles @@ -39,10 +39,10 @@ pub struct PowerScheduleMetadata { } /// The metadata for runs in the calibration stage. -impl PowerScheduleMetadata { - /// Creates a new [`struct@PowerScheduleMetadata`] +impl SchedulerMetadata { + /// Creates a new [`struct@SchedulerMetadata`] #[must_use] - pub fn new(strat: PowerSchedule) -> Self { + pub fn new(strat: Option) -> Self { Self { strat, exec_time: Duration::from_millis(0), @@ -56,7 +56,7 @@ impl PowerScheduleMetadata { /// The powerschedule strategy #[must_use] - pub fn strat(&self) -> PowerSchedule { + pub fn strat(&self) -> Option { self.strat } @@ -132,7 +132,6 @@ impl PowerScheduleMetadata { #[allow(missing_docs)] #[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq)] pub enum PowerSchedule { - RAND, EXPLORE, EXPLOIT, FAST, @@ -143,12 +142,8 @@ pub enum PowerSchedule { /// A corpus scheduler using power schedules #[derive(Clone, Debug)] -pub struct PowerQueueScheduler; - -impl Default for PowerQueueScheduler { - fn default() -> Self { - Self::new() - } +pub struct PowerQueueScheduler { + strat: PowerSchedule, } impl Scheduler for PowerQueueScheduler @@ -158,6 +153,10 @@ 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::() { + state.add_metadata::(SchedulerMetadata::new(Some(self.strat))); + } + let current_idx = *state.corpus().current(); let mut depth = match current_idx { @@ -166,19 +165,21 @@ where .get(parent_idx)? .borrow_mut() .metadata_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("PowerScheduleTestData not found".to_string()))? + .get_mut::() + .ok_or_else(|| { + Error::key_not_found("SchedulerTestcaseMetaData not found".to_string()) + })? .depth(), None => 0, }; - // Attach a `PowerScheduleTestData` to the queue entry. + // Attach a `SchedulerTestcaseMetaData` to the queue entry. depth += 1; state .corpus() .get(idx)? .borrow_mut() - .add_metadata(PowerScheduleTestcaseMetaData::new(depth)); + .add_metadata(SchedulerTestcaseMetaData::new(depth)); Ok(()) } @@ -191,9 +192,9 @@ where if *cur + 1 >= state.corpus().count() { let psmeta = state .metadata_mut() - .get_mut::() + .get_mut::() .ok_or_else(|| { - Error::key_not_found("PowerScheduleMetadata not found".to_string()) + Error::key_not_found("SchedulerMetadata not found".to_string()) })?; psmeta.set_queue_cycles(psmeta.queue_cycles() + 1); 0 @@ -204,6 +205,22 @@ where None => 0, }; *state.corpus_mut().current_mut() = Some(id); + + // Update the handicap + let mut testcase = state.corpus().get(id)?.borrow_mut(); + let tcmeta = testcase + .metadata_mut() + .get_mut::() + .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(id) } } @@ -212,7 +229,7 @@ where impl PowerQueueScheduler { /// Create a new [`PowerQueueScheduler`] #[must_use] - pub fn new() -> Self { - Self + pub fn new(strat: PowerSchedule) -> Self { + PowerQueueScheduler { strat } } } diff --git a/libafl/src/schedulers/testcase_score.rs b/libafl/src/schedulers/testcase_score.rs index 4c70d1db06..63ddda944a 100644 --- a/libafl/src/schedulers/testcase_score.rs +++ b/libafl/src/schedulers/testcase_score.rs @@ -1,12 +1,12 @@ //! The `TestcaseScore` is an evaluator providing scores of corpus items. use crate::{ bolts::{HasLen, HasRefCnt}, - corpus::{Corpus, PowerScheduleTestcaseMetaData, Testcase}, + corpus::{Corpus, SchedulerTestcaseMetaData, Testcase}, feedbacks::MapIndexesMetadata, inputs::Input, schedulers::{ minimizer::{IsFavoredMetadata, TopRatedsMetadata}, - powersched::{PowerSchedule, PowerScheduleMetadata}, + powersched::{PowerSchedule, SchedulerMetadata}, }, state::{HasCorpus, HasMetadata}, Error, @@ -79,44 +79,52 @@ where fn compute(entry: &mut Testcase, state: &S) -> Result { let psmeta = state .metadata() - .get::() - .ok_or_else(|| Error::key_not_found("PowerScheduleMetadata not found".to_string()))?; + .get::() + .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; - let fuzz_mu = if psmeta.strat() == PowerSchedule::COE { - let corpus = state.corpus(); - let mut n_paths = 0; - let mut v = 0.0; - let cur_index = state.corpus().current().unwrap(); - for idx in 0..corpus.count() { - let n_fuzz_entry = if cur_index == idx { - entry - .metadata() - .get::() - .ok_or_else(|| { - Error::key_not_found("PowerScheduleTestData not found".to_string()) - })? - .n_fuzz_entry() - } else { - corpus - .get(idx)? - .borrow() - .metadata() - .get::() - .ok_or_else(|| { - Error::key_not_found("PowerScheduleTestData not found".to_string()) - })? - .n_fuzz_entry() - }; - v += libm::log2(f64::from(psmeta.n_fuzz()[n_fuzz_entry])); - n_paths += 1; + let fuzz_mu = if let Some(strat) = psmeta.strat() { + if strat == PowerSchedule::COE { + let corpus = state.corpus(); + let mut n_paths = 0; + let mut v = 0.0; + let cur_index = state.corpus().current().unwrap(); + for idx in 0..corpus.count() { + let n_fuzz_entry = if cur_index == idx { + entry + .metadata() + .get::() + .ok_or_else(|| { + Error::key_not_found( + "SchedulerTestcaseMetaData not found".to_string(), + ) + })? + .n_fuzz_entry() + } else { + corpus + .get(idx)? + .borrow() + .metadata() + .get::() + .ok_or_else(|| { + Error::key_not_found( + "SchedulerTestcaseMetaData not found".to_string(), + ) + })? + .n_fuzz_entry() + }; + v += libm::log2(f64::from(psmeta.n_fuzz()[n_fuzz_entry])); + n_paths += 1; + } + + if n_paths == 0 { + return Err(Error::unknown(String::from("Queue state corrput"))); + } + + v /= f64::from(n_paths); + v + } else { + 0.0 } - - if n_paths == 0 { - return Err(Error::unknown(String::from("Queue state corrput"))); - } - - v /= f64::from(n_paths); - v } else { 0.0 }; @@ -133,9 +141,9 @@ where let favored = entry.has_metadata::(); let tcmeta = entry .metadata() - .get::() + .get::() .ok_or_else(|| { - Error::key_not_found("PowerScheduleTestcaseMetaData not found".to_string()) + Error::key_not_found("SchedulerTestcaseMetaData not found".to_string()) })?; if q_exec_us * 0.1 > avg_exec_us { @@ -191,81 +199,87 @@ 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::RAND => { - // Nothing happens in EXPLORE - } - PowerSchedule::EXPLOIT => { - factor = MAX_FACTOR; - } - PowerSchedule::COE => { - if libm::log2(f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()])) > fuzz_mu - && !favored - { - // Never skip favorites. - factor = 0.0; + if let Some(strat) = psmeta.strat() { + match strat { + PowerSchedule::EXPLORE => { + // Nothing happens in EXPLORE } - } - PowerSchedule::FAST => { - if entry.fuzz_level() != 0 { - let lg = libm::log2(f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()])); - - match lg { - f if f < 2.0 => { - factor = 4.0; - } - f if (2.0..4.0).contains(&f) => { - factor = 3.0; - } - f if (4.0..5.0).contains(&f) => { - factor = 2.0; - } - f if (6.0..7.0).contains(&f) => { - if !favored { - factor = 0.8; - } - } - f if (7.0..8.0).contains(&f) => { - if !favored { - factor = 0.6; - } - } - f if f >= 8.0 => { - if !favored { - factor = 0.4; - } - } - _ => { - factor = 1.0; - } - } - - if favored { - factor *= 1.15; + PowerSchedule::EXPLOIT => { + factor = MAX_FACTOR; + } + PowerSchedule::COE => { + if libm::log2(f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()])) > fuzz_mu + && !favored + { + // Never skip favorites. + factor = 0.0; } } - } - PowerSchedule::LIN => { - factor = (entry.fuzz_level() as f64) - / f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()] + 1); - } - PowerSchedule::QUAD => { - factor = ((entry.fuzz_level() * entry.fuzz_level()) as f64) - / f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()] + 1); + PowerSchedule::FAST => { + if entry.fuzz_level() != 0 { + let lg = libm::log2(f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()])); + + match lg { + f if f < 2.0 => { + factor = 4.0; + } + f if (2.0..4.0).contains(&f) => { + factor = 3.0; + } + f if (4.0..5.0).contains(&f) => { + factor = 2.0; + } + f if (6.0..7.0).contains(&f) => { + if !favored { + factor = 0.8; + } + } + f if (7.0..8.0).contains(&f) => { + if !favored { + factor = 0.6; + } + } + f if f >= 8.0 => { + if !favored { + factor = 0.4; + } + } + _ => { + factor = 1.0; + } + } + + if favored { + factor *= 1.15; + } + } + } + PowerSchedule::LIN => { + factor = (entry.fuzz_level() as f64) + / f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()] + 1); + } + PowerSchedule::QUAD => { + factor = ((entry.fuzz_level() * entry.fuzz_level()) as f64) + / f64::from(psmeta.n_fuzz()[tcmeta.n_fuzz_entry()] + 1); + } } } - if psmeta.strat() != PowerSchedule::EXPLORE { - if factor > MAX_FACTOR { - factor = MAX_FACTOR; - } + if let Some(strat) = psmeta.strat() { + if strat == PowerSchedule::EXPLORE { + if factor > MAX_FACTOR { + factor = MAX_FACTOR; + } - perf_score *= factor / POWER_BETA; + perf_score *= factor / POWER_BETA; + } } // Lower bound if the strat is not COE. - if psmeta.strat() == PowerSchedule::COE && perf_score < 1.0 { - perf_score = 1.0; + if let Some(strat) = psmeta.strat() { + if strat == PowerSchedule::COE && perf_score < 1.0 { + perf_score = 1.0; + } } // Upper bound @@ -299,16 +313,19 @@ where let mut weight = 1.0; let psmeta = state .metadata() - .get::() - .ok_or_else(|| Error::key_not_found("PowerScheduleMetadata not found".to_string()))?; + .get::() + .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; let tcmeta = entry .metadata() - .get::() - .ok_or_else(|| Error::key_not_found("PowerScheduleTestData not found".to_string()))?; + .get::() + .ok_or_else(|| { + Error::key_not_found("SchedulerTestcaseMetaData not found".to_string()) + })?; // This means that this testcase has never gone through the calibration stage before1, // In this case we'll just return the default weight + // This methoud is called in corpus's on_add() method. Fuzz_level is zero at that time. if entry.fuzz_level() == 0 || psmeta.cycles() == 0 { return Ok(weight); } @@ -324,15 +341,20 @@ where let q_bitmap_size = tcmeta.bitmap_size() as f64; - match psmeta.strat() { - PowerSchedule::FAST | PowerSchedule::COE | PowerSchedule::LIN | PowerSchedule::QUAD => { - let hits = psmeta.n_fuzz()[tcmeta.n_fuzz_entry()]; - if hits > 0 { - weight *= libm::log10(f64::from(hits)) + 1.0; + if let Some(strat) = psmeta.strat() { + match strat { + PowerSchedule::FAST + | PowerSchedule::COE + | PowerSchedule::LIN + | PowerSchedule::QUAD => { + let hits = psmeta.n_fuzz()[tcmeta.n_fuzz_entry()]; + if hits > 0 { + weight *= libm::log10(f64::from(hits)) + 1.0; + } } + // EXPLORE and EXPLOIT fall into this + _ => {} } - // EXPLORE and EXPLOIT fall into this - _ => {} } weight *= avg_exec_us / q_exec_us; diff --git a/libafl/src/schedulers/weighted.rs b/libafl/src/schedulers/weighted.rs index bc12f94945..b98811aeff 100644 --- a/libafl/src/schedulers/weighted.rs +++ b/libafl/src/schedulers/weighted.rs @@ -1,5 +1,5 @@ //! The queue corpus scheduler with weighted queue item selection from aflpp (`https://github.com/AFLplusplus/AFLplusplus/blob/1d4f1e48797c064ee71441ba555b29fc3f467983/src/afl-fuzz-queue.c#L32`) -//! This queue corpus scheduler needs calibration stage and the power schedule stage. +//! This queue corpus scheduler needs calibration stage. use alloc::{ string::{String, ToString}, @@ -8,10 +8,10 @@ use alloc::{ use crate::{ bolts::rands::Rand, - corpus::{Corpus, PowerScheduleTestcaseMetaData}, + corpus::{Corpus, SchedulerTestcaseMetaData}, inputs::Input, schedulers::{ - powersched::PowerScheduleMetadata, + powersched::SchedulerMetadata, testcase_score::{CorpusWeightTestcaseScore, TestcaseScore}, Scheduler, }, @@ -215,6 +215,10 @@ 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::() { + state.add_metadata(SchedulerMetadata::new(None)); + } + if !state.has_metadata::() { state.add_metadata(WeightedScheduleMetadata::new()); } @@ -227,21 +231,23 @@ where .get(parent_idx)? .borrow_mut() .metadata_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("PowerScheduleTestData not found".to_string()))? + .get_mut::() + .ok_or_else(|| { + Error::key_not_found("SchedulerTestcaseMetaData not found".to_string()) + })? .depth(), None => 0, }; - // Attach a `PowerScheduleTestData` to the queue entry. + // Attach a `SchedulerTestcaseMetaData` to the queue entry. depth += 1; state .corpus() .get(idx)? .borrow_mut() - .add_metadata(PowerScheduleTestcaseMetaData::new(depth)); + .add_metadata(SchedulerTestcaseMetaData::new(depth)); - // Recrate the alias table + // Recreate the alias table self.create_alias_table(state)?; Ok(()) } @@ -281,9 +287,9 @@ where if current_cycles > corpus_counts { let psmeta = state .metadata_mut() - .get_mut::() + .get_mut::() .ok_or_else(|| { - Error::key_not_found("PowerScheduleMetadata not found".to_string()) + Error::key_not_found("SchedulerMetadata not found".to_string()) })?; psmeta.set_queue_cycles(psmeta.queue_cycles() + 1); } diff --git a/libafl/src/stages/calibrate.rs b/libafl/src/stages/calibrate.rs index 4812bd9a04..bfa77e7ea1 100644 --- a/libafl/src/stages/calibrate.rs +++ b/libafl/src/stages/calibrate.rs @@ -3,14 +3,14 @@ use crate::{ bolts::current_time, bolts::tuples::MatchName, - corpus::{Corpus, PowerScheduleTestcaseMetaData}, + corpus::{Corpus, SchedulerTestcaseMetaData}, events::{EventFirer, LogSeverity}, executors::{Executor, ExitKind, HasObservers}, feedbacks::MapFeedbackState, fuzzer::Evaluator, inputs::Input, observers::{MapObserver, ObserversTuple}, - schedulers::powersched::PowerScheduleMetadata, + schedulers::powersched::SchedulerMetadata, stages::Stage, state::{HasClientPerfMonitor, HasCorpus, HasFeedbackStates, HasMetadata}, Error, @@ -160,13 +160,13 @@ where } }; - // If power schedule is used, update it - let use_powerschedule = state.has_metadata::() + // If weighted scheduler or powerscheduler is used, update it + let use_powerschedule = state.has_metadata::() && state .corpus() .get(corpus_idx)? .borrow() - .has_metadata::(); + .has_metadata::(); if use_powerschedule { let map = executor @@ -176,10 +176,7 @@ where let bitmap_size = map.count_bytes(); - let psmeta = state - .metadata_mut() - .get_mut::() - .unwrap(); + let psmeta = state.metadata_mut().get_mut::().unwrap(); let handicap = psmeta.queue_cycles(); psmeta.set_exec_time(psmeta.exec_time() + total_time); @@ -196,9 +193,9 @@ where let data = testcase .metadata_mut() - .get_mut::() + .get_mut::() .ok_or_else(|| { - Error::key_not_found("PowerScheduleTestData not found".to_string()) + Error::key_not_found("SchedulerTestcaseMetaData not found".to_string()) })?; data.set_bitmap_size(bitmap_size); diff --git a/libafl/src/stages/power.rs b/libafl/src/stages/power.rs index 484d524832..4b0fa8ad5b 100644 --- a/libafl/src/stages/power.rs +++ b/libafl/src/stages/power.rs @@ -4,17 +4,14 @@ use alloc::string::{String, ToString}; use core::{fmt::Debug, marker::PhantomData}; use crate::{ - bolts::rands::Rand, - corpus::{Corpus, PowerScheduleTestcaseMetaData}, + corpus::{Corpus, SchedulerTestcaseMetaData}, executors::{Executor, HasObservers}, fuzzer::Evaluator, inputs::Input, mutators::Mutator, observers::{MapObserver, ObserversTuple}, schedulers::{ - powersched::{PowerSchedule, PowerScheduleMetadata}, - testcase_score::CorpusPowerTestcaseScore, - TestcaseScore, + powersched::SchedulerMetadata, testcase_score::CorpusPowerTestcaseScore, TestcaseScore, }, stages::{MutationalStage, Stage}, state::{HasClientPerfMonitor, HasCorpus, HasMetadata, HasRand}, @@ -67,29 +64,8 @@ where #[allow(clippy::cast_sign_loss)] fn iterations(&self, state: &mut S, corpus_idx: usize) -> Result { // Update handicap - let use_random = state - .metadata_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("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 - .metadata_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("PowerScheduleTestcaseMetaData 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(score) } @@ -126,10 +102,8 @@ where let psmeta = state .metadata_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("PowerScheduleMetadata not found".to_string()) - })?; + .get_mut::() + .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; hash %= psmeta.n_fuzz().len(); // Update the path frequency @@ -141,9 +115,9 @@ where .get(idx)? .borrow_mut() .metadata_mut() - .get_mut::() + .get_mut::() .ok_or_else(|| { - Error::key_not_found("PowerScheduleTestData not found".to_string()) + Error::key_not_found("SchedulerTestcaseMetaData not found".to_string()) })? .set_n_fuzz_entry(hash); } @@ -194,10 +168,7 @@ where Z: Evaluator, { /// Creates a new [`PowerMutationalStage`] - pub fn new(state: &mut S, mutator: M, map_observer_name: &O, strat: PowerSchedule) -> Self { - if !state.has_metadata::() { - state.add_metadata::(PowerScheduleMetadata::new(strat)); - } + pub fn new(mutator: M, map_observer_name: &O) -> Self { Self { map_observer_name: map_observer_name.name().to_string(), mutator,