refactor corpus dir

This commit is contained in:
Andrea Fioraldi 2021-01-12 13:50:25 +01:00
parent b821d90186
commit 004c3e3233
5 changed files with 341 additions and 332 deletions

View File

@ -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<I, R>
where
I: Input,
R: Rand,
{
entries: Vec<RefCell<Testcase<I>>>,
pos: usize,
phantom: PhantomData<R>,
}
impl<I, R> HasTestcaseVec<I> for InMemoryCorpus<I, R>
where
I: Input,
R: Rand,
{
fn entries(&self) -> &[RefCell<Testcase<I>>] {
&self.entries
}
fn entries_mut(&mut self) -> &mut Vec<RefCell<Testcase<I>>> {
&mut self.entries
}
}
impl<I, R> Corpus<I, R> for InMemoryCorpus<I, R>
where
I: Input,
R: Rand,
{
/// Gets the next entry
#[inline]
fn next(&mut self, rand: &mut R) -> Result<(&RefCell<Testcase<I>>, 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<Testcase<I>>, usize) {
(self.get(self.pos), self.pos)
}
}
impl<I, R> InMemoryCorpus<I, R>
where
I: Input,
R: Rand,
{
pub fn new() -> Self {
Self {
entries: vec![],
pos: 0,
phantom: PhantomData,
}
}
}

View File

