FRET-LibAFL/libafl_qemu/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,
};
use serde::{Deserialize, Serialize};
use std::{cell::UnsafeCell, cmp::max};
use libafl::bolts::tuples::Named;

use crate::{
    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;

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

/// 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")
    }
}

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

/// A simple observer, just overlooking the runtime of the target.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct QemuClockObserver {
    name: String,
    last_runtime: 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(),
            last_runtime: 0,
        }
    }

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

impl<S> Observer<S> for QemuClockObserver
where
    S: UsesInput,
{
    fn pre_exec(&mut self, _state: &mut S, _input: &S::Input) -> Result<(), Error> {
        self.last_runtime=0;
        Ok(())
    }

    fn post_exec(&mut self, _state: &mut S, _input: &S::Input, _exit_kind: &ExitKind) -> Result<(), Error> {
        unsafe { self.last_runtime = emu::icount_get_raw() };
        Ok(())
    }
}

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

impl Default for QemuClockObserver {
    fn default() -> Self {
        Self {
            name: String::from("clock"),
            last_runtime: 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)
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ClockFeedback {
    exec_time: Option<u64>,
    name: String,
}

impl<S> Feedback<S> for ClockFeedback
where
    S: UsesInput + HasClientPerfMonitor,
{
    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(observer.last_runtime());
        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() = match self.exec_time {
            Some(s) => Some(Duration::from_nanos(s << 3)),    //emulated time is << 3, real time more like * 360
            None => None,
        };
        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 ClockFeedback {
    #[inline]
    fn name(&self) -> &str {
        self.name.as_str()
    }
}

impl ClockFeedback {
    /// Creates a new [`ClockFeedback`], deciding if the value of a [`TimeObserver`] 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 [`TimeObserver`] 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")
    }
}