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

@@ -38,12 +38,11 @@ let mut shmem_buf = shmem.as_mut_slice();
 
 Here we make a shared memory region; `shmem`, and write this to environmental variable `__AFL_SHM_ID`. Then the instrumented binary, or the forkserver, finds this shared memory region (from the aforementioned env var) to record its coverage. On your fuzzer side, you can pass this shmem map to your `Observer` to obtain coverage feedbacks combined with any `Feedback`.
 
-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.
+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>