From d8739d4295a68bb315a2b6d8d377c15b2090bad5 Mon Sep 17 00:00:00 2001 From: "Matt." Date: Mon, 31 Mar 2025 20:08:58 -0300 Subject: [PATCH] Use context management for temporary breakpoints (#2820) --- pwndbg/aglib/next.py | 26 ++++++++++---------------- pwndbg/aglib/shellcode.py | 32 ++++++++++++++++---------------- pwndbg/commands/start.py | 4 ++-- pwndbg/dbg/__init__.py | 13 +++++++++++++ 4 files changed, 41 insertions(+), 34 deletions(-) diff --git a/pwndbg/aglib/next.py b/pwndbg/aglib/next.py index 55a92593f..63e686318 100644 --- a/pwndbg/aglib/next.py +++ b/pwndbg/aglib/next.py @@ -106,20 +106,17 @@ async def break_next_branch(ec: pwndbg.dbg_mod.ExecutionController, address=None proc = pwndbg.dbg.selected_inferior() if ins: - bp = proc.break_at(BreakpointLocation(ins.address), internal=True) - await ec.cont(bp) - bp.remove() + with proc.break_at(BreakpointLocation(ins.address), internal=True) as bp: + await ec.cont(bp) return ins async def break_next_interrupt(ec: pwndbg.dbg_mod.ExecutionController, address=None): ins = next_int(address) - proc = pwndbg.dbg.selected_inferior() if ins: - bp = proc.break_at(BreakpointLocation(ins.address), internal=True) - await ec.cont(bp) - bp.remove() + with proc.break_at(BreakpointLocation(ins.address), internal=True) as bp: + await ec.cont(bp) return ins @@ -186,9 +183,8 @@ async def break_on_next_matching_instruction( # Only set breakpoints at a different PC location, otherwise we # will continue until we hit a breakpoint that's not related to # this opeeration, or the program halts. - bp = proc.break_at(BreakpointLocation(ins.address), internal=True) - await ec.cont(bp) - bp.remove() + with proc.break_at(BreakpointLocation(ins.address), internal=True) as bp: + await ec.cont(bp) return ins else: # We don't want to be spinning in place, nudge execution forward @@ -201,9 +197,8 @@ async def break_on_next_matching_instruction( if nb is not None: if nb.address != pwndbg.aglib.regs.pc: # Stop right at the next branch instruction. - bp = proc.break_at(BreakpointLocation(nb.address), internal=True) - await ec.cont(bp) - bp.remove() + with proc.break_at(BreakpointLocation(nb.address), internal=True) as bp: + await ec.cont(bp) else: # Nudge execution so we take the branch we're on top of. pass @@ -257,6 +252,5 @@ async def break_on_next(ec: pwndbg.dbg_mod.ExecutionController, address=None) -> ins = pwndbg.aglib.disasm.one(address) proc = pwndbg.dbg.selected_inferior() - bp = proc.break_at(BreakpointLocation(ins.address + ins.size), internal=True) - await ec.cont(bp) - bp.remove() + with proc.break_at(BreakpointLocation(ins.address + ins.size), internal=True) as bp: + await ec.cont(bp) diff --git a/pwndbg/aglib/shellcode.py b/pwndbg/aglib/shellcode.py index 3238f2099..6f7ca67a2 100644 --- a/pwndbg/aglib/shellcode.py +++ b/pwndbg/aglib/shellcode.py @@ -124,22 +124,22 @@ async def exec_shellcode( # Execute. target_address = starting_address + len(blob) - bp = pwndbg.dbg.selected_inferior().break_at(BreakpointLocation(target_address), internal=True) - while True: - try: - await ec.cont(bp) - break - except CancelledError: - if disable_breakpoints: - # We probably hit another breakpoint, but in this mode we're - # supposed to ignore any breakpoints that aren't the one we put - # at the end of the range, so just retry. - continue - - # We hit an external break, and we haven't been told to ignore it. - raise - - bp.remove() + with pwndbg.dbg.selected_inferior().break_at( + BreakpointLocation(target_address), internal=True + ) as bp: + while True: + try: + await ec.cont(bp) + break + except CancelledError: + if disable_breakpoints: + # We probably hit another breakpoint, but in this mode we're + # supposed to ignore any breakpoints that aren't the one we put + # at the end of the range, so just retry. + continue + + # We hit an external break, and we haven't been told to ignore it. + raise # Restore the state of the context skip. if pwndbg.dbg.is_gdblib_available(): diff --git a/pwndbg/commands/start.py b/pwndbg/commands/start.py index 236d0e6f5..891255fe1 100644 --- a/pwndbg/commands/start.py +++ b/pwndbg/commands/start.py @@ -41,8 +41,8 @@ def breakpoint_at_entry(): bp = proc.break_at(BreakpointLocation(addr), internal=True) async def ctrl(ec: pwndbg.dbg_mod.ExecutionController): - await ec.cont(bp) - bp.remove() + with bp: + await ec.cont(bp) proc.dispatch_execution_controller(ctrl) diff --git a/pwndbg/dbg/__init__.py b/pwndbg/dbg/__init__.py index c3754f53c..3523a0189 100644 --- a/pwndbg/dbg/__init__.py +++ b/pwndbg/dbg/__init__.py @@ -104,6 +104,10 @@ class Arch: class StopPoint: """ The handle to either an insalled breakpoint or watchpoint. + + May be used in a `with` statement, in which case the stop point is + automatically removed at the end of the statement. This allows for easy + implementation of temporary breakpoints. """ def remove(self) -> None: @@ -118,6 +122,15 @@ class StopPoint: """ raise NotImplementedError() + def __enter__(self) -> StopPoint: + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + """ + Automatic breakpoint removal. + """ + self.remove() + class BreakpointLocation: """