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

associated type for some traits

Browse Source
This commit is contained in:
Andrea Fioraldi 2020-11-02 19:12:28 +01:00
parent e42f13bba6
commit 0f4d4d9795
8 changed files with 226 additions and 142 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

132
src/corpus/mod.rs
View File

@ -6,29 +6,30 @@ use crate::inputs::Input;
use crate::AflError;
use std::path::PathBuf;
use std::marker::PhantomData;
pub trait HasEntriesVec<InputT: Input> {
pub trait HasEntriesVec<I> where I: Input {
/// Get the entries vector field
fn entries(&self) -> &Vec<Box<Testcase<InputT>>>;
fn entries(&self) -> &Vec<Box<Testcase<I>>>;
/// Get the entries vector field (mutable)
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<InputT>>>;
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>;
}
/// Corpus with all current testcases
pub trait Corpus<InputT: Input, RandT: Rand> : HasEntriesVec<InputT> + HasRand<RandT> {
pub trait Corpus<I> : HasEntriesVec<I> + HasRand where I: Input {
/// Returns the number of elements
fn count(&self) -> usize {
self.entries().len()
}
/// Add an entry to the corpus
fn add(&mut self, mut entry: Box<Testcase<InputT>>) {
fn add(&mut self, mut entry: Box<Testcase<I>>) {
self.entries_mut().push(entry);
}
/// Removes an entry from the corpus, returning it if it was present.
fn remove(&mut self, entry: &Testcase<InputT>) -> Option<Box<Testcase<InputT>>> {
fn remove(&mut self, entry: &Testcase<I>) -> Option<Box<Testcase<I>>> {
let mut i: usize = 0;
let mut found = false;
for x in self.entries() {
@ -45,46 +46,49 @@ pub trait Corpus<InputT: Input, RandT: Rand> : HasEntriesVec<InputT> + HasRand<R
}
/// Gets a random entry
fn random_entry(&mut self) -> Result<&Box<Testcase<InputT>>, AflError> {
let id = self.rand_mut().below(self.entries().len() as u64) as usize;
fn random_entry(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
let len = { self.entries().len() };
let id = self.rand_mut().below(len as u64) as usize;
Ok(self.entries_mut().get_mut(id).unwrap())
}
/// Gets the next entry (random by default)
fn get(&mut self) -> Result<&Box<Testcase<InputT>>, AflError> {
fn get(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
self.random_entry()
}
}
pub struct InMemoryCorpus<'a, InputT: Input, RandT: Rand> {
rand: &'a mut RandT,
entries: Vec<Box<Testcase<InputT>>>,
pub struct InMemoryCorpus<'a, I, R> where I: Input, R: Rand {
rand: &'a mut R,
entries: Vec<Box<Testcase<I>>>
}
impl<InputT: Input, RandT: Rand> HasEntriesVec<InputT> for InMemoryCorpus<'_, InputT, RandT> {
fn entries(&self) -> &Vec<Box<Testcase<InputT>>> {
impl<I, R> HasEntriesVec<I> for InMemoryCorpus<'_, I, R> where I: Input, R: Rand {
fn entries(&self) -> &Vec<Box<Testcase<I>>> {
&self.entries
}
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<InputT>>>{
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{
&mut self.entries
}
}
impl<InputT: Input, RandT: Rand> HasRand<RandT> for InMemoryCorpus<'_, InputT, RandT> {
fn rand(&self) -> &Box<dyn Rand> {
impl<I, R> HasRand for InMemoryCorpus<'_, I, R> where I: Input, R: Rand {
type R = R;
fn rand(&self) -> &Self::R {
&self.rand
}
fn rand_mut(&mut self) -> &mut Box<dyn Rand> {
fn rand_mut(&mut self) -> &mut Self::R {
&mut self.rand
}
}
impl<InputT: Input, RandT: Rand> Corpus<InputT, RandT> for InMemoryCorpus<'_, InputT, RandT> {
impl<I, R> Corpus<I> for InMemoryCorpus<'_, I, R> where I: Input, R: Rand {
// Just use the default implementation
}
impl<InputT: Input, RandT: Rand> InMemoryCorpus<'_, InputT, RandT> {
pub fn new<'a>(rand: &'a mut RandT) -> Self {
impl<'a, I, R> InMemoryCorpus<'a, I, R> where I: Input, R: Rand {
pub fn new(rand: &'a mut R) -> Self {
InMemoryCorpus {
rand: rand,
entries: vec![],
@ -92,38 +96,40 @@ impl<InputT: Input, RandT: Rand> InMemoryCorpus<'_, InputT, RandT> {
}
}
pub struct OnDiskCorpus<'a, InputT: Input, RandT: Rand> {
rand: &'a mut RandT,
entries: Vec<Box<Testcase<InputT>>>,
pub struct OnDiskCorpus<'a, I, R> where I: Input, R: Rand {
rand: &'a mut R,
entries: Vec<Box<Testcase<I>>>,
dir_path: PathBuf,
}
impl<InputT: Input, RandT: Rand> HasEntriesVec<InputT> for OnDiskCorpus<'_, InputT, RandT> {
fn entries(&self) -> &Vec<Box<Testcase<InputT>>> {
impl<I, R> HasEntriesVec<I> for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
fn entries(&self) -> &Vec<Box<Testcase<I>>> {
&self.entries
}
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<InputT>>>{
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{
&mut self.entries
}
}
impl<InputT: Input, RandT: Rand> HasRand<RandT> for OnDiskCorpus<'_, InputT, RandT> {
fn rand(&self) -> &Box<dyn Rand> {
impl<I, R> HasRand for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
type R = R;
fn rand(&self) -> &Self::R {
&self.rand
}
fn rand_mut(&mut self) -> &mut Box<dyn Rand> {
fn rand_mut(&mut self) -> &mut Self::R {
&mut self.rand
}
}
impl<InputT: Input, RandT: Rand> Corpus<InputT, RandT> for OnDiskCorpus<'_, InputT, RandT> {
impl<I, R> Corpus<I> for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
/// Add an entry and save it to disk
fn add(&mut self, mut entry: Box<Testcase<InputT>>) {
if entry.filename() == None {
fn add(&mut self, mut entry: Box<Testcase<I>>) {
if *entry.filename() == None {
// TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL)
let filename = &(String::from("id:") + &self.entries.len().to_string());
let filename = self.dir_path.join(filename);
entry.filename_mut() = filename;
*entry.filename_mut() = Some(filename);
}
self.entries.push(entry);
}
@ -131,8 +137,8 @@ impl<InputT: Input, RandT: Rand> Corpus<InputT, RandT> for OnDiskCorpus<'_, Inpu
// TODO save and remove files, cache, etc..., ATM use just InMemoryCorpus
}
impl<InputT: Input, RandT: Rand> OnDiskCorpus<'_, InputT, RandT> {
pub fn new<'a>(rand: &'a mut RandT, dir_path: PathBuf) -> Self {
impl<'a, I, R> OnDiskCorpus<'a, I, R> where I: Input, R: Rand {
pub fn new(rand: &'a mut R, dir_path: PathBuf) -> Self {
OnDiskCorpus {
dir_path: dir_path,
entries: vec![],
@ -142,50 +148,72 @@ impl<InputT: Input, RandT: Rand> OnDiskCorpus<'_, InputT, RandT> {
}
/// A Queue-like corpus, wrapping an existing Corpus instance
pub struct QueueCorpus<'a, InputT: Input, RandT: Rand, CorpusT: Corpus<InputT, RandT>> {
corpus: CorpusT,
pub struct QueueCorpus<I, C> where I: Input, C: Corpus<I> {
corpus: C,
phantom: PhantomData<I>,
pos: usize,
cycles: u64,
}
impl<'a, InputT: Input, RandT: Rand, CorpusT: Corpus<InputT, RandT>> Corpus<InputT, RandT> for QueueCorpus<'_, InputT, RandT, CorpusT> {
impl<'a, I, C> HasEntriesVec<I> for QueueCorpus<I, C> where I: Input, C: Corpus<I> {
fn entries(&self) -> &Vec<Box<Testcase<I>>> {
self.corpus.entries()
}
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{
self.corpus.entries_mut()
}
}
impl<'a, I, C> HasRand for QueueCorpus<I, C> where I: Input, C: Corpus<I> {
type R = C::R;
fn rand(&self) -> &Self::R {
self.corpus.rand()
}
fn rand_mut(&mut self) -> &mut Self::R {
self.corpus.rand_mut()
}
}
impl<'a, I, C> Corpus<I> for QueueCorpus<I, C> where I: Input, C: Corpus<I> {
/// Returns the number of elements
fn count(&self) -> usize {
self.corpus.count()
}
fn add(&mut self, entry: Box<Testcase<InputT>>) {
fn add(&mut self, entry: Box<Testcase<I>>) {
self.corpus.add(entry);
}
/// Removes an entry from the corpus, returning it if it was present.
fn remove(&mut self, entry: &Testcase<InputT>) -> Option<Box<Testcase<InputT>>> {
fn remove(&mut self, entry: &Testcase<I>) -> Option<Box<Testcase<I>>> {
self.corpus.remove(entry)
}
/// Gets a random entry
fn random_entry(&mut self) -> Result<&Box<Testcase<InputT>>, AflError> {
fn random_entry(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
self.corpus.random_entry()
}
/// Gets the next entry
fn get(&mut self) -> Result<&Box<Testcase<InputT>>, AflError> {
if self.count() == 0 {
fn get(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
if self.corpus.count() == 0 {
return Err(AflError::Empty("Testcases".to_string()));
}
self.pos = self.pos + 1;
if self.pos >= self.count() {
if self.pos >= self.corpus.count() {
self.cycles = self.cycles + 1;
self.pos = 0;
}
Ok(self.corpus.entries.get_mut(self.pos).unwrap())
Ok(self.corpus.entries_mut().get_mut(self.pos).unwrap())
}
}
impl<'a, InputT: Input, RandT: Rand, CorpusT: Corpus<InputT, RandT>> QueueCorpus<'_, InputT, RandT, CorpusT> {
pub fn new(corpus: CorpusT) -> Self {
QueueCorpus {
impl<'a, I, C> QueueCorpus<I, C> where I: Input, C: Corpus<I> {
pub fn new(corpus: C) -> Self {
QueueCorpus::<I, C> {
corpus: corpus,
phantom: PhantomData,
cycles: 0,
pos: 0,
}
@ -217,10 +245,10 @@ mod tests {
let mut q = QueueCorpus::new(OnDiskCorpus::new(&mut rand, PathBuf::from("fancy/path")));
let i = Box::new(BytesInput::new(vec![0; 4]));
let mut t = Box::new(Testcase::new(i));
t.set_filename(PathBuf::from("fancyfile"));
*t.filename_mut() = Some(PathBuf::from("fancyfile"));
q.add(t);
let filename = q.get().unwrap().get_filename().unwrap().to_owned();
assert_eq!(filename, q.get().unwrap().get_filename().unwrap().to_owned());
let filename = q.get().unwrap().filename().as_ref().unwrap().to_owned();
assert_eq!(filename, q.get().unwrap().filename().as_ref().unwrap().to_owned());
assert_eq!(filename, PathBuf::from("fancy/path/fancyfile"));
}
}

38
src/corpus/testcase.rs
View File

@ -7,15 +7,15 @@ use std::path::PathBuf;
pub trait TestcaseMetadata {}
/*
pub trait TestcaseTrait<InputT: Input> {
pub trait TestcaseTrait<I: Input> {
/// Make sure to return a valid input instance loading it from disk if not in memory
fn load_input(&mut self) -> Result<&Box<InputT>, AflError>;
fn load_input(&mut self) -> Result<&Box<I>, AflError>;
/// Get the input, if any
fn input(&self) -> &Option<Box<InputT>>;
fn input(&self) -> &Option<Box<I>>;
/// Get the input, if any (mutable)
fn input_mut(&mut self) -> &mut Option<Box<InputT>>;
fn input_mut(&mut self) -> &mut Option<Box<I>>;
/// Get the filename, if any
fn filename(&self) -> &Option<PathBuf>;
@ -29,38 +29,44 @@ pub trait TestcaseTrait<InputT: Input> {
*/
#[derive(Default)]
pub struct Testcase<InputT: Input> {
input: Option<Box<InputT>>,
pub struct Testcase<I> where I: Input {
input: Option<Box<I>>,
filename: Option<PathBuf>,
metadatas: HashMap<String, Box<dyn TestcaseMetadata>>,
}
impl<InputT: Input> Testcase<InputT> {
impl<I> Testcase<I> where I: Input {
/// Make sure to return a valid input instance loading it from disk if not in memory
pub fn load_input(&mut self) -> Result<&Box<InputT>, AflError> {
pub fn load_input(&mut self) -> Result<&Box<I>, AflError> {
// TODO: Implement cache to disk
self.input.as_ref().ok_or(AflError::NotImplemented("load_input".to_string()))
}
/// Get the input, if any
pub fn input(&self) -> &Option<Box<InputT>> {
pub fn input(&self) -> &Option<Box<I>> {
&self.input
}
/// Get the input, if any (mutable)
pub fn input_mut(&mut self) -> &mut Option<Box<InputT>> {
pub fn input_mut(&mut self) -> &mut Option<Box<I>> {
&mut self.input
}
/// Set the input
pub fn set_input(&mut self, input: Option<Box<I>>) {
self.input = input;
}
/// Get the filename, if any
pub fn filename(&self) -> &Option<PathBuf> {
&self.filename
}
/// Get the filename, if any (mutable)
pub fn filename_mut(&mut self, filename: PathBuf) -> &mut Option<PathBuf> {
pub fn filename_mut(&mut self) -> &mut Option<PathBuf> {
&mut self.filename
}
/// Set the filename
pub fn set_filename(&mut self, filename: Option<PathBuf>) {
self.filename = filename;
}
/// Get all the metadatas into an HashMap
pub fn metadatas(&mut self) -> &mut HashMap<String, Box<dyn TestcaseMetadata>> {
@ -68,7 +74,7 @@ impl<InputT: Input> Testcase<InputT> {
}
/// Create a new DefaultTestcase instace given an input
pub fn new(input: Box<InputT>) -> Self {
pub fn new(input: Box<I>) -> Self {
Testcase {
input: Some(input),
filename: None,
@ -77,10 +83,10 @@ impl<InputT: Input> Testcase<InputT> {
}
/// Create a new DefaultTestcase instace given an input and a filename
pub fn new_with_filename(input: Box<InputT>, filename: &PathBuf) -> Self {
pub fn new_with_filename(input: Box<I>, filename: PathBuf) -> Self {
Testcase {
input: Some(input),
filename: filename,
filename: Some(filename),
metadatas: HashMap::default(),
}
}

11
src/engines/aflengine.rs
View File

@ -1,18 +1,19 @@
use std::time;
use crate::engines::Engine;
use crate::inputs::Input;
use crate::executors::Executor;
use crate::feedbacks::Feedback;
use crate::monitors::Monitor;
use crate::stages::Stage;
use crate::utils::Rand;
pub struct AflEngine<'a> {
pub struct AflEngine<'a, InputT: Input> {
pub rand: &'a mut dyn Rand,
pub feedbacks: Vec<Box<dyn Feedback>>,
pub feedbacks: Vec<Box<dyn Feedback<InputT>>>,
pub stages: Vec<Box<dyn Stage>>,
pub current_stage: &'a Box<dyn Stage>,
pub stages: Vec<Box<dyn Stage<InputT>>>,
pub current_stage: &'a Box<dyn Stage<InputT>>,
pub executor: Box<dyn Executor>,
@ -25,4 +26,4 @@ pub struct AflEngine<'a> {
pub monitors: Vec<Box<dyn Monitor>>,
}
impl Engine<'_> for AflEngine<'_> {}
impl<InputT: Input> Engine<'_> for AflEngine<'_, InputT> {}

11
src/feedbacks/mod.rs
View File

@ -1,15 +1,22 @@
extern crate num;
use crate::corpus::Testcase;
use crate::inputs::Input;
use crate::executors::Executor;
use crate::observers::MapObserver;
use num::Integer;
use std::marker::PhantomData;
pub trait Feedback<InputT: Input> {
/// is_interesting should return the "Interestingness" from 0 to 255 (percent times 2.55)
fn is_interesting(&mut self, executor: &dyn Executor, entry: &Testcase<InputT>) -> u8;
}
/*
pub trait Feedback {
/// is_interesting should return the "Interestingness" from 0 to 255 (percent times 2.55)
fn is_interesting(&mut self, executor: &dyn Executor, entry: &dyn Testcase) -> u8;
fn is_interesting(&mut self, executor: &dyn Executor, entry: &Testcase) -> u8;
}
pub trait Reducer<T: Integer + Copy + 'static> {
@ -69,3 +76,5 @@ impl<'a, MapT: Integer + Copy + 'static, ReducerT: Reducer<MapT>> MapFeedback<Ma
type MaxMapFeedback<MapT> = MapFeedback<MapT, dyn MaxReducer<MapT>>;
#[allow(dead_code)]
type MinMapFeedback<MapT> = MapFeedback<MapT, dyn MinReducer<MapT>>;
*/

22
src/mutators/mod.rs
View File

@ -1,20 +1,30 @@
use crate::inputs::Input;
use crate::utils::{Rand, HasRand};
use crate::utils::HasRand;
use crate::corpus::Corpus;
use crate::AflError;
pub mod scheduled;
pub use scheduled::{ComposedByMutations, ScheduledMutator, HavocBytesMutator};
pub trait Mutator<InputT : Input, RandT: Rand> : HasRand<RandT> {
pub trait HasOptionCorpus<I> where I: Input {
type C : Corpus<I>;
/// Get the associated corpus, if any
fn corpus(&self) -> &Option<Box<Self::C>>;
/// Get the associated corpus, if any (mutable)
fn corpus_mut(&mut self) -> &mut Option<Box<Self::C>>;
/// Set the associated corpus
fn set_corpus(&mut self, corpus: Option<Box<Self::C>>);
}
pub trait Mutator<I> : HasRand + HasOptionCorpus<I> where I: Input {
/// Mutate a given input
fn mutate(&mut self, input: &mut InputT, stage_idx: i32) -> Result<(), AflError>;
fn mutate(&mut self, input: &mut I, stage_idx: i32) -> Result<(), AflError>;
/// Post-process given the outcome of the execution
fn post_exec(&mut self, _is_interesting: bool, _stage_idx: i32) -> Result<(), AflError> {
Ok(())
}
/// Get the associated corpus, if any
fn corpus(&mut self) -> &mut Option<Box<dyn Corpus<InputT, RandT>>>;
}

142
src/mutators/scheduled.rs
View File

@ -1,4 +1,4 @@
use crate::mutators::Mutator;
use crate::mutators::{HasOptionCorpus, Mutator};
use crate::utils::{Rand, HasRand};
use crate::corpus::Corpus;
use crate::inputs::{Input, HasBytesVec};
@ -7,27 +7,27 @@ use crate::AflError;
use std::marker::PhantomData;
/// The generic function type that identifies mutations
type MutationFunction<MutatorT, InputT> = fn(&mut MutatorT, &mut InputT) -> Result<(), AflError>;
type MutationFunction<M, I> = fn(&mut M, &mut I) -> Result<(), AflError>;
pub trait ComposedByMutations<InputT : Input> {
pub trait ComposedByMutations<I> where I: Input {
/// Get a mutation by index
fn mutation_by_idx(&self, index: usize) -> Result<MutationFunction<Self, InputT>, AflError>;
fn mutation_by_idx(&self, index: usize) -> Result<MutationFunction<Self, I>, AflError>;
/// Get the number of mutations
fn mutations_count(&self) -> usize;
/// Add a mutation
fn add_mutation(&mut self, mutation: MutationFunction<Self, InputT>);
fn add_mutation(&mut self, mutation: MutationFunction<Self, I>);
}
pub trait ScheduledMutator<InputT : Input, RandT: Rand>: Mutator<InputT, RandT> + ComposedByMutations<InputT> {
/// Computer the number of iterations used to apply stacked mutations
fn iterations(&mut self, _input: &InputT) -> u64 {
pub trait ScheduledMutator<I>: Mutator<I> + ComposedByMutations<I> where I: Input {
/// Compute the number of iterations used to apply stacked mutations
fn iterations(&mut self, _input: &I) -> u64 {
1 << (1 + self.rand_mut().below(7))
}
/// Get the next mutation to apply
fn schedule(&mut self, _input: &InputT) -> Result<MutationFunction<Self, InputT>, AflError> {
fn schedule(&mut self, _input: &I) -> Result<MutationFunction<Self, I>, AflError> {
let count = self.mutations_count() as u64;
if count == 0 {
return Err(AflError::Empty("no mutations".to_string()));
@ -41,7 +41,7 @@ pub trait ScheduledMutator<InputT : Input, RandT: Rand>: Mutator<InputT, RandT>
/// New default implementation for mutate
/// Implementations must forward mutate() to this method
fn scheduled_mutate(&mut self, input: &mut InputT, _stage_idx: i32) -> Result<(), AflError> {
fn scheduled_mutate(&mut self, input: &mut I, _stage_idx: i32) -> Result<(), AflError> {
let num = self.iterations(input);
for _ in 0..num {
self.schedule(input)?(self, input)?;
@ -50,23 +50,47 @@ pub trait ScheduledMutator<InputT : Input, RandT: Rand>: Mutator<InputT, RandT>
}
}
pub struct DefaultScheduledMutator<InputT : Input, RandT: Rand> {
rand: Box<RandT>,
corpus: Option<Box<dyn Corpus<InputT, RandT>>>,
mutations: Vec<MutationFunction<Self, InputT>>
pub struct DefaultScheduledMutator<'a, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
rand: &'a mut R,
corpus: Option<Box<C>>,
mutations: Vec<MutationFunction<Self, I>>
}
impl<InputT : Input, RandT: Rand> HasRand<RandT> for DefaultScheduledMutator<InputT, RandT> {
fn rand(&self) -> &Box<dyn Rand> {
impl<'a, I, R, C> HasRand for DefaultScheduledMutator<'_, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
type R = R;
fn rand(&self) -> &Self::R {
&self.rand
}
fn rand_mut(&mut self) -> &mut Box<dyn Rand> {
fn rand_mut(&mut self) -> &mut Self::R {
&mut self.rand
}
}
impl<InputT : Input, RandT: Rand> ComposedByMutations<InputT> for DefaultScheduledMutator<InputT, RandT> {
fn mutation_by_idx(&self, index: usize) -> Result<MutationFunction<Self, InputT>, AflError> {
impl<I, R, C> HasOptionCorpus<I> for DefaultScheduledMutator<'_, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
type C = C;
fn corpus(&self) -> &Option<Box<Self::C>> {
&self.corpus
}
fn corpus_mut(&mut self) -> &mut Option<Box<Self::C>> {
&mut self.corpus
}
fn set_corpus(&mut self, corpus: Option<Box<Self::C>>) {
self.corpus = corpus
}
}
impl<'a, I, R, C> Mutator<I> for DefaultScheduledMutator<'_, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
fn mutate(&mut self, input: &mut I, _stage_idx: i32) -> Result<(), AflError> {
self.scheduled_mutate(input, _stage_idx)
}
}
impl<'a, I, R, C> ComposedByMutations<I> for DefaultScheduledMutator<'_, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
fn mutation_by_idx(&self, index: usize) -> Result<MutationFunction<Self, I>, AflError> {
if index >= self.mutations.len() {
return Err(AflError::Unknown("oob".to_string()));
}
@ -77,28 +101,18 @@ impl<InputT : Input, RandT: Rand> ComposedByMutations<InputT> for DefaultSchedul
self.mutations.len()
}
fn add_mutation(&mut self, mutation: MutationFunction<Self, InputT>) {
fn add_mutation(&mut self, mutation: MutationFunction<Self, I>) {
self.mutations.push(mutation)
}
}
impl<InputT : Input, RandT: Rand> ScheduledMutator<InputT> for DefaultScheduledMutator<InputT, RandT> {
impl<'a, I, R, C> ScheduledMutator<I> for DefaultScheduledMutator<'_, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
// Just use the default methods
}
impl<InputT : Input, RandT: Rand> Mutator<InputT, RandT> for DefaultScheduledMutator<InputT, RandT> {
fn mutate(&mut self, input: &mut InputT, _stage_idx: i32) -> Result<(), AflError> {
self.scheduled_mutate(input, _stage_idx)
}
fn corpus(&mut self) -> &mut Option<Box<dyn Corpus>> {
&mut self.corpus
}
}
impl<InputT : Input, RandT: Rand> DefaultScheduledMutator<InputT> {
impl<'a, I, R, C> DefaultScheduledMutator<'a, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
/// Create a new DefaultScheduledMutator instance without mutations and corpus
pub fn new(rand: Box<RandT>) -> Self {
pub fn new(rand: &'a mut R) -> Self {
DefaultScheduledMutator {
rand: rand,
corpus: None,
@ -107,7 +121,7 @@ impl<InputT : Input, RandT: Rand> DefaultScheduledMutator<InputT> {
}
/// Create a new DefaultScheduledMutator instance specifying mutations and corpus too
pub fn new_all(rand: Box<RandT>, corpus: Option<Box<dyn Corpus>>, mutations: Vec<MutationFunction<Self, InputT>>) -> Self {
pub fn new_all(rand: &'a mut R, corpus: Option<Box<C>>, mutations: Vec<MutationFunction<Self, I>>) -> Self {
DefaultScheduledMutator {
rand: rand,
corpus: corpus,
@ -117,56 +131,70 @@ impl<InputT : Input, RandT: Rand> DefaultScheduledMutator<InputT> {
}
/// Bitflip mutation for inputs with a bytes vector
pub fn mutation_bitflip<MutatorT: Mutator<InputT, RandT>, InputT: Input + HasBytesVec, RandT: Rand>(mutator: &mut MutatorT, input: &mut InputT) -> Result<(), AflError> {
pub fn mutation_bitflip<M, I>(mutator: &mut M, input: &mut I) -> Result<(), AflError> where M: Mutator<I>, I: Input + HasBytesVec {
let bit = mutator.rand_mut().below(input.bytes().len() as u64) as usize;
input.bytes_mut()[bit >> 3] ^= (128 >> (bit & 7)) as u8;
Ok(())
}
/// Schedule some selected byte level mutations given a ScheduledMutator type
pub struct HavocBytesMutator<InputT: Input + HasBytesVec, RandT: Rand, ScheduledMutatorT: ScheduledMutator<InputT, RandT>> {
scheduled: ScheduledMutatorT,
_phantom: PhantomData<InputT>
pub struct HavocBytesMutator<I, S> where I: Input + HasBytesVec, S: ScheduledMutator<I> {
scheduled: S,
phantom: PhantomData<I>
}
impl<InputT: Input + HasBytesVec, RandT: Rand, ScheduledMutatorT: ScheduledMutator<InputT, RandT>> Mutator<InputT, RandT> for HavocBytesMutator<InputT, RandT, ScheduledMutatorT> {
fn mutate(&mut self, input: &mut InputT, stage_idx: i32) -> Result<(), AflError> {
self.scheduled.mutate(input, stage_idx)
}
impl<I, S> HasRand for HavocBytesMutator<I, S> where I: Input + HasBytesVec, S: ScheduledMutator<I> {
type R = S::R;
fn corpus(&mut self) -> &mut Option<Box<dyn Corpus>> {
self.scheduled.corpus()
}
}
impl<InputT: Input + HasBytesVec, RandT: Rand, ScheduledMutatorT: ScheduledMutator<InputT, RandT>> HasRand<RandT> for HavocBytesMutator<InputT, RandT, ScheduledMutatorT> {
fn rand(&self) -> &Box<dyn Rand> {
fn rand(&self) -> &Self::R {
self.scheduled.rand()
}
fn rand_mut(&mut self) -> &mut Box<dyn Rand> {
fn rand_mut(&mut self) -> &mut Self::R {
self.scheduled.rand_mut()
}
}
impl<InputT: Input + HasBytesVec, RandT: Rand, ScheduledMutatorT: ScheduledMutator<InputT, RandT>> HavocBytesMutator<InputT, RandT, ScheduledMutatorT> {
impl<I, S> HasOptionCorpus<I> for HavocBytesMutator<I, S> where I: Input + HasBytesVec, S: ScheduledMutator<I> {
type C = S::C;
fn corpus(&self) -> &Option<Box<Self::C>> {
self.scheduled.corpus()
}
fn corpus_mut(&mut self) -> &mut Option<Box<Self::C>> {
self.scheduled.corpus_mut()
}
fn set_corpus(&mut self, corpus: Option<Box<Self::C>>) {
self.scheduled.set_corpus(corpus)
}
}
impl<I, S> Mutator<I> for HavocBytesMutator<I, S> where I: Input + HasBytesVec, S: ScheduledMutator<I> {
fn mutate(&mut self, input: &mut I, stage_idx: i32) -> Result<(), AflError> {
self.scheduled.mutate(input, stage_idx)
}
}
impl<I, S> HavocBytesMutator<I, S> where I: Input + HasBytesVec, S: ScheduledMutator<I> {
/// Create a new HavocBytesMutator instance given a ScheduledMutator to wrap
pub fn new(mut scheduled: ScheduledMutatorT) -> Self {
pub fn new(mut scheduled: S) -> Self {
scheduled.add_mutation(mutation_bitflip);
HavocBytesMutator {
scheduled: scheduled,
_phantom: PhantomData
phantom: PhantomData
}
}
}
impl<InputT: Input + HasBytesVec, RandT: Rand> HavocBytesMutator<InputT, RandT, DefaultScheduledMutator<InputT, RandT>> {
impl<'a, I, R, C> HavocBytesMutator<I, DefaultScheduledMutator<'a, I, R, C>> where I: Input + HasBytesVec, R: Rand, C: Corpus<I> {
/// Create a new HavocBytesMutator instance wrapping DefaultScheduledMutator
pub fn new_default(rand: Box<RandT>) -> Self {
let mut scheduled = DefaultScheduledMutator::new(rand);
pub fn new_default(rand: &'a mut R) -> Self {
let mut scheduled = DefaultScheduledMutator::<'a, I, R, C>::new(rand);
scheduled.add_mutation(mutation_bitflip);
HavocBytesMutator {
scheduled: scheduled,
_phantom: PhantomData
phantom: PhantomData
}
}
}

4
src/stages/mod.rs
View File

@ -2,6 +2,6 @@ use crate::corpus::Testcase;
use crate::inputs::Input;
use crate::AflError;
/// Stages
pub trait Stage {
fn perform(&mut self, input: &dyn Input, entry: &mut Testcase) -> Result<(), AflError>;
pub trait Stage<InputT: Input> {
fn perform(&mut self, input: &dyn Input, entry: &mut Testcase<InputT>) -> Result<(), AflError>;
}

8
src/utils.rs
View File

@ -40,12 +40,14 @@ pub trait Rand: Debug {
}
/// Has a Rand box field
pub trait HasRand<RandT: Rand> {
pub trait HasRand {
type R : Rand;
/// Get the hold Rand instance
fn rand(&self) -> &Box<RandT>;
fn rand(&self) -> &Self::R;
/// Get the hold Rand instance (mutable)
fn rand_mut(&mut self) -> &mut Box<RandT>;
fn rand_mut(&mut self) -> &mut Self::R;
}
const HASH_CONST: u64 = 0xa5b35705;