Integrate Chunk class into malloc_chunk command

pwndbg/commands/heap.py

@@ -15,6 +15,7 @@ from pwndbg.color import message
 from pwndbg.commands.config import extend_value_with_default
 from pwndbg.commands.config import get_config_parameters
 from pwndbg.commands.config import print_row
+from pwndbg.heap.ptmalloc import Chunk
 
 
 def read_chunk(addr):
@@ -278,39 +279,33 @@ parser.add_argument(
 @pwndbg.commands.OnlyWhenHeapIsInitialized
 def malloc_chunk(addr, fake=False, verbose=False, simple=False):
     """Print a malloc_chunk struct's contents."""
-    # points to the real start of the chunk
+    chunk = Chunk(addr)
-    cursor = int(addr)
 
     allocator = pwndbg.heap.current
     ptr_size = allocator.size_sz
 
-    size_field = pwndbg.gdblib.memory.u(cursor + allocator.chunk_key_offset("size"))
-    real_size = size_field & ~allocator.malloc_align_mask
-
     headers_to_print = []  # both state (free/allocated) and flags
     fields_to_print = set()  # in addition to addr and size
-    out_fields = "Addr: {}\n".format(M.get(cursor))
+    out_fields = "Addr: {}\n".format(M.get(chunk.address))
 
     if fake:
         headers_to_print.append(message.on("Fake chunk"))
         verbose = True  # print all fields for fake chunks
 
     if simple:
-        chunk = read_chunk(cursor)
-
         if not headers_to_print:
-            headers_to_print.append(message.hint(M.get(cursor)))
+            headers_to_print.append(message.hint(M.get(chunk.address)))
 
         out_fields = ""
         verbose = True
     else:
-        arena = allocator.get_arena_for_chunk(cursor)
+        arena = allocator.get_arena_for_chunk(chunk.address)
         arena_address = None
         is_top = False
         if not fake and arena:
             arena_address = arena.address
             top_chunk = arena["top"]
-            if cursor == top_chunk:
+            if chunk.address == top_chunk:
                 headers_to_print.append(message.off("Top chunk"))
                 is_top = True
 
@@ -322,35 +317,35 @@ def malloc_chunk(addr, fake=False, verbose=False, simple=False):
             if allocator.has_tcache():
                 tcachebins = allocator.tcachebins(None)
 
-            if real_size in fastbins.keys() and cursor in fastbins[real_size]:
+            if chunk.size in fastbins.keys() and chunk.address in fastbins[chunk.size]:
                 headers_to_print.append(message.on("Free chunk (fastbins)"))
                 if not verbose:
                     fields_to_print.add("fd")
 
-            elif real_size in smallbins.keys() and cursor in bin_addrs(
+            elif chunk.size in smallbins.keys() and chunk.address in bin_addrs(
-                smallbins[real_size], "smallbins"
+                smallbins[chunk.size], "smallbins"
             ):
                 headers_to_print.append(message.on("Free chunk (smallbins)"))
                 if not verbose:
                     fields_to_print.update(["fd", "bk"])
 
-            elif real_size >= list(largebins.items())[0][0] and cursor in bin_addrs(
+            elif chunk.size >= list(largebins.items())[0][0] and chunk.address in bin_addrs(
-                largebins[(list(largebins.items())[allocator.largebin_index(real_size) - 64][0])],
+                largebins[(list(largebins.items())[allocator.largebin_index(chunk.size) - 64][0])],
                 "largebins",
             ):
                 headers_to_print.append(message.on("Free chunk (largebins)"))
                 if not verbose:
                     fields_to_print.update(["fd", "bk", "fd_nextsize", "bk_nextsize"])
 
-            elif cursor in bin_addrs(unsortedbin["all"], "unsortedbin"):
+            elif chunk.address in bin_addrs(unsortedbin["all"], "unsortedbin"):
                 headers_to_print.append(message.on("Free chunk (unsortedbin)"))
                 if not verbose:
                     fields_to_print.update(["fd", "bk"])
 
             elif (
                 allocator.has_tcache()
-                and real_size in tcachebins.keys()
+                and chunk.size in tcachebins.keys()
-                and cursor + ptr_size * 2 in bin_addrs(tcachebins[real_size], "tcachebins")
+                and chunk.address + ptr_size * 2 in bin_addrs(tcachebins[chunk.size], "tcachebins")
             ):
                 headers_to_print.append(message.on("Free chunk (tcache)"))
                 if not verbose:
@@ -362,9 +357,9 @@ def malloc_chunk(addr, fake=False, verbose=False, simple=False):
     if verbose:
         fields_to_print.update(["prev_size", "size", "fd", "bk", "fd_nextsize", "bk_nextsize"])
     else:
-        out_fields += "Size: 0x{:02x}\n".format(size_field)
+        out_fields += "Size: 0x{:02x}\n".format(chunk.size_field)
 
-    prev_inuse, is_mmapped, non_main_arena = allocator.chunk_flags(size_field)
+    prev_inuse, is_mmapped, non_main_arena = allocator.chunk_flags(chunk.size_field)
     if prev_inuse:
         headers_to_print.append(message.hint("PREV_INUSE"))
     if is_mmapped:
@@ -376,7 +371,7 @@ def malloc_chunk(addr, fake=False, verbose=False, simple=False):
     for field_to_print in fields_ordered:
         if field_to_print in fields_to_print:
             out_fields += message.system(field_to_print) + ": 0x{:02x}\n".format(
-                pwndbg.gdblib.memory.u(cursor + allocator.chunk_key_offset(field_to_print))
+                getattr(chunk, field_to_print)
             )
 
     print(" | ".join(headers_to_print) + "\n" + out_fields)