Forkserver: Add file input support (#880)

* make use of clap derive in forkserver_simple

* (re)introduce use_shmem_testcase flag to ForkserverExecutor

* set use_shmem_testcase flag automatically based on forkserver handshake

* remove illegal_state and just .unwrap instead as the None case is unreachable

* fix: removed pub method

* cargo fmt

* remove illegal_state #2 and just .unwrap instead as the None case is unreachable

* change shmem unwrap to unwrap_unchecked

* fix double mut

* removed @@ warning

docs/src/core_concepts/executor.md

@@ -40,10 +40,9 @@ Here we make a shared memory region; `shmem`, and write this to environmental va
 
 Another feature of the `ForkserverExecutor` to mention is the shared memory testcases. In normal cases, the mutated input is passed between the forkserver and the instrumented binary via `.cur_input` file. You can improve your forkserver fuzzer's performance by passing the input with shared memory.  
 
+If the target is configured to use shared memory testcases, the `ForkserverExecutor` will notice this during the handshake and will automatically set up things accordingly.
 See AFL++'s [_documentation_](https://github.com/AFLplusplus/AFLplusplus/blob/stable/instrumentation/README.persistent_mode.md#5-shared-memory-fuzzing) or the fuzzer example in `forkserver_simple/src/program.c` for reference.
 
-It is very simple, when you call `ForkserverExecutor::new()` with `use_shmem_testcase` true, the `ForkserverExecutor` sets things up and your harness can just fetch the input from `__AFL_FUZZ_TESTCASE_BUF`
-
 ## InprocessForkExecutor
 
 Finally, we'll talk about the `InProcessForkExecutor`.

fuzzers/forkserver_simple/Cargo.toml

@@ -16,6 +16,6 @@ codegen-units = 1
 opt-level = 3
 
 [dependencies]
-libafl = { path = "../../libafl/" }
+libafl = { path = "../../libafl/", features = ["std", "derive"] }
-clap = { version = "3.2", features = ["default"] }
+clap = { version = "4.0", features = ["derive"] }
 nix = "0.25"

fuzzers/forkserver_simple/src/main.rs

@@ -1,7 +1,7 @@
 use core::time::Duration;
 use std::path::PathBuf;
 
-use clap::{Arg, Command};
+use clap::{self, Parser};
 #[cfg(not(target_vendor = "apple"))]
 use libafl::bolts::shmem::StdShMemProvider;
 #[cfg(target_vendor = "apple")]
@@ -30,53 +30,68 @@ use libafl::{
 };
 use nix::sys::signal::Signal;
 
+/// The commandline args this fuzzer accepts
+#[derive(Debug, Parser)]
+#[command(
+    name = "forkserver_simple",
+    about = "This is a simple example fuzzer to fuzz a executable instrumented by afl-cc.",
+    author = "tokatoka <tokazerkje@outlook.com>"
+)]
+struct Opt {
+    #[arg(
+        help = "The instrumented binary we want to fuzz",
+        name = "EXEC",
+        required = true
+    )]
+    executable: String,
+
+    #[arg(
+        help = "The directory to read initial inputs from ('seeds')",
+        name = "INPUT_DIR",
+        required = true
+    )]
+    in_dir: PathBuf,
+
+    #[arg(
+        help = "Timeout for each individual execution, in milliseconds",
+        short = 't',
+        long = "timeout",
+        default_value = "1200"
+    )]
+    timeout: u64,
+
+    #[arg(
+        help = "If not set, the child's stdout and stderror will be redirected to /dev/null",
+        short = 'd',
+        long = "debug-child",
+        default_value = "false"
+    )]
+    debug_child: bool,
+
+    #[arg(
+        help = "Arguments passed to the target",
+        name = "arguments",
+        num_args(1..),
+        allow_hyphen_values = true,
+        default_value = "[].to_vec()"
+    )]
+    arguments: Vec<String>,
+
+    #[arg(
+        help = "Signal used to stop child",
+        short = 's',
+        long = "signal",
+        value_parser = str::parse::<Signal>,
+        default_value = "SIGKILL"
+    )]
+    signal: Signal,
+}
+
 #[allow(clippy::similar_names)]
 pub fn main() {
-    let res = Command::new("forkserver_simple")
+    let opt = Opt::parse();
-        .about("Example Forkserver fuzer")
-        .arg(
-            Arg::new("executable")
-                .help("The instrumented binary we want to fuzz")
-                .required(true)
-                .takes_value(true),
-        )
-        .arg(
-            Arg::new("in")
-                .help("The directory to read initial inputs from ('seeds')")
-                .required(true)
-                .takes_value(true),
-        )
-        .arg(
-            Arg::new("timeout")
-                .help("Timeout for each individual execution, in milliseconds")
-                .short('t')
-                .long("timeout")
-                .default_value("1200"),
-        )
-        .arg(
-            Arg::new("debug_child")
-                .help("If not set, the child's stdout and stderror will be redirected to /dev/null")
-                .short('d')
-                .long("debug-child"),
-        )
-        .arg(
-            Arg::new("arguments")
-                .help("Arguments passed to the target")
-                .multiple_values(true)
-                .takes_value(true),
-        )
-        .arg(
-            Arg::new("signal")
-                .help("Signal used to stop child")
-                .short('s')
-                .long("signal")
-                .default_value("SIGKILL"),
-        )
-        .get_matches();
 
-    let corpus_dirs = vec![PathBuf::from(
+    let corpus_dirs: Vec<PathBuf> = [opt.in_dir].to_vec();
-        res.get_one::<String>("in").unwrap().to_string(),
-    )];
 
     const MAP_SIZE: usize = 65536;
 
@@ -151,17 +166,14 @@ pub fn main() {
     let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective);
 
     // If we should debug the child
-    let debug_child = res.is_present("debug_child");
+    let debug_child = opt.debug_child;
 
     // Create the executor for the forkserver
-    let args = match res.values_of("arguments") {
+    let args = opt.arguments;
-        Some(vec) => vec.map(|s| s.to_string()).collect::<Vec<String>>().to_vec(),
-        None => [].to_vec(),
-    };
 
     let mut tokens = Tokens::new();
     let forkserver = ForkserverExecutor::builder()
-        .program(res.get_one::<String>("executable").unwrap())
+        .program(opt.executable)
         .debug_child(debug_child)
         .shmem_provider(&mut shmem_provider)
         .autotokens(&mut tokens)
@@ -171,17 +183,8 @@ pub fn main() {
 
     let mut executor = TimeoutForkserverExecutor::with_signal(
         forkserver,
-        Duration::from_millis(
+        Duration::from_millis(opt.timeout),
-            res.get_one::<String>("timeout")
+        opt.signal,
-                .unwrap()
-                .to_string()
-                .parse()
-                .expect("Could not parse timeout in milliseconds"),
-        ),
-        res.get_one::<String>("signal")
-            .unwrap()
-            .parse::<Signal>()
-            .unwrap(),
     )
     .expect("Failed to create the executor.");
 

libafl/src/executors/forkserver.rs

@@ -369,6 +369,9 @@ pub trait HasForkserver {
 
     /// The map of the fuzzer, mutable
     fn shmem_mut(&mut self) -> &mut Option<<<Self as HasForkserver>::SP as ShMemProvider>::ShMem>;
+
+    /// Whether testcases are expected in shared memory
+    fn uses_shmem_testcase(&self) -> bool;
 }
 
 /// The timeout forkserver executor that wraps around the standard forkserver executor and sets a timeout before each run.
@@ -417,21 +420,19 @@ where
 
         let last_run_timed_out = self.executor.forkserver().last_run_timed_out();
 
-        match &mut self.executor.shmem_mut() {
+        if self.executor.uses_shmem_testcase() {
-            Some(shmem) => {
+            let shmem = unsafe { self.executor.shmem_mut().as_mut().unwrap_unchecked() };
-                let target_bytes = input.target_bytes();
+            let target_bytes = input.target_bytes();
-                let size = target_bytes.as_slice().len();
+            let size = target_bytes.as_slice().len();
-                let size_in_bytes = size.to_ne_bytes();
+            let size_in_bytes = size.to_ne_bytes();
-                // The first four bytes tells the size of the shmem.
+            // The first four bytes tells the size of the shmem.
-                shmem.as_mut_slice()[..4].copy_from_slice(&size_in_bytes[..4]);
+            shmem.as_mut_slice()[..4].copy_from_slice(&size_in_bytes[..4]);
-                shmem.as_mut_slice()[SHMEM_FUZZ_HDR_SIZE..(SHMEM_FUZZ_HDR_SIZE + size)]
+            shmem.as_mut_slice()[SHMEM_FUZZ_HDR_SIZE..(SHMEM_FUZZ_HDR_SIZE + size)]
-                    .copy_from_slice(target_bytes.as_slice());
+                .copy_from_slice(target_bytes.as_slice());
-            }
+        } else {
-            None => {
+            self.executor
-                self.executor
+                .input_file_mut()
-                    .input_file_mut()
+                .write_buf(input.target_bytes().as_slice())?;
-                    .write_buf(input.target_bytes().as_slice())?;
-            }
         }
 
         let send_len = self
@@ -503,6 +504,7 @@ where
     target: OsString,
     args: Vec<OsString>,
     input_file: InputFile,
+    uses_shmem_testcase: bool,
     forkserver: Forkserver,
     observers: OT,
     map: Option<SP::ShMem>,
@@ -521,6 +523,7 @@ where
             .field("target", &self.target)
             .field("args", &self.args)
             .field("input_file", &self.input_file)
+            .field("use_shmem_testcase", &self.uses_shmem_testcase)
             .field("forkserver", &self.forkserver)
             .field("observers", &self.observers)
             .field("map", &self.map)
@@ -572,6 +575,7 @@ pub struct ForkserverExecutorBuilder<'a, SP> {
     envs: Vec<(OsString, OsString)>,
     debug_child: bool,
     use_stdin: bool,
+    uses_shmem_testcase: bool,
     is_persistent: bool,
     is_deferred_frksrv: bool,
     autotokens: Option<&'a mut Tokens>,
@@ -651,6 +655,7 @@ impl<'a, SP> ForkserverExecutorBuilder<'a, SP> {
             if (status & FS_OPT_SHDMEM_FUZZ == FS_OPT_SHDMEM_FUZZ) && map.is_some() {
                 println!("Using SHARED MEMORY FUZZING feature.");
                 send_status |= FS_OPT_SHDMEM_FUZZ;
+                self.uses_shmem_testcase = true;
             }
 
             if (status & FS_OPT_AUTODICT == FS_OPT_AUTODICT) && self.autotokens.is_some() {
@@ -704,6 +709,7 @@ impl<'a, SP> ForkserverExecutorBuilder<'a, SP> {
             target,
             args: self.arguments.clone(),
             input_file,
+            uses_shmem_testcase: self.uses_shmem_testcase,
             forkserver,
             observers,
             map,
@@ -765,6 +771,7 @@ impl<'a> ForkserverExecutorBuilder<'a, StdShMemProvider> {
             envs: vec![],
             debug_child: false,
             use_stdin: true,
+            uses_shmem_testcase: false,
             is_persistent: false,
             is_deferred_frksrv: false,
             autotokens: None,
@@ -882,6 +889,7 @@ impl<'a> ForkserverExecutorBuilder<'a, StdShMemProvider> {
             envs: self.envs,
             debug_child: self.debug_child,
             use_stdin: self.use_stdin,
+            uses_shmem_testcase: self.uses_shmem_testcase,
             is_persistent: self.is_persistent,
             is_deferred_frksrv: self.is_deferred_frksrv,
             autotokens: self.autotokens,
@@ -917,20 +925,18 @@ where
         let mut exit_kind = ExitKind::Ok;
 
         // Write to testcase
-        match &mut self.map {
+        if self.uses_shmem_testcase {
-            Some(map) => {
+            let map = unsafe { self.map.as_mut().unwrap_unchecked() };
-                let target_bytes = input.target_bytes();
+            let target_bytes = input.target_bytes();
-                let size = target_bytes.as_slice().len();
+            let size = target_bytes.as_slice().len();
-                let size_in_bytes = size.to_ne_bytes();
+            let size_in_bytes = size.to_ne_bytes();
-                // The first four bytes tells the size of the shmem.
+            // The first four bytes tells the size of the shmem.
-                map.as_mut_slice()[..SHMEM_FUZZ_HDR_SIZE]
+            map.as_mut_slice()[..SHMEM_FUZZ_HDR_SIZE]
-                    .copy_from_slice(&size_in_bytes[..SHMEM_FUZZ_HDR_SIZE]);
+                .copy_from_slice(&size_in_bytes[..SHMEM_FUZZ_HDR_SIZE]);
-                map.as_mut_slice()[SHMEM_FUZZ_HDR_SIZE..(SHMEM_FUZZ_HDR_SIZE + size)]
+            map.as_mut_slice()[SHMEM_FUZZ_HDR_SIZE..(SHMEM_FUZZ_HDR_SIZE + size)]
-                    .copy_from_slice(target_bytes.as_slice());
+                .copy_from_slice(target_bytes.as_slice());
-            }
+        } else {
-            None => {
+            self.input_file.write_buf(input.target_bytes().as_slice())?;
-                self.input_file.write_buf(input.target_bytes().as_slice())?;
-            }
         }
 
         let send_len = self
@@ -1061,6 +1067,11 @@ where
     fn shmem_mut(&mut self) -> &mut Option<SP::ShMem> {
         &mut self.map
     }
+
+    #[inline]
+    fn uses_shmem_testcase(&self) -> bool {
+        self.uses_shmem_testcase
+    }
 }
 
 impl<E> UsesState for TimeoutForkserverExecutor<E>