From 2b092f40fa7e9f045006b23174bf653e9b4c1153 Mon Sep 17 00:00:00 2001
From: Patrick Gersch <gersch.patrick@gmail.com>
Date: Mon, 26 Dec 2022 11:20:30 +0100
Subject: [PATCH] SimpleMonitor optionally displays user_monitor stats (#970)

* Adding with_user_monitor() to SimpleMonitor

* Satisfy clippy

* Satisfy fmt and pylibafl

* Fix leading whitespace
---
 fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs | 17 ++++++++-------
 libafl/src/monitors/mod.rs                | 25 ++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs b/fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs
index 74ba09d992..29c2d08e44 100644
--- a/fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs
+++ b/fuzzers/fuzzbench_fork_qemu/src/fuzzer.rs
@@ -195,13 +195,16 @@ fn fuzz(
     let file_null = File::open("/dev/null")?;
 
     // 'While the stats are state, they are usually used in the broker - which is likely never restarted
-    let monitor = SimpleMonitor::new(|s| {
-        #[cfg(unix)]
-        writeln!(&mut stdout_cpy, "{}", s).unwrap();
-        #[cfg(windows)]
-        println!("{s}");
-        writeln!(log.borrow_mut(), "{:?} {}", current_time(), s).unwrap();
-    });
+    let monitor = SimpleMonitor::with_user_monitor(
+        |s| {
+            #[cfg(unix)]
+            writeln!(&mut stdout_cpy, "{s}").unwrap();
+            #[cfg(windows)]
+            println!("{s}");
+            writeln!(log.borrow_mut(), "{:?} {s}", current_time()).unwrap();
+        },
+        true,
+    );
 
     let mut shmem_provider = StdShMemProvider::new()?;
 
diff --git a/libafl/src/monitors/mod.rs b/libafl/src/monitors/mod.rs
index d88bad7004..4f62d84ca5 100644
--- a/libafl/src/monitors/mod.rs
+++ b/libafl/src/monitors/mod.rs
@@ -16,6 +16,7 @@ pub use prometheus::PrometheusMonitor;
 #[cfg(feature = "std")]
 pub mod disk;
 use alloc::{fmt::Debug, string::String, vec::Vec};
+use core::fmt::Write;
 use core::{fmt, time::Duration};
 
 #[cfg(feature = "std")]
@@ -386,6 +387,7 @@ where
 {
     print_fn: F,
     start_time: Duration,
+    print_user_monitor: bool,
     client_stats: Vec<ClientStats>,
 }
 
@@ -421,7 +423,7 @@ where
     }
 
     fn display(&mut self, event_msg: String, sender_id: u32) {
-        let fmt = format!(
+        let mut fmt = format!(
             "[{} #{}] run time: {}, clients: {}, corpus: {}, objectives: {}, executions: {}, exec/sec: {}",
             event_msg,
             sender_id,
@@ -432,6 +434,14 @@ where
             self.total_execs(),
             self.execs_per_sec_pretty()
         );
+
+        if self.print_user_monitor {
+            let client = self.client_stats_mut_for(sender_id);
+            for (key, val) in &client.user_monitor {
+                write!(fmt, ", {key}: {val}").unwrap();
+            }
+        }
+
         (self.print_fn)(fmt);
 
         // Only print perf monitor if the feature is enabled
@@ -459,6 +469,7 @@ where
         Self {
             print_fn,
             start_time: current_time(),
+            print_user_monitor: false,
             client_stats: vec![],
         }
     }
@@ -468,6 +479,17 @@ where
         Self {
             print_fn,
             start_time,
+            print_user_monitor: false,
+            client_stats: vec![],
+        }
+    }
+
+    /// Creates the monitor that also prints the user monitor
+    pub fn with_user_monitor(print_fn: F, print_user_monitor: bool) -> Self {
+        Self {
+            print_fn,
+            start_time: current_time(),
+            print_user_monitor,
             client_stats: vec![],
         }
     }
@@ -971,6 +993,7 @@ pub mod pybind {
                 inner: SimpleMonitor {
                     print_fn: Box::new(closure),
                     start_time: self.inner.start_time,
+                    print_user_monitor: false,
                     client_stats: self.inner.client_stats.clone(),
                 },
                 print_fn: self.print_fn.clone(),