From c748fecbe26e96101cdcc1e093e7137d5928819d Mon Sep 17 00:00:00 2001
From: Alwin Berger <alwin.berger@tu-dortmund.de>
Date: Mon, 19 Dec 2022 13:13:38 +0100
Subject: [PATCH] add last api callsite to system state

---
 fuzzers/FRET/src/fuzzer.rs              |  9 +++-
 fuzzers/FRET/src/systemstate/helpers.rs | 63 +++++++++++++++++++++----
 2 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/fuzzers/FRET/src/fuzzer.rs b/fuzzers/FRET/src/fuzzer.rs
index f4d5ed9fab..ce1f1447a9 100644
--- a/fuzzers/FRET/src/fuzzer.rs
+++ b/fuzzers/FRET/src/fuzzer.rs
@@ -103,6 +103,13 @@ pub fn fuzz() {
     // let svh = elf
     //     .resolve_symbol("vPortEnterCritical", 0)
     //     .expect("Symbol vPortEnterCritical not found");
+    let app_start = elf
+        .resolve_symbol("__APP_CODE_START__", 0)
+        .expect("Symbol __APP_CODE_START__ not found");
+    let app_end = elf
+        .resolve_symbol("__APP_CODE_END__", 0)
+        .expect("Symbol __APP_CODE_END__ not found");
+    let app_range = app_start..app_end;
 
     let breakpoint = elf
         .resolve_symbol(
@@ -204,7 +211,7 @@ pub fn fuzz() {
         let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective);
         let mut hooks = QemuHooks::new(&emu,
             tuple_list!(QemuEdgeCoverageHelper::default(),QemuStateRestoreHelper::new(),
-        QemuSystemStateHelper::new(svh,curr_tcb_pointer,task_queue_addr,0)));
+        QemuSystemStateHelper::new(svh,curr_tcb_pointer,task_queue_addr,0,app_range)));
 
         // Create a QEMU in-process executor
         let executor = QemuExecutor::new(
diff --git a/fuzzers/FRET/src/systemstate/helpers.rs b/fuzzers/FRET/src/systemstate/helpers.rs
index 55085f93c6..e749592c10 100644
--- a/fuzzers/FRET/src/systemstate/helpers.rs
+++ b/fuzzers/FRET/src/systemstate/helpers.rs
@@ -1,4 +1,6 @@
+use std::cell::UnsafeCell;
 use std::io::Write;
+use std::ops::Range;
 use libafl::prelude::UsesInput;
 use libafl_qemu::GuestAddr;
 use libafl_qemu::QemuHooks;
@@ -35,6 +37,7 @@ pub struct QemuSystemStateHelper {
     tcb_addr: u32,
     ready_queues: u32,
     input_counter: u32,
+    app_range: Range<u32>,
 }
 
 impl QemuSystemStateHelper {
@@ -44,12 +47,14 @@ impl QemuSystemStateHelper {
         tcb_addr: u32,
         ready_queues: u32,
         input_counter: u32,
+        app_range: Range<u32>,
     ) -> Self {
         QemuSystemStateHelper {
             kerneladdr,
             tcb_addr: tcb_addr,
             ready_queues: ready_queues,
             input_counter: input_counter,
+            app_range,
         }
     }
 }
@@ -62,7 +67,8 @@ where
     where
         QT: QemuHelperTuple<S>,
     {
-        _hooks.instruction(self.kerneladdr, exec_syscall_hook::<QT, S>, false)
+        _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>));
     }
 }
 
@@ -97,14 +103,16 @@ where
     };
     systemstate.current_tcb = freertos::emu_lookup::lookup(emulator,curr_tcb_addr);
 
-    // unsafe {
-    //     match SAVED_JUMP.take() {
-    //         Some(s) => {
-    //             systemstate.last_pc = Some(s.0);
-    //         },
-    //         None => (),
-    //     }
-    // }
+    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 {
@@ -138,4 +146,41 @@ where
     }
 
     unsafe { CURRENT_SYSTEMSTATE_VEC.push(systemstate); }
+}
+
+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!("*");
+    }
 }
\ No newline at end of file