corpus and testcase added
This commit is contained in:
parent
bbaa0829a2
commit
445ae1e122
186
src/corpus/mod.rs
Normal file
186
src/corpus/mod.rs
Normal 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");
|
||||
|
||||
}
|
||||
}
|
@ -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<'_> {
|
||||
|
||||
}
|
@ -1,5 +1,5 @@
|
||||
pub mod aflengine;
|
||||
|
||||
pub trait Engine {
|
||||
pub trait Engine<'a> {
|
||||
|
||||
}
|
23
src/utils.rs
23
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);
|
||||
|
Loading…
x
Reference in New Issue
Block a user