add fuzz_with_reset macro

lrubasze · lrubasze · commit f3425f8a1804 · 2026-04-10T13:36:27.000+02:00
diff --git a/README.md b/README.md
@@ -52,6 +52,38 @@ environment variable `AFL_NO_CFG_FUZZING` to `1` when building.
 [AFLplusplus]: https://aflplus.plus/
 [rust]: https://www.rust-lang.org
 
+## Resettable State (`fuzz_with_reset!`)
+
+AFL++ persistent mode runs the fuzz target in a loop. Static initialization (e.g., `OnceLock`, `lazy_static`, `once_cell::Lazy`) only executes on the first iteration — subsequent iterations skip those code paths, causing AFL's stability metric to drop.
+
+Use `fuzz_with_reset!` to provide a reset closure that clears static state after each iteration.
+
+Note: the example uses `Mutex<Option<T>>` instead of `OnceLock`/`OnceCell` because those types do not support resetting out-of-the-box.
+
+```rust
+use std::sync::Mutex;
+
+static CACHE: Mutex<Option<Vec<u8>>> = Mutex::new(None);
+
+fn main() {
+    afl::fuzz_with_reset!(|data: &[u8]| {
+        let mut cache = CACHE.lock().unwrap();
+        if cache.is_none() {
+            *cache = Some(data.to_vec());
+        }
+        drop(cache);
+        // ... fuzz logic ...
+    }, || {
+        // Reset closure: called after each successful iteration
+        *CACHE.lock().unwrap() = None;
+    });
+}
+```
+
+A `fuzz_with_reset_nohook!` variant is also available (like `fuzz_nohook!`, it does not override the panic hook).
+
+See [`afl/examples/reset_demo.rs`](afl/examples/reset_demo.rs) for a complete example.
+
 ## IJON
 
 If you want to use [IJON](https://github.com/AFLplusplus/AFLplusplus/blob/stable/docs/IJON.md) - helping fuzzer coverage through code annotation - then
diff --git a/afl/examples/reset_demo.rs b/afl/examples/reset_demo.rs
@@ -0,0 +1,42 @@
+// Demonstrates how `fuzz_with_reset!` improves AFL++ persistent mode stability
+// when using static state.
+//
+// Setup:
+//   `cargo run -p cargo-afl -- afl build --example reset_demo --manifest-path afl/Cargo.toml`
+//   `mkdir -p /tmp/afl-input && echo "test" > /tmp/afl-input/seed`
+//
+// Without reset (low stability):
+//   `AFL_NO_UI=1 cargo run -p cargo-afl -- afl fuzz \
+//     -i /tmp/afl-input -o /tmp/afl-out-bad -V 15 target/debug/examples/reset_demo`
+//
+// With reset (high stability):
+//   `USE_RESET=1 AFL_NO_UI=1 cargo run -p cargo-afl -- afl fuzz \
+//     -i /tmp/afl-input -o /tmp/afl-out-reset -V 15 target/debug/examples/reset_demo`
+//
+// Compare stability:
+//   `grep stability /tmp/afl-out-bad/default/fuzzer_stats /tmp/afl-out-reset/default/fuzzer_stats`
+
+use std::sync::Mutex;
+
+static CACHE: Mutex<Option<Vec<u8>>> = Mutex::new(None);
+
+fn main() {
+    if std::env::var("USE_RESET").is_ok() {
+        afl::fuzz_with_reset!(|data: &[u8]| { fuzz_body(data) }, || {
+            *CACHE.lock().unwrap() = None;
+        });
+    } else {
+        afl::fuzz!(|data: &[u8]| {
+            fuzz_body(data);
+        });
+    }
+}
+
+fn fuzz_body(data: &[u8]) {
+    let mut cache = CACHE.lock().unwrap();
+    if cache.is_none() {
+        *cache = Some(data.to_vec());
+    }
+    drop(cache);
+    assert!(!(data.len() > 2 && data[0] == b'x'), "crash");
+}
diff --git a/afl/src/lib.rs b/afl/src/lib.rs
@@ -242,9 +242,39 @@ pub static mut __afl_sharedmem_fuzzing: i32 = 1;
 /// });
 /// # }
 /// ```
-pub fn fuzz<F>(hook: bool, mut closure: F)
+pub fn fuzz<F>(hook: bool, closure: F)
 where
     F: FnMut(&[u8]) + std::panic::RefUnwindSafe,
+{
+    fuzz_with_reset(hook, closure, || {});
+}
+
+/// Like [`fuzz()`], but calls a `reset` closure after each successful iteration.
+///
+/// This is useful when the fuzz target uses static state (e.g., `OnceLock`, `lazy_static`)
+/// that must be cleared between iterations in AFL++ persistent mode. Without resetting,
+/// code paths that run only on the first iteration cause AFL's stability metric to drop.
+///
+/// ```rust,no_run
+/// # extern crate afl;
+/// # use afl::fuzz_with_reset;
+/// # use std::sync::Mutex;
+/// # static CACHE: Mutex<Option<Vec<u8>>> = Mutex::new(None);
+/// # fn main() {
+/// fuzz_with_reset(true, |data| {
+///     let mut cache = CACHE.lock().unwrap();
+///     if cache.is_none() {
+///         *cache = Some(data.to_vec());
+///     }
+/// }, || {
+///     *CACHE.lock().unwrap() = None;
+/// });
+/// # }
+/// ```
+pub fn fuzz_with_reset<F, R>(hook: bool, mut closure: F, mut reset: R)
+where
+    F: FnMut(&[u8]) + std::panic::RefUnwindSafe,
+    R: FnMut(),
 {
     // this marker strings needs to be in the produced executable for
     // afl-fuzz to detect `persistent mode` and `defered mode`
@@ -281,8 +311,9 @@ where
     unsafe { __afl_manual_init() };
 
     if unsafe { __afl_fuzz_ptr.is_null() } {
-        // in-memory testcase delivery is not enabled
-        // get buffer from AFL++ through stdin
+        // In-memory testcase delivery is not enabled; the target is not running
+        // in a persistent loop, so `reset` is not needed here.
+        // Get buffer from AFL++ through stdin.
         let result = io::stdin().read_to_end(&mut input);
         if result.is_err() {
             return;
@@ -322,7 +353,9 @@ where
                 // process before the stack frames are unwinded.
                 std::process::abort();
             }
+
             input.clear();
+            reset();
         }
     }
 }
