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):