Fix TUI (#3151)

Author: tokatoka

libafl/src/monitors/stats/manager.rs

@@ -230,7 +230,7 @@ impl ClientStatsManager {
     #[must_use]
     pub fn item_geometry(&self) -> ItemGeometry {
         let mut total_item_geometry = ItemGeometry::new();
-        if self.client_stats().len() < 2 {
+        if self.client_stats.is_empty() {
             return total_item_geometry;
         }
         let mut ratio_a: u64 = 0;
@@ -240,23 +240,15 @@ impl ClientStatsManager {
             .iter()
             .filter(|(_, client)| client.enabled())
         {
-            let afl_stats = client
-                .get_user_stats("AflStats")
-                .map_or("None".to_string(), ToString::to_string);
+            let afl_stats = client.get_user_stats("AflStats");
             let stability = client.get_user_stats("stability").map_or(
                 UserStats::new(UserStatsValue::Ratio(0, 100), AggregatorOps::Avg),
                 Clone::clone,
             );
 
-            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);
+            if let Some(stat) = afl_stats {
+                let stats = stat.to_string();
+                let afl_stats_json: Value = serde_json::from_str(stats.as_str()).unwrap();
                 total_item_geometry.pending +=
                     afl_stats_json["pending"].as_u64().unwrap_or_default();
                 total_item_geometry.pend_fav +=
@@ -277,6 +269,7 @@ impl ClientStatsManager {
         } else {
             Some((ratio_a as f64) / (ratio_b as f64))
         };
+
         total_item_geometry
     }
 }

libafl/src/monitors/tui/ui.rs

@@ -1,5 +1,5 @@
 //! The UI-specific parts of [`super::TuiMonitor`]
-use alloc::{string::ToString, sync::Arc, vec::Vec};
+use alloc::{sync::Arc, vec::Vec};
 use core::cmp::{max, min};
 use std::sync::RwLock;
 
@@ -27,13 +27,27 @@ pub struct TuiUi {
     enhanced_graphics: bool,
     show_logs: bool,
     client_idx: usize,
-    clients: usize,
+    clients: Vec<usize>,
     charts_tab_idx: usize,
     graph_data: Vec<(f64, f64)>,
 
     pub should_quit: bool,
 }
 
+fn next_larger(sorted: &[usize], value: usize) -> Option<usize> {
+    if let Some(index) = sorted.iter().position(|x| *x > value) {
+        return Some(index);
+    }
+    None
+}
+
+fn next_smaller(sorted: &[usize], value: usize) -> Option<usize> {
+    if let Some(index) = sorted.iter().rposition(|x| *x < value) {
+        return Some(index);
+    }
+    None
+}
+
 impl TuiUi {
     #[must_use]
     pub fn new(title: String, enhanced_graphics: bool) -> Self {
@@ -68,20 +82,31 @@ impl TuiUi {
     }
 
     pub fn on_right(&mut self) {
-        if self.clients > 0 && self.client_idx < self.clients - 1 {
-            self.client_idx += 1;
+        if let Some(idx) = next_larger(&self.clients, self.client_idx) {
+            self.client_idx = self.clients[idx];
         }
     }
 
     pub fn on_left(&mut self) {
-        if self.client_idx > 0 {
-            self.client_idx -= 1;
+        if let Some(idx) = next_smaller(&self.clients, self.client_idx) {
+            self.client_idx = self.clients[idx];
         }
     }
 
     /// Draw the current TUI context
     pub fn draw(&mut self, f: &mut Frame, app: &Arc<RwLock<TuiContext>>) {
-        self.clients = app.read().unwrap().clients_num;
+        let new = app.read().unwrap().clients_num;
+        if new != self.clients.len() {
+            // get the list of all clients
+            let mut all: Vec<usize> = app.read().unwrap().clients.keys().copied().collect();
+            all.sort_unstable();
+
+            // move the current client to the first one
+            self.client_idx = all[0];
+
+            // move the vector holding all clients ids
+            self.clients = all;
+        }
 
         let body = Layout::default()
             .constraints(if self.show_logs {
@@ -417,7 +442,7 @@ impl TuiUi {
         let empty_geometry: ItemGeometry = ItemGeometry::new();
         let item_geometry: &ItemGeometry = if is_overall {
             &tui_context.total_item_geometry
-        } else if self.clients == 0 {
+        } else if self.clients.is_empty() {
             &empty_geometry
         } else {
             let clients = &tui_context.clients;
@@ -494,7 +519,7 @@ impl TuiUi {
         let empty_timing: ProcessTiming = ProcessTiming::new();
         let tup: (Duration, &ProcessTiming) = if is_overall {
             (tui_context.start_time, &tui_context.total_process_timing)
-        } else if self.clients == 0 {
+        } else if self.clients.is_empty() {
             (current_time(), &empty_timing)
         } else {
             let clients = &tui_context.clients;
@@ -570,11 +595,9 @@ impl TuiUi {
             vec![
                 Row::new(vec![
                     Cell::from(Span::raw("clients")),
-                    Cell::from(Span::raw(format!("{}", self.clients))),
+                    Cell::from(Span::raw(format!("{}", self.clients.len()))),
                     Cell::from(Span::raw("total execs")),
                     Cell::from(Span::raw(format_big_number(app.total_execs))),
-                    Cell::from(Span::raw("map density")),
-                    Cell::from(Span::raw(app.total_map_density.to_string())),
                 ]),
                 Row::new(vec![
                     Cell::from(Span::raw("solutions")),
@@ -683,14 +706,6 @@ impl TuiUi {
                             .map_or(0, |x| x.executions),
                     ))),
                 ]),
-                Row::new(vec![
-                    Cell::from(Span::raw("map density")),
-                    Cell::from(Span::raw(
-                        app.clients
-                            .get(&self.client_idx)
-                            .map_or("0%".to_string(), |x| x.map_density.to_string()),
-                    )),
-                ]),
             ]
         };
 

libafl/src/stages/afl_stats.rs

@@ -410,7 +410,7 @@ where
             "{{\
                 \"pending\":{},\
                 \"pending_fav\":{},\
-                \"own_finds:\"{},\
+                \"own_finds\":{},\
                 \"imported\":{}\
             }}",
             stats.pending_total, stats.pending_favs, stats.corpus_found, stats.corpus_imported