FRET-LibAFL/fuzzers/FRET/src/systemstate/helpers.rs

use std::cell::UnsafeCell;
use std::io::Write;
use std::ops::Range;
use libafl::prelude::UsesInput;
use libafl_qemu::Emulator;
use libafl_qemu::GuestAddr;
use libafl_qemu::QemuHooks;
use libafl_qemu::edges::QemuEdgesMapMetadata;
use libafl_qemu::emu;
use libafl_qemu::hooks;
use crate::systemstate::RawFreeRTOSSystemState;
use crate::systemstate::CURRENT_SYSTEMSTATE_VEC;
use crate::systemstate::NUM_PRIOS;
use super::freertos::TCB_t;
use super::freertos::rtos_struct::List_Item_struct;
use super::freertos::rtos_struct::*;
use super::freertos;

use libafl_qemu::{
    helper::{QemuHelper, QemuHelperTuple},
    // edges::SAVED_JUMP,
};

//============================= Struct definitions

pub static mut INTR_OFFSET : Option<u64> = None;
pub static mut INTR_DONE : bool = true;

// only used when inputs are injected
pub static mut NEXT_INPUT : Vec<u8> = Vec::new();

//============================= Qemu Helper

/// A Qemu Helper with reads FreeRTOS specific structs from Qemu whenever certain syscalls occur, also inject inputs
#[derive(Debug)]
pub struct QemuSystemStateHelper {
    kerneladdr: u32,
    tcb_addr: u32,
    ready_queues: u32,
    input_counter: Option<u64>,
    app_range: Range<u32>,
}

impl QemuSystemStateHelper {
    #[must_use]
    pub fn new(
        kerneladdr: u32,
        tcb_addr: u32,
        ready_queues: u32,
        input_counter: Option<u64>,
        app_range: Range<u32>,
    ) -> Self {
        QemuSystemStateHelper {
            kerneladdr,
            tcb_addr: tcb_addr,
            ready_queues: ready_queues,
            input_counter: input_counter,
            app_range,
        }
    }
}

