From 1c85c3af13b6a8e2128bab91b18722dfe9767a36 Mon Sep 17 00:00:00 2001 From: Dominik Maier Date: Fri, 5 Apr 2024 14:23:56 +0200 Subject: [PATCH] Add option to enabled/disable client stats and fix #1771 (#2001) * Add option to enabled/disable client stats and fix #1771 * more fix * fix map_density * even more fix * remove need for vec in Aggregator::aggregate * fix json weirdness - remove individual clients (is that all right? ) * Make pretty --- libafl/src/monitors/disk.rs | 5 ++-- libafl/src/monitors/mod.rs | 47 +++++++++++++++++++++---------- libafl/src/monitors/multi.rs | 4 +-- libafl/src/monitors/prometheus.rs | 4 +-- libafl/src/monitors/tui/mod.rs | 42 +++++++++++---------------- 5 files changed, 55 insertions(+), 47 deletions(-) diff --git a/libafl/src/monitors/disk.rs b/libafl/src/monitors/disk.rs index 4419f97d0f..f977eb1da5 100644 --- a/libafl/src/monitors/disk.rs +++ b/libafl/src/monitors/disk.rs @@ -72,7 +72,7 @@ executions = {} exec_sec = {} ", format_duration_hms(&(cur_time - self.start_time())), - self.client_stats().len(), + self.client_stats_count(), self.corpus_size(), self.objective_size(), self.total_execs(), @@ -211,12 +211,11 @@ where let line = json!({ "run_time": current_time() - self.base.start_time(), - "clients": self.base.client_stats().len(), + "clients": self.client_stats_count(), "corpus": self.base.corpus_size(), "objectives": self.base.objective_size(), "executions": self.base.total_execs(), "exec_sec": self.base.execs_per_sec(), - "clients": &self.client_stats().get(1..) }); writeln!(&file, "{line}").expect("Unable to write JSON to file"); } diff --git a/libafl/src/monitors/mod.rs b/libafl/src/monitors/mod.rs index 508767588e..c37f051e60 100644 --- a/libafl/src/monitors/mod.rs +++ b/libafl/src/monitors/mod.rs @@ -62,22 +62,20 @@ impl Aggregator { /// takes the key and the ref to clients stats then aggregate them all. fn aggregate(&mut self, name: &str, client_stats: &[ClientStats]) { - let mut gather = vec![]; + let mut gather = client_stats + .iter() + .filter_map(|client| client.user_monitor.get(name)); - for client in client_stats { - if let Some(x) = client.user_monitor.get(name) { - gather.push(x); - } - } + let gather_count = gather.clone().count(); - let (mut init, op) = match gather.first() { + let (mut init, op) = match gather.next() { Some(x) => (x.value().clone(), x.aggregator_op().clone()), _ => { return; } }; - for item in gather.iter().skip(1) { + for item in gather { match op { AggregatorOps::None => { // Nothing @@ -112,7 +110,7 @@ impl Aggregator { if let AggregatorOps::Avg = op { // if avg then divide last. - init = match init.stats_div(gather.len()) { + init = match init.stats_div(gather_count) { Some(x) => x, _ => { return; @@ -340,6 +338,8 @@ fn prettify_float(value: f64) -> String { /// A simple struct to keep track of client monitor #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct ClientStats { + /// If this client is enabled. This is set to `true` the first time we see this client. + pub enabled: bool, // monitor (maybe we need a separated struct?) /// The corpus size for this client pub corpus_size: u64, @@ -506,6 +506,14 @@ pub trait Monitor { .fold(0_u64, |acc, x| acc + x.corpus_size) } + /// Count the number of enabled client stats + fn client_stats_count(&self) -> usize { + self.client_stats() + .iter() + .filter(|client| client.enabled) + .count() + } + /// Amount of elements in the objectives (combined for all children) fn objective_size(&self) -> u64 { self.client_stats() @@ -538,14 +546,23 @@ pub trait Monitor { /// The client monitor for a specific id, creating new if it doesn't exist fn client_stats_insert(&mut self, client_id: ClientId) { - let client_stat_count = self.client_stats().len(); - for _ in client_stat_count..(client_id.0 + 1) as usize { + let total_client_stat_count = self.client_stats().len(); + for _ in total_client_stat_count..=(client_id.0) as usize { self.client_stats_mut().push(ClientStats { - last_window_time: current_time(), - start_time: current_time(), + enabled: false, + last_window_time: Duration::from_secs(0), + start_time: Duration::from_secs(0), ..ClientStats::default() }); } + let new_stat = self.client_stats_mut_for(client_id); + if !new_stat.enabled { + let timestamp = current_time(); + // I have never seen this man in my life + new_stat.start_time = timestamp; + new_stat.last_window_time = timestamp; + new_stat.enabled = true; + } } /// Get mutable reference to client stats @@ -673,7 +690,7 @@ impl Monitor for SimplePrintingMonitor { event_msg, sender_id.0, format_duration_hms(&(current_time() - self.start_time)), - self.client_stats().len(), + self.client_stats_count(), self.corpus_size(), self.objective_size(), self.total_execs(), @@ -749,7 +766,7 @@ where event_msg, sender_id.0, format_duration_hms(&(current_time() - self.start_time)), - self.client_stats().len(), + self.client_stats_count(), self.corpus_size(), self.objective_size(), self.total_execs(), diff --git a/libafl/src/monitors/multi.rs b/libafl/src/monitors/multi.rs index 9ef49fe328..cf338781f3 100644 --- a/libafl/src/monitors/multi.rs +++ b/libafl/src/monitors/multi.rs @@ -75,7 +75,7 @@ where "[{}] (GLOBAL) run time: {}, clients: {}, corpus: {}, objectives: {}, executions: {}, exec/sec: {}", head, format_duration_hms(&(current_time() - self.start_time)), - self.client_stats().len(), + self.client_stats_count(), self.corpus_size(), self.objective_size(), self.total_execs(), @@ -106,7 +106,7 @@ where #[cfg(feature = "introspection")] { // Print the client performance monitor. Skip the Client 0 which is the broker - for (i, client) in self.client_stats.iter().skip(1).enumerate() { + for (i, client) in self.client_stats.iter().filter(|x| x.enabled).enumerate() { let fmt = format!("Client {:03}:\n{}", i + 1, client.introspection_monitor); (self.print_fn)(&fmt); } diff --git a/libafl/src/monitors/prometheus.rs b/libafl/src/monitors/prometheus.rs index a3c1475802..63132b5896 100644 --- a/libafl/src/monitors/prometheus.rs +++ b/libafl/src/monitors/prometheus.rs @@ -142,7 +142,7 @@ where stat: String::new(), }) .set(run_time.try_into().unwrap()); // run time in seconds, which can be converted to a time format by Grafana or similar - let total_clients = self.client_stats().len().try_into().unwrap(); // convert usize to u64 (unlikely that # of clients will be > 2^64 -1...) + let total_clients = self.client_stats_count().try_into().unwrap(); // convert usize to u64 (unlikely that # of clients will be > 2^64 -1...) self.clients_count .get_or_create(&Labels { client: sender_id.0, @@ -156,7 +156,7 @@ where event_msg, sender_id.0, format_duration_hms(&(current_time() - self.start_time)), - self.client_stats().len(), + self.client_stats_count(), self.corpus_size(), self.objective_size(), self.total_execs(), diff --git a/libafl/src/monitors/tui/mod.rs b/libafl/src/monitors/tui/mod.rs index 5373e9f275..d407e0ad9d 100644 --- a/libafl/src/monitors/tui/mod.rs +++ b/libafl/src/monitors/tui/mod.rs @@ -1,6 +1,7 @@ //! Monitor based on ratatui use alloc::{boxed::Box, string::ToString}; +use core::cmp; use std::{ collections::VecDeque, fmt::Write as _, @@ -336,12 +337,16 @@ pub struct TuiMonitor { } impl Monitor for TuiMonitor { - /// the client monitor, mutable + /// The client monitor, mutable + /// This also includes disabled "padding" clients. + /// Results should be filtered by `.enabled`. fn client_stats_mut(&mut self) -> &mut Vec { &mut self.client_stats } - /// the client monitor + /// The client monitor + /// This also includes disabled "padding" clients. + /// Results should be filtered by `.enabled`. fn client_stats(&self) -> &[ClientStats] { &self.client_stats } @@ -419,8 +424,8 @@ impl Monitor for TuiMonitor { #[cfg(feature = "introspection")] { - // Print the client performance monitor. Skip the Client 0 which is the broker - for (i, client) in self.client_stats.iter().skip(1).enumerate() { + // Print the client performance monitor. Skip the Client IDs that have never sent anything. + for (i, client) in self.client_stats.iter().filter(|x| x.enabled).enumerate() { self.context .write() .unwrap() @@ -484,25 +489,12 @@ impl TuiMonitor { } fn map_density(&self) -> String { - if self.client_stats.len() < 2 { - return "0%".to_string(); - } - let mut max_map_density = self - .client_stats() - .get(1) - .unwrap() - .get_user_stats("edges") - .map_or("0%".to_string(), ToString::to_string); - - for client in self.client_stats().iter().skip(2) { - let client_map_density = client - .get_user_stats("edges") - .map_or(String::new(), ToString::to_string); - if client_map_density > max_map_density { - max_map_density = client_map_density; - } - } - max_map_density + self.client_stats() + .iter() + .filter(|client| client.enabled) + .filter_map(|client| client.get_user_stats("edges")) + .map(ToString::to_string) + .fold("0%".to_string(), cmp::max) } fn item_geometry(&self) -> ItemGeometry { @@ -512,7 +504,7 @@ impl TuiMonitor { } let mut ratio_a: u64 = 0; let mut ratio_b: u64 = 0; - for client in self.client_stats().iter().skip(1) { + for client in self.client_stats().iter().filter(|client| client.enabled) { let afl_stats = client .get_user_stats("AflStats") .map_or("None".to_string(), ToString::to_string); @@ -555,7 +547,7 @@ impl TuiMonitor { if self.client_stats.len() > 1 { let mut new_path_time = Duration::default(); let mut new_objectives_time = Duration::default(); - for client in self.client_stats().iter().skip(1) { + for client in self.client_stats().iter().filter(|client| client.enabled) { new_path_time = client.last_corpus_time.max(new_path_time); new_objectives_time = client.last_objective_time.max(new_objectives_time); }