Syx snapshot fixes and tuning

accel/tcg/cputlb.c

@@ -47,6 +47,10 @@
 #endif
 #include "tcg/tcg-ldst.h"
 #include "tcg/oversized-guest.h"
+//// --- Begin LibAFL code ---
+//#define CONFIG_DEBUG_SYX
+#include "libafl/syx-snapshot/syx-snapshot.h"
+//// --- End LibAFL code ---
 
 /* DEBUG defines, enable DEBUG_TLB_LOG to log to the CPU_LOG_MMU target */
 /* #define DEBUG_TLB */
@@ -90,10 +94,12 @@ QEMU_BUILD_BUG_ON(NB_MMU_MODES > 16);
 #define ALL_MMUIDX_BITS ((1 << NB_MMU_MODES) - 1)
 
 //// --- Begin LibAFL code ---
-
-// void syx_snapshot_dirty_list_add(hwaddr paddr);
-void syx_snapshot_dirty_list_add_hostaddr(void* host_addr);
-
+// Use this snippet multiple times below
+#define SYX_SNAPSHOT_DIRTY_LIST_ADD_HOSTADDR_PROBE(dbg, access_type, addr, entry_full, phost) { \
+    if (access_type == MMU_DATA_STORE && !(flags & (TLB_MMIO | TLB_DISCARD_WRITE))) { \
+        SYX_DEBUG("%s %llx %llx\n", dbg, addr, addr+ (entry_full)->xlat_section); \
+        syx_snapshot_dirty_list_add_hostaddr((phost)); \
+    }}\
 //// --- End LibAFL code ---
 
 static inline size_t tlb_n_entries(CPUTLBDescFast *fast)
@@ -450,6 +456,15 @@ void tlb_flush_all_cpus_synced(CPUState *src_cpu)
     tlb_flush_by_mmuidx_all_cpus_synced(src_cpu, ALL_MMUIDX_BITS);
 }
 
+void tlb_flush_all_cpus(void) 
+{
+    const run_on_cpu_func fn = tlb_flush_by_mmuidx_async_work;
+
+    tlb_debug("mmu_idx: 0x%"PRIx16"\n", ALL_MMUIDX_BITS);
+
+    flush_all_helper(NULL, fn, RUN_ON_CPU_HOST_INT(ALL_MMUIDX_BITS));
+}
+
 static bool tlb_hit_page_mask_anyprot(CPUTLBEntry *tlb_entry,
                                       vaddr page, vaddr mask)
 {
@@ -1406,13 +1421,7 @@ static int probe_access_internal(CPUState *cpu, vaddr addr,
 
     /* Everything else is RAM. */
     *phost = (void *)((uintptr_t)addr + entry->addend);
-//// --- Begin LibAFL code ---
 
-    if (access_type == MMU_DATA_STORE) {
-        syx_snapshot_dirty_list_add_hostaddr(*phost);
-    }
-
-//// --- End LibAFL code ---
     return flags;
 }
 
@@ -1430,6 +1439,10 @@ int probe_access_full(CPUArchState *env, vaddr addr, int size,
         int dirtysize = size == 0 ? 1 : size;
         notdirty_write(env_cpu(env), addr, dirtysize, *pfull, retaddr);
         flags &= ~TLB_NOTDIRTY;
+
+        //// --- Begin LibAFL code ---
+        SYX_SNAPSHOT_DIRTY_LIST_ADD_HOSTADDR_PROBE("probe_access_full", access_type, addr, *pfull, *phost);
+        //// --- End LibAFL code ---
     }
 
     return flags;
@@ -1454,6 +1467,10 @@ int probe_access_full_mmu(CPUArchState *env, vaddr addr, int size,
         int dirtysize = size == 0 ? 1 : size;
         notdirty_write(env_cpu(env), addr, dirtysize, *pfull, 0);
         flags &= ~TLB_NOTDIRTY;
+
+        //// --- Begin LibAFL code ---
+        SYX_SNAPSHOT_DIRTY_LIST_ADD_HOSTADDR_PROBE("probe_access_full_mmu", access_type, addr, *pfull, *phost);
+        //// --- End LibAFL code ---
     }
 
     return flags;
@@ -1477,6 +1494,10 @@ int probe_access_flags(CPUArchState *env, vaddr addr, int size,
         int dirtysize = size == 0 ? 1 : size;
         notdirty_write(env_cpu(env), addr, dirtysize, full, retaddr);
         flags &= ~TLB_NOTDIRTY;
+
+        //// --- Begin LibAFL code ---
+        SYX_SNAPSHOT_DIRTY_LIST_ADD_HOSTADDR_PROBE("probe_access_flags", access_type, addr, full, *phost);
+        //// --- End LibAFL code ---
     }
 
     return flags;
@@ -1512,6 +1533,10 @@ void *probe_access(CPUArchState *env, vaddr addr, int size,
         /* Handle clean RAM pages.  */
         if (flags & TLB_NOTDIRTY) {
             notdirty_write(env_cpu(env), addr, size, full, retaddr);
+
+            //// --- Begin LibAFL code ---
+            SYX_SNAPSHOT_DIRTY_LIST_ADD_HOSTADDR_PROBE("probe_access", access_type, addr, full, host);
+            //// --- End LibAFL code ---
         }
     }
 
@@ -1728,6 +1753,13 @@ static void mmu_watch_or_dirty(CPUState *cpu, MMULookupPageData *data,
     if (flags & TLB_NOTDIRTY) {
         notdirty_write(cpu, addr, size, full, ra);
         flags &= ~TLB_NOTDIRTY;
+
+        //// --- Begin LibAFL code ---
+        if (!(flags & (TLB_MMIO | TLB_DISCARD_WRITE))) {
+            SYX_DEBUG("mmu_watch_or_dirty %llx %llx\n", addr, addr+full->xlat_section);
+            syx_snapshot_dirty_list_add_hostaddr(data->haddr);
+        }
+        //// --- End LibAFL code ---
     }
     data->flags = flags;
 }
@@ -1772,14 +1804,6 @@ static bool mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
             l->memop ^= MO_BSWAP;
         }
 
