From da8f17f29e7d2f0b0b682cdb05ee510f8aec977b Mon Sep 17 00:00:00 2001
From: Dominik Maier <domenukk@gmail.com>
Date: Mon, 14 Oct 2024 11:06:21 +0200
Subject: [PATCH] Add AFLStats to QEMU_Launcher TUI (#2611)

* Remove prelude from qemu_launcher

* QEMU_Launcher: Add stats to TUI

* More tui
---
 fuzzers/binary_only/qemu_launcher/src/client.rs  |  2 +-
 .../binary_only/qemu_launcher/src/instance.rs    | 16 +++++++++++-----
 fuzzers/binary_only/qemu_launcher/src/options.rs |  2 +-
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/fuzzers/binary_only/qemu_launcher/src/client.rs b/fuzzers/binary_only/qemu_launcher/src/client.rs
index b702ab2ddd..abe0488911 100644
--- a/fuzzers/binary_only/qemu_launcher/src/client.rs
+++ b/fuzzers/binary_only/qemu_launcher/src/client.rs
@@ -34,7 +34,7 @@ pub struct Client<'a> {
     options: &'a FuzzerOptions,
 }
 
-impl<'a> Client<'a> {
+impl Client<'_> {
     pub fn new(options: &FuzzerOptions) -> Client {
         Client { options }
     }
diff --git a/fuzzers/binary_only/qemu_launcher/src/instance.rs b/fuzzers/binary_only/qemu_launcher/src/instance.rs
index 438dadcb1a..1a0fd277e9 100644
--- a/fuzzers/binary_only/qemu_launcher/src/instance.rs
+++ b/fuzzers/binary_only/qemu_launcher/src/instance.rs
@@ -1,5 +1,5 @@
 use core::{fmt::Debug, ptr::addr_of_mut};
-use std::{fs, marker::PhantomData, ops::Range, process};
+use std::{fs, marker::PhantomData, ops::Range, process, time::Duration};
 
 #[cfg(feature = "simplemgr")]
 use libafl::events::SimpleEventManager;
@@ -23,8 +23,8 @@ use libafl::{
         powersched::PowerSchedule, IndexesLenTimeMinimizerScheduler, PowerQueueScheduler,
     },
     stages::{
-        calibrate::CalibrationStage, power::StdPowerMutationalStage, ShadowTracingStage,
-        StagesTuple, StdMutationalStage,
+        calibrate::CalibrationStage, power::StdPowerMutationalStage, AflStatsStage, IfStage,
+        ShadowTracingStage, StagesTuple, StdMutationalStage,
     },
     state::{HasCorpus, StdState, UsesState},
     Error, HasMetadata, NopFuzzer,
@@ -73,7 +73,7 @@ pub struct Instance<'a, M: Monitor> {
     phantom: PhantomData<M>,
 }
 
-impl<'a, M: Monitor> Instance<'a, M> {
+impl<M: Monitor> Instance<'_, M> {
     #[allow(clippy::similar_names)] // elf != self
     fn coverage_filter(&self, qemu: Qemu) -> Result<StdAddressFilter, Error> {
         /* Conversion is required on 32-bit targets, but not on 64-bit ones */
@@ -107,6 +107,7 @@ impl<'a, M: Monitor> Instance<'a, M> {
         }
     }
 
+    #[allow(clippy::too_many_lines)]
     pub fn run<ET>(&mut self, modules: ET, state: Option<ClientState>) -> Result<(), Error>
     where
         ET: EmulatorModuleTuple<ClientState> + Debug,
@@ -135,6 +136,11 @@ impl<'a, M: Monitor> Instance<'a, M> {
 
         let calibration = CalibrationStage::new(&map_feedback);
 
+        let stats_stage = IfStage::new(
+            |_, _, _, _| Ok(self.options.tui),
+            tuple_list!(AflStatsStage::new(Duration::from_secs(5))),
+        );
+
         // Feedback to rate the interestingness of an input
         // This one is composed by two Feedbacks in OR
         let mut feedback = feedback_or!(
@@ -268,7 +274,7 @@ impl<'a, M: Monitor> Instance<'a, M> {
                 StdPowerMutationalStage::new(mutator);
 
             // The order of the stages matter!
-            let mut stages = tuple_list!(calibration, tracing, i2s, power);
+            let mut stages = tuple_list!(calibration, tracing, i2s, power, stats_stage);
 
             self.fuzz(&mut state, &mut fuzzer, &mut executor, &mut stages)
         } else {
diff --git a/fuzzers/binary_only/qemu_launcher/src/options.rs b/fuzzers/binary_only/qemu_launcher/src/options.rs
index 32bc8a6294..56c25fed27 100644
--- a/fuzzers/binary_only/qemu_launcher/src/options.rs
+++ b/fuzzers/binary_only/qemu_launcher/src/options.rs
@@ -199,7 +199,7 @@ impl FuzzerOptions {
             let mut cmd = FuzzerOptions::command();
             cmd.error(
                 ErrorKind::ValueValidation,
-                format!("The `drcov` option is only supported with `rerun_input`."),
+                "The `drcov` option is only supported with `rerun_input`.".to_string(),
             )
             .exit();
         }