Merge branch 'feat/reg_save_restore' of https://github.com/rhelmot/unicorn into rhelmot-feat/reg_save_restore

bindings/python/sample_x86.py

@@ -312,6 +312,56 @@ def test_i386_inout():
         print("ERROR: %s" % e)
 
 
+def test_i386_reg_save():
+    print("Save/restore registers in opaque blob")
+    address = 0
+    code = '\x40'  # inc eax
+    try:
+        # Initialize emulator
+        mu = Uc(UC_ARCH_X86, UC_MODE_32)
+
+        # map 8KB memory for this emulation
+        mu.mem_map(address, 8 * 1024, UC_PROT_ALL)
+
+        # write machine code to be emulated to memory
+        mu.mem_write(address, code)
+
+        print(">>> set eax to 1")
+        mu.reg_write(UC_X86_REG_EAX, 1)
+
+        print(">>> execute 'inc eax'")
+        mu.emu_start(address, address+1)
+
+        print(">>> save the register state")
+        saved_regs = mu.regstate_save()
+
+        print(">>> execute 'inc eax'")
+        mu.emu_start(address, address+1)
+
+        print(">>> assert eax == 3")
+        assert mu.reg_read(UC_X86_REG_EAX) == 3
+
+        print(">>> restore the register state")
+        mu.regstate_restore(saved_regs)
+
+        print(">>> assert eax == 2")
+        assert mu.reg_read(UC_X86_REG_EAX) == 2
+
+        print(">>> execute 'inc eax'")
+        mu.emu_start(address, address+1)
+
+        print(">>> assert eax == 3")
+        assert mu.reg_read(UC_X86_REG_EAX) == 3
+
+        print(">>> restore the register state")
+        mu.regstate_restore(saved_regs)
+
+        print(">>> assert eax == 2")
+        assert mu.reg_read(UC_X86_REG_EAX) == 2
+
+    except UcError as e:
+        print("ERROR: %s" % e)
+
 def test_x86_64():
     print("Emulate x86_64 code")
     try:
@@ -483,6 +533,8 @@ if __name__ == '__main__':
     print("=" * 20)
     test_i386_inout()
     print("=" * 20)
+    test_i386_reg_save()
+    print("=" * 20)
     test_x86_64()
     print("=" * 20)
     test_x86_64_syscall()

bindings/python/unicorn/unicorn.py