-        //// --- Begin LibAFL code ---
-
-        // TODO: Does not work?
-        // if (type == MMU_DATA_STORE) {
-            syx_snapshot_dirty_list_add_hostaddr(l->page[0].haddr);
-        // }
-
-        //// --- End LibAFL code ---
 
     } else {
         /* Finish compute of page crossing. */
@@ -1803,15 +1827,6 @@ static bool mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
             mmu_watch_or_dirty(cpu, &l->page[1], type, ra);
         }
 
-        //// --- Begin LibAFL code ---
-
-        // if (type == MMU_DATA_STORE) {
-            syx_snapshot_dirty_list_add_hostaddr(l->page[0].haddr);
-            syx_snapshot_dirty_list_add_hostaddr(l->page[1].haddr);
-        // }
-
-        //// --- End LibAFL code ---
-
         /*
          * Since target/sparc is the only user of TLB_BSWAP, and all
          * Sparc accesses are aligned, any treatment across two pages
@@ -1907,14 +1922,12 @@ static void *atomic_mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
     hostaddr = (void *)((uintptr_t)addr + tlbe->addend);
     full = &cpu->neg.tlb.d[mmu_idx].fulltlb[index];
 
-    //// --- Begin LibAFL code ---
-
-    syx_snapshot_dirty_list_add_hostaddr(hostaddr);
-
-    //// --- End LibAFL code ---
-
     if (unlikely(tlb_addr & TLB_NOTDIRTY)) {
         notdirty_write(cpu, addr, size, full, retaddr);
+        //// --- Begin LibAFL code ---
+        SYX_DEBUG("atomic_mmu_lookup %llx %llx\n", addr, addr+full->xlat_section);
+        syx_snapshot_dirty_list_add_hostaddr(hostaddr);
+        //// --- End LibAFL code ---
     }
 
     if (unlikely(tlb_addr & TLB_FORCE_SLOW)) {

block/block-backend.c

@@ -29,9 +29,7 @@
 #include "migration/misc.h"
 
 //// --- Begin LibAFL code ---
-#ifdef CONFIG_SOFTMMU
 #include "libafl/syx-snapshot/syx-snapshot.h"
-#endif
 //// --- End LibAFL code ---
 
 /* Number of coroutines to reserve per attached device model */
@@ -1650,18 +1648,14 @@ static void coroutine_fn blk_aio_read_entry(void *opaque)
     assert(qiov->size == acb->bytes);
 
 //// --- Begin LibAFL code ---
