add test for heap bins

.gitignore

@@ -4,6 +4,7 @@ __pycache__/
 
 # C extensions
 *.so
+!tests/binaries/glibcs/**/*.so
 
 # Distribution / packaging
 .Python

tests/binaries/glibcs/2.33/ld-linux-x86-64.so.2

@@ -0,0 +1 @@
+ld-2.33.so

tests/binaries/glibcs/2.33/libc.so.6

@@ -0,0 +1 @@
+libc-2.33.so

tests/binaries/glibcs/README.md

@@ -0,0 +1,4 @@
+## glibcs used by test
+|     name     | version |         pkgversion         |       docker tag/id       |                              sha256                              |
+| ------------ | ------- | -------------------------- | ------------------------- | ---------------------------------------------------------------- |
+| libc-2.33.so |  2.33   | Ubuntu GLIBC 2.33-0ubuntu5 | ubuntu:21.04/de6f83bfe0b6 | 86ca990a4719b1d4ed8f56e9c6c373e33ad8a40a85fb262cc9ac94ab67feaed0 |

tests/binaries/heap_bins.c

@@ -0,0 +1,131 @@
+#include <stdio.h>
+#include <stdint.h>
+#include <malloc.h>
+
+#define PADDING_SIZE 0x10
+#define TCACHE_SIZE 0x20
+#define TCACHE_COUNT 0x7
+#define FASTBIN_SIZE TCACHE_SIZE
+#define FASTBIN_COUNT 0x4
+#define SMALLBIN_SIZE 0x200
+#define SMALLBIN_COUNT 0x3
+#define LARGEBIN_SIZE 0x500
+#define LARGEBIN_COUNT 0x3
+
+// for export only
+const size_t padding_size = PADDING_SIZE;
+const size_t tcache_size = TCACHE_SIZE;
+const size_t tcache_count = TCACHE_COUNT;
+const size_t fastbin_size = FASTBIN_SIZE;
+const size_t fastbin_count = FASTBIN_COUNT;
+const size_t smallbin_size = SMALLBIN_SIZE;
+const size_t smallbin_count = SMALLBIN_COUNT;
+const size_t largebin_size = LARGEBIN_SIZE;
+const size_t largebin_count = LARGEBIN_COUNT;
+
+int break_id = 0;
+void *tcache[TCACHE_COUNT];
+void *fastbin[FASTBIN_COUNT];
+void *smallbin[SMALLBIN_COUNT + TCACHE_COUNT];
+void *largebin[LARGEBIN_COUNT];
+
+void breakpoint()
+{
+    fprintf(stderr, "Breakpoint #%d\n", ++break_id);
+    return;
+}
+
+void alloc_chunks()
+{
+    void *padding;
+    for (int i = 0; i < TCACHE_COUNT; i++)
+        tcache[i] = malloc(TCACHE_SIZE);
+    for (int i = 0; i < FASTBIN_COUNT; i++)
+        fastbin[i] = malloc(FASTBIN_SIZE);
+    for (int i = 0; i < SMALLBIN_COUNT + TCACHE_COUNT; i++)
+    {
+        smallbin[i] = malloc(SMALLBIN_SIZE);
+        // prevent consolidate
+        padding = malloc(PADDING_SIZE);
+    }
+    for (int i = 0; i < LARGEBIN_COUNT; i++)
+    {
+        largebin[i] = malloc(LARGEBIN_SIZE);
+        // prevent consolidate
+        padding = malloc(PADDING_SIZE);
+    }
+    breakpoint();
+    return;
+}
+
+void tcache_test()
+{
+    for (int i = 0; i < TCACHE_COUNT; i++)
+        free(tcache[i]);
+    breakpoint();
+    return;
+}
+
+void fastbin_test()
+{
+    // tcache is already full, so freed chunk will be put into fastbin
+    for (int i = 0; i < FASTBIN_COUNT; i++)
+        free(fastbin[i]);
+    breakpoint();
+    return;
+}
+
+void unsortedbin_test()
+{
+    for (int i = 0; i < SMALLBIN_COUNT + TCACHE_COUNT; i++)
+        free(smallbin[i]);
+    breakpoint();
+    return;
+}
+
+void smallbin_test()
+{
+    void *tmp;
+    // trigger unsortedbin consolidate
+    tmp = malloc(SMALLBIN_SIZE + 0x10);
+    breakpoint();
+    return;
+}
+
+void largebin_test()
+{
+    void *tmp;
+    for (int i = 0; i < LARGEBIN_COUNT; i++)
+        free(largebin[i]);
+    tmp = malloc(LARGEBIN_SIZE + 0x10);
+    breakpoint();
+    return;
+}
+
+void breakchains()
+{
+    *((uint64_t *)(smallbin[TCACHE_COUNT])) = 0xdeadbeef;
+    *((uint64_t *)(largebin[0]) + 1) = 0xdeadbeef;
+    breakpoint();
+}
+
+void initial()
+{
+    setbuf(stdin, NULL);
+    setbuf(stdout, NULL);
+    setbuf(stderr, NULL);
+    return;
+}
+
+int main()
+{
+    initial();
+    alloc_chunks();
+    tcache_test();
+    fastbin_test();
+    unsortedbin_test();
+    smallbin_test();
+    largebin_test();
+    breakchains();
+    return 0;
+}

