From e914cc9c1481210fdd83c001ac07ff26abe3ae02 Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Wed, 10 Nov 2021 16:00:15 +0100 Subject: [PATCH] add ForkserverBytesCoverageSugar --- libafl_sugar/src/forkserver.rs | 308 +++++++++++++++++++++++++++++++++ libafl_sugar/src/lib.rs | 5 + 2 files changed, 313 insertions(+) create mode 100644 libafl_sugar/src/forkserver.rs diff --git a/libafl_sugar/src/forkserver.rs b/libafl_sugar/src/forkserver.rs new file mode 100644 index 0000000000..10601b423e --- /dev/null +++ b/libafl_sugar/src/forkserver.rs @@ -0,0 +1,308 @@ +use typed_builder::TypedBuilder; + +use std::{fs, net::SocketAddr, path::PathBuf, time::Duration}; + +use libafl::{ + bolts::{ + current_nanos, + launcher::Launcher, + rands::StdRand, + shmem::{ShMem, ShMemProvider, StdShMemProvider}, + tuples::{tuple_list, Merge}, + }, + corpus::{ + CachedOnDiskCorpus, Corpus, IndexesLenTimeMinimizerCorpusScheduler, OnDiskCorpus, + QueueCorpusScheduler, + }, + events::EventConfig, + executors::{ForkserverExecutor, TimeoutForkserverExecutor}, + feedback_or, feedback_or_fast, + feedbacks::{CrashFeedback, MapFeedbackState, MaxMapFeedback, TimeFeedback, TimeoutFeedback}, + fuzzer::{Fuzzer, StdFuzzer}, + generators::RandBytesGenerator, + mutators::scheduled::{havoc_mutations, tokens_mutations, StdScheduledMutator}, + mutators::token_mutations::Tokens, + observers::{ConstMapObserver, HitcountsMapObserver, TimeObserver}, + stages::StdMutationalStage, + state::{HasCorpus, HasMetadata, StdState}, + stats::MultiStats, + Error, +}; + +use crate::{CORPUS_CACHE_SIZE, DEFAULT_TIMEOUT_SECS}; + +pub const MAP_SIZE: usize = 65536; + +#[derive(TypedBuilder)] +pub struct ForkserverBytesCoverageSugar<'a> { + /// Laucher configuration (default is random) + #[builder(default = None, setter(strip_option))] + configuration: Option, + /// Timeout of the executor + #[builder(default = None, setter(strip_option))] + timeout: Option, + /// Input directories + input_dirs: &'a [PathBuf], + /// Output directory + output_dir: PathBuf, + /// Dictionary + #[builder(default = None, setter(strip_option))] + tokens_file: Option, + // Flag if use CmpLog + //#[builder(default = false)] + //use_cmplog: bool, + #[builder(default = 1337_u16)] + broker_port: u16, + /// The list of cores to run on + cores: &'a [usize], + /// The `ip:port` address of another broker to connect our new broker to for multi-machine + /// clusters. + #[builder(default = None, setter(strip_option))] + remote_broker_addr: Option, + /// Path to program to execute + program: String, + /// Arguments of the program to execute + arguments: &'a [String], + #[builder(default = false)] + /// Use shared mem testcase delivery + shmem_testcase: bool, + // Coverage map size + //#[builder(default = MAP_SIZE)] + //map_size: usize, +} + +#[allow(clippy::similar_names)] +impl<'a> ForkserverBytesCoverageSugar<'a> { + #[allow(clippy::too_many_lines, clippy::similar_names)] + pub fn run(&mut self) { + let conf = match self.configuration.as_ref() { + Some(name) => EventConfig::from_name(name), + None => EventConfig::AlwaysUnique, + }; + + let timeout = Duration::from_secs(self.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS)); + + let mut out_dir = self.output_dir.clone(); + if fs::create_dir(&out_dir).is_err() { + println!("Out dir at {:?} already exists.", &out_dir); + assert!( + out_dir.is_dir(), + "Out dir at {:?} is not a valid directory!", + &out_dir + ); + } + let mut crashes = out_dir.clone(); + crashes.push("crashes"); + out_dir.push("queue"); + + let shmem_provider = StdShMemProvider::new().expect("Failed to init shared memory"); + + let stats = MultiStats::new(|s| println!("{}", s)); + + let mut run_client = |state: Option>, mut mgr, _core_id| { + // Coverage map shared between target and fuzzer + let mut shmem = StdShMemProvider::new() + .expect("Failed to init shared memory") + .new_map(MAP_SIZE) + .unwrap(); + shmem.write_to_env("__AFL_SHM_ID").unwrap(); + let shmem_map = shmem.map_mut(); + + // Create an observation channel using the coverage map + let edges_observer = unsafe { + HitcountsMapObserver::new(ConstMapObserver::<_, MAP_SIZE>::new_from_ptr( + "shared_mem", + shmem_map.as_mut_ptr(), + )) + }; + + // Create an observation channel to keep track of the execution time + let time_observer = TimeObserver::new("time"); + + // The state of the edges feedback. + let feedback_state = MapFeedbackState::with_observer(&edges_observer); + + // Feedback to rate the interestingness of an input + // This one is composed by two Feedbacks in OR + let feedback = feedback_or!( + // New maximization map feedback linked to the edges observer and the feedback state + MaxMapFeedback::new_tracking(&feedback_state, &edges_observer, true, false), + // Time feedback, this one does not need a feedback state + TimeFeedback::new_with_observer(&time_observer) + ); + + // A feedback to choose if an input is a solution or not + let objective = feedback_or_fast!(CrashFeedback::new(), TimeoutFeedback::new()); + + // If not restarting, create a State from scratch + let mut state = state.unwrap_or_else(|| { + StdState::new( + // RNG + StdRand::with_seed(current_nanos()), + // Corpus that will be evolved, we keep a part in memory for performance + CachedOnDiskCorpus::new(out_dir.clone(), CORPUS_CACHE_SIZE).unwrap(), + // Corpus in which we store solutions (crashes in this example), + // on disk so the user can get them after stopping the fuzzer + OnDiskCorpus::new(crashes.clone()).unwrap(), + // States of the feedbacks. + // They are the data related to the feedbacks that you want to persist in the State. + tuple_list!(feedback_state), + ) + }); + + // Create a dictionary if not existing + if let Some(tokens_file) = &self.tokens_file { + if state.metadata().get::().is_none() { + state.add_metadata(Tokens::from_tokens_file(tokens_file)?); + } + } + + // A minimization+queue policy to get testcasess from the corpus + let scheduler = + IndexesLenTimeMinimizerCorpusScheduler::new(QueueCorpusScheduler::new()); + + // A fuzzer with feedbacks and a corpus scheduler + let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); + + // Create the executor for an in-process function with one observer for edge coverage and one for the execution time + let mut executor = TimeoutForkserverExecutor::new( + ForkserverExecutor::new( + self.program.clone(), + self.arguments, + self.shmem_testcase, + tuple_list!(edges_observer, time_observer), + ) + .expect("Failed to create the executor."), + timeout, + ) + .expect("Failed to create the executor."); + + // In case the corpus is empty (on first run), reset + if state.corpus().count() < 1 { + if self.input_dirs.is_empty() { + // Generator of printable bytearrays of max size 32 + let mut generator = RandBytesGenerator::new(32); + + // Generate 8 initial inputs + state + .generate_initial_inputs( + &mut fuzzer, + &mut executor, + &mut generator, + &mut mgr, + 8, + ) + .expect("Failed to generate the initial corpus"); + println!( + "We imported {} inputs from the generator.", + state.corpus().count() + ); + } else { + println!("Loading from {:?}", &self.input_dirs); + // Load from disk + state + .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, self.input_dirs) + .unwrap_or_else(|_| { + panic!("Failed to load initial corpus at {:?}", &self.input_dirs); + }); + println!("We imported {} inputs from disk.", state.corpus().count()); + } + } + + if self.tokens_file.is_some() { + // Setup a basic mutator + let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations())); + let mutational = StdMutationalStage::new(mutator); + + // The order of the stages matter! + let mut stages = tuple_list!(mutational); + fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?; + } else { + // Setup a basic mutator + let mutator = StdScheduledMutator::new(havoc_mutations()); + let mutational = StdMutationalStage::new(mutator); + + // The order of the stages matter! + let mut stages = tuple_list!(mutational); + fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?; + } + + Ok(()) + }; + + let launcher = Launcher::builder() + .shmem_provider(shmem_provider) + .configuration(conf) + .stats(stats) + .run_client(&mut run_client) + .cores(self.cores) + .broker_port(self.broker_port) + .remote_broker_addr(self.remote_broker_addr); + #[cfg(unix)] + let launcher = launcher.stdout_file(Some("/dev/null")); + match launcher.build().launch() { + Ok(()) => (), + Err(Error::ShuttingDown) => println!("\nFuzzing stopped by user. Good Bye."), + Err(err) => panic!("Fuzzingg failed {:?}", err), + } + } +} + +/* +#[cfg(feature = "python")] +pub mod pybind { + use crate::inmemory; + use pyo3::prelude::*; + use pyo3::types::PyBytes; + use std::path::PathBuf; + + #[pyclass(unsendable)] + struct InMemoryBytesCoverageSugar { + input_dirs: Vec, + output_dir: PathBuf, + broker_port: u16, + cores: Vec, + } + + #[pymethods] + impl InMemoryBytesCoverageSugar { + #[new] + fn new( + input_dirs: Vec, + output_dir: PathBuf, + broker_port: u16, + cores: Vec, + ) -> Self { + Self { + input_dirs, + output_dir, + broker_port, + cores, + } + } + + #[allow(clippy::needless_pass_by_value)] + pub fn run(&self, harness: PyObject) { + inmemory::InMemoryBytesCoverageSugar::builder() + .input_dirs(&self.input_dirs) + .output_dir(self.output_dir.clone()) + .broker_port(self.broker_port) + .cores(&self.cores) + .harness(|buf| { + Python::with_gil(|py| -> PyResult<()> { + let args = (PyBytes::new(py, buf),); // TODO avoid copy + harness.call1(py, args)?; + Ok(()) + }) + .unwrap(); + }) + .build() + .run(); + } + } + + pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { + m.add_class::()?; + Ok(()) + } +}*/ diff --git a/libafl_sugar/src/lib.rs b/libafl_sugar/src/lib.rs index d480b5861f..6e16c960b5 100644 --- a/libafl_sugar/src/lib.rs +++ b/libafl_sugar/src/lib.rs @@ -8,6 +8,11 @@ pub mod qemu; #[cfg(target_os = "linux")] pub use qemu::QemuBytesCoverageSugar; +#[cfg(target_family = "unix")] +pub mod forkserver; +#[cfg(target_family = "unix")] +pub use forkserver::ForkserverBytesCoverageSugar; + pub const DEFAULT_TIMEOUT_SECS: u64 = 1200; pub const CORPUS_CACHE_SIZE: usize = 4096;