Update snapshots

include/io/channel-buffer.h

@@ -42,6 +42,10 @@ struct QIOChannelBuffer {
     size_t usage;    /* Current size of data */
     size_t offset;   /* Offset for future I/O ops */
     uint8_t *data;
+    
+    //// --- Begin LibAFL code ---
+    bool internal_allocation;
+    //// --- End LibAFL code ---
 };
 
 
@@ -56,4 +60,11 @@ struct QIOChannelBuffer {
 QIOChannelBuffer *
 qio_channel_buffer_new(size_t capacity);
 
+//// --- Begin LibAFL code ---
+
+QIOChannelBuffer *
+qio_channel_buffer_new_external(uint8_t* buf, size_t capacity, size_t usage);
+
+//// --- End LibAFL code ---
+
 #endif /* QIO_CHANNEL_BUFFER_H */

io/channel-buffer.c

@@ -35,16 +35,49 @@ qio_channel_buffer_new(size_t capacity)
     if (capacity) {
         ioc->data = g_new0(uint8_t, capacity);
         ioc->capacity = capacity;
+        
+        //// --- Begin LibAFL code ---
+        ioc->internal_allocation = true;
+        //// --- End LibAFL code ---
     }
 
     return ioc;
 }
 
+//// --- Begin LibAFL code ---
+
+QIOChannelBuffer *
+qio_channel_buffer_new_external(uint8_t* buf, size_t capacity, size_t usage) {
+    assert(usage <= capacity);
+    assert(buf != NULL);
+
+    QIOChannelBuffer *ioc;
+
+    ioc = QIO_CHANNEL_BUFFER(object_new(TYPE_QIO_CHANNEL_BUFFER));
+
+    ioc->data = buf;
+    ioc->capacity = capacity;
+    ioc->usage = usage;
+    ioc->internal_allocation = false;
+
+    return ioc;
+}
+
+//// --- End LibAFL code ---
 
 static void qio_channel_buffer_finalize(Object *obj)
 {
     QIOChannelBuffer *ioc = QIO_CHANNEL_BUFFER(obj);
+
+    //// --- Begin LibAFL code ---
+    
+    if (ioc->internal_allocation) {
         g_free(ioc->data);
+    }
+
+    //// --- End LibAFL code ---
+    // g_free(ioc->data);
+
     ioc->capacity = ioc->usage = ioc->offset = 0;
 }
 
