unicornafl: Add UAF chcker to loader (#2009)

* impl uaf chcker

By adding a list of freed chunks, add the chunk to the list during free, check whether the allocated block is in the freed chunk list during malloc, and if so, remove the chunk from the freed chunk list, in __check_mem_access check whether the address is in the freed chunk list. This enables the detection of uaf.

* make uaf_check be configruable

1 files changed, 22 insertions, 1 deletions

diff --git a/unicorn_mode/helper_scripts/unicorn_loader.py b/unicorn_mode/helper_scripts/unicorn_loader.py
index d0995f83..a83e7000 100644
--- a/unicorn_mode/helper_scripts/unicorn_loader.py
+++ b/unicorn_mode/helper_scripts/unicorn_loader.py
@@ -87,9 +87,10 @@ class UnicornSimpleHeap(object):
 
     _uc = None  # Unicorn engine instance to interact with
     _chunks = []  # List of all known chunks
+    _chunks_freed = [] # List of all freed chunks
     _debug_print = False  # True to print debug information
 
-    def __init__(self, uc, debug_print=False):
+    def __init__(self, uc, debug_print=False， uaf_check=False):
         self._uc = uc
         self._debug_print = debug_print
 
@@ -111,6 +112,13 @@ class UnicornSimpleHeap(object):
             try:
                 self._uc.mem_map(addr, total_chunk_size, UC_PROT_READ | UC_PROT_WRITE)
                 chunk = self.HeapChunk(addr, total_chunk_size, size)
+
+                if self.uaf_check:
+                    for chunk_freed in self._chunks_freed:
+                        if chunk_freed.is_buffer_in_chunk(chunk.data_addr, 1):
+                            self._chunks_freed.remove(chunk_freed)
+                            break
+
                 if self._debug_print:
                     print(
                         "Allocating 0x{0:x}-byte chunk @ 0x{1:016x}".format(
@@ -164,6 +172,10 @@ class UnicornSimpleHeap(object):
                         )
                     )
                 self._uc.mem_unmap(chunk.actual_addr, chunk.total_size)
+
+                if self.uaf_check:
+                    self._chunks_freed.append(chunk)
+                    
                 self._chunks.remove(chunk)
                 return True
         # Freed an object that doesn't exist. Maybe 'dobule-free' or 'invalid free' vulnerability here.
@@ -187,6 +199,15 @@ class UnicornSimpleHeap(object):
                     # Force a memory-based crash
                     uc.force_crash(UcError(UC_ERR_READ_PROT))
 
+        if self.uaf_check:
+            for chunk in self._chunks_freed:
+                if address >= chunk.actual_addr and (
+                    (address + size) <= (chunk.actual_addr + chunk.total_size)
+                ):
+                    if chunk.is_buffer_in_chunk(address, size):
+                        print("Use-after-free @ 0x{0:016x}".format(address))
+                        uc.force_crash(UcError(UC_ERR_FETCH_UNMAPPED))
+
 
 # ---------------------------
 # ---- Loading function