Make Stats stage send stats again (#2830)

* Make Stats stage send stats again

* re-remove stats mod

* clp, fmt

* clip

fuzzers/baby/baby_fuzzer_custom_executor/src/main.rs

@@ -18,7 +18,7 @@ use libafl::{
     mutators::{havoc_mutations::havoc_mutations, scheduled::StdScheduledMutator},
     observers::StdMapObserver,
     schedulers::QueueScheduler,
-    stages::mutational::StdMutationalStage,
+    stages::{mutational::StdMutationalStage, AflStatsStage, CalibrationStage},
     state::{HasCorpus, HasExecutions, StdState, UsesState},
 };
 use libafl_bolts::{current_nanos, nonzero, rands::StdRand, tuples::tuple_list, AsSlice};
@@ -86,6 +86,12 @@ pub fn main() {
     // Feedback to rate the interestingness of an input
     let mut feedback = MaxMapFeedback::new(&observer);
 
+    let calibration_stage = CalibrationStage::new(&feedback);
+    let stats_stage = AflStatsStage::builder()
+        .map_observer(&observer)
+        .build()
+        .unwrap();
+
     // A feedback to choose if an input is a solution or not
     let mut objective = feedback_and_fast!(
         // Look for crashes.
@@ -151,7 +157,11 @@ pub fn main() {
 
     // Setup a mutational stage with a basic bytes mutator
     let mutator = StdScheduledMutator::new(havoc_mutations());
-    let mut stages = tuple_list!(StdMutationalStage::new(mutator));
+    let mut stages = tuple_list!(
+        calibration_stage,
+        StdMutationalStage::new(mutator),
+        stats_stage
+    );
 
     fuzzer
         .fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)

libafl/Cargo.toml

@@ -39,6 +39,7 @@ default = [
   "regex",
   "serdeany_autoreg",
   "libafl_bolts/xxh3",
+  "tui_monitor",
 ]
 document-features = ["dep:document-features"]
 
@@ -200,6 +201,11 @@ nautilus = [
   "regex",
 ]
 
+[[example]]
+name = "tui_mock"
+path = "./examples/tui_mock/main.rs"
+required-features = ["std", "tui_monitor"]
+
 [build-dependencies]
 rustversion = "1.0.17"
 

libafl/examples/tui_mock/main.rs

@@ -0,0 +1,20 @@
+//! An example for TUI that uses the TUI without any real data.
+//! This is mainly to fix the UI without having to run a real fuzzer.
+
+use std::{thread::sleep, time::Duration};
+
+use libafl::monitors::{tui::TuiMonitor, ClientStats, Monitor};
+use libafl_bolts::ClientId;
+
+pub fn main() {
+    let mut monitor = TuiMonitor::builder().build();
+
+    let client_stats = ClientStats {
+        corpus_size: 1024,
+        executions: 512,
+        ..ClientStats::default()
+    };
+
+    monitor.display("Test", ClientId(0));
+    sleep(Duration::from_secs(10));
+}

libafl/src/monitors/mod.rs

@@ -313,15 +313,15 @@ impl fmt::Display for UserStatsValue {
 /// Prettifies float values for human-readable output
 fn prettify_float(value: f64) -> String {
     let (value, suffix) = match value {
-        value if value >= 1000000.0 => (value / 1000000.0, "M"),
+        value if value >= 1_000_000.0 => (value / 1_000_000.0, "M"),
-        value if value >= 1000.0 => (value / 1000.0, "k"),
+        value if value >= 1_000.0 => (value / 1_000.0, "k"),
         value => (value, ""),
     };
     match value {
-        value if value >= 1000000.0 => {
+        value if value >= 1_000_000.0 => {
             format!("{value:.2}{suffix}")
         }
-        value if value >= 1000.0 => {
+        value if value >= 1_000.0 => {
             format!("{value:.1}{suffix}")
         }
         value if value >= 100.0 => {

libafl/src/stages/afl_stats.rs

@@ -24,9 +24,10 @@ use serde::{Deserialize, Serialize};
 use crate::feedbacks::{CRASH_FEEDBACK_NAME, TIMEOUT_FEEDBACK_NAME};
 use crate::{
     corpus::{Corpus, HasCurrentCorpusId, SchedulerTestcaseMetadata, Testcase},
-    events::EventFirer,
+    events::{Event, EventFirer},
     executors::HasObservers,
     inputs::UsesInput,
+    monitors::{AggregatorOps, UserStats, UserStatsValue},
     mutators::Tokens,
     observers::MapObserver,
     schedulers::{minimizer::IsFavoredMetadata, HasQueueCycles},
@@ -35,6 +36,7 @@ use crate::{
     std::string::ToString,
     Error, HasMetadata, HasNamedMetadata, HasScheduler,
 };
+
 /// AFL++'s default stats update interval
 pub const AFL_FUZZER_STATS_UPDATE_INTERVAL_SECS: u64 = 60;
 
@@ -76,7 +78,7 @@ libafl_bolts::impl_serdeany!(FuzzTime);
 #[derive(Debug, Clone)]
 pub struct AflStatsStage<C, E, EM, O, S, Z> {
     map_observer_handle: Handle<C>,
-    stats_file_path: PathBuf,
+    stats_file_path: Option<PathBuf>,
     plot_file_path: Option<PathBuf>,
     start_time: u64,
     // the number of testcases that have been fuzzed
@@ -238,7 +240,7 @@ pub struct AFLPlotData<'a> {
 impl<C, E, EM, O, S, Z> Stage<E, EM, S, Z> for AflStatsStage<C, E, EM, O, S, Z>
 where
     E: HasObservers,
-    EM: EventFirer,
+    EM: EventFirer<State = S>,
     Z: HasScheduler<<S::Corpus as Corpus>::Input, S>,
     S: HasImported
         + HasCorpus
@@ -260,7 +262,7 @@ where
         fuzzer: &mut Z,
         executor: &mut E,
         state: &mut S,
-        _manager: &mut EM,
+        manager: &mut EM,
     ) -> Result<(), Error> {
         let Some(corpus_idx) = state.current_corpus_id()? else {
             return Err(Error::illegal_state(
@@ -365,6 +367,7 @@ where
             execs_since_crash: total_executions - self.execs_at_last_objective,
             exec_timeout: self.exec_timeout,
             slowest_exec_ms: self.slowest_exec.as_millis(),
+            // TODO: getting rss_mb may take some extra millis, so might make sense to make this optional
             #[cfg(unix)]
             peak_rss_mb: peak_rss_mb_child_processes()?,
             #[cfg(not(unix))]
@@ -398,10 +401,36 @@ where
             saved_crashes: &stats.saved_crashes,
             execs_done: &stats.execs_done,
         };
-        self.write_fuzzer_stats(&stats)?;
+        self.maybe_write_fuzzer_stats(&stats)?;
         if self.plot_file_path.is_some() {
             self.write_plot_data(&plot_data)?;
         }
+
+        drop(testcase);
+
+        // We construct this simple json by hand to squeeze out some extra speed.
+        let json = format!(
+            "{{\
+                \"pending\":{},\
+                \"pending_fav\":{},\
+                \"own_finds:\"{},\
+                \"imported\":{}\
+            }}",
+            stats.pending_total, stats.pending_favs, stats.corpus_found, stats.corpus_imported
+        );
+
+        manager.fire(
+            state,
+            Event::UpdateUserStats {
+                name: Cow::Borrowed("AflStats"),
+                value: UserStats::new(
+                    UserStatsValue::String(Cow::Owned(json)),
+                    AggregatorOps::None,
+                ),
+                phantom: PhantomData,
+            },
+        )?;
+
         Ok(())
     }
 
@@ -428,15 +457,17 @@ where
         AflStatsStageBuilder::new()
     }
 
-    fn write_fuzzer_stats(&self, stats: &AFLFuzzerStats) -> Result<(), Error> {
+    /// Writes a stats file, if a `stats_file_path` is set.
-        let tmp_file = self
+    fn maybe_write_fuzzer_stats(&self, stats: &AFLFuzzerStats) -> Result<(), Error> {
-            .stats_file_path
+        if let Some(stats_file_path) = &self.stats_file_path {
+            let tmp_file = stats_file_path
                 .parent()
                 .expect("fuzzer_stats file must have a parent!")
                 .join(".fuzzer_stats_tmp");
             std::fs::write(&tmp_file, stats.to_string())?;
-        _ = std::fs::copy(&tmp_file, &self.stats_file_path)?;
+            _ = std::fs::copy(&tmp_file, stats_file_path)?;
             std::fs::remove_file(tmp_file)?;
+        }
         Ok(())
     }
 
@@ -557,8 +588,8 @@ impl Display for AFLPlotData<'_> {
     }
 }
 impl AFLPlotData<'_> {
-    fn get_header() -> String {
+    fn header() -> &'static str {
-        "# relative_time, cycles_done, cur_item, corpus_count, pending_total, pending_favs, total_edges, saved_crashes, saved_hangs, max_depth, execs_per_sec, execs_done, edges_found".to_string()
+        "# relative_time, cycles_done, cur_item, corpus_count, pending_total, pending_favs, total_edges, saved_crashes, saved_hangs, max_depth, execs_per_sec, execs_done, edges_found"
     }
 }
 impl Display for AFLFuzzerStats<'_> {
@@ -736,10 +767,10 @@ where
             // check if it contains any data
             let file = File::open(path)?;
             if BufReader::new(file).lines().next().is_none() {
-                std::fs::write(path, AFLPlotData::get_header())?;
+                std::fs::write(path, AFLPlotData::header())?;
             }
         } else {
-            std::fs::write(path, AFLPlotData::get_header())?;
+            std::fs::write(path, AFLPlotData::header())?;
         }
         Ok(())
     }
@@ -757,19 +788,17 @@ where
     /// No `MapObserver` supplied to the builder
     /// No `stats_file_path` provieded
     pub fn build(self) -> Result<AflStatsStage<C, E, EM, O, S, Z>, Error> {
-        if self.stats_file_path.is_none() {
-            return Err(Error::illegal_argument("Must set `stats_file_path`"));
-        }
-        let stats_file_path = self.stats_file_path.unwrap();
         if self.map_observer_handle.is_none() {
             return Err(Error::illegal_argument("Must set `map_observer`"));
         }
         if let Some(ref plot_file) = self.plot_file_path {
             Self::create_plot_data_file(plot_file)?;
         }
-        Self::create_fuzzer_stats_file(&stats_file_path)?;
+        if let Some(stats_file_path) = &self.stats_file_path {
+            Self::create_fuzzer_stats_file(stats_file_path)?;
+        }
         Ok(AflStatsStage {
-            stats_file_path,
+            stats_file_path: self.stats_file_path,
             plot_file_path: self.plot_file_path,
             map_observer_handle: self.map_observer_handle.unwrap(),
             start_time: current_time().as_secs(),

libafl/src/stages/calibrate.rs

@@ -27,6 +27,14 @@ use crate::{
     Error, HasMetadata, HasNamedMetadata,
 };
 
+/// AFL++'s `CAL_CYCLES_FAST` + 1
+const CAL_STAGE_START: usize = 4;
+/// AFL++'s `CAL_CYCLES` + 1
+const CAL_STAGE_MAX: usize = 8;
+
+/// Default name for `CalibrationStage`; derived from AFL++
+pub const CALIBRATION_STAGE_NAME: &str = "calibration";
+
 /// The metadata to keep unstable entries
 /// Formula is same as AFL++: number of unstable entries divided by the number of filled entries.
 #[cfg_attr(
@@ -69,8 +77,6 @@ impl Default for UnstableEntriesMetadata {
     }
 }
 
-/// Default name for `CalibrationStage`; derived from AFL++
-pub const CALIBRATION_STAGE_NAME: &str = "calibration";
 /// The calibration stage will measure the average exec time and the target's stability for this input.
 #[derive(Clone, Debug)]
 pub struct CalibrationStage<C, E, O, OT, S> {
@@ -83,9 +89,6 @@ pub struct CalibrationStage<C, E, O, OT, S> {
     phantom: PhantomData<(E, O, OT, S)>,
 }
 
-const CAL_STAGE_START: usize = 4; // AFL++'s CAL_CYCLES_FAST + 1
-const CAL_STAGE_MAX: usize = 8; // AFL++'s CAL_CYCLES + 1
-
 impl<C, E, EM, O, OT, S, Z> Stage<E, EM, S, Z> for CalibrationStage<C, E, O, OT, S>
 where
     E: Executor<EM, <S::Corpus as Corpus>::Input, S, Z> + HasObservers<Observers = OT>,