From 6cb2be4408f7d4429f4aaf5913b0e2b97a9cd991 Mon Sep 17 00:00:00 2001
From: Alwin Berger <alwin.berger@tu-dortmund.de>
Date: Thu, 20 Jun 2024 10:15:01 +0200
Subject: [PATCH] add TimeProbMassScheduler

---
 fuzzers/FRET/src/fuzzer.rs |  6 +++---
 fuzzers/FRET/src/worst.rs  | 30 +++++++++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/fuzzers/FRET/src/fuzzer.rs b/fuzzers/FRET/src/fuzzer.rs
index 24239a77c9..54c0bef3fc 100644
--- a/fuzzers/FRET/src/fuzzer.rs
+++ b/fuzzers/FRET/src/fuzzer.rs
@@ -14,7 +14,7 @@ edges::{self, edges_map_mut_ptr, QemuEdgeCoverageHelper, MAX_EDGES_FOUND}, elf::
 };
 use rand::{SeedableRng, StdRng, Rng};
 use crate::{
-clock::{ClockTimeFeedback, IcHist, QemuClockIncreaseFeedback, QemuClockObserver, FUZZ_START_TIMESTAMP}, mutational::{input_bytes_to_interrupt_times, InterruptShiftStage, MINIMUM_INTER_ARRIVAL_TIME}, qemustate::QemuStateRestoreHelper, systemstate::{self, feedbacks::{DumpSystraceFeedback, NovelSystemStateFeedback, SystraceErrorFeedback}, helpers::{QemuSystemStateHelper, ISR_SYMBOLS}, observers::QemuSystemStateObserver, schedulers::{GenerationScheduler, LongestTraceScheduler}, stg::{stg_map_mut_slice, GraphMaximizerCorpusScheduler, STGEdge, STGNode, StgFeedback, MAX_STG_NUM}}, worst::{AlwaysTrueFeedback, ExecTimeIncFeedback, TimeMaximizerCorpusScheduler, TimeStateMaximizerCorpusScheduler} 
+clock::{ClockTimeFeedback, IcHist, QemuClockIncreaseFeedback, QemuClockObserver, FUZZ_START_TIMESTAMP}, mutational::{input_bytes_to_interrupt_times, InterruptShiftStage, MINIMUM_INTER_ARRIVAL_TIME}, qemustate::QemuStateRestoreHelper, systemstate::{self, feedbacks::{DumpSystraceFeedback, NovelSystemStateFeedback, SystraceErrorFeedback}, helpers::{QemuSystemStateHelper, ISR_SYMBOLS}, observers::QemuSystemStateObserver, schedulers::{GenerationScheduler, LongestTraceScheduler}, stg::{stg_map_mut_slice, GraphMaximizerCorpusScheduler, STGEdge, STGNode, StgFeedback, MAX_STG_NUM}}, worst::{AlwaysTrueFeedback, ExecTimeIncFeedback, TimeMaximizerCorpusScheduler, TimeProbMassScheduler, TimeStateMaximizerCorpusScheduler} 
 };
 use std::time::{SystemTime, UNIX_EPOCH};
 use clap::{Parser, Subcommand};
@@ -583,9 +583,9 @@ let mut run_client = |state: Option<_>, mut mgr, _core_id| {
         #[cfg(not(any(feature = "sched_afl", feature = "sched_stg", feature = "sched_genetic")))]
         let scheduler = QueueScheduler::new();  // fallback
         #[cfg(feature = "sched_afl",)]
-        let scheduler = TimeMaximizerCorpusScheduler::new(QueueScheduler::new());
+        let scheduler = TimeMaximizerCorpusScheduler::new(TimeProbMassScheduler::new());
         #[cfg(feature = "sched_stg")]
-        let scheduler = LongestTraceScheduler::new(GraphMaximizerCorpusScheduler::new(&stg_coverage_observer,QueueScheduler::new()));
+        let scheduler = GraphMaximizerCorpusScheduler::new(&stg_coverage_observer,TimeProbMassScheduler::new());
         #[cfg(feature = "sched_genetic")]
         let scheduler = GenerationScheduler::new();
 
diff --git a/fuzzers/FRET/src/worst.rs b/fuzzers/FRET/src/worst.rs
index 4db3ac1cde..77783d0120 100644
--- a/fuzzers/FRET/src/worst.rs
+++ b/fuzzers/FRET/src/worst.rs
@@ -6,7 +6,7 @@ use libafl::feedbacks::MapIndexesMetadata;
 use libafl::corpus::Testcase;
 use libafl::prelude::{UsesInput};
 use core::marker::PhantomData;
-use libafl::schedulers::{MinimizerScheduler, TestcaseScore};
+use libafl::schedulers::{MinimizerScheduler, ProbabilitySamplingScheduler, TestcaseScore};
 use std::path::PathBuf;
 use std::fs;
 use hashbrown::{HashMap};
@@ -408,4 +408,32 @@ where
             name: Cow::from("AlwaysTrueFeedback".to_string())
         }
     }
+}
+
+
+//=========================== Probability Mass Scheduler
+
+pub type TimeProbMassScheduler<S> =
+    ProbabilitySamplingScheduler<TimeProbFactor<S>, S>;
+
+#[derive(Debug, Clone)]
+pub struct TimeProbFactor<S>
+where
+    S: HasCorpus + HasMetadata,
+    S::Input: HasLen,
+{
+    phantom: PhantomData<S>,
+}
+
+impl<S> TestcaseScore<S> for TimeProbFactor<S>
+where
+    S: HasCorpus + HasMetadata,
+    S::Input: HasLen,
+{
+    fn compute(_state: &S, entry: &mut Testcase<<S as UsesInput>::Input>) -> Result<f64, Error> {
+        // TODO maybe enforce entry.exec_time().is_some()
+        let et = entry.exec_time().expect("testcase.exec_time is needed for scheduler");
+        let tns : i64 = et.as_nanos().try_into().expect("failed to convert time");
+        Ok((tns as f64)/1000.0) //microseconds
+    }
 }
\ No newline at end of file