Add support for performance profiling (#1413)

Co-authored-by: Gulshan Singh <gsgx@google.com>

gdbinit.py

@@ -1,9 +1,18 @@
+import cProfile
 import glob
 import locale
 import sys
+import time
 from os import environ
 from os import path
 
+_profiler = cProfile.Profile()
+
+_start_time = None
+if environ.get("PWNDBG_PROFILE") == "1":
+    _start_time = time.time()
+    _profiler.enable()
+
 # Allow users to use packages from a virtualenv
 # That's not 100% supported, but they do it on their own,
 # so we will warn them if the GDB's Python is not virtualenv's Python
@@ -73,3 +82,9 @@ if encoding != "UTF-8":
 environ["PWNLIB_NOTERM"] = "1"
 
 import pwndbg  # noqa: F401
+import pwndbg.profiling
+
+pwndbg.profiling.init(_profiler, _start_time)
+if environ.get("PWNDBG_PROFILE") == "1":
+    pwndbg.profiling.profiler.stop("pwndbg-load.pstats")
+    pwndbg.profiling.profiler.start()

profiling/print_stats.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python3
+
+import pstats
+from argparse import ArgumentParser
+
+
+def parse_args():
+    parser = ArgumentParser(description="Print the profiling stats from a pstats file")
+    parser.add_argument("-n", "--num", type=int, help="number of stats to print")
+    parser.add_argument("-f", "--filter", help="regex to match against each entry")
+    parser.add_argument(
+        "--no-strip", action="store_true", help="print the entire file path for each entry"
+    )
+    parser.add_argument("file", help="pstats file to parse")
+    parser.add_argument(
+        "-s",
+        "--sort",
+        choices=["calls", "ncalls", "cumulative", "time"],
+        default="cumulative",
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    s = pstats.Stats(args.file)
+    if not args.no_strip:
+        s.strip_dirs()
+    s.sort_stats(args.sort).print_stats(args.filter, args.num)
+
+
+if __name__ == "__main__":
+    main()

pwndbg/gdblib/prompt.py

@@ -1,4 +1,5 @@
 import re
+from os import environ
 
 import gdb
 
@@ -6,6 +7,7 @@ import pwndbg.decorators
 import pwndbg.gdblib.events
 import pwndbg.gdblib.functions
 import pwndbg.lib.memoize
+import pwndbg.profiling
 from pwndbg.color import disable_colors
 from pwndbg.color import message
 from pwndbg.lib.tips import get_tip_of_the_day
@@ -45,6 +47,10 @@ def initial_hook(*a):
     pwndbg.decorators.first_prompt = True
 
     prompt_hook(*a)
+
+    if environ.get("PWNDBG_PROFILE") == "1":
+        pwndbg.profiling.profiler.stop("pwndbg-first-prompt.pstats")
+
     gdb.prompt_hook = prompt_hook
 
 

pwndbg/profiling.py

@@ -0,0 +1,35 @@
+import cProfile
+import time
+from typing import Optional
+
+profiler = None
+
+
+def init(p: cProfile.Profile, _start_time: Optional[float]):
+    global profiler
+    profiler = Profiler(p)
+    profiler._start_time = _start_time
+
+
+class Profiler:
+    def __init__(self, p: cProfile.Profile):
+        self._profiler = p
+        self._start_time = None
+
+    def print_time_elapsed(self):
+        assert self._start_time is not None
+        return print("Time Elapsed:", time.time() - self._start_time)
+
+    def start(self):
+        self._start_time = time.time()
+        self._profiler.enable()
+
+    def stop(self, filename=None):
+        if not filename:
+            filename = f"pwndbg-{int(time.time())}.pstats"
+
+        self.print_time_elapsed()
+        self._profiler.disable()
+        self._start_time = None
+
+        self._profiler.dump_stats(filename)