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

Use the new metadata() function in more places (#1155)

Browse Source
This commit is contained in:
Matheus Baptistella 2023-03-17 12:03:57 -03:00 committed by GitHub
parent fd95560512
commit 104c170ade
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 92 additions and 278 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
libafl/src/mutators/gramatron.rs
View File

@ -13,7 +13,7 @@ use crate::{
inputs::{GramatronInput, Terminal}, inputs::{GramatronInput, Terminal},
mutators::{MutationResult, Mutator}, mutators::{MutationResult, Mutator},
random_corpus_id, random_corpus_id,
state::{HasCorpus, HasMetadata, HasRand}, state::{HasCorpus, HasMetadata, HasRand, HasTestcase},
Error, Error,
}; };
@ -99,7 +99,7 @@ pub struct GramatronSpliceMutator;
impl<S> Mutator<S::Input, S> for GramatronSpliceMutator impl<S> Mutator<S::Input, S> for GramatronSpliceMutator
where where
S: HasRand + HasCorpus<Input = GramatronInput> + HasMetadata, S: HasRand + HasCorpus<Input = GramatronInput> + HasMetadata + HasTestcase,
{ {
fn mutate( fn mutate(
&mut self, &mut self,
@ -117,7 +117,7 @@ where
let rand_num = state.rand_mut().next() as usize; let rand_num = state.rand_mut().next() as usize;
let mut other_testcase = state.corpus().get(idx)?.borrow_mut(); let mut other_testcase = state.testcase_mut(idx)?;
other_testcase.load_input()?; // Preload the input other_testcase.load_input()?; // Preload the input
if !other_testcase.has_metadata::<GramatronIdxMapMetadata>() { if !other_testcase.has_metadata::<GramatronIdxMapMetadata>() {

158
libafl/src/schedulers/ecofuzz.rs
View File

@ -10,7 +10,7 @@ use crate::{
inputs::UsesInput, inputs::UsesInput,
observers::{MapObserver, ObserversTuple}, observers::{MapObserver, ObserversTuple},
schedulers::{powersched::SchedulerMetadata, testcase_score::TestcaseScore, Scheduler}, schedulers::{powersched::SchedulerMetadata, testcase_score::TestcaseScore, Scheduler},
state::{HasCorpus, HasExecutions, HasMetadata, HasRand, UsesState}, state::{HasCorpus, HasExecutions, HasMetadata, HasRand, HasTestcase, UsesState},
Error, Error,
}; };
@ -77,7 +77,7 @@ pub struct EcoScheduler<O, S> {
impl<O, S> EcoScheduler<O, S> impl<O, S> EcoScheduler<O, S>
where where
S: HasCorpus + HasMetadata + HasRand + HasExecutions, S: HasCorpus + HasMetadata + HasRand + HasExecutions + HasTestcase,
O: MapObserver, O: MapObserver,
{ {
/// Create a new [`EcoScheduler`] without any power schedule /// Create a new [`EcoScheduler`] without any power schedule
@ -101,30 +101,22 @@ where
let count = state.corpus().count(); let count = state.corpus().count();
let (last_mutation_num, last_corpus_count) = { let (last_mutation_num, last_corpus_count) = {
let meta = state let meta = state.metadata::<EcoMetadata>()?;
.metadata_map()
.get::<EcoMetadata>()
.ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?;
(meta.last_mutation_num, meta.last_corpus_count) (meta.last_mutation_num, meta.last_corpus_count)
}; };
let computed_score = { let computed_score = {
let mut testcase = state.corpus().get(id)?.borrow_mut(); let mut testcase = state.testcase_mut(id)?;
let meta = testcase.metadata_mut::<EcoTestcaseMetadata>()?;
let meta = testcase
.metadata_map_mut()
.get_mut::<EcoTestcaseMetadata>()
.ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))?;
// Set was_fuzzed for the old current // Set was_fuzzed for the old current
meta.last_found = count - last_corpus_count; meta.last_found = count - last_corpus_count;
meta.last_energy = meta.mutation_num - last_mutation_num; meta.last_energy = meta.mutation_num - last_mutation_num;
meta.computed_score meta.computed_score
}; };
let meta = state let meta = state.metadata_mut::<EcoMetadata>()?;
.metadata_map_mut()
.get_mut::<EcoMetadata>()
.ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?;
let mut regret = meta.last_find_iteration as f64 / computed_score; let mut regret = meta.last_find_iteration as f64 / computed_score;
if regret == 0.0 { if regret == 0.0 {
@ -151,9 +143,7 @@ where
fn first_iteration(state: &mut S) -> Result<(), Error> { fn first_iteration(state: &mut S) -> Result<(), Error> {
let count = state.corpus().count(); let count = state.corpus().count();
state state
.metadata_map_mut() .metadata_mut::<EcoMetadata>()?
.get_mut::<EcoMetadata>()
.ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?
.initial_corpus_count .initial_corpus_count
.get_or_insert(count); .get_or_insert(count);
Ok(()) Ok(())
@ -163,7 +153,7 @@ where
fn schedule(state: &mut S) -> Result<CorpusId, Error> { fn schedule(state: &mut S) -> Result<CorpusId, Error> {
let mut selection = None; let mut selection = None;
for id in state.corpus().ids() { for id in state.corpus().ids() {
let was_fuzzed = state.corpus().get(id)?.borrow().scheduled_count() > 0; let was_fuzzed = state.testcase(id)?.scheduled_count() > 0;
if was_fuzzed { if was_fuzzed {
selection = Some(id); selection = Some(id);
break; break;
@ -171,33 +161,19 @@ where
} }
for id in state.corpus().ids() { for id in state.corpus().ids() {
let was_fuzzed = state.corpus().get(id)?.borrow().scheduled_count() > 0; let was_fuzzed = state.testcase(id)?.scheduled_count() > 0;
if was_fuzzed { if was_fuzzed {
state state.metadata_mut::<EcoMetadata>()?.state = EcoState::Exploration;
.metadata_map_mut()
.get_mut::<EcoMetadata>()
.ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?
.state = EcoState::Exploration;
return Ok(selection.expect("Error in the algorithm, this cannot be None")); return Ok(selection.expect("Error in the algorithm, this cannot be None"));
} }
} }
state state.metadata_mut::<EcoMetadata>()?.state = EcoState::Exploitation;
.metadata_map_mut()
.get_mut::<EcoMetadata>()
.ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?
.state = EcoState::Exploitation;
let mut cur = state.corpus().first(); let mut cur = state.corpus().first();
while let Some(id) = cur { while let Some(id) = cur {
let testcase_state = state let testcase_state = state.testcase(id)?.metadata::<EcoTestcaseMetadata>()?.state;
.corpus()
.get(id)?
.borrow()
.metadata_map()
.get::<EcoTestcaseMetadata>()
.ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))?
.state;
if testcase_state != EcoState::Exploitation { if testcase_state != EcoState::Exploitation {
break; break;
} }
@ -207,14 +183,8 @@ where
if cur.is_none() { if cur.is_none() {
for id in state.corpus().ids() { for id in state.corpus().ids() {
state state
.corpus() .testcase_mut(id)?
.get(id)? .metadata_mut::<EcoTestcaseMetadata>()?
.borrow_mut()
.metadata_map_mut()
.get_mut::<EcoTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("EcoTestcaseMetadata not found".to_string())
})?
.state = EcoState::None; .state = EcoState::None;
} }
@ -223,20 +193,13 @@ where
let mut selection = cur.unwrap(); let mut selection = cur.unwrap();
let mut selection_meta = state let mut selection_meta = state
.corpus() .testcase(selection)?
.get(selection)? .metadata::<EcoTestcaseMetadata>()?
.borrow()
.metadata_map()
.get::<EcoTestcaseMetadata>()
.ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))?
.clone(); .clone();
for id in state.corpus().ids() { for id in state.corpus().ids() {
let testcase = state.corpus().get(id)?.borrow(); let testcase = state.testcase(id)?;
let meta = testcase let meta = testcase.metadata::<EcoTestcaseMetadata>()?;
.metadata_map()
.get::<EcoTestcaseMetadata>()
.ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))?;
if meta.exec_by_mutation if meta.exec_by_mutation
* selection_meta.mutation_num * selection_meta.mutation_num
@ -262,7 +225,7 @@ where
impl<O, S> Scheduler for EcoScheduler<O, S> impl<O, S> Scheduler for EcoScheduler<O, S>
where where
S: HasCorpus + HasMetadata + HasRand + HasExecutions, S: HasCorpus + HasMetadata + HasRand + HasExecutions + HasTestcase,
O: MapObserver, O: MapObserver,
{ {
/// Add an entry to the corpus and return its index /// Add an entry to the corpus and return its index
@ -271,14 +234,8 @@ where
let mut depth = match current_idx { let mut depth = match current_idx {
Some(parent_idx) => state Some(parent_idx) => state
.corpus() .testcase_mut(parent_idx)?
.get(parent_idx)? .metadata_mut::<SchedulerTestcaseMetadata>()?
.borrow_mut()
.metadata_map_mut()
.get_mut::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?
.depth(), .depth(),
None => 0, None => 0,
}; };
@ -286,7 +243,7 @@ where
// Attach a `SchedulerTestcaseMetadata` to the queue entry. // Attach a `SchedulerTestcaseMetadata` to the queue entry.
depth += 1; depth += 1;
{ {
let mut testcase = state.corpus().get(idx)?.borrow_mut(); let mut testcase = state.testcase_mut(idx)?;
testcase.add_metadata(SchedulerTestcaseMetadata::with_n_fuzz_entry( testcase.add_metadata(SchedulerTestcaseMetadata::with_n_fuzz_entry(
depth, depth,
self.last_hash, self.last_hash,
@ -295,16 +252,11 @@ where
} }
// Add the testcase metadata for this scheduler // Add the testcase metadata for this scheduler
state state
.corpus() .testcase_mut(idx)?
.get(idx)?
.borrow_mut()
.add_metadata(EcoTestcaseMetadata::default()); .add_metadata(EcoTestcaseMetadata::default());
let executions = *state.executions(); let executions = *state.executions();
let meta = state let meta = state.metadata_mut::<EcoMetadata>()?;
.metadata_map_mut()
.get_mut::<EcoMetadata>()
.ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?;
let last_find_iteration = executions - meta.last_executions + 1; let last_find_iteration = executions - meta.last_executions + 1;
meta.last_find_iteration = last_find_iteration; meta.last_find_iteration = last_find_iteration;
@ -327,10 +279,7 @@ where
let mut hash = observer.hash() as usize; let mut hash = observer.hash() as usize;
let psmeta = state let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerMetadata>()
.ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?;
hash %= psmeta.n_fuzz().len(); hash %= psmeta.n_fuzz().len();
// Update the path frequency // Update the path frequency
@ -338,34 +287,18 @@ where
if let Some(id) = *state.corpus().current() { if let Some(id) = *state.corpus().current() {
state state
.corpus() .testcase_mut(id)?
.get(id)? .metadata_mut::<EcoTestcaseMetadata>()?
.borrow_mut()
.metadata_map_mut()
.get_mut::<EcoTestcaseMetadata>()
.ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))?
.mutation_num += 1; .mutation_num += 1;
let entry = state let entry = state
.corpus() .testcase(id)?
.get(id)? .metadata::<SchedulerTestcaseMetadata>()?
.borrow()
.metadata_map()
.get::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?
.n_fuzz_entry(); .n_fuzz_entry();
if entry == hash { if entry == hash {
state state
.corpus() .testcase_mut(id)?
.get(id)? .metadata_mut::<EcoTestcaseMetadata>()?
.borrow_mut()
.metadata_map_mut()
.get_mut::<EcoTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("EcoTestcaseMetadata not found".to_string())
})?
.exec_by_mutation += 1; .exec_by_mutation += 1;
} }
} }
@ -386,20 +319,13 @@ where
self.set_current_scheduled(state, Some(id))?; self.set_current_scheduled(state, Some(id))?;
let mutation_num = state let mutation_num = state
.corpus() .testcase_mut(id)?
.get(id)? .metadata_mut::<EcoTestcaseMetadata>()?
.borrow_mut()
.metadata_map_mut()
.get_mut::<EcoTestcaseMetadata>()
.ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))?
.mutation_num; .mutation_num;
let count = state.corpus().count(); let count = state.corpus().count();
let executions = *state.executions(); let executions = *state.executions();
let meta = state let meta = state.metadata_mut::<EcoMetadata>()?;
.metadata_map_mut()
.get_mut::<EcoMetadata>()
.ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?;
meta.last_mutation_num = mutation_num; meta.last_mutation_num = mutation_num;
meta.last_corpus_count = count; meta.last_corpus_count = count;
// TODO in theory it should be assigned at the beginning of the mutational stage // TODO in theory it should be assigned at the beginning of the mutational stage
@ -418,7 +344,7 @@ where
let current_idx = *state.corpus().current(); let current_idx = *state.corpus().current();
if let Some(idx) = current_idx { if let Some(idx) = current_idx {
let mut testcase = state.corpus().get(idx)?.borrow_mut(); let mut testcase = state.testcase_mut(idx)?;
let scheduled_count = testcase.scheduled_count(); let scheduled_count = testcase.scheduled_count();
// increase scheduled count, this was fuzz_level in afl // increase scheduled count, this was fuzz_level in afl
@ -452,17 +378,11 @@ where
} }
let (cur_state, rate) = { let (cur_state, rate) = {
let meta = state let meta = state.metadata::<EcoMetadata>()?;
.metadata_map()
.get::<EcoMetadata>()
.ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?;
(meta.state, meta.rate) (meta.state, meta.rate)
}; };
let meta = entry let meta = entry.metadata_mut::<EcoTestcaseMetadata>()?;
.metadata_map_mut()
.get_mut::<EcoTestcaseMetadata>()
.ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))?;
if cur_state == EcoState::Exploitation { if cur_state == EcoState::Exploitation {
meta.state = EcoState::Exploitation; meta.state = EcoState::Exploitation;

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

@ -13,7 +13,7 @@ use crate::{
inputs::UsesInput, inputs::UsesInput,
observers::{MapObserver, ObserversTuple}, observers::{MapObserver, ObserversTuple},
schedulers::{RemovableScheduler, Scheduler}, schedulers::{RemovableScheduler, Scheduler},
state::{HasCorpus, HasMetadata, UsesState}, state::{HasCorpus, HasMetadata, HasTestcase, UsesState},
Error, Error,
}; };
@ -180,7 +180,7 @@ where
impl<O, S> RemovableScheduler for PowerQueueScheduler<O, S> impl<O, S> RemovableScheduler for PowerQueueScheduler<O, S>
where where
S: HasCorpus + HasMetadata, S: HasCorpus + HasMetadata + HasTestcase,
O: MapObserver, O: MapObserver,
{ {
#[allow(clippy::cast_precision_loss)] #[allow(clippy::cast_precision_loss)]
@ -190,12 +190,7 @@ where
idx: CorpusId, idx: CorpusId,
prev: &Testcase<<Self::State as UsesInput>::Input>, prev: &Testcase<<Self::State as UsesInput>::Input>,
) -> Result<(), Error> { ) -> Result<(), Error> {
let prev_meta = prev let prev_meta = prev.metadata::<SchedulerTestcaseMetadata>()?;
.metadata_map()
.get::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?;
// Next depth is + 1 // Next depth is + 1
let prev_depth = prev_meta.depth() + 1; let prev_depth = prev_meta.depth() + 1;
@ -205,10 +200,7 @@ where
let prev_bitmap_size = prev_meta.bitmap_size(); let prev_bitmap_size = prev_meta.bitmap_size();
let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64); let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64);
let psmeta = state let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerMetadata>()
.ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?;
// We won't add new one because it'll get added when it gets executed in calirbation next time. // 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_exec_time(psmeta.exec_time() - prev_total_time);
@ -218,9 +210,7 @@ where
psmeta.set_bitmap_entries(psmeta.bitmap_entries() - 1); psmeta.set_bitmap_entries(psmeta.bitmap_entries() - 1);
state state
.corpus() .testcase_mut(idx)?
.get(idx)?
.borrow_mut()
.add_metadata(SchedulerTestcaseMetadata::new(prev_depth)); .add_metadata(SchedulerTestcaseMetadata::new(prev_depth));
Ok(()) Ok(())
} }
@ -238,22 +228,14 @@ where
) )
})?; })?;
let prev_meta = prev let prev_meta = prev.metadata::<SchedulerTestcaseMetadata>()?;
.metadata_map()
.get::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?;
// Use these to adjust `SchedulerMetadata` // Use these to adjust `SchedulerMetadata`
let (prev_total_time, prev_cycles) = prev_meta.cycle_and_time(); let (prev_total_time, prev_cycles) = prev_meta.cycle_and_time();
let prev_bitmap_size = prev_meta.bitmap_size(); let prev_bitmap_size = prev_meta.bitmap_size();
let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64); let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64);
let psmeta = state let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerMetadata>()
.ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?;
psmeta.set_exec_time(psmeta.exec_time() - prev_total_time); psmeta.set_exec_time(psmeta.exec_time() - prev_total_time);
psmeta.set_cycles(psmeta.cycles() - (prev_cycles as u64)); psmeta.set_cycles(psmeta.cycles() - (prev_cycles as u64));
@ -267,7 +249,7 @@ where
impl<O, S> Scheduler for PowerQueueScheduler<O, S> impl<O, S> Scheduler for PowerQueueScheduler<O, S>
where where
S: HasCorpus + HasMetadata, S: HasCorpus + HasMetadata + HasTestcase,
O: MapObserver, O: MapObserver,
{ {
/// Add an entry to the corpus and return its index /// Add an entry to the corpus and return its index
@ -276,14 +258,8 @@ where
let mut depth = match current_idx { let mut depth = match current_idx {
Some(parent_idx) => state Some(parent_idx) => state
.corpus() .testcase(parent_idx)?
.get(parent_idx)? .metadata::<SchedulerTestcaseMetadata>()?
.borrow()
.metadata_map()
.get::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?
.depth(), .depth(),
None => 0, None => 0,
}; };
@ -293,7 +269,7 @@ where
// Attach a `SchedulerTestcaseMetadata` to the queue entry. // Attach a `SchedulerTestcaseMetadata` to the queue entry.
depth += 1; depth += 1;
let mut testcase = state.corpus().get(idx)?.borrow_mut(); let mut testcase = state.testcase_mut(idx)?;
testcase.add_metadata(SchedulerTestcaseMetadata::with_n_fuzz_entry( testcase.add_metadata(SchedulerTestcaseMetadata::with_n_fuzz_entry(
depth, depth,
self.last_hash, self.last_hash,
@ -317,10 +293,7 @@ where
let mut hash = observer.hash() as usize; let mut hash = observer.hash() as usize;
let psmeta = state let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerMetadata>()
.ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?;
hash %= psmeta.n_fuzz().len(); hash %= psmeta.n_fuzz().len();
// Update the path frequency // Update the path frequency
@ -340,12 +313,7 @@ where
if let Some(next) = state.corpus().next(*cur) { if let Some(next) = state.corpus().next(*cur) {
next next
} else { } else {
let psmeta = state let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerMetadata not found".to_string())
})?;
psmeta.set_queue_cycles(psmeta.queue_cycles() + 1); psmeta.set_queue_cycles(psmeta.queue_cycles() + 1);
state.corpus().first().unwrap() state.corpus().first().unwrap()
} }
@ -367,18 +335,13 @@ where
let current_idx = *state.corpus().current(); let current_idx = *state.corpus().current();
if let Some(idx) = current_idx { if let Some(idx) = current_idx {
let mut testcase = state.corpus().get(idx)?.borrow_mut(); let mut testcase = state.testcase_mut(idx)?;
let scheduled_count = testcase.scheduled_count(); let scheduled_count = testcase.scheduled_count();
// increase scheduled count, this was fuzz_level in afl // increase scheduled count, this was fuzz_level in afl
testcase.set_scheduled_count(scheduled_count + 1); testcase.set_scheduled_count(scheduled_count + 1);
let tcmeta = testcase let tcmeta = testcase.metadata_mut::<SchedulerTestcaseMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?;
if tcmeta.handicap() >= 4 { if tcmeta.handicap() >= 4 {
tcmeta.set_handicap(tcmeta.handicap() - 4); tcmeta.set_handicap(tcmeta.handicap() - 4);

47
libafl/src/schedulers/testcase_score.rs
View File

@ -66,10 +66,7 @@ where
clippy::cast_lossless clippy::cast_lossless
)] )]
fn compute(entry: &mut Testcase<S::Input>, state: &S) -> Result<f64, Error> { fn compute(entry: &mut Testcase<S::Input>, state: &S) -> Result<f64, Error> {
let psmeta = state let psmeta = state.metadata::<SchedulerMetadata>()?;
.metadata_map()
.get::<SchedulerMetadata>()
.ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?;
let fuzz_mu = if let Some(strat) = psmeta.strat() { let fuzz_mu = if let Some(strat) = psmeta.strat() {
if strat == PowerSchedule::COE { if strat == PowerSchedule::COE {
@ -80,25 +77,13 @@ where
for idx in corpus.ids() { for idx in corpus.ids() {
let n_fuzz_entry = if cur_index == idx { let n_fuzz_entry = if cur_index == idx {
entry entry
.metadata_map() .metadata::<SchedulerTestcaseMetadata>()?
.get::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found(
"SchedulerTestcaseMetadata not found".to_string(),
)
})?
.n_fuzz_entry() .n_fuzz_entry()
} else { } else {
corpus corpus
.get(idx)? .get(idx)?
.borrow() .borrow()
.metadata_map() .metadata::<SchedulerTestcaseMetadata>()?
.get::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found(
"SchedulerTestcaseMetadata not found".to_string(),
)
})?
.n_fuzz_entry() .n_fuzz_entry()
}; };
v += libm::log2(f64::from(psmeta.n_fuzz()[n_fuzz_entry])); v += libm::log2(f64::from(psmeta.n_fuzz()[n_fuzz_entry]));
@ -128,12 +113,7 @@ where
let avg_bitmap_size = psmeta.bitmap_size() / psmeta.bitmap_entries(); let avg_bitmap_size = psmeta.bitmap_size() / psmeta.bitmap_entries();
let favored = entry.has_metadata::<IsFavoredMetadata>(); let favored = entry.has_metadata::<IsFavoredMetadata>();
let tcmeta = entry let tcmeta = entry.metadata::<SchedulerTestcaseMetadata>()?;
.metadata_map()
.get::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?;
if q_exec_us * 0.1 > avg_exec_us { if q_exec_us * 0.1 > avg_exec_us {
perf_score = 10.0; perf_score = 10.0;
@ -295,17 +275,9 @@ where
#[allow(clippy::cast_precision_loss, clippy::cast_lossless)] #[allow(clippy::cast_precision_loss, clippy::cast_lossless)]
fn compute(entry: &mut Testcase<S::Input>, state: &S) -> Result<f64, Error> { fn compute(entry: &mut Testcase<S::Input>, state: &S) -> Result<f64, Error> {
let mut weight = 1.0; let mut weight = 1.0;
let psmeta = state let psmeta = state.metadata::<SchedulerMetadata>()?;
.metadata_map()
.get::<SchedulerMetadata>()
.ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?;
let tcmeta = entry let tcmeta = entry.metadata::<SchedulerTestcaseMetadata>()?;
.metadata_map()
.get::<SchedulerTestcaseMetadata>()
.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, // This means that this testcase has never gone through the calibration stage before1,
// In this case we'll just return the default weight // In this case we'll just return the default weight
@ -349,12 +321,7 @@ where
None => 0.0, None => 0.0,
}; };
let avg_top_size = state let avg_top_size = state.metadata::<TopRatedsMetadata>()?.map().len() as f64;
.metadata_map()
.get::<TopRatedsMetadata>()
.ok_or_else(|| Error::key_not_found("TopRatedsMetadata not found".to_string()))?
.map()
.len() as f64;
weight *= 1.0 + (tc_ref / avg_top_size); weight *= 1.0 + (tc_ref / avg_top_size);
if favored { if favored {

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

@ -18,7 +18,7 @@ use crate::{
testcase_score::{CorpusWeightTestcaseScore, TestcaseScore}, testcase_score::{CorpusWeightTestcaseScore, TestcaseScore},
RemovableScheduler, Scheduler, RemovableScheduler, Scheduler,
}, },
state::{HasCorpus, HasMetadata, HasRand, UsesState}, state::{HasCorpus, HasMetadata, HasRand, HasTestcase, UsesState},
Error, Error,
}; };
@ -208,12 +208,7 @@ where
alias_probability.insert(*s_arr.get(&n_s).unwrap(), 1.0); alias_probability.insert(*s_arr.get(&n_s).unwrap(), 1.0);
} }
let wsmeta = state let wsmeta = state.metadata_mut::<WeightedScheduleMetadata>()?;
.metadata_map_mut()
.get_mut::<WeightedScheduleMetadata>()
.ok_or_else(|| {
Error::key_not_found("WeigthedScheduleMetadata not found".to_string())
})?;
// Update metadata // Update metadata
wsmeta.set_alias_probability(alias_probability); wsmeta.set_alias_probability(alias_probability);
@ -233,7 +228,7 @@ impl<F, O, S> RemovableScheduler for WeightedScheduler<F, O, S>
where where
F: TestcaseScore<S>, F: TestcaseScore<S>,
O: MapObserver, O: MapObserver,
S: HasCorpus + HasMetadata + HasRand, S: HasCorpus + HasMetadata + HasRand + HasTestcase,
{ {
#[allow(clippy::cast_precision_loss)] #[allow(clippy::cast_precision_loss)]
fn on_remove( fn on_remove(
@ -248,22 +243,14 @@ where
) )
})?; })?;
let prev_meta = prev let prev_meta = prev.metadata::<SchedulerTestcaseMetadata>()?;
.metadata_map()
.get::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?;
// Use these to adjust `SchedulerMetadata` // Use these to adjust `SchedulerMetadata`
let (prev_total_time, prev_cycles) = prev_meta.cycle_and_time(); let (prev_total_time, prev_cycles) = prev_meta.cycle_and_time();
let prev_bitmap_size = prev_meta.bitmap_size(); let prev_bitmap_size = prev_meta.bitmap_size();
let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64); let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64);
let psmeta = state let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerMetadata>()
.ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?;
psmeta.set_exec_time(psmeta.exec_time() - prev_total_time); psmeta.set_exec_time(psmeta.exec_time() - prev_total_time);
psmeta.set_cycles(psmeta.cycles() - (prev_cycles as u64)); psmeta.set_cycles(psmeta.cycles() - (prev_cycles as u64));
@ -281,12 +268,7 @@ where
idx: CorpusId, idx: CorpusId,
prev: &crate::corpus::Testcase<<Self::State as UsesInput>::Input>, prev: &crate::corpus::Testcase<<Self::State as UsesInput>::Input>,
) -> Result<(), Error> { ) -> Result<(), Error> {
let prev_meta = prev let prev_meta = prev.metadata::<SchedulerTestcaseMetadata>()?;
.metadata_map()
.get::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?;
// Next depth is + 1 // Next depth is + 1
let prev_depth = prev_meta.depth() + 1; let prev_depth = prev_meta.depth() + 1;
@ -296,10 +278,7 @@ where
let prev_bitmap_size = prev_meta.bitmap_size(); let prev_bitmap_size = prev_meta.bitmap_size();
let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64); let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64);
let psmeta = state let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerMetadata>()
.ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?;
// We won't add new one because it'll get added when it gets executed in calirbation next time. // 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_exec_time(psmeta.exec_time() - prev_total_time);
@ -321,7 +300,7 @@ impl<F, O, S> Scheduler for WeightedScheduler<F, O, S>
where where
F: TestcaseScore<S>, F: TestcaseScore<S>,
O: MapObserver, O: MapObserver,
S: HasCorpus + HasMetadata + HasRand, S: HasCorpus + HasMetadata + HasRand + HasTestcase,
{ {
/// Add an entry to the corpus and return its index /// Add an entry to the corpus and return its index
fn on_add(&mut self, state: &mut S, idx: CorpusId) -> Result<(), Error> { fn on_add(&mut self, state: &mut S, idx: CorpusId) -> Result<(), Error> {
@ -329,14 +308,8 @@ where
let mut depth = match current_idx { let mut depth = match current_idx {
Some(parent_idx) => state Some(parent_idx) => state
.corpus() .testcase_mut(parent_idx)?
.get(parent_idx)? .metadata_mut::<SchedulerTestcaseMetadata>()?
.borrow_mut()
.metadata_map_mut()
.get_mut::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?
.depth(), .depth(),
None => 0, None => 0,
}; };
@ -375,10 +348,7 @@ where
let mut hash = observer.hash() as usize; let mut hash = observer.hash() as usize;
let psmeta = state let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerMetadata>()
.ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?;
hash %= psmeta.n_fuzz().len(); hash %= psmeta.n_fuzz().len();
// Update the path frequency // Update the path frequency
@ -400,12 +370,7 @@ where
// Choose a random value between 0.000000000 and 1.000000000 // Choose a random value between 0.000000000 and 1.000000000
let probability = state.rand_mut().between(0, 1000000000) as f64 / 1000000000_f64; let probability = state.rand_mut().between(0, 1000000000) as f64 / 1000000000_f64;
let wsmeta = state let wsmeta = state.metadata_mut::<WeightedScheduleMetadata>()?;
.metadata_map_mut()
.get_mut::<WeightedScheduleMetadata>()
.ok_or_else(|| {
Error::key_not_found("WeigthedScheduleMetadata not found".to_string())
})?;
let current_cycles = wsmeta.runs_in_current_cycle(); let current_cycles = wsmeta.runs_in_current_cycle();
@ -424,12 +389,7 @@ where
// Update depth // Update depth
if current_cycles > corpus_counts { if current_cycles > corpus_counts {
let psmeta = state let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerMetadata not found".to_string())
})?;
psmeta.set_queue_cycles(psmeta.queue_cycles() + 1); psmeta.set_queue_cycles(psmeta.queue_cycles() + 1);
} }
@ -447,18 +407,13 @@ where
let current_idx = *state.corpus().current(); let current_idx = *state.corpus().current();
if let Some(idx) = current_idx { if let Some(idx) = current_idx {
let mut testcase = state.corpus().get(idx)?.borrow_mut(); let mut testcase = state.testcase_mut(idx)?;
let scheduled_count = testcase.scheduled_count(); let scheduled_count = testcase.scheduled_count();
// increase scheduled count, this was fuzz_level in afl // increase scheduled count, this was fuzz_level in afl
testcase.set_scheduled_count(scheduled_count + 1); testcase.set_scheduled_count(scheduled_count + 1);
let tcmeta = testcase let tcmeta = testcase.metadata_mut::<SchedulerTestcaseMetadata>()?;
.metadata_map_mut()
.get_mut::<SchedulerTestcaseMetadata>()
.ok_or_else(|| {
Error::key_not_found("SchedulerTestcaseMetadata not found".to_string())
})?;
if tcmeta.handicap() >= 4 { if tcmeta.handicap() >= 4 {
tcmeta.set_handicap(tcmeta.handicap() - 4); tcmeta.set_handicap(tcmeta.handicap() - 4);

15
libafl/src/stages/sync.rs
View File

@ -17,7 +17,10 @@ use crate::{
fuzzer::{Evaluator, EvaluatorObservers, ExecutionProcessor}, fuzzer::{Evaluator, EvaluatorObservers, ExecutionProcessor},
inputs::{Input, InputConverter, UsesInput}, inputs::{Input, InputConverter, UsesInput},
stages::Stage, stages::Stage,
state::{HasClientPerfMonitor, HasCorpus, HasExecutions, HasMetadata, HasRand, UsesState}, state::{
HasClientPerfMonitor, HasCorpus, HasExecutions, HasMetadata, HasRand, HasTestcase,
UsesState,
},
Error, Error,
}; };
@ -234,7 +237,13 @@ where
impl<E, EM, IC, ICB, DI, S, SP, Z> Stage<E, EM, Z> for SyncFromBrokerStage<IC, ICB, DI, S, SP> impl<E, EM, IC, ICB, DI, S, SP, Z> Stage<E, EM, Z> for SyncFromBrokerStage<IC, ICB, DI, S, SP>
where where
EM: UsesState<State = S> + EventFirer, EM: UsesState<State = S> + EventFirer,
S: UsesInput + HasClientPerfMonitor + HasExecutions + HasCorpus + HasRand + HasMetadata, S: UsesInput
+ HasClientPerfMonitor
+ HasExecutions
+ HasCorpus
+ HasRand
+ HasMetadata
+ HasTestcase,
SP: ShMemProvider, SP: ShMemProvider,
E: HasObservers<State = S> + Executor<EM, Z>, E: HasObservers<State = S> + Executor<EM, Z>,
for<'a> E::Observers: Deserialize<'a>, for<'a> E::Observers: Deserialize<'a>,
@ -262,7 +271,7 @@ where
last_id.map_or_else(|| state.corpus().first(), |id| state.corpus().next(id)); last_id.map_or_else(|| state.corpus().first(), |id| state.corpus().next(id));
while let Some(id) = cur_id { while let Some(id) = cur_id {
let input = state.corpus().get(id)?.borrow_mut().load_input()?.clone(); let input = state.testcase_mut(id)?.load_input()?.clone();
self.client.fire( self.client.fire(
state, state,

2
libafl/src/state/mod.rs
View File

@ -228,7 +228,7 @@ pub trait HasTestcase: HasCorpus {
/// To get mutable testcase /// To get mutable testcase
fn testcase_mut( fn testcase_mut(
&mut self, &self,
id: CorpusId, id: CorpusId,
) -> Result<RefMut<Testcase<<Self as UsesInput>::Input>>, Error> { ) -> Result<RefMut<Testcase<<Self as UsesInput>::Input>>, Error> {
Ok(self.corpus().get(id)?.borrow_mut()) Ok(self.corpus().get(id)?.borrow_mut())