From 104c170ade67d752300b400529c8e26bd9229e14 Mon Sep 17 00:00:00 2001 From: Matheus Baptistella <63082492+matheusbaptistella@users.noreply.github.com> Date: Fri, 17 Mar 2023 12:03:57 -0300 Subject: [PATCH] Use the new metadata() function in more places (#1155) --- libafl/src/mutators/gramatron.rs | 6 +- libafl/src/schedulers/ecofuzz.rs | 158 ++++++------------------ libafl/src/schedulers/powersched.rs | 67 +++------- libafl/src/schedulers/testcase_score.rs | 47 ++----- libafl/src/schedulers/weighted.rs | 75 +++-------- libafl/src/stages/sync.rs | 15 ++- libafl/src/state/mod.rs | 2 +- 7 files changed, 92 insertions(+), 278 deletions(-) diff --git a/libafl/src/mutators/gramatron.rs b/libafl/src/mutators/gramatron.rs index aacc086b22..330096aa85 100644 --- a/libafl/src/mutators/gramatron.rs +++ b/libafl/src/mutators/gramatron.rs @@ -13,7 +13,7 @@ use crate::{ inputs::{GramatronInput, Terminal}, mutators::{MutationResult, Mutator}, random_corpus_id, - state::{HasCorpus, HasMetadata, HasRand}, + state::{HasCorpus, HasMetadata, HasRand, HasTestcase}, Error, }; @@ -99,7 +99,7 @@ pub struct GramatronSpliceMutator; impl Mutator for GramatronSpliceMutator where - S: HasRand + HasCorpus + HasMetadata, + S: HasRand + HasCorpus + HasMetadata + HasTestcase, { fn mutate( &mut self, @@ -117,7 +117,7 @@ where 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 if !other_testcase.has_metadata::() { diff --git a/libafl/src/schedulers/ecofuzz.rs b/libafl/src/schedulers/ecofuzz.rs index 3ff6821112..0c826a90df 100644 --- a/libafl/src/schedulers/ecofuzz.rs +++ b/libafl/src/schedulers/ecofuzz.rs @@ -10,7 +10,7 @@ use crate::{ inputs::UsesInput, observers::{MapObserver, ObserversTuple}, schedulers::{powersched::SchedulerMetadata, testcase_score::TestcaseScore, Scheduler}, - state::{HasCorpus, HasExecutions, HasMetadata, HasRand, UsesState}, + state::{HasCorpus, HasExecutions, HasMetadata, HasRand, HasTestcase, UsesState}, Error, }; @@ -77,7 +77,7 @@ pub struct EcoScheduler { impl EcoScheduler where - S: HasCorpus + HasMetadata + HasRand + HasExecutions, + S: HasCorpus + HasMetadata + HasRand + HasExecutions + HasTestcase, O: MapObserver, { /// Create a new [`EcoScheduler`] without any power schedule @@ -101,30 +101,22 @@ where let count = state.corpus().count(); let (last_mutation_num, last_corpus_count) = { - let meta = state - .metadata_map() - .get::() - .ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?; + let meta = state.metadata::()?; (meta.last_mutation_num, meta.last_corpus_count) }; let computed_score = { - let mut testcase = state.corpus().get(id)?.borrow_mut(); + let mut testcase = state.testcase_mut(id)?; + + let meta = testcase.metadata_mut::()?; - let meta = testcase - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))?; // Set was_fuzzed for the old current meta.last_found = count - last_corpus_count; meta.last_energy = meta.mutation_num - last_mutation_num; meta.computed_score }; - let meta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?; + let meta = state.metadata_mut::()?; let mut regret = meta.last_find_iteration as f64 / computed_score; if regret == 0.0 { @@ -151,9 +143,7 @@ where fn first_iteration(state: &mut S) -> Result<(), Error> { let count = state.corpus().count(); state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))? + .metadata_mut::()? .initial_corpus_count .get_or_insert(count); Ok(()) @@ -163,7 +153,7 @@ where fn schedule(state: &mut S) -> Result { let mut selection = None; 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 { selection = Some(id); break; @@ -171,33 +161,19 @@ where } 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 { - state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))? - .state = EcoState::Exploration; + state.metadata_mut::()?.state = EcoState::Exploration; return Ok(selection.expect("Error in the algorithm, this cannot be None")); } } - state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))? - .state = EcoState::Exploitation; + state.metadata_mut::()?.state = EcoState::Exploitation; let mut cur = state.corpus().first(); while let Some(id) = cur { - let testcase_state = state - .corpus() - .get(id)? - .borrow() - .metadata_map() - .get::() - .ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))? - .state; + let testcase_state = state.testcase(id)?.metadata::()?.state; + if testcase_state != EcoState::Exploitation { break; } @@ -207,14 +183,8 @@ where if cur.is_none() { for id in state.corpus().ids() { state - .corpus() - .get(id)? - .borrow_mut() - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("EcoTestcaseMetadata not found".to_string()) - })? + .testcase_mut(id)? + .metadata_mut::()? .state = EcoState::None; } @@ -223,20 +193,13 @@ where let mut selection = cur.unwrap(); let mut selection_meta = state - .corpus() - .get(selection)? - .borrow() - .metadata_map() - .get::() - .ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))? + .testcase(selection)? + .metadata::()? .clone(); for id in state.corpus().ids() { - let testcase = state.corpus().get(id)?.borrow(); - let meta = testcase - .metadata_map() - .get::() - .ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))?; + let testcase = state.testcase(id)?; + let meta = testcase.metadata::()?; if meta.exec_by_mutation * selection_meta.mutation_num @@ -262,7 +225,7 @@ where impl Scheduler for EcoScheduler where - S: HasCorpus + HasMetadata + HasRand + HasExecutions, + S: HasCorpus + HasMetadata + HasRand + HasExecutions + HasTestcase, O: MapObserver, { /// Add an entry to the corpus and return its index @@ -271,14 +234,8 @@ where let mut depth = match current_idx { Some(parent_idx) => state - .corpus() - .get(parent_idx)? - .borrow_mut() - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })? + .testcase_mut(parent_idx)? + .metadata_mut::()? .depth(), None => 0, }; @@ -286,7 +243,7 @@ where // Attach a `SchedulerTestcaseMetadata` to the queue entry. 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( depth, self.last_hash, @@ -295,16 +252,11 @@ where } // Add the testcase metadata for this scheduler state - .corpus() - .get(idx)? - .borrow_mut() + .testcase_mut(idx)? .add_metadata(EcoTestcaseMetadata::default()); let executions = *state.executions(); - let meta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?; + let meta = state.metadata_mut::()?; let last_find_iteration = executions - meta.last_executions + 1; meta.last_find_iteration = last_find_iteration; @@ -327,10 +279,7 @@ where let mut hash = observer.hash() as usize; - let psmeta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; + let psmeta = state.metadata_mut::()?; hash %= psmeta.n_fuzz().len(); // Update the path frequency @@ -338,34 +287,18 @@ where if let Some(id) = *state.corpus().current() { state - .corpus() - .get(id)? - .borrow_mut() - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))? + .testcase_mut(id)? + .metadata_mut::()? .mutation_num += 1; let entry = state - .corpus() - .get(id)? - .borrow() - .metadata_map() - .get::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })? + .testcase(id)? + .metadata::()? .n_fuzz_entry(); if entry == hash { state - .corpus() - .get(id)? - .borrow_mut() - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("EcoTestcaseMetadata not found".to_string()) - })? + .testcase_mut(id)? + .metadata_mut::()? .exec_by_mutation += 1; } } @@ -386,20 +319,13 @@ where self.set_current_scheduled(state, Some(id))?; let mutation_num = state - .corpus() - .get(id)? - .borrow_mut() - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))? + .testcase_mut(id)? + .metadata_mut::()? .mutation_num; let count = state.corpus().count(); let executions = *state.executions(); - let meta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?; + let meta = state.metadata_mut::()?; meta.last_mutation_num = mutation_num; meta.last_corpus_count = count; // 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(); 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(); // increase scheduled count, this was fuzz_level in afl @@ -452,17 +378,11 @@ where } let (cur_state, rate) = { - let meta = state - .metadata_map() - .get::() - .ok_or_else(|| Error::key_not_found("EcoMetadata not found".to_string()))?; + let meta = state.metadata::()?; (meta.state, meta.rate) }; - let meta = entry - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("EcoTestcaseMetadata not found".to_string()))?; + let meta = entry.metadata_mut::()?; if cur_state == EcoState::Exploitation { meta.state = EcoState::Exploitation; diff --git a/libafl/src/schedulers/powersched.rs b/libafl/src/schedulers/powersched.rs index b0257f3a72..6a40d0ff8e 100644 --- a/libafl/src/schedulers/powersched.rs +++ b/libafl/src/schedulers/powersched.rs @@ -13,7 +13,7 @@ use crate::{ inputs::UsesInput, observers::{MapObserver, ObserversTuple}, schedulers::{RemovableScheduler, Scheduler}, - state::{HasCorpus, HasMetadata, UsesState}, + state::{HasCorpus, HasMetadata, HasTestcase, UsesState}, Error, }; @@ -180,7 +180,7 @@ where impl RemovableScheduler for PowerQueueScheduler where - S: HasCorpus + HasMetadata, + S: HasCorpus + HasMetadata + HasTestcase, O: MapObserver, { #[allow(clippy::cast_precision_loss)] @@ -190,12 +190,7 @@ where idx: CorpusId, prev: &Testcase<::Input>, ) -> Result<(), Error> { - let prev_meta = prev - .metadata_map() - .get::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })?; + let prev_meta = prev.metadata::()?; // Next depth is + 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_log = libm::log2(prev_bitmap_size as f64); - let psmeta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; + 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); @@ -218,9 +210,7 @@ where psmeta.set_bitmap_entries(psmeta.bitmap_entries() - 1); state - .corpus() - .get(idx)? - .borrow_mut() + .testcase_mut(idx)? .add_metadata(SchedulerTestcaseMetadata::new(prev_depth)); Ok(()) } @@ -238,22 +228,14 @@ where ) })?; - let prev_meta = prev - .metadata_map() - .get::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })?; + 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_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; + let psmeta = state.metadata_mut::()?; psmeta.set_exec_time(psmeta.exec_time() - prev_total_time); psmeta.set_cycles(psmeta.cycles() - (prev_cycles as u64)); @@ -267,7 +249,7 @@ where impl Scheduler for PowerQueueScheduler where - S: HasCorpus + HasMetadata, + S: HasCorpus + HasMetadata + HasTestcase, O: MapObserver, { /// Add an entry to the corpus and return its index @@ -276,14 +258,8 @@ where let mut depth = match current_idx { Some(parent_idx) => state - .corpus() - .get(parent_idx)? - .borrow() - .metadata_map() - .get::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })? + .testcase(parent_idx)? + .metadata::()? .depth(), None => 0, }; @@ -293,7 +269,7 @@ where // Attach a `SchedulerTestcaseMetadata` to the queue entry. 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( depth, self.last_hash, @@ -317,10 +293,7 @@ where let mut hash = observer.hash() as usize; - let psmeta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; + let psmeta = state.metadata_mut::()?; hash %= psmeta.n_fuzz().len(); // Update the path frequency @@ -340,12 +313,7 @@ where if let Some(next) = state.corpus().next(*cur) { next } else { - let psmeta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("SchedulerMetadata not found".to_string()) - })?; + let psmeta = state.metadata_mut::()?; psmeta.set_queue_cycles(psmeta.queue_cycles() + 1); state.corpus().first().unwrap() } @@ -367,18 +335,13 @@ where let current_idx = *state.corpus().current(); 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(); // increase scheduled count, this was fuzz_level in afl testcase.set_scheduled_count(scheduled_count + 1); - let tcmeta = testcase - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })?; + let tcmeta = testcase.metadata_mut::()?; if tcmeta.handicap() >= 4 { tcmeta.set_handicap(tcmeta.handicap() - 4); diff --git a/libafl/src/schedulers/testcase_score.rs b/libafl/src/schedulers/testcase_score.rs index d79999b71a..2319d36f55 100644 --- a/libafl/src/schedulers/testcase_score.rs +++ b/libafl/src/schedulers/testcase_score.rs @@ -66,10 +66,7 @@ where clippy::cast_lossless )] fn compute(entry: &mut Testcase, state: &S) -> Result { - let psmeta = state - .metadata_map() - .get::() - .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; + let psmeta = state.metadata::()?; let fuzz_mu = if let Some(strat) = psmeta.strat() { if strat == PowerSchedule::COE { @@ -80,25 +77,13 @@ where for idx in corpus.ids() { let n_fuzz_entry = if cur_index == idx { entry - .metadata_map() - .get::() - .ok_or_else(|| { - Error::key_not_found( - "SchedulerTestcaseMetadata not found".to_string(), - ) - })? + .metadata::()? .n_fuzz_entry() } else { corpus .get(idx)? .borrow() - .metadata_map() - .get::() - .ok_or_else(|| { - Error::key_not_found( - "SchedulerTestcaseMetadata not found".to_string(), - ) - })? + .metadata::()? .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 favored = entry.has_metadata::(); - let tcmeta = entry - .metadata_map() - .get::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })?; + let tcmeta = entry.metadata::()?; if q_exec_us * 0.1 > avg_exec_us { perf_score = 10.0; @@ -295,17 +275,9 @@ where #[allow(clippy::cast_precision_loss, clippy::cast_lossless)] fn compute(entry: &mut Testcase, state: &S) -> Result { let mut weight = 1.0; - let psmeta = state - .metadata_map() - .get::() - .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; + let psmeta = state.metadata::()?; - let tcmeta = entry - .metadata_map() - .get::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })?; + let tcmeta = entry.metadata::()?; // This means that this testcase has never gone through the calibration stage before1, // In this case we'll just return the default weight @@ -349,12 +321,7 @@ where None => 0.0, }; - let avg_top_size = state - .metadata_map() - .get::() - .ok_or_else(|| Error::key_not_found("TopRatedsMetadata not found".to_string()))? - .map() - .len() as f64; + let avg_top_size = state.metadata::()?.map().len() as f64; weight *= 1.0 + (tc_ref / avg_top_size); if favored { diff --git a/libafl/src/schedulers/weighted.rs b/libafl/src/schedulers/weighted.rs index 84859bd035..d09280c5ed 100644 --- a/libafl/src/schedulers/weighted.rs +++ b/libafl/src/schedulers/weighted.rs @@ -18,7 +18,7 @@ use crate::{ testcase_score::{CorpusWeightTestcaseScore, TestcaseScore}, RemovableScheduler, Scheduler, }, - state::{HasCorpus, HasMetadata, HasRand, UsesState}, + state::{HasCorpus, HasMetadata, HasRand, HasTestcase, UsesState}, Error, }; @@ -208,12 +208,7 @@ where alias_probability.insert(*s_arr.get(&n_s).unwrap(), 1.0); } - let wsmeta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("WeigthedScheduleMetadata not found".to_string()) - })?; + let wsmeta = state.metadata_mut::()?; // Update metadata wsmeta.set_alias_probability(alias_probability); @@ -233,7 +228,7 @@ impl RemovableScheduler for WeightedScheduler where F: TestcaseScore, O: MapObserver, - S: HasCorpus + HasMetadata + HasRand, + S: HasCorpus + HasMetadata + HasRand + HasTestcase, { #[allow(clippy::cast_precision_loss)] fn on_remove( @@ -248,22 +243,14 @@ where ) })?; - let prev_meta = prev - .metadata_map() - .get::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })?; + 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_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; + let psmeta = state.metadata_mut::()?; psmeta.set_exec_time(psmeta.exec_time() - prev_total_time); psmeta.set_cycles(psmeta.cycles() - (prev_cycles as u64)); @@ -281,12 +268,7 @@ where idx: CorpusId, prev: &crate::corpus::Testcase<::Input>, ) -> Result<(), Error> { - let prev_meta = prev - .metadata_map() - .get::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })?; + let prev_meta = prev.metadata::()?; // Next depth is + 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_log = libm::log2(prev_bitmap_size as f64); - let psmeta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; + 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); @@ -321,7 +300,7 @@ impl Scheduler for WeightedScheduler where F: TestcaseScore, O: MapObserver, - S: HasCorpus + HasMetadata + HasRand, + S: HasCorpus + HasMetadata + HasRand + HasTestcase, { /// Add an entry to the corpus and return its index fn on_add(&mut self, state: &mut S, idx: CorpusId) -> Result<(), Error> { @@ -329,14 +308,8 @@ where let mut depth = match current_idx { Some(parent_idx) => state - .corpus() - .get(parent_idx)? - .borrow_mut() - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })? + .testcase_mut(parent_idx)? + .metadata_mut::()? .depth(), None => 0, }; @@ -375,10 +348,7 @@ where let mut hash = observer.hash() as usize; - let psmeta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| Error::key_not_found("SchedulerMetadata not found".to_string()))?; + let psmeta = state.metadata_mut::()?; hash %= psmeta.n_fuzz().len(); // Update the path frequency @@ -400,12 +370,7 @@ where // Choose a random value between 0.000000000 and 1.000000000 let probability = state.rand_mut().between(0, 1000000000) as f64 / 1000000000_f64; - let wsmeta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("WeigthedScheduleMetadata not found".to_string()) - })?; + let wsmeta = state.metadata_mut::()?; let current_cycles = wsmeta.runs_in_current_cycle(); @@ -424,12 +389,7 @@ where // Update depth if current_cycles > corpus_counts { - let psmeta = state - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("SchedulerMetadata not found".to_string()) - })?; + let psmeta = state.metadata_mut::()?; psmeta.set_queue_cycles(psmeta.queue_cycles() + 1); } @@ -447,18 +407,13 @@ where let current_idx = *state.corpus().current(); 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(); // increase scheduled count, this was fuzz_level in afl testcase.set_scheduled_count(scheduled_count + 1); - let tcmeta = testcase - .metadata_map_mut() - .get_mut::() - .ok_or_else(|| { - Error::key_not_found("SchedulerTestcaseMetadata not found".to_string()) - })?; + let tcmeta = testcase.metadata_mut::()?; if tcmeta.handicap() >= 4 { tcmeta.set_handicap(tcmeta.handicap() - 4); diff --git a/libafl/src/stages/sync.rs b/libafl/src/stages/sync.rs index a661184472..9bf821aeb2 100644 --- a/libafl/src/stages/sync.rs +++ b/libafl/src/stages/sync.rs @@ -17,7 +17,10 @@ use crate::{ fuzzer::{Evaluator, EvaluatorObservers, ExecutionProcessor}, inputs::{Input, InputConverter, UsesInput}, stages::Stage, - state::{HasClientPerfMonitor, HasCorpus, HasExecutions, HasMetadata, HasRand, UsesState}, + state::{ + HasClientPerfMonitor, HasCorpus, HasExecutions, HasMetadata, HasRand, HasTestcase, + UsesState, + }, Error, }; @@ -234,7 +237,13 @@ where impl Stage for SyncFromBrokerStage where EM: UsesState + EventFirer, - S: UsesInput + HasClientPerfMonitor + HasExecutions + HasCorpus + HasRand + HasMetadata, + S: UsesInput + + HasClientPerfMonitor + + HasExecutions + + HasCorpus + + HasRand + + HasMetadata + + HasTestcase, SP: ShMemProvider, E: HasObservers + Executor, for<'a> E::Observers: Deserialize<'a>, @@ -262,7 +271,7 @@ where last_id.map_or_else(|| state.corpus().first(), |id| state.corpus().next(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( state, diff --git a/libafl/src/state/mod.rs b/libafl/src/state/mod.rs index d87a6a9c3c..488fafdf15 100644 --- a/libafl/src/state/mod.rs +++ b/libafl/src/state/mod.rs @@ -228,7 +228,7 @@ pub trait HasTestcase: HasCorpus { /// To get mutable testcase fn testcase_mut( - &mut self, + &self, id: CorpusId, ) -> Result::Input>>, Error> { Ok(self.corpus().get(id)?.borrow_mut())