FRET-LibAFL/fuzzers/FRET/src/clock.rs

use hashbrown::{hash_map::Entry, HashMap};
use libafl::{
    bolts::{
        current_nanos,
        rands::StdRand,
        tuples::{tuple_list},
    },
    executors::{ExitKind},
    fuzzer::{StdFuzzer},
    inputs::{BytesInput, HasTargetBytes},
    observers::{Observer,VariableMapObserver},
    state::{StdState, HasNamedMetadata},
    Error,
    observers::ObserversTuple, prelude::UsesInput, impl_serdeany,
};
use serde::{Deserialize, Serialize};
use std::{cell::UnsafeCell, cmp::max, env, fs::OpenOptions, io::Write, time::Instant};
use libafl::bolts::tuples::Named;

use libafl_qemu::{
    emu,
    emu::Emulator,
    executor::QemuExecutor,
    helper::{QemuHelper, QemuHelperTuple, QemuInstrumentationFilter},
};
use libafl::events::EventFirer;
use libafl::state::HasClientPerfMonitor;
use libafl::inputs::Input;
use libafl::feedbacks::Feedback;
use libafl::SerdeAny;
use libafl::state::HasMetadata;
use libafl::corpus::testcase::Testcase;
use core::{fmt::Debug, time::Duration};
// use libafl::feedbacks::FeedbackState;
// use libafl::state::HasFeedbackStates;
use libafl::bolts::tuples::MatchName;
use std::time::{SystemTime, UNIX_EPOCH};

pub static mut FUZZ_START_TIMESTAMP : SystemTime = UNIX_EPOCH;

//========== Metadata
#[derive(Debug, SerdeAny, Serialize, Deserialize)]
pub struct QemuIcountMetadata {
    runtime: u64,
}

/// Metadata for [`QemuClockIncreaseFeedback`]
#[derive(Debug, Serialize, Deserialize, SerdeAny)]
pub struct MaxIcountMetadata {
    pub max_icount_seen: u64,
    pub name: String,
}

// impl FeedbackState for MaxIcountMetadata
// {
//     fn reset(&mut self) -> Result<(), Error> {
//         self.max_icount_seen = 0;
//         Ok(())
//     }
// }

impl Named for MaxIcountMetadata
{
    #[inline]
    fn name(&self) -> &str {
        self.name.as_str()
    }
}

impl MaxIcountMetadata
{
    /// Create new `MaxIcountMetadata`
    #[must_use]
    pub fn new(name: &'static str) -> Self {
        Self {
            max_icount_seen: 0,
            name: name.to_string(),
        }
    }
}

impl Default for MaxIcountMetadata {
    fn default() -> Self {
        Self::new("MaxClock")
    }
}

/// A piece of metadata tracking all icounts
#[derive(Debug, SerdeAny, Serialize, Deserialize)]
pub struct IcHist (pub Vec<(u64, u128)>, pub (u64,u128));

//========== Observer

/// A simple observer, just overlooking the runtime of the target.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct QemuClockObserver {
    name: String,
    start_tick: u64,
    end_tick: u64,
}

impl QemuClockObserver {
    /// Creates a new [`QemuClockObserver`] with the given name.
    #[must_use]
    pub fn new(name: &'static str) -> Self {
        Self {
            name: name.to_string(),
            start_tick: 0,
            end_tick: 0,
        }
    }

    /// Gets the runtime for the last execution of this target.
    #[must_use]
    pub fn last_runtime(&self) -> u64 {
        self.end_tick - self.start_tick
    }
}

impl<S> Observer<S> for QemuClockObserver
where
    S: UsesInput + HasMetadata,
{
    fn pre_exec(&mut self, _state: &mut S, _input: &S::Input) -> Result<(), Error> {
        // Only remember the pre-run ticks if presistent mode ist used
        #[cfg(not(feature = "snapshot_restore"))]
        unsafe {
            self.start_tick=emu::icount_get_raw();
            self.end_tick=self.start_tick;
        }
        // unsafe {
        //     println!("clock pre {}",emu::icount_get_raw());
        // }
        Ok(())
    }

    fn post_exec(&mut self, _state: &mut S, _input: &S::Input, _exit_kind: &ExitKind) -> Result<(), Error> {
        unsafe { self.end_tick = emu::icount_get_raw() };
        // println!("clock post {}", self.end_tick);
        // println!("Number of Ticks: {} <- {} {}",self.end_tick - self.start_tick, self.end_tick, self.start_tick);
        let metadata =_state.metadata_mut();
        let hist = metadata.get_mut::<IcHist>();
        let timestamp = SystemTime::now().duration_since(unsafe {FUZZ_START_TIMESTAMP}).unwrap().as_millis();
        match hist {
            None => {
                metadata.insert(IcHist(vec![(self.end_tick - self.start_tick, timestamp)],
                (self.end_tick - self.start_tick, timestamp)));
            }
            Some(v) => {
                v.0.push((self.end_tick - self.start_tick, timestamp));
                if (v.1.0 < self.end_tick-self.start_tick) {
                    v.1 = (self.end_tick - self.start_tick, timestamp);
                }
                if v.0.len() >= 100 {
                    if let Ok(td) = env::var("TIME_DUMP") {
                        let mut file = OpenOptions::new()
                            .read(true)
                            .write(true)
                            .create(true)
                            .append(true)
                            .open(td).expect("Could not open timedump");
                        let newv : Vec<(u64, u128)> = Vec::with_capacity(100);
                        for i in std::mem::replace(&mut v.0, newv).into_iter() {
                            writeln!(file, "{},{}", i.0, i.1).expect("Write to dump failed");
                        }
                    } else {
                        // If we don't write out values we don't need to remember them at all
                        v.0.clear();
                    }
                }
            }
        }
        Ok(())
    }
}

