Fix restore causing unicorn cpu exception

qiling/core.py

@@ -804,15 +804,15 @@ def restore(self, saved_states=None, snapshot=None):
             with open(snapshot, "rb") as load_state:
                 saved_states = pickle.load(load_state)
 
+        if "mem" in saved_states:
+            self.mem.restore(saved_states["mem"])
+
         if "cpu_context" in saved_states:
             self.arch.context_restore(saved_states["cpu_context"])
 
         if "reg" in saved_states:
             self.reg.restore(saved_states["reg"])
 
-        if "mem" in saved_states:
-            self.mem.restore(saved_states["mem"])
-
         if "fd" in saved_states:
             self.os.fd.restore(saved_states["fd"])
 

tests/test_elf.py

@@ -60,6 +60,48 @@ def stop(ql, *args, **kw):
 
         del ql
 
+    def _test_elf_linux_x86_snapshot_restore_common(self, reg=False, ctx=False):
+        rootfs = "../examples/rootfs/x86_linux"
+        cmdline = ["../examples/rootfs/x86_linux/bin/x86_hello"]
+        snapshot = os.path.join(rootfs, 'snapshot_restore_reg_ctx.snapshot')
+
+        ql = Qiling(cmdline, rootfs, verbose=QL_VERBOSE.DEBUG)
+
+        X86BASE = int(ql.profile.get("OS32", "load_address"), 16)
+        hook_address = X86BASE + 0x542 # call printf
+
+        def dump(ql):
+            nonlocal snapshot
+            nonlocal reg
+            nonlocal ctx
+            ql.save(reg=reg, cpu_context=ctx, os_context=True, loader=True, snapshot=snapshot)
+            ql.emu_stop()
+        ql.hook_address(dump, hook_address)
+
+        ql.run()
+
+        # make sure that the ending PC is the same as the hook address because dump stops the emulater
+        assert ql.reg.arch_pc == hook_address, f"0x{ql.reg.arch_pc:x} != 0x{hook_address:x}"
+        del ql
+
+        ql = Qiling(cmdline, rootfs, verbose=QL_VERBOSE.DEBUG)
+        ql.restore(snapshot=snapshot)
+
+        # ensure that the starting PC is same as the PC we stopped on when taking the snapshot
+        assert ql.reg.arch_pc == hook_address, f"0x{ql.reg.arch_pc:x} != 0x{hook_address:x}"
+
+        ql.run(begin=hook_address)
+        del ql
+
+    def test_elf_linux_x86_snapshot_restore_reg(self):
+        self._test_elf_linux_x86_snapshot_restore_common(reg=True, ctx=False)
+
+    def test_elf_linux_x86_snapshot_restore_ctx(self):
+        self._test_elf_linux_x86_snapshot_restore_common(reg=False, ctx=True)
+
+    def test_elf_linux_x86_snapshot_restore_reg_ctx(self):
+        self._test_elf_linux_x86_snapshot_restore_common(reg=True, ctx=True)
+
     PARAMS_PUTS = {'s': STRING}
 
     def test_elf_linux_x8664(self):