Skip to content

Commit

Permalink
Improve Flexibility of DumpToDiskStage (#2753)
Browse files Browse the repository at this point in the history 
* fixing empty multipart name

* fixing clippy

* improve flexibility of DumpToDiskStage

* adding note to MIGRATION.md
  • Loading branch information
@riesentoaster
riesentoaster committed Dec 11, 2024
1 parent faac9a7 commit 2a14db9
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 37 deletions.
3 changes: 2 additions & 1 deletion MIGRATION.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,4 +6,5 @@
- `MmapShMemProvider::new_shmem_persistent` has been removed in favour of `MmapShMem::persist`. You probably want to do something like this: `let shmem = MmapShMemProvider::new()?.new_shmem(size)?.persist()?;`

# 0.14.1 -> 0.14.2
- `MmapShMem::new` and `MmapShMemProvider::new_shmem_with_id` now take `AsRef<Path>` instead of a byte array for the filename/id.
- `MmapShMem::new` and `MmapShMemProvider::new_shmem_with_id` now take `AsRef<Path>` instead of a byte array for the filename/id.
- The closure passed to a `DumpToDiskStage` now provides the `Testcase` instead of just the `Input`.
128 changes: 92 additions & 36 deletions libafl/src/stages/dump.rs
View file
Original file line number Diff line number Diff line change
@@ -1,14 +1,20 @@
//! The [`DumpToDiskStage`] is a stage that dumps the corpus and the solutions to disk to e.g. allow AFL to sync
use alloc::{string::String, vec::Vec};
use alloc::vec::Vec;
use core::{clone::Clone, marker::PhantomData};
use std::{fs, fs::File, io::Write, path::PathBuf};
use std::{
fs::{self, File},
io::Write,
path::{Path, PathBuf},
string::{String, ToString},
};

use libafl_bolts::impl_serdeany;
use serde::{Deserialize, Serialize};

use crate::{
corpus::{Corpus, CorpusId},
corpus::{Corpus, CorpusId, Testcase},
inputs::Input,
stages::Stage,
state::{HasCorpus, HasRand, HasSolutions, UsesState},
Error, HasMetadata,
Expand All @@ -29,23 +35,26 @@ impl_serdeany!(DumpToDiskMetadata);

/// The [`DumpToDiskStage`] is a stage that dumps the corpus and the solutions to disk
#[derive(Debug)]
pub struct DumpToDiskStage<CB, EM, Z> {
pub struct DumpToDiskStage<CB1, CB2, EM, Z> {
solutions_dir: PathBuf,
corpus_dir: PathBuf,
to_bytes: CB,
to_bytes: CB1,
generate_filename: CB2,
phantom: PhantomData<(EM, Z)>,
}

impl<CB, EM, Z> UsesState for DumpToDiskStage<CB, EM, Z>
impl<CB1, CB2, EM, Z> UsesState for DumpToDiskStage<CB1, CB2, EM, Z>
where
EM: UsesState,
{
type State = EM::State;
}

impl<CB, E, EM, Z> Stage<E, EM, Z> for DumpToDiskStage<CB, EM, Z>
impl<CB1, CB2, P, E, EM, Z> Stage<E, EM, Z> for DumpToDiskStage<CB1, CB2, EM, Z>
where
CB: FnMut(&Self::Input, &Self::State) -> Vec<u8>,
CB1: FnMut(&Testcase<Self::Input>, &Self::State) -> Vec<u8>,
CB2: FnMut(&Testcase<Self::Input>, &CorpusId) -> P,
P: AsRef<Path>,
EM: UsesState,
E: UsesState<State = Self::State>,
Z: UsesState<State = Self::State>,
Expand Down Expand Up @@ -77,16 +86,63 @@ where
}
}

impl<CB, EM, Z> DumpToDiskStage<CB, EM, Z>
/// Implementation for `DumpToDiskStage` with a default `generate_filename` function.
impl<CB1, EM, Z> DumpToDiskStage<CB1, fn(&Testcase<EM::Input>, &CorpusId) -> String, EM, Z>
where
EM: UsesState,
Z: UsesState,
<EM as UsesState>::State: HasCorpus + HasSolutions + HasRand + HasMetadata,
<<EM as UsesState>::State as HasCorpus>::Corpus: Corpus<Input = EM::Input>,
<<EM as UsesState>::State as HasSolutions>::Solutions: Corpus<Input = EM::Input>,
{
/// Create a new [`DumpToDiskStage`] with a default `generate_filename` function.
pub fn new<A, B>(to_bytes: CB1, corpus_dir: A, solutions_dir: B) -> Result<Self, Error>
where
A: Into<PathBuf>,
B: Into<PathBuf>,
{
Self::new_with_custom_filenames(
to_bytes,
Self::generate_filename, // This is now of type `fn(&Testcase<EM::Input>, &CorpusId) -> String`
corpus_dir,
solutions_dir,
)
}

/// Default `generate_filename` function.
#[allow(clippy::trivially_copy_pass_by_ref)]
fn generate_filename(testcase: &Testcase<EM::Input>, id: &CorpusId) -> String {
[
Some(id.0.to_string()),
testcase.filename().clone(),
testcase
.input()
.as_ref()
.map(|t| t.generate_name(Some(*id))),
]
.iter()
.flatten()
.map(String::as_str)
.collect::<Vec<_>>()
.join("-")
}
}

impl<CB1, CB2, EM, Z> DumpToDiskStage<CB1, CB2, EM, Z>
where
EM: UsesState,
Z: UsesState,
<EM as UsesState>::State: HasCorpus + HasSolutions + HasRand + HasMetadata,
<<EM as UsesState>::State as HasCorpus>::Corpus: Corpus<Input = EM::Input>,
<<EM as UsesState>::State as HasSolutions>::Solutions: Corpus<Input = EM::Input>,
{
/// Create a new [`DumpToDiskStage`]
pub fn new<A, B>(to_bytes: CB, corpus_dir: A, solutions_dir: B) -> Result<Self, Error>
/// Create a new [`DumpToDiskStage`] with a custom `generate_filename` function.
pub fn new_with_custom_filenames<A, B>(
to_bytes: CB1,
generate_filename: CB2,
corpus_dir: A,
solutions_dir: B,
) -> Result<Self, Error>
where
A: Into<PathBuf>,
B: Into<PathBuf>,
Expand All @@ -102,7 +158,7 @@ where
}
let solutions_dir = solutions_dir.into();
if let Err(e) = fs::create_dir(&solutions_dir) {
if !corpus_dir.is_dir() {
if !solutions_dir.is_dir() {
return Err(Error::os_error(
e,
format!("Error creating directory {solutions_dir:?}"),
Expand All @@ -111,19 +167,27 @@ where
}
Ok(Self {
to_bytes,
generate_filename,
solutions_dir,
corpus_dir,
phantom: PhantomData,
})
}

#[inline]
fn dump_state_to_disk(&mut self, state: &mut <Self as UsesState>::State) -> Result<(), Error>
fn dump_state_to_disk<P: AsRef<Path>>(
&mut self,
state: &mut <Self as UsesState>::State,
) -> Result<(), Error>
where
CB: FnMut(
&<<<EM as UsesState>::State as HasCorpus>::Corpus as Corpus>::Input,
CB1: FnMut(
&Testcase<<<<EM as UsesState>::State as HasCorpus>::Corpus as Corpus>::Input>,
&<EM as UsesState>::State,
) -> Vec<u8>,
CB2: FnMut(
&Testcase<<<<EM as UsesState>::State as HasCorpus>::Corpus as Corpus>::Input>,
&CorpusId,
) -> P,
{
let (mut corpus_id, mut solutions_id) =
if let Some(meta) = state.metadata_map().get::<DumpToDiskMetadata>() {
Expand All @@ -138,37 +202,29 @@ where
while let Some(i) = corpus_id {
let mut testcase = state.corpus().get(i)?.borrow_mut();
state.corpus().load_input_into(&mut testcase)?;
let bytes = (self.to_bytes)(testcase.input().as_ref().unwrap(), state);

let fname = self.corpus_dir.join(format!(
"id_{i}_{}",
testcase
.filename()
.as_ref()
.map_or_else(|| "unnamed", String::as_str)
));
let bytes = (self.to_bytes)(&testcase, state);

let fname = self
.corpus_dir
.join((self.generate_filename)(&testcase, &i));
let mut f = File::create(fname)?;
drop(f.write_all(&bytes));

corpus_id = state.corpus().next(i);
}

while let Some(current_id) = solutions_id {
let mut testcase = state.solutions().get(current_id)?.borrow_mut();
while let Some(i) = solutions_id {
let mut testcase = state.solutions().get(i)?.borrow_mut();
state.solutions().load_input_into(&mut testcase)?;
let bytes = (self.to_bytes)(testcase.input().as_ref().unwrap(), state);

let fname = self.solutions_dir.join(format!(
"id_{current_id}_{}",
testcase
.filename()
.as_ref()
.map_or_else(|| "unnamed", String::as_str)
));
let bytes = (self.to_bytes)(&testcase, state);

let fname = self
.solutions_dir
.join((self.generate_filename)(&testcase, &i));
let mut f = File::create(fname)?;
drop(f.write_all(&bytes));

solutions_id = state.solutions().next(current_id);
solutions_id = state.solutions().next(i);
}

state.add_metadata(DumpToDiskMetadata {
Expand Down

0 comments on commit 2a14db9

Please sign in to comment.