Don't use top chunk heuristics (#712)

The original code that searched for an arena's top chunk via the heuristic method of iterating through each chunk on a heap was written before the OnlyWithLibcDebugSyms decorator was added to all heap-related commands. This decorator makes the heuristic method unnecessary as the 'top' member of a heap's arena will always be accessible.

Remove the get_top_chunk_addr() function as it only uses the heuristic approach. Use symbols in top_chunk() and vis_heap_chunks() rather than heuristics. Use get_arena_for_chunk() in vis_heap_chunks() to find an arena, originally this used get_arena() with no arguments which always returned the main arena. Fix get_arena_for_chunk(), which misidentified a chunk as a malloc_state struct rather than a malloc_chunk struct, and use read_chunk() rather than memory.poi() which didn't account for the size/mchunk_size symbols used in different versions of GLIBC.

pwndbg/commands/heap.py

@@ -55,7 +55,7 @@ def format_bin(bins, verbose=False, offset=None):
 
         if not verbose and (chain_fd == [0] and not count) and not is_chain_corrupted:
             continue
-        
+
         if bins_type == 'tcachebins':
             limit = 8
             if count <= 7:
@@ -86,34 +86,6 @@ def format_bin(bins, verbose=False, offset=None):
 
     return result
 
-def get_top_chunk_addr():
-    main_heap  = pwndbg.heap.current
-    heap_region = main_heap.get_heap_boundaries()
-	
-    if not heap_region:
-        print(message.error('Could not find the heap'))
-        return
-
-    heap_start = heap_region.vaddr
-    heap_end   = heap_start + heap_region.memsz
-
-    # If we don't know where the main_arena struct is, just iterate
-    # through all the heap objects until we hit the last one
-    last_addr = None
-    addr = heap_start
-    while addr < heap_end:
-        chunk = read_chunk(addr)
-        size = int(chunk['size'])
-
-        # Clear the bottom 3 bits
-        size &= ~7
-        if size == 0:
-            break
-        last_addr = addr
-        addr += size
-    address = last_addr
-    return address
-
 parser = argparse.ArgumentParser()
 parser.description = "Prints out chunks starting from the address specified by `addr`."
 parser.add_argument("addr", nargs="?", type=int, default=None, help="The address of the heap.")
@@ -237,8 +209,8 @@ def top_chunk(addr=None):
     """
     main_heap   = pwndbg.heap.current
     main_arena  = main_heap.get_arena(addr)
-    address = get_top_chunk_addr()
-    
+    address = main_arena['top']
+
     return malloc_chunk(address)
 
 
@@ -507,10 +479,10 @@ vis_heap_chunks_parser.add_argument('--naive', '-n', help='Don\'t use end-of-hea
 def vis_heap_chunks(address=None, count=None, naive=None):
     address = int(address) if address else pwndbg.heap.current.get_heap_boundaries().vaddr
     main_heap = pwndbg.heap.current
-    main_arena = main_heap.get_arena()
+    main_arena = main_heap.get_arena_for_chunk(address) if address else main_heap.main_arena
+
+    top_chunk = main_arena['top']
 
-    top_chunk = get_top_chunk_addr()
- 
 
     unpack = pwndbg.arch.unpack
 

pwndbg/heap/ptmalloc.py

@@ -350,7 +350,7 @@ class Heap(pwndbg.heap.heap.BaseHeap):
 
 
     def get_arena_for_chunk(self,addr):
-        chunk = pwndbg.memory.poi(self.malloc_state,addr)
+        chunk = pwndbg.commands.heap.read_chunk(addr)
         _,_,nm = self.chunk_flags(chunk['size'])
         if nm:
             r=self.get_arena(arena_addr=self.get_heap(addr)['ar_ptr'])