diff --git a/afl/src/corpus/inmemory.rs b/afl/src/corpus/inmemory.rs new file mode 100644 index 0000000000..b4cf3248ba --- /dev/null +++ b/afl/src/corpus/inmemory.rs @@ -0,0 +1,72 @@ +use alloc::vec::Vec; +use core::{cell::RefCell, marker::PhantomData}; +use serde::{Deserialize, Serialize}; + +use crate::{ + corpus::Corpus, corpus::HasTestcaseVec, corpus::Testcase, inputs::Input, utils::Rand, AflError, +}; + +/// A corpus handling all important fuzzing in memory. +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(bound = "I: serde::de::DeserializeOwned")] +pub struct InMemoryCorpus +where + I: Input, + R: Rand, +{ + entries: Vec>>, + pos: usize, + phantom: PhantomData, +} + +impl HasTestcaseVec for InMemoryCorpus +where + I: Input, + R: Rand, +{ + fn entries(&self) -> &[RefCell>] { + &self.entries + } + fn entries_mut(&mut self) -> &mut Vec>> { + &mut self.entries + } +} + +impl Corpus for InMemoryCorpus +where + I: Input, + R: Rand, +{ + /// Gets the next entry + #[inline] + fn next(&mut self, rand: &mut R) -> Result<(&RefCell>, usize), AflError> { + if self.count() == 0 { + Err(AflError::Empty("No entries in corpus".to_owned())) + } else { + let len = { self.entries().len() }; + let id = rand.below(len as u64) as usize; + self.pos = id; + Ok((self.get(id), id)) + } + } + + /// Returns the testacase we currently use + #[inline] + fn current_testcase(&self) -> (&RefCell>, usize) { + (self.get(self.pos), self.pos) + } +} + +impl InMemoryCorpus +where + I: Input, + R: Rand, +{ + pub fn new() -> Self { + Self { + entries: vec![], + pos: 0, + phantom: PhantomData, + } + } +} diff --git a/afl/src/corpus/mod.rs b/afl/src/corpus/mod.rs index 85aa9dc0d1..222e5dcf62 100644 --- a/afl/src/corpus/mod.rs +++ b/afl/src/corpus/mod.rs @@ -1,12 +1,19 @@ pub mod testcase; pub use testcase::Testcase; -use alloc::{borrow::ToOwned, vec::Vec}; -use core::{cell::RefCell, marker::PhantomData, ptr}; -use serde::{Deserialize, Serialize}; +pub mod inmemory; +pub use inmemory::InMemoryCorpus; #[cfg(feature = "std")] -use std::path::PathBuf; +pub mod ondisk; +#[cfg(feature = "std")] +pub use ondisk::OnDiskCorpus; + +pub mod queue; +pub use queue::QueueCorpus; + +use alloc::vec::Vec; +use core::{cell::RefCell, ptr}; use crate::{inputs::Input, utils::Rand, AflError}; @@ -84,27 +91,6 @@ where } } - /// Returns the testcase for the given idx, with loaded input - /*fn load_testcase(&mut self, idx: usize) -> Result<(), AflError> { - let testcase = self.get(idx); - // Ensure testcase is loaded - match testcase.input() { - None => { - let new_testcase = match testcase.filename() { - Some(filename) => Testcase::load_from_disk(filename)?, - None => { - return Err(AflError::IllegalState( - "Neither input, nor filename specified for testcase".into(), - )) - } - }; - - self.replace(idx, new_testcase)?; - } - _ => (), - } - Ok(()) - }*/ // TODO: IntoIter /// Gets the next entry fn next(&mut self, rand: &mut R) -> Result<(&RefCell>, usize), AflError>; @@ -112,309 +98,3 @@ where /// Returns the testacase we currently use fn current_testcase(&self) -> (&RefCell>, usize); } - -/// A corpus handling all important fuzzing in memory. -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(bound = "I: serde::de::DeserializeOwned")] -pub struct InMemoryCorpus -where - I: Input, - R: Rand, -{ - entries: Vec>>, - pos: usize, - phantom: PhantomData, -} - -impl HasTestcaseVec for InMemoryCorpus -where - I: Input, - R: Rand, -{ - fn entries(&self) -> &[RefCell>] { - &self.entries - } - fn entries_mut(&mut self) -> &mut Vec>> { - &mut self.entries - } -} - -impl Corpus for InMemoryCorpus -where - I: Input, - R: Rand, -{ - /// Gets the next entry - #[inline] - fn next(&mut self, rand: &mut R) -> Result<(&RefCell>, usize), AflError> { - if self.count() == 0 { - Err(AflError::Empty("No entries in corpus".to_owned())) - } else { - let len = { self.entries().len() }; - let id = rand.below(len as u64) as usize; - self.pos = id; - Ok((self.get(id), id)) - } - } - - /// Returns the testacase we currently use - #[inline] - fn current_testcase(&self) -> (&RefCell>, usize) { - (self.get(self.pos), self.pos) - } -} - -impl InMemoryCorpus -where - I: Input, - R: Rand, -{ - pub fn new() -> Self { - Self { - entries: vec![], - pos: 0, - phantom: PhantomData, - } - } -} - -/// A corpus able to store testcases to dis, and load them from disk, when they are being used. -#[cfg(feature = "std")] -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(bound = "I: serde::de::DeserializeOwned")] -pub struct OnDiskCorpus -where - I: Input, - R: Rand, -{ - entries: Vec>>, - dir_path: PathBuf, - pos: usize, - phantom: PhantomData, -} - -#[cfg(feature = "std")] -impl HasTestcaseVec for OnDiskCorpus -where - I: Input, - R: Rand, -{ - #[inline] - fn entries(&self) -> &[RefCell>] { - &self.entries - } - #[inline] - fn entries_mut(&mut self) -> &mut Vec>> { - &mut self.entries - } -} - -#[cfg(feature = "std")] -impl Corpus for OnDiskCorpus -where - I: Input, - R: Rand, -{ - /// Add an entry and save it to disk - fn add(&mut self, mut entry: Testcase) -> usize { - match entry.filename() { - None => { - // TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL) - let filename = self.dir_path.join(format!("id_{}", &self.entries.len())); - let filename_str = filename.to_str().expect("Invalid Path"); - entry.set_filename(filename_str.into()); - } - _ => {} - } - self.entries.push(RefCell::new(entry)); - self.entries.len() - 1 - } - - #[inline] - fn current_testcase(&self) -> (&RefCell>, usize) { - (self.get(self.pos), self.pos) - } - - /// Gets the next entry - #[inline] - fn next(&mut self, rand: &mut R) -> Result<(&RefCell>, usize), AflError> { - if self.count() == 0 { - Err(AflError::Empty("No entries in corpus".to_owned())) - } else { - let len = { self.entries().len() }; - let id = rand.below(len as u64) as usize; - self.pos = id; - Ok((self.get(id), id)) - } - } - - // TODO save and remove files, cache, etc..., ATM use just InMemoryCorpus -} - -#[cfg(feature = "std")] -impl OnDiskCorpus -where - I: Input, - R: Rand, -{ - pub fn new(dir_path: PathBuf) -> Self { - Self { - dir_path: dir_path, - entries: vec![], - pos: 0, - phantom: PhantomData, - } - } -} - -/// A Queue-like corpus, wrapping an existing Corpus instance -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(bound = "I: serde::de::DeserializeOwned")] -pub struct QueueCorpus -where - C: Corpus, - I: Input, - R: Rand, -{ - corpus: C, - pos: usize, - cycles: u64, - phantom: PhantomData<(I, R)>, -} - -impl HasTestcaseVec for QueueCorpus -where - C: Corpus, - I: Input, - R: Rand, -{ - #[inline] - fn entries(&self) -> &[RefCell>] { - self.corpus.entries() - } - #[inline] - fn entries_mut(&mut self) -> &mut Vec>> { - self.corpus.entries_mut() - } -} - -impl Corpus for QueueCorpus -where - C: Corpus, - I: Input, - R: Rand, -{ - /// Returns the number of elements - #[inline] - fn count(&self) -> usize { - self.corpus.count() - } - - #[inline] - fn add(&mut self, entry: Testcase) -> usize { - self.corpus.add(entry) - } - - /// Removes an entry from the corpus, returning it if it was present. - #[inline] - fn remove(&mut self, entry: &Testcase) -> Option> { - self.corpus.remove(entry) - } - - /// Gets a random entry - #[inline] - fn random_entry(&self, rand: &mut R) -> Result<(&RefCell>, usize), AflError> { - self.corpus.random_entry(rand) - } - - /// Returns the testacase we currently use - #[inline] - fn current_testcase(&self) -> (&RefCell>, usize) { - (self.get(self.pos - 1), self.pos - 1) - } - - /// Gets the next entry - #[inline] - fn next(&mut self, _rand: &mut R) -> Result<(&RefCell>, usize), AflError> { - self.pos += 1; - if self.corpus.count() == 0 { - return Err(AflError::Empty("Corpus".to_owned())); - } - if self.pos > self.corpus.count() { - // TODO: Always loop or return informational error? - self.pos = 1; - self.cycles += 1; - } - Ok((&self.corpus.entries()[self.pos - 1], self.pos - 1)) - } -} - -impl QueueCorpus -where - C: Corpus, - I: Input, - R: Rand, -{ - pub fn new(corpus: C) -> Self { - Self { - corpus: corpus, - phantom: PhantomData, - cycles: 0, - pos: 0, - } - } - - #[inline] - pub fn cycles(&self) -> u64 { - self.cycles - } - - #[inline] - pub fn pos(&self) -> usize { - self.pos - } -} - -#[cfg(test)] -#[cfg(feature = "std")] -mod tests { - - use std::path::PathBuf; - - use crate::{ - corpus::{Corpus, OnDiskCorpus, QueueCorpus, Testcase}, - inputs::bytes::BytesInput, - utils::StdRand, - }; - - #[test] - fn test_queuecorpus() { - let mut rand = StdRand::new(0); - let mut q = QueueCorpus::new(OnDiskCorpus::::new(PathBuf::from( - "fancy/path", - ))); - let t = Testcase::with_filename(BytesInput::new(vec![0 as u8; 4]), "fancyfile".into()); - q.add(t); - let filename = q - .next(&mut rand) - .unwrap() - .0 - .borrow() - .filename() - .as_ref() - .unwrap() - .to_owned(); - assert_eq!( - filename, - q.next(&mut rand) - .unwrap() - .0 - .borrow() - .filename() - .as_ref() - .unwrap() - .to_owned() - ); - assert_eq!(filename, "fancyfile"); - } -} diff --git a/afl/src/corpus/ondisk.rs b/afl/src/corpus/ondisk.rs new file mode 100644 index 0000000000..e18bc81aa8 --- /dev/null +++ b/afl/src/corpus/ondisk.rs @@ -0,0 +1,99 @@ +use alloc::vec::Vec; +use core::{cell::RefCell, marker::PhantomData}; +use serde::{Deserialize, Serialize}; + +#[cfg(feature = "std")] +use std::path::PathBuf; + +use crate::{ + corpus::Corpus, corpus::HasTestcaseVec, corpus::Testcase, inputs::Input, utils::Rand, AflError, +}; + +/// A corpus able to store testcases to disk, and load them from disk, when they are being used. +#[cfg(feature = "std")] +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(bound = "I: serde::de::DeserializeOwned")] +pub struct OnDiskCorpus +where + I: Input, + R: Rand, +{ + entries: Vec>>, + dir_path: PathBuf, + pos: usize, + phantom: PhantomData, +} + +#[cfg(feature = "std")] +impl HasTestcaseVec for OnDiskCorpus +where + I: Input, + R: Rand, +{ + #[inline] + fn entries(&self) -> &[RefCell>] { + &self.entries + } + #[inline] + fn entries_mut(&mut self) -> &mut Vec>> { + &mut self.entries + } +} + +#[cfg(feature = "std")] +impl Corpus for OnDiskCorpus +where + I: Input, + R: Rand, +{ + /// Add an entry and save it to disk + fn add(&mut self, mut entry: Testcase) -> usize { + match entry.filename() { + None => { + // TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL) + let filename = self.dir_path.join(format!("id_{}", &self.entries.len())); + let filename_str = filename.to_str().expect("Invalid Path"); + entry.set_filename(filename_str.into()); + } + _ => {} + } + self.entries.push(RefCell::new(entry)); + self.entries.len() - 1 + } + + #[inline] + fn current_testcase(&self) -> (&RefCell>, usize) { + (self.get(self.pos), self.pos) + } + + /// Gets the next entry + #[inline] + fn next(&mut self, rand: &mut R) -> Result<(&RefCell>, usize), AflError> { + if self.count() == 0 { + Err(AflError::Empty("No entries in corpus".to_owned())) + } else { + let len = { self.entries().len() }; + let id = rand.below(len as u64) as usize; + self.pos = id; + Ok((self.get(id), id)) + } + } + + // TODO save and remove files, cache, etc..., ATM use just InMemoryCorpus +} + +#[cfg(feature = "std")] +impl OnDiskCorpus +where + I: Input, + R: Rand, +{ + pub fn new(dir_path: PathBuf) -> Self { + Self { + dir_path: dir_path, + entries: vec![], + pos: 0, + phantom: PhantomData, + } + } +} diff --git a/afl/src/corpus/queue.rs b/afl/src/corpus/queue.rs new file mode 100644 index 0000000000..0781a69308 --- /dev/null +++ b/afl/src/corpus/queue.rs @@ -0,0 +1,159 @@ +use alloc::vec::Vec; +use core::{cell::RefCell, marker::PhantomData}; +use serde::{Deserialize, Serialize}; + +use crate::{ + corpus::Corpus, corpus::HasTestcaseVec, corpus::Testcase, inputs::Input, utils::Rand, AflError, +}; + +/// A Queue-like corpus, wrapping an existing Corpus instance +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(bound = "I: serde::de::DeserializeOwned")] +pub struct QueueCorpus +where + C: Corpus, + I: Input, + R: Rand, +{ + corpus: C, + pos: usize, + cycles: u64, + phantom: PhantomData<(I, R)>, +} + +impl HasTestcaseVec for QueueCorpus +where + C: Corpus, + I: Input, + R: Rand, +{ + #[inline] + fn entries(&self) -> &[RefCell>] { + self.corpus.entries() + } + #[inline] + fn entries_mut(&mut self) -> &mut Vec>> { + self.corpus.entries_mut() + } +} + +impl Corpus for QueueCorpus +where + C: Corpus, + I: Input, + R: Rand, +{ + /// Returns the number of elements + #[inline] + fn count(&self) -> usize { + self.corpus.count() + } + + #[inline] + fn add(&mut self, entry: Testcase) -> usize { + self.corpus.add(entry) + } + + /// Removes an entry from the corpus, returning it if it was present. + #[inline] + fn remove(&mut self, entry: &Testcase) -> Option> { + self.corpus.remove(entry) + } + + /// Gets a random entry + #[inline] + fn random_entry(&self, rand: &mut R) -> Result<(&RefCell>, usize), AflError> { + self.corpus.random_entry(rand) + } + + /// Returns the testacase we currently use + #[inline] + fn current_testcase(&self) -> (&RefCell>, usize) { + (self.get(self.pos - 1), self.pos - 1) + } + + /// Gets the next entry + #[inline] + fn next(&mut self, _rand: &mut R) -> Result<(&RefCell>, usize), AflError> { + self.pos += 1; + if self.corpus.count() == 0 { + return Err(AflError::Empty("Corpus".to_owned())); + } + if self.pos > self.corpus.count() { + // TODO: Always loop or return informational error? + self.pos = 1; + self.cycles += 1; + } + Ok((&self.corpus.entries()[self.pos - 1], self.pos - 1)) + } +} + +impl QueueCorpus +where + C: Corpus, + I: Input, + R: Rand, +{ + pub fn new(corpus: C) -> Self { + Self { + corpus: corpus, + phantom: PhantomData, + cycles: 0, + pos: 0, + } + } + + #[inline] + pub fn cycles(&self) -> u64 { + self.cycles + } + + #[inline] + pub fn pos(&self) -> usize { + self.pos + } +} + +#[cfg(test)] +#[cfg(feature = "std")] +mod tests { + + use std::path::PathBuf; + + use crate::{ + corpus::{Corpus, OnDiskCorpus, QueueCorpus, Testcase}, + inputs::bytes::BytesInput, + utils::StdRand, + }; + + #[test] + fn test_queuecorpus() { + let mut rand = StdRand::new(0); + let mut q = QueueCorpus::new(OnDiskCorpus::::new(PathBuf::from( + "fancy/path", + ))); + let t = Testcase::with_filename(BytesInput::new(vec![0 as u8; 4]), "fancyfile".into()); + q.add(t); + let filename = q + .next(&mut rand) + .unwrap() + .0 + .borrow() + .filename() + .as_ref() + .unwrap() + .to_owned(); + assert_eq!( + filename, + q.next(&mut rand) + .unwrap() + .0 + .borrow() + .filename() + .as_ref() + .unwrap() + .to_owned() + ); + assert_eq!(filename, "fancyfile"); + } +} diff --git a/afl/src/corpus/testcase.rs b/afl/src/corpus/testcase.rs index d389b3c880..0fa53464d9 100644 --- a/afl/src/corpus/testcase.rs +++ b/afl/src/corpus/testcase.rs @@ -21,7 +21,6 @@ where filename: Option, /// Accumulated fitness from all the feedbacks fitness: u32, - // TODO find a way to use TypeId /// Map of metadatas associated with this testcase metadatas: SerdeAnyMap, }