tests/binaries/makefile

@@ -28,7 +28,9 @@ else
 CFLAGS		  +=	  -O1
 endif
 
+PWD=$(shell pwd)
 GLIBC=/glibc_versions/2.29/tcache_x64
+GLIBC_2_33=$(PWD)/glibcs/2.33
 
 .PHONY : all clean
 
@@ -73,6 +75,12 @@ heap_bugs: heap_bugs.c
 	-Wl,--dynamic-linker=${GLIBC}/elf/ld.so \
 	-g -o heap_bugs.out heap_bugs.c
 
+heap_bins: heap_bins.c
+	${CC} \
+	-Wl,-rpath=${GLIBC_2_33} \
+	-Wl,--dynamic-linker=${GLIBC_2_33}/ld-linux-x86-64.so.2 \
+	-g -O0 -o heap_bins.out heap_bins.c
+
 clean :
 	@echo "[+] Cleaning stuff"
 	@rm -f $(COMPILED) $(LINKED) $(COMPILED_ASM) $(LINKED_ASM) $(COMPILED_GO)

tests/test_heap_bins.py

@@ -0,0 +1,132 @@
+import gdb
+
+import pwndbg.heap
+import pwndbg.memory
+import pwndbg.symbol
+import pwndbg.vmmap
+import tests
+
+BINARY = tests.binaries.get('heap_bins.out')
+
+
+def test_heap_bins(start_binary):
+    """
+    Tests pwndbg heap bins commands
+    """
+    start_binary(BINARY)
+    gdb.execute('set context-output /dev/null')
+    gdb.execute('b breakpoint', to_string=True)
+
+    # check if all bins are empty at first
+    gdb.execute('continue')
+    allocator = pwndbg.heap.current
+
+    addr = pwndbg.symbol.address("tcache_size")
+    tcache_size = allocator._request2size(pwndbg.memory.u64(addr))
+    addr = pwndbg.symbol.address("tcache_count")
+    tcache_count = pwndbg.memory.u64(addr)
+    addr = pwndbg.symbol.address("fastbin_size")
+    fastbin_size = allocator._request2size(pwndbg.memory.u64(addr))
+    addr = pwndbg.symbol.address("fastbin_count")
+    fastbin_count = pwndbg.memory.u64(addr)
+    addr = pwndbg.symbol.address("smallbin_size")
+    smallbin_size = allocator._request2size(pwndbg.memory.u64(addr))
+    addr = pwndbg.symbol.address("smallbin_count")
+    smallbin_count = pwndbg.memory.u64(addr)
+    addr = pwndbg.symbol.address("largebin_size")
+    largebin_size = allocator._request2size(pwndbg.memory.u64(addr))
+    addr = pwndbg.symbol.address("largebin_count")
+    largebin_count = pwndbg.memory.u64(addr)
+
+    result = allocator.tcachebins()
+    assert result['type'] == 'tcachebins'
+    assert tcache_size in result
+    assert result[tcache_size][1] == 0 and len(result[tcache_size][0]) == 1
+
+    result = allocator.fastbins()
+    assert result['type'] == 'fastbins'
+    assert fastbin_size in result
+    assert len(result[fastbin_size]) == 1
+
+    result = allocator.unsortedbin()
+    assert result['type'] == 'unsortedbin'
+    assert len(result['all'][0]) == 1
+    assert not result['all'][2]
+
+    result = allocator.smallbins()
+    assert result['type'] == 'smallbins'
+    assert smallbin_size in result
+    assert len(result[smallbin_size][0]) == 1 and len(result[smallbin_size][1]) == 1
+    assert not result[smallbin_size][2]
+
+    result = allocator.largebins()
+    assert result['type'] == 'largebins'
+    largebin_size = list(result.items())[allocator.largebin_index(largebin_size) - 64][0]
+    assert largebin_size in result
+    assert len(result[largebin_size][0]) == 1 and len(result[largebin_size][1]) == 1
+    assert not result[largebin_size][2]
+
+    # check tcache
+    gdb.execute('continue')
+
+    result = allocator.tcachebins()
+    assert result['type'] == 'tcachebins'
+    assert tcache_size in result
+    assert result[tcache_size][1] == tcache_count and len(result[tcache_size][0]) == tcache_count + 1
+    for addr in result[tcache_size][0][:-1]:
+        assert pwndbg.vmmap.find(addr)
+
+    # check fastbin
+    gdb.execute('continue')
+
+    result = allocator.fastbins()
+    assert result['type'] == 'fastbins'
+    assert (fastbin_size in result) and (len(result[fastbin_size]) == fastbin_count + 1)
+    for addr in result[fastbin_size][:-1]:
+        assert pwndbg.vmmap.find(addr)
+
+    # check unsortedbin
+    gdb.execute('continue')
+
+    result = allocator.unsortedbin()
+    assert result['type'] == 'unsortedbin'
+    assert len(result['all'][0]) == smallbin_count + 2 and len(result['all'][1]) == smallbin_count + 2
+    assert not result['all'][2]
+    for addr in result['all'][0][:-1]:
+        assert pwndbg.vmmap.find(addr)
+    for addr in result['all'][1][:-1]:
+        assert pwndbg.vmmap.find(addr)
+
+    # check smallbins
+    gdb.execute('continue')
+
+    result = allocator.smallbins()
+    assert result['type'] == 'smallbins'
+    assert len(result[smallbin_size][0]) == smallbin_count + 2 and len(result[smallbin_size][1]) == smallbin_count + 2
+    assert not result[smallbin_size][2]
+    for addr in result[smallbin_size][0][:-1]:
+        assert pwndbg.vmmap.find(addr)
+    for addr in result[smallbin_size][1][:-1]:
+        assert pwndbg.vmmap.find(addr)
+
+    # check largebins
+    gdb.execute('continue')
+
+    result = allocator.largebins()
+    assert result['type'] == 'largebins'
+    assert len(result[largebin_size][0]) == largebin_count + 2 and len(result[largebin_size][1]) == largebin_count + 2
+    assert not result[largebin_size][2]
+    for addr in result[largebin_size][0][:-1]:
+        assert pwndbg.vmmap.find(addr)
+    for addr in result[largebin_size][1][:-1]:
+        assert pwndbg.vmmap.find(addr)
+
+    # check corrupted
+    gdb.execute('continue')
+    result = allocator.smallbins()
+    assert result['type'] == 'smallbins'
+    assert result[smallbin_size][2]
+
+    result = allocator.largebins()
+    assert result['type'] == 'largebins'
+    assert result[largebin_size][2]