Add Chunk class & template_heap_command

pwndbg/commands/heap.py

@@ -95,6 +95,43 @@ def format_bin(bins, verbose=False, offset=None):
     return result
 
 
+parser = argparse.ArgumentParser()
+parser.description = ("Template heap command. You can ignore this.")
+parser.add_argument("addr", type=int, help="Address of a chunk header.")
+
+
+@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.OnlyWhenRunning
+@pwndbg.commands.OnlyWithResolvedHeapSyms
+@pwndbg.commands.OnlyWhenHeapIsInitialized
+def template_heap_command(addr):
+    """Template heap command with example uses of pwndbg's heap classes."""
+
+    # `addr` is a gdb.Value of one of several types, depending on how the user invoked this command
+    # e.g. template_heap_command <address> vs. template_heap_command <symbol>
+
+    # The `Chunk` class abstracts away many heap & gdb module internals.
+    chunk = pwndbg.heap.ptmalloc.Chunk(addr)
+
+    print(f"chunk.address: 0x{chunk.address:02x}")
+
+    # Be aware that if a chunk field is unreadable (e.g. a fake chunk straddling an unmapped page boundary) it will be None, always check.
+    if chunk.prev_size is not None:
+        print(f"chunk.prev_size: 0x{chunk.prev_size:02x}")
+
+    if chunk.size is not None:
+        print(f"chunk.size: 0x{chunk.size:02x}")
+
+    if chunk.flags is not None:
+        print(f"chunk.flags: {chunk.flags}")
+
+    if chunk.fd is not None:
+        print(f"chunk.fd: 0x{chunk.fd:02x}")
+
+    if chunk.bk is not None:
+        print(f"chunk.bk: 0x{chunk.bk:02x}")
+
+
 parser = argparse.ArgumentParser()
 parser.description = (
     "Iteratively print chunks on a heap, default to the current thread's active heap."

pwndbg/heap/ptmalloc.py

@@ -32,6 +32,83 @@ def heap_for_ptr(ptr):
     return ptr & ~(HEAP_MAX_SIZE - 1)
 
 
+class Chunk:
+    def __init__(self, addr):
+        self._gdbValue = pwndbg.gdblib.memory.poi(pwndbg.heap.current.malloc_chunk, addr)
+        self.address = int(self._gdbValue.address)
+        self._prev_size = None
+        self._size = None
+        self._flags = None
+        self._fd = None
+        self._bk = None
+        # TODO fd_nextsize, bk_nextsize, key, REVEAL_PTR etc.
+
+    @property
+    def prev_size(self):
+        if self._prev_size is None:
+            if "mchunk_prev_size" in (field.name for field in self._gdbValue.type.fields()):
+                prev_size_field_name = "mchunk_prev_size"
+            else:
+                prev_size_field_name = "prev_size"
+
+            try:
+                self._prev_size = int(self._gdbValue[prev_size_field_name])
+            except gdb.MemoryError:
+                pass
+
+        return self._prev_size
+
+    @property
+    def size(self):
+        if self._size is None:
+            if "mchunk_size" in (field.name for field in self._gdbValue.type.fields()):
+                size_field_name = "mchunk_size"
+            else:
+                size_field_name = "size"
+
+            try:
+                self._size = int(self._gdbValue[size_field_name])
+            except gdb.MemoryError:
+                pass
+
+        return self._size
+
+    @property
+    def flags(self):
+        if self._flags is None:
+            sz = self.size
+            if sz is not None:
+                self._flags = {
+                    "non_main_arena": bool(sz & ptmalloc.NON_MAIN_ARENA),
+                    "is_mmapped": bool(sz & ptmalloc.IS_MMAPPED),
+                    "prev_inuse": bool(sz & ptmalloc.PREV_INUSE),
+                }
+
+        return self._flags
+
+    @property
+    def fd(self):
+        if self._fd is None:
+            try:
+                self._fd = int(self._gdbValue["fd"])
+            except gdb.MemoryError:
+                pass
+
+        return self._fd
+
+    @property
+    def bk(self):
+        if self._bk is None:
+            try:
+                self._bk = int(self._gdbValue["bk"])
+            except gdb.MemoryError:
+                pass
+
+        return self._bk
+
+    # TODO Other useful methods e.g. next_chunk(), __iter__, __str__
+
+
 class Arena:
     def __init__(self, addr, heaps):
         self.addr = addr