impl<S> QemuHelper<S> for QemuSystemStateHelper
where
    S: UsesInput,
{
    fn first_exec<QT>(&self, _hooks: &QemuHooks<'_, QT, S>)
    where
        QT: QemuHelperTuple<S>,
    {
        _hooks.instruction(self.kerneladdr, exec_syscall_hook::<QT, S>, false);
        _hooks.jmps(Some(gen_jmp_is_syscall::<QT, S>), Some(trace_api_call::<QT, S>));
    }

    // TODO: refactor duplicate code
    fn pre_exec(&mut self, _emulator: &Emulator, _input: &S::Input) {
        unsafe {
            CURRENT_SYSTEMSTATE_VEC.clear();
            let p = LAST_API_CALL.with(|x| x.get());
            *p = None;
        }
    }

    fn post_exec(&mut self, emulator: &Emulator, _input: &S::Input) {
        trigger_collection(emulator, self)
    }
}

#[inline]
fn trigger_collection(emulator: &Emulator, h: &QemuSystemStateHelper) {
    let listbytes : u32 = u32::try_from(std::mem::size_of::<freertos::List_t>()).unwrap();
    let mut systemstate = RawFreeRTOSSystemState::default();
    unsafe {
        // TODO: investigate why can_do_io is not set sometimes, as this is just a workaround
        let c = emulator.current_cpu().unwrap();
        let can_do_io = (*c.raw_ptr()).can_do_io;
        (*c.raw_ptr()).can_do_io = 1;
        systemstate.qemu_tick = emu::icount_get_raw();
        (*c.raw_ptr()).can_do_io = can_do_io;
    }
    let mut buf : [u8; 4] = [0,0,0,0];
    match h.input_counter {
        Some(s) => unsafe { emulator.read_phys_mem(s, &mut buf); },
        None => (),
    };
    systemstate.input_counter = u32::from_le_bytes(buf);

    let curr_tcb_addr : freertos::void_ptr = freertos::emu_lookup::lookup(emulator, h.tcb_addr);
    if curr_tcb_addr == 0 {
        return;
    };
    systemstate.current_tcb = freertos::emu_lookup::lookup(emulator,curr_tcb_addr);

    unsafe {
        LAST_API_CALL.with(|x|
            match *x.get() {
                Some(s) => {
                    systemstate.last_pc = Some(s.0 as u64);
                },
                None => (),
            }
        );
    }
    // println!("{:?}",std::str::from_utf8(&current_tcb.pcTaskName));

    for i in 0..NUM_PRIOS {
        let target : u32 = listbytes*u32::try_from(i).unwrap()+h.ready_queues;
        systemstate.prio_ready_lists[i] = freertos::emu_lookup::lookup(emulator, target);
        // println!("List at {}: {:?}",target, systemstate.prio_ready_lists[i]);
        let mut next_index = systemstate.prio_ready_lists[i].pxIndex;
        for _j in 0..systemstate.prio_ready_lists[i].uxNumberOfItems {
            // always jump over the xListEnd marker
            if (target..target+listbytes).contains(&next_index) {
                let next_item : freertos::MiniListItem_t = freertos::emu_lookup::lookup(emulator, next_index);
                let new_next_index=next_item.pxNext;
                systemstate.dumping_ground.insert(next_index,List_MiniItem_struct(next_item));
                next_index = new_next_index;
            }
            let next_item : freertos::ListItem_t = freertos::emu_lookup::lookup(emulator, next_index);
            // println!("Item at {}: {:?}",next_index,next_item);
            assert_eq!(next_item.pvContainer,target);
            let new_next_index=next_item.pxNext;
            let next_tcb : TCB_t= freertos::emu_lookup::lookup(emulator,next_item.pvOwner);
            // println!("TCB at {}: {:?}",next_item.pvOwner,next_tcb);
            systemstate.dumping_ground.insert(next_item.pvOwner,TCB_struct(next_tcb.clone()));
            systemstate.dumping_ground.insert(next_index,List_Item_struct(next_item));
            next_index=new_next_index;
        }
        // Handle edge case where the end marker was not included yet
        if (target..target+listbytes).contains(&next_index) {
            let next_item : freertos::MiniListItem_t = freertos::emu_lookup::lookup(emulator, next_index);
            systemstate.dumping_ground.insert(next_index,List_MiniItem_struct(next_item));
        }
    }

    unsafe { CURRENT_SYSTEMSTATE_VEC.push(systemstate); }
}

pub fn exec_syscall_hook<QT, S>(
    hooks: &mut QemuHooks<'_, QT, S>,
    _state: Option<&mut S>,
    _pc: u32,
)
where
    S: UsesInput,
    QT: QemuHelperTuple<S>,
{
    let emulator = hooks.emulator();
    let h = hooks.helpers().match_first_type::<QemuSystemStateHelper>().expect("QemuSystemHelper not found in helper tupel");
    trigger_collection(emulator, h);
}

thread_local!(static LAST_API_CALL : UnsafeCell<Option<(GuestAddr,GuestAddr)>> = UnsafeCell::new(None));

pub fn gen_jmp_is_syscall<QT, S>(
    hooks: &mut QemuHooks<'_, QT, S>,
    _state: Option<&mut S>,
    src: GuestAddr,
    dest: GuestAddr,
) -> Option<u64>
where
    S: UsesInput,
    QT: QemuHelperTuple<S>,
{
    if let Some(h) = hooks.helpers().match_first_type::<QemuSystemStateHelper>() {
        if h.app_range.contains(&src) && !h.app_range.contains(&dest) {
            // println!("New jmp {:x} {:x}", src, dest);
            return Some(1);
        }
    }
    return None;
}

pub fn trace_api_call<QT, S>(
    _hooks: &mut QemuHooks<'_, QT, S>,
    _state: Option<&mut S>,
    src: GuestAddr, dest: GuestAddr, id: u64
)
where
    S: UsesInput,
    QT: QemuHelperTuple<S>,
{
    unsafe {
        let p = LAST_API_CALL.with(|x| x.get());
        *p = Some((src,dest));
        // print!("*");
    }
}