Change save/restore API to have explicit alloc/free functions. Python has explicit save/update functions.

bindings/python/unicorn/unicorn.py

@@ -80,6 +80,7 @@ def _setup_prototype(lib, fname, restype, *argtypes):
 
 ucerr = ctypes.c_int
 uc_engine = ctypes.c_void_p
+uc_context = ctypes.c_void_p
 uc_hook_h = ctypes.c_size_t
 
 _setup_prototype(_uc, "uc_version", ctypes.c_uint, ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_int))
@@ -100,8 +101,10 @@ _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_context_save", ctypes.c_voidp, uc_engine, ctypes.c_voidp)
-_setup_prototype(_uc, "uc_context_restore", None, uc_engine, ctypes.c_voidp)
+_setup_prototype(_uc, "uc_context_alloc", ucerr, uc_engine, ctypes.POINTER(uc_context))
+_setup_prototype(_uc, "uc_context_free", ucerr, uc_context)
+_setup_prototype(_uc, "uc_context_save", ucerr, uc_engine, uc_context)
+_setup_prototype(_uc, "uc_context_restore", ucerr, uc_engine, uc_context)
 _setup_prototype(_uc, "free", None, ctypes.c_voidp)
 
 # uc_hook_add is special due to variable number of arguments
@@ -443,27 +446,36 @@ class Uc(object):
             raise UcError(status)
         h = 0
 
-    def context_save(self, store=None):
-        if store is None:
-            ptr = ctypes.cast(0, ctypes.c_voidp)
-            return _ActivePointer(_uc.uc_context_save(self._uch, ptr))
-        elif type(store) is _ActivePointer:
-            _uc.uc_context_save(self._uch, store.pointer)
-            return store
-        else:
-            raise TypeError("Bad register store %s" % repr(store))
+    def context_save(self):
+        ptr = ctypes.cast(0, ctypes.c_voidp)
+        status = _uc.uc_context_alloc(self._uch, ctypes.byref(ptr))
+        if status != uc.UC_ERR_OK:
+            raise UcError(status)
 
-    def context_restore(self, store):
-        if type(store) is not _ActivePointer:
-            raise TYpeError("Bad register store %s" % repr(store))
-        _uc.uc_context_restore(self._uch, store.pointer)
+        status = _uc.uc_context_save(self._uch, ptr)
+        if status != uc.UC_ERR_OK:
+            raise UcError(status)
 
-class _ActivePointer(object):
+        return SavedContext(ptr)
+
+    def context_update(self, context):
+        status = _uc.uc_context_save(self._uch, context.pointer)
+        if status != uc.UC_ERR_OK:
+            raise UcError(status)
+
+    def context_restore(self, context):
+        status = _uc.uc_context_restore(self._uch, context.pointer)
+        if status != uc.UC_ERR_OK:
+            raise UcError(status)
+
+class SavedContext(object):
     def __init__(self, pointer):
         self.pointer = pointer
 
     def __del__(self):
-        _uc.free(self.pointer)
+        status = _uc.uc_context_free(self.pointer)
+        if status != uc.UC_ERR_OK:
+            raise UcError(status)
 
 # print out debugging info
 def debug():

include/unicorn/unicorn.h

@@ -273,6 +273,15 @@ typedef enum uc_query_type {
     UC_QUERY_PAGE_SIZE,
 } uc_query_type;
 
+// Metadata stub for the variable-size cpu context used with uc_context_*()
+typedef struct uc_context {
+   uc_arch arch;
+   uc_mode mode;
+   size_t size;
+   bool used;
+   char data[0];
+} uc_context;
+
 /*
  Return combined API version & major and minor version numbers.
 
@@ -624,23 +633,46 @@ 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);
 
+/*
+ Allocate a region that can be used with uc_context_{save,restore} to perform
+ quick save/rollback of the CPU context, which includes registers and some
+ internal metadata. Contexts may not be shared across engine instances with
+ differing arches or modes.
+
+ @uc: handle returned by uc_open()
+ @context: pointer to a uc_engine*. This will be updated with the pointer to
+   the new context on successful return of this function.
+
+ @return UC_ERR_OK on success, or other value on failure (refer to uc_err enum
+   for detailed error).
+*/
+UNICORN_EXPORT
+uc_err uc_context_alloc(uc_engine *uc, uc_context **context);
+
+/*
+ Free the resource allocated by uc_context_alloc.
+
+ @context: handle returned by uc_context_alloc()
+
+ @return UC_ERR_OK on success, or other value on failure (refer to uc_err enum
+   for detailed error).
+*/
+UNICORN_EXPORT
+uc_err uc_context_free(uc_context *context);
+
 /*
  Save a copy of the internal CPU context.
  This API should be used to efficiently make or update a saved copy of the
  internal CPU state.
 
  @uc: handle returned by uc_open()
- @buffer: pointer to the region to store the context in. The first call to
-   this function should pass NULL in this parameter, so a region of the
-   appropriate size for the current architecture can be allocated. Further calls
-   to this function may pass in the return value of previous calls.
+ @context: handle returned by uc_context_alloc()
 
- @return a pointer to the region the context was 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.
+ @return UC_ERR_OK on success, or other value on failure (refer to uc_err enum
+   for detailed error).
 */
 UNICORN_EXPORT
