add number arg for 'heap chunk' command (#745)

* add number arg for 'heap chunk' command

* add doc for heap chunk

* simplify code for heap chunk command

* add testcase for heap chunk

* fix testcase

Co-authored-by: shenyuan <junchaoluan@gmail.com>

docs/commands/heap.md

@@ -60,6 +60,20 @@ Glibc's behavior. To be able to view unaligned chunks as well, you can disable
 this with the `--allow-unaligned` flag. Note that this might result in
 incorrect output.
 
+
+There is an optional `number` argument, to specify the number of chunks printed by this command. To do so, simply provide the `--number` argument:
+
+```
+gef➤ heap chunk --number 6 0x4e5400
+Chunk(addr=0x4e5400, size=0xd0, flags=PREV_INUSE)
+Chunk(addr=0x4e54d0, size=0x1a0, flags=PREV_INUSE)
+Chunk(addr=0x4e5670, size=0x200, flags=PREV_INUSE)
+Chunk(addr=0x4e5870, size=0xbc0, flags=PREV_INUSE)
+Chunk(addr=0x4e6430, size=0x330, flags=PREV_INUSE)
+Chunk(addr=0x4e6760, size=0x4c0, flags=PREV_INUSE)
+
+```
+
 ### `heap arenas` command ###
 
 Multi-threaded programs have different arenas, and the knowledge of the

gef.py

@@ -1217,6 +1217,8 @@ def flags_as_string(self):
         flags = []
         if self.has_p_bit():
             flags.append(Color.colorify("PREV_INUSE", "red bold"))
+        else:
+            flags.append(Color.colorify("! PREV_INUSE", "green bold"))
         if self.has_m_bit():
             flags.append(Color.colorify("IS_MMAPPED", "red bold"))
         if self.has_n_bit():
@@ -7001,13 +7003,13 @@ class GlibcHeapChunkCommand(GenericCommand):
     See https://github.com/sploitfun/lsploits/blob/master/glibc/malloc/malloc.c#L1123."""
 
     _cmdline_ = "heap chunk"
-    _syntax_  = "{:s} [-h] [--allow-unaligned] address".format(_cmdline_)
+    _syntax_  = "{:s} [-h] [--allow-unaligned] [--number] address".format(_cmdline_)
 
     def __init__(self):
         super().__init__(complete=gdb.COMPLETE_LOCATION)
         return
 
-    @parse_arguments({"address": ""}, {"--allow-unaligned": True})
+    @parse_arguments({"address": ""}, {"--allow-unaligned": True, "--number": 1})
     @only_if_gdb_running
     def do_invoke(self, *args, **kwargs):
         args = kwargs["arguments"]
@@ -7016,12 +7018,26 @@ def do_invoke(self, *args, **kwargs):
             self.usage()
             return
 
-        if get_glibc_arena() is None:
-            return
-
         addr = parse_address(args.address)
-        chunk = GlibcChunk(addr, allow_unaligned=args.allow_unaligned)
-        gef_print(chunk.psprint())
+        current_chunk = GlibcChunk(addr, allow_unaligned=args.allow_unaligned)
+
+        if args.number > 1:
+            for _ in range(args.number):
+                if current_chunk.size == 0:
+                    break
+
+                gef_print(str(current_chunk))
+                next_chunk_addr = current_chunk.get_next_chunk_addr()
+                if not Address(value=next_chunk_addr).valid:
+                    break
+
+                next_chunk = current_chunk.get_next_chunk()
+                if next_chunk is None:
+                    break
+
+                current_chunk = next_chunk
+        else:
+            gef_print(current_chunk.psprint())
         return
 
 

tests/runtests.py

@@ -242,6 +242,13 @@ def test_cmd_heap_chunk(self):
         res = gdb_run_silent_cmd(cmd, target=target)
         self.assertNoException(res)
         self.assertIn("NON_MAIN_ARENA flag: ", res)
+
+        cmd = "heap chunk --number 2 p1"
+        self.assertFailIfInactiveSession(gdb_run_cmd(cmd, target=target))
+        res = gdb_run_silent_cmd(cmd, target=target)
+        self.assertNoException(res)
+        heap_lines = [x for x in res.splitlines() if x.startswith("Chunk(addr=")]
+        self.assertTrue(len(heap_lines) == 2)
         return
 
     def test_cmd_heap_chunks(self):