Add unit-test for find_fake_fast command

.github/workflows/tests.yml

@@ -32,11 +32,12 @@ jobs:
         echo 'Installed packages:'
         python -m pip freeze
 
-    # We use `sudo` for `attachp` command tests
+    # We to set `kernel.yama.ptrace_scope=0` for `attachp` command tests
     - name: Run tests
       run: |
         mkdir .cov
-        PWNDBG_GITHUB_ACTIONS_TEST_RUN=1 sudo --preserve-env ./tests.sh
+        sudo sysctl -w kernel.yama.ptrace_scope=0
+        PWNDBG_GITHUB_ACTIONS_TEST_RUN=1 ./tests.sh
 
     - name: Process coverage data
       if: matrix.os == 'ubuntu-22.04'

pwndbg/color/backtrace.py

@@ -1,6 +1,6 @@
 import pwndbg.color.theme as theme
-import pwndbg.gdblib.config as config
 from pwndbg.color import generateColorFunction
+from pwndbg.gdblib import config
 
 config_prefix = theme.add_param("backtrace-prefix", "►", "prefix for current backtrace label")
 config_prefix_color = theme.add_color_param(

pwndbg/color/chain.py

@@ -1,6 +1,6 @@
 import pwndbg.color.theme as theme
-import pwndbg.gdblib.config as config
 from pwndbg.color import generateColorFunction
+from pwndbg.gdblib import config
 
 config_arrow_color = theme.add_color_param(
     "chain-arrow-color", "normal", "color of chain formatting (arrow)"

pwndbg/color/context.py

@@ -1,6 +1,6 @@
 import pwndbg.color.theme as theme
-import pwndbg.gdblib.config as config
 from pwndbg.color import generateColorFunction
+from pwndbg.gdblib import config
 
 config_prefix_color = theme.add_color_param(
     "code-prefix-color", "none", "color for 'context code' command (prefix marker)"

pwndbg/color/disasm.py

@@ -6,10 +6,10 @@ import pwndbg.color.memory as M
 import pwndbg.color.syntax_highlight as H
 import pwndbg.color.theme as theme
 import pwndbg.disasm.jump
-import pwndbg.gdblib.config as config
 from pwndbg.color import generateColorFunction
 from pwndbg.color import ljust_colored
 from pwndbg.color.message import on
+from pwndbg.gdblib import config
 
 capstone_branch_groups = set((capstone.CS_GRP_CALL, capstone.CS_GRP_JUMP))
 

pwndbg/color/enhance.py

@@ -1,6 +1,6 @@
 import pwndbg.color.theme as theme
-import pwndbg.gdblib.config as config
 from pwndbg.color import generateColorFunction
+from pwndbg.gdblib import config
 
 config_integer_color = theme.add_color_param(
     "enhance-integer-value-color", "none", "color of value enhance (integer)"

pwndbg/color/hexdump.py

@@ -1,6 +1,6 @@
 import pwndbg.color.theme as theme
-import pwndbg.gdblib.config as config
 from pwndbg.color import generateColorFunction
+from pwndbg.gdblib import config
 
 config_normal = theme.add_color_param(
     "hexdump-normal-color", "none", "color for hexdump command (normal bytes)"

pwndbg/color/memory.py

@@ -1,8 +1,8 @@
 import pwndbg.color.theme as theme
-import pwndbg.gdblib.config as config
 import pwndbg.gdblib.vmmap
 from pwndbg.color import generateColorFunction
 from pwndbg.color import normal
+from pwndbg.gdblib import config
 
 config_stack = theme.add_color_param("memory-stack-color", "yellow", "color for stack memory")
 config_heap = theme.add_color_param("memory-heap-color", "blue", "color for heap memory")

pwndbg/color/nearpc.py

@@ -1,6 +1,6 @@
 import pwndbg.color.theme as theme
-import pwndbg.gdblib.config as config
 from pwndbg.color import generateColorFunction
+from pwndbg.gdblib import config
 
 config_symbol = theme.add_color_param(
     "nearpc-symbol-color", "normal", "color for nearpc command (symbol)"

pwndbg/color/telescope.py

@@ -1,6 +1,6 @@
 import pwndbg.color.theme as theme
-import pwndbg.gdblib.config as config
 from pwndbg.color import generateColorFunction
+from pwndbg.gdblib import config
 
 offset_color = theme.add_color_param(
     "telescope-offset-color", "normal", "color of the telescope command (offset prefix)"

pwndbg/commands/heap.py

@@ -1,6 +1,5 @@
 import argparse
 import ctypes
-import struct
 
 import gdb
 
@@ -10,6 +9,7 @@ import pwndbg.commands
 import pwndbg.gdblib.config
 import pwndbg.gdblib.typeinfo
 import pwndbg.glibc
+import pwndbg.lib.heap.helpers
 from pwndbg.color import generateColorFunction
 from pwndbg.color import message
 from pwndbg.commands.config import extend_value_with_default
@@ -633,21 +633,8 @@ def find_fake_fast(addr, size=None, align=False):
 
     print(C.banner("FAKE CHUNKS"))
     step = malloc_alignment if align else 1
-    for i in range(0, len(mem), step):
-        candidate = mem[i : i + psize]
-        if len(candidate) == psize:
-            value = struct.unpack(fmt, candidate)[0]
-
-            # Clear any flags
-            value &= ~0xF
-
-            if value < min_fast:
-                continue
-
-            # The value must be less than or equal to the max size we're looking
-            # for, but still be able to reach the target address
-            if value <= size and i + value >= size:
-                malloc_chunk(start + i - psize, fake=True)
+    for offset in pwndbg.lib.heap.helpers.find_fastbin_size(mem, size, step):
+        malloc_chunk(start + offset, fake=True)
 
 
 pwndbg.gdblib.config.add_param(

pwndbg/lib/heap/helpers.py

@@ -0,0 +1,25 @@
+import struct
+
+import pwndbg.gdblib.arch
+
+
+def find_fastbin_size(mem: bytes, max_size: int, step: int):
+    psize = pwndbg.gdblib.arch.ptrsize
+    min_fast = 4 * psize
+    fmt = {"little": "<", "big": ">"}[pwndbg.gdblib.arch.endian] + {4: "I", 8: "Q"}[psize]
+
+    for i in range(0, len(mem), step):
+        candidate = mem[i : i + psize]
+        if len(candidate) == psize:
+            value = struct.unpack(fmt, candidate)[0]
+
+            # Clear any flags
+            value &= ~0xF
+
+            if value < min_fast:
+                continue
+
+            # The value must be less than or equal to the max size we're looking
+            # for, but still be able to reach the target address
+            if value <= max_size and i + value >= max_size:
+                yield i - psize

tests.sh

@@ -1,4 +1,7 @@
 #!/bin/bash
 
-cd tests/gdb-tests
-./tests.sh $@
+# Run integration tests
+(cd tests/gdb-tests && ./tests.sh $@)
+
+# Run unit tests
+coverage run -m pytest tests/unit-tests

tests/unit-tests/test_find_fastbin_size.py

@@ -0,0 +1,80 @@
+import sys
+from unittest.mock import MagicMock
+
+import pytest
+from pwnlib.util.packing import p64
+
+# Replace `pwndbg.commands` module with a mock to prevent import errors, as well
+# as the `load_commands` function
+module_name = "pwndbg.commands"
+module = MagicMock(__name__=module_name, load_commands=lambda: None)
+sys.modules[module_name] = module
+
+# Load the mocks for the `gdb` and `gdblib` modules
+import mocks.gdb
+import mocks.gdblib  # noqa: F401
+
+# We must import the function under test after all the mocks are imported
+from pwndbg.lib.heap.helpers import find_fastbin_size
+
+
+def setup_mem(max_size, offsets):
+    buf = bytearray([0] * max_size)
+    for offset, value in offsets.items():
+        buf[offset : offset + 8] = p64(value)
+
+    return buf
+
+
+def test_too_small():
+    max_size = 0x80
+    offsets = {
+        0x8: 0x10,
+    }
+    buf = setup_mem(max_size, offsets)
+    with pytest.raises(StopIteration):
+        next(find_fastbin_size(buf, max_size, 1))
+
+    with pytest.raises(StopIteration):
+        next(find_fastbin_size(buf, max_size, 8))
+
+
+def test_normal():
+    max_size = 0x20
+    offsets = {
+        0x8: 0x20,
+    }
+    buf = setup_mem(max_size, offsets)
+    assert 0x0 == next(find_fastbin_size(buf, max_size, 1))
+    assert 0x0 == next(find_fastbin_size(buf, max_size, 8))
+
+
+def test_nozero_flags():
+    max_size = 0x20
+    offsets = {
+        0x8: 0x2F,
+    }
+    buf = setup_mem(max_size, offsets)
+    assert 0x0 == next(find_fastbin_size(buf, max_size, 1))
+    assert 0x0 == next(find_fastbin_size(buf, max_size, 8))
+
+
+def test_normal():
+    max_size = 0x20
+    offsets = {
+        0x8: 0x20,
+    }
+    buf = setup_mem(max_size, offsets)
+    assert 0x0 == next(find_fastbin_size(buf, max_size, 1))
+    assert 0x0 == next(find_fastbin_size(buf, max_size, 8))
+
+
+def test_unaligned():
+    max_size = 0x20
+    offsets = {
+        0x9: 0x20,
+    }
+    buf = setup_mem(max_size, offsets)
+    assert 0x1 == next(find_fastbin_size(buf, max_size, 1))
+    with pytest.raises(StopIteration):
+        next(find_fastbin_size(buf, max_size, 8))