move snapshot code as well.

AFLplusplus · Aug 13, 2024 · 31ee26f · 31ee26f
1 parent 63f75da
commit 31ee26f
Show file tree

Hide file tree

Showing 6 changed files with 106 additions and 93 deletions.
diff --git a/accel/tcg/tcg-runtime.c b/accel/tcg/tcg-runtime.c
@@ -39,99 +39,6 @@
 
 #include "libafl/exit.h"
 
-#ifndef CONFIG_USER_ONLY
-
-#include "sysemu/runstate.h"
-#include "migration/snapshot.h"
-#include "qapi/error.h"
-#include "qemu/error-report.h"
-#include "qemu/main-loop.h"
-#include "hw/core/cpu.h"
-#include "sysemu/hw_accel.h"
-#include <stdlib.h>
-#include <string.h>
-
-void libafl_save_qemu_snapshot(char *name, bool sync);
-void libafl_load_qemu_snapshot(char *name, bool sync);
-
-static void save_snapshot_cb(void* opaque)
-{
-    char* name = (char*)opaque;
-    Error *err = NULL;
-    if(!save_snapshot(name, true, NULL, false, NULL, &err)) {
-        error_report_err(err);
-        error_report("Could not save snapshot");
-    }
-    free(opaque);
-}
-
-void libafl_save_qemu_snapshot(char *name, bool sync)
-{
-    // use snapshots synchronously, use if main loop is not running
-    if (sync) {
-        //TODO: eliminate this code duplication
-        //by passing a heap-allocated buffer from rust to c,
-        //which c needs to free
-        Error *err = NULL;
-        if(!save_snapshot(name, true, NULL, false, NULL, &err)) {
-            error_report_err(err);
-            error_report("Could not save snapshot");
-        }
-        return;
-    }
-    char* name_buffer = malloc(strlen(name)+1);
-    strcpy(name_buffer, name);
-    aio_bh_schedule_oneshot_full(qemu_get_aio_context(), save_snapshot_cb, (void*)name_buffer, "save_snapshot");
-}
-
-static void load_snapshot_cb(void* opaque)
-{
-    char* name = (char*)opaque;
-    Error *err = NULL;
-
-    int saved_vm_running = runstate_is_running();
-    vm_stop(RUN_STATE_RESTORE_VM);
-
-    bool loaded = load_snapshot(name, NULL, false, NULL, &err);
-
-    if(!loaded) {
-        error_report_err(err);
-        error_report("Could not load snapshot");
-    }
-    if (loaded && saved_vm_running) {
-        vm_start();
-    }
-    free(opaque);
-}
-
-void libafl_load_qemu_snapshot(char *name, bool sync)
-{
-    // use snapshots synchronously, use if main loop is not running
-    if (sync) {
-        //TODO: see libafl_save_qemu_snapshot
-        Error *err = NULL;
-
-        int saved_vm_running = runstate_is_running();
-        vm_stop(RUN_STATE_RESTORE_VM);
-
-        bool loaded = load_snapshot(name, NULL, false, NULL, &err);
-
-        if(!loaded) {
-            error_report_err(err);
-            error_report("Could not load snapshot");
-        }
-        if (loaded && saved_vm_running) {
-            vm_start();
-        }
-        return;
-    }
-    char* name_buffer = malloc(strlen(name)+1);
-    strcpy(name_buffer, name);
-    aio_bh_schedule_oneshot_full(qemu_get_aio_context(), load_snapshot_cb, (void*)name_buffer, "load_snapshot");
-}
-
-#endif
-
 void HELPER(libafl_qemu_handle_breakpoint)(CPUArchState *env, uint64_t pc)
 {
     CPUState* cpu = env_cpu(env);

diff --git a/include/libafl/cpu.h b/include/libafl/cpu.h
@@ -24,6 +24,8 @@ int libafl_qemu_num_regs(CPUState* cpu);
 void libafl_flush_jit(void);
 void libafl_breakpoint_invalidate(CPUState* cpu, target_ulong pc);
 
+#ifdef CONFIG_USER_ONLY
 int libafl_qemu_main(void);
 int libafl_qemu_run(void);
 void libafl_set_qemu_env(CPUArchState* env);
+#endif
diff --git a/include/libafl/qemu_snapshot.h b/include/libafl/qemu_snapshot.h
@@ -0,0 +1,6 @@
+#pragma once
+
+#include "qemu/osdep.h"
+
+void libafl_save_qemu_snapshot(char* name, bool sync);
+void libafl_load_qemu_snapshot(char* name, bool sync);
diff --git a/libafl/cpu.c b/libafl/cpu.c
@@ -1,6 +1,9 @@
 #include "qemu/osdep.h"
+
+#ifdef CONFIG_USER_ONLY
 #include "qemu.h"
 #include "user-internals.h"
+#endif
 
 #include "exec/gdbstub.h"
 #include "exec/cpu-defs.h"
@@ -16,7 +19,10 @@
 int gdb_write_register(CPUState* cpu, uint8_t* mem_buf, int reg);
 
 static __thread GByteArray* libafl_qemu_mem_buf = NULL;
+
+#ifdef CONFIG_USER_ONLY
 static __thread CPUArchState* libafl_qemu_env;
+#endif
 
 #ifndef CONFIG_USER_ONLY
 uint8_t* libafl_paddr2host(CPUState* cpu, hwaddr addr, bool is_write)
@@ -137,6 +143,7 @@ void libafl_flush_jit(void)
     CPU_FOREACH(cpu) { tb_flush(cpu); }
 }
 
