Fix AFLStatsStage (#3371)

* refactor

* type puzzle

Author: tokatoka

crates/libafl/src/common/nautilus/README.md

@@ -1,4 +1,4 @@
-# Nautilus 2.0 LibAFL Mutator
+# Nautilus 2.0 `LibAFL` Mutator
 
 Nautilus is a coverage guided, grammar-based mutator. You can use it to improve your test coverage and find more bugs. By specifying the grammar of semi-valid inputs, Nautilus is able to perform complex mutation and to uncover more interesting test cases. Many of the ideas behind the original fuzzer are documented in a paper published at NDSS 2019.
 
@@ -7,7 +7,7 @@ Nautilus is a coverage guided, grammar-based mutator. You can use it to improve
 </p>
 
 Version 2.0 has added many improvements to this early prototype.
-Features from version 2.0 we support in LibAFL:
+Features from version 2.0 we support in `LibAFL`:
 
 * Support for grammars specified in python
 * Support for non-context free grammars using python scripts to generate inputs from the structure

crates/libafl/src/common/nautilus/grammartec/context.rs

@@ -123,7 +123,7 @@ impl Context {
         *self
             .names_to_nt_id
             .get(nt)
-            .expect(&("no such nonterminal: ".to_owned() + nt))
+            .unwrap_or_else(|| panic!("{}", ("no such nonterminal: ".to_owned() + nt)))
     }
 
     #[must_use]

crates/libafl/src/common/nautilus/mod.rs

@@ -1,4 +1,4 @@
-//! LibAFL version of the [`Nautilus`](https://github.com/nautilus-fuzz/nautilus) grammar fuzzer
+//! `LibAFL` version of the [`Nautilus`](https://github.com/nautilus-fuzz/nautilus) grammar fuzzer
 #![doc = include_str!("README.md")]
 
 #[allow(missing_docs)]

crates/libafl/src/common/nautilus/regex_mutator/mod.rs

@@ -11,7 +11,7 @@ pub struct RegexScript {
 
 impl RegexScript {
     pub fn new<R: Rand>(rand: &mut R) -> Self {
-        let len = if rand.next() % 256 == 0 {
+        let len = if rand.next().is_multiple_of(256) {
             rand.next() % 0xffff
         } else {
             let len = 1 << (rand.next() % 8);

crates/libafl/src/executors/command.rs

@@ -236,6 +236,7 @@ pub struct PTraceCommandConfigurator {
 
 #[cfg(all(feature = "intel_pt", target_os = "linux"))]
 impl CommandConfigurator<Pid> for PTraceCommandConfigurator {
+    #[allow(unreachable_code)]
     fn spawn_child(&mut self, target_bytes: OwnedSlice<'_, u8>) -> Result<Pid, Error> {
         use nix::{
             sys::{

crates/libafl/src/stages/afl_stats.rs

@@ -1,6 +1,10 @@
 //! Stage to compute and report AFL++ stats
 use alloc::{borrow::Cow, string::String, vec::Vec};
-use core::{fmt::Display, marker::PhantomData, time::Duration};
+use core::{
+    fmt::{Debug, Display},
+    marker::PhantomData,
+    time::Duration,
+};
 use std::{
     fs::{File, OpenOptions},
     io::{BufRead, BufReader, Write},
@@ -14,7 +18,7 @@ use libafl_bolts::{
     Named,
     core_affinity::CoreId,
     current_time,
-    tuples::{Handle, Handled, MatchNameRef},
+    tuples::{Handle, Handled, MatchName},
 };
 use serde::{Deserialize, Serialize};
 
@@ -25,6 +29,7 @@ use crate::{
     corpus::{Corpus, HasCurrentCorpusId, SchedulerTestcaseMetadata, Testcase},
     events::{Event, EventFirer, EventWithStats},
     executors::HasObservers,
+    feedbacks::MapFeedbackMetadata,
     monitors::stats::{AggregatorOps, UserStats, UserStatsValue},
     mutators::Tokens,
     observers::MapObserver,
@@ -73,8 +78,9 @@ libafl_bolts::impl_serdeany!(FuzzTime);
 /// The [`AflStatsStage`] is a Stage that calculates and writes
 /// AFL++'s `fuzzer_stats` and `plot_data` information.
 #[derive(Debug, Clone)]
-pub struct AflStatsStage<C, E, EM, I, O, S, Z> {
+pub struct AflStatsStage<C, I, O> {
     map_observer_handle: Handle<C>,
+    map_name: Cow<'static, str>,
     stats_file_path: Option<PathBuf>,
     plot_file_path: Option<PathBuf>,
     start_time: u64,
@@ -112,7 +118,7 @@ pub struct AflStatsStage<C, E, EM, I, O, S, Z> {
     autotokens_enabled: bool,
     /// The core we are bound to
     core_id: CoreId,
-    phantom_data: PhantomData<(E, EM, I, O, S, Z)>,
+    phantom: PhantomData<(I, O)>,
 }
 
 /// AFL++'s `fuzzer_stats`
@@ -190,9 +196,9 @@ pub struct AflFuzzerStats<'a> {
     /// TODO
     cpu_affinity: usize,
     /// how many edges have been found
-    edges_found: u64,
+    edges_found: usize,
     /// Size of our edges map
-    total_edges: u64,
+    total_edges: usize,
     /// how many edges are non-deterministic
     var_byte_count: usize,
     /// TODO:
@@ -224,20 +230,21 @@ pub struct AFLPlotData<'a> {
     pending_total: &'a usize,
     pending_favs: &'a usize,
     /// Note: renamed `map_size` -> `total_edges` for consistency with `fuzzer_stats`
-    total_edges: &'a u64,
+    total_edges: &'a usize,
     saved_crashes: &'a u64,
     saved_hangs: &'a u64,
     max_depth: &'a u64,
     execs_per_sec: &'a u64,
     /// Note: renamed `total_execs` -> `execs_done` for consistency with `fuzzer_stats`
     execs_done: &'a u64,
-    edges_found: &'a u64,
+    edges_found: &'a usize,
 }
 
-impl<C, E, EM, I, O, S, Z> Stage<E, EM, S, Z> for AflStatsStage<C, E, EM, I, O, S, Z>
+impl<C, E, EM, I, O, S, Z> Stage<E, EM, S, Z> for AflStatsStage<C, I, O>
 where
     C: AsRef<O> + Named,
     E: HasObservers,
+    <E as HasObservers>::Observers: MatchName,
     EM: EventFirer<I, S>,
     Z: HasScheduler<I, S>,
     S: HasImported
@@ -248,8 +255,8 @@ where
         + HasNamedMetadata
         + Stoppable
         + HasCurrentCorpusId,
-    E::Observers: MatchNameRef,
     O: MapObserver,
+    for<'de> <O as MapObserver>::Entry: Serialize + Deserialize<'de> + 'static + Debug,
     C: AsRef<O> + Named,
     Z::Scheduler: HasQueueCycles,
 {
@@ -299,13 +306,16 @@ where
         self.maybe_update_cycles(queue_cycles);
         self.maybe_update_cycles_wo_finds(queue_cycles);
 
+        let map_feedback = state
+            .named_metadata_map()
+            .get::<MapFeedbackMetadata<<O as MapObserver>::Entry>>(&self.map_name)
+            .unwrap();
+
+        let filled_entries_in_map = map_feedback.num_covered_map_indexes;
         let observers = executor.observers();
-        let map_observer = observers
-            .get(&self.map_observer_handle)
-            .ok_or_else(|| Error::key_not_found("invariant: MapObserver not found".to_string()))?
-            .as_ref();
-        let filled_entries_in_map = map_observer.count_bytes();
-        let map_size = map_observer.usable_count();
+        let map = observers[&self.map_observer_handle].as_ref();
+        let map_size = map.usable_count();
+
         // Since we do not calibrate when using `QueueScheduler`; we cannot calculate unstable entries.
         let unstable_entries_in_map = state
             .metadata_map()
@@ -370,7 +380,7 @@ where
             #[cfg(not(unix))]
             peak_rss_mb: 0, // TODO for Windows
             cpu_affinity: self.core_id.0,
-            total_edges: map_size as u64,
+            total_edges: map_size,
             edges_found: filled_entries_in_map,
             var_byte_count: unstable_entries_in_map,
             havoc_expansion: 0, // TODO
@@ -435,7 +445,7 @@ where
     }
 }
 
-impl<C, E, EM, I, O, S, Z> Restartable<S> for AflStatsStage<C, E, EM, I, O, S, Z> {
+impl<C, I, O, S> Restartable<S> for AflStatsStage<C, I, O> {
     fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
         Ok(true)
     }
@@ -445,17 +455,13 @@ impl<C, E, EM, I, O, S, Z> Restartable<S> for AflStatsStage<C, E, EM, I, O, S, Z
     }
 }
 
-impl<C, E, EM, I, O, S, Z> AflStatsStage<C, E, EM, I, O, S, Z>
+impl<C, I, O> AflStatsStage<C, I, O>
 where
-    E: HasObservers,
-    EM: EventFirer<I, S>,
-    S: HasImported + HasMetadata + HasExecutions,
-    C: AsRef<O> + Named,
-    O: MapObserver,
+    C: Named,
 {
     /// Builder for `AflStatsStage`
     #[must_use]
-    pub fn builder() -> AflStatsStageBuilder<C, E, EM, I, O, S, Z> {
+    pub fn builder() -> AflStatsStageBuilder<C, I, O> {
         AflStatsStageBuilder::new()
     }
 
@@ -514,7 +520,10 @@ where
     }
 
     #[cfg(feature = "track_hit_feedbacks")]
-    fn maybe_update_last_crash(&mut self, testcase: &Testcase<I>, state: &S) {
+    fn maybe_update_last_crash<S>(&mut self, testcase: &Testcase<I>, state: &S)
+    where
+        S: HasExecutions,
+    {
         #[cfg(feature = "track_hit_feedbacks")]
         if testcase
             .hit_objectives()
@@ -526,7 +535,10 @@ where
     }
 
     #[cfg(feature = "track_hit_feedbacks")]
-    fn maybe_update_last_hang(&mut self, testcase: &Testcase<I>, state: &S) {
+    fn maybe_update_last_hang<S>(&mut self, testcase: &Testcase<I>, state: &S)
+    where
+        S: HasExecutions,
+    {
         if testcase
             .hit_objectives()
             .contains(&Cow::Borrowed(TIMEOUT_FEEDBACK_NAME))
@@ -558,7 +570,7 @@ where
 
     #[expect(clippy::cast_precision_loss)]
     #[expect(clippy::unused_self)]
-    fn calculate_stability(&self, unstable_entries: usize, filled_entries: u64) -> f64 {
+    fn calculate_stability(&self, unstable_entries: usize, filled_entries: usize) -> f64 {
         ((filled_entries as f64 - unstable_entries as f64) / filled_entries as f64) * 100.0
     }
 }
@@ -648,28 +660,25 @@ pub fn get_run_cmdline() -> Cow<'static, str> {
 
 /// The Builder for `AflStatsStage`
 #[derive(Debug)]
-pub struct AflStatsStageBuilder<C, E, EM, I, O, S, Z> {
+pub struct AflStatsStageBuilder<C, I, O> {
     stats_file_path: Option<PathBuf>,
     plot_file_path: Option<PathBuf>,
     core_id: Option<CoreId>,
     map_observer_handle: Option<Handle<C>>,
+    map_name: Option<String>,
     uses_autotokens: bool,
     report_interval: Duration,
     dict_count: usize,
     exec_timeout: u64,
     banner: String,
     version: String,
     target_mode: String,
-    phantom_data: PhantomData<(E, EM, I, O, S, Z)>,
+    phantom_data: PhantomData<(I, O)>,
 }
 
-impl<C, E, EM, I, O, S, Z> AflStatsStageBuilder<C, E, EM, I, O, S, Z>
+impl<C, I, O> AflStatsStageBuilder<C, I, O>
 where
-    C: AsRef<O> + Named,
-    E: HasObservers,
-    EM: EventFirer<I, S>,
-    O: MapObserver,
-    S: HasImported + HasMetadata + HasExecutions,
+    C: Named,
 {
     fn new() -> Self {
         Self {
@@ -678,6 +687,7 @@ where
             plot_file_path: None,
             core_id: None,
             map_observer_handle: None,
+            map_name: None,
             uses_autotokens: false,
             dict_count: 0,
             exec_timeout: 0,
@@ -718,6 +728,17 @@ where
         self.map_observer_handle = Some(map_observer.handle());
         self
     }
+
+    /// map name to check the filled count
+    #[must_use]
+    pub fn map_name<F>(mut self, map_feedback: &F) -> Self
+    where
+        F: Named,
+    {
+        self.map_name = Some(map_feedback.name().to_string());
+        self
+    }
+
     /// If we use autotokens provided by the target
     #[must_use]
     pub fn uses_autotokens(mut self, uses: bool) -> Self {
@@ -782,10 +803,13 @@ where
     /// No `MapObserver` supplied to the builder
     /// No `stats_file_path` provieded
     #[allow(clippy::type_complexity)]
-    pub fn build(self) -> Result<AflStatsStage<C, E, EM, I, O, S, Z>, Error> {
+    pub fn build(self) -> Result<AflStatsStage<C, I, O>, Error> {
         if self.map_observer_handle.is_none() {
             return Err(Error::illegal_argument("Must set `map_observer`"));
         }
+        let Some(map_name) = self.map_name else {
+            return Err(Error::illegal_argument("Must set `map_name`"));
+        };
         if let Some(ref plot_file) = self.plot_file_path {
             Self::create_plot_data_file(plot_file)?;
         }
@@ -794,6 +818,7 @@ where
         }
         Ok(AflStatsStage {
             stats_file_path: self.stats_file_path,
+            map_name: Cow::Owned(map_name),
             plot_file_path: self.plot_file_path,
             map_observer_handle: self.map_observer_handle.unwrap(),
             start_time: current_time().as_secs(),
@@ -820,7 +845,7 @@ where
             dict_count: self.dict_count,
             core_id: self.core_id.unwrap_or(CoreId(0)),
             autotokens_enabled: self.uses_autotokens,
-            phantom_data: PhantomData,
+            phantom: PhantomData,
         })
     }
 }

crates/libafl_bolts/src/rands/mod.rs

@@ -567,9 +567,9 @@ pub mod pybind {
     #[pyclass(unsendable, name = "StdRand")]
     #[expect(clippy::unsafe_derive_deserialize)]
     #[derive(Serialize, Deserialize, Debug, Clone)]
-    /// Python class for StdRand
+    /// Python class for `StdRand`
     pub struct PythonStdRand {
-        /// Rust wrapped StdRand object
+        /// Rust wrapped `StdRand` object
         pub inner: StdRand,
     }
 

fuzzers/baby/baby_fuzzer_custom_executor/src/main.rs

@@ -88,6 +88,7 @@ pub fn main() {
     let calibration_stage = CalibrationStage::new(&feedback);
     let stats_stage = AflStatsStage::builder()
         .map_observer(&observer)
+        .map_name(&feedback)
         .build()
         .unwrap();
 

fuzzers/binary_only/qemu_launcher/src/instance.rs

@@ -188,12 +188,14 @@ where
             |_, _, _, _| Ok(self.options.tui),
             tuple_list!(AflStatsStage::builder()
                 .map_observer(&edges_observer)
+                .map_name(&map_feedback)
                 .build()?),
         );
         let stats_stage_cmplog = IfStage::new(
             |_, _, _, _| Ok(self.options.tui),
             tuple_list!(AflStatsStage::builder()
                 .map_observer(&edges_observer)
+                .map_name(&map_feedback)
                 .build()?),
         );
 

fuzzers/forkserver/libafl-fuzz/src/fuzzer.rs

@@ -186,6 +186,7 @@ define_run_client!(state, mgr, fuzzer_dir, core_id, opt, is_main_node, {
         .core_id(core_id)
         .report_interval(Duration::from_secs(opt.stats_interval))
         .map_observer(&edges_observer)
+        .map_name(&map_feedback)
         .uses_autotokens(!opt.no_autodict)
         .tokens(&tokens)
         .banner(opt.executable.display().to_string())