SYS-OSS/FRET-LibAFL
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Change when weighted scheduler make new table & don't update global metadata in on_remove/on_replace (#2011)

Browse Source 
* fix

* backtick
This commit is contained in:
Dongjia "toka" Zhang 2024-04-07 17:24:27 +02:00 committed by GitHub
parent 527b892c1d
commit 159da0d311
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 33 additions and 111 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

74
libafl/src/schedulers/mod.rs
View File

@ -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<Testcase<<Self::State as UsesInput>::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::<SchedulerTestcaseMetadata>()?;
// 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::<SchedulerMetadata>()?;
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<<Self::State as UsesInput>::Input>,
) -> Result<(), Error> {
let prev_meta = prev.metadata::<SchedulerTestcaseMetadata>()?;
// 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::<SchedulerMetadata>()?;
// 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<O, S>: Scheduler
where

31
libafl/src/schedulers/powersched.rs
View File

@ -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<O, S> HasAFLRemovableScheduler for PowerQueueScheduler<O, S>
impl<O, S> RemovableScheduler for PowerQueueScheduler<O, S>
where
S: State + HasTestcase + HasMetadata + HasCorpus,
O: MapObserver,
{
}
impl<O, S> RemovableScheduler for PowerQueueScheduler<O, S>
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<Testcase<<Self::State as UsesInput>::Input>>,
_state: &mut Self::State,
_idx: CorpusId,
_prev: &Option<Testcase<<Self::State as UsesInput>::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<<Self::State as UsesInput>::Input>,
_state: &mut Self::State,
_idx: CorpusId,
_prev: &Testcase<<Self::State as UsesInput>::Input>,
) -> Result<(), Error> {
self.on_replace_metadata(state, idx, prev)
Ok(())
}
}

39
libafl/src/schedulers/weighted.rs
View File

@ -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<F, O, S> {
table_invalidated: bool,
strat: Option<PowerSchedule>,
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<F, O, S> HasAFLRemovableScheduler for WeightedScheduler<F, O, S>
where
F: TestcaseScore<S>,
S: State + HasTestcase + HasMetadata + HasCorpus + HasRand,
O: MapObserver,
{
}
impl<F, O, S> RemovableScheduler for WeightedScheduler<F, O, S>
where
F: TestcaseScore<S>,
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<Testcase<<Self::State as UsesInput>::Input>>,
_state: &mut Self::State,
_idx: CorpusId,
_prev: &Option<Testcase<<Self::State as UsesInput>::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<<Self::State as UsesInput>::Input>,
_state: &mut Self::State,
_idx: CorpusId,
_prev: &Testcase<<Self::State as UsesInput>::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<OT>(
@ -301,6 +300,10 @@ where
#[allow(clippy::similar_names, clippy::cast_precision_loss)]
fn next(&mut self, state: &mut S) -> Result<CorpusId, Error> {
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(