-void *uc_context_save(uc_engine *uc, void *buffer);
+uc_err uc_context_save(uc_engine *uc, uc_context *context);
 
 /*
  Restore the current CPU context from a saved copy.
@@ -648,10 +680,13 @@ void *uc_context_save(uc_engine *uc, void *buffer);
  state saved by uc_context_save().
 
  @uc: handle returned by uc_open()
- @buffer: pointer returned by uc_context_save()
+ @buffer: handle returned by uc_context_alloc that has been used with uc_context_save
+
+ @return UC_ERR_OK on success, or other value on failure (refer to uc_err enum
+   for detailed error).
 */
 UNICORN_EXPORT
-void uc_context_restore(uc_engine *uc, void *buffer);
+uc_err uc_context_restore(uc_engine *uc, uc_context *context);
 
 #ifdef __cplusplus
 }

tests/unit/test_x86.c

@@ -742,6 +742,7 @@ static void test_x86_16(void **state)
 static void test_i386_reg_save(void **state)
 {
     uc_engine *uc;
+    uc_context *saved_regs;
 
     static const uint64_t address = 0;
     static const uint8_t code[] = {
@@ -764,8 +765,11 @@ static void test_i386_reg_save(void **state)
     // step one instruction
     uc_assert_success(uc_emu_start(uc, address, address+1, 0, 0));
 
+    // grab a buffer to use for state saving
+    uc_assert_success(uc_context_alloc(uc, &saved_regs));
+
     // save the state
-    void *saved_regs = uc_context_save(uc, NULL);
+    uc_assert_success(uc_context_save(uc, saved_regs));
 
     // step one instruction
     uc_assert_success(uc_emu_start(uc, address, address+1, 0, 0));
@@ -796,7 +800,7 @@ static void test_i386_reg_save(void **state)
     assert_int_equal(eax, 2);
 
     // clean up;
-    free(saved_regs);
+    uc_context_free(saved_regs);
     uc_assert_success(uc_close(uc));
 }
 /******************************************************************************/

uc.c

@@ -1174,20 +1174,47 @@ size_t cpu_regs_size(uc_arch arch, uc_mode mode)
 }
 
 UNICORN_EXPORT
-void *uc_context_save(uc_engine *uc, void *buffer)
+uc_err uc_context_alloc(uc_engine *uc, uc_context **context)
 {
-    size_t sz = cpu_regs_size(uc->arch, uc->mode);
-    if (!buffer) {
-        buffer = malloc(sz);
+    size_t size = cpu_regs_size(uc->arch, uc->mode);
+    *context = malloc(size + sizeof(uc_context));
+    if (*context) {
+        (*context)->size = size;
+        (*context)->arch = uc->arch;
+        (*context)->mode = uc->mode;
+        (*context)->used = false;
+        return UC_ERR_OK;
+    } else {
+        return UC_ERR_NOMEM;
     }
-
-    memcpy(buffer, uc->cpu->env_ptr, sz);
-    return buffer;
 }
 
 UNICORN_EXPORT
-void uc_context_restore(uc_engine *uc, void *buffer)
+uc_err uc_context_free(uc_context *context)
 {
-    size_t sz = cpu_regs_size(uc->arch, uc->mode);
-    memcpy(uc->cpu->env_ptr, buffer, sz);
+    free(context);
+    return UC_ERR_OK;
+}
+
+UNICORN_EXPORT
+uc_err uc_context_save(uc_engine *uc, uc_context *context)
+{
+    if (context->arch != uc->arch || context->mode != uc->mode) {
+        return UC_ERR_ARG;
+    } else {
+        memcpy(context->data, uc->cpu->env_ptr, context->size);
+        context->used = true;
+        return UC_ERR_OK;
+    }
+}
+
+UNICORN_EXPORT
+uc_err uc_context_restore(uc_engine *uc, uc_context *context)
+{
+    if (context->arch != uc->arch || context->mode != uc->mode || !context->used) {
+        return UC_ERR_ARG;
+    } else {
+        memcpy(uc->cpu->env_ptr, context->data, context->size);
+        return UC_ERR_OK;
+    }
 }