pass Testcase inside Rc+RefCell

This commit is contained in:
Andrea Fioraldi 2020-11-04 15:19:59 +01:00
parent fdf566bda4
commit 2e711360c1
11 changed files with 145 additions and 57 deletions

View File

@ -1,5 +1,5 @@
pub mod testcase; pub mod testcase;
pub use testcase::Testcase; pub use testcase::{Testcase, TestcaseMetadata};
use crate::utils::{Rand, HasRand}; use crate::utils::{Rand, HasRand};
use crate::inputs::Input; use crate::inputs::Input;
@ -7,13 +7,15 @@ use crate::AflError;
use std::path::PathBuf; use std::path::PathBuf;
use std::marker::PhantomData; use std::marker::PhantomData;
use std::cell::RefCell;
use std::rc::Rc;
pub trait HasEntriesVec<I> where I: Input { pub trait HasEntriesVec<I> where I: Input {
/// Get the entries vector field /// Get the entries vector field
fn entries(&self) -> &Vec<Box<Testcase<I>>>; fn entries(&self) -> &Vec<Rc<RefCell<Testcase<I>>>>;
/// Get the entries vector field (mutable) /// Get the entries vector field (mutable)
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>; fn entries_mut(&mut self) -> &mut Vec<Rc<RefCell<Testcase<I>>>>;
} }
/// Corpus with all current testcases /// Corpus with all current testcases
@ -25,17 +27,17 @@ pub trait Corpus<I> : HasEntriesVec<I> + HasRand where I: Input {
/// Add an entry to the corpus /// Add an entry to the corpus
#[allow(unused_mut)] #[allow(unused_mut)]
fn add(&mut self, mut entry: Box<Testcase<I>>) { fn add(&mut self, mut entry: Rc<RefCell<Testcase<I>>>) {
self.entries_mut().push(entry); self.entries_mut().push(entry);
} }
/// Removes an entry from the corpus, returning it if it was present. /// Removes an entry from the corpus, returning it if it was present.
fn remove(&mut self, entry: &Testcase<I>) -> Option<Box<Testcase<I>>> { fn remove(&mut self, entry: &Testcase<I>) -> Option<Rc<RefCell<Testcase<I>>>> {
let mut i: usize = 0; let mut i: usize = 0;
let mut found = false; let mut found = false;
for x in self.entries() { for x in self.entries() {
i = i + 1; i = i + 1;
if x.as_ref() as *const _ == entry as *const _ { if &*x.borrow() as *const _ == entry as *const _ { // TODO check if correct
found = true; found = true;
break; break;
} }
@ -47,28 +49,28 @@ pub trait Corpus<I> : HasEntriesVec<I> + HasRand where I: Input {
} }
/// Gets a random entry /// Gets a random entry
fn random_entry(&mut self) -> Result<&Box<Testcase<I>>, AflError> { fn random_entry(&mut self) -> Result<&Rc<RefCell<Testcase<I>>>, AflError> {
let len = { self.entries().len() }; let len = { self.entries().len() };
let id = self.rand_mut().below(len as u64) as usize; let id = self.rand_mut().below(len as u64) as usize;
Ok(self.entries_mut().get_mut(id).unwrap()) Ok(self.entries_mut().get_mut(id).unwrap())
} }
/// Gets the next entry (random by default) /// Gets the next entry (random by default)
fn get(&mut self) -> Result<&Box<Testcase<I>>, AflError> { fn get(&mut self) -> Result<&Rc<RefCell<Testcase<I>>>, AflError> {
self.random_entry() self.random_entry()
} }
} }
pub struct InMemoryCorpus<'a, I, R> where I: Input, R: Rand { pub struct InMemoryCorpus<'a, I, R> where I: Input, R: Rand {
rand: &'a mut R, rand: &'a mut R,
entries: Vec<Box<Testcase<I>>> entries: Vec<Rc<RefCell<Testcase<I>>>>
} }
impl<I, R> HasEntriesVec<I> for InMemoryCorpus<'_, I, R> where I: Input, R: Rand { impl<I, R> HasEntriesVec<I> for InMemoryCorpus<'_, I, R> where I: Input, R: Rand {
fn entries(&self) -> &Vec<Box<Testcase<I>>> { fn entries(&self) -> &Vec<Rc<RefCell<Testcase<I>>>> {
&self.entries &self.entries
} }
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{ fn entries_mut(&mut self) -> &mut Vec<Rc<RefCell<Testcase<I>>>>{
&mut self.entries &mut self.entries
} }
} }
@ -99,15 +101,15 @@ impl<'a, I, R> InMemoryCorpus<'a, I, R> where I: Input, R: Rand {
pub struct OnDiskCorpus<'a, I, R> where I: Input, R: Rand { pub struct OnDiskCorpus<'a, I, R> where I: Input, R: Rand {
rand: &'a mut R, rand: &'a mut R,
entries: Vec<Box<Testcase<I>>>, entries: Vec<Rc<RefCell<Testcase<I>>>>,
dir_path: PathBuf, dir_path: PathBuf,
} }
impl<I, R> HasEntriesVec<I> for OnDiskCorpus<'_, I, R> where I: Input, R: Rand { impl<I, R> HasEntriesVec<I> for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
fn entries(&self) -> &Vec<Box<Testcase<I>>> { fn entries(&self) -> &Vec<Rc<RefCell<Testcase<I>>>> {
&self.entries &self.entries
} }
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{ fn entries_mut(&mut self) -> &mut Vec<Rc<RefCell<Testcase<I>>>>{
&mut self.entries &mut self.entries
} }
} }
@ -125,12 +127,12 @@ impl<I, R> HasRand for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
impl<I, R> Corpus<I> for OnDiskCorpus<'_, I, R> where I: Input, R: Rand { impl<I, R> Corpus<I> for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
/// Add an entry and save it to disk /// Add an entry and save it to disk
fn add(&mut self, mut entry: Box<Testcase<I>>) { fn add(&mut self, entry: Rc<RefCell<Testcase<I>>>) {
if *entry.filename() == None { if *entry.borrow().filename() == None {
// TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL) // 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 = &(String::from("id:") + &self.entries.len().to_string());
let filename = self.dir_path.join(filename); let filename = self.dir_path.join(filename);
*entry.filename_mut() = Some(filename); *entry.borrow_mut().filename_mut() = Some(filename);
} }
self.entries.push(entry); self.entries.push(entry);
} }
@ -157,10 +159,10 @@ pub struct QueueCorpus<I, C> where I: Input, C: Corpus<I> {
} }
impl<'a, I, C> HasEntriesVec<I> for QueueCorpus<I, C> where I: Input, C: Corpus<I> { impl<'a, I, C> HasEntriesVec<I> for QueueCorpus<I, C> where I: Input, C: Corpus<I> {
fn entries(&self) -> &Vec<Box<Testcase<I>>> { fn entries(&self) -> &Vec<Rc<RefCell<Testcase<I>>>> {
self.corpus.entries() self.corpus.entries()
} }
fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{ fn entries_mut(&mut self) -> &mut Vec<Rc<RefCell<Testcase<I>>>>{
self.corpus.entries_mut() self.corpus.entries_mut()
} }
} }
@ -182,22 +184,22 @@ impl<'a, I, C> Corpus<I> for QueueCorpus<I, C> where I: Input, C: Corpus<I> {
self.corpus.count() self.corpus.count()
} }
fn add(&mut self, entry: Box<Testcase<I>>) { fn add(&mut self, entry: Rc<RefCell<Testcase<I>>>) {
self.corpus.add(entry); self.corpus.add(entry);
} }
/// Removes an entry from the corpus, returning it if it was present. /// Removes an entry from the corpus, returning it if it was present.
fn remove(&mut self, entry: &Testcase<I>) -> Option<Box<Testcase<I>>> { fn remove(&mut self, entry: &Testcase<I>) -> Option<Rc<RefCell<Testcase<I>>>> {
self.corpus.remove(entry) self.corpus.remove(entry)
} }
/// Gets a random entry /// Gets a random entry
fn random_entry(&mut self) -> Result<&Box<Testcase<I>>, AflError> { fn random_entry(&mut self) -> Result<&Rc<RefCell<Testcase<I>>>, AflError> {
self.corpus.random_entry() self.corpus.random_entry()
} }
/// Gets the next entry /// Gets the next entry
fn get(&mut self) -> Result<&Box<Testcase<I>>, AflError> { fn get(&mut self) -> Result<&Rc<RefCell<Testcase<I>>>, AflError> {
if self.corpus.count() == 0 { if self.corpus.count() == 0 {
return Err(AflError::Empty("Testcases".to_string())); return Err(AflError::Empty("Testcases".to_string()));
} }
@ -238,18 +240,19 @@ mod tests {
use crate::utils::Xoshiro256StarRand; use crate::utils::Xoshiro256StarRand;
use std::path::PathBuf; use std::path::PathBuf;
use std::cell::RefCell;
use std::rc::Rc;
#[test] #[test]
fn test_queuecorpus() { fn test_queuecorpus() {
let mut rand = Xoshiro256StarRand::new(); let mut rand = Xoshiro256StarRand::new();
let mut q = QueueCorpus::new(OnDiskCorpus::new(&mut rand, PathBuf::from("fancy/path"))); let mut q = QueueCorpus::new(OnDiskCorpus::new(&mut rand, PathBuf::from("fancy/path")));
let i = Box::new(BytesInput::new(vec![0; 4])); let i = BytesInput::new(vec![0; 4]);
let mut t = Box::new(Testcase::new(i)); let t = Rc::new(RefCell::new(Testcase::new_with_filename(i, PathBuf::from("fancyfile"))));
*t.filename_mut() = Some(PathBuf::from("fancyfile"));
q.add(t); q.add(t);
let filename = q.get().unwrap().filename().as_ref().unwrap().to_owned(); let filename = q.get().unwrap().borrow().filename().as_ref().unwrap().to_owned();
assert_eq!(filename, q.get().unwrap().filename().as_ref().unwrap().to_owned()); assert_eq!(filename, q.get().unwrap().borrow().filename().as_ref().unwrap().to_owned());
assert_eq!(filename, PathBuf::from("fancy/path/fancyfile")); assert_eq!(filename, PathBuf::from("fancy/path/fancyfile"));
} }
} }

View File

@ -9,13 +9,13 @@ pub trait TestcaseMetadata {}
/* /*
pub trait TestcaseTrait<I: Input> { pub trait TestcaseTrait<I: Input> {
/// Make sure to return a valid input instance loading it from disk if not in memory /// Make sure to return a valid input instance loading it from disk if not in memory
fn load_input(&mut self) -> Result<&Box<I>, AflError>; fn load_input(&mut self) -> Result<&I, AflError>;
/// Get the input, if any /// Get the input, if any
fn input(&self) -> &Option<Box<I>>; fn input(&self) -> &Option<I>;
/// Get the input, if any (mutable) /// Get the input, if any (mutable)
fn input_mut(&mut self) -> &mut Option<Box<I>>; fn input_mut(&mut self) -> &mut Option<I>;
/// Get the filename, if any /// Get the filename, if any
fn filename(&self) -> &Option<PathBuf>; fn filename(&self) -> &Option<PathBuf>;
@ -30,28 +30,31 @@ pub trait TestcaseTrait<I: Input> {
#[derive(Default)] #[derive(Default)]
pub struct Testcase<I> where I: Input { pub struct Testcase<I> where I: Input {
input: Option<Box<I>>, input: Option<I>, // TODO remove box
filename: Option<PathBuf>, filename: Option<PathBuf>,
metadatas: HashMap<String, Box<dyn TestcaseMetadata>>, metadatas: HashMap<String, Box<dyn TestcaseMetadata>>,
} }
impl<I> Testcase<I> where I: Input { impl<I> Testcase<I> where I: Input {
/// Make sure to return a valid input instance loading it from disk if not in memory /// Make sure to return a valid input instance loading it from disk if not in memory
pub fn load_input(&mut self) -> Result<&Box<I>, AflError> { pub fn load_input(&mut self) -> Result<&I, AflError> {
// TODO: Implement cache to disk // TODO: Implement cache to disk
self.input.as_ref().ok_or(AflError::NotImplemented("load_input".to_string())) match self.input.as_ref() {
Some(i) => Ok(i),
None => Err(AflError::NotImplemented("load_input".to_string()))
}
} }
/// Get the input, if any /// Get the input, if any
pub fn input(&self) -> &Option<Box<I>> { pub fn input(&self) -> &Option<I> {
&self.input &self.input
} }
/// Get the input, if any (mutable) /// Get the input, if any (mutable)
pub fn input_mut(&mut self) -> &mut Option<Box<I>> { pub fn input_mut(&mut self) -> &mut Option<I> {
&mut self.input &mut self.input
} }
/// Set the input /// Set the input
pub fn set_input(&mut self, input: Option<Box<I>>) { pub fn set_input(&mut self, input: Option<I>) {
self.input = input; self.input = input;
} }
@ -74,7 +77,7 @@ impl<I> Testcase<I> where I: Input {
} }
/// Create a new DefaultTestcase instace given an input /// Create a new DefaultTestcase instace given an input
pub fn new(input: Box<I>) -> Self { pub fn new(input: I) -> Self {
Testcase { Testcase {
input: Some(input), input: Some(input),
filename: None, filename: None,
@ -83,7 +86,7 @@ impl<I> Testcase<I> where I: Input {
} }
/// Create a new DefaultTestcase instace given an input and a filename /// Create a new DefaultTestcase instace given an input and a filename
pub fn new_with_filename(input: Box<I>, filename: PathBuf) -> Self { pub fn new_with_filename(input: I, filename: PathBuf) -> Self {
Testcase { Testcase {
input: Some(input), input: Some(input),
filename: Some(filename), filename: Some(filename),

View File

@ -7,7 +7,7 @@ use crate::feedbacks::Feedback;
use crate::monitors::Monitor; use crate::monitors::Monitor;
use crate::stages::Stage; use crate::stages::Stage;
use crate::utils::Rand; use crate::utils::Rand;
/*
pub struct AflEngine<'a, I: Input> { pub struct AflEngine<'a, I: Input> {
pub rand: &'a mut dyn Rand, pub rand: &'a mut dyn Rand,
pub feedbacks: Vec<Box<dyn Feedback<I>>>, pub feedbacks: Vec<Box<dyn Feedback<I>>>,
@ -27,3 +27,4 @@ pub struct AflEngine<'a, I: Input> {
} }
impl<I: Input> Engine<'_> for AflEngine<'_, I> {} impl<I: Input> Engine<'_> for AflEngine<'_, I> {}
*/

View File

@ -1,3 +1,14 @@
pub mod aflengine; pub mod aflengine;
pub trait Engine<'a> {} use crate::AflError;
use crate::inputs::Input;
use crate::corpus::testcase::Testcase;
use std::cell::RefCell;
use std::rc::Rc;
pub trait Engine<'a, I> where I: Input {
fn execute(&mut self, input: &mut I, entry: Rc<RefCell<Testcase<I>>>) -> Result<bool, AflError>;
}

