corpus and testcase added

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<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");
+        
+    }
+}

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<dyn Rand>,
+    pub rand: &'a mut dyn Rand,
-    //pub feedback: Vec<Box<dyn Feedback>>,
+    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<'_> {
 
 }

src/engines/mod.rs

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

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