Syx Snapshot rework

- Most of the tables are now GHashtable instances
- Snapshot correctness checking
- Simplified API
- More callbacks to catch more dirty pages

accel/tcg/cputlb.c

@@ -1472,6 +1472,13 @@ 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;
 }
 
@@ -1819,10 +1826,10 @@ static bool mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
 
         //// --- Begin LibAFL code ---
 
-        // TODO: check if the second condition solves faulty dirty address report
-        if (type == MMU_DATA_STORE && !(flags & (TLB_INVALID_MASK | TLB_MMIO))) {
+        // TODO: Does not work?
+        // if (type == MMU_DATA_STORE) {
             syx_snapshot_dirty_list_add_hostaddr(l->page[0].haddr);
-        }
+        // }
 
         //// --- End LibAFL code ---
 
@@ -1850,10 +1857,10 @@ static bool mmu_lookup(CPUState *cpu, vaddr addr, MemOpIdx oi,
 
         //// --- Begin LibAFL code ---
 
-        if (type == MMU_DATA_STORE) {
+        // 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 ---
 

include/exec/ramblock.h

@@ -36,6 +36,9 @@ struct RAMBlock {
     uint32_t flags;
     /* Protected by iothread lock.  */
     char idstr[256];
+//// --- Begin LibAFL code ---
+    guint idstr_hash;
+//// --- End LibAFL code ---
     /* RCU-enabled, writes protected by the ramlist lock */
     QLIST_ENTRY(RAMBlock) next;
     QLIST_HEAD(, RAMBlockNotifier) ramblock_notifiers;

libafl_extras/syx-misc.h

@@ -1,5 +1,6 @@
 #pragma once
-//#ifdef QEMU_SYX
+
+#include "qemu/error-report.h"
 
 #define SYX_PRINTF(format, ...)     fprintf(stderr, ("[QEMU-SYX] " format), ##__VA_ARGS__)
 
@@ -12,5 +13,3 @@
 #define SYX_WARNING(format, ...)   warn_report(("[QEMU-SYX] " format), ##__VA_ARGS__)
 
 #define SYX_ERROR(format, ...)   error_report(("[QEMU-SYX] " format), ##__VA_ARGS__)
-
-//#endif

libafl_extras/syx-snapshot/device-save.c

@@ -18,7 +18,7 @@ extern int vmstate_save(QEMUFile *f, SaveStateEntry *se, JSONWriter *vmdesc);
 extern void save_section_footer(QEMUFile *f, SaveStateEntry *se);
 
 // iothread must be locked
-device_save_state_t* device_save_all(void) {
+DeviceSaveState* device_save_all(void) {
     return device_save_kind(DEVICE_SNAPSHOT_ALL, NULL);
 }
 
@@ -32,8 +32,8 @@ static int is_in_list(char* str, char** list) {
     return 0;
 }
 
-device_save_state_t* device_save_kind(device_snapshot_kind_t kind, char** names) {
-    device_save_state_t* dss = g_new0(device_save_state_t, 1);
+DeviceSaveState* device_save_kind(DeviceSnapshotKind kind, char** names) {
+    DeviceSaveState* dss = g_new0(DeviceSaveState, 1);
     SaveStateEntry *se;
 
     dss->kind = DEVICE_SAVE_KIND_FULL;
@@ -85,7 +85,7 @@ device_save_state_t* device_save_kind(device_snapshot_kind_t kind, char** names)
     return dss;
 }
 
-void device_restore_all(device_save_state_t* dss) {
+void device_restore_all(DeviceSaveState* dss) {
     assert(dss->save_buffer != NULL);
 
     QIOChannelBuffer* bioc = qio_channel_buffer_new_external(dss->save_buffer, QEMU_FILE_RAM_LIMIT, dss->save_buffer_size);
@@ -103,7 +103,7 @@ void device_restore_all(device_save_state_t* dss) {
     qemu_fclose(f);
 }
 
-void device_free_all(device_save_state_t* dss) {
+void device_free_all(DeviceSaveState* dss) {
     g_free(dss->save_buffer);
 }
 

libafl_extras/syx-snapshot/device-save.h

@@ -4,23 +4,23 @@
 
 #define DEVICE_SAVE_KIND_FULL   0
 
-typedef struct device_save_state_s {
+typedef struct DeviceSaveState {
     uint8_t kind;
     uint8_t* save_buffer;
     size_t save_buffer_size;
-} device_save_state_t;
+} DeviceSaveState;
 
 // Type of device snapshot
-typedef enum device_snapshot_kind_e {
+typedef enum DeviceSnapshotKind {
     DEVICE_SNAPSHOT_ALL,
     DEVICE_SNAPSHOT_ALLOWLIST,
     DEVICE_SNAPSHOT_DENYLIST
-} device_snapshot_kind_t;
+} DeviceSnapshotKind;
 
-device_save_state_t* device_save_all(void);
-device_save_state_t* device_save_kind(device_snapshot_kind_t kind, char** names);
+DeviceSaveState* device_save_all(void);
+DeviceSaveState* device_save_kind(DeviceSnapshotKind kind, char** names);
 
-void device_restore_all(device_save_state_t* device_save_state);
-void device_free_all(device_save_state_t* dss);
+void device_restore_all(DeviceSaveState* device_save_state);
+void device_free_all(DeviceSaveState* dss);
 
 char** device_list_all(void);

libafl_extras/syx-snapshot/syx-snapshot.c

@@ -1,36 +1,85 @@
 #include "qemu/osdep.h"
-#include "sysemu/sysemu.h"
-#include "cpu.h"
 #include "qemu/main-loop.h"
-#include "migration/qemu-file.h"
+#include "sysemu/sysemu.h"
 #include "migration/vmstate.h"
-#include "migration/savevm.h"
-#include "memory.h"
+#include "cpu.h"
 
-#include "exec/ram_addr.h"
-/// Physical to host memory in system memory.
 #include "exec/ramlist.h"
-#include "exec/address-spaces.h"
+#include "exec/ram_addr.h"
 #include "exec/exec-all.h"
 
-#include "sysemu/block-backend.h"
-#include "migration/register.h"
-
-// #include "target/i386/cpu.h"
-
 #include "syx-snapshot.h"
-#include "channel-buffer-writeback.h"
 #include "device-save.h"
 
 #define SYX_SNAPSHOT_LIST_INIT_SIZE      4096
 #define SYX_SNAPSHOT_LIST_GROW_FACTOR    2
+#define TARGET_NEXT_PAGE_ADDR(p) \
+    ((typeof(p))(((uintptr_t) p + TARGET_PAGE_SIZE) & TARGET_PAGE_MASK))
 
-syx_snapshot_state_t syx_snapshot_state = {0};
+
+SyxSnapshotState syx_snapshot_state = {0};
 static MemoryRegion* mr_to_enable = NULL;
 
-void syx_snapshot_init(void) {
-    //syx_snapshot_init_params_t* params = (syx_snapshot_init_params_t*) opaque;
-    //uint64_t page_size = params->page_size;
+static void destroy_ramblock_snapshot(gpointer root_snapshot);
+static void syx_snapshot_dirty_list_flush(SyxSnapshot* snapshot);
+
+static void rb_save_dirty_addr_to_table(gpointer offset_within_rb, gpointer unused, gpointer rb_dirty_list_to_page_args_ptr);
+static void rb_dirty_list_to_dirty_pages(gpointer rb_idstr_hash, gpointer rb_dirty_list_hash_table_ptr, gpointer rbs_dirty_pages_ptr);
+static inline void syx_snapshot_dirty_list_add_internal(RAMBlock* rb, ram_addr_t offset);
+static void empty_rb_dirty_list(gpointer rb_idstr_hash, gpointer rb_dirty_list_hash_table_ptr, gpointer user_data);
+static void destroy_snapshot_dirty_page_list(gpointer snapshot_dirty_page_list_ptr);
+
+static void root_restore_rb_page(gpointer offset_within_rb, gpointer unused, gpointer root_restore_args_ptr);
+static void root_restore_rb(gpointer rb_idstr_hash, gpointer rb_dirty_pages_hash_table_ptr, gpointer snapshot_ptr);
+static void root_restore_check_memory_rb(gpointer rb_idstr_hash, gpointer rb_dirty_pages_hash_table_ptr, gpointer snapshot_ptr);
+
+static SyxSnapshotIncrement* syx_snapshot_increment_free(SyxSnapshotIncrement* increment);
+
+static RAMBlock* ramblock_lookup(gpointer rb_idstr_hash)
+{
+    RAMBlock* block;
+    RAMBLOCK_FOREACH(block) {
+        if (rb_idstr_hash == GINT_TO_POINTER(block->idstr_hash)) {
+            return block;
+        }
+    }
+
+    return NULL;
+}
+
+// Root snapshot API
+static SyxSnapshotRoot syx_snapshot_root_new(DeviceSnapshotKind kind, char** devices);
+static void syx_snapshot_root_free(SyxSnapshotRoot* root);
+
+struct rb_dirty_list_to_page_args {
+    RAMBlock* rb;
+    SyxSnapshotDirtyPageList* dirty_page_list;
+    uint64_t* table_idx;
+};
+
+struct rb_page_root_restore_args {
+    RAMBlock* rb;
+    SyxSnapshotRAMBlock* snapshot_rb;
+};
+
+struct rb_increment_restore_args {
+    SyxSnapshot* snapshot;
+    SyxSnapshotIncrement* increment;
+};
+
+struct rb_page_increment_restore_args {
+    RAMBlock* rb;
+    SyxSnapshot* snapshot;
+    SyxSnapshotIncrement* increment;
+};
+
+struct rb_check_memory_args {
+    SyxSnapshot* snapshot; // IN
+    uint64_t nb_inconsistent_pages; // OUT
+};
+
+void syx_snapshot_init(void)
+{
     uint64_t page_size = TARGET_PAGE_SIZE;
 
     syx_snapshot_state.page_size = page_size;
@@ -41,92 +90,103 @@ void syx_snapshot_init(void) {
     syx_snapshot_state.is_enabled = false;
 }
 
-uint64_t syx_snapshot_handler(CPUState* cpu, uint32_t cmd, target_ulong target_opaque) {
-    return (uint64_t) -1;
-}
+SyxSnapshot* syx_snapshot_new(bool track, DeviceSnapshotKind kind, char** devices)
+{
+    SyxSnapshot* snapshot = g_new0(SyxSnapshot, 1);
 
-syx_snapshot_t* syx_snapshot_create(bool track, device_snapshot_kind_t kind, char** devices) {
-    syx_snapshot_t* snapshot = g_new0(syx_snapshot_t, 1);
-
-    snapshot->root_snapshot = syx_snapshot_root_create(kind, devices);
+    snapshot->root_snapshot = syx_snapshot_root_new(kind, devices);
     snapshot->last_incremental_snapshot = NULL;
-    snapshot->dirty_list = syx_snapshot_dirty_list_create();
+    snapshot->rbs_dirty_list = g_hash_table_new_full(g_direct_hash, g_direct_equal, NULL, (GDestroyNotify) g_hash_table_remove_all);
 
     if (track) {
         syx_snapshot_track(&syx_snapshot_state.tracked_snapshots, snapshot);
     }
 
+#ifdef CONFIG_DEBUG_TCG
+    SYX_PRINTF("[Snapshot Creation] Checking snapshot memory consistency\n");
+    g_hash_table_foreach(snapshot->rbs_dirty_list, root_restore_check_memory_rb, snapshot);
+    SYX_PRINTF("[Snapshot Creation] Memory is consistent.\n");
+#endif
+
     syx_snapshot_state.is_enabled = true;
 
     return snapshot;
 }
 
-void syx_snapshot_free(syx_snapshot_t* snapshot) {
-    syx_snapshot_increment_t* increment = snapshot->last_incremental_snapshot;
+void syx_snapshot_free(SyxSnapshot* snapshot)
+{
+    SyxSnapshotIncrement* increment = snapshot->last_incremental_snapshot;
 
     while (increment != NULL) {
         increment = syx_snapshot_increment_free(increment);
     }
 
-    syx_snapshot_dirty_list_free(&snapshot->dirty_list);
+    g_hash_table_remove_all(snapshot->rbs_dirty_list);
+
     syx_snapshot_root_free(&snapshot->root_snapshot);
 
     g_free(snapshot);
 }
 
-syx_snapshot_root_t syx_snapshot_root_create(device_snapshot_kind_t kind, char** devices) {
-    syx_snapshot_root_t root = {0};
+static void destroy_ramblock_snapshot(gpointer root_snapshot)
+{
+    SyxSnapshotRAMBlock* snapshot_rb = root_snapshot;
+
+    g_free(snapshot_rb->ram);
+    g_free(snapshot_rb);
+}
+
+static SyxSnapshotRoot syx_snapshot_root_new(DeviceSnapshotKind kind, char** devices)
+{
+    SyxSnapshotRoot root = {0};
 
     RAMBlock* block;
-    uint64_t nb_blocks = 0;
-    device_save_state_t* dss = device_save_kind(kind, devices);
+    RAMBlock* inner_block;
+    DeviceSaveState* dss = device_save_kind(kind, devices);
 
-    RAMBLOCK_FOREACH(block) {
-        nb_blocks++;
-    }
-
-    root.ram_blocks = g_new0(syx_snapshot_ramblock_t, nb_blocks);
-    root.nb_ram_blocks = nb_blocks;
+    root.rbs_snapshot = g_hash_table_new_full(g_direct_hash, g_direct_equal, NULL, destroy_ramblock_snapshot);
     root.dss = dss;
 
-    uint64_t ram_block_idx = 0;
     RAMBLOCK_FOREACH(block) {
-        syx_snapshot_ramblock_t* snapshot_ram_block = &root.ram_blocks[ram_block_idx];
-        strcpy(snapshot_ram_block->idstr, block->idstr);
-        snapshot_ram_block->used_length = block->used_length;
+        RAMBLOCK_FOREACH(inner_block) {
+            if (block != inner_block && inner_block->idstr_hash == block->idstr_hash) {
+                SYX_ERROR("Hash collision detected on RAMBlocks %s and %s, snapshotting will not work correctly.", inner_block->idstr, block->idstr);
+                exit(1);
+            }
+        }
 
-        snapshot_ram_block->ram = g_new(uint8_t, block->used_length);
-        memcpy(snapshot_ram_block->ram, block->host, block->used_length);
+        SyxSnapshotRAMBlock* snapshot_rb = g_new(SyxSnapshotRAMBlock, 1);
+        snapshot_rb->used_length = block->used_length;
+        snapshot_rb->ram = g_new(uint8_t, block->used_length);
+        memcpy(snapshot_rb->ram, block->host, block->used_length);
 
-        ram_block_idx++;
+        g_hash_table_insert(root.rbs_snapshot, GINT_TO_POINTER(block->idstr_hash), snapshot_rb);
     }
-    assert(ram_block_idx == nb_blocks);
 
     return root;
 }
 
-void syx_snapshot_root_free(syx_snapshot_root_t* root) {
-    for (uint64_t i = 0; i < root->nb_ram_blocks; ++i) {
-        g_free(root->ram_blocks[i].ram);
-    }
-
-    g_free(root->ram_blocks);
+static void syx_snapshot_root_free(SyxSnapshotRoot* root)
+{
+    g_hash_table_destroy(root->rbs_snapshot);
 }
 
-syx_snapshot_tracker_t syx_snapshot_tracker_init(void) {
-    syx_snapshot_tracker_t tracker = {
+SyxSnapshotTracker syx_snapshot_tracker_init(void)
+{
+    SyxSnapshotTracker tracker = {
         .length = 0,
         .capacity = SYX_SNAPSHOT_LIST_INIT_SIZE,
-        .tracked_snapshots = g_new(syx_snapshot_t*, SYX_SNAPSHOT_LIST_INIT_SIZE)
+        .tracked_snapshots = g_new(SyxSnapshot*, SYX_SNAPSHOT_LIST_INIT_SIZE)
     };
 
     return tracker;
 }
 
-void syx_snapshot_track(syx_snapshot_tracker_t* tracker, syx_snapshot_t* snapshot) {
+void syx_snapshot_track(SyxSnapshotTracker* tracker, SyxSnapshot* snapshot)
+{
     if (tracker->length == tracker->capacity) {
         tracker->capacity *= SYX_SNAPSHOT_LIST_GROW_FACTOR;
-        tracker->tracked_snapshots = g_realloc(tracker->tracked_snapshots, tracker->capacity * sizeof(syx_snapshot_t*));
+        tracker->tracked_snapshots = g_realloc(tracker->tracked_snapshots, tracker->capacity * sizeof(SyxSnapshot*));
     }
 
     assert(tracker->length < tracker->capacity);
@@ -135,7 +195,8 @@ void syx_snapshot_track(syx_snapshot_tracker_t* tracker, syx_snapshot_t* snapsho
     tracker->length++;
 }
 
-void syx_snapshot_stop_track(syx_snapshot_tracker_t* tracker, syx_snapshot_t* snapshot) {
+void syx_snapshot_stop_track(SyxSnapshotTracker* tracker, SyxSnapshot* snapshot)
+{
     for (uint64_t i = 0; i < tracker->length; ++i) {
         if (tracker->tracked_snapshots[i] == snapshot) {
             for (uint64_t j = i + i; j < tracker->length; ++j) {
@@ -150,217 +211,203 @@ void syx_snapshot_stop_track(syx_snapshot_tracker_t* tracker, syx_snapshot_t* sn
     abort();
 }
 
-void syx_snapshot_increment_push(syx_snapshot_t* snapshot, device_snapshot_kind_t kind, char** devices) {
-    syx_snapshot_increment_t* increment = g_new0(syx_snapshot_increment_t, 1);
+static void rb_save_dirty_addr_to_table(gpointer offset_within_rb, gpointer unused, gpointer rb_dirty_list_to_page_args_ptr)
+{
+    struct rb_dirty_list_to_page_args* args = rb_dirty_list_to_page_args_ptr;
+    RAMBlock* rb = args->rb;
+    SyxSnapshotDirtyPage* dirty_page = &args->dirty_page_list->dirty_pages[*args->table_idx];
+    dirty_page->offset_within_rb = (ram_addr_t) offset_within_rb;
+
+    memcpy((gpointer) dirty_page->data, rb->host + (ram_addr_t) offset_within_rb, syx_snapshot_state.page_size);
+
+    *args->table_idx += 1;
+}
+
+static void rb_dirty_list_to_dirty_pages(gpointer rb_idstr_hash, gpointer rb_dirty_list_hash_table_ptr, gpointer rbs_dirty_pages_ptr)
+{
+    GHashTable* rbs_dirty_pages = rbs_dirty_pages_ptr;
+    GHashTable* rb_dirty_list = rb_dirty_list_hash_table_ptr;
+
+    RAMBlock* rb = ramblock_lookup(rb_idstr_hash);
+
+    if (rb) {
+            SyxSnapshotDirtyPageList* dirty_page_list = g_new(SyxSnapshotDirtyPageList, 1);
+            dirty_page_list->length = g_hash_table_size(rb_dirty_list);
+            dirty_page_list->dirty_pages = g_new(SyxSnapshotDirtyPage, dirty_page_list->length);
+
+            uint64_t* ctr = g_new0(uint64_t, 1);
+
+            struct rb_dirty_list_to_page_args dirty_list_to_page_args = {
+                .rb = rb,
+                .table_idx = ctr,
+                .dirty_page_list = dirty_page_list
+            };
+
+            g_hash_table_foreach(rbs_dirty_pages, rb_save_dirty_addr_to_table, &dirty_list_to_page_args);
+
+            g_free(dirty_list_to_page_args.table_idx);
+    } else {
+        SYX_ERROR("Impossible to find RAMBlock with pages marked as dirty.");
+    }
+}
+
+static void destroy_snapshot_dirty_page_list(gpointer snapshot_dirty_page_list_ptr)
+{
+    SyxSnapshotDirtyPageList* snapshot_dirty_page_list = snapshot_dirty_page_list_ptr;
+
+    for (uint64_t i = 0; i < snapshot_dirty_page_list->length; ++i) {
+        g_free(snapshot_dirty_page_list->dirty_pages[i].data);
+    }
+
+    g_free(snapshot_dirty_page_list->dirty_pages);
+    g_free(snapshot_dirty_page_list);
+}
+
+void syx_snapshot_increment_push(SyxSnapshot* snapshot, DeviceSnapshotKind kind, char** devices)
+{
+    SyxSnapshotIncrement* increment = g_new0(SyxSnapshotIncrement, 1);
     increment->parent = snapshot->last_incremental_snapshot;
     snapshot->last_incremental_snapshot = increment;
 
-    increment->dirty_page_list = syx_snapshot_dirty_list_to_dirty_page_list(&snapshot->dirty_list);
+    increment->rbs_dirty_pages = g_hash_table_new_full(g_direct_hash, g_direct_equal, NULL, destroy_snapshot_dirty_page_list);
+    g_hash_table_foreach(snapshot->rbs_dirty_list, rb_dirty_list_to_dirty_pages, increment->rbs_dirty_pages);
     increment->dss = device_save_kind(kind, devices);
 
-    syx_snapshot_dirty_list_flush(&snapshot->dirty_list);
+    g_hash_table_remove_all(snapshot->rbs_dirty_list);
 }
 
-static syx_snapshot_dirty_page_t* get_dirty_page_from_addr_rec(syx_snapshot_increment_t* increment, hwaddr addr) {
+static SyxSnapshotDirtyPage* get_dirty_page_from_addr_rec(SyxSnapshotIncrement* increment, RAMBlock* rb, ram_addr_t offset_within_rb)
+{
     if (increment == NULL) {
         return NULL;
     }
 
-    for (uint64_t i = 0; i < increment->dirty_page_list.length; ++i) {
-        if (increment->dirty_page_list.dirty_pages[i].addr == addr) {
-            return &increment->dirty_page_list.dirty_pages[i];
+    SyxSnapshotDirtyPageList* dpl = g_hash_table_lookup(increment->rbs_dirty_pages, GINT_TO_POINTER(rb->idstr_hash));
+
+    if (dpl) {
+        for (uint64_t i = 0; i < dpl->length; ++i) {
+            if (dpl->dirty_pages[i].offset_within_rb == offset_within_rb) {
+                return &dpl->dirty_pages[i];
+            }
         }
     }
 
-    return  get_dirty_page_from_addr_rec(increment->parent, addr);
+    return  get_dirty_page_from_addr_rec(increment->parent, rb, offset_within_rb);
 }
 
-static syx_snapshot_ramblock_t* find_ramblock(syx_snapshot_root_t* root, char* idstr) {
-    for (size_t i = 0; i < root->nb_ram_blocks; i++) {
-        if (!strcmp(idstr, root->ram_blocks[i].idstr)) {
-            return &root->ram_blocks[i];
-        }
-    }
 
-    return NULL;
-}
+static void restore_dirty_page_to_increment(gpointer offset_within_rb, gpointer _unused, gpointer args_ptr) {
+    struct rb_page_increment_restore_args* args = args_ptr;
+    RAMBlock* rb = args->rb;
+    SyxSnapshot* snapshot = args->snapshot;
+    SyxSnapshotIncrement* increment = args->increment;
+    ram_addr_t offset = (ram_addr_t) offset_within_rb;
 
-static void restore_page_from_root(syx_snapshot_root_t* root, MemoryRegion* mr, hwaddr addr) {
-    MemoryRegionSection mr_section = memory_region_find(mr, addr, syx_snapshot_state.page_size);
+    SyxSnapshotDirtyPage* dp = get_dirty_page_from_addr_rec(increment, rb, offset);
 
-    if (mr_section.size == 0) {
-        assert(mr_section.mr == NULL);
-
-        SYX_PRINTF("Did not found a memory region while restoring the address %p from root snapshot.\n", (void*) addr);
-        return;
-    }
-
-    if (mr_section.mr->ram) {
-        syx_snapshot_ramblock_t* ram_block = find_ramblock(root, mr_section.mr->ram_block->idstr);
-        assert(ram_block != NULL);
-        assert(!strcmp(mr_section.mr->ram_block->idstr, ram_block->idstr));
-
-        memcpy(mr_section.mr->ram_block->host + mr_section.offset_within_region,
-                ram_block->ram + mr_section.offset_within_region, syx_snapshot_state.page_size);
+    if (dp) {
+        memcpy(rb->host + offset, dp->data, syx_snapshot_state.page_size);
     } else {
-        if (!strcmp(mr_section.mr->name, "tseg-blackhole")) {
-            assert(mr_to_enable == NULL);
-            memory_region_set_enabled(mr_section.mr, false);
-            mr_to_enable = mr_section.mr;
-            restore_page_from_root(root, mr, addr);
-        } else {
-            // SYX_WARNING("[SYX SNAPSHOT RESTORE] Ram section not found for page @ 0x%lx (mr %s...)\n", addr, mr_section.mr->name);
-        }
-    }
-}
-//
-static void restore_page(MemoryRegion* mr, syx_snapshot_dirty_page_t* page) {
-    MemoryRegionSection mr_section = memory_region_find(mr, page->addr, syx_snapshot_state.page_size);
-    assert(mr_section.size != 0 && mr_section.mr != NULL);
-    assert(mr_section.mr->ram);
+        SyxSnapshotRAMBlock* rrb = g_hash_table_lookup(snapshot->root_snapshot.rbs_snapshot, GINT_TO_POINTER(rb->idstr_hash));
+        assert(rrb);
 
-    memcpy(mr_section.mr->ram_block->host + mr_section.offset_within_region, page->data, syx_snapshot_state.page_size);
-}
-
-static void restore_to_last_increment(syx_snapshot_t* snapshot, MemoryRegion* mr) {
-    syx_snapshot_increment_t* increment = snapshot->last_incremental_snapshot;
-    syx_snapshot_dirty_list_t* dirty_list = &snapshot->dirty_list;
-
-    for (uint64_t i = 0; i < dirty_list->length; ++i) {
-        syx_snapshot_dirty_page_t* dirty_page = get_dirty_page_from_addr_rec(increment, dirty_list->dirty_addr[i]);
-        if (dirty_page == NULL) {
-            restore_page_from_root(&snapshot->root_snapshot, mr, dirty_list->dirty_addr[i]);
-        } else {
-            restore_page(mr, dirty_page);
-        }
+        memcpy(rb->host + offset, rrb->ram, syx_snapshot_state.page_size);
     }
 }
 
-void syx_snapshot_increment_pop(syx_snapshot_t* snapshot) {
-    MemoryRegion* system_mr = get_system_memory();
-    syx_snapshot_increment_t* last_increment = snapshot->last_incremental_snapshot;
+static void restore_rb_to_increment(gpointer rb_idstr_hash, gpointer rb_dirty_pages_hash_table_ptr, gpointer args_ptr) {
+    struct rb_increment_restore_args* args = args_ptr;
+    GHashTable* rb_dirty_pages_hash_table = rb_dirty_pages_hash_table_ptr;
 
-    restore_to_last_increment(snapshot, system_mr);
+    RAMBlock* rb = ramblock_lookup(rb_idstr_hash);
+    struct rb_page_increment_restore_args page_args = {
+        .snapshot = args->snapshot,
+        .increment = args->increment,
+        .rb = rb
+    };
+
+    g_hash_table_foreach(rb_dirty_pages_hash_table, restore_dirty_page_to_increment, &page_args);
+}
+
+static void restore_to_increment(SyxSnapshot* snapshot, SyxSnapshotIncrement* increment)
+{
+    struct rb_increment_restore_args args = {
+        .snapshot = snapshot,
+        .increment = increment
+    };
+
+    g_hash_table_foreach(snapshot->rbs_dirty_list, restore_rb_to_increment, &args);
+}
+
+void syx_snapshot_increment_pop(SyxSnapshot* snapshot)
+{
+    SyxSnapshotIncrement* last_increment = snapshot->last_incremental_snapshot;
 
     device_restore_all(last_increment->dss);
-
-    syx_snapshot_dirty_list_flush(&snapshot->dirty_list);
+    restore_to_increment(snapshot, last_increment);
 
     snapshot->last_incremental_snapshot = last_increment->parent;
     syx_snapshot_increment_free(last_increment);
+
+    syx_snapshot_dirty_list_flush(snapshot);
 }
 
-void syx_snapshot_increment_restore_last(syx_snapshot_t* snapshot) {
-    MemoryRegion* system_mr = get_system_memory();
-    syx_snapshot_increment_t* last_increment = snapshot->last_incremental_snapshot;
-
-    restore_to_last_increment(snapshot, system_mr);
+void syx_snapshot_increment_restore_last(SyxSnapshot* snapshot)
+{
+    SyxSnapshotIncrement* last_increment = snapshot->last_incremental_snapshot;
 
     device_restore_all(last_increment->dss);
+    restore_to_increment(snapshot, last_increment);
 
-    syx_snapshot_dirty_list_flush(&snapshot->dirty_list);
+    syx_snapshot_dirty_list_flush(snapshot);
 }
 
-syx_snapshot_increment_t* syx_snapshot_increment_free(syx_snapshot_increment_t* increment) {
-    syx_snapshot_increment_t* parent = increment->parent;
-
-    for (uint64_t i = 0; i < increment->dirty_page_list.length; ++i) {
-        g_free(increment->dirty_page_list.dirty_pages[i].data);
-    }
-
+static SyxSnapshotIncrement* syx_snapshot_increment_free(SyxSnapshotIncrement* increment)
+{
+    SyxSnapshotIncrement* parent_increment = increment->parent;
+    g_hash_table_destroy(increment->rbs_dirty_pages);
     device_free_all(increment->dss);
     g_free(increment);
-
-    return parent;
+    return parent_increment;
 }
 
-syx_snapshot_dirty_list_t syx_snapshot_dirty_list_create(void) {
-    syx_snapshot_dirty_list_t dirty_list = {
-        .length = 0,
-        .capacity = SYX_SNAPSHOT_LIST_INIT_SIZE,
-        .dirty_addr = g_new(hwaddr, SYX_SNAPSHOT_LIST_INIT_SIZE)
-    };
-
-    return dirty_list;
+static void syx_snapshot_dirty_list_flush(SyxSnapshot* snapshot)
+{
+    g_hash_table_foreach(snapshot->rbs_dirty_list, empty_rb_dirty_list, (gpointer) snapshot);
 }
 
-void syx_snapshot_dirty_list_free(syx_snapshot_dirty_list_t* dirty_list) {
-    g_free(dirty_list->dirty_addr);
-}
-
-static inline syx_snapshot_dirty_page_t* syx_snapshot_save_page_from_addr(MemoryRegion* mr, hwaddr addr) {
-    syx_snapshot_dirty_page_t* dirty_page = g_new(syx_snapshot_dirty_page_t, 1);
-    
-    dirty_page->addr = addr;
-    dirty_page->data = g_new(uint8_t, syx_snapshot_state.page_size);
-    
-    MemoryRegionSection mr_section = memory_region_find(mr, addr, syx_snapshot_state.page_size);
-
-    assert(mr_section.size != 0 && mr_section.mr != NULL);
-
-    if (!mr_section.mr->ram) {
-        return NULL;
-    }
-    
-    memcpy(dirty_page->data, mr_section.mr->ram_block->host + mr_section.offset_within_region, syx_snapshot_state.page_size);
-    return dirty_page;
-}
-
-syx_snapshot_dirty_page_list_t syx_snapshot_dirty_list_to_dirty_page_list(syx_snapshot_dirty_list_t* dirty_list) {
-    syx_snapshot_dirty_page_list_t dirty_page_list = {
-        .length = dirty_list->length,
-        .dirty_pages = g_new(syx_snapshot_dirty_page_t, dirty_list->length)
-    };
-    MemoryRegion* system_mr = get_system_memory();
-
-    uint64_t real_len = 0;
-    for (uint64_t i = 0; i < dirty_page_list.length; ++i) {
-        syx_snapshot_dirty_page_t* page = syx_snapshot_save_page_from_addr(system_mr, dirty_list->dirty_addr[i]);
-        if (page == NULL) {
-            continue;
-        }
-        dirty_page_list.dirty_pages[real_len] = *page;
-        real_len++;
-        g_free(page);
-    }
-
-    dirty_page_list.length = real_len;
-
-    return dirty_page_list;
-}
-
-static inline void syx_snapshot_dirty_list_add_internal(hwaddr paddr) {
-    paddr &= syx_snapshot_state.page_mask;
+static inline void syx_snapshot_dirty_list_add_internal(RAMBlock* rb, ram_addr_t offset)
+{
+    assert((offset & syx_snapshot_state.page_mask) == offset); // offsets should always be page-aligned.
 
     for (uint64_t i = 0; i < syx_snapshot_state.tracked_snapshots.length; ++i) {
-        syx_snapshot_dirty_list_t* dirty_list = &syx_snapshot_state.tracked_snapshots.tracked_snapshots[i]->dirty_list;
+        SyxSnapshot* snapshot = syx_snapshot_state.tracked_snapshots.tracked_snapshots[i];
 
-        // Avoid adding already marked addresses
-        for (uint64_t j = 0; j < dirty_list->length; ++j) {
-            if (dirty_list->dirty_addr[j] == paddr) {
-                goto next_snapshot_dirty_list;
-            }
+        GHashTable* rb_dirty_list = g_hash_table_lookup(snapshot->rbs_dirty_list, GINT_TO_POINTER(rb->idstr_hash));
+
+        if (unlikely(!rb_dirty_list)) {
+#ifdef SYX_SNAPSHOT_DEBUG
+            printf("rb_dirty_list did not exit, creating...\n");
+#endif
+            rb_dirty_list = g_hash_table_new(g_direct_hash, g_direct_equal);
+            g_hash_table_insert(snapshot->rbs_dirty_list, GINT_TO_POINTER(rb->idstr_hash), rb_dirty_list);
         }
 
-        if (dirty_list->length == dirty_list->capacity) {
-            //SYX_PRINTF("[DL %lu] Reallocation %lu ---> %lu\n", i, dirty_list->length, dirty_list->length * SYX_SNAPSHOT_LIST_GROW_FACTOR);
-            dirty_list->capacity *= SYX_SNAPSHOT_LIST_GROW_FACTOR;
-            dirty_list->dirty_addr = g_realloc(dirty_list->dirty_addr, dirty_list->capacity * sizeof(hwaddr));
+        if (g_hash_table_add(rb_dirty_list, GINT_TO_POINTER(offset))) {
+#ifdef SYX_SNAPSHOT_DEBUG
+            SYX_PRINTF("[%s] Marking offset 0x%lx as dirty\n", rb->idstr, offset);
+#endif
         }
-
-        dirty_list->dirty_addr[dirty_list->length] = paddr;
-        dirty_list->length++;
-
-        // SYX_PRINTF("[DL %lu] Added dirty page @addr 0x%lx. dirty list len: %lu\n", i, paddr, dirty_list->length);
-    next_snapshot_dirty_list:;
     }
 }
 
-bool syx_snapshot_is_enabled(void) {
+bool syx_snapshot_is_enabled(void)
+{
     return syx_snapshot_state.is_enabled;
 }
 
 /*
+// TODO: Check if using this method is better for performances.
 // The implementation is pretty bad, it would be nice to store host addr directly for
 // the memcopy happening later on.
 __attribute__((target("no-3dnow,no-sse,no-mmx"),no_caller_saved_registers)) void syx_snapshot_dirty_list_add_tcg_target(uint64_t dummy, void* host_addr) {
@@ -372,19 +419,17 @@ __attribute__((target("no-3dnow,no-sse,no-mmx"),no_caller_saved_registers)) void
     ram_addr_t offset;
     RAMBlock* rb = qemu_ram_block_from_host((void*) host_addr, true, &offset);
 
-    assert(rb);
+    if (!rb) {
+        return;
+    }
 
-    hwaddr paddr = rb->mr->addr + offset;
-    // If this assert is ever false, please understand why
-    // qemu_ram_block_from_host result with true as second
-    // param would not be host page aligned.
-    assert(paddr == (paddr & syx_snapshot_state.page_mask));
-
-    syx_snapshot_dirty_list_add_internal(paddr);
+    syx_snapshot_dirty_list_add_internal(rb, offset);
 }
 */
 
-void syx_snapshot_dirty_list_add_hostaddr(void* host_addr) {
+// host_addr should be page-aligned.
+void syx_snapshot_dirty_list_add_hostaddr(void* host_addr)
+{
     // early check to know whether we should log the page access or not
     if (!syx_snapshot_is_enabled()) {
         return;
@@ -393,59 +438,167 @@ void syx_snapshot_dirty_list_add_hostaddr(void* host_addr) {
     ram_addr_t offset;
     RAMBlock* rb = qemu_ram_block_from_host((void*) host_addr, true, &offset);
 
-#ifdef CONFIG_DEBUG_TCG
-    assert(rb);
-#else
+#ifdef SYX_SNAPSHOT_DEBUG
+    SYX_PRINTF("Should mark offset 0x%lx as dirty\n", offset);
+#endif
+
     if (!rb) {
         return;
     }
+
+    syx_snapshot_dirty_list_add_internal(rb, offset);
+}
+
+void syx_snapshot_dirty_list_add_hostaddr_range(void* host_addr, uint64_t len)
+{
+    assert(len < INT64_MAX);
+    int64_t len_signed = (int64_t) len;
+
+    syx_snapshot_dirty_list_add_hostaddr(QEMU_ALIGN_PTR_DOWN(host_addr, syx_snapshot_state.page_size));
+    void* next_page_addr = TARGET_NEXT_PAGE_ADDR(host_addr);
+    assert(next_page_addr > host_addr);
+    assert(QEMU_PTR_IS_ALIGNED(next_page_addr, TARGET_PAGE_SIZE));
+
+    int64_t len_to_next_page = next_page_addr - host_addr;
+
+    host_addr += len_to_next_page;
+    len_signed -= len_to_next_page;
+
+    while(len_signed > 0) {
+        assert(QEMU_PTR_IS_ALIGNED(host_addr, TARGET_PAGE_SIZE));
+
+        syx_snapshot_dirty_list_add_hostaddr(host_addr);
+        len_signed -= TARGET_PAGE_SIZE;
+    }
+}
+
+static void empty_rb_dirty_list(gpointer _rb_idstr_hash, gpointer rb_dirty_list_hash_table_ptr, gpointer _user_data)
+{
+    GHashTable* rb_dirty_hash_table = rb_dirty_list_hash_table_ptr;
+    g_hash_table_remove_all(rb_dirty_hash_table);
+}
+
+static void root_restore_rb_page(gpointer offset_within_rb, gpointer _unused, gpointer root_restore_args_ptr)
+{
+    struct rb_page_root_restore_args* args = root_restore_args_ptr;
+    RAMBlock* rb = args->rb;
+    SyxSnapshotRAMBlock* snapshot_rb = args->snapshot_rb;
+
+
+    // safe cast because ram_addr_t is also an alias to void*
+    void* host_rb_restore = rb->host + (ram_addr_t) offset_within_rb;
+    void* host_snapshot_rb_restore = (gpointer) snapshot_rb->ram + (ram_addr_t) offset_within_rb;
+
+#ifdef SYX_SNAPSHOT_DEBUG
+    SYX_PRINTF("\t[%s] Restore at offset 0x%lx of size %lu...\n", rb->idstr, (uint64_t) offset_within_rb, syx_snapshot_state.page_size);
 #endif
 
-    hwaddr paddr = rb->mr->addr + offset;
-    // If this assert is ever false, please understand why
-    // qemu_ram_block_from_host result with true as second
-    // param would not be host page aligned.
-    assert(paddr == (paddr & syx_snapshot_state.page_mask));
-
-    syx_snapshot_dirty_list_add_internal(paddr);
+    memcpy(host_rb_restore, host_snapshot_rb_restore, syx_snapshot_state.page_size);
+    //TODO: manage special case of TSEG.
 }
 
+static void root_restore_rb(gpointer rb_idstr_hash, gpointer rb_dirty_pages_hash_table_ptr, gpointer snapshot_ptr)
+{
+    SyxSnapshot* snapshot = snapshot_ptr;
+    GHashTable* rb_dirty_pages_hash_table = rb_dirty_pages_hash_table_ptr;
+    RAMBlock* rb = ramblock_lookup(rb_idstr_hash);
 
+    if (rb) {
+        SyxSnapshotRAMBlock* snapshot_ramblock = g_hash_table_lookup(snapshot->root_snapshot.rbs_snapshot, rb_idstr_hash);
 
-void syx_snapshot_dirty_list_add(hwaddr paddr) {
-    if (!syx_snapshot_is_enabled()) {
-        return;
-    }
+        struct rb_page_root_restore_args root_restore_args = {
+            .rb = rb,
+            .snapshot_rb = snapshot_ramblock
+        };
 
-    syx_snapshot_dirty_list_add_internal(paddr);
-}
+#ifdef CONFIG_DEBUG_TCG
+        SYX_PRINTF("Restoring RB %s...\n", rb->idstr);
+#endif
 
-inline void syx_snapshot_dirty_list_flush(syx_snapshot_dirty_list_t* dirty_list) {
-    dirty_list->length = 0;
-}
+        g_hash_table_foreach(rb_dirty_pages_hash_table, root_restore_rb_page, &root_restore_args);
 
-static void syx_snapshot_restore_root_from_dirty_list(syx_snapshot_root_t* root, MemoryRegion* mr, syx_snapshot_dirty_list_t* dirty_list) {
-    for (size_t i = 0; i < dirty_list->length; ++i) {
-        restore_page_from_root(root, mr, dirty_list->dirty_addr[i]);
+#ifdef CONFIG_DEBUG_TCG
+        SYX_PRINTF("Finished to restore RB %s\n", rb->idstr);
+#endif
+    } else {
+        SYX_ERROR("Saved RAMBlock not found.");
+        exit(1);
     }
 }
 
-void syx_snapshot_root_restore(syx_snapshot_t* snapshot) {
+static void root_restore_check_memory_rb(gpointer rb_idstr_hash, gpointer rb_dirty_pages_hash_table_ptr, gpointer check_memory_args_ptr)
+{
+    struct rb_check_memory_args* args = check_memory_args_ptr;
+    SyxSnapshot* snapshot = args->snapshot;
+    RAMBlock* rb = ramblock_lookup(rb_idstr_hash);
+
+    args->nb_inconsistent_pages = 0;
+    if (rb) {
+        SYX_PRINTF("Checking memory consistency of %s... ", rb->idstr);
+        SyxSnapshotRAMBlock* rb_snapshot = g_hash_table_lookup(snapshot->root_snapshot.rbs_snapshot, rb_idstr_hash);
+        assert(rb_snapshot);
+
+        assert(rb->used_length == rb_snapshot->used_length);
+
+        for (uint64_t i = 0; i < rb->used_length; i += syx_snapshot_state.page_size) {
+            if (memcmp(rb->host + i, rb_snapshot->ram + i, syx_snapshot_state.page_size) != 0) {
+                SYX_ERROR("\nFound incorrect page at offset 0x%lx\n", i);
+                for (uint64_t j = 0; j < syx_snapshot_state.page_size; j++) {
+                    if (*(rb->host + i + j) != *(rb_snapshot->ram + i + j)) {
+                        SYX_ERROR("\t- byte at address 0x%lx differs\n", i + j);
+                    }
+                }
+                args->nb_inconsistent_pages++;
+            }
+
+        }
+
+        if (args->nb_inconsistent_pages > 0) {
+            SYX_ERROR("[%s] Found %lu page %s.\n", rb->idstr, args->nb_inconsistent_pages, args->nb_inconsistent_pages > 1 ? "inconsistencies" : "inconsistency");
+        } else {
+            SYX_PRINTF("OK.\n");
+        }
+    } else {
+        SYX_ERROR("RB not found...\n");
+        exit(1);
+    }
+}
+
+uint64_t syx_snapshot_check_memory_consistency(SyxSnapshot* snapshot)
+{
+    struct rb_check_memory_args args = {
+        .snapshot = snapshot
+    };
+    g_hash_table_foreach(snapshot->rbs_dirty_list, root_restore_check_memory_rb, &args);
+    return args.nb_inconsistent_pages;
+}
+
+void syx_snapshot_root_restore(SyxSnapshot* snapshot)
+{
+    // health check.
+    CPUState* cpu;
+    CPU_FOREACH(cpu) {
+        assert(cpu->stopped);
+    }
+
 	bool must_unlock_iothread = false;
-    MemoryRegion* system_mr = get_system_memory();
 
 	if (!qemu_mutex_iothread_locked()) {
 		qemu_mutex_lock_iothread();
 		must_unlock_iothread = true;
 	}
 
-    syx_snapshot_restore_root_from_dirty_list(&snapshot->root_snapshot, system_mr, &snapshot->dirty_list);
+    // In case, we first restore devices if there is a modification of memory layout
+    device_restore_all(snapshot->root_snapshot.dss);
+
+    g_hash_table_foreach(snapshot->rbs_dirty_list, root_restore_rb, snapshot);
+
     if (mr_to_enable) {
         memory_region_set_enabled(mr_to_enable, true);
         mr_to_enable = NULL;
     }
-    device_restore_all(snapshot->root_snapshot.dss);
-    syx_snapshot_dirty_list_flush(&snapshot->dirty_list);
+
+    syx_snapshot_dirty_list_flush(snapshot);
 
 	if (must_unlock_iothread) {
 		qemu_mutex_unlock_iothread();

libafl_extras/syx-snapshot/syx-snapshot.h

@@ -1,87 +1,79 @@
+/*
+ * SYX Snapshot
+ *
+ * A speed-oriented snapshot mechanism.
+ *
+ * TODO: complete documentation.
+ */
+
 #pragma once
 #include "qemu/osdep.h"
 #include "qom/object.h"
-#include "device-save.h"
 #include "sysemu/sysemu.h"
+
+#include "device-save.h"
 #include "../syx-misc.h"
 
 /**
  * Saved ramblock
  */
-typedef struct syx_snapshot_ramblock_s {
+typedef struct SyxSnapshotRAMBlock {
     uint8_t* ram; // RAM block
     uint64_t used_length; // Length of the ram block
-    char idstr[256]; // Unique string identifier for the ramblock
-} syx_snapshot_ramblock_t;
+} SyxSnapshotRAMBlock;
 
 /**
  * A root snapshot representation. 
  */
-typedef struct syx_snapshot_root_s {
-    syx_snapshot_ramblock_t* ram_blocks;
-    uint64_t nb_ram_blocks;
-
-    device_save_state_t* dss;
-} syx_snapshot_root_t;
-
-/**
- * A snapshot's dirty list. It only stores dirty addresses
- * (without data). It is the developer's responsibility to
- * to effectively save dirty pages when it is necessary.
- */
-typedef struct syx_snapshot_dirty_list_s {
-    // Dirty pages since the last snapshot
-    // Only physical addresses are stored at this point
-    // Better if a few addresses are marked
-    hwaddr* dirty_addr;
-    uint64_t length;
-    uint64_t capacity;
-} syx_snapshot_dirty_list_t;
+typedef struct SyxSnapshotRoot {
+    GHashTable* rbs_snapshot; // hash map: H(rb) -> SyxSnapshotRAMBlock
+    DeviceSaveState* dss;
+} SyxSnapshotRoot;
 
 /**
  * A list of dirty pages with their old data.
  */
-typedef struct syx_snapshot_dirty_page_s {
-    hwaddr addr;
+typedef struct SyxSnapshotDirtyPage {
+    ram_addr_t offset_within_rb;
     uint8_t* data;
-} syx_snapshot_dirty_page_t;
+} SyxSnapshotDirtyPage;
 
-typedef struct syx_snapshot_dirty_page_list_s {
-    syx_snapshot_dirty_page_t* dirty_pages;
+typedef struct SyxSnapshotDirtyPageList {
+    SyxSnapshotDirtyPage* dirty_pages;
     uint64_t length;
-} syx_snapshot_dirty_page_list_t;
+} SyxSnapshotDirtyPageList;
 
 /**
  * A snapshot increment. It is used to quickly
  * save a VM state.
  */
-typedef struct syx_snapshot_increment_s {
+typedef struct SyxSnapshotIncrement {
     // Back to root snapshot if NULL
-    struct syx_snapshot_increment_s* parent;
+    struct SyxSnapshotIncrement* parent;
 
-    device_save_state_t* dss;
+    DeviceSaveState* dss;
 
-    syx_snapshot_dirty_page_list_t dirty_page_list;
-} syx_snapshot_increment_t;
+    GHashTable* rbs_dirty_pages; // hash map: H(rb) -> SyxSnapshotDirtyPageList
+} SyxSnapshotIncrement;
 
 /**
  * A snapshot. It is the main object used in this API to
- * handle snapshoting.
+ * handle snapshotting.
  */
-typedef struct syx_snapshot_s {
-    syx_snapshot_root_t root_snapshot;
-    syx_snapshot_increment_t* last_incremental_snapshot;
+typedef struct SyxSnapshot {
+    SyxSnapshotRoot root_snapshot;
+    SyxSnapshotIncrement* last_incremental_snapshot;
 
-    syx_snapshot_dirty_list_t dirty_list;
-} syx_snapshot_t;
+    GHashTable* rbs_dirty_list; // hash map: H(rb) -> GHashTable(offset_within_ramblock). Filled lazily.
+} SyxSnapshot;
 
-typedef struct syx_snapshot_tracker_s {
-    syx_snapshot_t** tracked_snapshots;
+typedef struct SyxSnapshotTracker {
+    SyxSnapshot** tracked_snapshots;
     uint64_t length;
     uint64_t capacity;
-} syx_snapshot_tracker_t;
+} SyxSnapshotTracker;
 
-typedef struct syx_snapshot_state_s {
+typedef struct SyxSnapshotState {
     bool is_enabled;
 
     uint64_t page_size;
@@ -89,53 +81,38 @@ typedef struct syx_snapshot_state_s {
  
     // Actively tracked snapshots. Their dirty lists will
     // be updated at each dirty access
-    syx_snapshot_tracker_t tracked_snapshots;
-} syx_snapshot_state_t;
+    SyxSnapshotTracker tracked_snapshots;
+} SyxSnapshotState;
 
-//
-// Namespace API's functions
-//
 
-//void syx_snapshot_init(void* opaque);
 void syx_snapshot_init(void);
-uint64_t syx_snapshot_handler(CPUState* cpu, uint32_t cmd, target_ulong target_opaque);
 
 
 //
 // Snapshot API
 //
 
-syx_snapshot_t* syx_snapshot_create(bool track, device_snapshot_kind_t kind, char** devices);
-void syx_snapshot_free(syx_snapshot_t* snapshot);
-// void syx_snapshot_load(syx_snapshot_t* snapshot);
+SyxSnapshot* syx_snapshot_new(bool track, DeviceSnapshotKind kind, char** devices);
+void syx_snapshot_free(SyxSnapshot* snapshot);
+void syx_snapshot_root_restore(SyxSnapshot* snapshot);
+uint64_t syx_snapshot_check_memory_consistency(SyxSnapshot* snapshot);
 
+// Push the current RAM state and saves it
+void syx_snapshot_increment_push(SyxSnapshot* snapshot, DeviceSnapshotKind kind, char** devices);
 
-//
-// Root snapshot API
-//
+// Restores the last push. Restores the root snapshot if no incremental snapshot is present.
+void syx_snapshot_increment_pop(SyxSnapshot* snapshot);
 
-syx_snapshot_root_t syx_snapshot_root_create(device_snapshot_kind_t kind, char** devices);
-void syx_snapshot_root_restore(syx_snapshot_t* snapshot);
-void syx_snapshot_root_free(syx_snapshot_root_t* root);
+void syx_snapshot_increment_restore_last(SyxSnapshot* snapshot);
 
 
 //
 // Snapshot tracker API
 //
 
-syx_snapshot_tracker_t syx_snapshot_tracker_init(void);
-void syx_snapshot_track(syx_snapshot_tracker_t* tracker, syx_snapshot_t* snapshot);
-void syx_snapshot_stop_track(syx_snapshot_tracker_t* tracker, syx_snapshot_t* snapshot);
-
-
-//
-// Snapshot increment API
-//
-
-void syx_snapshot_increment_push(syx_snapshot_t* snapshot, device_snapshot_kind_t kind, char** devices);
-void syx_snapshot_increment_pop(syx_snapshot_t* snapshot);
-void syx_snapshot_increment_restore_last(syx_snapshot_t* snapshot);
-syx_snapshot_increment_t* syx_snapshot_increment_free(syx_snapshot_increment_t* increment);
+SyxSnapshotTracker syx_snapshot_tracker_init(void);
+void syx_snapshot_track(SyxSnapshotTracker* tracker, SyxSnapshot* snapshot);
+void syx_snapshot_stop_track(SyxSnapshotTracker* tracker, SyxSnapshot* snapshot);
 
 
 //
@@ -149,19 +126,15 @@ bool syx_snapshot_is_enabled(void);
 // Dirty list API
 //
 
-syx_snapshot_dirty_list_t syx_snapshot_dirty_list_create(void);
-void syx_snapshot_dirty_list_free(syx_snapshot_dirty_list_t* dirty_list);
-syx_snapshot_dirty_page_list_t syx_snapshot_dirty_list_to_dirty_page_list(syx_snapshot_dirty_list_t* dirty_list);
-void syx_snapshot_dirty_list_flush(syx_snapshot_dirty_list_t* dirty_list);
-
 void syx_snapshot_dirty_list_add_hostaddr(void* host_addr);
+void syx_snapshot_dirty_list_add_hostaddr_range(void* host_addr, uint64_t len);
 
 /**
  * @brief Add a dirty physical address to the list
  * 
  * @param paddr The physical address to add
  */
-void syx_snapshot_dirty_list_add(hwaddr paddr);
+void syx_snapshot_dirty_list_add_paddr(hwaddr paddr);
 
 /**
  * @brief Same as syx_snapshot_dirty_list_add. The difference

system/physmem.c

@@ -1613,6 +1613,11 @@ void qemu_ram_set_idstr(RAMBlock *new_block, const char *name, DeviceState *dev)
         }
     }
     pstrcat(new_block->idstr, sizeof(new_block->idstr), name);
+//// --- Begin LibAFL code ---
+
+    new_block->idstr_hash = g_str_hash(new_block->idstr);
+
+//// --- End LibAFL code ---
 
     RCU_READ_LOCK_GUARD();
     RAMBLOCK_FOREACH(block) {
@@ -2885,6 +2890,12 @@ 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,
                                                            MemTxAttrs attrs,
@@ -2912,6 +2923,11 @@ static inline MemTxResult address_space_write_rom_internal(AddressSpace *as,
             switch (type) {
             case WRITE_DATA:
                 memcpy(ram_ptr, buf, l);
+//// --- Begin LibAFL code ---
+
+                syx_snapshot_dirty_list_add_hostaddr_range(ram_ptr, l);
+
+//// --- End LibAFL code ---
                 invalidate_and_set_dirty(mr, addr1, l);
                 break;
             case FLUSH_CACHE: