Repository: pwndbg/pwndbg

Issue:
vis_heap_chunks() output is inaccurate
The following screenshot shows the vis_heap_chunks() command output when introspecting the default heap, which has 3 0x20-sized chunks allocated on it.

![vis_heap_chunks errors](https://user-images.githubusercontent.com/16000770/74610386-acce6c00-50ea-11ea-8411-1837aac693c4.png)

The first command incorrectly shows the quadword at address 0x405040 as belonging to the 3rd (magenta) chunk, when it actually belongs to the 2nd (cyan) chunk. This issue is only present in the dev branch. The output also shows the quadword at address 0x405060 as belonging to the top chunk, when it belongs to the 3rd chunk. This issue is present in both dev & stable branches.

The second command does not show the last quadword of the yellow chunk's user data, despite it being correctly marked as belonging to the yellow chunk in the first command output. This issue is present in both dev & stable branches.


Discussion:
I have a PR ready, let me know if you want it submitted and against which branch.
Can you give some more context so the issue can be reproduced? E.g. an example minimal working binary and your libc version? 

Also, we accept PRs to the `dev` branch.
Here's a minimal working binary, it just requests 3 minimum-sized chunks.
```
#include <stdlib.h>
void main(void) {
    malloc(1);
    malloc(1);
    malloc(1);
}
```
You can link against any version of GLIBC, ignoring the main thread's tcache in versions >= 2.26. The issue only concerns allocated chunks near the top chunk and how they're coloured by vis_heap_chunks().

Allow the binary to make its 3 requests under GDB then issue the `vis_heap_chunks` command without any arguments. Pwndbg will print output similar to the screenshot in the original issue comment, just without the chunk contents ("AAAA" etc) and possibly preceded by a tcache chunk depending on which version of GLIBC you linked against.

Resolving patch:
diff --git a/pwndbg/commands/heap.py b/pwndbg/commands/heap.py
--- a/pwndbg/commands/heap.py
+++ b/pwndbg/commands/heap.py
@@ -86,6 +86,7 @@ def format_bin(bins, verbose=False, offset=None):
 
     return result
 
+
 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.")
@@ -467,51 +468,52 @@ def find_fake_fast(addr, size):
 
 
 vis_heap_chunks_parser = argparse.ArgumentParser(description='Visualize heap chunks at the specified address')
-vis_heap_chunks_parser.add_argument('count', nargs='?', default=10,
-                    help='Number of chunks to visualize')
+vis_heap_chunks_parser.add_argument('count', type=lambda n:max(int(n, 0),1), nargs='?', default=10, help='Number of chunks to visualize')
 vis_heap_chunks_parser.add_argument('address', help='Start address', nargs='?', default=None)
-vis_heap_chunks_parser.add_argument('--naive', '-n', help='Don\'t use end-of-heap heuristics', action='store_true', default=False)
+vis_heap_chunks_parser.add_argument('--naive', '-n', help='Attempt to keep printing beyond the top chunk', action='store_true', default=False)
 
 @pwndbg.commands.ArgparsedCommand(vis_heap_chunks_parser)
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithLibcDebugSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
 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
+    heap_region = main_heap.get_heap_boundaries(address)
     main_arena = main_heap.get_arena_for_chunk(address) if address else main_heap.main_arena
 
     top_chunk = main_arena['top']
+    ptr_size = main_heap.size_sz
 
+    # Build a list of addresses that delimit each chunk.
+    chunk_delims = []
+    cursor = int(address) if address else heap_region.start
 
-    unpack = pwndbg.arch.unpack
-
-    cells_map = {}
-    chunk_id = 0
-    ptr_size = pwndbg.arch.ptrsize
-    while chunk_id < count:
-        prev_size = unpack(pwndbg.memory.read(address, ptr_size))
-        current_size = unpack(pwndbg.memory.read(address+ptr_size, ptr_size))
-        real_size = current_size & ~main_heap.malloc_align_mask
-        prev_inuse = current_size & 1
-        stop_addr = address + real_size
+    for _ in range(count + 1):
+        # Don't read beyond the heap mapping if --naive or corrupted heap.
+        if cursor not in heap_region:
+            chunk_delims.append(heap_region.end)
+            break
 
-        while address < stop_addr and (naive or address < top_chunk):
-            assert address not in cells_map
-            cells_map[address] = chunk_id
-            address += ptr_size
+        size_field = pwndbg.memory.u(cursor + ptr_size)
+        real_size = size_field & ~main_heap.malloc_align_mask
+        prev_inuse = main_heap.chunk_flags(size_field)[0]
 
-        if prev_inuse and (naive or address != top_chunk):
-            cells_map[address - real_size] -= 1
+        # Don't repeatedly operate on the same address (e.g. chunk size of 0).
+        if cursor in chunk_delims or cursor + ptr_size in chunk_delims:
+            break
 
-        chunk_id += 1
+        if prev_inuse:
+            chunk_delims.append(cursor + ptr_size)
+        else:
+            chunk_delims.append(cursor)
 
-        # we reached top chunk, add it's metadata and break
-        if address >= top_chunk:
-            cells_map[address] = chunk_id
-            cells_map[address+ptr_size] = chunk_id
+        if (cursor == top_chunk and not naive) or (cursor == heap_region.end - ptr_size*2):
+            chunk_delims.append(cursor + ptr_size*2)
             break
 
+        cursor += real_size
+
+    # Build the output buffer, changing color at each chunk delimiter.
     # TODO: maybe print free chunks in bold or underlined
     color_funcs = [
         generateColorFunction("yellow"),
@@ -521,7 +523,6 @@ def vis_heap_chunks(address=None, count=None, naive=None):
         generateColorFunction("blue"),
     ]
 
-    addrs = sorted(cells_map.keys())
     bin_collections = [
         pwndbg.heap.current.fastbins(None),
         pwndbg.heap.current.unsortedbin(None),
@@ -535,30 +536,32 @@ def vis_heap_chunks(address=None, count=None, naive=None):
     out = ''
     asc = ''
     labels = []
+    cursor = int(address) if address else heap_region.start
+
+    for c, stop in enumerate(chunk_delims):
+        color_func = color_funcs[c % len(color_funcs)]
 
-    for addr in addrs:
-        if printed % 2 == 0:
-            out += "\n0x%x" % addr
+        while cursor != stop:
+            if printed % 2 == 0:
+                out += "\n0x%x" % cursor
 
-        cell = unpack(pwndbg.memory.read(addr, ptr_size))
-        cell_hex = '\t0x{:0{n}x}'.format(cell, n=ptr_size*2)
+            cell = pwndbg.memory.u(cursor)
+            cell_hex = '\t0x{:0{n}x}'.format(cell, n=ptr_size*2)
 
-        chunk_idx = cells_map[addr]
-        color_func_idx = chunk_idx % len(color_funcs)
-        color_func = color_funcs[color_func_idx]
+            out += color_func(cell_hex)
+            printed += 1
 
-        out += color_func(cell_hex)
-        printed += 1
+            labels.extend(bin_labels(cursor, bin_collections))
+            if cursor == top_chunk:
+                labels.append('Top chunk')
 
-        labels.extend(bin_labels(addr, bin_collections))
-        if addr == top_chunk:
-            labels.append('Top chunk')
+            asc += bin_ascii(pwndbg.memory.read(cursor, ptr_size))
+            if printed % 2 == 0:
+                out += '\t' + color_func(asc) + ('\t <-- ' + ', '.join(labels) if len(labels) else '')
+                asc = ''
+                labels = []
 
-        asc += bin_ascii(pwndbg.memory.read(addr, ptr_size))
-        if printed % 2 == 0:
-            out += '\t' + color_func(asc) + ('\t <-- ' + ', '.join(labels) if len(labels) else '')
-            asc = ''
-            labels = []
+            cursor += ptr_size
 
     print(out)
 
diff --git a/pwndbg/heap/ptmalloc.py b/pwndbg/heap/ptmalloc.py
--- a/pwndbg/heap/ptmalloc.py
+++ b/pwndbg/heap/ptmalloc.py
@@ -372,8 +372,8 @@ def get_heap_boundaries(self, addr=None):
         adresses inside it or a fake Page for the containing heap for non-main arenas.
         """
         page = pwndbg.memory.Page(0, 0, 0, 0)
-        brk = self.get_region()
-        if addr is None or brk.vaddr < addr < brk.vaddr + brk.memsz:
+        brk = self.get_region(addr)
+        if brk == self.get_region():
             # Occasionally, the [heap] vm region and the actual start of the heap are
             # different, e.g. [heap] starts at 0x61f000 but mp_.sbrk_base is 0x620000.
             # Return an adjusted Page object if this is the case.
@@ -382,13 +382,8 @@ def get_heap_boundaries(self, addr=None):
                 page.vaddr = sbrk_base
                 page.memsz = brk.memsz - (sbrk_base - brk.vaddr)
                 return page
-            else:
-                return brk
-        else:
-            page.vaddr = heap_for_ptr(addr)
-            heap = self.get_heap(page.vaddr)
-            page.memsz = int(heap['size'])
-            return page
+
+        return brk
 
 
     def get_region(self, addr=None):


Test patch: