From 6d5e574d9d0effa4711795ceaebf693ad38b08b0 Mon Sep 17 00:00:00 2001
From: k4lizen <124312252+k4lizen@users.noreply.github.com>
Date: Sun, 11 Aug 2024 22:44:22 +0200
Subject: [PATCH] better unmapped memory handling in vis_heap_chunk (#2343)

---
 pwndbg/commands/heap.py        | 33 +++++++++++++++++++++++++--------
 pwndbg/gdblib/heap/ptmalloc.py |  3 +++
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/pwndbg/commands/heap.py b/pwndbg/commands/heap.py
index e5f1bbb93..7d804cd6a 100644
--- a/pwndbg/commands/heap.py
+++ b/pwndbg/commands/heap.py
@@ -990,6 +990,10 @@ def vis_heap_chunks(
         addr = count
         count = pwndbg.config.default_visualize_chunk_number
 
+    if addr is not None and not pwndbg.gdblib.memory.is_readable_address(int(addr)):
+        print(message.error("The provided address is not readable."))
+        return
+
     if addr is not None:
         cursor = int(addr)
         heap_region = Heap(cursor)
@@ -1011,6 +1015,7 @@ def vis_heap_chunks(
     chunk = Chunk(cursor)
 
     chunk_id = 0
+    reached_mapping_end = False
     while True:
         if not all_chunks and chunk_id == count + 1:
             break
@@ -1018,6 +1023,7 @@ def vis_heap_chunks(
         # Don't read beyond the heap mapping if --beyond_top or corrupted heap.
         if cursor not in heap_region:
             chunk_delims.append(heap_region.end)
+            reached_mapping_end = True
             break
 
         # Don't repeatedly operate on the same address (e.g. chunk size of 0).
@@ -1029,8 +1035,13 @@ def vis_heap_chunks(
         else:
             chunk_delims.append(cursor)
 
-        if (chunk.is_top_chunk and not beyond_top) or (cursor == heap_region.end - ptr_size * 2):
+        if chunk.is_top_chunk and not beyond_top:
+            chunk_delims.append(cursor + ptr_size * 2)
+            break
+
+        if cursor == heap_region.end - ptr_size * 2:
             chunk_delims.append(cursor + ptr_size * 2)
+            reached_mapping_end = True
             break
 
         cursor += chunk.real_size
@@ -1047,12 +1058,15 @@ def vis_heap_chunks(
         generateColorFunction("blue"),
     ]
 
-    bin_collections = [
-        allocator.fastbins(arena.address),
-        allocator.unsortedbin(arena.address),
-        allocator.smallbins(arena.address),
-        allocator.largebins(arena.address),
-    ]
+    bin_collections = []
+    if arena is not None:
+        # Heap() Case 4; fake/mmapped chunk
+        bin_collections = [
+            allocator.fastbins(arena.address),
+            allocator.unsortedbin(arena.address),
+            allocator.smallbins(arena.address),
+            allocator.largebins(arena.address),
+        ]
     if allocator.has_tcache():
         # Only check for tcache entries belonging to the current thread,
         # it's difficult (impossible?) to find all the thread caches for a
@@ -1109,7 +1123,7 @@ def vis_heap_chunks(
             printed += 1
 
             labels.extend(bin_labels_map.get(cursor, []))
-            if cursor == arena.top:
+            if arena is not None and cursor == arena.top:
                 labels.append("Top chunk")
 
             asc += bin_ascii(data)
@@ -1122,6 +1136,9 @@ def vis_heap_chunks(
 
     print(out)
 
+    if reached_mapping_end:
+        print(f"Reached end of memory mapping ({hex(heap_region.end)}).")
+
     if has_huge_chunk and pwndbg.config.max_visualize_chunk_size == 0:
         print(
             message.warn(
diff --git a/pwndbg/gdblib/heap/ptmalloc.py b/pwndbg/gdblib/heap/ptmalloc.py
index 5fda1b16c..c90adfa58 100644
--- a/pwndbg/gdblib/heap/ptmalloc.py
+++ b/pwndbg/gdblib/heap/ptmalloc.py
@@ -489,6 +489,9 @@ class Heap:
             self._gdbValue = None
         else:
             heap_region = allocator.get_region(addr)
+            if heap_region is None:
+                raise ValueError(f"Cannot build heap object on an unmapped address ({hex(addr)})")
+
             heap_info = allocator.get_heap(addr)
             try:
                 ar_ptr = int(heap_info["ar_ptr"])