Fix context saving (#1335)

* Fix context size

* Make UcContext convertible to bytes and picklable

Fix when updaing context

* Test context pickling

* Fix double free when the context is pickled from bytes

bindings/python/sample_x86.py

@@ -4,7 +4,7 @@
 from __future__ import print_function
 from unicorn import *
 from unicorn.x86_const import *
-
+import pickle
 
 X86_CODE32 = b"\x41\x4a\x66\x0f\xef\xc1" # INC ecx; DEC edx; PXOR xmm0, xmm1
 X86_CODE32_LOOP = b"\x41\x4a\xeb\xfe" # INC ecx; DEC edx; JMP self-loop
@@ -453,11 +453,17 @@ def test_i386_context_save():
         print(">>> Saving CPU context")
         saved_context = mu.context_save()
 
+        print(">>> Pickling CPU context")
+        pickled_saved_context = pickle.dumps(saved_context)
+
         print(">>> Running emulation for the second time")
         mu.emu_start(address, address+1)
         print(">>> Emulation done. Below is the CPU context")
         print(">>> EAX = 0x%x" %(mu.reg_read(UC_X86_REG_EAX)))
 
+        print(">>> Unpickling CPU context")
+        saved_context = pickle.loads(pickled_saved_context)
+
         print(">>> CPU context restored. Below is the CPU context")
         mu.context_restore(saved_context)
         print(">>> EAX = 0x%x" %(mu.reg_read(UC_X86_REG_EAX)))

bindings/python/unicorn/unicorn.py

@@ -604,7 +604,7 @@ class Uc(object):
         return context
 
     def context_update(self, context):
-        status = _uc.uc_context_save(self._uch, context)
+        status = _uc.uc_context_save(self._uch, context.context)
         if status != uc.UC_ERR_OK:
             raise UcError(status)
 
@@ -628,16 +628,40 @@ class Uc(object):
             _uc.uc_free(regions)
 
 
-class UcContext(ctypes.Structure):
+class UcContext:
     def __init__(self, h):
-        self.context = uc_context()
+        self._context = uc_context()
-
+        self._size = _uc.uc_context_size(h)
-        status = _uc.uc_context_alloc(h, ctypes.byref(self.context))
+        self._to_free = True
+        status = _uc.uc_context_alloc(h, ctypes.byref(self._context))
         if status != uc.UC_ERR_OK:
             raise UcError(status)
     
+    @property
+    def context(self):
+        return self._context
+
+    @property
+    def size(self):
+        return self._size
+
+    # Make UcContext picklable
+    def __getstate__(self):
+        return (bytes(self), self.size)
+    
+    def __setstate__(self, state):
+        self._size = state[1]
+        self._context = ctypes.cast(ctypes.create_string_buffer(state[0], self._size), uc_context)
+        # __init__ won'e be invoked, so we are safe to set it here.
+        self._to_free = False
+        
+    def __bytes__(self):
+        return ctypes.string_at(self.context, self.size)
+
     def __del__(self):
-        _uc.uc_free(self.context)
+        # We need this property since we shouldn't free it if the object is constructed from pickled bytes.
+        if self._to_free:
+            _uc.uc_free(self._context)
 
 
 # print out debugging info

uc.c

@@ -1321,12 +1321,12 @@ UNICORN_EXPORT
 uc_err uc_context_alloc(uc_engine *uc, uc_context **context)
 {
     struct uc_context **_context = context;
-    size_t size = cpu_context_size(uc->arch, uc->mode);
+    size_t size = uc_context_size(uc);
 
     *_context = malloc(size);
     if (*_context) {
         (*_context)->jmp_env_size = sizeof(*uc->cpu->jmp_env);
-        (*_context)->context_size = size - sizeof(uc_context) - (*_context)->jmp_env_size;
+        (*_context)->context_size = cpu_context_size(uc->arch, uc->mode);
         return UC_ERR_OK;
     } else {
         return UC_ERR_NOMEM;