@@ -142,7 +175,14 @@ static int qio_channel_buffer_close(QIOChannel *ioc,
 {
     QIOChannelBuffer *bioc = QIO_CHANNEL_BUFFER(ioc);
 
+    //// --- Begin LibAFL code ---
+    
+    if (bioc->internal_allocation) {
         g_free(bioc->data);
+    }
+
+    //// --- End LibAFL code ---
+    //g_free(bioc->data);
     bioc->data = NULL;
     bioc->capacity = bioc->usage = bioc->offset = 0;
 

libafl_extras/meson.build

@@ -1,4 +1,5 @@
 specific_ss.add(when: 'CONFIG_SOFTMMU', if_true: [files(
   'syx-snapshot/device-save.c',
   'syx-snapshot/syx-snapshot.c',
+  'syx-snapshot/channel-buffer-writeback.c',
 )])

libafl_extras/syx-snapshot/channel-buffer-writeback.c

@@ -0,0 +1,292 @@
+#include "qemu/osdep.h"
+#include "migration/qemu-file.h"
+#include "channel-buffer-writeback.h"
+#include "../syx-misc.h"
+
+QIOChannelBufferWriteback* qio_channel_buffer_writeback_new(size_t capacity, uint8_t* writeback_buf, size_t writeback_buf_capacity, size_t* writeback_buf_usage) {
+    assert(writeback_buf != NULL);
+    assert(writeback_buf_usage != NULL);
+
+    QIOChannelBufferWriteback *ioc;
+
+    ioc = QIO_CHANNEL_BUFFER_WRITEBACK(object_new(TYPE_QIO_CHANNEL_BUFFER_WRITEBACK));
+
+    assert(writeback_buf != NULL);
+
+    if (capacity) {
+        ioc->data = g_new0(uint8_t, capacity);
+        ioc->capacity = capacity;
+        ioc->internal_allocation = true;
+    }
+
+    ioc->writeback_buf = writeback_buf;
+    ioc->writeback_buf_capacity = writeback_buf_capacity;
+    ioc->writeback_buf_usage = writeback_buf_usage;
+
+    return ioc;
+}
+
+QIOChannelBufferWriteback*
+qio_channel_buffer_writeback_new_external(uint8_t* buf, size_t capacity, size_t usage, uint8_t* writeback_buf, size_t writeback_buf_capacity, size_t* writeback_buf_usage) {
+    assert(buf != NULL);
+    assert(usage <= capacity);
+    assert(writeback_buf != NULL);
+    assert(writeback_buf_usage != NULL);
+
+    QIOChannelBufferWriteback *ioc;
+
+    ioc = QIO_CHANNEL_BUFFER_WRITEBACK(object_new(TYPE_QIO_CHANNEL_BUFFER_WRITEBACK));
+
+    assert(writeback_buf != NULL);
+
+    ioc->data = buf;
+    ioc->capacity = capacity;
+    ioc->usage = usage;
+    ioc->internal_allocation = false;
+
+    ioc->writeback_buf = writeback_buf;
+    ioc->writeback_buf_capacity = writeback_buf_capacity;
+    ioc->writeback_buf_usage = writeback_buf_usage;
+
+    return ioc;
+}
+
+
+static void qio_channel_buffer_writeback_finalize(Object *obj) {
+    QIOChannelBufferWriteback* bwioc = QIO_CHANNEL_BUFFER_WRITEBACK(obj);
+
+    assert(bwioc->writeback_buf_capacity >= bwioc->usage);
+
+    if (bwioc->writeback_buf) {
+        memcpy(bwioc->writeback_buf, bwioc->data, bwioc->usage);
+        *(bwioc->writeback_buf_usage) = bwioc->usage;
+    }
+
+    if (bwioc->internal_allocation) {
+        g_free(bwioc->data);
+        bwioc->data = NULL;
+        bwioc->capacity = bwioc->usage = bwioc->offset = 0;
+    }
+}
+
+
+static ssize_t qio_channel_buffer_writeback_readv(QIOChannel *ioc,
+                                                    const struct iovec *iov,
+                                                    size_t niov,
+                                                    int **fds,
+                                                    size_t *nfds,
+                                                    int flags,
+                                                    Error **errp)
+{
+    QIOChannelBufferWriteback* bwioc = QIO_CHANNEL_BUFFER_WRITEBACK(ioc);
+    ssize_t ret = 0;
+    size_t i;
+
+    for (i = 0; i < niov; i++) {
+        size_t want = iov[i].iov_len;
+
+        if (bwioc->offset >= bwioc->usage) {
+            break;
+        }
+
+        if ((bwioc->offset + want) > bwioc->usage)  {
+            want = bwioc->usage - bwioc->offset;
+        }
+
+        memcpy(iov[i].iov_base, bwioc->data + bwioc->offset, want);
+        ret += want;
+        bwioc->offset += want;
+    }
+
+    return ret;
+}
+
+static ssize_t qio_channel_buffer_writeback_writev(QIOChannel* ioc,
+                                    const struct iovec *iov,
+                                    size_t niov,
+                                    int *fds,
+                                    size_t nfds,
+                                    int flags,
+                                    Error **errp)
+{
+    QIOChannelBufferWriteback* bwioc = QIO_CHANNEL_BUFFER_WRITEBACK(ioc);
+    ssize_t ret = 0;
+    size_t i;
+    size_t towrite = 0;
+
+    for (i = 0; i < niov; i++) {
+        towrite += iov[i].iov_len;
+    }
+
+    assert(bwioc->offset + towrite <= bwioc->capacity);
+    assert(bwioc->offset <= bwioc->usage);
+
+    for (i = 0; i < niov; i++) {
+        memcpy(bwioc->data + bwioc->offset,
+               iov[i].iov_base,
+               iov[i].iov_len);
+        bwioc->offset += iov[i].iov_len;
+        bwioc->usage += iov[i].iov_len;
+        ret += iov[i].iov_len;
+    }
+
+    return ret;
+}
+
+static int qio_channel_buffer_writeback_set_blocking(QIOChannel *ioc G_GNUC_UNUSED,
+                                           bool enabled G_GNUC_UNUSED,
+                                           Error **errp G_GNUC_UNUSED)
+{
+    return 0;
+}
+
+static off_t qio_channel_buffer_writeback_seek(QIOChannel *ioc,
+                                     off_t offset,
+                                     int whence,
+                                     Error **errp)
+{
+    QIOChannelBufferWriteback *bwioc = QIO_CHANNEL_BUFFER_WRITEBACK(ioc);
+    off_t new_pos;
+
+    switch(whence) {
+        case SEEK_SET:
+            new_pos = offset;
+            break;
+        case SEEK_CUR:
+            new_pos = (off_t) bwioc->offset + offset;
+            break;
+        case SEEK_END:
+            new_pos = (off_t) bwioc->usage + offset;
+            break;
+        default:
+            assert(false);
+    }
+
+    assert(new_pos >= 0 && new_pos <= bwioc->usage);
+
+    bwioc->offset = new_pos;
+
+    return new_pos;
+}
+
+static int qio_channel_buffer_writeback_close(QIOChannel *ioc,
+                                              Error **errp)
+{
+    QIOChannelBufferWriteback* bwioc = QIO_CHANNEL_BUFFER_WRITEBACK(ioc);
+
+    assert(bwioc->writeback_buf_capacity >= bwioc->usage);
+
+    if (bwioc->writeback_buf) {
+        memcpy(bwioc->writeback_buf, bwioc->data, bwioc->usage);
+        *(bwioc->writeback_buf_usage) = bwioc->usage;
+    }
+
+    if (bwioc->internal_allocation) {
+        g_free(bwioc->data);
+        bwioc->data = NULL;
+        bwioc->capacity = bwioc->usage = bwioc->offset = 0;
+    }
+
+    return 0;
+}
+
+typedef struct QIOChannelBufferWritebackSource QIOChannelBufferWritebackSource;
+struct QIOChannelBufferWritebackSource {
+    GSource parent;
+    QIOChannelBufferWriteback *bioc;
+    GIOCondition condition;
+};
+
+static gboolean
+qio_channel_buffer_writeback_source_prepare(GSource *source,
+                                  gint *timeout)
+{
+    QIOChannelBufferWritebackSource *bsource = (QIOChannelBufferWritebackSource *)source;
+
+    *timeout = -1;
+
+    return (G_IO_IN | G_IO_OUT) & bsource->condition;
+}
+
+static gboolean
+qio_channel_buffer_writeback_source_check(GSource *source)
+{
+    QIOChannelBufferWritebackSource *bsource = (QIOChannelBufferWritebackSource *)source;
+
+    return (G_IO_IN | G_IO_OUT) & bsource->condition;
+}
+
+static gboolean
+qio_channel_buffer_writeback_source_dispatch(GSource *source,
+                                   GSourceFunc callback,
+                                   gpointer user_data)
+{
+    QIOChannelFunc func = (QIOChannelFunc)callback;
+    QIOChannelBufferWritebackSource *bsource = (QIOChannelBufferWritebackSource *)source;
+
+    return (*func)(QIO_CHANNEL(bsource->bioc),
+                   ((G_IO_IN | G_IO_OUT) & bsource->condition),
+                   user_data);
+}
+
+static void
+qio_channel_buffer_writeback_source_finalize(GSource *source)
+{
+    QIOChannelBufferWritebackSource *ssource = (QIOChannelBufferWritebackSource *)source;
+
+    object_unref(OBJECT(ssource->bioc));
+}
+
+GSourceFuncs qio_channel_buffer_writeback_source_funcs = {
+        qio_channel_buffer_writeback_source_prepare,
+        qio_channel_buffer_writeback_source_check,
+        qio_channel_buffer_writeback_source_dispatch,
+        qio_channel_buffer_writeback_source_finalize
+};
+
+static GSource *qio_channel_buffer_writeback_create_watch(QIOChannel *ioc,
+                                                GIOCondition condition)
+{
+    QIOChannelBufferWriteback *bioc = QIO_CHANNEL_BUFFER_WRITEBACK(ioc);
+    QIOChannelBufferWritebackSource *ssource;
+    GSource *source;
+
+    source = g_source_new(&qio_channel_buffer_writeback_source_funcs,
+                          sizeof(QIOChannelBufferWritebackSource));
+    ssource = (QIOChannelBufferWritebackSource *)source;
+
+    ssource->bioc = bioc;
+    object_ref(OBJECT(bioc));
+
+    ssource->condition = condition;
+
+    return source;
+}
+
+static void qio_channel_buffer_writeback_class_init(ObjectClass *klass,
+                                          void *class_data G_GNUC_UNUSED)
+{
+    QIOChannelClass *ioc_klass = QIO_CHANNEL_CLASS(klass);
+
+    ioc_klass->io_writev = qio_channel_buffer_writeback_writev;
+    ioc_klass->io_readv = qio_channel_buffer_writeback_readv;
+    ioc_klass->io_set_blocking = qio_channel_buffer_writeback_set_blocking;
+    ioc_klass->io_seek = qio_channel_buffer_writeback_seek;
+    ioc_klass->io_close = qio_channel_buffer_writeback_close;
+    ioc_klass->io_create_watch = qio_channel_buffer_writeback_create_watch;
+}
+
+static const TypeInfo qio_channel_buffer_writeback_info = {
+        .parent = TYPE_QIO_CHANNEL,
+        .name = TYPE_QIO_CHANNEL_BUFFER_WRITEBACK,
+        .instance_size = sizeof(QIOChannelBufferWriteback),
+        .instance_finalize = qio_channel_buffer_writeback_finalize,
+        .class_init = qio_channel_buffer_writeback_class_init,
+};
+
+static void qio_channel_buffer_writeback_register_types(void)
+{
+    type_register_static(&qio_channel_buffer_writeback_info);
+}
+
+type_init(qio_channel_buffer_writeback_register_types);

libafl_extras/syx-snapshot/channel-buffer-writeback.h

@@ -0,0 +1,38 @@
+#pragma once
+
+#include "qemu/osdep.h"
+#include "migration/qemu-file.h"
+
+#include "io/channel.h"
+#include "qom/object.h"
+
+#define QEMU_FILE_RAM_LIMIT (32 * 1024 * 1024)
+
+#define TYPE_QIO_CHANNEL_BUFFER_WRITEBACK "qio-channel-buffer-writeback"
+OBJECT_DECLARE_SIMPLE_TYPE(QIOChannelBufferWriteback, QIO_CHANNEL_BUFFER_WRITEBACK)
+
+struct QIOChannelBufferWriteback {
+    QIOChannel parent;
+
+    size_t capacity;
+    size_t usage;
+    size_t offset;
+    uint8_t* data;
+
+    uint8_t* writeback_buf;
+    size_t writeback_buf_capacity;
+    size_t* writeback_buf_usage;
+
+    bool internal_allocation;
+};
+
+QIOChannelBufferWriteback* qio_channel_buffer_writeback_new(size_t capacity, uint8_t* writeback_buf, size_t writeback_buf_capacity, size_t* writeback_buf_usage);
+
+/**
+ * qio_channel_buffer_new_external:
+ * @buf: the buffer used
+ * @capacity: the total capacity of the underlying buffer
+ * @usage: The size actually used by the buffer
+ */
+QIOChannelBufferWriteback*
+qio_channel_buffer_writeback_new_external(uint8_t* buf, size_t capacity, size_t usage, uint8_t* writeback_buf, size_t writeback_buf_capacity, size_t* writeback_buf_usage);

libafl_extras/syx-snapshot/device-save.c

@@ -1,65 +1,21 @@
 #include "qemu/osdep.h"
 #include "device-save.h"
 #include "migration/qemu-file.h"
+#include "io/channel-buffer.h"
+#include "channel-buffer-writeback.h"
 #include "migration/vmstate.h"
 #include "qemu/main-loop.h"
-#include "libafl_extras/syx-misc.h"
+#include "../syx-misc.h"
 
 #include "migration/savevm.h"
 
-#define QEMU_FILE_RAM_LIMIT (32 * 1024 * 1024)
-
-///// From migration/savevm.c
-
-#include "qapi/qapi-commands-migration.h"
-#include "migration/vmstate.h"
-#include "migration/register.h"
-#include "qemu/uuid.h"
-
-typedef struct CompatEntry {
-    char idstr[256];
-    int instance_id;
-} CompatEntry;
-
-typedef struct SaveStateEntry {
-    QTAILQ_ENTRY(SaveStateEntry) entry;
-    char idstr[256];
-    uint32_t instance_id;
-    int alias_id;
-    int version_id;
-    /* version id read from the stream */
-    int load_version_id;
-    int section_id;
-    /* section id read from the stream */
-    int load_section_id;
-    const SaveVMHandlers *ops;
-    const VMStateDescription *vmsd;
-    void *opaque;
-    CompatEntry *compat;
-    int is_ram;
-} SaveStateEntry;
-
-typedef struct SaveState {
-    QTAILQ_HEAD(, SaveStateEntry) handlers;
-    SaveStateEntry *handler_pri_head[MIG_PRI_MAX + 1];
-    int global_section_id;
-    uint32_t len;
-    const char *name;
-    uint32_t target_page_bits;
-    uint32_t caps_count;
-    MigrationCapability *capabilities;
-    QemuUUID uuid;
-} SaveState;
-
-///// End migration/savevm.c
-
 int libafl_restoring_devices;
 
 extern SaveState savevm_state;
 
-void save_section_header(QEMUFile *f, SaveStateEntry *se, uint8_t section_type);
-void save_section_footer(QEMUFile *f, SaveStateEntry *se);
-int vmstate_save(QEMUFile *f, SaveStateEntry *se, JSONWriter *vmdesc);
+extern void save_section_header(QEMUFile *f, SaveStateEntry *se, uint8_t section_type);
+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) {
@@ -67,9 +23,12 @@ device_save_state_t* device_save_all(void) {
     SaveStateEntry *se;
 
     dss->kind = DEVICE_SAVE_KIND_FULL;
-    dss->save_buffer = qio_channel_buffer_new(QEMU_FILE_RAM_LIMIT);
+    dss->save_buffer = g_new(uint8_t, QEMU_FILE_RAM_LIMIT);
 
-    QEMUFile* f = qemu_file_new_output(QIO_CHANNEL(dss->save_buffer));
+    QIOChannelBufferWriteback* wbioc = qio_channel_buffer_writeback_new(QEMU_FILE_RAM_LIMIT, dss->save_buffer, QEMU_FILE_RAM_LIMIT, &dss->save_buffer_size);
+    QIOChannel* ioc = QIO_CHANNEL(wbioc);
+
+    QEMUFile* f = qemu_file_new_output(ioc);
     
     QTAILQ_FOREACH(se, &savevm_state.handlers, entry) {
         int ret;
@@ -87,7 +46,7 @@ device_save_state_t* device_save_all(void) {
             continue;
         }
 
-        // printf("Saving section %p %s...\n", se, se->idstr);
+        // SYX_PRINTF("Saving section %s...\n", se->idstr);
 
         save_section_header(f, se, QEMU_VM_SECTION_FULL);
 
@@ -105,48 +64,29 @@ device_save_state_t* device_save_all(void) {
 
     qemu_put_byte(f, QEMU_VM_EOF);
 
-    qemu_fflush(f);
-
-    // fclose will call io_close and free device_save_state->save_buffer, don't do that
-    //qemu_fclose(f);
+    qemu_fclose(f);
 
     return dss;
 }
 
 void device_restore_all(device_save_state_t* dss) {
-	  bool must_unlock_iothread = false;
+    assert(dss->save_buffer != NULL);
 
-	  Error* errp = NULL;
-	  qio_channel_io_seek(QIO_CHANNEL(dss->save_buffer), 0, SEEK_SET, &errp);
+    QIOChannelBuffer* bioc = qio_channel_buffer_new_external(dss->save_buffer, QEMU_FILE_RAM_LIMIT, dss->save_buffer_size);
+    QIOChannel* ioc = QIO_CHANNEL(bioc);
 
-    if(!dss->save_file) {
-      dss->save_file = qemu_file_new_input(QIO_CHANNEL(dss->save_buffer));
-    }
-    
-	  if (!qemu_mutex_iothread_locked()) {
-		  qemu_mutex_lock_iothread();
-		  must_unlock_iothread = true;
-	  }
+    QEMUFile* f = qemu_file_new_input(ioc);
     
     int save_libafl_restoring_devices = libafl_restoring_devices;
 	  libafl_restoring_devices = 1;
     
-    qemu_load_device_state(dss->save_file);
+    qemu_load_device_state(f);
     
     libafl_restoring_devices = save_libafl_restoring_devices;
 
-	  if (must_unlock_iothread) {
-		  qemu_mutex_unlock_iothread();
-	  }
-	  
-    // qemu_fclose(f);
+    qemu_fclose(f);
 }
 
 void device_free_all(device_save_state_t* dss) {
-    // g_free(dss->save_buffer);
-    Error* errp = NULL;
-    qio_channel_close(QIO_CHANNEL(dss->save_buffer), &errp);
-    object_unref(OBJECT(dss->save_buffer));
-    if (dss->save_file)
-      qemu_fclose(dss->save_file);
+    g_free(dss->save_buffer);
 }

libafl_extras/syx-snapshot/device-save.h

@@ -1,14 +1,13 @@
 #pragma once
 
 #include "qemu/osdep.h"
-#include "io/channel-buffer.h"
 
 #define DEVICE_SAVE_KIND_FULL   0
 
 typedef struct device_save_state_s {
     uint8_t kind;
-    QIOChannelBuffer* save_buffer;
-    QEMUFile* save_file;
+    uint8_t* save_buffer;
+    size_t save_buffer_size;
 } device_save_state_t;
 
 device_save_state_t* device_save_all(void);

libafl_extras/syx-snapshot/syx-snapshot.c

@@ -1,5 +1,4 @@
 #include "qemu/osdep.h"
-//#include "qemu-common.h"
 #include "sysemu/sysemu.h"
 #include "cpu.h"
 #include "qemu/main-loop.h"
@@ -9,6 +8,7 @@
 #include "memory.h"
 
 #include "exec/ram_addr.h"
+/// Physical to host memory in system memory.
 #include "exec/ramlist.h"
 #include "exec/address-spaces.h"
 #include "exec/exec-all.h"
@@ -16,59 +16,21 @@
 #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"
 
-///// From migration/savevm.c
-
-#include "qapi/qapi-commands-migration.h"
-#include "migration/vmstate.h"
-#include "migration/register.h"
-#include "qemu/uuid.h"
-
-typedef struct CompatEntry {
-    char idstr[256];
-    int instance_id;
-} CompatEntry;
-
-typedef struct SaveStateEntry {
-    QTAILQ_ENTRY(SaveStateEntry) entry;
-    char idstr[256];
-    uint32_t instance_id;
-    int alias_id;
-    int version_id;
-    /* version id read from the stream */
-    int load_version_id;
-    int section_id;
-    /* section id read from the stream */
-    int load_section_id;
-    const SaveVMHandlers *ops;
-    const VMStateDescription *vmsd;
-    void *opaque;
-    CompatEntry *compat;
-    int is_ram;
-} SaveStateEntry;
-
-typedef struct SaveState {
-    QTAILQ_HEAD(, SaveStateEntry) handlers;
-    SaveStateEntry *handler_pri_head[MIG_PRI_MAX + 1];
-    int global_section_id;
-    uint32_t len;
-    const char *name;
-    uint32_t target_page_bits;
-    uint32_t caps_count;
-    MigrationCapability *capabilities;
-    QemuUUID uuid;
-} SaveState;
-
-///// End migration/savevm.c
-
 #define SYX_SNAPSHOT_LIST_INIT_SIZE      4096
 #define SYX_SNAPSHOT_LIST_GROW_FACTOR    2
 
 syx_snapshot_state_t 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;
     uint64_t page_size = TARGET_PAGE_SIZE;
 
     syx_snapshot_state.page_size = page_size;
@@ -79,6 +41,10 @@ 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;
+}
+
 syx_snapshot_t* syx_snapshot_create(bool track) {
     syx_snapshot_t* snapshot = g_new0(syx_snapshot_t, 1);
 
@@ -184,7 +150,7 @@ void syx_snapshot_stop_track(syx_snapshot_tracker_t* tracker, syx_snapshot_t* sn
     abort();
 }
 
-void syx_snapshot_increment_push(syx_snapshot_t* snapshot) {
+void syx_snapshot_increment_push(syx_snapshot_t* snapshot, CPUState* cpu) {
     syx_snapshot_increment_t* increment = g_new0(syx_snapshot_increment_t, 1);
     increment->parent = snapshot->last_incremental_snapshot;
     snapshot->last_incremental_snapshot = increment;
@@ -220,12 +186,12 @@ static syx_snapshot_ramblock_t* find_ramblock(syx_snapshot_root_t* root, char* i
 }
 
 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); // memory_region_find is quite slow
+    MemoryRegionSection mr_section = memory_region_find(mr, addr, syx_snapshot_state.page_size);
 
     if (mr_section.size == 0) {
         assert(mr_section.mr == NULL);
 
-        SYX_WARNING("Did not found a memory region while restoring the address %p from root snapshot.\n", (void*) addr);
+        SYX_PRINTF("Did not found a memory region while restoring the address %p from root snapshot.\n", (void*) addr);
         return;
     }
 
@@ -236,9 +202,18 @@ static void restore_page_from_root(syx_snapshot_root_t* root, MemoryRegion* mr,
 
         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);
+    } 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);
@@ -363,18 +338,21 @@ static inline void syx_snapshot_dirty_list_add_internal(hwaddr paddr) {
         // Avoid adding already marked addresses
         for (uint64_t j = 0; j < dirty_list->length; ++j) {
             if (dirty_list->dirty_addr[j] == paddr) {
-                continue;
+                goto next_snapshot_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));
         }
 
         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:;
     }
 }
 
@@ -382,6 +360,30 @@ bool syx_snapshot_is_enabled(void) {
     return syx_snapshot_state.is_enabled;
 }
 
+/*
+// 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) {
+    // early check to know whether we should log the page access or not
+    if (!syx_snapshot_is_enabled()) {
+        return;
+    }
+
+    ram_addr_t offset;
+    RAMBlock* rb = qemu_ram_block_from_host((void*) host_addr, true, &offset);
+
+    assert(rb);
+    
+    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);
+}
+*/
+
 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()) {
@@ -403,6 +405,7 @@ void syx_snapshot_dirty_list_add_hostaddr(void* host_addr) {
 }
 
 
+
 void syx_snapshot_dirty_list_add(hwaddr paddr) {
     if (!syx_snapshot_is_enabled()) {
         return;
@@ -422,8 +425,23 @@ static void syx_snapshot_restore_root_from_dirty_list(syx_snapshot_root_t* root,
 }
 
 void syx_snapshot_root_restore(syx_snapshot_t* snapshot) {
+	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);
+    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);
+ 
+	if (must_unlock_iothread) {
+		qemu_mutex_unlock_iothread();
+	}
 }

libafl_extras/syx-snapshot/syx-snapshot.h

@@ -1,12 +1,16 @@
 #pragma once
-//#ifdef QEMU_SYX
-
 #include "qemu/osdep.h"
 #include "qom/object.h"
 #include "device-save.h"
-//#include "qemu-common.h"
 #include "sysemu/sysemu.h"
-#include "libafl_extras/syx-misc.h"
+#include "../syx-misc.h"
+
+/**
+ * SYX Snapshot parameters
+ */
+typedef struct syx_snapshot_init_params_s {
+    uint64_t page_size;
+} syx_snapshot_init_params_t;
 
 /**
  * Saved ramblock
@@ -99,7 +103,10 @@ typedef struct syx_snapshot_state_s {
 // 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
@@ -132,7 +139,7 @@ void syx_snapshot_stop_track(syx_snapshot_tracker_t* tracker, syx_snapshot_t* sn
 // Snapshot increment API
 //
 
-void syx_snapshot_increment_push(syx_snapshot_t* snapshot);
+void syx_snapshot_increment_push(syx_snapshot_t* snapshot, CPUState* cpu);
 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);
@@ -154,6 +161,7 @@ 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);
 
 /**
  * @brief Add a dirty physical address to the list
@@ -161,6 +169,17 @@ void syx_snapshot_dirty_list_flush(syx_snapshot_dirty_list_t* dirty_list);
  * @param paddr The physical address to add
  */
 void syx_snapshot_dirty_list_add(hwaddr paddr);
-void syx_snapshot_dirty_list_add_hostaddr(void* host_addr);
 
-//#endif
+/**
+ * @brief Same as syx_snapshot_dirty_list_add. The difference
+ * being that it has been specially compiled for full context
+ * saving so that it can be called from anywhere, even in
+ * extreme environments where SystemV ABI is not respected.
+ * It was created with tcg-target.inc.c environment in
+ * mind.
+ * 
+ * @param dummy A dummy argument. it is to comply with
+ *              tcg-target.inc.c specific environment.
+ * @param host_addr The host address where the dirty page is located.
+ */
+void syx_snapshot_dirty_list_add_tcg_target(uint64_t dummy, void* host_addr);

migration/savevm.c

@@ -196,40 +196,11 @@ const VMStateInfo vmstate_info_timer = {
 };
 
 
-typedef struct CompatEntry {
-    char idstr[256];
-    int instance_id;
-} CompatEntry;
+//// --- Begin LibAFL code ---
 
-typedef struct SaveStateEntry {
-    QTAILQ_ENTRY(SaveStateEntry) entry;
-    char idstr[256];
-    uint32_t instance_id;
-    int alias_id;
-    int version_id;
-    /* version id read from the stream */
-    int load_version_id;
-    int section_id;
-    /* section id read from the stream */
-    int load_section_id;
-    const SaveVMHandlers *ops;
-    const VMStateDescription *vmsd;
-    void *opaque;
-    CompatEntry *compat;
-    int is_ram;
-} SaveStateEntry;
+// definitions of CompatEntry SaveStateEntry and SaveState were here
 
-typedef struct SaveState {
-    QTAILQ_HEAD(, SaveStateEntry) handlers;
-    SaveStateEntry *handler_pri_head[MIG_PRI_MAX + 1];
-    int global_section_id;
-    uint32_t len;
-    const char *name;
-    uint32_t target_page_bits;
-    uint32_t caps_count;
-    MigrationCapability *capabilities;
-    QemuUUID uuid;
-} SaveState;
+//// --- End LibAFL code ---
 
 /* static */ SaveState savevm_state = {
     .handlers = QTAILQ_HEAD_INITIALIZER(savevm_state.handlers),

migration/savevm.h

@@ -29,6 +29,50 @@
 #define QEMU_VM_COMMAND              0x08
 #define QEMU_VM_SECTION_FOOTER       0x7e
 
+//// --- Begin LibAFL code ---
+
+#include "migration.h"
+#include "migration/vmstate.h"
+#include "migration/register.h"
+#include "qemu/uuid.h"
+
+typedef struct CompatEntry {
+    char idstr[256];
+    int instance_id;
+} CompatEntry;
+
+typedef struct SaveStateEntry {
+    QTAILQ_ENTRY(SaveStateEntry) entry;
+    char idstr[256];
+    uint32_t instance_id;
+    int alias_id;
+    int version_id;
+    /* version id read from the stream */
+    int load_version_id;
+    int section_id;
+    /* section id read from the stream */
+    int load_section_id;
+    const SaveVMHandlers *ops;
+    const VMStateDescription *vmsd;
+    void *opaque;
+    CompatEntry *compat;
+    int is_ram;
+} SaveStateEntry;
+
+typedef struct SaveState {
+    QTAILQ_HEAD(, SaveStateEntry) handlers;
+    SaveStateEntry *handler_pri_head[MIG_PRI_MAX + 1];
+    int global_section_id;
+    uint32_t len;
+    const char *name;
+    uint32_t target_page_bits;
+    uint32_t caps_count;
+    MigrationCapability *capabilities;
+    QemuUUID uuid;
+} SaveState;
+
+//// --- End LibAFL code ---
+
 bool qemu_savevm_state_blocked(Error **errp);
 void qemu_savevm_non_migratable_list(strList **reasons);
 void qemu_savevm_state_setup(QEMUFile *f);