extend python forkserver api (#500)

* initial attempt at api extension; untested

* updated/tested on forkserver_simple
This commit is contained in:
epi 2022-01-28 02:43:21 -06:00 committed by GitHub
parent 4e3e31df4e
commit 78bbe034a1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -17,7 +17,7 @@ use libafl::{
CachedOnDiskCorpus, Corpus, IndexesLenTimeMinimizerCorpusScheduler, OnDiskCorpus, CachedOnDiskCorpus, Corpus, IndexesLenTimeMinimizerCorpusScheduler, OnDiskCorpus,
QueueCorpusScheduler, QueueCorpusScheduler,
}, },
events::EventConfig, events::{EventConfig, EventRestarter, LlmpRestartingEventManager},
executors::{ForkserverExecutor, TimeoutForkserverExecutor}, executors::{ForkserverExecutor, TimeoutForkserverExecutor},
feedback_or, feedback_or_fast, feedback_or, feedback_or_fast,
feedbacks::{CrashFeedback, MapFeedbackState, MaxMapFeedback, TimeFeedback, TimeoutFeedback}, feedbacks::{CrashFeedback, MapFeedbackState, MaxMapFeedback, TimeFeedback, TimeoutFeedback},
@ -44,18 +44,18 @@ pub struct ForkserverBytesCoverageSugar<'a, const MAP_SIZE: usize> {
#[builder(default = None, setter(strip_option))] #[builder(default = None, setter(strip_option))]
configuration: Option<String>, configuration: Option<String>,
/// Timeout of the executor /// Timeout of the executor
#[builder(default = None, setter(strip_option))] #[builder(default = None)]
timeout: Option<u64>, timeout: Option<u64>,
/// Input directories /// Input directories
input_dirs: &'a [PathBuf], input_dirs: &'a [PathBuf],
/// Output directory /// Output directory
output_dir: PathBuf, output_dir: PathBuf,
/// Dictionary /// Dictionary
#[builder(default = None, setter(strip_option))] #[builder(default = None)]
tokens_file: Option<PathBuf>, tokens_file: Option<PathBuf>,
// Flag if use CmpLog // Flag if use CmpLog
//#[builder(default = false)] #[builder(default = None)]
//use_cmplog: bool, use_cmplog: Option<bool>,
#[builder(default = 1337_u16)] #[builder(default = 1337_u16)]
broker_port: u16, broker_port: u16,
/// The list of cores to run on /// The list of cores to run on
@ -74,6 +74,9 @@ pub struct ForkserverBytesCoverageSugar<'a, const MAP_SIZE: usize> {
#[builder(default = false)] #[builder(default = false)]
/// Print target program output /// Print target program output
debug_output: bool, debug_output: bool,
/// Fuzz `iterations` number of times, instead of indefinitely; implies use of `fuzz_loop_for`
#[builder(default = None)]
iterations: Option<u64>,
} }
#[allow(clippy::similar_names)] #[allow(clippy::similar_names)]
@ -86,6 +89,10 @@ impl<'a, const MAP_SIZE: usize> ForkserverBytesCoverageSugar<'a, MAP_SIZE> {
None => EventConfig::AlwaysUnique, None => EventConfig::AlwaysUnique,
}; };
if self.use_cmplog.unwrap_or(false) {
println!("[WARNING] use of cmplog not currently supported, use_cmplog ignored.");
}
let timeout = Duration::from_secs(self.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS)); let timeout = Duration::from_secs(self.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS));
let mut out_dir = self.output_dir.clone(); let mut out_dir = self.output_dir.clone();
@ -106,7 +113,9 @@ impl<'a, const MAP_SIZE: usize> ForkserverBytesCoverageSugar<'a, MAP_SIZE> {
let monitor = MultiMonitor::new(|s| println!("{}", s)); let monitor = MultiMonitor::new(|s| println!("{}", s));
let mut run_client = |state: Option<StdState<_, _, _, _, _>>, mut mgr, _core_id| { let mut run_client = |state: Option<StdState<_, _, _, _, _>>,
mut mgr: LlmpRestartingEventManager<_, _, _, _>,
_core_id| {
// Coverage map shared between target and fuzzer // Coverage map shared between target and fuzzer
let mut shmem = shmem_provider_client.new_shmem(MAP_SIZE).unwrap(); let mut shmem = shmem_provider_client.new_shmem(MAP_SIZE).unwrap();
shmem.write_to_env("__AFL_SHM_ID").unwrap(); shmem.write_to_env("__AFL_SHM_ID").unwrap();
@ -230,7 +239,20 @@ impl<'a, const MAP_SIZE: usize> ForkserverBytesCoverageSugar<'a, MAP_SIZE> {
// The order of the stages matter! // The order of the stages matter!
let mut stages = tuple_list!(mutational); let mut stages = tuple_list!(mutational);
fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
if let Some(iters) = self.iterations {
fuzzer.fuzz_loop_for(
&mut stages,
&mut executor,
&mut state,
&mut mgr,
iters,
)?;
mgr.on_restart(&mut state)?;
std::process::exit(0);
} else {
fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
}
} else { } else {
// Setup a basic mutator // Setup a basic mutator
let mutator = StdScheduledMutator::new(havoc_mutations()); let mutator = StdScheduledMutator::new(havoc_mutations());
@ -238,7 +260,20 @@ impl<'a, const MAP_SIZE: usize> ForkserverBytesCoverageSugar<'a, MAP_SIZE> {
// The order of the stages matter! // The order of the stages matter!
let mut stages = tuple_list!(mutational); let mut stages = tuple_list!(mutational);
fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
if let Some(iters) = self.iterations {
fuzzer.fuzz_loop_for(
&mut stages,
&mut executor,
&mut state,
&mut mgr,
iters,
)?;
mgr.on_restart(&mut state)?;
std::process::exit(0);
} else {
fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
}
} }
Ok(()) Ok(())
@ -277,6 +312,10 @@ pub mod pybind {
output_dir: PathBuf, output_dir: PathBuf,
broker_port: u16, broker_port: u16,
cores: Cores, cores: Cores,
use_cmplog: Option<bool>,
iterations: Option<u64>,
tokens_file: Option<PathBuf>,
timeout: Option<u64>,
} }
#[pymethods] #[pymethods]
@ -288,12 +327,20 @@ pub mod pybind {
output_dir: PathBuf, output_dir: PathBuf,
broker_port: u16, broker_port: u16,
cores: Vec<usize>, cores: Vec<usize>,
use_cmplog: Option<bool>,
iterations: Option<u64>,
tokens_file: Option<PathBuf>,
timeout: Option<u64>,
) -> Self { ) -> Self {
Self { Self {
input_dirs, input_dirs,
output_dir, output_dir,
broker_port, broker_port,
cores: cores.into(), cores: cores.into(),
use_cmplog,
iterations,
tokens_file,
timeout,
} }
} }
@ -307,6 +354,10 @@ pub mod pybind {
.cores(&self.cores) .cores(&self.cores)
.program(program) .program(program)
.arguments(&arguments) .arguments(&arguments)
.use_cmplog(self.use_cmplog)
.timeout(self.timeout)
.tokens_file(self.tokens_file.clone())
.iterations(self.iterations)
.build() .build()
.run(); .run();
} }