Fixes to TUI monitor if main thread panics (#699)

* Trying to fix the tui if the main thread panics

* cargo fmt

* Prettifying code

fuzzers/baby_fuzzer/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2021"
 
 [features]
 default = ["std"]
+tui = []
 std = []
 
 [profile.dev]

fuzzers/baby_fuzzer/src/main.rs

@@ -3,6 +3,10 @@ use std::path::PathBuf;
 #[cfg(windows)]
 use std::ptr::write_volatile;
 
+#[cfg(feature = "tui")]
+use libafl::monitors::tui::TuiMonitor;
+#[cfg(not(feature = "tui"))]
+use libafl::monitors::SimpleMonitor;
 use libafl::{
     bolts::{current_nanos, rands::StdRand, tuples::tuple_list, AsSlice},
     corpus::{InMemoryCorpus, OnDiskCorpus},
@@ -12,7 +16,6 @@ use libafl::{
     fuzzer::{Fuzzer, StdFuzzer},
     generators::RandPrintablesGenerator,
     inputs::{BytesInput, HasTargetBytes},
-    monitors::SimpleMonitor,
     mutators::scheduled::{havoc_mutations, StdScheduledMutator},
     observers::StdMapObserver,
     schedulers::QueueScheduler,
@@ -85,7 +88,10 @@ pub fn main() {
     .unwrap();
 
     // The Monitor trait define how the fuzzer stats are displayed to the user
+    #[cfg(not(feature = "tui"))]
     let mon = SimpleMonitor::new(|s| println!("{}", s));
+    #[cfg(feature = "tui")]
+    let mon = TuiMonitor::new(String::from("Baby Fuzzer"), false);
 
     // The event manager handle the various events generated during the fuzzing loop
     // such as the notification of the addition of a new item to the corpus

libafl/src/monitors/tui/mod.rs

@@ -1,6 +1,7 @@
 //! Monitor based on tui-rs
 
 use crossterm::{
+    cursor::{EnableBlinking, Show},
     event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode},
     execute,
     terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
@@ -12,6 +13,7 @@ use std::{
     collections::VecDeque,
     fmt::Write,
     io::{self, BufRead},
+    panic,
     string::String,
     sync::{Arc, RwLock},
     thread,
@@ -27,6 +29,8 @@ use crate::{
     monitors::{ClientStats, Monitor, UserStats},
 };
 
+use alloc::boxed::Box;
+
 mod ui;
 use ui::TuiUI;
 
@@ -361,6 +365,22 @@ fn run_tui_thread(
 
         let mut last_tick = Instant::now();
         let mut cnt = 0;
+
+        // Catching panics when the main thread dies
+        let old_hook = panic::take_hook();
+        panic::set_hook(Box::new(move |panic_info| {
+            disable_raw_mode().unwrap();
+            execute!(
+                io::stdout(),
+                LeaveAlternateScreen,
+                DisableMouseCapture,
+                Show,
+                EnableBlinking,
+            )
+            .unwrap();
+            old_hook(panic_info);
+        }));
+
         loop {
             // to avoid initial ui glitches
             if cnt < 8 {

libafl/src/monitors/tui/ui.rs

@@ -62,16 +62,20 @@ impl TuiUI {
     //pub fn on_down(&mut self) {}
 
     pub fn on_right(&mut self) {
-        // never 0
-        self.clients_idx = 1 + self.clients_idx % (self.clients - 1);
+        if self.clients != 0 {
+            // clients_idx never 0
+            self.clients_idx = 1 + self.clients_idx % (self.clients - 1);
+        }
     }
 
     pub fn on_left(&mut self) {
-        // never 0
-        if self.clients_idx == 1 {
-            self.clients_idx = self.clients - 1;
-        } else {
-            self.clients_idx = 1 + (self.clients_idx - 2) % (self.clients - 1);
+        if self.clients != 0 {
+            // clients_idx never 0
+            if self.clients_idx == 1 {
+                self.clients_idx = self.clients - 1;
+            } else {
+                self.clients_idx = 1 + (self.clients_idx - 2) % (self.clients - 1);
+            }
         }
     }