From 71a51a83f346532c7f6dbac11809b8001af9818c Mon Sep 17 00:00:00 2001
From: Romain Malmain <romain.malmain@pm.me>
Date: Tue, 19 Dec 2023 17:03:29 +0100
Subject: [PATCH] Added internal QEMU exit.

---
 libafl_extras/exit.c | 11 +++++++++++
 libafl_extras/exit.h | 15 +++++++++++++--
 system/runstate.c    | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/libafl_extras/exit.c b/libafl_extras/exit.c
index 9349480bdd..b76550cf26 100644
--- a/libafl_extras/exit.c
+++ b/libafl_extras/exit.c
@@ -97,6 +97,17 @@ CPUState* libafl_last_exit_cpu(void)
     return NULL;
 }
 
+void libafl_exit_request_internal(CPUState* cpu, uint64_t pc, ShutdownCause cause, int signal)
+{
+    last_exit_reason.kind = INTERNAL;
+    last_exit_reason.data.internal.cause = cause;
+    last_exit_reason.data.internal.signal = signal;
+
+    last_exit_reason.cpu = cpu;
+    last_exit_reason.next_pc = pc;
+    expected_exit = true;
+}
+
 void libafl_exit_request_sync_backdoor(CPUState* cpu, target_ulong pc)
 {
     last_exit_reason.kind = SYNC_BACKDOOR;
diff --git a/libafl_extras/exit.h b/libafl_extras/exit.h
index 290a8e3408..b5e647bff2 100644
--- a/libafl_extras/exit.h
+++ b/libafl_extras/exit.h
@@ -19,21 +19,31 @@ int libafl_qemu_set_breakpoint(target_ulong pc);
 int libafl_qemu_remove_breakpoint(target_ulong pc);
 
 enum libafl_exit_reason_kind {
-    BREAKPOINT = 0,
-    SYNC_BACKDOOR = 1,
+    INTERNAL = 0,
+    BREAKPOINT = 1,
+    SYNC_BACKDOOR = 2,
 };
 
+// A breakpoint has been triggered.
 struct libafl_exit_reason_breakpoint {
     target_ulong addr;
 };
 
+// A synchronous exit has been triggered.
 struct libafl_exit_reason_sync_backdoor { };
 
+// QEMU exited on its own for some reason.
+struct libafl_exit_reason_internal {
+    ShutdownCause cause;
+    int signal; // valid if cause == SHUTDOWN_CAUSE_HOST_SIGNAL
+};
+
 struct libafl_exit_reason {
     enum libafl_exit_reason_kind kind;
     CPUState* cpu; // CPU that triggered an exit.
     vaddr next_pc; // The PC that should be stored in the CPU when re-entering.
     union {
+        struct libafl_exit_reason_internal internal;
         struct libafl_exit_reason_breakpoint breakpoint; // kind == BREAKPOINT
         struct libafl_exit_reason_sync_backdoor backdoor; // kind == SYNC_BACKDOOR
     } data;
@@ -47,6 +57,7 @@ void libafl_exit_signal_vm_start(void);
 bool libafl_exit_asap(void);
 void libafl_sync_exit_cpu(void);
 
+void libafl_exit_request_internal(CPUState* cpu, uint64_t pc, ShutdownCause cause, int signal);
 void libafl_exit_request_sync_backdoor(CPUState* cpu, target_ulong pc);
 void libafl_exit_request_breakpoint(CPUState* cpu, target_ulong pc);
 struct libafl_exit_reason* libafl_get_exit_reason(void);
diff --git a/system/runstate.c b/system/runstate.c
index eb962f40e1..2bc3013319 100644
--- a/system/runstate.c
+++ b/system/runstate.c
@@ -70,6 +70,10 @@ static RunState current_run_state = RUN_STATE_PRELAUNCH;
 static RunState vmstop_requested = RUN_STATE__MAX;
 static QemuMutex vmstop_lock;
 
+//// --- Begin LibAFL code ---
+void libafl_exit_request_internal(CPUState* cpu, uint64_t pc, ShutdownCause cause, int signal);
+//// --- End LibAFL code ---
+
 typedef struct {
     RunState from;
     RunState to;
@@ -582,6 +586,16 @@ void qemu_system_reset_request(ShutdownCause reason)
     } else {
         reset_requested = reason;
     }
+
+//// --- Begin LibAFL code ---
+    if (current_cpu) {
+        CPUClass *cc = CPU_GET_CLASS(current_cpu);
+        libafl_exit_request_internal(current_cpu, cc->get_pc(current_cpu), shutdown_requested, -1);
+    } else {
+        libafl_exit_request_internal(NULL, 0, shutdown_requested, -1);
+    }
+//// --- End LibAFL code ---
+
     cpu_stop_current();
     qemu_notify_event();
 }
@@ -662,6 +676,16 @@ void qemu_system_killed(int signal, pid_t pid)
      * we are in a signal handler.
      */
     shutdown_requested = SHUTDOWN_CAUSE_HOST_SIGNAL;
+
+//// --- Begin LibAFL code ---
+    if (current_cpu) {
+        CPUClass *cc = CPU_GET_CLASS(current_cpu);
+        libafl_exit_request_internal(current_cpu, cc->get_pc(current_cpu), shutdown_requested, signal);
+    } else {
+        libafl_exit_request_internal(NULL, 0, shutdown_requested, signal);
+    }
+//// --- End LibAFL code ---
+
     qemu_notify_event();
 }
 
@@ -677,6 +701,16 @@ void qemu_system_shutdown_request(ShutdownCause reason)
     trace_qemu_system_shutdown_request(reason);
     replay_shutdown_request(reason);
     shutdown_requested = reason;
+
+//// --- Begin LibAFL code ---
+    if (current_cpu) {
+        CPUClass *cc = CPU_GET_CLASS(current_cpu);
+        libafl_exit_request_internal(current_cpu, cc->get_pc(current_cpu), shutdown_requested, -1);
+    } else {
+        libafl_exit_request_internal(NULL, 0, shutdown_requested, -1);
+    }
+//// --- End LibAFL code ---
+
     qemu_notify_event();
 }