fix exec_shellcode deadlock (#3028)

* fix exec_shellcode deadlock

* typo

pwndbg/aglib/shellcode.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 import contextlib
 from asyncio import CancelledError
+from typing import Iterator
 
 import pwnlib.asm
 import pwnlib.shellcraft
@@ -22,9 +23,6 @@ import pwndbg.aglib.vmmap
 from pwndbg.dbg import BreakpointLocation
 from pwndbg.dbg import ExecutionController
 
-if pwndbg.dbg.is_gdblib_available():
-    import pwndbg.gdblib.prompt
-
 
 def _get_syscall_return_value():
     """
@@ -46,7 +44,6 @@ async def exec_syscall(
     arg3=None,
     arg4=None,
     arg5=None,
-    disable_breakpoints=False,
 ):
     """
     Tries executing the given syscall in the context of the inferior.
@@ -60,41 +57,12 @@ async def exec_syscall(
     async with exec_shellcode(
         ec,
         syscall_bin,
-        restore_context=True,
-        disable_breakpoints=disable_breakpoints,
     ):
         return _get_syscall_return_value()
 
 
-@contextlib.asynccontextmanager
-async def exec_shellcode(
-    ec: ExecutionController, blob, restore_context=True, disable_breakpoints=False
-):
-    """
-    Tries executing the given blob of machine code in the current context of the
-    inferior, optionally restoring the values of the registers as they were
-    before the shellcode ran, as a means to allow for execution of the inferior
-    to continue uninterrupted. The value of the program counter is always
-    restored.
-
-    Additionally, the caller may specify an object to be called before the
-    context is restored, so that information stored in the registers after the
-    shellcode finishes can be retrieved. The return value of that call will be
-    returned by this function.
-
-    # Safety
-    Seeing as this function injects code directly into the inferior and runs it,
-    the caller must be careful to inject code that will (1) terminate and (2)
-    not cause the inferior to misbehave. Otherwise, it is fairly easy to crash
-    or currupt the memory in the inferior.
-    """
-
-    register_set = pwndbg.lib.regs.reg_sets[pwndbg.aglib.arch.name]
-    preserve_set = register_set.gpr + register_set.args + (register_set.pc, register_set.stack)
-
-    registers = {reg: pwndbg.aglib.regs[reg] for reg in preserve_set}
-    starting_address = registers[register_set.pc]
-
+@contextlib.contextmanager
+def _ctx_code(starting_address: int, blob) -> Iterator[None]:
     # Make sure the blob fits in the rest of the space we have in this page.
     #
     # NOTE: Technically, we could actually use anything from the whole page to
@@ -116,47 +84,76 @@ async def exec_shellcode(
     existing_code = pwndbg.aglib.memory.read(starting_address, len(blob))
     pwndbg.aglib.memory.write(starting_address, blob)
 
-    # The continue we use here will trigger an event that would get the context
-    # prompt to show, regardless of the circumstances. We don't want that, so
-    # we preserve the state of the context skip.
-    if pwndbg.dbg.is_gdblib_available():
-        would_skip_context = pwndbg.gdblib.prompt.context_shown
+    try:
+        yield
+    finally:
+        pwndbg.aglib.memory.write(starting_address, existing_code)
+
+
+@contextlib.contextmanager
+def _ctx_registers() -> Iterator[int]:
+    register_set = pwndbg.lib.regs.reg_sets[pwndbg.aglib.arch.name]
+    preserve_set = register_set.gpr + register_set.args + (register_set.pc, register_set.stack)
+
+    uncached_regs = pwndbg.dbg.selected_frame().regs()
+    registers = {reg: int(uncached_regs.by_name(reg)) for reg in preserve_set}
+    starting_address = registers[register_set.pc]
 
-    # Execute.
-    target_address = starting_address + len(blob)
+    try:
+        yield starting_address
+    finally:
+        # Restore the code and the program counter and, if requested, the rest of
+        # the registers.
+        setattr(pwndbg.aglib.regs, register_set.pc, starting_address)
+        for reg, val in registers.items():
+            setattr(pwndbg.aglib.regs, reg, val)
+
+
+async def _execute_until_addr(ec: ExecutionController, target_address: int) -> None:
     with pwndbg.dbg.selected_inferior().break_at(
         BreakpointLocation(target_address), internal=True
     ) as bp:
         while True:
             try:
-                await ec.cont(bp)
+                await ec.cont_selected_thread(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 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
+    assert pwndbg.dbg.selected_frame().pc() == target_address, "Target address is incorrect"
 
-    # Restore the state of the context skip.
-    if pwndbg.dbg.is_gdblib_available():
-        pwndbg.gdblib.prompt.context_shown = would_skip_context
 
-    # Make sure we're in the right place.
-    assert pwndbg.aglib.regs.pc == target_address
+@contextlib.asynccontextmanager
+async def exec_shellcode(ec: ExecutionController, blob):
+    """
+    Tries executing the given blob of machine code in the current context of the
+    inferior, optionally restoring the values of the registers as they were
+    before the shellcode ran, as a means to allow for execution of the inferior
+    to continue uninterrupted. The value of the program counter is always
+    restored.
 
-    try:
-        # Give the caller a chance to collect information from the environment
-        # before any of the context gets restored.
-        yield
-    finally:
-        # Restore the code and the program counter and, if requested, the rest of
-        # the registers.
-        pwndbg.aglib.memory.write(starting_address, existing_code)
-        setattr(pwndbg.aglib.regs, register_set.pc, starting_address)
-        if restore_context:
-            for reg, val in registers.items():
-                setattr(pwndbg.aglib.regs, reg, val)
+    Additionally, the caller may specify an object to be called before the
+    context is restored, so that information stored in the registers after the
+    shellcode finishes can be retrieved. The return value of that call will be
+    returned by this function.
+
+    # Safety
+    Seeing as this function injects code directly into the inferior and runs it,
+    the caller must be careful to inject code that will (1) terminate and (2)
+    not cause the inferior to misbehave. Otherwise, it is fairly easy to crash
+    or currupt the memory in the inferior.
+    """
+
+    with _ctx_registers() as starting_address:
+        target_address = starting_address + len(blob)
+        with _ctx_code(starting_address, blob):
+            try:
+                with pwndbg.dbg.ctx_suspend_events(pwndbg.dbg_mod.EventType.SUSPEND_ALL):
+                    await _execute_until_addr(ec, target_address)
+
+                yield
+            finally:
+                pass

pwndbg/commands/hijack_fd.py

@@ -129,9 +129,7 @@ async def exec_shellcode_with_stack(ec: pwndbg.dbg_mod.ExecutionController, blob
     original_stack = pwndbg.aglib.memory.read(stack_start, stack_size)
 
     try:
-        async with pwndbg.aglib.shellcode.exec_shellcode(
-            ec, blob, restore_context=True, disable_breakpoints=True
-        ):
+        async with pwndbg.aglib.shellcode.exec_shellcode(ec, blob):
             stack_diff_size = stack_start_diff - pwndbg.aglib.regs.sp
 
             # Make sure stack is not corrupted somehow

pwndbg/dbg/__init__.py

@@ -309,12 +309,28 @@ class ExecutionController:
         Throws `CancelledError` if a breakpoint or watchpoint is hit, the program
         exits, or if any other unexpected event that diverts execution happens
         while fulfulling the step.
+
+        FIXME GDB:
+        On GDB `stepi` will execute other threads. On LLDB not.
+        Please use `set scheduler-locking step`
         """
         raise NotImplementedError()
 
     def cont(self, until: StopPoint) -> Awaitable[None]:
         """
         Continues execution until the given breakpoint or whatchpoint is hit.
+        Continues execution on all threads.
+
+        Throws `CancelledError` if a breakpoint or watchpoint is hit that is not
+        the one given in `until`, the program exits, or if any other unexpected
+        event happens.
+        """
+        raise NotImplementedError()
+
+    def cont_selected_thread(self, until: StopPoint) -> Awaitable[None]:
+        """
+        Continues execution on single thread until the given breakpoint or whatchpoint is hit.
+        Continues execution on selected thread.
 
         Throws `CancelledError` if a breakpoint or watchpoint is hit that is not
         the one given in `until`, the program exits, or if any other unexpected

pwndbg/dbg/gdb/__init__.py

@@ -985,6 +985,7 @@ class GDBProcess(pwndbg.dbg_mod.Process):
 class GDBExecutionController(pwndbg.dbg_mod.ExecutionController):
     @override
     async def single_step(self):
+        # TODO: disable GDB ugly output
         gdb.execute("si")
 
         # Check if the program stopped because of the step we just took. If it
@@ -995,6 +996,7 @@ class GDBExecutionController(pwndbg.dbg_mod.ExecutionController):
 
     @override
     async def cont(self, until: pwndbg.dbg_mod.StopPoint):
+        # TODO: disable GDB ugly output
         gdb.execute("continue")
 
         # Check if the program stopped because of the breakpoint we were given,
@@ -1006,6 +1008,23 @@ class GDBExecutionController(pwndbg.dbg_mod.ExecutionController):
         ):
             raise CancelledError()
 
+    @override
+    async def cont_selected_thread(self, until: pwndbg.dbg_mod.StopPoint):
+        from pwndbg.gdblib.scheduler import lock_scheduler
+
+        with lock_scheduler():
+            # TODO: disable GDB ugly output
+            gdb.execute("continue")
+
+        # Check if the program stopped because of the breakpoint we were given,
+        # and, just like for the single step, propagate a cancellation error if
+        # it stopped for any other reason.
+        assert isinstance(until, GDBStopPoint)
+        if f"It stopped at breakpoint {until.inner.number}" not in gdb.execute(
+            "info program", to_string=True
+        ):
+            raise CancelledError()
+
 
 # Like in LLDB, we only need a single instance of the execution controller.
 EXECUTION_CONTROLLER = GDBExecutionController()

pwndbg/dbg/lldb/__init__.py

@@ -702,9 +702,11 @@ class YieldContinue:
     """
 
     target: LLDBStopPoint
+    selected_thread: bool
 
-    def __init__(self, target: LLDBStopPoint):
+    def __init__(self, target: LLDBStopPoint, selected_thread: bool = False):
         self.target = target
+        self.selected_thread = selected_thread
 
 
 class YieldSingleStep:
@@ -729,6 +731,11 @@ class LLDBExecutionController(pwndbg.dbg_mod.ExecutionController):
         assert isinstance(target, LLDBStopPoint)
         return OneShotAwaitable(YieldContinue(target))
 
+    @override
+    def cont_selected_thread(self, target: pwndbg.dbg_mod.StopPoint) -> Awaitable[None]:
+        assert isinstance(target, LLDBStopPoint)
+        return OneShotAwaitable(YieldContinue(target, selected_thread=True))
+
 
 # Our execution controller doesn't need to change between uses, as all the state
 # associated with it resides further up, in the Pwndbg CLI, so we can just share

pwndbg/dbg/lldb/repl/proc.py

@@ -363,9 +363,25 @@ class ProcessDriver:
                     continue
 
             elif isinstance(step, YieldContinue):
+                threads_suspended = []
+                if step.selected_thread:
+                    thread = self.process.GetSelectedThread()
+                    for t in self.process.threads:
+                        if t.id == thread.id:
+                            continue
+                        if not t.is_suspended:
+                            t.Suspend()
+                            threads_suspended.append(t)
+
                 # Continue the process and wait for the next stop-like event.
                 self.process.Continue()
                 healthy, event = self._run_until_next_stop()
+
+                if threads_suspended:
+                    for t in threads_suspended:
+                        if t.IsValid():
+                            t.Resume()
+
                 if not healthy:
                     # The process exited, Cancel the execution controller.
                     exception = CancelledError()

pwndbg/gdblib/scheduler.py

@@ -22,8 +22,10 @@ def lock_scheduler() -> Iterator[None]:
     old_config = gdb.parameter("scheduler-locking")
     if old_config != "on":
         gdb.execute("set scheduler-locking on")
-        yield
-        gdb.execute(f"set scheduler-locking {old_config}")
+        try:
+            yield
+        finally:
+            gdb.execute(f"set scheduler-locking {old_config}")
     else:
         yield