-#ifdef CONFIG_SOFTMMU
     if (!syx_snapshot_cow_cache_read_entry(rwco->blk, rwco->offset, acb->bytes, qiov, 0, rwco->flags)) {
-#endif
 //// --- End LibAFL code ---
-       rwco->ret = blk_co_do_preadv_part(rwco->blk, rwco->offset, acb->bytes, qiov,
+        rwco->ret = blk_co_do_preadv_part(rwco->blk, rwco->offset, acb->bytes, qiov,
                                       0, rwco->flags);
 //// --- Begin LibAFL code ---
-#ifdef CONFIG_SOFTMMU
     } else {
-       rwco->ret = 0;
+        rwco->ret = 0;
     }
-#endif
 //// --- End LibAFL code ---
 
     blk_aio_complete(acb);
@@ -1676,17 +1670,14 @@ static void coroutine_fn blk_aio_write_entry(void *opaque)
     assert(!qiov || qiov->size == acb->bytes);
 
 //// --- Begin LibAFL code ---
-#ifdef CONFIG_SOFTMMU
-   if (!syx_snapshot_cow_cache_write_entry(rwco->blk, rwco->offset, acb->bytes, qiov, 0, rwco->flags)) {
-#endif
+    if (!syx_snapshot_cow_cache_write_entry(rwco->blk, rwco->offset, acb->bytes, qiov, 0, rwco->flags)) {
 //// --- End LibAFL code ---
         rwco->ret = blk_co_do_pwritev_part(rwco->blk, rwco->offset, acb->bytes, qiov, 0, rwco->flags);
 //// --- Begin LibAFL code ---
-#ifdef CONFIG_SOFTMMU
-   } else {
-       rwco->ret = 0;
-   }
-#endif
+    } else {
+        SYX_DEBUG("Write to COW cache: 0x%llx 0x%llx\n", rwco->offset, acb->bytes);
+        rwco->ret = 0;
+    }
 //// --- End LibAFL code ---
 
     blk_aio_complete(acb);

cpu-target.c

@@ -73,13 +73,9 @@ static int cpu_common_post_load(void *opaque, int version_id)
     //tb_flush(cpu);
 
 //// --- Begin LibAFL code ---
-
-    // flushing the TBs every restore makes it really slow
-    // TODO handle writes to X code with specific calls to tb_invalidate_phys_addr
-    if (!libafl_devices_is_restoring()) {
-        tb_flush(cpu);
-    }
-
+    // Only invalidate per CPU virtual JMP cache
+    // Note: Global TB cache will be invalidated by SYX snapshot code
+    tcg_flush_jmp_cache(cpu);
 //// --- End LibAFL code ---
 
     return 0;

hmp-commands.hx

@@ -1859,3 +1859,43 @@ SRST
   List event channels in the guest
 ERST
 #endif
+
+
+    {
+        .name       = "syx-snapshot-new",
+        .args_type  = "",
+        .params     = "",
+        .help       = "create a new snapshot and store it in the static variable. Use syx-snapshot-init BEFORE.",
+        .cmd        = hmp_syx_snapshot_new,
+    },
+
+SRST
+``syx-snapshot-new``
+  Create a new snapshot with tracking enabled, using DEVICE_SNAPSHOT_ALL.
+ERST
+
+    {
+        .name       = "syx-snapshot-root-restore",
+        .args_type  = "",
+        .params     = "",
+        .help       = "restore the root snapshot from the static variable",
+        .cmd        = hmp_syx_snapshot_root_restore,
+    },
+
+SRST
+``syx-snapshot-root-restore``
+  Restore the root snapshot stored in the static variable.
+ERST
+
+    {
+        .name       = "syx-snapshot-init",
+        .args_type  = "",
+        .params     = "",
+        .help       = "Use BEFORE starting emulation",
+        .cmd        = hmp_syx_snapshot_init,
+    },
+
+SRST
+``syx-snapshot-identified``
+  Init syx.
+ERST

include/exec/exec-all.h

@@ -110,6 +110,9 @@ void tlb_flush_all_cpus_synced(CPUState *src_cpu);
  */
 void tlb_flush_page_by_mmuidx(CPUState *cpu, vaddr addr,
                               uint16_t idxmap);
+
+void tlb_flush_all_cpus(void);
+
 /**
  * tlb_flush_page_by_mmuidx_all_cpus_synced:
  * @cpu: Originating CPU of the flush

include/exec/ram_addr.h

@@ -325,6 +325,7 @@ static inline void cpu_physical_memory_set_dirty_range(ram_addr_t start,
                                   offset, next - page);
             }
 
+
             page = next;
             idx++;
             offset = 0;

include/libafl/system.h

@@ -4,8 +4,15 @@
 #include "gdbstub/enums.h"
 #include "sysemu/accel-ops.h"
 #include "sysemu/cpus.h"
+#include "sysemu/block-backend.h"
 
 int libafl_qemu_set_hw_breakpoint(vaddr addr);
 int libafl_qemu_remove_hw_breakpoint(vaddr addr);
 
 void libafl_qemu_init(int argc, char** argv);
+
+/** Write to a block device with aio API
+ * The same way the guest would, 
+ * thus this writes to the Syx COW cache (if it is initialized)
+ */
+int libafl_blk_write(BlockBackend *blk, void *buf, int64_t offset, int64_t sz);

include/monitor/hmp.h

@@ -179,4 +179,8 @@ void hmp_info_mtree(Monitor *mon, const QDict *qdict);
 void hmp_info_cryptodev(Monitor *mon, const QDict *qdict);
 void hmp_dumpdtb(Monitor *mon, const QDict *qdict);
 
+void hmp_syx_snapshot_new(Monitor *mon, const QDict *qdict);
+void hmp_syx_snapshot_root_restore(Monitor *mon, const QDict *qdict);
+void hmp_syx_snapshot_init(Monitor *mon, const QDict *qdict);
+
 #endif

libafl/blkdev.c

@@ -0,0 +1,40 @@
+
+#include "qemu/osdep.h"
+
+#include "qapi/error.h"
+#include "qapi/qmp/qdict.h"
+#include "qemu/option.h"
+#include "qemu/main-loop.h"
+#include "block/qdict.h"
+#include "libafl/system.h"
+
+
+
+#define NOT_DONE 0x7fffffff
+
+static void blk_rw_done(void *opaque, int ret)
+{
+    *(int *)opaque = ret;
+}
+
+int libafl_blk_write(BlockBackend *blk, void *buf, int64_t offset, int64_t sz)
+{
+	void *pattern_buf = NULL;
+	QEMUIOVector qiov;
+	int async_ret = NOT_DONE;
+
+	qemu_iovec_init(&qiov, 1);
+	qemu_iovec_add(&qiov, buf, sz);
+
+	blk_aio_pwritev(blk, offset, &qiov, 0, blk_rw_done, &async_ret);
+	while (async_ret == NOT_DONE) {
+		main_loop_wait(false);
+	}
+
+	//printf("async_ret: %d\n", async_ret);
+	//g_assert(async_ret == 0);
+
+	g_free(pattern_buf);
+	qemu_iovec_destroy(&qiov);
+    return async_ret;
+}

libafl/meson.build

@@ -22,10 +22,12 @@ specific_ss.add(files(
 specific_ss.add(when : 'CONFIG_SOFTMMU', if_true : [files(
                                                         'system.c',
                                                         'qemu_snapshot.c',
+                                                        'blkdev.c',
                                                         'syx-snapshot/device-save.c',
                                                         'syx-snapshot/syx-snapshot.c',
                                                         'syx-snapshot/syx-cow-cache.c',
                                                         'syx-snapshot/channel-buffer-writeback.c',
+                                                        'syx-snapshot/syx-hmp.c',
                                                     )])
 
 specific_ss.add(when : 'CONFIG_USER_ONLY', if_true : [files(

libafl/syx-snapshot/syx-cow-cache.c

@@ -76,6 +76,7 @@ static bool read_chunk_from_cache_layer_device(SyxCowCacheDevice* sccd,
 
     // cache hit
     if (found) {
+        printf("[SYX] cached chunk found: %llx %lx\n", blk_offset, g_array_get_element_size(sccd->data));
         void* data_position_ptr =
             g_array_element_ptr(sccd->data, GPOINTER_TO_UINT(data_position));
         assert(qemu_iovec_from_buf(qiov, qiov_offset, data_position_ptr,
@@ -217,9 +218,6 @@ void syx_cow_cache_read_entry(SyxCowCache* scc, BlockBackend* blk,
     size_t qiov_offset = 0;
     uint64_t chunk_size = 0;
 
-    // printf("[%s] Read 0x%zx bytes @addr %lx\n", blk_name(blk), qiov->size,
-    // offset);
-
     // First read the backing block device normally.
     assert(blk_co_preadv(blk, offset, bytes, qiov, flags) >= 0);