From 445ae1e122523ff3c56206e2fa9d3a7b9645f156 Mon Sep 17 00:00:00 2001 From: Dominik Maier Date: Mon, 26 Oct 2020 05:40:47 +0100 Subject: [PATCH] corpus and testcase added --- src/corpus/mod.rs | 186 +++++++++++++++++++++++++++++++++++++++ src/engines/aflengine.rs | 10 +-- src/engines/mod.rs | 2 +- src/utils.rs | 23 +++-- 4 files changed, 206 insertions(+), 15 deletions(-) create mode 100644 src/corpus/mod.rs diff --git a/src/corpus/mod.rs b/src/corpus/mod.rs new file mode 100644 index 0000000000..974bcf7313 --- /dev/null +++ b/src/corpus/mod.rs @@ -0,0 +1,186 @@ +use std::fmt::Debug; +use std::collections::HashMap; +use crate::AflError; +use crate::utils::Rand; +use crate::inputs::Input; + +pub trait TestcaseMetadata: Debug { + +} + +pub trait Testcase: Debug { + + fn load_input(&mut self) -> Result<&Box, AflError>; + fn is_on_disk(&self) -> bool; + fn get_filename(&self) -> &str; + fn get_metadatas(&mut self) -> &mut HashMap>; + +} + +/// Corpus with all current testcases +pub trait Corpus: Debug { + + /// Returns the number of elements + fn count(&self) -> usize; + + fn add(&mut self, entry: Box); + + /// Removes an entry from the corpus, returning it if it was present. + fn remove(&mut self, entry: &dyn Testcase) -> Option>; + + /// Gets a random entry + fn random_entry(&mut self) -> Result<&Box, AflError>; + + /// Gets the next entry + fn get(&mut self) -> Result<&Box, AflError>; + +} + +/// A queue-like corpus +pub trait Queue: Corpus { + + fn get_cycles(&self) -> u64; + + fn get_pos(&self) -> usize; + +} + +#[derive(Debug)] +pub struct DefaultQueue<'a> { + rand: &'a mut dyn Rand, + pos: usize, + cycles: u64, + entries: Vec>, + dir_path: String + +} + +impl Corpus for DefaultQueue<'_> { + + /// Returns the number of elements + fn count(&self) -> usize { + self.entries.len() + } + + fn add(&mut self, entry: Box) { + self.entries.push(entry); + } + + /// Removes an entry from the corpus, returning it if it was present. + fn remove(&mut self, entry: &dyn Testcase) -> Option> { + let mut i: usize = 0; + let mut found = false; + for x in &self.entries { + i = i + 1; + if x.as_ref().get_filename() == entry.get_filename() { + found = true; + break; + } + } + if !found { + return None; + } + Some(self.entries.remove(i)) + } + + /// Gets a random entry + fn random_entry(&mut self) -> Result<&Box, AflError> { + let id = self.rand.below(self.entries.len() as u64) as usize; + Ok(self.entries.get_mut(id).unwrap()) + } + + /// Gets the next entry + fn get(&mut self) -> Result<&Box, AflError> { + self.pos = self.pos + 1; + if self.pos > self.entries.len() { + self.cycles = self.cycles + 1; + self.pos = 0; + } + Ok(self.entries.get_mut(self.pos).unwrap()) + } +} + +impl Queue for DefaultQueue<'_> { + + fn get_cycles(&self) -> u64 { + self.cycles + } + + fn get_pos(&self) -> usize { + self.pos + } + +} + +impl DefaultQueue<'_> { + + pub fn new<'a>(rand: &'a mut dyn Rand, dir_path: &str) -> DefaultQueue<'a> { + DefaultQueue{ + cycles: 0, + dir_path: dir_path.to_owned(), + entries: vec![], + pos: 0, + rand: rand, + } + } + +} + +#[derive(Debug, Default)] +struct SimpleTestcase { + + is_on_disk: bool, + filename: String, + metadatas: HashMap>, + +} + +impl Testcase for SimpleTestcase { + + fn load_input(&mut self) -> Result<&Box, AflError> { + // TODO: Implement + Err(AflError::Unknown) + } + + fn is_on_disk(&self) -> bool { + self.is_on_disk + } + + fn get_filename(&self) -> &str { + &self.filename + } + + fn get_metadatas(&mut self) -> &mut HashMap> { + &mut self.metadatas + } + +} + +impl SimpleTestcase { + fn new(filename: &str) -> Self { + SimpleTestcase{ + filename: filename.to_owned(), + is_on_disk: false, + metadatas: HashMap::default() + } + } +} + +#[cfg(test)] +mod tests { + use crate::corpus::Corpus; + use crate::corpus::DefaultQueue; + use crate::corpus::SimpleTestcase; + use crate::utils::AflRand; + + #[test] + fn test_defaultqueue() { + let mut rand = AflRand::new(); + let mut q = DefaultQueue::new(&mut rand, "fancy/path"); + q.add(Box::new(SimpleTestcase::new("fancyfile"))); + let filename = q.get().unwrap().get_filename().to_owned(); + assert_eq!(filename, q.get().unwrap().get_filename()); + assert_eq!(filename, "fancyfile"); + + } +} \ No newline at end of file diff --git a/src/engines/aflengine.rs b/src/engines/aflengine.rs index f035008e98..132a7d145d 100644 --- a/src/engines/aflengine.rs +++ b/src/engines/aflengine.rs @@ -7,13 +7,13 @@ use crate::executors::Executor; use crate::engines::Engine; use crate::monitors::Monitor; -pub struct AflEngine { +pub struct AflEngine<'a> { - //pub rand: Box, - //pub feedback: Vec>, + pub rand: &'a mut dyn Rand, + pub feedbacks: Vec>, pub stages: Vec>, - pub current_stage: Box, + pub current_stage: &'a Box, pub executor: Box, @@ -27,6 +27,6 @@ pub struct AflEngine { } -impl Engine for AflEngine { +impl Engine<'_> for AflEngine<'_> { } \ No newline at end of file diff --git a/src/engines/mod.rs b/src/engines/mod.rs index be929fe9b2..4a64ecaabc 100644 --- a/src/engines/mod.rs +++ b/src/engines/mod.rs @@ -1,5 +1,5 @@ pub mod aflengine; -pub trait Engine { +pub trait Engine<'a> { } \ No newline at end of file diff --git a/src/utils.rs b/src/utils.rs index 875b49000f..2109930efd 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,13 +1,12 @@ //! Utility functions for AFL +use std::fmt::Debug; use std::debug_assert; use xxhrs::{XXH3_64}; /// Ways to get random around here -pub trait Rand { +pub trait Rand: Debug { - // Returns a new, randomly seeded, rand - fn new() -> Self; // Sets the seed of this Rand fn set_seed(&mut self, seed: u64); // Gets the next 64 bit value @@ -58,12 +57,7 @@ pub struct AflRand { impl Rand for AflRand { - fn new() -> AflRand { - let mut ret: AflRand = Default::default(); - ret.set_seed(0); // TODO: Proper random seed? - ret - } fn set_seed(&mut self, seed: u64) { self.rand_seed[0] = XXH3_64::hash_with_seed(HASH_CONST, &seed.to_le_bytes()); @@ -94,6 +88,17 @@ impl Rand for AflRand { } +impl AflRand { + + pub fn new() -> AflRand { + + let mut ret: AflRand = Default::default(); + ret.set_seed(0); // TODO: Proper random seed? + ret + } + +} + /// Get the next higher power of two fn next_pow2(val: u64) -> u64 { // Early exit so we don't have to do a wrapping subtract; @@ -115,7 +120,7 @@ mod tests { #[test] fn test_rand() { - let mut rand: AflRand = Rand::new(); + let mut rand = AflRand::new(); assert_ne!(rand.next(), rand.next()); assert!(rand.below(100) < 100); assert_eq!(rand.below(1), 0);