diff --git a/libafl/src/monitors/disk_aggregate.rs b/libafl/src/monitors/disk_aggregate.rs new file mode 100644 index 0000000000..5f236930a1 --- /dev/null +++ b/libafl/src/monitors/disk_aggregate.rs @@ -0,0 +1,125 @@ +//! Monitors that log aggregated stats to disk. + +use alloc::vec::Vec; +use core::{ + fmt::{Debug, Formatter}, + time::Duration, +}; +use std::{fs::OpenOptions, io::Write, path::PathBuf}; + +use libafl_bolts::{current_time, ClientId}; +use serde_json::json; + +use crate::monitors::{Aggregator, ClientStats, Monitor}; + +/// A monitor that wraps another monitor and logs aggregated stats to a JSON file. +#[derive(Clone)] +pub struct OnDiskJsonAggregateMonitor { + base: M, + aggregator: Aggregator, + json_path: PathBuf, + last_update: Duration, + update_interval: Duration, +} + +impl Debug for OnDiskJsonAggregateMonitor +where + M: Debug, +{ + fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { + f.debug_struct("OnDiskJsonAggregateMonitor") + .field("base", &self.base) + .field("last_update", &self.last_update) + .field("update_interval", &self.update_interval) + .field("json_path", &self.json_path) + .finish_non_exhaustive() + } +} + +impl Monitor for OnDiskJsonAggregateMonitor +where + M: Monitor, +{ + fn client_stats_mut(&mut self) -> &mut Vec { + self.base.client_stats_mut() + } + + fn client_stats(&self) -> &[ClientStats] { + self.base.client_stats() + } + + fn set_start_time(&mut self, time: Duration) { + self.base.set_start_time(time); + } + + fn start_time(&self) -> Duration { + self.base.start_time() + } + + fn aggregate(&mut self, name: &str) { + self.aggregator.aggregate(name, self.base.client_stats()); + self.base.aggregate(name); + } + + fn display(&mut self, event_msg: &str, sender_id: ClientId) { + // First let the base monitor handle its display + self.base.display(event_msg, sender_id); + + // Write JSON stats if update interval has elapsed + let cur_time = current_time(); + if cur_time - self.last_update >= self.update_interval { + self.last_update = cur_time; + + let file = OpenOptions::new() + .append(true) + .create(true) + .open(&self.json_path) + .expect("Failed to open JSON logging file"); + + let mut json_value = json!({ + "run_time": (cur_time - self.start_time()).as_secs(), + "clients": self.client_stats_count(), + "corpus": self.corpus_size(), + "objectives": self.objective_size(), + "executions": self.total_execs(), + "exec_sec": self.execs_per_sec(), + }); + + // Add all aggregated values directly to the root + if let Some(obj) = json_value.as_object_mut() { + obj.extend( + self.aggregator + .aggregated + .iter() + .map(|(k, v)| (k.clone(), json!(v))), + ); + } + + writeln!(&file, "{json_value}").expect("Unable to write JSON to file"); + } + } +} + +impl OnDiskJsonAggregateMonitor { + /// Creates a new [`OnDiskJsonAggregateMonitor`] + pub fn new

(base: M, json_path: P) -> Self + where + P: Into, + { + Self::with_interval(json_path, base, Duration::from_secs(10)) + } + + /// Creates a new [`OnDiskJsonAggregateMonitor`] with custom update interval + pub fn with_interval

(json_path: P, base: M, update_interval: Duration) -> Self + where + P: Into, + { + Self { + base, + aggregator: Aggregator::new(), + json_path: json_path.into(), + last_update: current_time() - update_interval, + update_interval, + } + } +} diff --git a/libafl/src/monitors/mod.rs b/libafl/src/monitors/mod.rs index 1f94a7f34a..bfc55084bc 100644 --- a/libafl/src/monitors/mod.rs +++ b/libafl/src/monitors/mod.rs @@ -3,24 +3,36 @@ pub mod multi; pub use multi::MultiMonitor; +#[cfg(feature = "std")] +pub mod disk; +#[cfg(feature = "std")] +pub use disk::{OnDiskJsonMonitor, OnDiskTomlMonitor}; + +#[cfg(feature = "std")] +pub mod disk_aggregate; +#[cfg(feature = "std")] +pub use disk_aggregate::OnDiskJsonAggregateMonitor; + #[cfg(all(feature = "tui_monitor", feature = "std"))] pub mod tui; +#[cfg(all(feature = "tui_monitor", feature = "std"))] +pub use tui::TuiMonitor; #[cfg(all(feature = "prometheus_monitor", feature = "std"))] pub mod prometheus; -use alloc::string::ToString; -#[cfg(all(feature = "prometheus_monitor", feature = "std"))] -pub use prometheus::PrometheusMonitor; -#[cfg(feature = "std")] -pub mod disk; -use alloc::{borrow::Cow, fmt::Debug, string::String, vec::Vec}; +use alloc::{ + borrow::Cow, + fmt::Debug, + string::{String, ToString}, + vec::Vec, +}; use core::{fmt, fmt::Write, time::Duration}; -#[cfg(feature = "std")] -pub use disk::{OnDiskJsonMonitor, OnDiskTomlMonitor}; use hashbrown::HashMap; use libafl_bolts::{current_time, format_duration_hms, ClientId}; +#[cfg(all(feature = "prometheus_monitor", feature = "std"))] +pub use prometheus::PrometheusMonitor; use serde::{Deserialize, Serialize}; #[cfg(feature = "afl_exec_sec")]