From bcef414231733327b163d7e634d92d1ab5bed73f Mon Sep 17 00:00:00 2001
From: BAYET <corentin.bayet@gmail.com>
Date: Tue, 14 Jan 2020 09:02:45 -0500
Subject: [PATCH] Handle serialization of cpu context save (#1129)

* Handle the cpu context save in a more pythonic way, so the context can be serialized and reuse in an other process using the same emulator architecture and modes

* Fix type error ; mistakes a size_t uint64_t ; breaks in 32bit...

Backports commit 8987ad0fffadd16669aa3b402e7e8aaab70ad700 from qemu
---
 bindings/python/unicorn/unicorn.py | 35 +++++++++++++++---------------
 include/unicorn/unicorn.h          | 11 ++++++++++
 uc.c                               |  6 +++++
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/bindings/python/unicorn/unicorn.py b/bindings/python/unicorn/unicorn.py
index 3289b537..8bfa640b 100644
--- a/bindings/python/unicorn/unicorn.py
+++ b/bindings/python/unicorn/unicorn.py
@@ -138,6 +138,7 @@ _setup_prototype(_uc, "uc_context_alloc", ucerr, uc_engine, ctypes.POINTER(uc_co
 _setup_prototype(_uc, "uc_free", ucerr, ctypes.c_void_p)
 _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, "uc_context_size", ctypes.c_size_t, uc_engine)
 _setup_prototype(_uc, "uc_mem_regions", ucerr, uc_engine, ctypes.POINTER(ctypes.POINTER(_uc_mem_region)), ctypes.POINTER(ctypes.c_uint32))
 
 # uc_hook_add is special due to variable number of arguments
@@ -582,24 +583,23 @@ class Uc(object):
         h = 0
 
     def context_save(self):
-        ptr = ctypes.cast(0, ctypes.c_voidp)
-        status = _uc.uc_context_alloc(self._uch, ctypes.byref(ptr))
+        size = _uc.uc_context_size(self._uch)
+
+        context = context_factory(size)
+
+        status = _uc.uc_context_save(self._uch, ctypes.byref(context))
         if status != uc.UC_ERR_OK:
             raise UcError(status)
 
-        status = _uc.uc_context_save(self._uch, ptr)
-        if status != uc.UC_ERR_OK:
-            raise UcError(status)
-
-        return SavedContext(ptr)
+        return ctypes.string_at(ctypes.byref(context), ctypes.sizeof(context))
 
     def context_update(self, context):
-        status = _uc.uc_context_save(self._uch, context.pointer)
+        status = _uc.uc_context_save(self._uch, context)
         if status != uc.UC_ERR_OK:
             raise UcError(status)
 
     def context_restore(self, context):
-        status = _uc.uc_context_restore(self._uch, context.pointer)
+        status = _uc.uc_context_restore(self._uch, context)
         if status != uc.UC_ERR_OK:
             raise UcError(status)
 
@@ -618,14 +618,15 @@ class Uc(object):
             _uc.uc_free(regions)
 
 
-class SavedContext(object):
-    def __init__(self, pointer):
-        self.pointer = pointer
-
-    def __del__(self):
-        status = _uc.uc_free(self.pointer)
-        if status != uc.UC_ERR_OK:
-            raise UcError(status)
+def context_factory(size):
+    class SavedContext(ctypes.Structure):
+        _fields_ = [
+            ('size', ctypes.c_size_t),
+            ('data', ctypes.c_char*size)
+            ]
+    ctxt = SavedContext()
+    ctxt.size = size
+    return ctxt
 
 # print out debugging info
 def debug():
diff --git a/include/unicorn/unicorn.h b/include/unicorn/unicorn.h
index d6789f9a..ae8a6dcf 100644
--- a/include/unicorn/unicorn.h
+++ b/include/unicorn/unicorn.h
@@ -729,6 +729,17 @@ uc_err uc_context_save(uc_engine *uc, uc_context *context);
 UNICORN_EXPORT
 uc_err uc_context_restore(uc_engine *uc, uc_context *context);
 
+/*
+  Return the size needed to store the cpu context. Can be used to allocate a buffer
+  to contain the cpu context and directly call uc_context_save.
+
+  @uc: handle returned by uc_open()
+
+  @return the size for needed to store the cpu context as as size_t.
+*/
+UNICORN_EXPORT
+size_t uc_context_size(uc_engine *uc);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/uc.c b/uc.c
index 82583712..3a743dff 100644
--- a/uc.c
+++ b/uc.c
@@ -1304,6 +1304,12 @@ uc_err uc_free(void *mem)
     return UC_ERR_OK;
 }
 
+UNICORN_EXPORT
+size_t uc_context_size(uc_engine *uc)
+{
+    return cpu_context_size(uc->arch, uc->mode);
+}
+
 UNICORN_EXPORT
 uc_err uc_context_save(uc_engine *uc, uc_context *context)
 {