Fix issues in the Heap Tracker (#2604)

* Allow reentrant memory management calls in the heap tracker

* Defer deletion of GDB breakpoints and improve handling of `free(0)` and `realloc(..., 0)`

* Display instances of `free(0)` in the output of the heap tracker

pwndbg/gdblib/ptmalloc2_tracking.py

@@ -50,6 +50,7 @@ that were not made explicit.
 from __future__ import annotations
 
 from typing import Dict
+from typing import List
 
 import gdb
 from sortedcontainers import SortedDict
@@ -179,6 +180,18 @@ class Chunk:
         self.flags = flags
 
 
+# GDB doesn't like having its breakpoints deleted during stop handlers, so we
+# defer deletion until the next stop event.
+DEFERED_DELETE: List[gdb.Breakpoint] = []
+
+
+@pwndbg.dbg.event_handler(pwndbg.dbg_mod.EventType.STOP)
+def _delete_defered():
+    for entry in DEFERED_DELETE:
+        entry.delete()
+    DEFERED_DELETE.clear()
+
+
 class Tracker:
     def __init__(self) -> None:
         self.free_chunks: SortedDict[int, Chunk] = SortedDict()
@@ -209,7 +222,9 @@ class Tracker:
 
         # Make sure we're not doing anything wrong.
         if thread in self.memory_management_calls:
-            assert self.memory_management_calls[thread]
+            assert self.memory_management_calls[
+                thread
+            ], "exit_memory_management_calls assert failed"
 
         self.memory_management_calls[thread] = False
 
@@ -238,7 +253,9 @@ class Tracker:
 
                 lo_heap = pwndbg.aglib.heap.ptmalloc.Heap(lo_addr)
                 hi_heap = pwndbg.aglib.heap.ptmalloc.Heap(hi_addr - 1)
-                assert lo_heap.arena is not None and hi_heap.arena is not None
+                assert (
+                    lo_heap.arena is not None and hi_heap.arena is not None
+                ), "malloc assert failed"
 
                 # TODO: Can this ever actually fail in real world use?
                 #
@@ -258,20 +275,25 @@ class Tracker:
                 # than to let it become a bug.
                 #
                 # [0]: https://sourceware.org/glibc/wiki/MallocInternals
-                assert lo_heap.start == hi_heap.start and lo_heap.end == hi_heap.end
+                assert (
+                    lo_heap.start == hi_heap.start and lo_heap.end == hi_heap.end
+                ), "malloc assert start failed"
 
                 # Remove all of our old handlers.
                 for i in reversed(range(lo_i, hi_i)):
                     addr, ch = self.free_chunks.popitem(index=i)
 
-                    self.free_watchpoints[addr].delete()
+                    self.free_watchpoints[addr].enabled = False
+                    DEFERED_DELETE.append(self.free_watchpoints[addr])
                     del self.free_watchpoints[addr]
 
                 # Add new handlers in their place. We scan over all of the chunks in
                 # the heap in the range of affected chunks, and add the ones that
                 # are free.
                 allocator = pwndbg.aglib.heap.current
-                assert isinstance(allocator, pwndbg.aglib.heap.ptmalloc.GlibcMemoryAllocator)
+                assert isinstance(
+                    allocator, pwndbg.aglib.heap.ptmalloc.GlibcMemoryAllocator
+                ), "malloc allocator assert failed"
                 bins_list = [
                     allocator.fastbins(lo_heap.arena.address),
                     allocator.smallbins(lo_heap.arena.address),
@@ -342,6 +364,11 @@ class MallocEnterBreakpoint(gdb.Breakpoint):
     def stop(self) -> bool:
         pwndbg.lib.cache.clear_cache("stop")
         requested_size = pwndbg.arguments.argument(0)
+        if self.tracker.is_performing_memory_management():
+            # This call was made from inside another memory management call.
+            # Ignore it.
+            return False
+
         self.tracker.enter_memory_management(MALLOC_NAME)
         AllocExitBreakpoint(self.tracker, requested_size, f"malloc({requested_size})")
         return False
@@ -358,6 +385,10 @@ class CallocEnterBreakpoint(gdb.Breakpoint):
         num_elements = pwndbg.arguments.argument(0)
         element_size = pwndbg.arguments.argument(1)
         requested_size = element_size * num_elements
+        if self.tracker.is_performing_memory_management():
+            # This call was made from inside another memory management call.
+            # Ignore it.
+            return False
 
         self.tracker.enter_memory_management(CALLOC_NAME)
         AllocExitBreakpoint(self.tracker, requested_size, f"calloc({num_elements}, {element_size})")
@@ -427,8 +458,28 @@ class ReallocEnterBreakpoint(gdb.Breakpoint):
 
         freed_pointer = pwndbg.arguments.argument(0)
         requested_size = pwndbg.arguments.argument(1)
+        if self.tracker.is_performing_memory_management():
+            # This call was made from inside another memory management call.
+            # Ignore it.
+            return False
 
         self.tracker.enter_memory_management(REALLOC_NAME)
+
+        if requested_size == 0:
+            # There's no right way to handle realloc(..., 0). C23 says it's
+            # undefined behavior, and prior versions say it's implementation-
+            # defined. Either way, print a warning and do nothing.
+            print(
+                message.warn(
+                    f"[-] realloc({self.freed_pointer:#x}, {requested_size}) ignored, as realloc(0, ...) is implementation defined"
+                )
+            )
+            return False
+
+        if freed_pointer == 0:
+            # Treat this realloc same as malloc
+            AllocExitBreakpoint(self.tracker, requested_size, f"realloc(0x0, {requested_size})")
+        else:
             ReallocExitBreakpoint(self.tracker, freed_pointer, requested_size)
         return False
 
@@ -478,7 +529,7 @@ class ReallocExitBreakpoint(gdb.FinishBreakpoint):
         return False
 
     def out_of_scope(self) -> None:
-        print(message.warn(f"warning: could not follow free request for chunk {self.ptr:#x}"))
+        print(message.warn(f"warning: could not follow free request for chunk {self.freed_ptr:#x}"))
         self.tracker.exit_memory_management()
 
 
@@ -490,6 +541,14 @@ class FreeEnterBreakpoint(gdb.Breakpoint):
     def stop(self) -> bool:
         pwndbg.lib.cache.clear_cache("stop")
         ptr = pwndbg.arguments.argument(0)
+        if self.tracker.is_performing_memory_management():
+            # This call was made from inside another memory management call.
+            # Ignore it.
+            return False
+        if ptr == 0:
+            # free(0) is a no-op.
+            print("[*] free(0x0)")
+            return False
 
         self.tracker.enter_memory_management(FREE_NAME)
         FreeExitBreakpoint(self.tracker, ptr)