@@ -100,6 +100,9 @@ _setup_prototype(_uc, "uc_mem_map_ptr", ucerr, uc_engine, ctypes.c_uint64, ctype
 _setup_prototype(_uc, "uc_mem_unmap", ucerr, uc_engine, ctypes.c_uint64, ctypes.c_size_t)
 _setup_prototype(_uc, "uc_mem_protect", ucerr, uc_engine, ctypes.c_uint64, ctypes.c_size_t, ctypes.c_uint32)
 _setup_prototype(_uc, "uc_query", ucerr, uc_engine, ctypes.c_uint32, ctypes.POINTER(ctypes.c_size_t))
+_setup_prototype(_uc, "uc_regstate_save", ctypes.c_voidp, uc_engine, ctypes.c_voidp)
+_setup_prototype(_uc, "uc_regstate_restore", None, uc_engine, ctypes.c_voidp)
+_setup_prototype(_uc, "free", None, ctypes.c_voidp)
 
 # uc_hook_add is special due to variable number of arguments
 _uc.uc_hook_add = _uc.uc_hook_add
@@ -440,6 +443,27 @@ class Uc(object):
             raise UcError(status)
         h = 0
 
+    def regstate_save(self, store=None):
+        if store is None:
+            ptr = ctypes.cast(0, ctypes.c_voidp)
+            return _ActivePointer(_uc.uc_regstate_save(self._uch, ptr))
+        elif type(store) is _ActivePointer:
+            _uc.uc_regstate_save(self._uch, store.pointer)
+            return store
+        else:
+            raise TypeError("Bad register store %s" % repr(store))
+
+    def regstate_restore(self, store):
+        if type(store) is not _ActivePointer:
+            raise TYpeError("Bad register store %s" % repr(store))
+        _uc.uc_regstate_restore(self._uch, store.pointer)
+
+class _ActivePointer(object):
+    def __init__(self, pointer):
+        self.pointer = pointer
+
+    def __del__(self):
+        _uc.free(self.pointer)
 
 # print out debugging info
 def debug():

include/unicorn/unicorn.h

@@ -624,6 +624,35 @@ uc_err uc_mem_protect(uc_engine *uc, uint64_t address, size_t size, uint32_t per
 UNICORN_EXPORT
 uc_err uc_mem_regions(uc_engine *uc, uc_mem_region **regions, uint32_t *count);
 
+/*
+ Save a copy of the current state's registers
+ This API should be used to efficiently make or update a saved copy of the
+ state's registers.
+
+ @uc: handle returned by uc_open()
+ @buffer: pointer to the region to store the registers in. The first call to
+   this function should pass NULL in this parameter, so a region of the
+   appropriate size for the current architecure can be allocated. Further calls
+   to this function may pass in the return value of previous calls.
+
+ @return a pointer to the region the registers were saved in. If buffer was
+   NULL, this is a newly allocated region, otherwise it is the same as buffer.
+   Any allocation performed by this function must be freed by the user.
+*/
+UNICORN_EXPORT
+void *uc_regstate_save(uc_engine *uc, void *buffer);
+
+/*
+ Restore the current state's registers from a saved copy
+ This API should be used to roll the CPU register state back to a previous
+ state saved by uc_regstate_save().
+
+ @uc: handle returned by uc_open()
+ @buffer: pointer returned by uc_regstate_save()
+*/
+UNICORN_EXPORT
+void uc_regstate_restore(uc_engine *uc, void *buffer);
+
 #ifdef __cplusplus
 }
 #endif

qemu/target-arm/unicorn.h

@@ -19,4 +19,7 @@ void arm_uc_init(struct uc_struct* uc);
 __attribute__ ((visibility ("default")))
 void arm64_uc_init(struct uc_struct* uc);
 
+extern const int ARM_REGS_STORAGE_SIZE;
+extern const int ARM64_REGS_STORAGE_SIZE;
+
 #endif

qemu/target-arm/unicorn_aarch64.c

@@ -10,6 +10,8 @@
 #include "uc_priv.h"
 
 
+const int ARM64_REGS_STORAGE_SIZE = offsetof(CPUARMState, tlb_table);
+
 static void arm64_set_pc(struct uc_struct *uc, uint64_t address)
 {
     ((CPUARMState *)uc->current_cpu->env_ptr)->pc = address;

qemu/target-arm/unicorn_arm.c

@@ -10,6 +10,8 @@
 #include "uc_priv.h"
 
 
+const int ARM_REGS_STORAGE_SIZE = offsetof(CPUARMState, tlb_table);
+
 static void arm_set_pc(struct uc_struct *uc, uint64_t address)
 {
     ((CPUARMState *)uc->current_cpu->env_ptr)->pc = address;

qemu/target-i386/unicorn.c

@@ -12,6 +12,8 @@
 #include "uc_priv.h"
 
 
+const int X86_REGS_STORAGE_SIZE = offsetof(CPUX86State, tlb_table);
+
 static void x86_set_pc(struct uc_struct *uc, uint64_t address)
 {
     ((CPUX86State *)uc->current_cpu->env_ptr)->eip = address;

qemu/target-i386/unicorn.h

@@ -12,4 +12,6 @@ void x86_reg_reset(struct uc_struct *uc);
 
 void x86_uc_init(struct uc_struct* uc);
 int x86_uc_machine_init(struct uc_struct *uc);
+
+extern const int X86_REGS_STORAGE_SIZE;
 #endif

qemu/target-m68k/unicorn.c

@@ -10,6 +10,8 @@
 #include "uc_priv.h"
 
 
+const int M68K_REGS_STORAGE_SIZE = offsetof(CPUM68KState, tlb_table);
+
 static void m68k_set_pc(struct uc_struct *uc, uint64_t address)
 {
     ((CPUM68KState *)uc->current_cpu->env_ptr)->pc = address;

qemu/target-m68k/unicorn.h

@@ -12,4 +12,5 @@ void m68k_reg_reset(struct uc_struct *uc);
 
 void m68k_uc_init(struct uc_struct* uc);
 
+extern const int M68K_REGS_STORAGE_SIZE;
 #endif

qemu/target-mips/unicorn.c

@@ -9,6 +9,14 @@
 #include "unicorn_common.h"
 #include "uc_priv.h"
 
+// prevent the lines from being compiled twice
+#ifdef TARGET_WORDS_BIGENDIAN
+#ifdef TARGET_MIPS64
+const int MIPS64_REGS_STORAGE_SIZE = offsetof(CPUMIPSState, tlb_table);
+#else // MIPS32
+const int MIPS_REGS_STORAGE_SIZE = offsetof(CPUMIPSState, tlb_table);
+#endif
+#endif
 
 static uint64_t mips_mem_redirect(uint64_t address)
 {

qemu/target-mips/unicorn.h

@@ -15,4 +15,7 @@ void mipsel_uc_init(struct uc_struct* uc);
 void mips64_uc_init(struct uc_struct* uc);
 void mips64el_uc_init(struct uc_struct* uc);
 
+extern const int MIPS_REGS_STORAGE_SIZE;
+extern const int MIPS64_REGS_STORAGE_SIZE;
+
 #endif

qemu/target-sparc/unicorn.c

@@ -10,6 +10,8 @@
 #include "uc_priv.h"
 
 
+const int SPARC_REGS_STORAGE_SIZE = offsetof(CPUSPARCState, tlb_table);
+
 static bool sparc_stop_interrupt(int intno)
 {
     switch(intno) {

qemu/target-sparc/unicorn.h

@@ -13,4 +13,7 @@ void sparc_reg_reset(struct uc_struct *uc);
 void sparc_uc_init(struct uc_struct* uc);
 void sparc64_uc_init(struct uc_struct* uc);
 
+extern const int SPARC_REGS_STORAGE_SIZE;
+extern const int SPARC64_REGS_STORAGE_SIZE;
+
 #endif

qemu/target-sparc/unicorn64.c

@@ -10,6 +10,8 @@
 #include "uc_priv.h"
 
 
+const int SPARC64_REGS_STORAGE_SIZE = offsetof(CPUSPARCState, tlb_table);
+
 static bool sparc_stop_interrupt(int intno)
 {
     switch(intno) {

tests/unit/test_x86.c

@@ -739,6 +739,68 @@ static void test_x86_16(void **state)
 
 /******************************************************************************/
 
+static void test_i386_reg_save(void **state)
+{
+    uc_engine *uc;
+
+    static const uint64_t address = 0;
+    static const uint8_t code[] = {
+        0x40       // inc eax
+    };
+    int32_t eax = 1;
+
+    // Initialize emulator
+    uc_assert_success(uc_open(UC_ARCH_X86, UC_MODE_32, &uc));
+
+    // map 8KB memory for this emulation
+    uc_assert_success(uc_mem_map(uc, address, 8 * 1024, UC_PROT_ALL));
+
+    // write machine code to be emulated to memory
+    uc_assert_success(uc_mem_write(uc, address, code, sizeof(code)));
+
+    // set eax to 1
+    uc_assert_success(uc_reg_write(uc, UC_X86_REG_EAX, &eax));
+
+    // step one instruction
+    uc_assert_success(uc_emu_start(uc, address, address+1, 0, 0));
+
+    // save the state
+    void *saved_regs = uc_regstate_save(uc, NULL);
+
+    // step one instruction
+    uc_assert_success(uc_emu_start(uc, address, address+1, 0, 0));
+
+    // check that eax == 3
+    uc_assert_success(uc_reg_read(uc, UC_X86_REG_EAX, &eax));
+    assert_int_equal(eax, 3);
+
+    // restore the state
+    uc_regstate_restore(uc, saved_regs);
+
+    // check that eax == 2
+    uc_assert_success(uc_reg_read(uc, UC_X86_REG_EAX, &eax));
+    assert_int_equal(eax, 2);
+
+    // step one instruction
+    uc_assert_success(uc_emu_start(uc, address, address+1, 0, 0));
+
+    // check that eax == 3
+    uc_assert_success(uc_reg_read(uc, UC_X86_REG_EAX, &eax));
+    assert_int_equal(eax, 3);
+
+    // restore the state
+    uc_regstate_restore(uc, saved_regs);
+
+    // check that eax == 2
+    uc_assert_success(uc_reg_read(uc, UC_X86_REG_EAX, &eax));
+    assert_int_equal(eax, 2);
+
+    // clean up;
+    free(saved_regs);
+    uc_assert_success(uc_close(uc));
+}
+/******************************************************************************/
+
 int main(void) {
     const struct CMUnitTest tests[] = {
         cmocka_unit_test(test_i386),
@@ -748,6 +810,7 @@ int main(void) {
         cmocka_unit_test(test_i386_invalid_mem_read),
         cmocka_unit_test(test_i386_invalid_mem_write),
         cmocka_unit_test(test_i386_jump_invalid),
+        cmocka_unit_test(test_i386_reg_save),
 
         cmocka_unit_test(test_x86_64),
         cmocka_unit_test(test_x86_64_syscall),

uc.c

@@ -1155,3 +1155,39 @@ uc_err uc_query(uc_engine *uc, uc_query_type type, size_t *result)
 
     return UC_ERR_OK;
 }
+
+size_t cpu_regs_size(uc_arch arch, uc_mode mode);
+size_t cpu_regs_size(uc_arch arch, uc_mode mode)
+{
+    // each of these constants is defined by offsetof(CPUXYZState, tlb_table)
+    // tbl_table is the first entry in the CPU_COMMON macro, so it marks the end
+    // of the interesting CPU registers
+    switch (arch) {
+        case UC_ARCH_M68K:  return M68K_REGS_STORAGE_SIZE;
+        case UC_ARCH_X86:   return X86_REGS_STORAGE_SIZE;
+        case UC_ARCH_ARM:   return ARM_REGS_STORAGE_SIZE;
+        case UC_ARCH_ARM64: return ARM64_REGS_STORAGE_SIZE;
+        case UC_ARCH_MIPS:  return mode & UC_MODE_MIPS64 ? MIPS64_REGS_STORAGE_SIZE : MIPS_REGS_STORAGE_SIZE;
+        case UC_ARCH_SPARC: return mode & UC_MODE_SPARC64 ? SPARC64_REGS_STORAGE_SIZE : SPARC_REGS_STORAGE_SIZE;
+        default: return 0;
+    }
+}
+
+UNICORN_EXPORT
+void *uc_regstate_save(uc_engine *uc, void *buffer)
+{
+    size_t sz = cpu_regs_size(uc->arch, uc->mode);
+    if (!buffer) {
+        buffer = malloc(sz);
+    }
+
+    memcpy(buffer, first_cpu->env_ptr, sz);
+    return buffer;
+}
+
+UNICORN_EXPORT
+void uc_regstate_restore(uc_engine *uc, void *buffer)
+{
+    size_t sz = cpu_regs_size(uc->arch, uc->mode);
+    memcpy(first_cpu->env_ptr, buffer, sz);
+}