Add multithreaded malloc_chunk tests (#1277)

* Add reset_on_thread decorator

* Apply reset_on_thread to Heap.multithreaded

* Add multithreaded malloc_chunk tests

* Clarify comment in C source

* Clarify expected thread number with assert in test

pwndbg/gdblib/events.py

@@ -112,6 +112,7 @@ registered: Dict[Any, List[Callable]] = {
     gdb.events.new_objfile: [],
     gdb.events.stop: [],
     gdb.events.start: [],
+    gdb.events.new_thread: [],
     gdb.events.before_prompt: [],  # The real event might not exist, but we wrap it
 }
 
@@ -184,6 +185,10 @@ def start(func):
     return connect(func, gdb.events.start, "start")
 
 
+def thread(func):
+    return connect(func, gdb.events.new_thread, "thread")
+
+
 before_prompt = partial(connect, event_handler=gdb.events.before_prompt, name="before_prompt")
 
 

pwndbg/gdblib/hooks.py

@@ -7,6 +7,7 @@ from pwndbg.lib.memoize import reset_on_objfile
 from pwndbg.lib.memoize import reset_on_prompt
 from pwndbg.lib.memoize import reset_on_start
 from pwndbg.lib.memoize import reset_on_stop
+from pwndbg.lib.memoize import reset_on_thread
 from pwndbg.lib.memoize import while_running
 
 # TODO: Combine these `update_*` hook callbacks into one method
@@ -60,6 +61,11 @@ def memoize_on_exit():
     reset_on_exit._reset()
 
 
+@pwndbg.gdblib.events.thread
+def memoize_on_new_thread():
+    reset_on_thread._reset()
+
+
 def init():
     """Calls all GDB hook functions that need to be called when GDB/pwndbg
     itself is loaded, as opposed to when an actual hook event occurs

pwndbg/gdblib/symbol.py

@@ -242,6 +242,7 @@ def address(symbol: str) -> int:
 
 
 @pwndbg.lib.memoize.reset_on_objfile
+@pwndbg.lib.memoize.reset_on_thread
 def static_linkage_symbol_address(symbol):
     if isinstance(symbol, int):
         return symbol

pwndbg/heap/ptmalloc.py

@@ -432,6 +432,7 @@ class Heap(pwndbg.heap.heap.BaseHeap):
 
     @property
     @pwndbg.lib.memoize.reset_on_objfile
+    @pwndbg.lib.memoize.reset_on_thread
     def multithreaded(self):
         """Is malloc operating within a multithreaded environment."""
         addr = pwndbg.gdblib.symbol.address("__libc_multiple_threads")

pwndbg/lib/memoize.py

@@ -151,6 +151,18 @@ class reset_on_cont(memoize):
     _reset = __reset_on_cont
 
 
+class reset_on_thread(memoize):
+    caches = []  # type: List[reset_on_thread]
+    kind = "thread"
+
+    @staticmethod
+    def __reset_on_thread() -> None:
+        for obj in reset_on_thread.caches:
+            obj.clear()
+
+    _reset = __reset_on_thread
+
+
 class while_running(memoize):
     caches = []  # type: List[while_running]
     kind = "running"

tests/binaries/heap_malloc_chunk.c

@@ -12,6 +12,8 @@
 #define mem2chunk(mem) ((void*)(mem) - CHUNK_HDR_SZ)
 
 void break_here(void) {}
+void configure_heap_layout(void);
+void* thread_func(void*);
 
 void* allocated_chunk = NULL;
 void* tcache_chunk = NULL;
@@ -21,6 +23,17 @@ void* large_chunk = NULL;
 void* unsorted_chunk = NULL;
 
 int main(void)
+{
+    configure_heap_layout();
+
+    break_here();
+
+    pthread_t thread;
+    pthread_create(&thread, NULL, thread_func, NULL);
+    pthread_join(thread, NULL);
+}
+
+void configure_heap_layout(void)
 {
     void* chunks[6] = {0};
 
@@ -67,12 +80,16 @@ int main(void)
     small_chunk = mem2chunk(before_remainder + 0x210);
     large_chunk = mem2chunk(large);
     unsorted_chunk = mem2chunk(unsorted);
+}
 
+void* thread_func(void* args)
+{
+    // Initialize a 2nd arena by allocating any size chunk.
+    malloc(0x18);
     break_here();
 
-    // Required for CI build to retrieve TLS variables.
-    // See:
-    // - https://github.com/pwndbg/pwndbg/pull/1086
-    // - https://sourceware.org/bugzilla/show_bug.cgi?id=24548
-    pthread_create(0,0,0,0);
+    configure_heap_layout();
+    break_here();
+
+    pthread_exit(NULL);
 }

tests/heap/test_heap.py

@@ -89,6 +89,37 @@ def test_malloc_chunk_command(start_binary):
     for name in chunk_types:
         assert results[name] == expected[name]
 
+    gdb.execute("continue")
+
+    # Print main thread's chunk from another thread
+    assert gdb.selected_thread().num == 2
+    results["large"] = gdb.execute("malloc_chunk large_chunk", to_string=True).splitlines()
+    expected = generate_expected_malloc_chunk_output(chunks)
+    assert results["large"] == expected["large"]
+
+    gdb.execute("continue")
+
+    # Test some non-main-arena chunks
+    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)
+    expected["allocated"][0] += " | NON_MAIN_ARENA"
+    expected["tcache"][0] += " | NON_MAIN_ARENA"
+    expected["fast"][0] += " | NON_MAIN_ARENA"
+
+    for name in chunk_types:
+        assert results[name] == expected[name]
+
+    # Print another thread's chunk from the main thread
+    gdb.execute("thread 1")
+    assert gdb.selected_thread().num == 1
+    results["large"] = gdb.execute("malloc_chunk large_chunk", to_string=True).splitlines()
+    assert results["large"] == expected["large"]
+
 
 def test_malloc_chunk_command_heuristic(start_binary):
     start_binary(HEAP_MALLOC_CHUNK)
@@ -110,6 +141,37 @@ def test_malloc_chunk_command_heuristic(start_binary):
     for name in chunk_types:
         assert results[name] == expected[name]
 
+    gdb.execute("continue")
+
+    # Print main thread's chunk from another thread
+    assert gdb.selected_thread().num == 2
+    results["large"] = gdb.execute("malloc_chunk large_chunk", to_string=True).splitlines()
+    expected = generate_expected_malloc_chunk_output(chunks)
+    assert results["large"] == expected["large"]
+
+    gdb.execute("continue")
+
+    # Test some non-main-arena chunks
+    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)
+    expected["allocated"][0] += " | NON_MAIN_ARENA"
+    expected["tcache"][0] += " | NON_MAIN_ARENA"
+    expected["fast"][0] += " | NON_MAIN_ARENA"
+
+    for name in chunk_types:
+        assert results[name] == expected[name]
+
+    # Print another thread's chunk from the main thread
+    gdb.execute("thread 1")
+    assert gdb.selected_thread().num == 1
+    results["large"] = gdb.execute("malloc_chunk large_chunk", to_string=True).splitlines()
+    assert results["large"] == expected["large"]
+
 
 class mock_for_heuristic:
     def __init__(self, mock_symbols=[], mock_all=False, mess_up_memory=False):