From cf114437355b6089d788703922664e27fd35b48f Mon Sep 17 00:00:00 2001
From: Lonny Wong <lonnywong@qq.com>
Date: Thu, 13 Oct 2022 19:45:56 +0800
Subject: [PATCH] add test for max-visualize-chunk-size

---
 pwndbg/commands/heap.py            |  2 +-
 tests/binaries/heap_vis.c          | 12 ++++++++++
 tests/heap/test_vis_heap_chunks.py | 35 ++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/pwndbg/commands/heap.py b/pwndbg/commands/heap.py
index 2d5511713..153f433da 100644
--- a/pwndbg/commands/heap.py
+++ b/pwndbg/commands/heap.py
@@ -767,7 +767,7 @@ def vis_heap_chunks(addr=None, count=None, naive=None, display_all=None):
 
     print(out)
 
-    if has_huge_chunk and max_visualize_chunk_size == 0:
+    if has_huge_chunk and pwndbg.config.max_visualize_chunk_size == 0:
         print(
             message.warn(
                 "You can try `set max-visualize-chunk-size 0x500` and re-run this command.\n"
diff --git a/tests/binaries/heap_vis.c b/tests/binaries/heap_vis.c
index b93eb037c..26d01d34b 100644
--- a/tests/binaries/heap_vis.c
+++ b/tests/binaries/heap_vis.c
@@ -1,5 +1,6 @@
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
 
 #include <pthread.h>
 
@@ -21,6 +22,17 @@ int main() {
 
     break_here();
 
+    allocs[3] = malloc(0x1000);
+    allocs[4] = malloc(0x2000);
+    free(allocs[3]);
+
+    break_here();
+
+    // mock overflow changing the chunk size
+    memset(allocs[0] - sizeof(void*), 'A', 8);
+
+    break_here();
+
     // We do not really need it for our test
     // but we need it so that our CI test pass can get TLS variables
     // See:
diff --git a/tests/heap/test_vis_heap_chunks.py b/tests/heap/test_vis_heap_chunks.py
index d91d73d53..73b1d0890 100644
--- a/tests/heap/test_vis_heap_chunks.py
+++ b/tests/heap/test_vis_heap_chunks.py
@@ -139,6 +139,7 @@ def test_vis_heap_chunk_command(start_binary):
     assert result_all2 == expected4_b
 
     del result_all2
+    del expected4_b
 
     ## Continue, so that alloc[1] is freed
     gdb.execute("continue")
@@ -186,3 +187,37 @@ def test_vis_heap_chunk_command(start_binary):
     expected_all3.append(vis_heap_line(suffix="\t <-- Top chunk"))
 
     assert result_all3 == expected_all3
+
+    del result_all3
+    del expected_all3
+
+    # Continue, malloc two large chunks and free one
+    gdb.execute("continue")
+
+    # Get default result without max-visualize-chunk-size setting
+    default_result = gdb.execute("vis_heap_chunk", to_string=True).splitlines()
+    assert len(default_result) > 0x300
+
+    # Set max display size to 100 (no "0x" for misalignment)
+    gdb.execute("set max-visualize-chunk-size 100")
+
+    omitted_result = gdb.execute("vis_heap_chunk", to_string=True).splitlines()
+    assert len(omitted_result) < 0x30
+    for omitted_line in omitted_result:
+        assert omitted_line in default_result or set(omitted_line) == {"."}
+
+    display_all_result = gdb.execute("vis_heap_chunk -a", to_string=True).splitlines()
+    assert display_all_result == default_result
+
+    del default_result
+    del omitted_result
+    del display_all_result
+
+    # Continue, mock overflow changing the chunk size
+    gdb.execute("continue")
+
+    overflow_result = gdb.execute("vis_heap_chunk", to_string=True)
+    assert "\t0x0000000000000000\t0x4141414141414141\t........AAAAAAAA" in overflow_result
+    assert len(overflow_result.splitlines()) < 0x500
+
+    del overflow_result