From 74621995104387133e6cab724e78cca4229bdb9d Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Thu, 6 Mar 2025 15:49:25 +0100
Subject: [PATCH 01/12] [add] HMP command for Syx snapshot

---
 hmp-commands.hx                        | 40 ++++++++++++++++++++++++++
 include/monitor/hmp.h                  |  4 +++
 libafl/meson.build                     |  1 +
 libafl/syx-snapshot/syx-snapshot-hmp.c | 34 ++++++++++++++++++++++
 4 files changed, 79 insertions(+)
 create mode 100644 libafl/syx-snapshot/syx-snapshot-hmp.c

diff --git a/hmp-commands.hx b/hmp-commands.hx
index 06746f0afc..56ff0f785d 100644
--- a/hmp-commands.hx
+++ b/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",
+        .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
\ No newline at end of file
diff --git a/include/monitor/hmp.h b/include/monitor/hmp.h
index ae116d9804..4a6db16ab2 100644
--- a/include/monitor/hmp.h
+++ b/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
diff --git a/libafl/meson.build b/libafl/meson.build
index b6d1dc52e0..c5bfa579d4 100644
--- a/libafl/meson.build
+++ b/libafl/meson.build
@@ -26,6 +26,7 @@ specific_ss.add(when : 'CONFIG_SOFTMMU', if_true : [files(
                                                         'syx-snapshot/syx-snapshot.c',
                                                         'syx-snapshot/syx-cow-cache.c',
                                                         'syx-snapshot/channel-buffer-writeback.c',
+                                                        'syx-snapshot/syx-snapshot-hmp.c',
                                                     )])
 
 specific_ss.add(when : 'CONFIG_USER_ONLY', if_true : [files(
diff --git a/libafl/syx-snapshot/syx-snapshot-hmp.c b/libafl/syx-snapshot/syx-snapshot-hmp.c
new file mode 100644
index 0000000000..51aa9be3fd
--- /dev/null
+++ b/libafl/syx-snapshot/syx-snapshot-hmp.c
@@ -0,0 +1,34 @@
+#include "libafl/syx-snapshot/syx-snapshot.h"
+#include "libafl/syx-snapshot/device-save.h"
+//#include "syx-snapshot/syx-snapshot-hmp.h"
+
+// Static snapshot variable
+static SyxSnapshot* current_snapshot = NULL;
+
+
+
+
+/**
+ * Create a new snapshot and store it in the static variable.
+ */
+void hmp_syx_snapshot_new(Monitor *mon, const QDict *qdict) {
+
+    if (current_snapshot != NULL) {
+        syx_snapshot_free(current_snapshot);
+    }
+    current_snapshot = syx_snapshot_new(true, false, DEVICE_SNAPSHOT_ALL, NULL);
+}
+
+void hmp_syx_snapshot_init(Monitor *mon, const QDict *qdict) {
+    syx_snapshot_init(false);
+}
+
+    /**
+     * Restore the root snapshot from the static variable.
+     */
+    void hmp_syx_snapshot_root_restore(Monitor * mon, const QDict* qdict)
+    {
+        if (current_snapshot != NULL) {
+            syx_snapshot_root_restore(current_snapshot);
+        }
+}
\ No newline at end of file

From 12ddfc30682f96bb42b0f1849bc4bcb39443c805 Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Fri, 7 Mar 2025 08:47:22 +0100
Subject: [PATCH 02/12] [add] debug messages to syx

---
 accel/tcg/cputlb.c                         | 14 +++++++-------
 include/libafl/syx-snapshot/syx-snapshot.h |  9 +++++++++
 system/physmem.c                           | 13 ++++++-------
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/accel/tcg/cputlb.c b/accel/tcg/cputlb.c
index 5e63b72118..bcea355005 100644
--- a/accel/tcg/cputlb.c
+++ b/accel/tcg/cputlb.c
@@ -47,6 +47,9 @@
 #endif
 #include "tcg/tcg-ldst.h"
 #include "tcg/oversized-guest.h"
+//// --- Begin LibAFL code ---
+#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 */
@@ -89,13 +92,6 @@ QEMU_BUILD_BUG_ON(sizeof(vaddr) > sizeof(run_on_cpu_data));
 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);
-
-//// --- End LibAFL code ---
-
 static inline size_t tlb_n_entries(CPUTLBDescFast *fast)
 {
     return (fast->mask >> CPU_TLB_ENTRY_BITS) + 1;
@@ -1409,6 +1405,7 @@ static int probe_access_internal(CPUState *cpu, vaddr addr,
 //// --- Begin LibAFL code ---
 
     if (access_type == MMU_DATA_STORE) {
+        DPRINTF_SYX("vaddr probe %llx\n", addr);
         syx_snapshot_dirty_list_add_hostaddr(*phost);
     }
 
@@ -1776,6 +1773,7 @@ static bool mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
 
         // TODO: Does not work?
         // if (type == MMU_DATA_STORE) {
+            DPRINTF_SYX("vaddr mmu_lookup %llx\n", addr);
             syx_snapshot_dirty_list_add_hostaddr(l->page[0].haddr);
         // }
 
@@ -1806,6 +1804,7 @@ static bool mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
         //// --- Begin LibAFL code ---
 
         // if (type == MMU_DATA_STORE) {
+            DPRINTF_SYX("vaddr crosspage %llx\n", addr);
             syx_snapshot_dirty_list_add_hostaddr(l->page[0].haddr);
             syx_snapshot_dirty_list_add_hostaddr(l->page[1].haddr);
         // }
@@ -1909,6 +1908,7 @@ static void *atomic_mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
 
     //// --- Begin LibAFL code ---
 
+    DPRINTF_SYX("vaddr atomic %llx\n", addr);
     syx_snapshot_dirty_list_add_hostaddr(hostaddr);
 
     //// --- End LibAFL code ---
diff --git a/include/libafl/syx-snapshot/syx-snapshot.h b/include/libafl/syx-snapshot/syx-snapshot.h
index 5a6cbbc85c..1b835e8cc0 100644
--- a/include/libafl/syx-snapshot/syx-snapshot.h
+++ b/include/libafl/syx-snapshot/syx-snapshot.h
@@ -18,6 +18,15 @@
 
 #include "libafl/syx-misc.h"
 
+#define DEBUG_SYX
+
+#ifdef DEBUG_SYX
+#define DPRINTF_SYX(fmt, ...)                                   \
+    do { printf("SYX: " fmt , ## __VA_ARGS__); } while (0)
+#else
+#define DPRINTF_SYX(fmt, ...) do {} while (0)
+#endif
+
 #define SYX_SNAPSHOT_COW_CACHE_DEFAULT_CHUNK_SIZE 64
 #define SYX_SNAPSHOT_COW_CACHE_DEFAULT_MAX_BLOCKS (1024 * 1024)
 
diff --git a/system/physmem.c b/system/physmem.c
index 52867a474f..e8a803874e 100644
--- a/system/physmem.c
+++ b/system/physmem.c
@@ -82,6 +82,10 @@
 #include <daxctl/libdaxctl.h>
 #endif
 
+//// --- Begin LibAFL code ---
+#include "libafl/syx-snapshot/syx-snapshot.h"
+//// --- End LibAFL code ---
+
 //#define DEBUG_SUBPAGE
 
 /* ram_list is read under rcu_read_lock()/rcu_read_unlock().  Writes
@@ -3018,11 +3022,6 @@ enum write_rom_type {
     FLUSH_CACHE,
 };
 
-//// --- Begin LibAFL code ---
-
-void syx_snapshot_dirty_list_add_hostaddr_range(void* host_addr, uint64_t len);
-
-//// --- End LibAFL code ---
 
 static inline MemTxResult address_space_write_rom_internal(AddressSpace *as,
                                                            hwaddr addr,
@@ -3052,10 +3051,10 @@ static inline MemTxResult address_space_write_rom_internal(AddressSpace *as,
             case WRITE_DATA:
                 memcpy(ram_ptr, buf, l);
 //// --- Begin LibAFL code ---
-
+                DPRINTF_SYX("hwaddr write_rom %llx\n", addr);
                 syx_snapshot_dirty_list_add_hostaddr_range(ram_ptr, l);
 
-//// --- End LibAFL code ---
+                //// --- End LibAFL code ---
                 invalidate_and_set_dirty(mr, addr1, l);
                 break;
             case FLUSH_CACHE:

From fdb2844c484fbc620cbae7468f13758b5e5d2028 Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Fri, 7 Mar 2025 17:02:56 +0100
Subject: [PATCH 03/12] [fix+add] debug syx and TLB

---
 accel/tcg/cputlb.c                         | 12 ++++++------
 include/libafl/syx-snapshot/syx-snapshot.h | 10 +---------
 system/physmem.c                           |  2 +-
 3 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/accel/tcg/cputlb.c b/accel/tcg/cputlb.c
index bcea355005..72c80ec621 100644
--- a/accel/tcg/cputlb.c
+++ b/accel/tcg/cputlb.c
@@ -52,7 +52,7 @@
 //// --- End LibAFL code ---
 
 /* DEBUG defines, enable DEBUG_TLB_LOG to log to the CPU_LOG_MMU target */
-/* #define DEBUG_TLB */
+#define DEBUG_TLB 
 /* #define DEBUG_TLB_LOG */
 
 #ifdef DEBUG_TLB
@@ -77,7 +77,7 @@
 } while (0)
 
 #define assert_cpu_is_self(cpu) do {                              \
-        if (DEBUG_TLB_GATE) {                                     \
+        if (DEBUG_TLB_GATE && false) {                                     \
             g_assert(!(cpu)->created || qemu_cpu_is_self(cpu));   \
         }                                                         \
     } while (0)
@@ -1405,7 +1405,7 @@ static int probe_access_internal(CPUState *cpu, vaddr addr,
 //// --- Begin LibAFL code ---
 
     if (access_type == MMU_DATA_STORE) {
-        DPRINTF_SYX("vaddr probe %llx\n", addr);
+        SYX_DEBUG("vaddr probe %llx\n", addr);
         syx_snapshot_dirty_list_add_hostaddr(*phost);
     }
 
@@ -1773,7 +1773,7 @@ static bool mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
 
         // TODO: Does not work?
         // if (type == MMU_DATA_STORE) {
-            DPRINTF_SYX("vaddr mmu_lookup %llx\n", addr);
+            SYX_DEBUG("vaddr mmu_lookup %llx %d\n", addr, type);
             syx_snapshot_dirty_list_add_hostaddr(l->page[0].haddr);
         // }
 
@@ -1804,7 +1804,7 @@ static bool mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
         //// --- Begin LibAFL code ---
 
         // if (type == MMU_DATA_STORE) {
-            DPRINTF_SYX("vaddr crosspage %llx\n", addr);
+            SYX_DEBUG("vaddr crosspage %llx\n", addr);
             syx_snapshot_dirty_list_add_hostaddr(l->page[0].haddr);
             syx_snapshot_dirty_list_add_hostaddr(l->page[1].haddr);
         // }
@@ -1908,7 +1908,7 @@ static void *atomic_mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
 
     //// --- Begin LibAFL code ---
 
-    DPRINTF_SYX("vaddr atomic %llx\n", addr);
+    SYX_DEBUG("vaddr atomic %llx\n", addr);
     syx_snapshot_dirty_list_add_hostaddr(hostaddr);
 
     //// --- End LibAFL code ---
diff --git a/include/libafl/syx-snapshot/syx-snapshot.h b/include/libafl/syx-snapshot/syx-snapshot.h
index 1b835e8cc0..b3dc5c3f80 100644
--- a/include/libafl/syx-snapshot/syx-snapshot.h
+++ b/include/libafl/syx-snapshot/syx-snapshot.h
@@ -16,17 +16,9 @@
 #include "device-save.h"
 #include "syx-cow-cache.h"
 
+#define CONFIG_DEBUG_SYX
 #include "libafl/syx-misc.h"
 
-#define DEBUG_SYX
-
-#ifdef DEBUG_SYX
-#define DPRINTF_SYX(fmt, ...)                                   \
-    do { printf("SYX: " fmt , ## __VA_ARGS__); } while (0)
-#else
-#define DPRINTF_SYX(fmt, ...) do {} while (0)
-#endif
-
 #define SYX_SNAPSHOT_COW_CACHE_DEFAULT_CHUNK_SIZE 64
 #define SYX_SNAPSHOT_COW_CACHE_DEFAULT_MAX_BLOCKS (1024 * 1024)
 
diff --git a/system/physmem.c b/system/physmem.c
index e8a803874e..570c906662 100644
--- a/system/physmem.c
+++ b/system/physmem.c
@@ -3051,7 +3051,7 @@ static inline MemTxResult address_space_write_rom_internal(AddressSpace *as,
             case WRITE_DATA:
                 memcpy(ram_ptr, buf, l);
 //// --- Begin LibAFL code ---
-                DPRINTF_SYX("hwaddr write_rom %llx\n", addr);
+                SYX_DEBUG("hwaddr write_rom %llx\n", addr);
                 syx_snapshot_dirty_list_add_hostaddr_range(ram_ptr, l);
 
                 //// --- End LibAFL code ---

From 5ea61582abc171e94557ad76da239c8c4886fd5c Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Fri, 7 Mar 2025 18:22:49 +0100
Subject: [PATCH 04/12] [add] flush TLB on new snapshot tracking

---
 accel/tcg/cputlb.c                 | 9 +++++++++
 include/exec/exec-all.h            | 3 +++
 libafl/syx-snapshot/syx-snapshot.c | 4 ++++
 3 files changed, 16 insertions(+)

diff --git a/accel/tcg/cputlb.c b/accel/tcg/cputlb.c
index 72c80ec621..cea5fc45ba 100644
--- a/accel/tcg/cputlb.c
+++ b/accel/tcg/cputlb.c
@@ -446,6 +446,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() 
+{
+    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)
 {
diff --git a/include/exec/exec-all.h b/include/exec/exec-all.h
index 2e4c4cc4b4..5ee47d3626 100644
--- a/include/exec/exec-all.h
+++ b/include/exec/exec-all.h
@@ -108,6 +108,9 @@ void tlb_flush_all_cpus_synced(CPUState *src_cpu);
  * Flush one page from the TLB of the specified CPU, for the specified
  * MMU indexes.
  */
+
+void tlb_flush_all_cpus();
+
 void tlb_flush_page_by_mmuidx(CPUState *cpu, vaddr addr,
                               uint16_t idxmap);
 /**
diff --git a/libafl/syx-snapshot/syx-snapshot.c b/libafl/syx-snapshot/syx-snapshot.c
index 7eb4ad79e8..a3337784d5 100644
--- a/libafl/syx-snapshot/syx-snapshot.c
+++ b/libafl/syx-snapshot/syx-snapshot.c
@@ -187,6 +187,10 @@ SyxSnapshot* syx_snapshot_new(bool track, bool is_active_bdrv_cache,
 
     if (track) {
         syx_snapshot_track(&syx_snapshot_state.tracked_snapshots, snapshot);
+
+        //make sure to catch all new writes
+        //with a filled TLB there might be missed writes from jitted TCG code
+        tlb_flush_all_cpus();
     }
 
     syx_snapshot_state.is_enabled = true;

From ee8ec01da35cf06b989ee9fb2b17f5ceb14c868b Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Fri, 14 Mar 2025 09:22:14 +0100
Subject: [PATCH 05/12] [fix] rename syx-snapshot-hmp.c -> syx-hmp.c

---
 libafl/meson.build                               |  2 +-
 .../{syx-snapshot-hmp.c => syx-hmp.c}            | 16 ++++++++--------
 2 files changed, 9 insertions(+), 9 deletions(-)
 rename libafl/syx-snapshot/{syx-snapshot-hmp.c => syx-hmp.c} (70%)

diff --git a/libafl/meson.build b/libafl/meson.build
index c5bfa579d4..c1967ffde2 100644
--- a/libafl/meson.build
+++ b/libafl/meson.build
@@ -26,7 +26,7 @@ specific_ss.add(when : 'CONFIG_SOFTMMU', if_true : [files(
                                                         'syx-snapshot/syx-snapshot.c',
                                                         'syx-snapshot/syx-cow-cache.c',
                                                         'syx-snapshot/channel-buffer-writeback.c',
-                                                        'syx-snapshot/syx-snapshot-hmp.c',
+                                                        'syx-snapshot/syx-hmp.c',
                                                     )])
 
 specific_ss.add(when : 'CONFIG_USER_ONLY', if_true : [files(
diff --git a/libafl/syx-snapshot/syx-snapshot-hmp.c b/libafl/syx-snapshot/syx-hmp.c
similarity index 70%
rename from libafl/syx-snapshot/syx-snapshot-hmp.c
rename to libafl/syx-snapshot/syx-hmp.c
index 51aa9be3fd..f27c1b63ba 100644
--- a/libafl/syx-snapshot/syx-snapshot-hmp.c
+++ b/libafl/syx-snapshot/syx-hmp.c
@@ -23,12 +23,12 @@ void hmp_syx_snapshot_init(Monitor *mon, const QDict *qdict) {
     syx_snapshot_init(false);
 }
 
-    /**
-     * Restore the root snapshot from the static variable.
-     */
-    void hmp_syx_snapshot_root_restore(Monitor * mon, const QDict* qdict)
-    {
-        if (current_snapshot != NULL) {
-            syx_snapshot_root_restore(current_snapshot);
-        }
+/**
+ * Restore the root snapshot from the static variable.
+ */
+void hmp_syx_snapshot_root_restore(Monitor * mon, const QDict* qdict)
+{
+    if (current_snapshot != NULL) {
+        syx_snapshot_root_restore(current_snapshot);
+    }
 }
\ No newline at end of file

From 5b0a3cc9e2c58a7fa334d2365155f8f182cb885a Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Fri, 14 Mar 2025 09:39:01 +0100
Subject: [PATCH 06/12] [add] libafl_blk_write. Also remove CONFIG_SOFTMMU from
 block-backend.c

The reason is that CONFIG_SOFTMMU is not available for the "block" source set.
And it is not required, because "block" will only be built for system mode and tools, which we can deactivate using configure.
---
 block/block-backend.c         | 24 +++++++--------------
 include/libafl/system.h       |  7 +++++++
 libafl/blkdev.c               | 39 +++++++++++++++++++++++++++++++++++
 libafl/meson.build            |  1 +
 libafl/syx-snapshot/syx-hmp.c |  4 ++--
 5 files changed, 57 insertions(+), 18 deletions(-)
 create mode 100644 libafl/blkdev.c

diff --git a/block/block-backend.c b/block/block-backend.c
index ff92137a24..031b269704 100644
--- a/block/block-backend.c
+++ b/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,15 @@ 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;
+        SYX_DEBUG("Read from COW cache: 0x%llx 0x%llx\n", rwco->offset, acb->bytes);
+        rwco->ret = 0;
     }
-#endif
 //// --- End LibAFL code ---
 
     blk_aio_complete(acb);
@@ -1676,17 +1671,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);
diff --git a/include/libafl/system.h b/include/libafl/system.h
index 9a6e9c6d4d..aae6a3f5c7 100644
--- a/include/libafl/system.h
+++ b/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)
+ */
+void libafl_blk_write(BlockBackend *blk, void *buf, int64_t offset, int64_t sz);
\ No newline at end of file
diff --git a/libafl/blkdev.c b/libafl/blkdev.c
new file mode 100644
index 0000000000..ec00f30601
--- /dev/null
+++ b/libafl/blkdev.c
@@ -0,0 +1,39 @@
+
+#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;
+}
+
+void 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);
+}
\ No newline at end of file
diff --git a/libafl/meson.build b/libafl/meson.build
index c1967ffde2..7df94b2072 100644
--- a/libafl/meson.build
+++ b/libafl/meson.build
@@ -22,6 +22,7 @@ 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',
diff --git a/libafl/syx-snapshot/syx-hmp.c b/libafl/syx-snapshot/syx-hmp.c
index f27c1b63ba..0786417f7d 100644
--- a/libafl/syx-snapshot/syx-hmp.c
+++ b/libafl/syx-snapshot/syx-hmp.c
@@ -16,11 +16,11 @@ void hmp_syx_snapshot_new(Monitor *mon, const QDict *qdict) {
     if (current_snapshot != NULL) {
         syx_snapshot_free(current_snapshot);
     }
-    current_snapshot = syx_snapshot_new(true, false, DEVICE_SNAPSHOT_ALL, NULL);
+    current_snapshot = syx_snapshot_new(true, true, DEVICE_SNAPSHOT_ALL, NULL);
 }
 
 void hmp_syx_snapshot_init(Monitor *mon, const QDict *qdict) {
-    syx_snapshot_init(false);
+    syx_snapshot_init(true);
 }
 
 /**

From 596b0a19580e1369932aa3790af0dcdb8a5ee6bf Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Sat, 5 Apr 2025 09:48:53 +0200
Subject: [PATCH 07/12] [fix+debug] Fix syx_snapshot_state.before_fuzz_cache
 handling. Add some DEBUG print statements

---
 accel/tcg/cputlb.c                         |  2 +-
 block/block-backend.c                      |  2 +-
 include/libafl/system.h                    |  2 +-
 include/libafl/syx-snapshot/syx-snapshot.h |  2 +-
 libafl/blkdev.c                            |  5 +++--
 libafl/syx-snapshot/syx-cow-cache.c        |  7 +++++--
 libafl/syx-snapshot/syx-snapshot.c         | 19 ++++++++++---------
 7 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/accel/tcg/cputlb.c b/accel/tcg/cputlb.c
index cea5fc45ba..14d478f322 100644
--- a/accel/tcg/cputlb.c
+++ b/accel/tcg/cputlb.c
@@ -52,7 +52,7 @@
 //// --- End LibAFL code ---
 
 /* DEBUG defines, enable DEBUG_TLB_LOG to log to the CPU_LOG_MMU target */
-#define DEBUG_TLB 
+//#define DEBUG_TLB 
 /* #define DEBUG_TLB_LOG */
 
 #ifdef DEBUG_TLB
diff --git a/block/block-backend.c b/block/block-backend.c
index 031b269704..4ded8ea19e 100644
--- a/block/block-backend.c
+++ b/block/block-backend.c
@@ -29,6 +29,7 @@
 #include "migration/misc.h"
 
 //// --- Begin LibAFL code ---
+#define CONFIG_DEBUG_SYX
 #include "libafl/syx-snapshot/syx-snapshot.h"
 //// --- End LibAFL code ---
 
@@ -1654,7 +1655,6 @@ static void coroutine_fn blk_aio_read_entry(void *opaque)
                                       0, rwco->flags);
 //// --- Begin LibAFL code ---
     } else {
-        SYX_DEBUG("Read from COW cache: 0x%llx 0x%llx\n", rwco->offset, acb->bytes);
         rwco->ret = 0;
     }
 //// --- End LibAFL code ---
diff --git a/include/libafl/system.h b/include/libafl/system.h
index aae6a3f5c7..d7f3323b25 100644
--- a/include/libafl/system.h
+++ b/include/libafl/system.h
@@ -15,4 +15,4 @@ void libafl_qemu_init(int argc, char** argv);
  * The same way the guest would, 
  * thus this writes to the Syx COW cache (if it is initialized)
  */
-void libafl_blk_write(BlockBackend *blk, void *buf, int64_t offset, int64_t sz);
\ No newline at end of file
+int libafl_blk_write(BlockBackend *blk, void *buf, int64_t offset, int64_t sz);
\ No newline at end of file
diff --git a/include/libafl/syx-snapshot/syx-snapshot.h b/include/libafl/syx-snapshot/syx-snapshot.h
index b3dc5c3f80..29027624fc 100644
--- a/include/libafl/syx-snapshot/syx-snapshot.h
+++ b/include/libafl/syx-snapshot/syx-snapshot.h
@@ -16,7 +16,7 @@
 #include "device-save.h"
 #include "syx-cow-cache.h"
 
-#define CONFIG_DEBUG_SYX
+//#define CONFIG_DEBUG_SYX
 #include "libafl/syx-misc.h"
 
 #define SYX_SNAPSHOT_COW_CACHE_DEFAULT_CHUNK_SIZE 64
diff --git a/libafl/blkdev.c b/libafl/blkdev.c
index ec00f30601..a2a1f7c1ec 100644
--- a/libafl/blkdev.c
+++ b/libafl/blkdev.c
@@ -17,7 +17,7 @@ static void blk_rw_done(void *opaque, int ret)
     *(int *)opaque = ret;
 }
 
-void libafl_blk_write(BlockBackend *blk, void *buf, int64_t offset, int64_t sz)
+int libafl_blk_write(BlockBackend *blk, void *buf, int64_t offset, int64_t sz)
 {
 	void *pattern_buf = NULL;
 	QEMUIOVector qiov;
@@ -31,9 +31,10 @@ void libafl_blk_write(BlockBackend *blk, void *buf, int64_t offset, int64_t sz)
 		main_loop_wait(false);
 	}
 
-	printf("async_ret: %d\n", async_ret);
+	//printf("async_ret: %d\n", async_ret);
 	//g_assert(async_ret == 0);
 
 	g_free(pattern_buf);
 	qemu_iovec_destroy(&qiov);
+    return async_ret;
 }
\ No newline at end of file
diff --git a/libafl/syx-snapshot/syx-cow-cache.c b/libafl/syx-snapshot/syx-cow-cache.c
index 10f5de0fe2..e765b98504 100644
--- a/libafl/syx-snapshot/syx-cow-cache.c
+++ b/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,8 +218,7 @@ 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);
+    printf("[SYX] [%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);
@@ -227,9 +227,12 @@ void syx_cow_cache_read_entry(SyxCowCache* scc, BlockBackend* blk,
     if (!QTAILQ_EMPTY(&scc->layers)) {
         for (; qiov_offset < qiov->size;
              blk_offset += chunk_size, qiov_offset += chunk_size) {
+            int i = 0;
             QTAILQ_FOREACH(layer, &scc->layers, next)
             {
                 chunk_size = layer->chunk_size;
+                //printf("[SYX] check cache layer %d\n", i);
+                i++;
                 if (read_chunk_from_cache_layer(layer, blk, qiov, qiov_offset,
                                                 blk_offset)) {
                     break;
diff --git a/libafl/syx-snapshot/syx-snapshot.c b/libafl/syx-snapshot/syx-snapshot.c
index a3337784d5..2d181528b4 100644
--- a/libafl/syx-snapshot/syx-snapshot.c
+++ b/libafl/syx-snapshot/syx-snapshot.c
@@ -173,23 +173,24 @@ SyxSnapshot* syx_snapshot_new(bool track, bool is_active_bdrv_cache,
     snapshot->rbs_dirty_list =
         g_hash_table_new_full(g_direct_hash, g_direct_equal, NULL,
                               (GDestroyNotify)g_hash_table_remove_all);
-    snapshot->bdrvs_cow_cache = syx_cow_cache_new();
-
+    
     if (is_active_bdrv_cache) {
-        syx_cow_cache_move(snapshot->bdrvs_cow_cache,
-                           &syx_snapshot_state.before_fuzz_cache);
-        syx_snapshot_state.active_bdrv_cache_snapshot = snapshot;
+        // we have cached writes from BEFORE fuzzing starts
+        snapshot->bdrvs_cow_cache = syx_snapshot_state.before_fuzz_cache;
+        syx_snapshot_state.before_fuzz_cache = NULL;
     } else {
-        syx_cow_cache_push_layer(snapshot->bdrvs_cow_cache,
-                                 SYX_SNAPSHOT_COW_CACHE_DEFAULT_CHUNK_SIZE,
-                                 SYX_SNAPSHOT_COW_CACHE_DEFAULT_MAX_BLOCKS);
+        snapshot->bdrvs_cow_cache = syx_cow_cache_new();
     }
+    syx_cow_cache_push_layer(snapshot->bdrvs_cow_cache,
+        SYX_SNAPSHOT_COW_CACHE_DEFAULT_CHUNK_SIZE,
+        SYX_SNAPSHOT_COW_CACHE_DEFAULT_MAX_BLOCKS);
+    syx_snapshot_state.active_bdrv_cache_snapshot = snapshot;
 
     if (track) {
         syx_snapshot_track(&syx_snapshot_state.tracked_snapshots, snapshot);
 
         //make sure to catch all new writes
-        //with a filled TLB there might be missed writes from jitted TCG code
+        //with a filled TLB there might be missed writes
         tlb_flush_all_cpus();
     }
 

From 80791c947ad6a4596d1d849179f9e1127e76a56c Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Sat, 5 Apr 2025 17:59:03 +0200
Subject: [PATCH 08/12] [fix] Mark RAM as notdirty before tracking new
 snapshot, or after snapshot restore

Also flush TLB and TB more specifically
---
 accel/tcg/cputlb.c                 | 11 ++++++-----
 cpu-target.c                       | 10 +++-------
 include/exec/exec-all.h            |  6 +++---
 include/exec/ram_addr.h            |  1 +
 libafl/syx-snapshot/syx-snapshot.c | 26 ++++++++++++++++++++++++++
 5 files changed, 39 insertions(+), 15 deletions(-)

diff --git a/accel/tcg/cputlb.c b/accel/tcg/cputlb.c
index 14d478f322..83dd7b0fa1 100644
--- a/accel/tcg/cputlb.c
+++ b/accel/tcg/cputlb.c
@@ -48,6 +48,7 @@
 #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 ---
 
@@ -446,7 +447,7 @@ 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 tlb_flush_all_cpus(void) 
 {
     const run_on_cpu_func fn = tlb_flush_by_mmuidx_async_work;
 
@@ -1781,10 +1782,10 @@ static bool mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
         //// --- Begin LibAFL code ---
 
         // TODO: Does not work?
-        // if (type == MMU_DATA_STORE) {
+        if (type == MMU_DATA_STORE && !(flags & (TLB_MMIO | TLB_DISCARD_WRITE))) {
             SYX_DEBUG("vaddr mmu_lookup %llx %d\n", addr, type);
             syx_snapshot_dirty_list_add_hostaddr(l->page[0].haddr);
-        // }
+        }
 
         //// --- End LibAFL code ---
 
@@ -1812,11 +1813,11 @@ static bool mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
 
         //// --- Begin LibAFL code ---
 
-        // if (type == MMU_DATA_STORE) {
+        if (type == MMU_DATA_STORE && !(flags & (TLB_MMIO | TLB_DISCARD_WRITE))) {
             SYX_DEBUG("vaddr crosspage %llx\n", addr);
             syx_snapshot_dirty_list_add_hostaddr(l->page[0].haddr);
             syx_snapshot_dirty_list_add_hostaddr(l->page[1].haddr);
-        // }
+        }
 
         //// --- End LibAFL code ---
 
diff --git a/cpu-target.c b/cpu-target.c
index 4347488cb3..3378a5df17 100644
--- a/cpu-target.c
+++ b/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;
diff --git a/include/exec/exec-all.h b/include/exec/exec-all.h
index 5ee47d3626..64b480a274 100644
--- a/include/exec/exec-all.h
+++ b/include/exec/exec-all.h
@@ -108,11 +108,11 @@ void tlb_flush_all_cpus_synced(CPUState *src_cpu);
  * Flush one page from the TLB of the specified CPU, for the specified
  * MMU indexes.
  */
-
-void tlb_flush_all_cpus();
-
 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
diff --git a/include/exec/ram_addr.h b/include/exec/ram_addr.h
index 891c44cf2d..439a1f3fe0 100644
--- a/include/exec/ram_addr.h
+++ b/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;
diff --git a/libafl/syx-snapshot/syx-snapshot.c b/libafl/syx-snapshot/syx-snapshot.c
index 2d181528b4..7a0078239f 100644
--- a/libafl/syx-snapshot/syx-snapshot.c
+++ b/libafl/syx-snapshot/syx-snapshot.c
@@ -8,6 +8,7 @@
 #include "exec/ramlist.h"
 #include "exec/ram_addr.h"
 #include "exec/exec-all.h"
+#include "exec/address-spaces.h"
 
 #include "libafl/syx-snapshot/syx-snapshot.h"
 #include "libafl/syx-snapshot/device-save.h"
@@ -98,6 +99,23 @@ static void root_restore_check_memory_rb(gpointer rb_idstr_hash,
 static SyxSnapshotIncrement*
 syx_snapshot_increment_free(SyxSnapshotIncrement* increment);
 
+//set all RAM as clear (not dirty)
+static void all_ram_notdirty(void) {
+    
+    MemoryRegion *sysmem, *subregion, *next;
+    sysmem = get_system_memory();
+    QTAILQ_FOREACH_SAFE(subregion, &sysmem->subregions, subregions_link,
+                        next)
+    {
+        if (subregion->ram) {
+            #ifdef SYX_SNAPSHOT_DEBUG
+            printf("memory_region_reset_dirty: %llx %llx\n", subregion->addr, subregion->size);
+            #endif
+            memory_region_reset_dirty(subregion, 0, subregion->size, DIRTY_MEMORY_MIGRATION);
+        }
+    }
+}
+
 static RAMBlock* ramblock_lookup(gpointer rb_idstr_hash)
 {
     RAMBlock* block;
@@ -192,6 +210,8 @@ SyxSnapshot* syx_snapshot_new(bool track, bool is_active_bdrv_cache,
         //make sure to catch all new writes
         //with a filled TLB there might be missed writes
         tlb_flush_all_cpus();
+
+        all_ram_notdirty();
     }
 
     syx_snapshot_state.is_enabled = true;
@@ -630,6 +650,11 @@ static void root_restore_rb_page(gpointer offset_within_rb, gpointer _unused,
     memcpy(host_rb_restore, host_snapshot_rb_restore,
            syx_snapshot_state.page_size);
     // TODO: manage special case of TSEG.
+
+    // Invalidate TBs
+    tb_invalidate_phys_range(rb->offset + (ram_addr_t)offset_within_rb,
+                             rb->offset + (ram_addr_t)offset_within_rb +
+                                 syx_snapshot_state.page_size);
 }
 
 static void root_restore_rb(gpointer rb_idstr_hash,
@@ -742,6 +767,7 @@ void syx_snapshot_root_restore(SyxSnapshot* snapshot)
     }
 
     syx_snapshot_dirty_list_flush(snapshot);
+    all_ram_notdirty();
 
     if (must_unlock_bql) {
         bql_unlock();

From 8d8f145f0c2f5e91f08124556f53909bed7a311c Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Tue, 8 Apr 2025 08:06:50 +0200
Subject: [PATCH 09/12] [fix] SYX: keep in line with notdirty_write

---
 accel/tcg/cputlb.c | 69 ++++++++++++++++++++++++----------------------
 1 file changed, 36 insertions(+), 33 deletions(-)

diff --git a/accel/tcg/cputlb.c b/accel/tcg/cputlb.c
index 83dd7b0fa1..485d85469b 100644
--- a/accel/tcg/cputlb.c
+++ b/accel/tcg/cputlb.c
@@ -1412,17 +1412,19 @@ 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_DEBUG("vaddr probe %llx\n", addr);
-        syx_snapshot_dirty_list_add_hostaddr(*phost);
-    }
 
-//// --- End LibAFL code ---
     return flags;
 }
 
+//// --- Begin LibAFL code ---
+// Use this snippet multiple times just 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 ---
+
 int probe_access_full(CPUArchState *env, vaddr addr, int size,
                       MMUAccessType access_type, int mmu_idx,
                       bool nonfault, void **phost, CPUTLBEntryFull **pfull,
@@ -1437,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;
@@ -1461,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;
@@ -1484,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_full_flags", access_type, addr, full, *phost);
+        //// --- End LibAFL code ---
     }
 
     return flags;
@@ -1519,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 ---
         }
     }
 
@@ -1735,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;
 }
@@ -1779,15 +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 && !(flags & (TLB_MMIO | TLB_DISCARD_WRITE))) {
-            SYX_DEBUG("vaddr mmu_lookup %llx %d\n", addr, type);
-            syx_snapshot_dirty_list_add_hostaddr(l->page[0].haddr);
-        }
-
-        //// --- End LibAFL code ---
 
     } else {
         /* Finish compute of page crossing. */
@@ -1811,16 +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 && !(flags & (TLB_MMIO | TLB_DISCARD_WRITE))) {
-            SYX_DEBUG("vaddr crosspage %llx\n", addr);
-            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
@@ -1916,15 +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_DEBUG("vaddr atomic %llx\n", addr);
-    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)) {

From 6c4e90476dfefc67f9021b7ff86537a9e3d6f734 Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Tue, 8 Apr 2025 08:07:23 +0200
Subject: [PATCH 10/12] [add] SYX: DMA flatview_write_continue_step

---
 system/physmem.c | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/system/physmem.c b/system/physmem.c
index 570c906662..af466e1727 100644
--- a/system/physmem.c
+++ b/system/physmem.c
@@ -83,6 +83,7 @@
 #endif
 
 //// --- Begin LibAFL code ---
+#define CONFIG_DEBUG_SYX
 #include "libafl/syx-snapshot/syx-snapshot.h"
 //// --- End LibAFL code ---
 
@@ -2807,6 +2808,17 @@ static MemTxResult flatview_write_continue_step(MemTxAttrs attrs,
                                                false, true);
 
         memmove(ram_ptr, buf, *l);
+
+        //// --- Begin LibAFL code ---
+        SYX_DEBUG("flatview_write_continue_step %llx %d\n", mr_addr, *l);
+        unsigned long hpage = (unsigned long)ram_ptr & TARGET_PAGE_MASK;
+
+        while (hpage < (unsigned long) ram_ptr + *l) {
+            syx_snapshot_dirty_list_add_hostaddr((void*)hpage);
+            hpage += 1<<TARGET_PAGE_BITS;
+        }
+        //// --- End LibAFL code ---
+
         invalidate_and_set_dirty(mr, mr_addr, *l);
 
         return MEMTX_OK;

From d2936b188b3224c57ddd4b9dea920ed8c59727ad Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Tue, 8 Apr 2025 08:35:27 +0200
Subject: [PATCH 11/12] [fix] Cleanup and remove debug printfs

---
 accel/tcg/cputlb.c                  | 8 ++++----
 block/block-backend.c               | 2 +-
 hmp-commands.hx                     | 2 +-
 libafl/syx-snapshot/syx-cow-cache.c | 2 +-
 system/physmem.c                    | 6 +++++-
 5 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/accel/tcg/cputlb.c b/accel/tcg/cputlb.c
index 485d85469b..23e45324de 100644
--- a/accel/tcg/cputlb.c
+++ b/accel/tcg/cputlb.c
@@ -53,7 +53,7 @@
 //// --- End LibAFL code ---
 
 /* DEBUG defines, enable DEBUG_TLB_LOG to log to the CPU_LOG_MMU target */
-//#define DEBUG_TLB 
+/* #define DEBUG_TLB */
 /* #define DEBUG_TLB_LOG */
 
 #ifdef DEBUG_TLB
@@ -78,7 +78,7 @@
 } while (0)
 
 #define assert_cpu_is_self(cpu) do {                              \
-        if (DEBUG_TLB_GATE && false) {                                     \
+        if (DEBUG_TLB_GATE) {                                     \
             g_assert(!(cpu)->created || qemu_cpu_is_self(cpu));   \
         }                                                         \
     } while (0)
@@ -1417,7 +1417,7 @@ static int probe_access_internal(CPUState *cpu, vaddr addr,
 }
 
 //// --- Begin LibAFL code ---
-// Use this snippet multiple times just below
+// 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); \
@@ -1496,7 +1496,7 @@ int probe_access_flags(CPUArchState *env, vaddr addr, int size,
         flags &= ~TLB_NOTDIRTY;
 
         //// --- Begin LibAFL code ---
-        SYX_SNAPSHOT_DIRTY_LIST_ADD_HOSTADDR_PROBE("probe_access_full_flags", access_type, addr, full, *phost);
+        SYX_SNAPSHOT_DIRTY_LIST_ADD_HOSTADDR_PROBE("probe_access_flags", access_type, addr, full, *phost);
         //// --- End LibAFL code ---
     }
 
diff --git a/block/block-backend.c b/block/block-backend.c
index 4ded8ea19e..23fbc3016a 100644
--- a/block/block-backend.c
+++ b/block/block-backend.c
@@ -29,7 +29,7 @@
 #include "migration/misc.h"
 
 //// --- Begin LibAFL code ---
-#define CONFIG_DEBUG_SYX
+//#define CONFIG_DEBUG_SYX
 #include "libafl/syx-snapshot/syx-snapshot.h"
 //// --- End LibAFL code ---
 
diff --git a/hmp-commands.hx b/hmp-commands.hx
index 56ff0f785d..5268756201 100644
--- a/hmp-commands.hx
+++ b/hmp-commands.hx
@@ -1865,7 +1865,7 @@ ERST
         .name       = "syx-snapshot-new",
         .args_type  = "",
         .params     = "",
-        .help       = "create a new snapshot and store it in the static variable",
+        .help       = "create a new snapshot and store it in the static variable. Use syx-snapshot-init BEFORE.",
         .cmd        = hmp_syx_snapshot_new,
     },
 
diff --git a/libafl/syx-snapshot/syx-cow-cache.c b/libafl/syx-snapshot/syx-cow-cache.c
index e765b98504..47232068fe 100644
--- a/libafl/syx-snapshot/syx-cow-cache.c
+++ b/libafl/syx-snapshot/syx-cow-cache.c
@@ -218,7 +218,7 @@ void syx_cow_cache_read_entry(SyxCowCache* scc, BlockBackend* blk,
     size_t qiov_offset = 0;
     uint64_t chunk_size = 0;
 
-    printf("[SYX] [%s] Read 0x%zx bytes @addr %lx\n", blk_name(blk), qiov->size, offset);
+    //printf("[SYX] [%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);
diff --git a/system/physmem.c b/system/physmem.c
index af466e1727..54ff810791 100644
--- a/system/physmem.c
+++ b/system/physmem.c
@@ -83,7 +83,7 @@
 #endif
 
 //// --- Begin LibAFL code ---
-#define CONFIG_DEBUG_SYX
+//#define CONFIG_DEBUG_SYX
 #include "libafl/syx-snapshot/syx-snapshot.h"
 //// --- End LibAFL code ---
 
@@ -2810,6 +2810,9 @@ static MemTxResult flatview_write_continue_step(MemTxAttrs attrs,
         memmove(ram_ptr, buf, *l);
 
         //// --- Begin LibAFL code ---
+	    /**
+	      * Should catch DMA writes by devices 
+	      * Not sure it is needed right now...
         SYX_DEBUG("flatview_write_continue_step %llx %d\n", mr_addr, *l);
         unsigned long hpage = (unsigned long)ram_ptr & TARGET_PAGE_MASK;
 
@@ -2817,6 +2820,7 @@ static MemTxResult flatview_write_continue_step(MemTxAttrs attrs,
             syx_snapshot_dirty_list_add_hostaddr((void*)hpage);
             hpage += 1<<TARGET_PAGE_BITS;
         }
+	    */
         //// --- End LibAFL code ---
 
         invalidate_and_set_dirty(mr, mr_addr, *l);

From 99f3ff6f586066a7d02edc6162613c96ed4bebb0 Mon Sep 17 00:00:00 2001
From: enisrat <84377236+enisrat@users.noreply.github.com>
Date: Sat, 17 May 2025 17:49:28 +0200
Subject: [PATCH 12/12] [fix] cleanup

---
 accel/tcg/cputlb.c                         | 18 +++++++++---------
 block/block-backend.c                      |  1 -
 include/libafl/syx-snapshot/syx-snapshot.h |  1 -
 libafl/syx-snapshot/syx-cow-cache.c        |  5 -----
 4 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/accel/tcg/cputlb.c b/accel/tcg/cputlb.c
index 23e45324de..aa26ab9d09 100644
--- a/accel/tcg/cputlb.c
+++ b/accel/tcg/cputlb.c
@@ -93,6 +93,15 @@ QEMU_BUILD_BUG_ON(sizeof(vaddr) > sizeof(run_on_cpu_data));
 QEMU_BUILD_BUG_ON(NB_MMU_MODES > 16);
 #define ALL_MMUIDX_BITS ((1 << NB_MMU_MODES) - 1)
 
+//// --- Begin LibAFL code ---
+// 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)
 {
     return (fast->mask >> CPU_TLB_ENTRY_BITS) + 1;
@@ -1416,15 +1425,6 @@ static int probe_access_internal(CPUState *cpu, vaddr addr,
     return flags;
 }
 
-//// --- Begin LibAFL code ---
-// 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 ---
-
 int probe_access_full(CPUArchState *env, vaddr addr, int size,
                       MMUAccessType access_type, int mmu_idx,
                       bool nonfault, void **phost, CPUTLBEntryFull **pfull,
diff --git a/block/block-backend.c b/block/block-backend.c
index 23fbc3016a..bf01315ee4 100644
--- a/block/block-backend.c
+++ b/block/block-backend.c
@@ -29,7 +29,6 @@
 #include "migration/misc.h"
 
 //// --- Begin LibAFL code ---
-//#define CONFIG_DEBUG_SYX
 #include "libafl/syx-snapshot/syx-snapshot.h"
 //// --- End LibAFL code ---
 
diff --git a/include/libafl/syx-snapshot/syx-snapshot.h b/include/libafl/syx-snapshot/syx-snapshot.h
index 29027624fc..5a6cbbc85c 100644
--- a/include/libafl/syx-snapshot/syx-snapshot.h
+++ b/include/libafl/syx-snapshot/syx-snapshot.h
@@ -16,7 +16,6 @@
 #include "device-save.h"
 #include "syx-cow-cache.h"
 
-//#define CONFIG_DEBUG_SYX
 #include "libafl/syx-misc.h"
 
 #define SYX_SNAPSHOT_COW_CACHE_DEFAULT_CHUNK_SIZE 64
diff --git a/libafl/syx-snapshot/syx-cow-cache.c b/libafl/syx-snapshot/syx-cow-cache.c
index 47232068fe..b73e9bc9ff 100644
--- a/libafl/syx-snapshot/syx-cow-cache.c
+++ b/libafl/syx-snapshot/syx-cow-cache.c
@@ -218,8 +218,6 @@ void syx_cow_cache_read_entry(SyxCowCache* scc, BlockBackend* blk,
     size_t qiov_offset = 0;
     uint64_t chunk_size = 0;
 
-    //printf("[SYX] [%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);
 
@@ -227,12 +225,9 @@ void syx_cow_cache_read_entry(SyxCowCache* scc, BlockBackend* blk,
     if (!QTAILQ_EMPTY(&scc->layers)) {
         for (; qiov_offset < qiov->size;
              blk_offset += chunk_size, qiov_offset += chunk_size) {
-            int i = 0;
             QTAILQ_FOREACH(layer, &scc->layers, next)
             {
                 chunk_size = layer->chunk_size;
-                //printf("[SYX] check cache layer %d\n", i);
-                i++;
                 if (read_chunk_from_cache_layer(layer, blk, qiov, qiov_offset,
                                                 blk_offset)) {
                     break;