Provide autoreleas_pool to restrict memory usage (#76)

sledgeh4w · May 12, 2024 · 89b382c · 89b382c
1 parent 547143f
commit 89b382c
Show file tree

Hide file tree

Showing 9 changed files with 118 additions and 106 deletions.
diff --git a/README.md b/README.md
@@ -33,7 +33,7 @@ import uuid
 from chomper import Chomper
 from chomper.const import ARCH_ARM64, OS_IOS
 
-# The system libraries will be automatically loaded from `rootfs_path` on iOS
+# For iOS, system libraries will be automatically loaded from `rootfs_path`
 emu = Chomper(
     arch=ARCH_ARM64,
     os_type=OS_IOS,
@@ -76,15 +76,17 @@ objc = ObjC(emu)
 
 emu.load_module("examples/ios/apps/cn.com.scal.sichuanair/zsch")
 
-# Construct NSString object
-a1 = objc.msg_send("NSString", "stringWithUTF8String:", "test")
+# Use this context manager to ensure that Objective-C objects can be automatically released
+with objc.autorelease_pool():
+    # Construct NSString object
+    a1 = objc.msg_send("NSString", "stringWithUTF8String:", "test")
 
-# Call ObjC method
-req_sign = objc.msg_send("ZSCHRSA", "getReqSign:", a1)
+    # Call Objective-C method
+    req_sign = objc.msg_send("ZSCHRSA", "getReqSign:", a1)
 
-# Convert NSString object to C string
-result_ptr = objc.msg_send(req_sign, "cStringUsingEncoding:", 4)
-result = emu.read_string(result_ptr)
+    # Convert NSString object to C string
+    result_ptr = objc.msg_send(req_sign, "cStringUsingEncoding:", 4)
+    result = emu.read_string(result_ptr)
 ```
 
 Emulate Android native libraries.
@@ -114,25 +116,5 @@ result_size = emu.call_address(libszstone.base + 0x2F1C8, a1, a2, a3)
 result = emu.read_bytes(a3, result_size)
 ```
 
-Hook instructions.
-
-```python
-def hook_code(uc, address, size, user_data):
-    pass
-
-symbol = emu.find_symbol("strlen")
-emu.add_hook(symbol.address, hook_code)
-```
-
-Trace instructions.
-
-```python
-# Trace all instructions
-emu = Chomper(arch=ARCH_ARM64, os_type=OS_ANDROID, trace_instr=True)
-
-# Trace instructions in this module
-emu.load_module("examples/android/rootfs/system/lib64/libc.so", trace_inst=True)
-```
-
 ## Examples
-[Here](https://github.com/sledgeh4w/chomper/tree/main/examples) are a fews examples of encryption emulations for security vendors.
+[Here](https://github.com/sledgeh4w/chomper/tree/main/examples) are some encryption emulation examples for security vendors.
diff --git a/examples/example_ios_ali_vmp_sign.py b/examples/example_ios_ali_vmp_sign.py
@@ -66,11 +66,12 @@ def main():
     app_key = "xPEj7uv0KuziQnXUyPIBNUjnDvvHuW09VOYFuLYBcY-jV6fgqmfy5B1y75_iSuRM5U2zNq7MRoR9N1F-UthTEgv-QBWk68gr95BrAySzWuDzt08FrkeBZWQCGyZ0iAybalYLOJEF7nkKBtmDGLewcw=="
     objc.msg_send(ali_tiger_tally_instance, "initialize:", pyobj2nsobj(emu, app_key))
 
-    encrypt_str = '{"biClassId":["2","3","4"]}'
-    encrypt_bytes = objc.msg_send(pyobj2nsobj(emu, encrypt_str), "dataUsingEncoding:", 1)
+    with objc.autorelease_pool():
+        encrypt_str = '{"biClassId":["2","3","4"]}'
+        encrypt_bytes = objc.msg_send(pyobj2nsobj(emu, encrypt_str), "dataUsingEncoding:", 1)
 
-    vmp_sign = objc.msg_send(ali_tiger_tally_instance, "vmpSign:", encrypt_bytes)
-    logger.info("vmp sign: %s", emu.read_string(objc.msg_send(vmp_sign, "cStringUsingEncoding:", 4)))
+        vmp_sign = objc.msg_send(ali_tiger_tally_instance, "vmpSign:", encrypt_bytes)
+        logger.info("vmpSign: %s", emu.read_string(objc.msg_send(vmp_sign, "cStringUsingEncoding:", 4)))
 
 
 if __name__ == "__main__":

diff --git a/examples/example_ios_bangbang.py b/examples/example_ios_bangbang.py
diff --git a/examples/example_ios_ijm.py b/examples/example_ios_ijm.py
@@ -42,17 +42,18 @@ def main():
     # Skip a special check of ijm
     emu.add_interceptor(czair.base + 0x1038F0004, hook_retval(1))
 
-    # Encrypt
-    encrypt_str = '{"biClassId":["2","3","4"]}'
-    encrypt_result = objc.msg_send("JMBox125", "JMBox167:JMBox501:", pyobj2nsobj(emu, encrypt_str), 1)
+    with objc.autorelease_pool():
+        # Encrypt
+        encrypt_str = '{"biClassId":["2","3","4"]}'
+        encrypt_result = objc.msg_send("JMBox125", "JMBox167:JMBox501:", pyobj2nsobj(emu, encrypt_str), 1)
 
-    logger.info("encrypt_result: %s", emu.read_string(objc.msg_send(encrypt_result, "cStringUsingEncoding:", 4)))
+        logger.info("Encrypt result: %s", emu.read_string(objc.msg_send(encrypt_result, "cStringUsingEncoding:", 4)))
 
-    # Decrypt
-    decrypt_str = "XKQYFMCP9Eb0IUzrQ9KaRRvTeFcYYyLcInrS/IWp6be1+VZa14GanCrzeb3DR45HW+XH0xiZLA5WUjUcXnlpM+CC6EtauUDUxCLap3QPWRyewLUosCB/ESHE7341DQca6lx5KFcP0XCkBpGlEKpACR5v7TwNBxc62auNBDvmEY422LTAUEEBrC8FDE+Y4DS2IJTLN6h9f7hdmQ4zUnY4cwyZXwgdIoH+bVuNy6TSw1JjQaFF/fLLHVZOQovrMcjtTpMZGr8xOSoW/+msiZzKwET3"
-    decrypt_result = objc.msg_send("JMBox125", "JMBox167:JMBox501:", pyobj2nsobj(emu, decrypt_str), 1)
+        # Decrypt
+        decrypt_str = "XKQYFMCP9Eb0IUzrQ9KaRRvTeFcYYyLcInrS/IWp6be1+VZa14GanCrzeb3DR45HW+XH0xiZLA5WUjUcXnlpM+CC6EtauUDUxCLap3QPWRyewLUosCB/ESHE7341DQca6lx5KFcP0XCkBpGlEKpACR5v7TwNBxc62auNBDvmEY422LTAUEEBrC8FDE+Y4DS2IJTLN6h9f7hdmQ4zUnY4cwyZXwgdIoH+bVuNy6TSw1JjQaFF/fLLHVZOQovrMcjtTpMZGr8xOSoW/+msiZzKwET3"
+        decrypt_result = objc.msg_send("JMBox125", "JMBox167:JMBox501:", pyobj2nsobj(emu, decrypt_str), 1)
 
-    logger.info("decrypt_result: %s", emu.read_string(objc.msg_send(decrypt_result, "cStringUsingEncoding:", 4)))
+        logger.info("Decrypt result: %s", emu.read_string(objc.msg_send(decrypt_result, "cStringUsingEncoding:", 4)))
 
 
 if __name__ == "__main__":

diff --git a/src/chomper/memory.py b/src/chomper/memory.py
@@ -110,7 +110,11 @@ def realloc(self, address: int, size: int) -> int:
             return address
 
         new_address = self.alloc(size)
-        self.uc.mem_write(new_address, bytes(self.uc.mem_read(address, block_size)))
+
+        data = self.uc.mem_read(address, block_size)
+        self.uc.mem_write(new_address, bytes(data))
+
+        self.free(address)
 
         return new_address
 

diff --git a/src/chomper/objc.py b/src/chomper/objc.py
@@ -1,4 +1,4 @@
-# from contextlib import contextmanager
+from contextlib import contextmanager
 from typing import Union
 
 
@@ -41,14 +41,14 @@ def msg_send(self, receiver: Union[int, str], sel: Union[int, str], *args) -> in
         receiver = self.get_class(receiver) if isinstance(receiver, str) else receiver
         sel = self.get_sel(sel) if isinstance(sel, str) else sel
 
-        str_ptrs = []
+        mem_ptrs = []
         new_args = []
 
         for arg in args:
             if isinstance(arg, str):
                 str_ptr = self.emu.create_string(arg)
 
-                str_ptrs.append(str_ptr)
+                mem_ptrs.append(str_ptr)
                 new_args.append(str_ptr)
 
             else:
@@ -58,18 +58,19 @@ def msg_send(self, receiver: Union[int, str], sel: Union[int, str], *args) -> in
             return self.emu.call_symbol("_objc_msgSend", receiver, sel, *new_args)
 
         finally:
-            for str_ptr in str_ptrs:
-                self.emu.free(str_ptr)
+            for mem_ptr in mem_ptrs:
+                self.emu.free(mem_ptr)
 
     def release(self, obj: int):
         """Release object."""
         self.emu.call_symbol("_objc_release", obj)
 
-    # @contextmanager
-    # def autorelease_pool(self):
-    #     """Ensure Objetive-C objects are automatically released."""
-    #     context = self.emu.call_symbol("_objc_autoreleasePoolPush")
-    #     try:
-    #         yield context
-    #     finally:
-    #         self.emu.call_symbol("_objc_autoreleasePoolPop", context)
+    @contextmanager
+    def autorelease_pool(self):
+        """Ensure Objetive-C objects can be automatically released."""
+        context = self.emu.call_symbol("_objc_autoreleasePoolPush")
+
+        try:
+            yield context
+        finally:
+            self.emu.call_symbol("_objc_autoreleasePoolPop", context)
diff --git a/src/chomper/os/ios/hooks.py b/src/chomper/os/ios/hooks.py
@@ -7,7 +7,7 @@
 
 from unicorn.unicorn import UC_HOOK_CODE_TYPE
 
-from chomper.utils import pyobj2nsobj, pyobj2cfobj
+from chomper.utils import pyobj2cfobj
 from chomper.objc import ObjC
 
 hooks: Dict[str, UC_HOOK_CODE_TYPE] = {}
@@ -436,7 +436,7 @@ def hook_mg_copy_answer(uc, address, size, user_data):
     key = emu.read_string(str_ptr)
 
     if key in emu.os.device_info:
-        return pyobj2nsobj(emu, emu.os.device_info[key])
+        return pyobj2cfobj(emu, emu.os.device_info[key])
 
     return 0
 

diff --git a/src/chomper/utils.py b/src/chomper/utils.py
@@ -23,6 +23,8 @@ def pyobj2nsobj(emu, obj: object) -> int:
     """
     objc = ObjC(emu)
 
+    mem_ptrs = []
+
     if isinstance(obj, dict):
         ns_obj = objc.msg_send("NSMutableDictionary", "dictionary")
 
@@ -43,14 +45,21 @@ def pyobj2nsobj(emu, obj: object) -> int:
         ns_obj = objc.msg_send("NSString", "stringWithUTF8String:", obj)
 
     elif isinstance(obj, bytes):
-        buffer = emu.create_buffer(len(obj))
-        emu.write_bytes(buffer, obj)
+        if obj:
+            buffer = emu.create_buffer(len(obj))
+            emu.write_bytes(buffer, obj)
+            mem_ptrs.append(buffer)
+        else:
+            buffer = 0
 
         ns_obj = objc.msg_send("NSData", "dataWithBytes:length:", buffer, len(obj))
 
     else:
         raise TypeError(f"Unsupported type: {type(obj)}")
 
+    for mem_ptr in mem_ptrs:
+        emu.free(mem_ptr)
+
     return ns_obj
 
 
@@ -62,7 +71,7 @@ def pyobj2cfobj(emu, obj: object) -> int:
     """
     cf_allocator_system_default = emu.find_symbol("___kCFAllocatorSystemDefault")
 
-    str_ptrs = []
+    mem_ptrs = []
     cf_strs = []
 
     if isinstance(obj, dict):
@@ -106,7 +115,7 @@ def pyobj2cfobj(emu, obj: object) -> int:
 
     elif isinstance(obj, str):
         str_ptr = emu.create_string(obj)
-        str_ptrs.append(str_ptr)
+        mem_ptrs.append(str_ptr)
 
         cf_obj = emu.call_symbol(
             "_CFStringCreateWithCString",
@@ -119,6 +128,7 @@ def pyobj2cfobj(emu, obj: object) -> int:
         if obj:
             buffer = emu.create_buffer(len(obj))
             emu.write_bytes(buffer, obj)
+            mem_ptrs.append(buffer)
         else:
             buffer = 0
 
@@ -132,8 +142,8 @@ def pyobj2cfobj(emu, obj: object) -> int:
     else:
         raise TypeError(f"Unsupported type: {type(obj)}")
 
-    for str_ptr in str_ptrs:
-        emu.free(str_ptr)
+    for mem_ptr in mem_ptrs:
+        emu.free(mem_ptr)
 
     for cf_str in cf_strs:
         emu.call_symbol("_CFRelease", cf_str)

diff --git a/tests/test_objc.py b/tests/test_objc.py
@@ -1,77 +1,89 @@
 def test_ns_string(emu_ios, objc):
-    string = objc.msg_send("NSString", "stringWithUTF8String:", "chomper")
+    with objc.autorelease_pool():
+        string = objc.msg_send("NSString", "stringWithUTF8String:", "chomper")
 
-    assert string
+        assert string
 
 
 def test_ns_array(emu_ios, objc):
-    array = objc.msg_send("NSMutableArray", "array")
+    with objc.autorelease_pool():
+        array = objc.msg_send("NSMutableArray", "array")
 
-    assert array
+        assert array
 
-    string = objc.msg_send("NSString", "stringWithUTF8String:", "chomper")
-    objc.msg_send(array, "addObject:", string)
+        string = objc.msg_send("NSString", "stringWithUTF8String:", "chomper")
+        objc.msg_send(array, "addObject:", string)
 
 
 def test_ns_dictionary(emu_ios, objc):
-    dictionary = objc.msg_send("NSMutableDictionary", "dictionary")
+    with objc.autorelease_pool():
+        dictionary = objc.msg_send("NSMutableDictionary", "dictionary")
 
-    assert dictionary
+        assert dictionary
 
-    string = objc.msg_send("NSString", "stringWithUTF8String:", "chomper")
-    objc.msg_send(dictionary, "setObject:forKey:", string, string)
+        string = objc.msg_send("NSString", "stringWithUTF8String:", "chomper")
+        objc.msg_send(dictionary, "setObject:forKey:", string, string)
 
 
 def test_ns_data(emu_ios, objc):
-    data_bytes = b"chomper"
+    with objc.autorelease_pool():
+        data_bytes = b"chomper"
 
-    buffer = emu_ios.create_buffer(len(data_bytes))
-    emu_ios.write_bytes(buffer, data_bytes)
+        buffer = emu_ios.create_buffer(len(data_bytes))
+        emu_ios.write_bytes(buffer, data_bytes)
 
-    data = objc.msg_send("NSData", "dataWithBytes:length:", buffer, len(data_bytes))
+        data = objc.msg_send("NSData", "dataWithBytes:length:", buffer, len(data_bytes))
 
-    assert data
+        assert data
 
 
 def test_ns_url(emu_ios, objc):
-    string = objc.msg_send("NSString", "stringWithUTF8String:", "https://google.com")
+    with objc.autorelease_pool():
+        string = objc.msg_send(
+            "NSString", "stringWithUTF8String:", "https://google.com"
+        )
 
-    url = objc.msg_send("NSURL", "alloc")
-    objc.msg_send(url, "initWithString:", string)
+        url = objc.msg_send("NSURL", "alloc")
+        objc.msg_send(url, "initWithString:", string)
 
-    assert url
+        assert url
 
 
 def test_ns_request(emu_ios, objc):
-    string = objc.msg_send("NSString", "stringWithUTF8String:", "https://google.com")
+    with objc.autorelease_pool():
+        string = objc.msg_send(
+            "NSString", "stringWithUTF8String:", "https://google.com"
+        )
 
-    url = objc.msg_send("NSURL", "alloc")
-    objc.msg_send(url, "initWithString:", string)
+        url = objc.msg_send("NSURL", "alloc")
+        objc.msg_send(url, "initWithString:", string)
 
-    request = objc.msg_send("NSMutableURLRequest", "requestWithURL:", url)
+        request = objc.msg_send("NSMutableURLRequest", "requestWithURL:", url)
 
-    assert request
+        assert request
 
 
 def test_ns_locale(emu_ios, objc):
-    preferred_languages = objc.msg_send("NSLocale", "preferredLanguages")
+    with objc.autorelease_pool():
+        preferred_languages = objc.msg_send("NSLocale", "preferredLanguages")
 
-    assert preferred_languages
+        assert preferred_languages
 
-    preferred_language = objc.msg_send(preferred_languages, "firstObject")
-    str_ptr = objc.msg_send(preferred_language, "cStringUsingEncoding:", 4)
+        preferred_language = objc.msg_send(preferred_languages, "firstObject")
+        str_ptr = objc.msg_send(preferred_language, "cStringUsingEncoding:", 4)
 
-    assert len(emu_ios.read_string(str_ptr)) > 0
+        assert len(emu_ios.read_string(str_ptr)) > 0
 
 
 def test_ns_user_defaults(emu_ios, objc):
-    user_defaults = objc.msg_send("NSUserDefaults", "standardUserDefaults")
+    with objc.autorelease_pool():
+        user_defaults = objc.msg_send("NSUserDefaults", "standardUserDefaults")
 
-    assert user_defaults
+        assert user_defaults
 
-    key = objc.msg_send("NSString", "stringWithUTF8String:", "AppleLocale")
+        key = objc.msg_send("NSString", "stringWithUTF8String:", "AppleLocale")
 
-    apple_locale = objc.msg_send(user_defaults, "stringForKey:", key)
-    str_ptr = objc.msg_send(apple_locale, "cStringUsingEncoding:", 4)
+        apple_locale = objc.msg_send(user_defaults, "stringForKey:", key)
+        str_ptr = objc.msg_send(apple_locale, "cStringUsingEncoding:", 4)
 
-    assert len(emu_ios.read_string(str_ptr)) > 0
+        assert len(emu_ios.read_string(str_ptr)) > 0