diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs
index 264494a0..45962acf 100644
--- a/bindings/rust/src/lib.rs
+++ b/bindings/rust/src/lib.rs
@@ -35,7 +35,6 @@ extern crate std;
 use alloc::boxed::Box;
 use alloc::rc::Rc;
 use alloc::vec::Vec;
-use bitflags::Flags;
 use core::cell::UnsafeCell;
 use core::ptr;
 
diff --git a/qemu/accel/tcg/cputlb.c b/qemu/accel/tcg/cputlb.c
index c120d27b..138fc25c 100644
--- a/qemu/accel/tcg/cputlb.c
+++ b/qemu/accel/tcg/cputlb.c
@@ -2156,7 +2156,7 @@ store_helper(CPUArchState *env, target_ulong addr, uint64_t val,
     }
 
     if (uc->snapshot_level && mr->ram && mr->priority < uc->snapshot_level) {
-        mr = memory_cow(uc, mr, addr & TARGET_PAGE_MASK, TARGET_PAGE_SIZE);
+        mr = memory_cow(uc, mr, paddr & TARGET_PAGE_MASK, TARGET_PAGE_SIZE);
         if (!mr) {
             uc->invalid_addr = paddr;
             uc->invalid_error = UC_ERR_NOMEM;
@@ -2164,7 +2164,7 @@ store_helper(CPUArchState *env, target_ulong addr, uint64_t val,
             return;
         }
         /* refill tlb after CoW */
-        tlb_fill(env_cpu(env), paddr, size, MMU_DATA_STORE,
+        tlb_fill(env_cpu(env), addr, size, MMU_DATA_STORE,
                  mmu_idx, retaddr);
         index = tlb_index(env, mmu_idx, addr);
         entry = tlb_entry(env, mmu_idx, addr);
diff --git a/tests/unit/test_mem.c b/tests/unit/test_mem.c
index 51378802..2bf197a6 100644
--- a/tests/unit/test_mem.c
+++ b/tests/unit/test_mem.c
@@ -315,6 +315,66 @@ static void test_snapshot(void)
     OK(uc_close(uc));
 }
 
+static bool test_snapshot_with_vtlb_callback(uc_engine *uc, uint64_t addr,
+                                             uc_mem_type type,
+                                             uc_tlb_entry *result,
+                                             void *user_data)
+{
+    result->paddr = addr - 0x400000000;
+    result->perms = UC_PROT_ALL;
+    return true;
+}
+
+static void test_snapshot_with_vtlb(void)
+{
+    uc_engine *uc;
+    uc_context *c0, *c1;
+    uint32_t mem;
+    uc_hook hook;
+
+    // mov eax, [0x2020]; inc eax; mov [0x2020], eax
+    char code[] = "\xA1\x20\x20\x00\x00\x04\x00\x00\x00\xFF\xC0\xA3\x20\x20\x00\x00\x04\x00\x00\x00";
+
+    OK(uc_open(UC_ARCH_X86, UC_MODE_64, &uc));
+
+    // Allocate contexts
+    OK(uc_context_alloc(uc, &c0));
+    OK(uc_context_alloc(uc, &c1));
+    OK(uc_ctl_context_mode(uc, UC_CTL_CONTEXT_MEMORY));
+
+
+    OK(uc_ctl_tlb_mode(uc, UC_TLB_VIRTUAL));
+    OK(uc_hook_add(uc, &hook, UC_HOOK_TLB_FILL, test_snapshot_with_vtlb_callback, NULL, 1, 0));
+
+    // Map physical memory
+    OK(uc_mem_map(uc, 0x1000, 0x1000, UC_PROT_EXEC | UC_PROT_READ));
+    OK(uc_mem_write(uc, 0x1000, code, sizeof(code) - 1));
+    OK(uc_mem_map(uc, 0x2000, 0x1000, UC_PROT_ALL));
+
+    // Initial context save
+    OK(uc_context_save(uc, c0));
+
+    OK(uc_emu_start(uc, 0x400000000 + 0x1000, 0x400000000 + 0x1000 + sizeof(code) - 1, 0, 0));
+    OK(uc_mem_read(uc, 0x2020, &mem, sizeof(mem)));
+    TEST_CHECK(mem == 1);
+    OK(uc_context_save(uc, c1));
+    OK(uc_emu_start(uc, 0x400000000 + 0x1000, 0x400000000 + 0x1000 + sizeof(code) - 1, 0, 0));
+    OK(uc_mem_read(uc, 0x2020, &mem, sizeof(mem)));
+    TEST_CHECK(mem == 2);
+    OK(uc_context_restore(uc, c1));
+    // TODO check mem
+    OK(uc_mem_read(uc, 0x2020, &mem, sizeof(mem)));
+    TEST_CHECK(mem == 1);
+    OK(uc_context_restore(uc, c0));
+    OK(uc_mem_read(uc, 0x2020, &mem, sizeof(mem)));
+    TEST_CHECK(mem == 0);
+    // TODO check mem
+
+    OK(uc_context_free(c0));
+    OK(uc_context_free(c1));
+    OK(uc_close(uc));
+}
+
 static void test_context_snapshot(void)
 {
     uc_engine *uc;
@@ -393,6 +453,7 @@ TEST_LIST = {{"test_map_correct", test_map_correct},
              {"test_mem_protect_remove_exec", test_mem_protect_remove_exec},
              {"test_mem_protect_mmio", test_mem_protect_mmio},
              {"test_snapshot", test_snapshot},
+             {"test_snapshot_with_vtlb", test_snapshot_with_vtlb},
              {"test_context_snapshot", test_context_snapshot},
              {"test_snapshot_unmap", test_snapshot_unmap},
              {NULL, NULL}};