Add `malloc_chunk` command tests for heuristic heap (#1234)

tests/test_heap.py

@@ -418,92 +418,99 @@ def test_find_fake_fast_command(start_binary):
     gdb.execute("find_fake_fast (void*)&fake_chunk+0x70")
 
 
-def test_malloc_chunk_command(start_binary):
+def generate_expected_malloc_chunk_output(chunks):
-    start_binary(HEAP_MALLOC_CHUNK)
+    expected = {}
-    gdb.execute("break break_here")
+    expected["allocated"] = [
-    gdb.execute("continue")
-
-    allocated_chunk = pwndbg.gdblib.memory.poi(
-        pwndbg.heap.current.malloc_chunk, gdb.lookup_symbol("allocated_chunk")[0].value()
-    )
-    tcache_chunk = pwndbg.gdblib.memory.poi(
-        pwndbg.heap.current.malloc_chunk, gdb.lookup_symbol("tcache_chunk")[0].value()
-    )
-    fast_chunk = pwndbg.gdblib.memory.poi(
-        pwndbg.heap.current.malloc_chunk, gdb.lookup_symbol("fast_chunk")[0].value()
-    )
-    small_chunk = pwndbg.gdblib.memory.poi(
-        pwndbg.heap.current.malloc_chunk, gdb.lookup_symbol("small_chunk")[0].value()
-    )
-    large_chunk = pwndbg.gdblib.memory.poi(
-        pwndbg.heap.current.malloc_chunk, gdb.lookup_symbol("large_chunk")[0].value()
-    )
-    unsorted_chunk = pwndbg.gdblib.memory.poi(
-        pwndbg.heap.current.malloc_chunk, gdb.lookup_symbol("unsorted_chunk")[0].value()
-    )
-
-    allocated_result = gdb.execute("malloc_chunk allocated_chunk", to_string=True).splitlines()
-    tcache_result = gdb.execute("malloc_chunk tcache_chunk", to_string=True).splitlines()
-    fast_result = gdb.execute("malloc_chunk fast_chunk", to_string=True).splitlines()
-    small_result = gdb.execute("malloc_chunk small_chunk", to_string=True).splitlines()
-    large_result = gdb.execute("malloc_chunk large_chunk", to_string=True).splitlines()
-    unsorted_result = gdb.execute("malloc_chunk unsorted_chunk", to_string=True).splitlines()
-
-    allocated_expected = [
         "Allocated chunk | PREV_INUSE",
-        f"Addr: {allocated_chunk.address}",
+        f"Addr: {chunks['allocated'].address}",
-        f"Size: 0x{int(allocated_chunk['mchunk_size' if 'mchunk_size' in (f.name for f in allocated_chunk.type.fields()) else 'size']):02x}",
+        f"Size: 0x{int(chunks['allocated']['mchunk_size' if 'mchunk_size' in (f.name for f in chunks['allocated'].type.fields()) else 'size']):02x}",
         "",
     ]
 
-    tcache_expected = [
+    expected["tcache"] = [
         f"Free chunk ({'tcache' if pwndbg.heap.current.has_tcache else 'fastbins'}) | PREV_INUSE",
-        f"Addr: {tcache_chunk.address}",
+        f"Addr: {chunks['tcache'].address}",
-        f"Size: 0x{int(tcache_chunk['mchunk_size' if 'mchunk_size' in (f.name for f in tcache_chunk.type.fields()) else 'size']):02x}",
+        f"Size: 0x{int(chunks['tcache']['mchunk_size' if 'mchunk_size' in (f.name for f in chunks['tcache'].type.fields()) else 'size']):02x}",
-        f"fd: 0x{int(tcache_chunk['fd']):02x}",
+        f"fd: 0x{int(chunks['tcache']['fd']):02x}",
         "",
     ]
 
-    fast_expected = [
+    expected["fast"] = [
         "Free chunk (fastbins) | PREV_INUSE",
-        f"Addr: {fast_chunk.address}",
+        f"Addr: {chunks['fast'].address}",
-        f"Size: 0x{int(fast_chunk['mchunk_size' if 'mchunk_size' in (f.name for f in fast_chunk.type.fields()) else 'size']):02x}",
+        f"Size: 0x{int(chunks['fast']['mchunk_size' if 'mchunk_size' in (f.name for f in chunks['fast'].type.fields()) else 'size']):02x}",
-        f"fd: 0x{int(fast_chunk['fd']):02x}",
+        f"fd: 0x{int(chunks['fast']['fd']):02x}",
         "",
     ]
 
-    small_expected = [
+    expected["small"] = [
         "Free chunk (smallbins) | PREV_INUSE",
-        f"Addr: {small_chunk.address}",
+        f"Addr: {chunks['small'].address}",
-        f"Size: 0x{int(small_chunk['mchunk_size' if 'mchunk_size' in (f.name for f in small_chunk.type.fields()) else 'size']):02x}",
+        f"Size: 0x{int(chunks['small']['mchunk_size' if 'mchunk_size' in (f.name for f in chunks['small'].type.fields()) else 'size']):02x}",
-        f"fd: 0x{int(small_chunk['fd']):02x}",
+        f"fd: 0x{int(chunks['small']['fd']):02x}",
-        f"bk: 0x{int(small_chunk['bk']):02x}",
+        f"bk: 0x{int(chunks['small']['bk']):02x}",
         "",
     ]
 
-    large_expected = [
+    expected["large"] = [
         "Free chunk (largebins) | PREV_INUSE",
-        f"Addr: {large_chunk.address}",
+        f"Addr: {chunks['large'].address}",
-        f"Size: 0x{int(large_chunk['mchunk_size' if 'mchunk_size' in (f.name for f in large_chunk.type.fields()) else 'size']):02x}",
+        f"Size: 0x{int(chunks['large']['mchunk_size' if 'mchunk_size' in (f.name for f in chunks['large'].type.fields()) else 'size']):02x}",
-        f"fd: 0x{int(large_chunk['fd']):02x}",
+        f"fd: 0x{int(chunks['large']['fd']):02x}",
-        f"bk: 0x{int(large_chunk['bk']):02x}",
+        f"bk: 0x{int(chunks['large']['bk']):02x}",
-        f"fd_nextsize: 0x{int(large_chunk['fd_nextsize']):02x}",
+        f"fd_nextsize: 0x{int(chunks['large']['fd_nextsize']):02x}",
-        f"bk_nextsize: 0x{int(large_chunk['bk_nextsize']):02x}",
+        f"bk_nextsize: 0x{int(chunks['large']['bk_nextsize']):02x}",
         "",
     ]
 
-    unsorted_expected = [
+    expected["unsorted"] = [
         "Free chunk (unsortedbin) | PREV_INUSE",
-        f"Addr: {unsorted_chunk.address}",
+        f"Addr: {chunks['unsorted'].address}",
-        f"Size: 0x{int(unsorted_chunk['mchunk_size' if 'mchunk_size' in (f.name for f in unsorted_chunk.type.fields()) else 'size']):02x}",
+        f"Size: 0x{int(chunks['unsorted']['mchunk_size' if 'mchunk_size' in (f.name for f in chunks['unsorted'].type.fields()) else 'size']):02x}",
-        f"fd: 0x{int(unsorted_chunk['fd']):02x}",
+        f"fd: 0x{int(chunks['unsorted']['fd']):02x}",
-        f"bk: 0x{int(unsorted_chunk['bk']):02x}",
+        f"bk: 0x{int(chunks['unsorted']['bk']):02x}",
         "",
     ]
 
-    assert allocated_result == allocated_expected
+    return expected
-    assert tcache_result == tcache_expected
+
-    assert fast_result == fast_expected
+
-    assert small_result == small_expected
+def test_malloc_chunk_command(start_binary):
-    assert large_result == large_expected
+    start_binary(HEAP_MALLOC_CHUNK)
-    assert unsorted_result == unsorted_expected
+    gdb.execute("break break_here")
+    gdb.execute("continue")
+
+    chunks = {}
+    results = {}
+    chunk_types = ["allocated", "tcache", "fast", "small", "large", "unsorted"]
+    for name in chunk_types:
+        chunks[name] = pwndbg.gdblib.memory.poi(
+            pwndbg.heap.current.malloc_chunk, gdb.lookup_symbol(f"{name}_chunk")[0].value()
+        )
+        results[name] = gdb.execute(f"malloc_chunk {name}_chunk", to_string=True).splitlines()
+
+    expected = generate_expected_malloc_chunk_output(chunks)
+
+    for name in chunk_types:
+        assert results[name] == expected[name]
+
+
+def test_malloc_chunk_command_heuristic(start_binary):
+    start_binary(HEAP_MALLOC_CHUNK)
+    gdb.execute("set resolve-heap-via-heuristic on")
+    gdb.execute("break break_here")
+    gdb.execute("continue")
+
+    chunks = {}
+    results = {}
+    chunk_types = ["allocated", "tcache", "fast", "small", "large", "unsorted"]
+    for name in chunk_types:
+        chunks[name] = pwndbg.heap.current.malloc_chunk(
+            gdb.lookup_symbol(f"{name}_chunk")[0].value()
+        )
+        results[name] = gdb.execute(f"malloc_chunk {name}_chunk", to_string=True).splitlines()
+
+    expected = generate_expected_malloc_chunk_output(chunks)
+
+    for name in chunk_types:
+        assert results[name] == expected[name]