From d2b7d134434cd2411681fd1b7a5e2d3bdbacc8c4 Mon Sep 17 00:00:00 2001
From: hedger <hedger@nanode.su>
Date: Wed, 11 Jan 2017 18:22:48 +0300
Subject: [PATCH] Fixed circular refs preventing Uc instances from being GC'd.
 Added a test case, requires `objgraph` module.

---
 bindings/python/unicorn/unicorn.py | 30 +++++++++++---
 tests/regress/leaked_refs.py       | 63 ++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+), 5 deletions(-)
 create mode 100644 tests/regress/leaked_refs.py

diff --git a/bindings/python/unicorn/unicorn.py b/bindings/python/unicorn/unicorn.py
index 93d67c0..df13bde 100644
--- a/bindings/python/unicorn/unicorn.py
+++ b/bindings/python/unicorn/unicorn.py
@@ -7,6 +7,7 @@ import pkg_resources
 import inspect
 import os.path
 import sys
+import weakref
 
 from . import x86_const, unicorn_const as uc
 
@@ -211,8 +212,27 @@ class uc_x86_xmm(ctypes.Structure):
         ("high_qword", ctypes.c_uint64),
     ]
 
+# Subclassing ref to allow property assignment.
+class UcRef(weakref.ref):
+    pass
+
+# This class tracks Uc instance destruction and releases handles.
+class UcCleanupManager(object):
+    def __init__(self):
+        self._refs = {}
+
+    def register(self, uc):
+        ref = UcRef(uc, self._finalizer)
+        ref._uch = uc._uch
+        self._refs[id(ref)] = ref
+
+    def _finalizer(self, ref):
+        del self._refs[id(ref)]
+        Uc.release_handle(ref._uch)
 
 class Uc(object):
+    _cleanup = UcCleanupManager()
+
     def __init__(self, arch, mode):
         # verify version compatibility with the core before doing anything
         (major, minor, _combined) = uc_version()
@@ -231,13 +251,13 @@ class Uc(object):
         self._callbacks = {}
         self._ctype_cbs = {}
         self._callback_count = 0
+        self._cleanup.register(self)
 
-    # destructor to be called automatically when object is destroyed.
-    def __del__(self):
-        if self._uch:
+    @staticmethod
+    def release_handle(uch):
+        if uch:
             try:
-                status = _uc.uc_close(self._uch)
-                self._uch = None
+                status = _uc.uc_close(uch)
                 if status != uc.UC_ERR_OK:
                     raise UcError(status)
             except:  # _uc might be pulled from under our feet
diff --git a/tests/regress/leaked_refs.py b/tests/regress/leaked_refs.py
new file mode 100644
index 0000000..263345d
--- /dev/null
+++ b/tests/regress/leaked_refs.py
@@ -0,0 +1,63 @@
+#!/usr/bin/python
+
+from __future__ import print_function
+
+import time
+
+from unicorn import *
+from unicorn.x86_const import *
+
+import objgraph
+
+import regress
+
+ADDRESS = 0x8048000
+STACK_ADDRESS = 0xffff000
+STACK_SIZE = 4096
+'''
+31 DB           xor ebx, ebx
+53              push ebx
+43              inc ebx
+53              push ebx
+6A 02           push 2
+6A 66           push 66h
+58              pop eax
+89 E1           mov ecx, esp
+CD 80           int 80h
+'''
+CODE = "\x31\xDB\x53\x43\x53\x6A\x02\x6A\x66\x58\x89\xE1\xCD\x80"
+EP = ADDRESS + 0x54
+
+def hook_code(mu, address, size, user_data):
+        print(">>> Tracing instruction at 0x%x, instruction size = %u" %(address, size))
+        
+def emu_loop():
+    emu = Uc(UC_ARCH_X86, UC_MODE_32)
+    emu.mem_map(ADDRESS, 0x1000)
+    emu.mem_write(EP, CODE)
+
+    emu.mem_map(STACK_ADDRESS, STACK_SIZE)
+    emu.reg_write(UC_X86_REG_ESP, STACK_ADDRESS + STACK_SIZE)
+    
+    i = emu.hook_add(UC_HOOK_CODE, hook_code, None)
+    emu.hook_del(i)
+    
+    emu.emu_start(EP, EP + len(CODE), count = 3)
+    print("EIP: 0x%x" % emu.reg_read(UC_X86_REG_EIP))
+
+def debugMem():
+    import gc
+    gc.collect()  # don't care about stuff that would be garbage collected properly
+    #print("Orphaned objects in gc.garbage:", gc.garbage)
+    assert(len(objgraph.by_type("Uc")) == 0)
+    #assert(len(objgraph.get_leaking_objects()) == 0)
+
+class EmuLoopReferenceTest(regress.RegressTest):
+    def runTest(self):
+        for i in range(5):
+            emu_loop()
+        debugMem()
+
+
+if __name__ == '__main__':
+    regress.main()
\ No newline at end of file