add experimental corpus pruning

This commit is contained in:
Alwin Berger 2024-10-29 18:08:18 +01:00
parent 4d0ec2e427
commit d88cefb894

View File

@ -4,8 +4,10 @@ use core::{marker::PhantomData, time::Duration};
use libafl_bolts::current_time; use libafl_bolts::current_time;
use itertools::Itertools;
use libafl::{ use libafl::{
corpus::{Corpus, HasCurrentCorpusId}, events::EventFirer, prelude::minimizer::TopRatedsMetadata, schedulers::minimizer::IsFavoredMetadata, stages::Stage, state::{HasCorpus, HasImported, UsesState}, Error, HasMetadata corpus::{Corpus, HasCurrentCorpusId}, events::EventFirer, prelude::{minimizer::TopRatedsMetadata, RemovableScheduler}, schedulers::minimizer::IsFavoredMetadata, stages::Stage, state::{HasCorpus, HasImported, UsesState}, Error, HasMetadata, HasScheduler
}; };
use libafl::{ use libafl::{
events::Event, events::Event,
@ -35,12 +37,13 @@ impl<E, EM, Z> Stage<E, EM, Z> for SchedulerStatsStage<E, EM, Z>
where where
E: UsesState, E: UsesState,
EM: EventFirer<State = Self::State>, EM: EventFirer<State = Self::State>,
Z: UsesState<State = Self::State>, Z: UsesState<State = Self::State> + HasScheduler,
<Z as HasScheduler>::Scheduler: RemovableScheduler,
Self::State: HasImported + HasCorpus + HasMetadata, Self::State: HasImported + HasCorpus + HasMetadata,
{ {
fn perform( fn perform(
&mut self, &mut self,
_fuzzer: &mut Z, fuzzer: &mut Z,
_executor: &mut E, _executor: &mut E,
state: &mut <Self as UsesState>::State, state: &mut <Self as UsesState>::State,
_manager: &mut EM, _manager: &mut EM,
@ -59,7 +62,7 @@ where
if cur.checked_sub(self.last_report_time).unwrap_or_default() > self.stats_report_interval { if cur.checked_sub(self.last_report_time).unwrap_or_default() > self.stats_report_interval {
if let Some(meta) = state.metadata_map().get::<TopRatedsMetadata>() { if let Some(meta) = state.metadata_map().get::<TopRatedsMetadata>() {
let kc = meta.map.keys().count(); let kc = meta.map.keys().count();
let mut v : Vec<_> = meta.map.values().collect(); let mut v : Vec<_> = meta.map.values().cloned().collect();
v.sort_unstable(); v.sort_unstable();
v.dedup(); v.dedup();
let vc = v.len(); let vc = v.len();
@ -90,6 +93,31 @@ where
self.imported_size self.imported_size
); );
self.last_report_time = cur; self.last_report_time = cur;
// Experimental pruning
#[cfg(any(feature = "sched_stg",feature = "sched_afl"))]
{
const PRUNE_THRESHOLD: usize = 200;
if state.corpus().count() > PRUNE_THRESHOLD*vc {
// println!("Pruning corpus, keeping {} / {}", PRUNE_THRESHOLD*vc);
let corpus = state.corpus_mut();
let currid = corpus.current();
let ids : Vec<_> = corpus.ids().filter_map(|x| {
let tc = corpus.get(x).unwrap().borrow();
let md = tc.metadata_map();
if md.get::<IsFavoredMetadata>().is_some() || &Some(x) == currid || v.contains(&&x) {
None
} else {
Some((x, tc.exec_time().clone()))
}
}).sorted_by_key(|x| x.1).take(usize::saturating_sub(corpus.count(),(PRUNE_THRESHOLD/2)*vc)).sorted_by_key(|x| x.0).unique().rev().collect();
for (cid, _) in ids {
let c = state.corpus_mut().remove(cid).unwrap();
fuzzer
.scheduler_mut()
.on_remove(state, cid, &Some(c))?;
}
}
}
} }
} }