@ -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<Testcase<I>>, usize), AflError>;
@ -112,309 +98,3 @@ where
/// Returns the testacase we currently use
fn current_testcase(&self) -> (&RefCell<Testcase<I>>, usize);
}
/// A corpus handling all important fuzzing in memory.
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(bound = "I: serde::de::DeserializeOwned")]
pub struct InMemoryCorpus<I, R>
where
I: Input,
R: Rand,
{
entries: Vec<RefCell<Testcase<I>>>,
pos: usize,
phantom: PhantomData<R>,
}
impl<I, R> HasTestcaseVec<I> for InMemoryCorpus<I, R>
where
I: Input,
R: Rand,
{
fn entries(&self) -> &[RefCell<Testcase<I>>] {
&self.entries
}
fn entries_mut(&mut self) -> &mut Vec<RefCell<Testcase<I>>> {
&mut self.entries
}
}
impl<I, R> Corpus<I, R> for InMemoryCorpus<I, R>
where
I: Input,
R: Rand,
{
/// Gets the next entry
#[inline]
fn next(&mut self, rand: &mut R) -> Result<(&RefCell<Testcase<I>>, 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<Testcase<I>>, usize) {
(self.get(self.pos), self.pos)
}
}
impl<I, R> InMemoryCorpus<I, R>
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<I, R>
where
I: Input,
R: Rand,
{
entries: Vec<RefCell<Testcase<I>>>,
dir_path: PathBuf,
pos: usize,
phantom: PhantomData<R>,
}
#[cfg(feature = "std")]
impl<I, R> HasTestcaseVec<I> for OnDiskCorpus<I, R>
where
I: Input,
R: Rand,
{
#[inline]
fn entries(&self) -> &[RefCell<Testcase<I>>] {
&self.entries
}
#[inline]
fn entries_mut(&mut self) -> &mut Vec<RefCell<Testcase<I>>> {
&mut self.entries
}
}
#[cfg(feature = "std")]
impl<I, R> Corpus<I, R> for OnDiskCorpus<I, R>
where
I: Input,
R: Rand,
{
/// Add an entry and save it to disk
fn add(&mut self, mut entry: Testcase<I>) -> 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<Testcase<I>>, usize) {
(self.get(self.pos), self.pos)
}
/// Gets the next entry
#[inline]
fn next(&mut self, rand: &mut R) -> Result<(&RefCell<Testcase<I>>, 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<I, R> OnDiskCorpus<I, R>
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<C, I, R>
where
C: Corpus<I, R>,
I: Input,
R: Rand,
{
corpus: C,
pos: usize,
cycles: u64,
phantom: PhantomData<(I, R)>,
}
impl<C, I, R> HasTestcaseVec<I> for QueueCorpus<C, I, R>
where
C: Corpus<I, R>,
I: Input,
R: Rand,
{
#[inline]
fn entries(&self) -> &[RefCell<Testcase<I>>] {
self.corpus.entries()
}
#[inline]
fn entries_mut(&mut self) -> &mut Vec<RefCell<Testcase<I>>> {
self.corpus.entries_mut()
}
}
impl<C, I, R> Corpus<I, R> for QueueCorpus<C, I, R>
where
C: Corpus<I, R>,
I: Input,
R: Rand,
{
/// Returns the number of elements
#[inline]
fn count(&self) -> usize {
self.corpus.count()
}
#[inline]
fn add(&mut self, entry: Testcase<I>) -> usize {
self.corpus.add(entry)
}
/// Removes an entry from the corpus, returning it if it was present.
#[inline]
fn remove(&mut self, entry: &Testcase<I>) -> Option<Testcase<I>> {
self.corpus.remove(entry)
}
/// Gets a random entry
#[inline]
fn random_entry(&self, rand: &mut R) -> Result<(&RefCell<Testcase<I>>, usize), AflError> {
self.corpus.random_entry(rand)
}
/// Returns the testacase we currently use
#[inline]
fn current_testcase(&self) -> (&RefCell<Testcase<I>>, usize) {
(self.get(self.pos - 1), self.pos - 1)
}
/// Gets the next entry
#[inline]
fn next(&mut self, _rand: &mut R) -> Result<(&RefCell<Testcase<I>>, 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<C, I, R> QueueCorpus<C, I, R>
where
C: Corpus<I, R>,
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::<BytesInput, StdRand>::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");
}
}

99
afl/src/corpus/ondisk.rs Normal file
View File

@ -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<I, R>
where
I: Input,
R: Rand,
{
entries: Vec<RefCell<Testcase<I>>>,
dir_path: PathBuf,
pos: usize,
phantom: PhantomData<R>,
}
#[cfg(feature = "std")]
impl<I, R> HasTestcaseVec<I> for OnDiskCorpus<I, R>
where
I: Input,
R: Rand,
{
#[inline]
fn entries(&self) -> &[RefCell<Testcase<I>>] {
&self.entries
}
#[inline]
fn entries_mut(&mut self) -> &mut Vec<RefCell<Testcase<I>>> {
&mut self.entries
}
}
#[cfg(feature = "std")]
impl<I, R> Corpus<I, R> for OnDiskCorpus<I, R>
where
I: Input,
R: Rand,
{
/// Add an entry and save it to disk
fn add(&mut self, mut entry: Testcase<I>) -> 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<Testcase<I>>, usize) {
(self.get(self.pos), self.pos)
}
/// Gets the next entry
#[inline]
fn next(&mut self, rand: &mut R) -> Result<(&RefCell<Testcase<I>>, 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<I, R> OnDiskCorpus<I, R>
where
I: Input,
R: Rand,
{
pub fn new(dir_path: PathBuf) -> Self {
Self {
dir_path: dir_path,
entries: vec![],
pos: 0,
phantom: PhantomData,
}
}
}

159
afl/src/corpus/queue.rs Normal file
View File

@ -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<C, I, R>
where
C: Corpus<I, R>,
I: Input,
R: Rand,
{
corpus: C,
pos: usize,
cycles: u64,
phantom: PhantomData<(I, R)>,
}
impl<C, I, R> HasTestcaseVec<I> for QueueCorpus<C, I, R>
where
C: Corpus<I, R>,
I: Input,
R: Rand,
{
#[inline]
fn entries(&self) -> &[RefCell<Testcase<I>>] {
self.corpus.entries()
}
#[inline]
fn entries_mut(&mut self) -> &mut Vec<RefCell<Testcase<I>>> {
self.corpus.entries_mut()
}
}
impl<C, I, R> Corpus<I, R> for QueueCorpus<C, I, R>
where
C: Corpus<I, R>,
I: Input,
R: Rand,
{
/// Returns the number of elements
#[inline]
fn count(&self) -> usize {
self.corpus.count()
}
#[inline]
fn add(&mut self, entry: Testcase<I>) -> usize {
self.corpus.add(entry)
}
/// Removes an entry from the corpus, returning it if it was present.
#[inline]
fn remove(&mut self, entry: &Testcase<I>) -> Option<Testcase<I>> {
self.corpus.remove(entry)
}
/// Gets a random entry
#[inline]
fn random_entry(&self, rand: &mut R) -> Result<(&RefCell<Testcase<I>>, usize), AflError> {
self.corpus.random_entry(rand)
}
/// Returns the testacase we currently use
#[inline]
fn current_testcase(&self) -> (&RefCell<Testcase<I>>, usize) {
(self.get(self.pos - 1), self.pos - 1)
}
/// Gets the next entry
#[inline]
fn next(&mut self, _rand: &mut R) -> Result<(&RefCell<Testcase<I>>, 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<C, I, R> QueueCorpus<C, I, R>
where
C: Corpus<I, R>,
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::<BytesInput, StdRand>::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");
}
}

View File

@ -21,7 +21,6 @@ where
filename: Option<String>,
/// Accumulated fitness from all the feedbacks
fitness: u32,
// TODO find a way to use TypeId
/// Map of metadatas associated with this testcase
metadatas: SerdeAnyMap,
}