impl Named for QemuClockObserver {
    #[inline]
    fn name(&self) -> &str {
        &self.name
    }
}

impl Default for QemuClockObserver {
    fn default() -> Self {
        Self {
            name: String::from("clock"),
            start_tick: 0,
            end_tick: 0,
        }
    }
}

//========== Feedback
/// Nop feedback that annotates execution time in the new testcase, if any
/// for this Feedback, the testcase is never interesting (use with an OR).
/// It decides, if the given [`QemuClockObserver`] value of a run is interesting.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ClockTimeFeedback {
    exec_time: Option<Duration>,
    name: String,
}

impl<S> Feedback<S> for ClockTimeFeedback
where
    S: UsesInput + HasClientPerfMonitor + HasMetadata,
{
    #[allow(clippy::wrong_self_convention)]
    fn is_interesting<EM, OT>(
        &mut self,
        _state: &mut S,
        _manager: &mut EM,
        _input: &S::Input,
        observers: &OT,
        _exit_kind: &ExitKind,
    ) -> Result<bool, Error>
    where
        EM: EventFirer<State = S>,
        OT: ObserversTuple<S>,
    {
        // TODO Replace with match_name_type when stable
        let observer = observers.match_name::<QemuClockObserver>(self.name()).unwrap();
        self.exec_time = Some(Duration::from_nanos(observer.last_runtime() << 4)); // Assume a somewhat realistic multiplier of clock, it does not matter
        Ok(false)
    }

    /// Append to the testcase the generated metadata in case of a new corpus item
    #[inline]
    fn append_metadata(
        &mut self,
        _state: &mut S,
        testcase: &mut Testcase<S::Input>,
    ) -> Result<(), Error> {
        *testcase.exec_time_mut() = self.exec_time;
        self.exec_time = None;
        Ok(())
    }

    /// Discard the stored metadata in case that the testcase is not added to the corpus
    #[inline]
    fn discard_metadata(&mut self, _state: &mut S, _input: &S::Input) -> Result<(), Error> {
        self.exec_time = None;
        Ok(())
    }
}

impl Named for ClockTimeFeedback {
    #[inline]
    fn name(&self) -> &str {
        self.name.as_str()
    }
}

impl ClockTimeFeedback {
    /// Creates a new [`ClockFeedback`], deciding if the value of a [`QemuClockObserver`] with the given `name` of a run is interesting.
    #[must_use]
    pub fn new(name: &'static str) -> Self {
        Self {
            exec_time: None,
            name: name.to_string(),
        }
    }

    /// Creates a new [`ClockFeedback`], deciding if the given [`QemuClockObserver`] value of a run is interesting.
    #[must_use]
    pub fn new_with_observer(observer: &QemuClockObserver) -> Self {
        Self {
            exec_time: None,
            name: observer.name().to_string(),
        }
    }
}

/// A [`Feedback`] rewarding increasing the execution cycles on Qemu.
#[derive(Debug)]
pub struct QemuClockIncreaseFeedback {
    name: String,
}

impl<S> Feedback<S> for QemuClockIncreaseFeedback
where
    S: UsesInput + HasNamedMetadata + HasClientPerfMonitor + Debug,
{
    fn is_interesting<EM, OT>(
        &mut self,
        state: &mut S,
        _manager: &mut EM,
        _input: &S::Input,
        _observers: &OT,
        _exit_kind: &ExitKind,
    ) -> Result<bool, Error>
    where
        EM: EventFirer<State = S>,
        OT: ObserversTuple<S>,
    {
        let observer = _observers.match_name::<QemuClockObserver>("clock")
            .expect("QemuClockObserver not found");
        let clock_state = state
            .named_metadata_mut()
            .get_mut::<MaxIcountMetadata>(&self.name)
            .unwrap();
        if observer.last_runtime() > clock_state.max_icount_seen {
            // println!("Clock improving {}",observer.last_runtime());
            clock_state.max_icount_seen = observer.last_runtime();
            return Ok(true);
        }
        Ok(false)
    }

    /// Append to the testcase the generated metadata in case of a new corpus item
    #[inline]
    fn append_metadata(&mut self, _state: &mut S, testcase: &mut Testcase<S::Input>) -> Result<(), Error> {
        // testcase.metadata_mut().insert(QemuIcountMetadata{runtime: self.last_runtime});
        Ok(())
    }

    /// Discard the stored metadata in case that the testcase is not added to the corpus
    #[inline]
    fn discard_metadata(&mut self, _state: &mut S, _input: &S::Input) -> Result<(), Error> {
        Ok(())
    }

}

impl Named for QemuClockIncreaseFeedback {
    #[inline]
    fn name(&self) -> &str {
        &self.name
    }
}

impl QemuClockIncreaseFeedback {
    /// Creates a new [`HitFeedback`]
    #[must_use]
    pub fn new(name: &'static str) -> Self {
        Self {name: String::from(name)}
    }
}

impl Default for QemuClockIncreaseFeedback {
    fn default() -> Self {
        Self::new("MaxClock")
    }
}