From 2a14db936c35ed9b054e84569783975dc5031d68 Mon Sep 17 00:00:00 2001 From: Valentin Huber Date: Sun, 8 Dec 2024 21:46:38 +0100 Subject: [PATCH] Improve Flexibility of DumpToDiskStage (#2753) * fixing empty multipart name * fixing clippy * improve flexibility of DumpToDiskStage * adding note to MIGRATION.md --- MIGRATION.md | 3 +- libafl/src/stages/dump.rs | 128 +++++++++++++++++++++++++++----------- 2 files changed, 94 insertions(+), 37 deletions(-) diff --git a/MIGRATION.md b/MIGRATION.md index 0f591564ae..e140eab2e7 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -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` instead of a byte array for the filename/id. \ No newline at end of file +- `MmapShMem::new` and `MmapShMemProvider::new_shmem_with_id` now take `AsRef` instead of a byte array for the filename/id. +- The closure passed to a `DumpToDiskStage` now provides the `Testcase` instead of just the `Input`. \ No newline at end of file diff --git a/libafl/src/stages/dump.rs b/libafl/src/stages/dump.rs index ecfc2c2a10..ee80fa6062 100644 --- a/libafl/src/stages/dump.rs +++ b/libafl/src/stages/dump.rs @@ -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, @@ -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 { +pub struct DumpToDiskStage { solutions_dir: PathBuf, corpus_dir: PathBuf, - to_bytes: CB, + to_bytes: CB1, + generate_filename: CB2, phantom: PhantomData<(EM, Z)>, } -impl UsesState for DumpToDiskStage +impl UsesState for DumpToDiskStage where EM: UsesState, { type State = EM::State; } -impl Stage for DumpToDiskStage +impl Stage for DumpToDiskStage where - CB: FnMut(&Self::Input, &Self::State) -> Vec, + CB1: FnMut(&Testcase, &Self::State) -> Vec, + CB2: FnMut(&Testcase, &CorpusId) -> P, + P: AsRef, EM: UsesState, E: UsesState, Z: UsesState, @@ -77,7 +86,49 @@ where } } -impl DumpToDiskStage +/// Implementation for `DumpToDiskStage` with a default `generate_filename` function. +impl DumpToDiskStage, &CorpusId) -> String, EM, Z> +where + EM: UsesState, + Z: UsesState, + ::State: HasCorpus + HasSolutions + HasRand + HasMetadata, + <::State as HasCorpus>::Corpus: Corpus, + <::State as HasSolutions>::Solutions: Corpus, +{ + /// Create a new [`DumpToDiskStage`] with a default `generate_filename` function. + pub fn new(to_bytes: CB1, corpus_dir: A, solutions_dir: B) -> Result + where + A: Into, + B: Into, + { + Self::new_with_custom_filenames( + to_bytes, + Self::generate_filename, // This is now of type `fn(&Testcase, &CorpusId) -> String` + corpus_dir, + solutions_dir, + ) + } + + /// Default `generate_filename` function. + #[allow(clippy::trivially_copy_pass_by_ref)] + fn generate_filename(testcase: &Testcase, 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::>() + .join("-") + } +} + +impl DumpToDiskStage where EM: UsesState, Z: UsesState, @@ -85,8 +136,13 @@ where <::State as HasCorpus>::Corpus: Corpus, <::State as HasSolutions>::Solutions: Corpus, { - /// Create a new [`DumpToDiskStage`] - pub fn new(to_bytes: CB, corpus_dir: A, solutions_dir: B) -> Result + /// Create a new [`DumpToDiskStage`] with a custom `generate_filename` function. + pub fn new_with_custom_filenames( + to_bytes: CB1, + generate_filename: CB2, + corpus_dir: A, + solutions_dir: B, + ) -> Result where A: Into, B: Into, @@ -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:?}"), @@ -111,6 +167,7 @@ where } Ok(Self { to_bytes, + generate_filename, solutions_dir, corpus_dir, phantom: PhantomData, @@ -118,12 +175,19 @@ where } #[inline] - fn dump_state_to_disk(&mut self, state: &mut ::State) -> Result<(), Error> + fn dump_state_to_disk>( + &mut self, + state: &mut ::State, + ) -> Result<(), Error> where - CB: FnMut( - &<<::State as HasCorpus>::Corpus as Corpus>::Input, + CB1: FnMut( + &Testcase<<<::State as HasCorpus>::Corpus as Corpus>::Input>, &::State, ) -> Vec, + CB2: FnMut( + &Testcase<<<::State as HasCorpus>::Corpus as Corpus>::Input>, + &CorpusId, + ) -> P, { let (mut corpus_id, mut solutions_id) = if let Some(meta) = state.metadata_map().get::() { @@ -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 {