From f5df117611db0e5b215028ed74d513b05f1ce568 Mon Sep 17 00:00:00 2001
From: theguy147 <37738506+theguy147@users.noreply.github.com>
Date: Wed, 22 Sep 2021 04:53:08 +0200
Subject: [PATCH] Add feature to print heap chunks of all arenas (#722)

* feat: print heap chunks for all arenas

* style: adapt looks of heap chunks for arenas

* refactor: make GlibcArena iterable
---
 docs/commands/heap.md | 10 +++++++++-
 gef.py                | 44 +++++++++++++++++++++++--------------------
 tests/runtests.py     |  4 ++--
 3 files changed, 35 insertions(+), 23 deletions(-)

diff --git a/docs/commands/heap.md b/docs/commands/heap.md
index 7df59769c..775af40a8 100644
--- a/docs/commands/heap.md
+++ b/docs/commands/heap.md
@@ -28,6 +28,14 @@ gef➤ heap chunks [arena_address]
 
 ![heap-chunks-arena](https://i.imgur.com/y1fybRx.png)
 
+In order to display the chunks of all the available arenas at once use
+
+```
+gef➤ heap chunks -a
+```
+
+![heap-chunks-all](https://i.imgur.com/pTjRJFo.png)
+
 Because usually the heap chunks are aligned to a certain number of bytes in
 memory GEF automatically re-aligns the chunks data start addresses to match
 Glibc's behavior. To be able to view unaligned chunks as well, you can disable
@@ -59,7 +67,7 @@ Multi-threaded programs have different arenas, and the knowledge of the
 to help you list all the arenas allocated in your program **at the moment you
 call the command**.
 
-![heap-arenas](https://i.imgur.com/ajbLiCF.png)
+![heap-arenas](https://i.imgur.com/RUTiADa.png)
 
 ### `heap set-arena` command ###
 
diff --git a/gef.py b/gef.py
index 1ed6bc88e..ff8a67e84 100644
--- a/gef.py
+++ b/gef.py
@@ -998,6 +998,12 @@ def __getattr__(self, item):
     def __int__(self):
         return self.__addr
 
+    def __iter__(self):
+        arena = self
+        while arena is not None:
+            yield arena
+            arena = arena.get_next()
+
     def fastbin(self, i):
         """Return head chunk in fastbinsY[i]."""
         addr = int(self.fastbinsY[i])
@@ -1056,8 +1062,11 @@ def get_heap_for_ptr(ptr):
         return ptr & ~(heap_max_size - 1)
 
     def __str__(self):
-        fmt = "Arena (base={:#x}, top={:#x}, last_remainder={:#x}, next={:#x}, next_free={:#x}, system_mem={:#x})"
-        return fmt.format(self.__addr, self.top, self.last_remainder, self.n, self.nfree, self.sysmem)
+        fmt = "{:s}(base={:#x}, top={:#x}, last_remainder={:#x}, next={:#x}, next_free={:#x}, system_mem={:#x})"
+        return fmt.format(
+            Color.colorify("Arena", "blue bold underline"),
+            self.__addr, self.top, self.last_remainder, self.n, self.nfree, self.sysmem
+        )
 
 
 class GlibcChunk:
@@ -1263,6 +1272,11 @@ def get_glibc_arena(addr=None):
         return None
 
 
+def get_glibc_arenas(addr=None, get_all=True):
+    arena = get_glibc_arena(addr)
+    return [arena for arena in iter(arena)] if get_all else [arena]
+
+
 def titlify(text, color=None, msg_color=None):
     """Print a centered title."""
     cols = get_terminal_size()[1]
@@ -6924,17 +6938,9 @@ class GlibcHeapArenaCommand(GenericCommand):
 
     @only_if_gdb_running
     def do_invoke(self, argv):
-        try:
-            arena = GlibcArena(__gef_current_arena__)
-        except gdb.error:
-            err("Could not find Glibc main arena")
-            return
-
-        while True:
-            gef_print("{}".format(arena))
-            arena = arena.get_next()
-            if arena is None:
-                break
+        arenas = get_glibc_arenas()
+        for arena in arenas:
+            gef_print(str(arena))
         return
 
 
@@ -6974,7 +6980,7 @@ class GlibcHeapChunksCommand(GenericCommand):
     the base address of a different arena can be passed"""
 
     _cmdline_ = "heap chunks"
-    _syntax_  = "{0} [-h] [--allow-unaligned] [arena_address]".format(_cmdline_)
+    _syntax_  = "{0} [-h] [--all] [--allow-unaligned] [arena_address]".format(_cmdline_)
     _example_ = "\n{0}\n{0} 0x555555775000".format(_cmdline_)
 
     def __init__(self):
@@ -6982,16 +6988,14 @@ def __init__(self):
         self.add_setting("peek_nb_byte", 16, "Hexdump N first byte(s) inside the chunk data (0 to disable)")
         return
 
-    @parse_arguments({"arena_address": ""}, {"--allow-unaligned": True})
+    @parse_arguments({"arena_address": ""}, {("--all", "-a"): True, "--allow-unaligned": True})
     @only_if_gdb_running
     def do_invoke(self, *args, **kwargs):
         args = kwargs["arguments"]
 
-        arena = get_glibc_arena(addr=args.arena_address)
-        if arena is None:
-            err("No valid arena")
-            return
-        self.dump_chunks_arena(arena, allow_unaligned=args.allow_unaligned)
+        arenas = get_glibc_arenas(addr=args.arena_address, get_all=args.all)
+        for arena in arenas:
+            self.dump_chunks_arena(arena, print_arena=args.all, allow_unaligned=args.allow_unaligned)
 
     def dump_chunks_arena(self, arena, print_arena=False, allow_unaligned=False):
         top_chunk_addr = arena.top
diff --git a/tests/runtests.py b/tests/runtests.py
index c1b4919cd..a1c6911a7 100755
--- a/tests/runtests.py
+++ b/tests/runtests.py
@@ -193,7 +193,7 @@ def test_cmd_heap_arenas(self):
         self.assertFailIfInactiveSession(gdb_run_cmd(cmd, target=target))
         res = gdb_start_silent_cmd(cmd, target=target)
         self.assertNoException(res)
-        self.assertIn("Arena (base=", res)
+        self.assertIn("Arena(base=", res)
         return
 
     def test_cmd_heap_set_arena(self):
@@ -202,7 +202,7 @@ def test_cmd_heap_set_arena(self):
         self.assertFailIfInactiveSession(gdb_run_cmd(cmd, target=target))
         res = gdb_run_silent_cmd(cmd, target=target, after=["heap arenas",])
         self.assertNoException(res)
-        self.assertIn("Arena (base=", res)
+        self.assertIn("Arena(base=", res)
         return
 
     def test_cmd_heap_chunk(self):