Skip to content

Commit

Permalink
Add OnDiskJsonAggregateMonitor (AFLplusplus#2845)
Browse files Browse the repository at this point in the history
* Add OnDiskJsonAggregateMonitor

* Fix formatting

* Remove unnecessary trait bounds

* Remove unnecessary import restriction

* Remove unnecessary imports

* fix formatting
  • Loading branch information
riesentoaster authored Jan 15, 2025
1 parent dc39aad commit ba09cb0
Show file tree
Hide file tree
Showing 2 changed files with 145 additions and 8 deletions.
125 changes: 125 additions & 0 deletions libafl/src/monitors/disk_aggregate.rs
Original file line number Diff line number Diff line change
@@ -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<M> {
base: M,
aggregator: Aggregator,
json_path: PathBuf,
last_update: Duration,
update_interval: Duration,
}

impl<M> Debug for OnDiskJsonAggregateMonitor<M>
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<M> Monitor for OnDiskJsonAggregateMonitor<M>
where
M: Monitor,
{
fn client_stats_mut(&mut self) -> &mut Vec<ClientStats> {
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<M> OnDiskJsonAggregateMonitor<M> {
/// Creates a new [`OnDiskJsonAggregateMonitor`]
pub fn new<P>(base: M, json_path: P) -> Self
where
P: Into<PathBuf>,
{
Self::with_interval(json_path, base, Duration::from_secs(10))
}

/// Creates a new [`OnDiskJsonAggregateMonitor`] with custom update interval
pub fn with_interval<P>(json_path: P, base: M, update_interval: Duration) -> Self
where
P: Into<PathBuf>,
{
Self {
base,
aggregator: Aggregator::new(),
json_path: json_path.into(),
last_update: current_time() - update_interval,
update_interval,
}
}
}
28 changes: 20 additions & 8 deletions libafl/src/monitors/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand Down

0 comments on commit ba09cb0

Please sign in to comment.