From 1e01ccc8fd717b067b697d5b7353a5d587f6a484 Mon Sep 17 00:00:00 2001
From: Resery <50428593+Resery@users.noreply.github.com>
Date: Tue, 27 Feb 2024 02:43:50 -0600
Subject: 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
---
 unicorn_mode/helper_scripts/unicorn_loader.py | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

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
-- 
cgit 1.4.1