fix: heap bins commands (#805)

* fix: heap bins commands

* review: implement suggestions from @hugsy

gef.py

@@ -1405,6 +1405,10 @@ def __str__(self) -> str:
                 f"last_remainder={self.last_remainder:#x}, next={self.n:#x}, next_free={self.nfree:#x}, "
                 f"system_mem={self.sysmem:#x})")
 
+    @property
+    def addr(self) -> int:
+        return int(self)
+
 
 class GlibcChunk:
     """Glibc chunk class. The default behavior (from_base=False) is to interpret the data starting at the memory
@@ -7307,22 +7311,27 @@ def __init__(self) -> None:
         super().__init__(complete=gdb.COMPLETE_LOCATION)
         return
 
+    @parse_arguments({"arena_address": ""}, {})
     @only_if_gdb_running
-    def do_invoke(self, argv: List[str]) -> None:
+    def do_invoke(self, *_: Any, **kwargs: Any) -> None:
         def fastbin_index(sz: int) -> int:
             return (sz >> 4) - 2 if SIZE_SZ == 8 else (sz >> 3) - 2
 
+        args = kwargs["arguments"]
+        if not gef.heap.main_arena:
+            err("Heap not initialized")
+            return
+
         SIZE_SZ = gef.arch.ptrsize
         MAX_FAST_SIZE = 80 * SIZE_SZ // 4
         NFASTBINS = fastbin_index(MAX_FAST_SIZE) - 1
 
-        arena = GlibcArena(f"*{argv[0]}") if len(argv) == 1 else gef.heap.main_arena
-
+        arena = GlibcArena(f"*{args.arena_address}") if args.arena_address else gef.heap.selected_arena
         if arena is None:
             err("Invalid Glibc arena")
             return
 
-        gef_print(titlify(f"Fastbins for arena {int(arena):#x}"))
+        gef_print(titlify(f"Fastbins for arena at {arena.addr:#x}"))
         for i in range(NFASTBINS):
             gef_print(f"Fastbins[idx={i:d}, size={(i+2)*SIZE_SZ*2:#x}] ", end="")
             chunk = arena.fastbin(i)
@@ -7368,15 +7377,16 @@ def __init__(self) -> None:
         super().__init__(complete=gdb.COMPLETE_LOCATION)
         return
 
+    @parse_arguments({"arena_address": ""}, {})
     @only_if_gdb_running
-    def do_invoke(self, argv: List[str]) -> None:
+    def do_invoke(self, *_: Any, **kwargs: Any) -> None:
+        args = kwargs["arguments"]
         if gef.heap.main_arena is None:
-            err("Invalid Glibc arena")
+            err("Heap not initialized")
             return
-
-        arena_addr = f"*{argv[0]}" if len(argv) == 1 else gef.heap.selected_arena
-        gef_print(titlify(f"Unsorted Bin for arena '{arena_addr!s}'"))
-        nb_chunk = GlibcHeapBinsCommand.pprint_bin(arena_addr, 0, "unsorted_")
+        arena_addr = args.arena_address if args.arena_address else f"{gef.heap.selected_arena.addr:#x}"
+        gef_print(titlify(f"Unsorted Bin for arena at {arena_addr}"))
+        nb_chunk = GlibcHeapBinsCommand.pprint_bin(f"*{arena_addr}", 0, "unsorted_")
         if nb_chunk >= 0:
             info(f"Found {nb_chunk:d} chunks in unsorted bin.")
         return
@@ -7393,17 +7403,19 @@ def __init__(self) -> None:
         super().__init__(complete=gdb.COMPLETE_LOCATION)
         return
 
+    @parse_arguments({"arena_address": ""}, {})
     @only_if_gdb_running
-    def do_invoke(self, argv: List[str]) -> None:
+    def do_invoke(self, *_: Any, **kwargs: Any) -> None:
+        args = kwargs["arguments"]
         if not gef.heap.main_arena:
             err("Heap not initialized")
             return
 
-        arena = GlibcArena(f"*{argv[0]}") if len(argv) == 1 else gef.heap.selected_arena
-        gef_print(titlify(f"Small Bins for arena '{arena!s}'"))
+        arena_addr = args.arena_address if args.arena_address else f"{gef.heap.selected_arena.addr:#x}"
+        gef_print(titlify(f"Small Bins for arena at {arena_addr}"))
         bins = {}
         for i in range(1, 63):
-            nb_chunk = GlibcHeapBinsCommand.pprint_bin(arena, i, "small_")
+            nb_chunk = GlibcHeapBinsCommand.pprint_bin(f"*{arena_addr}", i, "small_")
             if nb_chunk < 0:
                 break
             if nb_chunk > 0:
@@ -7423,17 +7435,19 @@ def __init__(self) -> None:
         super().__init__(complete=gdb.COMPLETE_LOCATION)
         return
 
+    @parse_arguments({"arena_address": ""}, {})
     @only_if_gdb_running
-    def do_invoke(self, argv: List[str]) -> None:
-        if gef.heap.main_arena  is None:
-            err("Invalid Glibc arena")
+    def do_invoke(self, *_: Any, **kwargs: Any) -> None:
+        args = kwargs["arguments"]
+        if gef.heap.main_arena is None:
+            err("Heap not initialized")
             return
 
-        arena_addr = f"*{argv[0]}" if len(argv) == 1 else gef.heap.selected_arena
-        gef_print(titlify(f"Large Bins for arena '{arena_addr!s}'"))
+        arena_addr = args.arena_address if args.arena_address else f"{gef.heap.selected_arena.addr:#x}"
+        gef_print(titlify(f"Large Bins for arena at {arena_addr}"))
         bins = {}
         for i in range(63, 126):
-            nb_chunk = GlibcHeapBinsCommand.pprint_bin(arena_addr, i, "large_")
+            nb_chunk = GlibcHeapBinsCommand.pprint_bin(f"*{arena_addr}", i, "large_")
             if nb_chunk < 0:
                 break
             if nb_chunk > 0:

tests/runtests.py

@@ -301,7 +301,7 @@ def test_cmd_heap_bins_large(self):
         self.assertIn("size=0x420", res)
 
     def test_cmd_heap_bins_non_main(self):
-        cmd = "python gdb.execute(f'heap bins fast {gef.heap.main_arena}')"
+        cmd = "python gdb.execute(f'heap bins fast {gef.heap.main_arena.addr}')"
         before = ["set environment GLIBC_TUNABLES glibc.malloc.tcache_count=0"]
         target = _target("heap-non-main")
         res = gdb_run_silent_cmd(cmd, before=before, target=target)