trying to sqeeze some speed

Author: domenukk

DEVELOPMENT.md

@@ -9,6 +9,13 @@ We currently develop with backwards compatible python3 in mind.
 Functions may be annotated inline `dev like(this: str) -> None:`.
 However to annotate types, we prefer to add type comments (PEP something) `like = this # type: str`.
 
+## Profiling
+
+For speed, ucf kills the python vm with the internal `os._exit` instead of doing a clean `exit`.
+This, however, trips profiling tools.
+
+In case you want to run a profiler (or other biz), run ucf with `UCF_DEBUG_CLEAN_SHUTDOWN=1`.
+
 ## Debugging
 
 There are different layers to be debugged.

ucf

@@ -9,10 +9,8 @@ import os
 from typing import Any, Callable, Iterable
 
 from unicorefuzz import configspec
-from unicorefuzz.angr_harness import AngrHarness
 from unicorefuzz.configspec import serialize_spec, UNICOREFUZZ_SPEC
 from unicorefuzz.harness import Harness
-from unicorefuzz.probe_wrapper import ProbeWrapper
 from unicorefuzz.unicorefuzz import Unicorefuzz
 
 
@@ -51,6 +49,8 @@ def wrap_probe(args: argparse.Namespace) -> None:
     Attach, break and forward memory from target
     Former probewrapper.py
     """
+    from unicorefuzz.probe_wrapper import ProbeWrapper
+
     ProbeWrapper(load_conf(args)).wrap_gdb_target()
 
 
@@ -69,6 +69,8 @@ def run_angr(args: argparse.Namespace) -> None:
     Drop the memory in the angr harness and start concolic execution
     Former angr-harness.py
     """
+    from unicorefuzz.angr_harness import AngrHarness
+
     AngrHarness(load_conf(args)).get_angry(args.input_file)
 
 

unicorefuzz/harness.py

@@ -3,14 +3,14 @@
 Main (Unicorn-)Harness, used alongside AFL.
 """
 import argparse
+import gc
 import os
 import sys
 import time
 from typing import Optional, Tuple, Dict, List
 
 from capstone import Cs
 from unicorn import *
-from unicorn.x86_const import *
 
 from unicorefuzz import x64utils
 from unicorefuzz.unicorefuzz import (
@@ -102,13 +102,21 @@ def harness(self, input_file: str, wait: bool, debug: bool, trace: bool) -> None
         :param debug: if we should enable unicorn debugger
         :param trace: trace or not
         """
+
+        # Exit without clean python vm shutdown:
+        # "The os._exit() function can be used if it is absolutely positively necessary to exit immediately"
+        # Many times faster!
+        # noinspection PyProtectedMember
+        exit_func = os._exit if not os.getenv("UCF_DEBUG_CLEAN_SHUTDOWN") else exit
+
         uc, entry, exits = self.uc_init(
             input_file, wait, trace, verbose=(debug or trace)
         )
         if debug:
             self.uc_debug(uc, entry_point=entry, exit_point=exits[0])
         else:
             self.uc_run(uc, entry, exits[0])
+            exit_func(0)
 
     def uc_init(
         self, input_file, wait: bool = False, trace: bool = False, verbose: bool = False
@@ -161,13 +169,25 @@ def uc_init(
         # On error: map memory, add exits.
         uc.hook_add(UC_HOOK_MEM_UNMAPPED, unicorn_debug_mem_invalid_access, self)
 
+        # import gc
+        # gc.collect()
+        if os.getenv("UCF_DEBUG_MEMORY"):
+            from pympler import muppy, summary
+
+            all_objects = muppy.get_objects()
+            sum1 = summary.summarize(all_objects)
+            summary.print_(sum1)
+
         # Last chance to hook before forkserver starts (if running as afl child)
         debug_sleep = os.getenv("UCF_DEBUG_SLEEP_BEFORE_FORK")
         if debug_sleep:
             print(
                 "[d] Sleeping. Forkserver will start in {} seconds.".format(debug_sleep)
             )
             time.sleep(float(debug_sleep))
+
+        gc.collect()
+
         # starts the afl forkserver. Won't fork if afl is not around.
         self.uc_start_forkserver(uc, exits)
 
@@ -237,10 +257,6 @@ def uc_run(self, uc: Uc, entry_point: int, exit_point: int) -> None:
                 )
             )
             self.force_crash(e)
-        # Exit without clean python vm shutdown:
-        # "The os._exit() function can be used if it is absolutely positively necessary to exit immediately"
-        # Many times faster!
-        os._exit(0)
 
     def map_known_mem(self, uc: Uc):
         """

unicorefuzz/unicorefuzz.py

@@ -78,16 +78,18 @@
     "armbe": ARMBE,
 }
 
+
 # emulate from @begin, and stop when reaching address @until
 def uc_forkserver_init(uc: Uc, exits: List[int]) -> None:
     import ctypes
-    from unicorn import unicorn, UC_ERR_OK, UcError
+    from unicorn import unicorn
 
     exit_count = len(exits)
     unicorn._uc.uc_afl_forkserver_init(
         uc._uch, ctypes.c_size_t(exit_count), (ctypes.c_uint64 * exit_count)(*exits)
     )
 
+
 def regs_from_unicorn(arch: Architecture) -> List[str]:
     """
     Get all (supported) registers of an arch from Unicorn constants