diff --git a/pwndbg/commands/heap.py b/pwndbg/commands/heap.py index cd294ee6..cfca6bbb 100644 --- a/pwndbg/commands/heap.py +++ b/pwndbg/commands/heap.py @@ -299,21 +299,16 @@ def malloc_chunk(addr, fake=False, verbose=False, simple=False): out_fields = "" verbose = True else: - arena = allocator.get_arena_for_chunk(chunk.address) - arena_address = None - is_top = False + arena = chunk.arena if not fake and arena: - arena_address = arena.address - top_chunk = arena["top"] - if chunk.address == top_chunk: + if chunk.is_top_chunk: headers_to_print.append(message.off("Top chunk")) - is_top = True - if not is_top: - fastbins = allocator.fastbins(arena_address) or {} - smallbins = allocator.smallbins(arena_address) or {} - largebins = allocator.largebins(arena_address) or {} - unsortedbin = allocator.unsortedbin(arena_address) or {} + if not chunk.is_top_chunk: + fastbins = allocator.fastbins(arena.address) or {} + smallbins = allocator.smallbins(arena.address) or {} + largebins = allocator.largebins(arena.address) or {} + unsortedbin = allocator.unsortedbin(arena.address) or {} if allocator.has_tcache(): tcachebins = allocator.tcachebins(None) diff --git a/pwndbg/heap/ptmalloc.py b/pwndbg/heap/ptmalloc.py index 43197b09..6184cadd 100644 --- a/pwndbg/heap/ptmalloc.py +++ b/pwndbg/heap/ptmalloc.py @@ -32,8 +32,26 @@ def heap_for_ptr(ptr): class Chunk: - def __init__(self, addr): - if type(pwndbg.heap.current.malloc_chunk) == gdb.Type: + __slots__ = ( + "_gdbValue", + "address", + "_prev_size", + "_size", + "_real_size", + "_flags", + "_non_main_arena", + "_is_mmapped", + "_prev_inuse", + "_fd", + "_bk", + "_fd_nextsize", + "_bk_nextsize", + "_arena", + "_is_top_chunk", + ) + + def __init__(self, addr, arena=None): + if isinstance(pwndbg.heap.current.malloc_chunk, gdb.Type): self._gdbValue = pwndbg.gdblib.memory.poi(pwndbg.heap.current.malloc_chunk, addr) else: self._gdbValue = pwndbg.heap.current.malloc_chunk(addr) @@ -49,8 +67,8 @@ class Chunk: self._bk = None self._fd_nextsize = None self._bk_nextsize = None - - # TODO key, REVEAL_PTR + self._arena = arena + self._is_top_chunk = None # Some chunk fields were renamed in GLIBC 2.25 master branch. def __match_renamed_field(self, field): @@ -60,7 +78,7 @@ class Chunk: } for field_name in field_renames[field]: - if field_name in (f.name for f in self._gdbValue.type.fields()): + if gdb.types.has_field(self._gdbValue.type, field_name): return field_name raise ValueError(f"Chunk field name did not match any of {field_renames[field]}.") @@ -177,22 +195,61 @@ class Chunk: return self._bk_nextsize - # TODO Other useful methods e.g. next_chunk(), __iter__, __str__ + @property + def arena(self): + if self._arena is None: + try: + ar_ptr = pwndbg.heap.current.get_heap(self.address)["ar_ptr"] + ar_ptr.fetch_lazy() + except Exception: + ar_ptr = None + if ar_ptr is not None and ar_ptr in (ar.address for ar in pwndbg.heap.current.arenas): + self._arena = Arena(ar_ptr) + else: + self._arena = Arena(pwndbg.heap.current.main_arena.address) + + return self._arena + + @property + def is_top_chunk(self): + if self._is_top_chunk is None: + ar = self.arena + if ar is not None and self.address == ar.top: + self._is_top_chunk = True + else: + self._is_top_chunk = False + + return self._is_top_chunk class Arena: - def __init__(self, addr, heaps): - self.addr = addr + __slots__ = ("_gdbValue", "address", "is_main_arena", "_top", "heaps") + + def __init__(self, addr, heaps=None): + if isinstance(pwndbg.heap.current.malloc_state, gdb.Type): + self._gdbValue = pwndbg.gdblib.memory.poi(pwndbg.heap.current.malloc_state, addr) + else: + self._gdbValue = pwndbg.heap.current.malloc_state(addr) + self.address = int(self._gdbValue.address) + self.is_main_arena = self.address == pwndbg.heap.current.main_arena.address + self._top = None self.heaps = heaps + @property + def top(self): + if self._top is None: + try: + self._top = int(self._gdbValue["top"]) + except gdb.MemoryError: + pass + + return self._top + def __str__(self): - res = [] prefix = "[%%%ds] " % (pwndbg.gdblib.arch.ptrsize * 2) prefix_len = len(prefix % ("")) - arena_name = ( - hex(self.addr) if self.addr != pwndbg.heap.current.main_arena.address else "main" - ) - res.append(message.hint(prefix % (arena_name)) + str(self.heaps[0])) + arena_name = "main" if self.is_main_arena else hex(self.address) + res = [message.hint(prefix % (arena_name)) + str(self.heaps[0])] for h in self.heaps[1:]: res.append(" " * prefix_len + str(h))