@@ -363,25 +396,72 @@ macro_rules! fuzz_nohook {
 
 #[doc(hidden)]
 #[macro_export]
-macro_rules! __fuzz {
-    ($hook:expr, |$buf:ident| $body:expr) => {
-        $crate::fuzz($hook, |$buf| $body);
+macro_rules! __reset_or_noop {
+    () => {
+        || {}
     };
-    ($hook:expr, |$buf:ident: &[u8]| $body:expr) => {
-        $crate::fuzz($hook, |$buf| $body);
+    ($reset:expr) => {
+        $reset
     };
-    ($hook:expr, |$buf:ident: $dty: ty| $body:expr) => {
-        $crate::fuzz($hook, |$buf| {
-            let $buf: $dty = {
-                let mut data = ::arbitrary::Unstructured::new($buf);
-                if let Ok(d) = ::arbitrary::Arbitrary::arbitrary(&mut data).map_err(|_| "") {
-                    d
-                } else {
-                    return;
-                }
-            };
+}
 
-            $body
-        });
+#[doc(hidden)]
+#[macro_export]
+macro_rules! __fuzz {
+    ($hook:expr, |$buf:ident| $body:expr $(, $reset:expr)?) => {
+        $crate::fuzz_with_reset($hook, |$buf| $body, $crate::__reset_or_noop!($($reset)?));
+    };
+    ($hook:expr, |$buf:ident: &[u8]| $body:expr $(, $reset:expr)?) => {
+        $crate::fuzz_with_reset($hook, |$buf| $body, $crate::__reset_or_noop!($($reset)?));
+    };
+    ($hook:expr, |$buf:ident: $dty: ty| $body:expr $(, $reset:expr)?) => {
+        $crate::fuzz_with_reset(
+            $hook,
+            |$buf| {
+                let $buf: $dty = {
+                    let mut data = ::arbitrary::Unstructured::new($buf);
+                    if let Ok(d) = ::arbitrary::Arbitrary::arbitrary(&mut data).map_err(|_| "") {
+                        d
+                    } else {
+                        return;
+                    }
+                };
+
+                $body
+            },
+            $crate::__reset_or_noop!($($reset)?),
+        );
     };
 }
+
+/// Like [`fuzz!`], but accepts a second closure that resets state after each iteration.
+///
+/// This is useful when the fuzz target uses static state (e.g., `OnceLock`, `lazy_static`)
+/// that must be cleared between iterations in AFL++ persistent mode.
+///
+/// ```rust,no_run
+/// # #[macro_use] extern crate afl;
+/// # use std::sync::Mutex;
+/// # static CACHE: Mutex<Option<Vec<u8>>> = Mutex::new(None);
+/// # fn main() {
+/// fuzz_with_reset!(|data: &[u8]| {
+///     let mut cache = CACHE.lock().unwrap();
+///     if cache.is_none() {
+///         *cache = Some(data.to_vec());
+///     }
+/// }, || {
+///     *CACHE.lock().unwrap() = None;
+/// });
+/// # }
+/// ```
+#[macro_export]
+macro_rules! fuzz_with_reset {
+    ( $($x:tt)* ) => { $crate::__fuzz!(true, $($x)*) }
+}
+
+/// Like [`fuzz_with_reset!`], but panics that are caught inside the fuzzed code are not turned
+/// into crashes.
+#[macro_export]
+macro_rules! fuzz_with_reset_nohook {
+    ( $($x:tt)* ) => { $crate::__fuzz!(false, $($x)*) }
+}
diff --git a/cargo-afl/tests/integration.rs b/cargo-afl/tests/integration.rs
@@ -100,32 +100,85 @@ fn integration_maze() {
     unreachable!();
 }
 
+#[test]
+fn integration_fuzz_with_reset() {
+    // Run without reset (expect low stability)
+    let dir_no_reset = fuzz_example_with_envs("reset_demo", 15, &[]);
+
+    // Run with reset (expect high stability)
+    let dir_with_reset = fuzz_example_with_envs("reset_demo", 15, &[("USE_RESET", "1")]);
+
+    let stability_no_reset = parse_stability(dir_no_reset.path());
+    let stability_with_reset = parse_stability(dir_with_reset.path());
+
+    // On Linux/x86_64 we observe ~95% stability, on macOS/aarch64 ~85%
+    // due to ARM's relaxed memory model affecting bitmap synchronization.
+    let min_stability_expected = 80.0;
+
+    assert!(
+        stability_no_reset < min_stability_expected,
+        "Stability without reset ({stability_no_reset}%) should be below {min_stability_expected}%"
+    );
+    assert!(
+        stability_with_reset > min_stability_expected,
+        "Stability with reset ({stability_with_reset}%) should be above {min_stability_expected}%"
+    );
+}
+
 fn fuzz_example(name: &str, should_crash: bool) {
-    let temp_dir = tempfile::TempDir::new().expect("Could not create temporary directory");
+    let temp_dir = fuzz_example_with_envs(name, 5, &[("AFL_BENCH_UNTIL_CRASH", "1")]);
     let temp_dir_path = temp_dir.path();
+    assert!(temp_dir_path.join("default").join("fuzzer_stats").is_file());
+    let crashes = std::fs::read_dir(temp_dir_path.join("default").join("crashes"))
+        .unwrap()
+        .count();
+    if should_crash {
+        assert!(crashes >= 1);
+    } else {
+        assert_eq!(0, crashes);
+    }
+}
+
+fn fuzz_example_with_envs(
+    name: &str,
+    timeout_secs: u32,
+    envs: &[(&str, &str)],
+) -> tempfile::TempDir {
+    let temp_dir = tempfile::TempDir::new().expect("Could not create temporary directory");
     let _: ExitStatus = process::Command::new(cargo_afl_path())
         .arg("afl")
         .arg("fuzz")
         .arg("-i")
         .arg(input_path())
         .arg("-o")
-        .arg(temp_dir_path)
-        .args(["-V", "5"]) // 5 seconds
+        .arg(temp_dir.path())
+        .args(["-V", &timeout_secs.to_string()])
         .arg(examples_path(name))
-        .env("AFL_BENCH_UNTIL_CRASH", "1")
         .env("AFL_NO_CRASH_README", "1")
         .env("AFL_NO_UI", "1")
+        .envs(envs.iter().copied())
         .stdout(process::Stdio::inherit())
         .stderr(process::Stdio::inherit())
         .status()
         .expect("Could not run cargo afl fuzz");
-    assert!(temp_dir_path.join("default").join("fuzzer_stats").is_file());
-    let crashes = std::fs::read_dir(temp_dir_path.join("default").join("crashes"))
-        .unwrap()
-        .count();
-    if should_crash {
-        assert!(crashes >= 1);
-    } else {
-        assert_eq!(0, crashes);
+    temp_dir
+}
+
+fn parse_stability(output_dir: &path::Path) -> f64 {
+    let stats_path = output_dir.join("default").join("fuzzer_stats");
+    let contents = std::fs::read_to_string(&stats_path)
+        .unwrap_or_else(|e| panic!("Failed to read {}: {e}", stats_path.display()));
+    for line in contents.lines() {
+        if let Some(value) = line.strip_prefix("stability") {
+            let value = value
+                .trim()
+                .trim_start_matches(':')
+                .trim()
+                .trim_end_matches('%');
+            return value
+                .parse()
+                .unwrap_or_else(|e| panic!("Failed to parse stability value '{value}': {e}"));
+        }
     }
+    panic!("No stability line found in {}", stats_path.display());
 }
