From f0ea79575df52fefda09ada833b4f598ad95563d Mon Sep 17 00:00:00 2001 From: Gulshan Singh Date: Thu, 13 Oct 2022 21:46:05 -0700 Subject: [PATCH] Add unit-test for find_fake_fast command --- .github/workflows/tests.yml | 5 +- pwndbg/color/backtrace.py | 2 +- pwndbg/color/chain.py | 2 +- pwndbg/color/context.py | 2 +- pwndbg/color/disasm.py | 2 +- pwndbg/color/enhance.py | 2 +- pwndbg/color/hexdump.py | 2 +- pwndbg/color/memory.py | 2 +- pwndbg/color/nearpc.py | 2 +- pwndbg/color/telescope.py | 2 +- pwndbg/commands/heap.py | 19 +---- pwndbg/lib/heap/__init__.py | 0 pwndbg/lib/heap/helpers.py | 25 +++++++ tests.sh | 7 +- tests/unit-tests/test_find_fastbin_size.py | 80 ++++++++++++++++++++++ 15 files changed, 125 insertions(+), 29 deletions(-) create mode 100644 pwndbg/lib/heap/__init__.py create mode 100644 pwndbg/lib/heap/helpers.py create mode 100644 tests/unit-tests/test_find_fastbin_size.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index bab24d95..6ccf003a 100644 --- a/.github/workflows/tests.yml +++ b/.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' diff --git a/pwndbg/color/backtrace.py b/pwndbg/color/backtrace.py index 151b97e3..2ee0fd70 100644 --- a/pwndbg/color/backtrace.py +++ b/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( diff --git a/pwndbg/color/chain.py b/pwndbg/color/chain.py index 7316a9be..0700916a 100644 --- a/pwndbg/color/chain.py +++ b/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)" diff --git a/pwndbg/color/context.py b/pwndbg/color/context.py index c282e6d2..5eb57a8b 100644 --- a/pwndbg/color/context.py +++ b/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)" diff --git a/pwndbg/color/disasm.py b/pwndbg/color/disasm.py index e49ec7af..c594ecbc 100644 --- a/pwndbg/color/disasm.py +++ b/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)) diff --git a/pwndbg/color/enhance.py b/pwndbg/color/enhance.py index 12f639a0..f254f2ad 100644 --- a/pwndbg/color/enhance.py +++ b/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)" diff --git a/pwndbg/color/hexdump.py b/pwndbg/color/hexdump.py index 2bc4c6d9..621329b8 100644 --- a/pwndbg/color/hexdump.py +++ b/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)" diff --git a/pwndbg/color/memory.py b/pwndbg/color/memory.py index 32937b4c..55adb5b7 100644 --- a/pwndbg/color/memory.py +++ b/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") diff --git a/pwndbg/color/nearpc.py b/pwndbg/color/nearpc.py index acfd02c7..f60549d5 100644 --- a/pwndbg/color/nearpc.py +++ b/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)" diff --git a/pwndbg/color/telescope.py b/pwndbg/color/telescope.py index 61c904bd..c87cab13 100644 --- a/pwndbg/color/telescope.py +++ b/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)" diff --git a/pwndbg/commands/heap.py b/pwndbg/commands/heap.py index d46273e7..b44285f0 100644 --- a/pwndbg/commands/heap.py +++ b/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( diff --git a/pwndbg/lib/heap/__init__.py b/pwndbg/lib/heap/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/pwndbg/lib/heap/helpers.py b/pwndbg/lib/heap/helpers.py new file mode 100644 index 00000000..b295ba48 --- /dev/null +++ b/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 diff --git a/tests.sh b/tests.sh index 9ed03941..e889035f 100755 --- a/tests.sh +++ b/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 diff --git a/tests/unit-tests/test_find_fastbin_size.py b/tests/unit-tests/test_find_fastbin_size.py new file mode 100644 index 00000000..53876429 --- /dev/null +++ b/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))