View File

@ -177,6 +177,7 @@ mod tests {
use crate::observers::Observer; use crate::observers::Observer;
use crate::AflError; use crate::AflError;
#[derive(Clone)]
struct NopInput {} struct NopInput {}
impl Input for NopInput { impl Input for NopInput {
fn serialize(&self) -> Result<&[u8], AflError> { fn serialize(&self) -> Result<&[u8], AflError> {

View File

@ -1,5 +1,4 @@
pub mod inmemory; pub mod inmemory;
pub use inmemory::{InMemoryExecutor};
use crate::inputs::Input; use crate::inputs::Input;
use crate::observers::Observer; use crate::observers::Observer;

View File

@ -5,10 +5,11 @@ use std::fs::File;
use std::io::Read; use std::io::Read;
use std::io::Write; use std::io::Write;
use std::path::PathBuf; use std::path::PathBuf;
use std::clone::Clone;
use crate::AflError; use crate::AflError;
pub trait Input { pub trait Input : Clone {
fn to_file(&self, path: &PathBuf) -> Result<(), AflError> { fn to_file(&self, path: &PathBuf) -> Result<(), AflError> {
let mut file = File::create(path)?; let mut file = File::create(path)?;
file.write_all(self.serialize()?)?; file.write_all(self.serialize()?)?;

View File

@ -4,7 +4,6 @@ use crate::corpus::Corpus;
use crate::AflError; use crate::AflError;
pub mod scheduled; pub mod scheduled;
pub use scheduled::{ComposedByMutations, ScheduledMutator, HavocBytesMutator};
pub trait HasOptionCorpus<I> where I: Input { pub trait HasOptionCorpus<I> where I: Input {
type C : Corpus<I>; type C : Corpus<I>;
@ -19,7 +18,7 @@ pub trait HasOptionCorpus<I> where I: Input {
fn set_corpus(&mut self, corpus: Option<Box<Self::C>>); fn set_corpus(&mut self, corpus: Option<Box<Self::C>>);
} }
pub trait Mutator<I> : HasRand + HasOptionCorpus<I> where I: Input { pub trait Mutator<I> : HasRand where I: Input {
/// Mutate a given input /// Mutate a given input
fn mutate(&mut self, input: &mut I, stage_idx: i32) -> Result<(), AflError>; fn mutate(&mut self, input: &mut I, stage_idx: i32) -> Result<(), AflError>;

View File

@ -20,7 +20,7 @@ pub trait ComposedByMutations<I> where I: Input {
fn add_mutation(&mut self, mutation: MutationFunction<Self, I>); fn add_mutation(&mut self, mutation: MutationFunction<Self, I>);
} }
pub trait ScheduledMutator<I>: Mutator<I> + ComposedByMutations<I> where I: Input { pub trait ScheduledMutator<I>: Mutator<I> + HasOptionCorpus<I> + ComposedByMutations<I> where I: Input {
/// Compute the number of iterations used to apply stacked mutations /// Compute the number of iterations used to apply stacked mutations
fn iterations(&mut self, _input: &I) -> u64 { fn iterations(&mut self, _input: &I) -> u64 {
1 << (1 + self.rand_mut().below(7)) 1 << (1 + self.rand_mut().below(7))

View File

@ -1,7 +1,19 @@
pub mod mutational;
use crate::corpus::Testcase; use crate::corpus::Testcase;
use crate::inputs::Input; use crate::inputs::Input;
use crate::engines::Engine;
use crate::AflError; use crate::AflError;
/// Stages
pub trait Stage<InputT: Input> { pub trait HasEngine<'a, I> where I: Input {
fn perform(&mut self, input: &dyn Input, entry: &mut Testcase<InputT>) -> Result<(), AflError>; type E : Engine<'a, I>;
fn engine(&self) -> &Self::E;
fn engine_mut(&mut self) -> &mut Self::E;
}
pub trait Stage<'a, I> : HasEngine<'a, I> where I: Input {
/// Run the stage
fn perform(&mut self, entry: &mut Testcase<I>) -> Result<(), AflError>;
} }

View File

@ -1,16 +1,74 @@
use std::Vec; use crate::AflError;
use crate::mutators::Mutator; use crate::mutators::Mutator;
use crate::inputs::Input; use crate::inputs::Input;
use crate::utils::{Rand, HasRand};
use crate::stages::{Stage, HasEngine};
use crate::corpus::testcase::Testcase;
use crate::engines::Engine;
pub struct MutationalStage { use std::cell::RefCell;
mutators: Vec<Box<dyn Mutator>>; use std::rc::Rc;
}
impl Stage for MutationalStage { pub trait MutationalStage<'a, I> : Stage<'a, I> + HasRand where I: Input {
fn mutators(&self) -> &Vec<Box<dyn Mutator<I, R = Self::R>>>;
fn Perform(&mut self, input: &Input, entry: &mut Entry) -> Result<(), AflError> { fn mutators_mut(&mut self) -> &mut Vec<Box<dyn Mutator<I, R = Self::R>>>;
// TODO: Implement me
Err(AflError::NotImplemented("Stage does not perform yet")); fn add_mutator(&mut self, mutator: Box<dyn Mutator<I, R = Self::R>>) {
self.mutators_mut().push(mutator);
} }
fn iterations(&mut self) -> usize {
1 + self.rand_mut().below(128) as usize
}
fn perform_mutational(&mut self, entry: Rc<RefCell<Testcase<I>>>) -> Result<(), AflError> {
let num = self.iterations();
let mut input = entry.borrow_mut().load_input()?.clone();
for i in 0..num {
for m in self.mutators_mut() {
m.mutate(&mut input, i as i32)?;
}
let interesting = self.engine_mut().execute(&mut input, entry.clone())?;
for m in self.mutators_mut() {
m.post_exec(interesting, i as i32)?;
}
input = entry.borrow_mut().load_input()?.clone();
}
Ok(())
}
}
pub struct DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> {
rand: &'a mut R,
engine: &'a mut E,
mutators: Vec<Box<dyn Mutator<I, R = R>>>
}
impl<'a, I, R, E> HasRand for DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> {
type R = R;
fn rand(&self) -> &Self::R {
&self.rand
}
fn rand_mut(&mut self) -> &mut Self::R {
&mut self.rand
}
}
impl<'a, I, R, E> HasEngine<'a, I> for DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> {
type E = E;
fn engine(&self) -> &Self::E {
self.engine
}
fn engine_mut(&mut self) -> &mut Self::E {
self.engine
}
} }