corpus and testcase added

This commit is contained in:
Dominik Maier 2020-10-26 05:40:47 +01:00
parent bbaa0829a2
commit 445ae1e122
4 changed files with 206 additions and 15 deletions

186
src/corpus/mod.rs Normal file
View File

@ -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<dyn Input>, AflError>;
fn is_on_disk(&self) -> bool;
fn get_filename(&self) -> &str;
fn get_metadatas(&mut self) -> &mut HashMap<String, Box<dyn TestcaseMetadata>>;
}
/// Corpus with all current testcases
pub trait Corpus: Debug {
/// Returns the number of elements
fn count(&self) -> usize;
fn add(&mut self, entry: Box<dyn Testcase>);
/// Removes an entry from the corpus, returning it if it was present.
fn remove(&mut self, entry: &dyn Testcase) -> Option<Box<dyn Testcase>>;
/// Gets a random entry
fn random_entry(&mut self) -> Result<&Box<dyn Testcase>, AflError>;
/// Gets the next entry
fn get(&mut self) -> Result<&Box<dyn Testcase>, 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<Box<dyn Testcase>>,
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<dyn Testcase>) {
self.entries.push(entry);
}
/// Removes an entry from the corpus, returning it if it was present.
fn remove(&mut self, entry: &dyn Testcase) -> Option<Box<dyn Testcase>> {
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<dyn Testcase>, 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<dyn Testcase>, 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<String, Box<dyn TestcaseMetadata>>,
}
impl Testcase for SimpleTestcase {
fn load_input(&mut self) -> Result<&Box<dyn Input>, 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<String, Box<dyn TestcaseMetadata>> {
&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");
}
}

View File

@ -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<dyn Rand>,
//pub feedback: Vec<Box<dyn Feedback>>,
pub rand: &'a mut dyn Rand,
pub feedbacks: Vec<Box<dyn Feedback>>,
pub stages: Vec<Box<dyn Stage>>,
pub current_stage: Box<dyn Stage>,
pub current_stage: &'a Box<dyn Stage>,
pub executor: Box<dyn Executor>,
@ -27,6 +27,6 @@ pub struct AflEngine {
}
impl Engine for AflEngine {
impl Engine<'_> for AflEngine<'_> {
}

View File

@ -1,5 +1,5 @@
pub mod aflengine;
pub trait Engine {
pub trait Engine<'a> {
}

View File

@ -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);