Use context management for temporary breakpoints (#2820)

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)

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():

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)
 

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:
     """