Fix SimpleRestartingMonitor after restart (#1669)

* Fix SimpleRestartingMonitor after restart

* a

* a

* a

* a

* a

* a

* a

* a

* a

* ci

---------

Co-authored-by: toka <tokazerkje@outlook.com>

libafl/src/events/simple.rs

@@ -9,9 +9,12 @@ use alloc::{
 use core::ffi::c_void;
 #[cfg(all(unix, feature = "std"))]
 use core::ptr::write_volatile;
-#[cfg(feature = "std")]
-use core::sync::atomic::{compiler_fence, Ordering};
 use core::{fmt::Debug, marker::PhantomData};
+#[cfg(feature = "std")]
+use core::{
+    sync::atomic::{compiler_fence, Ordering},
+    time::Duration,
+};
 
 #[cfg(all(feature = "std", any(windows, not(feature = "fork"))))]
 use libafl_bolts::os::startable_self;
@@ -28,12 +31,6 @@ use serde::{de::DeserializeOwned, Serialize};
 use super::{CustomBufEventResult, CustomBufHandlerFn, HasCustomBufHandlers, ProgressReporter};
 #[cfg(all(unix, feature = "std"))]
 use crate::events::{shutdown_handler, SHUTDOWN_SIGHANDLER_DATA};
-#[cfg(feature = "std")]
-use crate::{
-    corpus::Corpus,
-    monitors::SimplePrintingMonitor,
-    state::{HasCorpus, HasSolutions},
-};
 use crate::{
     events::{
         BrokerEventResult, Event, EventFirer, EventManager, EventManagerId, EventProcessor,
@@ -44,6 +41,11 @@ use crate::{
     state::{HasClientPerfMonitor, HasExecutions, HasLastReportTime, HasMetadata, UsesState},
     Error,
 };
+#[cfg(feature = "std")]
+use crate::{
+    monitors::{ClientStats, SimplePrintingMonitor},
+    state::{HasCorpus, HasSolutions},
+};
 
 /// The llmp connection from the actual fuzzer to the process supervising it
 const _ENV_FUZZER_SENDER: &str = "_AFL_ENV_FUZZER_SENDER";
@@ -343,6 +345,7 @@ where
 #[cfg(feature = "std")]
 impl<MT, S, SP> EventRestarter for SimpleRestartingEventManager<MT, S, SP>
 where
+    MT: Monitor,
     S: UsesInput + Serialize,
     SP: ShMemProvider,
 {
@@ -350,7 +353,11 @@ where
     fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         // First, reset the page to 0 so the next iteration can read read from the beginning of this page
         self.staterestorer.reset();
-        self.staterestorer.save(state)
+        self.staterestorer.save(&(
+            state,
+            self.simple_event_mgr.monitor.start_time(),
+            self.simple_event_mgr.monitor.client_stats(),
+        ))
     }
 
     fn send_exiting(&mut self) -> Result<(), Error> {
@@ -539,7 +546,7 @@ where
         };
 
         // If we're restarting, deserialize the old state.
-        let (state, mgr) = match staterestorer.restore::<S>()? {
+        let (state, mgr) = match staterestorer.restore::<(S, Duration, Vec<ClientStats>)>()? {
             None => {
                 log::info!("First run. Let's set it all up");
                 // Mgr to send and receive msgs from/to all other fuzzer instances
@@ -549,15 +556,14 @@ where
                 )
             }
             // Restoring from a previous run, deserialize state and corpus.
-            Some(state) => {
+            Some((state, start_time, clients_stats)) => {
                 log::info!("Subsequent run. Loaded previous state.");
                 // We reset the staterestorer, the next staterestorer and receiver (after crash) will reuse the page from the initial message.
                 staterestorer.reset();
 
-                // load the corpus size into monitor to still display the correct numbers after restart.
-                let client_stats = monitor.client_stats_mut_for(ClientId(0));
-                client_stats.update_corpus_size(state.corpus().count().try_into()?);
-                client_stats.update_objective_size(state.solutions().count().try_into()?);
+                // reload the state of the monitor to display the correct stats after restarts
+                monitor.set_start_time(start_time);
+                *monitor.client_stats_mut() = clients_stats;
 
                 (
                     Some(state),

libafl/src/monitors/disk.rs

@@ -39,10 +39,15 @@ where
     }
 
     /// Time this fuzzing run stated
-    fn start_time(&mut self) -> Duration {
+    fn start_time(&self) -> Duration {
         self.base.start_time()
     }
 
+    /// Set creation time
+    fn set_start_time(&mut self, time: Duration) {
+        self.base.set_start_time(time);
+    }
+
     fn display(&mut self, event_msg: String, sender_id: ClientId) {
         let cur_time = current_time();
 
@@ -184,10 +189,14 @@ where
         self.base.client_stats()
     }
 
-    fn start_time(&mut self) -> Duration {
+    fn start_time(&self) -> Duration {
         self.base.start_time()
     }
 
+    fn set_start_time(&mut self, time: Duration) {
+        self.base.set_start_time(time);
+    }
+
     fn display(&mut self, event_msg: String, sender_id: ClientId) {
         if (self.log_record)(&mut self.base) {
             let file = OpenOptions::new()

libafl/src/monitors/mod.rs

@@ -82,7 +82,7 @@ fn prettify_float(value: f64) -> String {
 }
 
 /// A simple struct to keep track of client monitor
-#[derive(Debug, Clone, Default, Serialize)]
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
 pub struct ClientStats {
     // monitor (maybe we need a separated struct?)
     /// The corpus size for this client
@@ -235,7 +235,10 @@ pub trait Monitor {
     fn client_stats(&self) -> &[ClientStats];
 
     /// Creation time
-    fn start_time(&mut self) -> Duration;
+    fn start_time(&self) -> Duration;
+
+    /// Set creation time
+    fn set_start_time(&mut self, time: Duration);
 
     /// Show the monitor to the user
     fn display(&mut self, event_msg: String, sender_id: ClientId);
@@ -311,10 +314,15 @@ impl Monitor for NopMonitor {
     }
 
     /// Time this fuzzing run stated
-    fn start_time(&mut self) -> Duration {
+    fn start_time(&self) -> Duration {
         self.start_time
     }
 
+    /// Time this fuzzing run stated
+    fn set_start_time(&mut self, time: Duration) {
+        self.start_time = time;
+    }
+
     fn display(&mut self, _event_msg: String, _sender_id: ClientId) {}
 }
 
@@ -375,10 +383,15 @@ impl Monitor for SimplePrintingMonitor {
     }
 
     /// Time this fuzzing run stated
-    fn start_time(&mut self) -> Duration {
+    fn start_time(&self) -> Duration {
         self.start_time
     }
 
+    /// Time this fuzzing run stated
+    fn set_start_time(&mut self, time: Duration) {
+        self.start_time = time;
+    }
+
     fn display(&mut self, event_msg: String, sender_id: ClientId) {
         let mut userstats = self.client_stats()[sender_id.0 as usize]
             .user_monitor
@@ -452,10 +465,15 @@ where
     }
 
     /// Time this fuzzing run stated
-    fn start_time(&mut self) -> Duration {
+    fn start_time(&self) -> Duration {
         self.start_time
     }
 
+    /// Set creation time
+    fn set_start_time(&mut self, time: Duration) {
+        self.start_time = time;
+    }
+
     fn display(&mut self, event_msg: String, sender_id: ClientId) {
         let mut fmt = format!(
             "[{} #{}] run time: {}, clients: {}, corpus: {}, objectives: {}, executions: {}, exec/sec: {}",
@@ -1113,8 +1131,13 @@ pub mod pybind {
         }
 
         /// Time this fuzzing run stated
-        fn start_time(&mut self) -> Duration {
-            unwrap_me_mut!(self.wrapper, m, { m.start_time() })
+        fn start_time(&self) -> Duration {
+            unwrap_me!(self.wrapper, m, { m.start_time() })
+        }
+
+        /// set start time
+        fn set_start_time(&mut self, time: Duration) {
+            unwrap_me_mut!(self.wrapper, m, { m.set_start_time(time) });
         }
 
         fn display(&mut self, event_msg: String, sender_id: ClientId) {

libafl/src/monitors/multi.rs

@@ -49,8 +49,13 @@ where
         &self.client_stats
     }
 
+    /// Set creation time
+    fn set_start_time(&mut self, time: Duration) {
+        self.start_time = time;
+    }
+
     /// Time this fuzzing run stated
-    fn start_time(&mut self) -> Duration {
+    fn start_time(&self) -> Duration {
         self.start_time
     }
 

libafl/src/monitors/prometheus.rs

@@ -88,10 +88,15 @@ where
     }
 
     /// Time this fuzzing run stated
-    fn start_time(&mut self) -> Duration {
+    fn start_time(&self) -> Duration {
         self.start_time
     }
 
+    /// Set creation time
+    fn set_start_time(&mut self, time: Duration) {
+        self.start_time = time;
+    }
+
     #[allow(clippy::cast_sign_loss)]
     fn display(&mut self, event_msg: String, sender_id: ClientId) {
         // Update the prometheus metrics

libafl/src/monitors/tui/mod.rs

@@ -346,10 +346,15 @@ impl Monitor for TuiMonitor {
     }
 
     /// Time this fuzzing run stated
-    fn start_time(&mut self) -> Duration {
+    fn start_time(&self) -> Duration {
         self.start_time
     }
 
+    /// Set creation time
+    fn set_start_time(&mut self, time: Duration) {
+        self.start_time = time;
+    }
+
     #[allow(clippy::cast_sign_loss)]
     fn display(&mut self, event_msg: String, sender_id: ClientId) {
         let cur_time = current_time();

libafl_targets/src/libfuzzer/observers/mod.rs

@@ -1,4 +1,5 @@
+/// oom observer
 #[cfg(feature = "libfuzzer_oom")]
-mod oom;
+pub mod oom;
 #[cfg(feature = "libfuzzer_oom")]
 pub use oom::*;