+#ifdef CONFIG_USER_ONLY
 __attribute__((weak)) int libafl_qemu_main(void)
 {
     libafl_qemu_run();
@@ -150,3 +157,4 @@ int libafl_qemu_run(void)
 }
 
 void libafl_set_qemu_env(CPUArchState* env) { libafl_qemu_env = env; }
+#endif
diff --git a/libafl/meson.build b/libafl/meson.build
@@ -19,6 +19,7 @@ specific_ss.add(files(
                 ))
 
 specific_ss.add(when : 'CONFIG_SOFTMMU', if_true : [files(
+                                                        'qemu_snapshot.c',
                                                         'syx-snapshot/device-save.c',
                                                         'syx-snapshot/syx-snapshot.c',
                                                         'syx-snapshot/syx-cow-cache.c',

diff --git a/libafl/qemu_snapshot.c b/libafl/qemu_snapshot.c
@@ -0,0 +1,89 @@
+#include "libafl/qemu_snapshot.h"
+
+#include "sysemu/runstate.h"
+#include "migration/snapshot.h"
+#include "qapi/error.h"
+#include "qemu/error-report.h"
+#include "qemu/main-loop.h"
+#include "hw/core/cpu.h"
+#include "sysemu/hw_accel.h"
+#include <stdlib.h>
+#include <string.h>
+
+static void save_snapshot_cb(void* opaque)
+{
+    char* name = (char*)opaque;
+    Error* err = NULL;
+    if (!save_snapshot(name, true, NULL, false, NULL, &err)) {
+        error_report_err(err);
+        error_report("Could not save snapshot");
+    }
+    free(opaque);
+}
+
+static void load_snapshot_cb(void* opaque)
+{
+    char* name = (char*)opaque;
+    Error* err = NULL;
+
+    int saved_vm_running = runstate_is_running();
+    vm_stop(RUN_STATE_RESTORE_VM);
+
+    bool loaded = load_snapshot(name, NULL, false, NULL, &err);
+
+    if (!loaded) {
+        error_report_err(err);
+        error_report("Could not load snapshot");
+    }
+    if (loaded && saved_vm_running) {
+        vm_start();
+    }
+    free(opaque);
+}
+
+void libafl_save_qemu_snapshot(char* name, bool sync)
+{
+    // use snapshots synchronously, use if main loop is not running
+    if (sync) {
+        // TODO: eliminate this code duplication
+        // by passing a heap-allocated buffer from rust to c,
+        // which c needs to free
+        Error* err = NULL;
+        if (!save_snapshot(name, true, NULL, false, NULL, &err)) {
+            error_report_err(err);
+            error_report("Could not save snapshot");
+        }
+        return;
+    }
+    char* name_buffer = malloc(strlen(name) + 1);
+    strcpy(name_buffer, name);
+    aio_bh_schedule_oneshot_full(qemu_get_aio_context(), save_snapshot_cb,
+                                 (void*)name_buffer, "save_snapshot");
+}
+
+void libafl_load_qemu_snapshot(char* name, bool sync)
+{
+    // use snapshots synchronously, use if main loop is not running
+    if (sync) {
+        // TODO: see libafl_save_qemu_snapshot
+        Error* err = NULL;
+
+        int saved_vm_running = runstate_is_running();
+        vm_stop(RUN_STATE_RESTORE_VM);
+
+        bool loaded = load_snapshot(name, NULL, false, NULL, &err);
+
+        if (!loaded) {
+            error_report_err(err);
+            error_report("Could not load snapshot");
+        }
+        if (loaded && saved_vm_running) {
+            vm_start();
+        }
+        return;
+    }
+    char* name_buffer = malloc(strlen(name) + 1);
+    strcpy(name_buffer, name);
+    aio_bh_schedule_oneshot_full(qemu_get_aio_context(), load_snapshot_cb,
+                                 (void*)name_buffer, "load_snapshot");
+}