From a97010cf151e948ab6fcff178cdb563b1d2e5d16 Mon Sep 17 00:00:00 2001
From: Glenn Washburn <development@efficientek.com>
Date: Fri, 20 Aug 2021 07:11:42 +0000
Subject: [PATCH] Fix issue where restoring a snapshot created with reg=True
 throws an exception

When restoring a snapshot created with reg=True a unicorn cpu exception can
be thrown when restoring the GS register. It looks like for at least 32bit
mode on X86 processors unicorn will validate memory structures pointed to by
loaded GS reg. If the memory is not loaded the validation will surely fail,
and thus the exception. The fix is then to load the memory maps first,
before loading the registers. Note, this validation does not happen when
restoring via cpu_context.

Included is a test case which will have errors if this fix is not applied.
---
 qiling/core.py    |  6 +++---
 tests/test_elf.py | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/qiling/core.py b/qiling/core.py
index ecb4ad4ff..8d2d0d31a 100644
--- a/qiling/core.py
+++ b/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"])
 
diff --git a/tests/test_elf.py b/tests/test_elf.py
index 0654f31e2..7c42342ad 100644
--- a/tests/test_elf.py
+++ b/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):