Restart loading initial inputs even after a crash/timeout (#1040)

* Track initial inputs loading

* libfuzzer libpng

* fuzzbench

* fix no_std

* fix no_std

* clippy

* fuzzers

fuzzers/forkserver_libafl_cc/src/main.rs

@@ -190,7 +190,7 @@ pub fn main() {
     .expect("Failed to create the executor.");
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &corpus_dirs)
             .unwrap_or_else(|err| {

fuzzers/forkserver_simple/src/main.rs

@@ -190,7 +190,7 @@ pub fn main() {
     .expect("Failed to create the executor.");
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &corpus_dirs)
             .unwrap_or_else(|err| {

fuzzers/frida_gdiplus/src/fuzzer.rs

@@ -198,7 +198,7 @@ unsafe fn fuzz(options: &FuzzerOptions) -> Result<(), Error> {
                 );
 
                 // In case the corpus is empty (on first run), reset
-                if state.corpus().count() < 1 {
+                if state.must_load_initial_inputs() {
                     state
                         .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &options.input)
                         .unwrap_or_else(|_| {
@@ -313,7 +313,7 @@ unsafe fn fuzz(options: &FuzzerOptions) -> Result<(), Error> {
                 );
 
                 // In case the corpus is empty (on first run), reset
-                if state.corpus().count() < 1 {
+                if state.must_load_initial_inputs() {
                     state
                         .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &options.input)
                         .unwrap_or_else(|_| {
@@ -443,7 +443,7 @@ unsafe fn fuzz(options: &FuzzerOptions) -> Result<(), Error> {
                 );
 
                 // In case the corpus is empty (on first run), reset
-                if state.corpus().count() < 1 {
+                if state.must_load_initial_inputs() {
                     state
                         .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &options.input)
                         .unwrap_or_else(|_| {

fuzzers/frida_libpng/src/fuzzer.rs

@@ -196,7 +196,7 @@ unsafe fn fuzz(options: &FuzzerOptions) -> Result<(), Error> {
                 );
 
                 // In case the corpus is empty (on first run), reset
-                if state.corpus().count() < 1 {
+                if state.must_load_initial_inputs() {
                     state
                         .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &options.input)
                         .unwrap_or_else(|_| {
@@ -311,7 +311,7 @@ unsafe fn fuzz(options: &FuzzerOptions) -> Result<(), Error> {
                 );
 
                 // In case the corpus is empty (on first run), reset
-                if state.corpus().count() < 1 {
+                if state.must_load_initial_inputs() {
                     state
                         .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &options.input)
                         .unwrap_or_else(|_| {
@@ -441,7 +441,7 @@ unsafe fn fuzz(options: &FuzzerOptions) -> Result<(), Error> {
                 );
 
                 // In case the corpus is empty (on first run), reset
-                if state.corpus().count() < 1 {
+                if state.must_load_initial_inputs() {
                     state
                         .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &options.input)
                         .unwrap_or_else(|_| {

fuzzers/fuzzbench/src/lib.rs

@@ -369,7 +369,7 @@ fn fuzz(
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &[seed_dir.clone()])
             .unwrap_or_else(|_| {

fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs

@@ -356,7 +356,7 @@ fn fuzz(
         }
     }
 
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &[seed_dir.clone()])
             .unwrap_or_else(|_| {

fuzzers/fuzzbench_qemu/src/fuzzer.rs

@@ -371,7 +371,7 @@ fn fuzz(
         }
     }
 
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &[seed_dir.clone()])
             .unwrap_or_else(|_| {

fuzzers/fuzzbench_text/src/lib.rs

@@ -430,7 +430,7 @@ fn fuzz_binary(
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &[seed_dir.clone()])
             .unwrap_or_else(|_| {
@@ -647,7 +647,7 @@ fn fuzz_text(
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &[seed_dir.clone()])
             .unwrap_or_else(|_| {

fuzzers/libafl_atheris/src/lib.rs

@@ -318,7 +318,7 @@ pub fn LLVMFuzzerRunDriver(
         let mut stages = tuple_list!(tracing, i2s, mutational);
 
         // In case the corpus is empty (on first run), reset
-        if state.corpus().count() < 1 {
+        if state.must_load_initial_inputs() {
             if input_dirs.is_empty() {
                 // Generator of printable bytearrays of max size 32
                 let mut generator = RandBytesGenerator::new(32);

fuzzers/libfuzzer_libmozjpeg/src/lib.rs

@@ -157,7 +157,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, corpus_dirs)
             .unwrap_or_else(|_| panic!("Failed to load initial corpus at {:?}", &corpus_dirs));

fuzzers/libfuzzer_libpng/src/lib.rs

@@ -191,7 +191,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, corpus_dirs)
             .unwrap_or_else(|_| panic!("Failed to load initial corpus at {:?}", &corpus_dirs));

fuzzers/libfuzzer_libpng_accounting/src/lib.rs

@@ -226,7 +226,7 @@ pub fn libafl_main() {
         }
 
         // In case the corpus is empty (on first run), reset
-        if state.corpus().count() < 1 {
+        if state.must_load_initial_inputs() {
             state
                 .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, &opt.input)
                 .unwrap_or_else(|e| {

fuzzers/libfuzzer_libpng_cmin/src/lib.rs

@@ -190,7 +190,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, corpus_dirs)
             .unwrap_or_else(|_| panic!("Failed to load initial corpus at {:?}", &corpus_dirs));

fuzzers/libfuzzer_libpng_ctx/src/lib.rs

@@ -217,7 +217,7 @@ pub fn libafl_main() {
         }
 
         // In case the corpus is empty (on first run), reset
-        if state.corpus().count() < 1 {
+        if state.must_load_initial_inputs() {
             state
                 .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, &opt.input)
                 .unwrap_or_else(|_| panic!("Failed to load initial corpus at {:?}", &opt.input));

fuzzers/libfuzzer_libpng_launcher/src/lib.rs

@@ -222,7 +222,7 @@ pub fn libafl_main() {
         }
 
         // In case the corpus is empty (on first run), reset
-        if state.corpus().count() < 1 {
+        if state.must_load_initial_inputs() {
             state
                 .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, &opt.input)
                 .unwrap_or_else(|_| panic!("Failed to load initial corpus at {:?}", &opt.input));

fuzzers/libfuzzer_reachability/src/lib.rs

@@ -136,7 +136,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, corpus_dirs)
             .unwrap_or_else(|_| panic!("Failed to load initial corpus at {:?}", corpus_dirs));

fuzzers/libfuzzer_stb_image/src/main.rs

@@ -142,7 +142,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, corpus_dirs)
             .unwrap_or_else(|_| panic!("Failed to load initial corpus at {:?}", corpus_dirs));

fuzzers/libfuzzer_stb_image_concolic/fuzzer/src/main.rs

@@ -179,7 +179,7 @@ fn fuzz(
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, corpus_dirs)
             .unwrap_or_else(|_| panic!("Failed to load initial corpus at {corpus_dirs:?}"));

fuzzers/libfuzzer_windows_asan/src/lib.rs

@@ -154,7 +154,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, corpus_dirs)
             .unwrap_or_else(|_| panic!("Failed to load initial corpus at {:?}", &corpus_dirs));

fuzzers/qemu_arm_launcher/src/fuzzer.rs

@@ -195,7 +195,7 @@ pub fn fuzz() {
         // Wrap the executor to keep track of the timeout
         let mut executor = TimeoutExecutor::new(executor, timeout);
 
-        if state.corpus().count() < 1 {
+        if state.must_load_initial_inputs() {
             state
                 .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &corpus_dirs)
                 .unwrap_or_else(|_| {

fuzzers/qemu_launcher/src/fuzzer.rs

@@ -176,7 +176,7 @@ pub fn fuzz() {
         // Wrap the executor to keep track of the timeout
         let mut executor = TimeoutExecutor::new(executor, timeout);
 
-        if state.corpus().count() < 1 {
+        if state.must_load_initial_inputs() {
             state
                 .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &corpus_dirs)
                 .unwrap_or_else(|_| {

fuzzers/qemu_systemmode/src/fuzzer.rs

@@ -204,7 +204,7 @@ pub fn fuzz() {
         // Wrap the executor to keep track of the timeout
         let mut executor = TimeoutExecutor::new(executor, timeout);
 
-        if state.corpus().count() < 1 {
+        if state.must_load_initial_inputs() {
             state
                 .load_initial_inputs(&mut fuzzer, &mut executor, &mut mgr, &corpus_dirs)
                 .unwrap_or_else(|_| {

fuzzers/tutorial/src/lib.rs

@@ -158,7 +158,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
     }
 
     // In case the corpus is empty (on first run), reset
-    if state.corpus().count() < 1 {
+    if state.must_load_initial_inputs() {
         state
             .load_initial_inputs(&mut fuzzer, &mut executor, &mut restarting_mgr, corpus_dirs)
             .unwrap_or_else(|_| panic!("Failed to load initial corpus at {:?}", &corpus_dirs));

libafl/src/state/mod.rs

@@ -5,6 +5,7 @@ use core::{fmt::Debug, marker::PhantomData, time::Duration};
 use std::{
     fs,
     path::{Path, PathBuf},
+    vec::Vec,
 };
 
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
@@ -194,6 +195,9 @@ pub struct StdState<I, C, R, SC> {
     /// Performance statistics for this fuzzer
     #[cfg(feature = "introspection")]
     introspection_monitor: ClientPerfMonitor,
+    #[cfg(feature = "std")]
+    /// Remaining initial inputs to load, if any
+    remaining_initial_files: Option<Vec<PathBuf>>,
     phantom: PhantomData<I>,
 }
 
@@ -347,23 +351,15 @@ where
     R: Rand,
     SC: Corpus<Input = <Self as UsesInput>::Input>,
 {
-    /// Loads inputs from a directory.
-    /// If `forced` is `true`, the value will be loaded,
-    /// even if it's not considered to be `interesting`.
-    pub fn load_from_directory<E, EM, Z>(
-        &mut self,
-        fuzzer: &mut Z,
-        executor: &mut E,
-        manager: &mut EM,
-        in_dir: &Path,
-        forced: bool,
-        loader: &mut dyn FnMut(&mut Z, &mut Self, &Path) -> Result<I, Error>,
-    ) -> Result<(), Error>
-    where
-        E: UsesState<State = Self>,
-        EM: UsesState<State = Self>,
-        Z: Evaluator<E, EM, State = Self>,
-    {
+    /// Decide if the state nust load the inputs
+    pub fn must_load_initial_inputs(&self) -> bool {
+        self.corpus().count() == 0
+            || (self.remaining_initial_files.is_some()
+                && !self.remaining_initial_files.as_ref().unwrap().is_empty())
+    }
+
+    /// List initial inputs from a directory.
+    fn visit_initial_directory(files: &mut Vec<PathBuf>, in_dir: &Path) -> Result<(), Error> {
         for entry in fs::read_dir(in_dir)? {
             let entry = entry?;
             let path = entry.path();
@@ -380,6 +376,47 @@ where
             let attr = attributes?;
 
             if attr.is_file() && attr.len() > 0 {
+                files.push(path);
+            } else if attr.is_dir() {
+                Self::visit_initial_directory(files, in_dir)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Loads initial inputs from the passed-in `in_dirs`.
+    /// If `forced` is true, will add all testcases, no matter what.
+    fn load_initial_inputs_custom<E, EM, Z>(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        manager: &mut EM,
+        in_dirs: &[PathBuf],
+        forced: bool,
+        loader: &mut dyn FnMut(&mut Z, &mut Self, &Path) -> Result<I, Error>,
+    ) -> Result<(), Error>
+    where
+        E: UsesState<State = Self>,
+        EM: EventFirer<State = Self>,
+        Z: Evaluator<E, EM, State = Self>,
+    {
+        if let Some(remaining) = self.remaining_initial_files.as_ref() {
+            // everything was loaded
+            if remaining.is_empty() {
+                return Ok(());
+            }
+        } else {
+            let mut files = vec![];
+            for in_dir in in_dirs {
+                Self::visit_initial_directory(&mut files, in_dir)?;
+            }
+
+            self.remaining_initial_files = Some(files);
+        }
+
+        // TODO option to shuffle the initial files
+        while let Some(path) = self.remaining_initial_files.as_mut().unwrap().pop() {
             println!("Loading file {:?} ...", &path);
             let input = loader(fuzzer, self, &path)?;
             if forced {
@@ -390,39 +427,8 @@ where
                     println!("File {:?} was not interesting, skipped.", &path);
                 }
             }
-            } else if attr.is_dir() {
-                self.load_from_directory(fuzzer, executor, manager, &path, forced, loader)?;
-            }
         }
 
-        Ok(())
-    }
-
-    /// Loads initial inputs from the passed-in `in_dirs`.
-    /// If `forced` is true, will add all testcases, no matter what.
-    fn load_initial_inputs_internal<E, EM, Z>(
-        &mut self,
-        fuzzer: &mut Z,
-        executor: &mut E,
-        manager: &mut EM,
-        in_dirs: &[PathBuf],
-        forced: bool,
-    ) -> Result<(), Error>
-    where
-        E: UsesState<State = Self>,
-        EM: EventFirer<State = Self>,
-        Z: Evaluator<E, EM, State = Self>,
-    {
-        for in_dir in in_dirs {
-            self.load_from_directory(
-                fuzzer,
-                executor,
-                manager,
-                in_dir,
-                forced,
-                &mut |_, _, path| I::from_file(path),
-            )?;
-        }
         manager.fire(
             self,
             Event::Log {
@@ -449,7 +455,14 @@ where
         EM: EventFirer<State = Self>,
         Z: Evaluator<E, EM, State = Self>,
     {
-        self.load_initial_inputs_internal(fuzzer, executor, manager, in_dirs, true)
+        self.load_initial_inputs_custom(
+            fuzzer,
+            executor,
+            manager,
+            in_dirs,
+            true,
+            &mut |_, _, path| I::from_file(path),
+        )
     }
 
     /// Loads initial inputs from the passed-in `in_dirs`.
@@ -465,7 +478,14 @@ where
         EM: EventFirer<State = Self>,
         Z: Evaluator<E, EM, State = Self>,
     {
-        self.load_initial_inputs_internal(fuzzer, executor, manager, in_dirs, false)
+        self.load_initial_inputs_custom(
+            fuzzer,
+            executor,
+            manager,
+            in_dirs,
+            false,
+            &mut |_, _, path| I::from_file(path),
+        )
     }
 }
 
@@ -574,6 +594,8 @@ where
             max_size: DEFAULT_MAX_SIZE,
             #[cfg(feature = "introspection")]
             introspection_monitor: ClientPerfMonitor::new(),
+            #[cfg(feature = "std")]
+            remaining_initial_files: None,
             phantom: PhantomData,
         };
         feedback.init_state(&mut state)?;

libafl_sugar/src/forkserver.rs

@@ -196,7 +196,7 @@ impl<'a, const MAP_SIZE: usize> ForkserverBytesCoverageSugar<'a, MAP_SIZE> {
             .expect("Failed to create the executor.");
 
             // In case the corpus is empty (on first run), reset
-            if state.corpus().count() < 1 {
+            if state.must_load_initial_inputs() {
                 if self.input_dirs.is_empty() {
                     // Generator of printable bytearrays of max size 32
                     let mut generator = RandBytesGenerator::new(32);

libafl_sugar/src/inmemory.rs

@@ -216,7 +216,7 @@ where
             );
 
             // In case the corpus is empty (on first run), reset
-            if state.corpus().count() < 1 {
+            if state.must_load_initial_inputs() {
                 if self.input_dirs.is_empty() {
                     // Generator of printable bytearrays of max size 32
                     let mut generator = RandBytesGenerator::new(32);

libafl_sugar/src/qemu.rs

@@ -232,7 +232,7 @@ where
                 let mut executor = ShadowExecutor::new(executor, tuple_list!(cmplog_observer));
 
                 // In case the corpus is empty (on first run), reset
-                if state.corpus().count() < 1 {
+                if state.must_load_initial_inputs() {
                     if self.input_dirs.is_empty() {
                         // Generator of printable bytearrays of max size 32
                         let mut generator = RandBytesGenerator::new(32);
@@ -335,7 +335,7 @@ where
                 let mut executor = TimeoutExecutor::new(executor, timeout);
 
                 // In case the corpus is empty (on first run), reset
-                if state.corpus().count() < 1 {
+                if state.must_load_initial_inputs() {
                     if self.input_dirs.is_empty() {
                         // Generator of printable bytearrays of max size 32
                         let mut generator = RandBytesGenerator::new(32);