implement the AFL-Style Tui (#1432)

* implement an AFL-Style TUI

* improve the tui/mod.rs according to the reviews

* fixing fmt manually

---------

Co-authored-by: toseven <[email protected]>
Co-authored-by: Dominik Maier <[email protected]>

Author: 3 people

libafl/src/monitors/mod.rs

@@ -87,12 +87,16 @@ pub struct ClientStats {
     // monitor (maybe we need a separated struct?)
     /// The corpus size for this client
     pub corpus_size: u64,
+    /// The time for the last update of the corpus size
+    pub last_corpus_time: Duration,
     /// The total executions for this client
     pub executions: u64,
     /// The number of executions of the previous state in case a client decrease the number of execution (e.g when restarting without saving the state)
     pub prev_state_executions: u64,
     /// The size of the objectives corpus for this client
     pub objective_size: u64,
+    /// The time for the last update of the objective size
+    pub last_objective_time: Duration,
     /// The last reported executions for this client
     #[cfg(feature = "afl_exec_sec")]
     pub last_window_executions: u64,
@@ -101,6 +105,8 @@ pub struct ClientStats {
     pub last_execs_per_sec: f64,
     /// The last time we got this information
     pub last_window_time: Duration,
+    /// the start time of the client
+    pub start_time: Duration,
     /// User-defined monitor
     pub user_monitor: HashMap<String, UserStats>,
     /// Client performance statistics
@@ -140,6 +146,7 @@ impl ClientStats {
     /// We got a new information about corpus size for this client, insert them.
     pub fn update_corpus_size(&mut self, corpus_size: u64) {
         self.corpus_size = corpus_size;
+        self.last_corpus_time = current_time();
     }
 
     /// We got a new information about objective corpus size for this client, insert them.
@@ -206,8 +213,9 @@ impl ClientStats {
         self.user_monitor.insert(name, value);
     }
 
+    #[must_use]
     /// Get a user-defined stat using the name
-    pub fn get_user_stats(&mut self, name: &str) -> Option<&UserStats> {
+    pub fn get_user_stats(&self, name: &str) -> Option<&UserStats> {
         self.user_monitor.get(name)
     }
 
@@ -275,6 +283,7 @@ pub trait Monitor {
         for _ in client_stat_count..(client_id.0 + 1) as usize {
             self.client_stats_mut().push(ClientStats {
                 last_window_time: current_time(),
+                start_time: current_time(),
                 ..ClientStats::default()
             });
         }

libafl/src/monitors/tui/mod.rs

@@ -1,6 +1,6 @@
 //! Monitor based on ratatui
 
-use alloc::boxed::Box;
+use alloc::{boxed::Box, string::ToString};
 use std::{
     collections::VecDeque,
     fmt::Write,
@@ -22,6 +22,7 @@ use crossterm::{
 use hashbrown::HashMap;
 use libafl_bolts::{current_time, format_duration_hms, ClientId};
 use ratatui::{backend::CrosstermBackend, Terminal};
+use serde_json::{self, Value};
 
 #[cfg(feature = "introspection")]
 use super::{ClientPerfMonitor, PerfFeature};
@@ -162,14 +163,53 @@ impl PerfTuiContext {
     }
 }
 
+#[derive(Debug, Default, Clone)]
+pub struct ProcessTiming {
+    pub client_start_time: Duration,
+    pub exec_speed: String,
+    pub last_new_entry: Duration,
+    pub last_saved_solution: Duration,
+}
+
+impl ProcessTiming {
+    fn new() -> Self {
+        Self {
+            exec_speed: "0".to_string(),
+            ..Default::default()
+        }
+    }
+}
+
+#[derive(Debug, Default, Clone)]
+pub struct ItemGeometry {
+    pub pending: u64,
+    pub pend_fav: u64,
+    pub own_finds: u64,
+    pub imported: u64,
+    pub stability: String,
+}
+
+impl ItemGeometry {
+    fn new() -> Self {
+        Self {
+            stability: "0%".to_string(),
+            ..Default::default()
+        }
+    }
+}
+
 #[derive(Debug, Default, Clone)]
 pub struct ClientTuiContext {
     pub corpus: u64,
     pub objectives: u64,
     pub executions: u64,
     /// Float value formatted as String
-    pub exec_sec: String,
+    pub map_density: String,
+
+    pub cycles_done: u64,
 
+    pub process_timing: ProcessTiming,
+    pub item_geometry: ItemGeometry,
     pub user_stats: HashMap<String, UserStats>,
 }
 
@@ -178,7 +218,47 @@ impl ClientTuiContext {
         self.corpus = client.corpus_size;
         self.objectives = client.objective_size;
         self.executions = client.executions;
-        self.exec_sec = exec_sec;
+        self.process_timing.client_start_time = client.start_time;
+        self.process_timing.last_new_entry = if client.last_corpus_time > client.start_time {
+            client.last_corpus_time - client.start_time
+        } else {
+            Duration::default()
+        };
+
+        self.process_timing.last_saved_solution = if client.last_objective_time > client.start_time
+        {
+            client.last_objective_time - client.start_time
+        } else {
+            Duration::default()
+        };
+
+        self.process_timing.exec_speed = exec_sec;
+
+        self.map_density = client
+            .get_user_stats("edges")
+            .map_or("0%".to_string(), ToString::to_string);
+
+        let default_json = serde_json::json!({
+            "pending": 0,
+            "pend_fav": 0,
+            "imported": 0,
+            "own_finds": 0,
+        });
+        let afl_stats = client
+            .get_user_stats("AflStats")
+            .map_or(default_json.to_string(), ToString::to_string);
+
+        let afl_stats_json: Value =
+            serde_json::from_str(afl_stats.as_str()).unwrap_or(default_json);
+        self.item_geometry.pending = afl_stats_json["pending"].as_u64().unwrap_or_default();
+        self.item_geometry.pend_fav = afl_stats_json["pend_fav"].as_u64().unwrap_or_default();
+        self.item_geometry.imported = afl_stats_json["imported"].as_u64().unwrap_or_default();
+        self.item_geometry.own_finds = afl_stats_json["own_finds"].as_u64().unwrap_or_default();
+
+        let stability = client
+            .get_user_stats("stability")
+            .map_or("0%".to_string(), ToString::to_string);
+        self.item_geometry.stability = stability;
 
         for (key, val) in &client.user_monitor {
             self.user_stats.insert(key.clone(), val.clone());
@@ -205,6 +285,14 @@ pub struct TuiContext {
     pub clients_num: usize,
     pub total_execs: u64,
     pub start_time: Duration,
+
+    pub total_map_density: String,
+    pub total_solutions: u64,
+    pub total_cycles_done: u64,
+    pub total_corpus_count: u64,
+
+    pub total_process_timing: ProcessTiming,
+    pub total_item_geometry: ItemGeometry,
 }
 
 impl TuiContext {
@@ -226,6 +314,13 @@ impl TuiContext {
             clients_num: 0,
             total_execs: 0,
             start_time,
+
+            total_map_density: "0%".to_string(),
+            total_solutions: 0,
+            total_cycles_done: 0,
+            total_corpus_count: 0,
+            total_item_geometry: ItemGeometry::new(),
+            total_process_timing: ProcessTiming::new(),
         }
     }
 }
@@ -264,14 +359,21 @@ impl Monitor for TuiMonitor {
             let execsec = self.execs_per_sec() as u64;
             let totalexec = self.total_execs();
             let run_time = cur_time - self.start_time;
+            let total_process_timing = self.process_timing();
 
             let mut ctx = self.context.write().unwrap();
+            ctx.total_process_timing = total_process_timing;
             ctx.corpus_size_timed.add(run_time, self.corpus_size());
             ctx.objective_size_timed
                 .add(run_time, self.objective_size());
             ctx.execs_per_sec_timed.add(run_time, execsec);
             ctx.total_execs = totalexec;
             ctx.clients_num = self.client_stats.len();
+            ctx.total_map_density = self.map_density();
+            ctx.total_solutions = self.objective_size();
+            ctx.total_cycles_done = 0;
+            ctx.total_corpus_count = self.corpus_size();
+            ctx.total_item_geometry = self.item_geometry();
         }
 
         let client = self.client_stats_mut_for(sender_id);
@@ -339,6 +441,91 @@ impl TuiMonitor {
             client_stats: vec![],
         }
     }
+
+    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
+    }
+
+    fn item_geometry(&self) -> ItemGeometry {
+        let mut total_item_geometry = ItemGeometry::new();
+        if self.client_stats.len() < 2 {
+            return total_item_geometry;
+        }
+        let mut ratio_a: u64 = 0;
+        let mut ratio_b: u64 = 0;
+        for client in self.client_stats().iter().skip(1) {
+            let afl_stats = client
+                .get_user_stats("AflStats")
+                .map_or("None".to_string(), ToString::to_string);
+            let stability = client
+                .get_user_stats("stability")
+                .map_or(&UserStats::Ratio(0, 100), |x| x);
+
+            if afl_stats != "None" {
+                let default_json = serde_json::json!({
+                    "pending": 0,
+                    "pend_fav": 0,
+                    "imported": 0,
+                    "own_finds": 0,
+                });
+                let afl_stats_json: Value =
+                    serde_json::from_str(afl_stats.as_str()).unwrap_or(default_json);
+                total_item_geometry.pending +=
+                    afl_stats_json["pending"].as_u64().unwrap_or_default();
+                total_item_geometry.pend_fav +=
+                    afl_stats_json["pend_fav"].as_u64().unwrap_or_default();
+                total_item_geometry.own_finds +=
+                    afl_stats_json["own_finds"].as_u64().unwrap_or_default();
+                total_item_geometry.imported +=
+                    afl_stats_json["imported"].as_u64().unwrap_or_default();
+            }
+
+            if let UserStats::Ratio(a, b) = stability {
+                ratio_a += a;
+                ratio_b += b;
+            }
+        }
+        total_item_geometry.stability = format!("{}%", ratio_a * 100 / ratio_b);
+        total_item_geometry
+    }
+
+    fn process_timing(&mut self) -> ProcessTiming {
+        let mut total_process_timing = ProcessTiming::new();
+        total_process_timing.exec_speed = self.execs_per_sec_pretty();
+        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) {
+                new_path_time = client.last_corpus_time.max(new_path_time);
+                new_objectives_time = client.last_objective_time.max(new_objectives_time);
+            }
+            if new_path_time > self.start_time {
+                total_process_timing.last_new_entry = new_path_time - self.start_time;
+            }
+            if new_objectives_time > self.start_time {
+                total_process_timing.last_saved_solution = new_objectives_time - self.start_time;
+            }
+        }
+        total_process_timing
+    }
 }
 
 fn run_tui_thread(context: Arc<RwLock<TuiContext>>, tick_rate: Duration, tui_ui: TuiUI) {