Trace function calls via lldb

lldb-trace.py

#!/usr/bin/python3

import json
import os
import re
import shlex
import subprocess
import sys
import time


write = sys.stdout.write
verbose = os.getenv("V")
pattern_match = lambda _x: True


def attach_pid(pid, pattern):
    import lldb

    debugger = lldb.debugger
    target = debugger.CreateTarget("")
    process = target.AttachToProcessWithID(lldb.SBListener(), pid, lldb.SBError())
    bp = set_breakpoints(target, pattern)
    return target, process, bp


def new_process(args, pattern):
    import lldb

    debugger = lldb.debugger
    target = debugger.CreateTarget(args[0])
    bp = set_breakpoints(target, pattern)
    launch_info = target.GetLaunchInfo()
    launch_info.SetArguments(args[1:], True)
    error = lldb.SBError()
    process = target.Launch(launch_info, error)
    return target, process, bp


def trace2(debugger, command, exe_ctx, result, internal_dict):
    # Update 'write' to write to $OUT.
    global write

    out = os.getenv("OUT")
    if out:
        f = open(out, "a", newline="\n")
        write = f.write

    pattern = os.getenv("PAT") or ".*"

    # command only contains args
    args = shlex.split(command)
    if len(args) >= 1 and args[0]:
        try:
            pid = int(args[0])
            target, process, bp = attach_pid(pid, pattern)
        except ValueError:
            # not a pid, treat as args
            target, process, bp = new_process(args, pattern)
    else:
        # used standalone in lldb
        target = exe_ctx.target
        process = exe_ctx.process
        bp = set_breakpoints(target, pattern)

    debug_write(f"BREAKPOINT SET ({bp.num_locations})\n")
    # count = 0
    # for loc in bp.locations:
    #     count += 1
    #     if count > 100:
    #         break
    #     print(loc)
    wait_for_completion(process)


def wait_for_completion(process):
    import lldb

    state = process.GetState()
    if state == lldb.eStateStopped:
        process.Continue()

    # Wait for the process to exit
    try:
        while True:
            state = process.GetState()
            if state == lldb.eStateExited:
                write(f"Process exited with status {process.GetExitStatus()}\n")
                break
            elif state == lldb.eStateStopped:
                # The process might hit other breakpoints; handle them if needed.
                debug_write(f"Process stopped. Breaking...\n")
                break

            # Sleep for a bit to avoid busy-waiting
            time.sleep(0.1)
    except KeyboardInterrupt:
        lldb.debugger.Terminate()


def set_breakpoints(target, pattern):
    import lldb

    global pattern_match

    # find the main module name
    try:
        module_name = str(target.modules[0].file)
    except IndexError:
        raise RuntimeError("no main module found")
    debug_write(f"MAIN MODULE: {module_name}\nBREAKPOINT PATTERN: {pattern}\n")
    bp = target.BreakpointCreateByRegex(pattern, module_name)
    assert bp.IsValid()
    pattern_match = re.compile(pattern).match
    bp.SetScriptCallbackFunction(f"{__name__}.breakpoint_callback")
    return bp


def breakpoint_callback(frame, bp_loc, extra_args, internal_dict):
    name = str(frame.name)
    if not pattern_match(name):
        debug_write(f"  IGNORE: {name}\n")
        # Disable noisy locations
        bp_loc.SetEnabled(False)
        return False
    thread = frame.thread
    # Figure out the stack depth
    frame_depth = len(thread.frames)
    prefix = " " * frame_depth
    # Print args. Render "empty" properly.
    args = frame.args
    if args.GetSize() == 0:
        args = ""
    else:
        args = str(args).replace("\n", " ")
        if len(args) > 120:
            args = args[:120] + "..."
    write(f"{thread.name}>{prefix}{frame.name}({args})\n")
    # Continue the process
    return False


def debug_write(message):
    if verbose:
        write(message)


def __lldb_init_module(debugger, internal_dict):
    debugger.HandleCommand(f"command script add -f {__name__}.trace2 trace2")


def main():
    import os, tempfile

    python_source = open(__file__, "rb").read()

    with tempfile.TemporaryDirectory(prefix="lldbscript") as dir:
        python_script_path = os.path.join(dir, "trace2.py")
        with open(python_script_path, "wb") as f:
            f.write(python_source)

        args = [
            os.getenv("LLDB") or "lldb",
            #"-b",
            #"--source-quietly",
            "-o",
            f"command script import {python_script_path}",
            "-o",
            f"trace2 {' '.join(map(shlex.quote, sys.argv[1:]))}",
        ]
        subprocess.run(args)


if __name__ == "__main__":
    main()