From 159da0d311065594ab59810cf4bfae018d40dadc Mon Sep 17 00:00:00 2001 From: "Dongjia \"toka\" Zhang" Date: Sun, 7 Apr 2024 17:24:27 +0200 Subject: [PATCH] Change when weighted scheduler make new table & don't update global metadata in on_remove/on_replace (#2011) * fix * backtick --- libafl/src/schedulers/mod.rs | 74 ----------------------------- libafl/src/schedulers/powersched.rs | 31 +++++------- libafl/src/schedulers/weighted.rs | 39 ++++++++------- 3 files changed, 33 insertions(+), 111 deletions(-) diff --git a/libafl/src/schedulers/mod.rs b/libafl/src/schedulers/mod.rs index e3664e8b35..192077cd4e 100644 --- a/libafl/src/schedulers/mod.rs +++ b/libafl/src/schedulers/mod.rs @@ -65,80 +65,6 @@ where } } -/// Define the metadata operations when removing testcase from AFL-style scheduler -pub trait HasAFLRemovableScheduler: RemovableScheduler -where - Self::State: HasCorpus + HasMetadata + HasTestcase, -{ - #[allow(clippy::cast_precision_loss)] - #[allow(clippy::cast_precision_loss)] - /// Adjusting metadata when removing the testcase - fn on_remove_metadata( - &mut self, - state: &mut Self::State, - _idx: CorpusId, - prev: &Option::Input>>, - ) -> Result<(), Error> { - let prev = prev.as_ref().ok_or_else(|| { - Error::illegal_argument( - "Power schedulers must be aware of the removed corpus entry for reweighting.", - ) - })?; - - let prev_meta = prev.metadata::()?; - - // Use these to adjust `SchedulerMetadata` - let (prev_total_time, prev_cycles) = prev_meta.cycle_and_time(); - let prev_bitmap_size = prev_meta.bitmap_size(); - let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64); - - let psmeta = state.metadata_mut::()?; - - psmeta.set_exec_time(psmeta.exec_time() - prev_total_time); - psmeta.set_cycles(psmeta.cycles() - (prev_cycles as u64)); - psmeta.set_bitmap_size(psmeta.bitmap_size() - prev_bitmap_size); - psmeta.set_bitmap_size_log(psmeta.bitmap_size_log() - prev_bitmap_size_log); - psmeta.set_bitmap_entries(psmeta.bitmap_entries() - 1); - - Ok(()) - } - - #[allow(clippy::cast_precision_loss)] - /// Adjusting metadata when replacing the corpus - fn on_replace_metadata( - &mut self, - state: &mut Self::State, - idx: CorpusId, - prev: &Testcase<::Input>, - ) -> Result<(), Error> { - let prev_meta = prev.metadata::()?; - - // Next depth is + 1 - let prev_depth = prev_meta.depth() + 1; - - // Use these to adjust `SchedulerMetadata` - let (prev_total_time, prev_cycles) = prev_meta.cycle_and_time(); - let prev_bitmap_size = prev_meta.bitmap_size(); - let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64); - - let psmeta = state.metadata_mut::()?; - - // We won't add new one because it'll get added when it gets executed in calirbation next time. - psmeta.set_exec_time(psmeta.exec_time() - prev_total_time); - psmeta.set_cycles(psmeta.cycles() - (prev_cycles as u64)); - psmeta.set_bitmap_size(psmeta.bitmap_size() - prev_bitmap_size); - psmeta.set_bitmap_size_log(psmeta.bitmap_size_log() - prev_bitmap_size_log); - psmeta.set_bitmap_entries(psmeta.bitmap_entries() - 1); - - state - .corpus() - .get(idx)? - .borrow_mut() - .add_metadata(SchedulerTestcaseMetadata::new(prev_depth)); - Ok(()) - } -} - /// Defines the common metadata operations for the AFL-style schedulers pub trait HasAFLSchedulerMetadata: Scheduler where diff --git a/libafl/src/schedulers/powersched.rs b/libafl/src/schedulers/powersched.rs index 8eade59576..da5142a741 100644 --- a/libafl/src/schedulers/powersched.rs +++ b/libafl/src/schedulers/powersched.rs @@ -12,9 +12,7 @@ use crate::{ corpus::{Corpus, CorpusId, HasTestcase, Testcase}, inputs::UsesInput, observers::{MapObserver, ObserversTuple}, - schedulers::{ - HasAFLRemovableScheduler, HasAFLSchedulerMetadata, RemovableScheduler, Scheduler, - }, + schedulers::{HasAFLSchedulerMetadata, RemovableScheduler, Scheduler}, state::{HasCorpus, HasMetadata, State, UsesState}, Error, }; @@ -186,34 +184,29 @@ where type State = S; } -impl HasAFLRemovableScheduler for PowerQueueScheduler +impl RemovableScheduler for PowerQueueScheduler where S: State + HasTestcase + HasMetadata + HasCorpus, O: MapObserver, { -} - -impl RemovableScheduler for PowerQueueScheduler -where - S: HasCorpus + HasMetadata + HasTestcase + State, - O: MapObserver, -{ + /// This will *NOT* neutralize the effect of this removed testcase from the global data such as `SchedulerMetadata` fn on_remove( &mut self, - state: &mut Self::State, - idx: CorpusId, - prev: &Option::Input>>, + _state: &mut Self::State, + _idx: CorpusId, + _prev: &Option::Input>>, ) -> Result<(), Error> { - self.on_remove_metadata(state, idx, prev) + Ok(()) } + /// This will *NOT* neutralize the effect of this removed testcase from the global data such as `SchedulerMetadata` fn on_replace( &mut self, - state: &mut Self::State, - idx: CorpusId, - prev: &Testcase<::Input>, + _state: &mut Self::State, + _idx: CorpusId, + _prev: &Testcase<::Input>, ) -> Result<(), Error> { - self.on_replace_metadata(state, idx, prev) + Ok(()) } } diff --git a/libafl/src/schedulers/weighted.rs b/libafl/src/schedulers/weighted.rs index 7a26628f76..9c4227894a 100644 --- a/libafl/src/schedulers/weighted.rs +++ b/libafl/src/schedulers/weighted.rs @@ -16,7 +16,7 @@ use crate::{ schedulers::{ powersched::{PowerSchedule, SchedulerMetadata}, testcase_score::{CorpusWeightTestcaseScore, TestcaseScore}, - HasAFLRemovableScheduler, HasAFLSchedulerMetadata, RemovableScheduler, Scheduler, + HasAFLSchedulerMetadata, RemovableScheduler, Scheduler, }, state::{HasCorpus, HasMetadata, HasRand, State, UsesState}, Error, @@ -93,6 +93,7 @@ libafl_bolts::impl_serdeany!(WeightedScheduleMetadata); /// A corpus scheduler using power schedules with weighted queue item selection algo. #[derive(Clone, Debug)] pub struct WeightedScheduler { + table_invalidated: bool, strat: Option, map_observer_name: String, last_hash: usize, @@ -121,6 +122,7 @@ where strat, map_observer_name: map_observer.name().to_string(), last_hash: 0, + table_invalidated: true, phantom: PhantomData, } } @@ -223,36 +225,32 @@ where type State = S; } -impl HasAFLRemovableScheduler for WeightedScheduler -where - F: TestcaseScore, - S: State + HasTestcase + HasMetadata + HasCorpus + HasRand, - O: MapObserver, -{ -} - impl RemovableScheduler for WeightedScheduler where F: TestcaseScore, O: MapObserver, S: HasCorpus + HasMetadata + HasRand + HasTestcase + State, { + /// This will *NOT* neutralize the effect of this removed testcase from the global data such as `SchedulerMetadata` fn on_remove( &mut self, - state: &mut Self::State, - idx: CorpusId, - prev: &Option::Input>>, + _state: &mut Self::State, + _idx: CorpusId, + _prev: &Option::Input>>, ) -> Result<(), Error> { - self.on_remove_metadata(state, idx, prev) + self.table_invalidated = true; + Ok(()) } + /// This will *NOT* neutralize the effect of this removed testcase from the global data such as `SchedulerMetadata` fn on_replace( &mut self, - state: &mut Self::State, - idx: CorpusId, - prev: &Testcase<::Input>, + _state: &mut Self::State, + _idx: CorpusId, + _prev: &Testcase<::Input>, ) -> Result<(), Error> { - self.on_replace_metadata(state, idx, prev) + self.table_invalidated = true; + Ok(()) } } @@ -284,7 +282,8 @@ where /// Called when a [`Testcase`] is added to the corpus fn on_add(&mut self, state: &mut S, idx: CorpusId) -> Result<(), Error> { self.on_add_metadata(state, idx)?; - self.create_alias_table(state) + self.table_invalidated = true; + Ok(()) } fn on_evaluation( @@ -301,6 +300,10 @@ where #[allow(clippy::similar_names, clippy::cast_precision_loss)] fn next(&mut self, state: &mut S) -> Result { + if self.table_invalidated { + self.create_alias_table(state)?; + self.table_invalidated = false; + } let corpus_counts = state.corpus().count(); if corpus_counts == 0 { Err(Error::empty(String::from(