From bdbd4ed7db0353e3b12011d33998a51b8c6c9966 Mon Sep 17 00:00:00 2001
From: neryaz <86964901+neryaz@users.noreply.github.com>
Date: Tue, 7 Mar 2023 23:58:39 +0200
Subject: [PATCH] Add print all chunks in vis_heap_chunks (#1604)

* Changing the arguments to vis_heap_chunks to be clearer

1. --native to --beyond_top
2. --display_all to --no_truncate

* Add print all chunks to vis_heap_chunks

* Preventing the use of the all_chunks argument together with the count argument in vis_heap_chunks

* Use linting for heap.py

* Fix test_vis_heap_chunks.py
According to cdd71a1d82906e7836a88b7ec12e834c130a8528 --display_all/-d moved to --no_truncate/-n

---------

Co-authored-by: Nerya Zadkani <nerya@tokagroup.com>
---
 pwndbg/commands/heap.py                       | 35 +++++++++++++------
 .../tests/heap/test_vis_heap_chunks.py        |  6 ++--
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/pwndbg/commands/heap.py b/pwndbg/commands/heap.py
index 71fb205f..814d25b7 100644
--- a/pwndbg/commands/heap.py
+++ b/pwndbg/commands/heap.py
@@ -697,7 +697,8 @@ parser = argparse.ArgumentParser(
 
 Default to the current arena's active heap.""",
 )
-parser.add_argument(
+group = parser.add_mutually_exclusive_group()
+group.add_argument(
     "count",
     nargs="?",
     type=lambda n: max(int(n, 0), 1),
@@ -706,26 +707,35 @@ parser.add_argument(
 )
 parser.add_argument("addr", nargs="?", default=None, help="Address of the first chunk.")
 parser.add_argument(
-    "--naive",
-    "-n",
+    "--beyond_top",
+    "-b",
     action="store_true",
     default=False,
     help="Attempt to keep printing beyond the top chunk.",
 )
 parser.add_argument(
-    "--display_all",
-    "-a",
+    "--no_truncate",
+    "-n",
     action="store_true",
     default=False,
     help="Display all the chunk contents (Ignore the `max-visualize-chunk-size` configuration).",
 )
+group.add_argument(
+    "--all_chunks",
+    "-a",
+    action="store_true",
+    default=False,
+    help=" Display all chunks (Ignore the default-visualize-chunk-number configuration).",
+)
 
 
 @pwndbg.commands.ArgparsedCommand(parser, category=CommandCategory.HEAP)
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
-def vis_heap_chunks(addr=None, count=None, naive=None, display_all=None) -> None:
+def vis_heap_chunks(
+    addr=None, count=None, beyond_top=None, no_truncate=None, all_chunks=None
+) -> None:
     """Visualize chunks on a heap, default to the current arena's active heap."""
     allocator = pwndbg.heap.current
 
@@ -749,8 +759,12 @@ def vis_heap_chunks(addr=None, count=None, naive=None, display_all=None) -> None
     cursor_backup = cursor
     chunk = Chunk(cursor)
 
-    for _ in range(count + 1):
-        # Don't read beyond the heap mapping if --naive or corrupted heap.
+    chunk_id = 0
+    while True:
+        if not all_chunks and chunk_id == count + 1:
+            break
+
+        # 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)
             break
@@ -764,12 +778,13 @@ def vis_heap_chunks(addr=None, count=None, naive=None, display_all=None) -> None
         else:
             chunk_delims.append(cursor)
 
-        if (chunk.is_top_chunk and not naive) or (cursor == heap_region.end - ptr_size * 2):
+        if (chunk.is_top_chunk and not beyond_top) or (cursor == heap_region.end - ptr_size * 2):
             chunk_delims.append(cursor + ptr_size * 2)
             break
 
         cursor += chunk.real_size
         chunk = Chunk(cursor)
+        chunk_id += 1
 
     # Build the output buffer, changing color at each chunk delimiter.
     # TODO: maybe print free chunks in bold or underlined
@@ -821,7 +836,7 @@ def vis_heap_chunks(addr=None, count=None, naive=None, display_all=None) -> None
         while cursor != stop:
             # skip the middle part of a huge chunk
             if (
-                not display_all
+                not no_truncate
                 and half_max_size > 0
                 and begin_addr + half_max_size <= cursor < end_addr - half_max_size
             ):
diff --git a/tests/gdb-tests/tests/heap/test_vis_heap_chunks.py b/tests/gdb-tests/tests/heap/test_vis_heap_chunks.py
index edf602c6..2f49f12a 100644
--- a/tests/gdb-tests/tests/heap/test_vis_heap_chunks.py
+++ b/tests/gdb-tests/tests/heap/test_vis_heap_chunks.py
@@ -216,12 +216,12 @@ def test_vis_heap_chunk_command(start_binary):
     for omitted_line in omitted_result:
         assert omitted_line in default_result or set(omitted_line) == {"."}
 
-    display_all_result = gdb.execute("vis_heap_chunk -a", to_string=True).splitlines()
-    assert display_all_result == default_result
+    no_truncate_result = gdb.execute("vis_heap_chunk -n", to_string=True).splitlines()
+    assert no_truncate_result == default_result
 
     del default_result
     del omitted_result
-    del display_all_result
+    del no_truncate_result
 
     # Continue, mock overflow changing the chunk size
     gdb.execute("continue")