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

Fix for #1881, ProbabilitySamplingScheduler seems to be inverted (#1884)

Browse Source 
* Fix for #1881, ProbabilitySamplingScheduler seems to be inverted

* Run 'cargo +nightly fmt'

* Added debug_assert as negative and infinite probability values would break the current implementation of next()

---------

Co-authored-by: Dan Blackwell <danblackwell@eduroam-int-dhcp-97-192-218.ucl.ac.uk>
Co-authored-by: Dan Blackwell <danblackwell@ML-W0NYVW4XCK.local>
Co-authored-by: Dongjia "toka" Zhang <tokazerkje@outlook.com>
This commit is contained in:
Dan Blackwell 2024-03-05 18:04:13 +00:00 committed by GitHub
parent e3f837d712
commit cd8003ee69
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 46 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
libafl/src/schedulers/probabilistic_sampling.rs
View File

@ -9,9 +9,9 @@ use libafl_bolts::rands::Rand;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::{ use crate::{
corpus::{Corpus, CorpusId, HasTestcase}, corpus::{Corpus, CorpusId, HasTestcase, Testcase},
inputs::UsesInput, inputs::UsesInput,
schedulers::{Scheduler, TestcaseScore}, schedulers::{RemovableScheduler, Scheduler, TestcaseScore},
state::{HasCorpus, HasMetadata, HasRand, State, UsesState}, state::{HasCorpus, HasMetadata, HasRand, State, UsesState},
Error, Error,
}; };
@ -74,23 +74,60 @@ where
#[allow(clippy::cast_precision_loss)] #[allow(clippy::cast_precision_loss)]
#[allow(clippy::unused_self)] #[allow(clippy::unused_self)]
pub fn store_probability(&self, state: &mut S, idx: CorpusId) -> Result<(), Error> { pub fn store_probability(&self, state: &mut S, idx: CorpusId) -> Result<(), Error> {
let factor = F::compute(state, &mut *state.corpus().get(idx)?.borrow_mut())?; let prob = F::compute(state, &mut *state.corpus().get(idx)?.borrow_mut())?;
if factor == 0.0 { debug_assert!(
return Err(Error::illegal_state( prob >= 0.0 && prob.is_finite(),
"Infinity probability calculated for probabilistic sampling scheduler", "scheduler probability is {prob}; to work correctly it must be >= 0.0 and finite"
)); );
}
let meta = state let meta = state
.metadata_map_mut() .metadata_map_mut()
.get_mut::<ProbabilityMetadata>() .get_mut::<ProbabilityMetadata>()
.unwrap(); .unwrap();
let prob = 1.0 / factor;
meta.map.insert(idx, prob); meta.map.insert(idx, prob);
meta.total_probability += prob; meta.total_probability += prob;
Ok(()) Ok(())
} }
} }
impl<F, S> RemovableScheduler for ProbabilitySamplingScheduler<F, S>
where
F: TestcaseScore<S>,
S: HasCorpus + HasMetadata + HasRand + HasTestcase + State,
{
fn on_remove(
&mut self,
state: &mut Self::State,
idx: CorpusId,
_testcase: &Option<Testcase<<Self::State as UsesInput>::Input>>,
) -> Result<(), Error> {
let meta = state
.metadata_map_mut()
.get_mut::<ProbabilityMetadata>()
.unwrap();
if let Some(prob) = meta.map.remove(&idx) {
meta.total_probability -= prob;
}
Ok(())
}
fn on_replace(
&mut self,
state: &mut Self::State,
idx: CorpusId,
_prev: &Testcase<<Self::State as UsesInput>::Input>,
) -> Result<(), Error> {
let meta = state
.metadata_map_mut()
.get_mut::<ProbabilityMetadata>()
.unwrap();
if let Some(prob) = meta.map.remove(&idx) {
meta.total_probability -= prob;
}
self.store_probability(state, idx)
}
}
impl<F, S> UsesState for ProbabilitySamplingScheduler<F, S> impl<F, S> UsesState for ProbabilitySamplingScheduler<F, S>
where where
S: State + HasTestcase, S: State + HasTestcase,