You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
pwndbg/profiling/benchmark.py

160 lines
3.5 KiB
Python

from typing import Callable
import gdb
import argparse
import pstats
import time
import cProfile
import pwndbg
import pwndbg.commands
import pwndbg.commands.context
import pwndbg.lib.cache
import pwndbg.aglib.regs
import pwndbg.aglib.vmmap
import pwndbg.commands.telescope
COUNT = 100
def run_benchmark(name: str, prefix: str, callback: Callable, count=COUNT) -> float:
"""
Return:
Average time to execute callback in seconds
"""
profiler = cProfile.Profile()
profiler.enable()
for _ in range(count):
callback()
profiler.disable()
pstats_file_name_base = f"{prefix}-{count}-{name}-{int(time.time())}"
profiler.dump_stats(f"{pstats_file_name_base}.pstats")
full_time = pstats.Stats(profiler).total_tt
print(f"Time elapsed: {full_time}. Average time: {full_time / count}")
return full_time / count
parser = argparse.ArgumentParser(
description="""
Benchmark contexts.
Uses cProfile to instrument the code.
Experimentally, this can cause the code to run ~2.5 times slower than it's real speed.
However, you can use it to find relative performance before/after changes, and find functions that take longer than they should.
""",
)
parser.add_argument("name", type=str, help="Name placed into output pstats filename")
@pwndbg.commands.Command(parser, category=pwndbg.commands.CommandCategory.DEV)
def benchmark_context(name: str):
# Context - clear cache
## Benchmark the `context` command, no caching
def run_with_clear():
pwndbg.commands.context.context.function()
pwndbg.lib.cache.clear_caches()
run_benchmark(
name,
"context-without-cache",
run_with_clear
)
pwndbg.lib.cache.clear_caches()
# Context - with cache
## Benchmark the `context` command with cache
run_benchmark(
name,
"context-with-cache",
# Bypass the decorator so cProfile/snakeviz sees things correctly
lambda: pwndbg.commands.context.context.function()
)
pwndbg.lib.cache.clear_caches()
# Step
def run_with_step():
gdb.execute("stepi")
# Bypass the decorator so cProfile can see function stack correctly
pwndbg.commands.context.context.function()
run_benchmark(
name,
"step",
run_with_step
)
# Regs
def regs():
gdb.execute("stepi")
pwndbg.commands.context.context_regs()
run_benchmark(
name,
"regs",
regs
)
# Disasm
def disasm():
gdb.execute("stepi")
pwndbg.commands.context.context_disasm()
run_benchmark(
name,
"disasm",
disasm
)
# Stack
def stack():
gdb.execute("stepi")
pwndbg.commands.context.context_stack()
run_benchmark(
name,
"stack",
stack
)
parser = argparse.ArgumentParser(
description="""
Benchmark telescoping the entire stack
""",
)
parser.add_argument("name", type=str, help="Name placed into output pstats filename")
@pwndbg.commands.Command(parser, category=pwndbg.commands.CommandCategory.DEV)
def benchmark_large_telescope(name: str):
# Telescope entire stack
stack_page = pwndbg.aglib.vmmap.find(pwndbg.aglib.regs.read_reg(pwndbg.aglib.regs.stack))
start = stack_page.start
len = stack_page.memsz
def print_all_stack():
pwndbg.commands.telescope.telescope(start, len // pwndbg.aglib.arch.ptrsize)
run_benchmark(
name,
"all-stack",
print_all_stack,
4
)