lyc
/
pwndbg
mirror of https://github.com/pwndbg/pwndbg.git
1
0
Fork 0
Code Releases Activity

Fix handling of O_NONBLOCK in IODriverPlainText (#3443)

Browse Source
pull/3460/head
Matt. 1 week ago committed by GitHub
parent 0daaf1889a
commit a3721aecaf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 569 additions and 248 deletions
Show Stats Download Patch File Download Diff File

46
pwndbg/dbg/lldb/__init__.py
Unescape View File

@ -2,6 +2,7 @@ from __future__ import annotations
import bisect import bisect
import collections import collections
import enum
import os import os
import random import random
import re import re
@ -1699,7 +1700,11 @@ class LLDBProcess(pwndbg.dbg_mod.Process):
_struct: lldb.SBStructuredData, _struct: lldb.SBStructuredData,
_internal, _internal,
) -> bool: ) -> bool:
return stop_handler(sp) try:
self.dbg.lldb_python_state_callback(LLDBPythonState.LLDB_STOP_HANDLER)
return stop_handler(sp)
finally:
self.dbg.lldb_python_state_callback(LLDBPythonState.PWNDBG)
sys.modules[self.dbg.module].__dict__[stop_handler_name] = handler sys.modules[self.dbg.module].__dict__[stop_handler_name] = handler
@ -1842,6 +1847,35 @@ class LLDBCommand(pwndbg.dbg_mod.CommandHandle):
self.command_name = command_name self.command_name = command_name
class LLDBPythonState(enum.Enum):
"""
State of LLDB Python execution.
Unlike in pwndbg-gdb, in pwndbg-lldb the responsibility of driving execution
of Python code forward is shared between Pwndbg and LLDB. Knowing which one
is in charge is crucial to the correct functioning of the Pwndbg REPL.
This class defines the different kinds of states we can be in.
"""
PWNDBG = 1
"Pwndbg is driving execution of Python code"
LLDB_COMMAND_HANDLER = 0
"Python code is executing from inside an LLDB command handler"
LLDB_STOP_HANDLER = 2
"Python code is executing from an LLDB breakpoint/watchpoint hook handler"
def _default_lldb_python_state_callback(_state: LLDBPythonState) -> None:
"""
Before being set by the Pwndbg REPL, LLDB instances need a sensible default
value for lldb_python_state_callback. This is it.
"""
raise RuntimeError("Pwndbg REPL failed to set lldb_python_state_callback")
class LLDB(pwndbg.dbg_mod.Debugger): class LLDB(pwndbg.dbg_mod.Debugger):
exec_states: List[lldb.SBExecutionState] exec_states: List[lldb.SBExecutionState]
@ -1867,8 +1901,8 @@ class LLDB(pwndbg.dbg_mod.Debugger):
# Relay used for exceptions originating from commands called through LLDB. # Relay used for exceptions originating from commands called through LLDB.
_exception_relay: BaseException | None _exception_relay: BaseException | None
in_lldb_command_handler: bool lldb_python_state_callback: Callable[[LLDBPythonState], None]
"Whether an LLDB command handler is currently running" "Callback to the REPL, used to notify it of LLDB driving Python code"
# temporarily suspend context output # temporarily suspend context output
should_suspend_ctx: bool should_suspend_ctx: bool
@ -1884,7 +1918,7 @@ class LLDB(pwndbg.dbg_mod.Debugger):
self.controllers = [] self.controllers = []
self._current_process_is_gdb_remote = False self._current_process_is_gdb_remote = False
self._exception_relay = None self._exception_relay = None
self.in_lldb_command_handler = False self.lldb_python_state_callback = _default_lldb_python_state_callback
self.should_suspend_ctx = False self.should_suspend_ctx = False
import pwndbg import pwndbg
@ -1959,12 +1993,12 @@ class LLDB(pwndbg.dbg_mod.Debugger):
def __call__(self, _, command, exe_context, result): def __call__(self, _, command, exe_context, result):
try: try:
debugger.exec_states.append(exe_context) debugger.exec_states.append(exe_context)
debugger.in_lldb_command_handler = True debugger.lldb_python_state_callback(LLDBPythonState.LLDB_COMMAND_HANDLER)
handler(debugger, command, True) handler(debugger, command, True)
except BaseException as e: except BaseException as e:
debugger._exception_relay = e debugger._exception_relay = e
finally: finally:
debugger.in_lldb_command_handler = False debugger.lldb_python_state_callback(LLDBPythonState.PWNDBG)
assert ( assert (
debugger.exec_states.pop() == exe_context debugger.exec_states.pop() == exe_context
), "Execution state mismatch on command handler" ), "Execution state mismatch on command handler"

186
pwndbg/dbg/lldb/repl/__init__.py
Unescape View File

@ -66,6 +66,7 @@ from pwndbg.color import message
from pwndbg.dbg import EventType from pwndbg.dbg import EventType
from pwndbg.dbg.lldb import LLDB from pwndbg.dbg.lldb import LLDB
from pwndbg.dbg.lldb import LLDBProcess from pwndbg.dbg.lldb import LLDBProcess
from pwndbg.dbg.lldb import LLDBPythonState
from pwndbg.dbg.lldb import OneShotAwaitable from pwndbg.dbg.lldb import OneShotAwaitable
from pwndbg.dbg.lldb.pset import InvalidParse from pwndbg.dbg.lldb.pset import InvalidParse
from pwndbg.dbg.lldb.pset import pget from pwndbg.dbg.lldb.pset import pget
@ -332,99 +333,116 @@ def run(
# This is the driver we're going to be using to handle the process. # This is the driver we're going to be using to handle the process.
relay = EventRelay(dbg) relay = EventRelay(dbg)
driver = ProcessDriver(debug=debug, event_handler=relay) with ProcessDriver(debug=debug, event_handler=relay) as driver:
# Set ourselves up to respond to changes in Python execution context.
# Set ourselves up to respond to SIGINT by interrupting the process if it is curr_state = LLDBPythonState.PWNDBG
# running, and doing nothing otherwise. in_lldb_command_handler = False
def handle_sigint(_sig, _frame):
driver.interrupt(in_lldb_command_handler=dbg.in_lldb_command_handler) def handle_lldb_python_state_change(new_state: LLDBPythonState) -> None:
nonlocal curr_state
signal.signal(signal.SIGINT, handle_sigint) nonlocal in_lldb_command_handler
match (curr_state, new_state):
show_greeting() case (LLDBPythonState.PWNDBG, LLDBPythonState.LLDB_STOP_HANDLER):
last_command = "" driver.pause_io_if_running()
case (LLDBPythonState.PWNDBG, LLDBPythonState.LLDB_COMMAND_HANDLER):
coroutine = controller(PwndbgController(), *args) driver.pause_io_if_running()
last_result: Any = None case (_, LLDBPythonState.PWNDBG):
last_exc: Exception | None = None driver.resume_io_if_running()
curr_state = new_state
while True:
try: dbg.lldb_python_state_callback = handle_lldb_python_state_change
# Execute the prompt hook.
dbg._fire_prompt_hook() # Set ourselves up to respond to SIGINT by interrupting the process if it is
# running, and doing nothing otherwise.
def handle_sigint(_sig, _frame):
driver.interrupt(in_lldb_command_handler=in_lldb_command_handler)
signal.signal(signal.SIGINT, handle_sigint)
show_greeting()
last_command = ""
coroutine = controller(PwndbgController(), *args)
last_result: Any = None
last_exc: Exception | None = None
while True:
try: try:
if last_exc is not None: # Execute the prompt hook.
action = coroutine.throw(last_exc) dbg._fire_prompt_hook()
else:
action = coroutine.send(last_result)
except StopIteration:
# Nothing else for us to do.
break
except asyncio.CancelledError:
# We requested a cancellation that wasn't overwritten.
break
finally:
last_exc = None
last_result = None
if isinstance(action, YieldInteractive):
if debug:
print("[-] REPL: Prompt next command from user interactively")
try: try:
if HAS_FZF: if last_exc is not None:
try: action = coroutine.throw(last_exc)
line = session.prompt(message=PROMPT)
except KeyboardInterrupt:
continue
else: else:
line = input(PROMPT) action = coroutine.send(last_result)
# If the input is empty (i.e., 'Enter'), use the previous command except StopIteration:
if line: # Nothing else for us to do.
last_command = line break
else: except asyncio.CancelledError:
line = last_command # We requested a cancellation that wasn't overwritten.
except EOFError: break
# Exit the REPL if there's nothing else to run. finally:
last_exc = asyncio.CancelledError() last_exc = None
continue last_result = None
if not exec_repl_command(line, sys.stdout, dbg, driver, relay): if isinstance(action, YieldInteractive):
last_exc = asyncio.CancelledError() if debug:
continue print("[-] REPL: Prompt next command from user interactively")
elif isinstance(action, YieldExecDirect): try:
if debug: if HAS_FZF:
print( try:
f"[-] REPL: Executing command '{action._command}' {'with' if action._capture else 'without'} output capture" line = session.prompt(message=PROMPT)
) except KeyboardInterrupt:
continue
else:
line = input(PROMPT)
# If the input is empty (i.e., 'Enter'), use the previous command
if line:
last_command = line
else:
line = last_command
except EOFError:
# Exit the REPL if there's nothing else to run.
last_exc = asyncio.CancelledError()
continue
if not exec_repl_command(line, sys.stdout, dbg, driver, relay):
last_exc = asyncio.CancelledError()
continue
elif isinstance(action, YieldExecDirect):
if debug:
print(
f"[-] REPL: Executing command '{action._command}' {'with' if action._capture else 'without'} output capture"
)
last_command = action._command last_command = action._command
if not action._prompt_silent: if not action._prompt_silent:
print(f"{PROMPT.value if HAS_FZF else PROMPT}{action._command}") print(f"{PROMPT.value if HAS_FZF else PROMPT}{action._command}")
try: try:
if action._capture: if action._capture:
with TextIOWrapper(BytesIO(), write_through=True) as output: with TextIOWrapper(BytesIO(), write_through=True) as output:
should_continue = exec_repl_command(
action._command, output, dbg, driver, relay
)
last_result = output.buffer.getvalue()
else:
should_continue = exec_repl_command( should_continue = exec_repl_command(
action._command, output, dbg, driver, relay action._command, sys.stdout, dbg, driver, relay
) )
last_result = output.buffer.getvalue() except BaseException as e:
else: last_exc = e
should_continue = exec_repl_command( continue
action._command, sys.stdout, dbg, driver, relay
) if not should_continue:
except BaseException as e: last_exc = asyncio.CancelledError()
last_exc = e continue
continue except UserCancelledError as e:
last_exc = e
if not should_continue:
last_exc = asyncio.CancelledError()
continue
except UserCancelledError as e:
last_exc = e
def exec_repl_command( def exec_repl_command(

266
pwndbg/dbg/lldb/repl/io.py
Unescape View File

@ -172,6 +172,12 @@ class IODriver:
""" """
raise NotImplementedError() raise NotImplementedError()
def close(self) -> None:
"""
Terminate this driver and release all resources associated with it.
"""
raise NotImplementedError()
def get_io_driver() -> IODriver: def get_io_driver() -> IODriver:
""" """
@ -197,6 +203,12 @@ class IODriverPlainText(IODriver):
in_thr: threading.Thread in_thr: threading.Thread
out_thr: threading.Thread out_thr: threading.Thread
stop_requested: threading.Event stop_requested: threading.Event
stop_fulfilled: threading.Semaphore
start_requested: threading.Semaphore
_closed: threading.Event
_running: bool
_stdout_nonblock_failed: bool
_stderr_nonblock_failed: bool
process: lldb.SBProcess process: lldb.SBProcess
@ -204,59 +216,104 @@ class IODriverPlainText(IODriver):
self.likely_output = threading.BoundedSemaphore(1) self.likely_output = threading.BoundedSemaphore(1)
self.process = None self.process = None
self.stop_requested = threading.Event() self.stop_requested = threading.Event()
self.start_requested = threading.BoundedSemaphore(2)
self.stop_fulfilled = threading.BoundedSemaphore(2)
self._closed = threading.Event()
self._running = False
self._stdout_nonblock_failed = False
self._stderr_nonblock_failed = False
assert self.start_requested.acquire()
assert self.start_requested.acquire()
assert self.stop_fulfilled.acquire()
assert self.stop_fulfilled.acquire()
self.in_thr = threading.Thread(target=self._handle_input)
self.out_thr = threading.Thread(target=self._handle_output)
self.in_thr.start()
self.out_thr.start()
@override @override
def stdio(self) -> Tuple[str | None, str | None, str | None]: def stdio(self) -> Tuple[str | None, str | None, str | None]:
return None, None, None return None, None, None
def _handle_input(self): def _handle_input(self):
while not self.stop_requested.is_set(): while not self._closed.is_set():
if SELECT_AVAILABLE: if not self.start_requested.acquire(blocking=True, timeout=1):
select.select([sys.stdin], [], [], 0.2) continue
try: while not self.stop_requested.is_set():
data = sys.stdin.read() if SELECT_AVAILABLE:
self.process.PutSTDIN(data) select.select([sys.stdin], [], [], 0.2)
except (BlockingIOError, TypeError):
# We have to check for TypeError here too, as, even though you try:
# *can* set stdin into nonblocking mode, it doesn't handle it data = sys.stdin.read()
# very gracefully. self.process.PutSTDIN(data)
# except (BlockingIOError, TypeError):
# See https://github.com/python/cpython/issues/57531 # We have to check for TypeError here too, as, even though you
# *can* set stdin into nonblocking mode, it doesn't handle it
# Ignore blocking errors, but wait for a little bit before # very gracefully.
# trying again if we don't have select(). #
if not SELECT_AVAILABLE: # See https://github.com/python/cpython/issues/57531
time.sleep(0.1)
# Ignore blocking errors, but wait for a little bit before
# trying again if we don't have select().
if not SELECT_AVAILABLE:
time.sleep(0.1)
self.stop_fulfilled.release()
def _handle_output(self): def _handle_output(self):
while not self.stop_requested.is_set(): while not self._closed.is_set():
# Try to acquire the semaphore. This will not succeed until the next if not self.start_requested.acquire(blocking=True, timeout=1):
# process output event is received by the event loop. continue
self.likely_output.acquire(timeout=0.2)
while not self.stop_requested.is_set():
# Don't actually stop ourselves, even if we can't acquire the # Try to acquire the semaphore. This will not succeed until the next
# semaphore. LLDB can be a little lazy with the standard output # process output event is received by the event loop.
# events, so we use the semaphore as way to respond much faster to self.likely_output.acquire(blocking=True, timeout=0.2)
# output than we otherwise would, but, even if we don't get an
# event, we should still read the output, albeit at a slower pace. # Don't actually stop ourselves, even if we can't acquire the
# semaphore. LLDB can be a little lazy with the standard output
# Copy everything out to standard outputs. # events, so we use the semaphore as way to respond much faster to
while True: # output than we otherwise would, but, even if we don't get an
stdout = self.process.GetSTDOUT(1024) # event, we should still read the output, albeit at a slower pace.
stderr = self.process.GetSTDERR(1024)
# Copy everything out to standard outputs.
if len(stdout) == 0 and len(stderr) == 0: stdout = b""
break stderr = b""
while True:
print(stdout, file=sys.stdout, end="") stdout += self.process.GetSTDOUT(1024).encode(sys.stdout.encoding)
print(stderr, file=sys.stderr, end="") stderr += self.process.GetSTDERR(1024).encode(sys.stderr.encoding)
sys.stdout.flush() if len(stdout) == 0 and len(stderr) == 0:
sys.stderr.flush() # Note that, even if we have pulled nothing new from LLDB,
# we still only exit the loop once we manage to push out
# both buffers in their entirety.
#
# This is consistent with the behavior of blocking on STDOUT
# and STDERR that we want, even if the underlying files are
# actually non-blocking.
break
# Crutially, we don't release the semaphore here. Releasing is the try:
# job of the on_output_event function. stdout = stdout[sys.stdout.buffer.write(stdout) :]
sys.stdout.buffer.flush()
except BlockingIOError as e:
# STDOUT is nonblocking at this point, and so writes may
# fail. We trim off however much we have managed to write
# from the buffer, and try again in the next iteration.
stdout = stdout[e.characters_written :]
try:
stderr = stderr[sys.stderr.buffer.write(stderr) :]
sys.stderr.buffer.flush()
except BlockingIOError as e:
# Same goes for STDERR as goes for STDOUT.
stderr = stderr[e.characters_written :]
# Crucially, we don't release the semaphore here. Releasing is the
# job of the on_output_event function.
self.stop_fulfilled.release()
@override @override
def on_output_event(self) -> None: def on_output_event(self) -> None:
@ -283,20 +340,83 @@ class IODriverPlainText(IODriver):
self.process = process self.process = process
self.stop_requested.clear() self.stop_requested.clear()
os.set_blocking(sys.stdin.fileno(), False) os.set_blocking(sys.stdin.fileno(), False)
self.in_thr = threading.Thread(target=self._handle_input)
self.out_thr = threading.Thread(target=self._handle_output) # Nonblocking output is NOT what we want, but in UNIX systems O_NONBLOCK
self.in_thr.start() # is set in the context of the so-called "open file description"[1][2],
self.out_thr.start() # rather than in the context of the file descriptor itself. So, these
# systems will helpfully - and silently, of course - propagate a change
# in blocking policy to all file descriptors that share the same open
# file description - such as ones created through F_DUPFD or dup(2).
#
# Since, in general, we can't know how STDIN, STDOUT and STDERR are
# related to each other ahead of time, and, more specifically, they
# often share the exact same open file description, we have to be able
# to gracefully handle the case in which setting O_NONBLOCK for STDIN
# will also necessarily set it for STDOUT and STDERR.
#
# The strategy this class elects to use, then, is to explicitly set all
# of them to the same blocking policy. While this doesn't solve the
# issue, it at least makes it so that it's not as surprising as it would
# be, otherwise. :)
#
# [1]: https://pubs.opengroup.org/onlinepubs/9799919799/
# [2]: https://linux.die.net/man/2/fcntl
try:
os.set_blocking(sys.stdout.fileno(), False)
except OSError:
# It's not guaranteed that sys.stdout is actually backed by a file,
# or that that file supports non-blocking operation. In fact, the
# Pwndbg CLI itself supports swapping out output the output streams
# as part of capturing command output.
#
# As such, we must also be able to gracefully handle this case.
self._stdout_nonblock_failed = True
try:
os.set_blocking(sys.stderr.fileno(), False)
except OSError:
# Same as above.
self._stderr_nonblock_failed = True
self.start_requested.release(2)
self._running = True
@override @override
def stop(self) -> None: def stop(self) -> None:
# Politely ask for the I/O processors to stop, and wait until they have # Politely ask for the I/O processors to stop, and wait until they have
# stopped on their own terms. # stopped on their own terms.
assert self._running, "Tried to stop an IODriverPlainText that is not running"
self.stop_requested.set() self.stop_requested.set()
self.in_thr.join() self.stop_fulfilled.acquire(blocking=True)
self.out_thr.join() self.stop_fulfilled.acquire(blocking=True)
os.set_blocking(sys.stdin.fileno(), True) os.set_blocking(sys.stdin.fileno(), True)
# See start()
try:
os.set_blocking(sys.stdout.fileno(), True)
except OSError:
if not self._stdout_nonblock_failed:
raise
try:
os.set_blocking(sys.stderr.fileno(), True)
except OSError:
if not self._stderr_nonblock_failed:
raise
self._stdout_nonblock_failed = False
self._stderr_nonblock_failed = False
self.process = None self.process = None
self._running = False
@override
def close(self) -> None:
if self._running:
self.stop()
self._closed.set()
self.in_thr.join()
self.out_thr.join()
def make_pty() -> Tuple[str, int] | None: def make_pty() -> Tuple[str, int] | None:
@ -361,6 +481,8 @@ class IODriverPseudoTerminal(IODriver):
io_thread: threading.Thread io_thread: threading.Thread
process: lldb.SBProcess process: lldb.SBProcess
termcontrol: OpportunisticTerminalControl termcontrol: OpportunisticTerminalControl
_stdout_nonblock_failed: bool
_stderr_nonblock_failed: bool
has_terminal_control: bool has_terminal_control: bool
@ -412,6 +534,9 @@ class IODriverPseudoTerminal(IODriver):
self.input_buffer = b"" self.input_buffer = b""
self.process = None self.process = None
self._stdout_nonblock_failed = False
self._stderr_nonblock_failed = False
@override @override
def stdio(self) -> Tuple[str | None, str | None, str | None]: def stdio(self) -> Tuple[str | None, str | None, str | None]:
return self.worker, self.worker, self.worker return self.worker, self.worker, self.worker
@ -440,8 +565,14 @@ class IODriverPseudoTerminal(IODriver):
data = os.read(self.manager, 1024) data = os.read(self.manager, 1024)
if len(data) == 0: if len(data) == 0:
break break
print(data.decode("utf-8"), end="")
sys.stdout.flush() while len(data) > 0:
try:
data = data[sys.stdout.buffer.write(data) :]
sys.stdout.buffer.flush()
except BlockingIOError as e:
data = data[e.characters_written :]
except IOError: except IOError:
pass pass
@ -453,6 +584,17 @@ class IODriverPseudoTerminal(IODriver):
self.stop_requested.clear() self.stop_requested.clear()
os.set_blocking(sys.stdin.fileno(), False) os.set_blocking(sys.stdin.fileno(), False)
# Same reasoning as IODriverPlainText applies here.
try:
os.set_blocking(sys.stdout.fileno(), False)
except OSError:
self._stdout_nonblock_failed = True
try:
os.set_blocking(sys.stderr.fileno(), False)
except OSError:
self._stderr_nonblock_failed = True
self.was_line_buffering = self.termcontrol.get_line_buffering() self.was_line_buffering = self.termcontrol.get_line_buffering()
self.was_echoing = self.termcontrol.get_echo() self.was_echoing = self.termcontrol.get_echo()
@ -470,9 +612,25 @@ class IODriverPseudoTerminal(IODriver):
self.io_thread.join() self.io_thread.join()
os.set_blocking(sys.stdin.fileno(), True) os.set_blocking(sys.stdin.fileno(), True)
# Same reasoning as IODriverPlainText applies here.
try:
os.set_blocking(sys.stdout.fileno(), True)
except OSError:
if not self._stdout_nonblock_failed:
raise
try:
os.set_blocking(sys.stderr.fileno(), True)
except OSError:
if not self._stderr_nonblock_failed:
raise
self.termcontrol.set_line_buffering(self.was_line_buffering) self.termcontrol.set_line_buffering(self.was_line_buffering)
self.termcontrol.set_echo(self.was_echoing) self.termcontrol.set_echo(self.was_echoing)
self._stdout_nonblock_failed = False
self._stderr_nonblock_failed = False
self.process = None self.process = None
@override @override
@ -489,3 +647,7 @@ class IODriverPseudoTerminal(IODriver):
# #
# TODO: Replace controlling PTY of the process once it is set up. # TODO: Replace controlling PTY of the process once it is set up.
pass pass
@override
def close(self) -> None:
pass

319
pwndbg/dbg/lldb/repl/proc.py
Unescape View File

@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import contextlib import contextlib
import enum
import sys import sys
from asyncio import CancelledError from asyncio import CancelledError
from typing import Any from typing import Any
@ -171,10 +172,27 @@ def _updates_scope_counter(target: str) -> Callable[[Callable[..., Any]], Any]:
return sub0 return sub0
class _IODriverState(enum.Enum):
STOPPED = 0
"The IODriver is not running, and must be entirely restarted"
RUNNING = 1
"The IODriver is running"
PAUSED = 2
"The IODriver is not running, but is alive and can be resumed"
class ProcessDriver: class ProcessDriver:
""" """
Drives the execution of a process, responding to its events and handling its Drives the execution of a process, responding to its events and handling its
I/O, and exposes a simple synchronous interface to the REPL interface. I/O, and exposes a simple synchronous interface to the REPL interface.
# IODriver State Machine
Because LLDB can make Python code from Pwndbg execute while an I/O driver is
active, and having the I/O driver active while Pwndbg is running leads to
all sorts of fun failure modes, we want to be able to pause it temporarily.
We, thus, use the states described in _IODriverState to keep track of what
operations may be performed on the current IODriver.
""" """
io: IODriver io: IODriver
@ -195,6 +213,9 @@ class ProcessDriver:
_in_run_coroutine: int _in_run_coroutine: int
"Nested scope counter for run_coroutine" "Nested scope counter for run_coroutine"
_io_driver_state: _IODriverState
"Whether the I/O driver is currently running"
def __init__(self, event_handler: EventHandler, debug=False): def __init__(self, event_handler: EventHandler, debug=False):
self.io = None self.io = None
self.process = None self.process = None
@ -205,20 +226,31 @@ class ProcessDriver:
self._pending_cancellation = False self._pending_cancellation = False
self._in_run_until_next_stop = 0 self._in_run_until_next_stop = 0
self._in_run_coroutine = 0 self._in_run_coroutine = 0
self._io_driver_state = _IODriverState.STOPPED
def __enter__(self) -> ProcessDriver:
return self
def __exit__(self, _exc_type, _exc_val, _exc_tb) -> None:
if self.io is not None:
self.io.close()
def debug_print(self, *args, **kwargs) -> None: def debug_print(self, *args, **kwargs) -> None:
if self.debug: if self.debug:
try: from io import StringIO
print("[*] ProcessDriver: ", end="")
print(*args, **kwargs) with StringIO() as out:
except BlockingIOError as e: print("[*] ProcessDriver: ", end="", file=out)
print(*args, file=out, **kwargs)
message = out.getvalue().encode(sys.stdout.encoding)
while len(message) > 0:
try: try:
# Try to inform the user of the error. message = message[sys.stdout.buffer.write(message) :]
print( sys.stdout.buffer.flush()
f"[-] ProcessDriver: Error after printing {e.characters_written} characters in the previous debug message. Information may be missing." except BlockingIOError as e:
) message = message[e.characters_written :]
except BlockingIOError:
pass
def has_process(self) -> bool: def has_process(self) -> bool:
""" """
@ -296,6 +328,63 @@ class ProcessDriver:
# Perform the called-provided action. # Perform the called-provided action.
interrupt() interrupt()
def _start_io_driver(self):
"""
Starts the IODriver handling I/O from the process.
"""
assert self._io_driver_state == _IODriverState.STOPPED
self.io.start(process=self.process)
self.debug_print(f"{self._io_driver_state} -> {_IODriverState.RUNNING}")
self._io_driver_state = _IODriverState.RUNNING
def _stop_io_driver(self):
"""
Stops the IODriver handling I/O from the process.
"""
assert (
self._io_driver_state == _IODriverState.PAUSED
or self._io_driver_state == _IODriverState.RUNNING
)
if self._io_driver_state == _IODriverState.PAUSED:
# Not invalid, but currently a NOP. See pause_io_if_running()
self._io_driver_state = _IODriverState.STOPPED
return
self.io.stop()
self.debug_print(f"{self._io_driver_state} -> {_IODriverState.STOPPED}")
self._io_driver_state = _IODriverState.STOPPED
def pause_io_if_running(self) -> None:
"""
Pauses the handling of process I/O if it is currently running.
"""
if self._io_driver_state != _IODriverState.RUNNING:
return
# Currently, pausing and stopping both stop the IODriver. Nonetheless,
# these operations must stay separate in the state machine, as they are
# meaningfully distinct in other ways.
self.io.stop()
self.debug_print(f"{self._io_driver_state} -> {_IODriverState.PAUSED}")
self._io_driver_state = _IODriverState.PAUSED
def resume_io_if_running(self) -> None:
"""
Resumes the handling of process I/O if it is currently running.
"""
if self._io_driver_state != _IODriverState.PAUSED:
return
assert self.process is not None
self.io.start(process=self.process)
self.debug_print(f"{self._io_driver_state} -> {_IODriverState.RUNNING}")
self._io_driver_state = _IODriverState.RUNNING
@_updates_scope_counter(target="_in_run_until_next_stop") @_updates_scope_counter(target="_in_run_until_next_stop")
def _run_until_next_stop( def _run_until_next_stop(
self, self,
@ -319,113 +408,126 @@ class ProcessDriver:
# If `only_if_started` is set, we defer the starting of the I/O driver # If `only_if_started` is set, we defer the starting of the I/O driver
# to the moment the start event is observed. Otherwise, we just start it # to the moment the start event is observed. Otherwise, we just start it
# immediately. # immediately.
io_started = False try:
if with_io and not only_if_started: if with_io and not only_if_started:
self.io.start(process=self.process) self._start_io_driver()
io_started = True
# Pick the first timeout value.
# Pick the first timeout value. timeout_time = first_timeout
timeout_time = first_timeout
# If `only_if_started` is not set, assume the process must have been
# If `only_if_started` is not set, assume the process must have been # started by a previous action and is running.
# started by a previous action and is running. running = not only_if_started
running = not only_if_started
result = None
result = None last_event = None
last_event = None delay_until_io_stopped: List[Callable[[], None]] = []
while True: while True:
event = lldb.SBEvent() event = lldb.SBEvent()
if not self.listener.WaitForEvent(timeout_time, event): if not self.listener.WaitForEvent(timeout_time, event):
self.debug_print(
f"Timed out after {timeout_time}s",
)
timeout_time = timeout
# If the process isn't running, we should stop.
if not running:
self.debug_print( self.debug_print(
"Waited too long for process to start running, giving up", f"Timed out after {timeout_time}s",
) )
result = _PollResultTimedOut(last_event) timeout_time = timeout
break
continue
last_event = event
if self.debug: # If the process isn't running, we should stop.
descr = lldb.SBStream() if not running:
if event.GetDescription(descr): self.debug_print(
self.debug_print(descr.GetData()) "Waited too long for process to start running, giving up",
else: )
self.debug_print(f"No description for {event}") result = _PollResultTimedOut(last_event)
if lldb.SBTarget.EventIsTargetEvent(event):
if event.GetType() == lldb.SBTarget.eBroadcastBitModulesLoaded:
# Notify the event handler that new modules got loaded in.
if fire_events:
self.eh.modules_loaded()
elif lldb.SBProcess.EventIsProcessEvent(event):
if (
event.GetType() == lldb.SBProcess.eBroadcastBitSTDOUT
or event.GetType() == lldb.SBProcess.eBroadcastBitSTDERR
):
# Notify the I/O driver that the process might have something
# new for it to consume.
self.io.on_output_event()
elif event.GetType() == lldb.SBProcess.eBroadcastBitStateChanged:
# The state of the process has changed.
new_state = lldb.SBProcess.GetStateFromEvent(event)
was_resumed = lldb.SBProcess.GetRestartedFromEvent(event)
if new_state == lldb.eStateStopped and not was_resumed:
# The process has stopped, so we're done processing events
# for the time being. Trigger the stopped event and return.
if fire_events:
self.eh.suspended(event)
result = _PollResultStopped(event)
break break
if new_state == lldb.eStateRunning or new_state == lldb.eStateStepping: continue
running = True last_event = event
# Start the I/O driver here if its start got deferred
# because of `only_if_started` being set.
if only_if_started and with_io:
self.io.start(process=self.process)
io_started = True
if ( if self.debug:
new_state == lldb.eStateExited descr = lldb.SBStream()
or new_state == lldb.eStateCrashed if event.GetDescription(descr):
or new_state == lldb.eStateDetached self.debug_print(descr.GetData())
): else:
# Nothing else for us to do here. Clear our internal self.debug_print(f"No description for {event}")
# references to the process, fire the exit event, and leave.
self.debug_print(
f"Process exited with state {new_state}",
)
self.process = None
self.listener = None
if lldb.SBTarget.EventIsTargetEvent(event):
if event.GetType() == lldb.SBTarget.eBroadcastBitModulesLoaded:
# Notify the event handler that new modules got loaded in.
if fire_events: if fire_events:
self.eh.exited() self.eh.modules_loaded()
if new_state == lldb.eStateExited: elif lldb.SBProcess.EventIsProcessEvent(event):
proc = lldb.SBProcess.GetProcessFromEvent(event) if (
desc = ( event.GetType() == lldb.SBProcess.eBroadcastBitSTDOUT
"" if not proc.exit_description else f" ({proc.exit_description})" or event.GetType() == lldb.SBProcess.eBroadcastBitSTDERR
):
# Notify the I/O driver that the process might have something
# new for it to consume.
self.io.on_output_event()
elif event.GetType() == lldb.SBProcess.eBroadcastBitStateChanged:
# The state of the process has changed.
new_state = lldb.SBProcess.GetStateFromEvent(event)
was_resumed = lldb.SBProcess.GetRestartedFromEvent(event)
if new_state == lldb.eStateStopped and not was_resumed:
# The process has stopped, so we're done processing events
# for the time being. Trigger the stopped event and return.
if fire_events:
self.eh.suspended(event)
result = _PollResultStopped(event)
break
if new_state == lldb.eStateRunning or new_state == lldb.eStateStepping:
running = True
# Start the I/O driver here if its start got deferred
# because of `only_if_started` being set.
if only_if_started and with_io:
self._start_io_driver()
if (
new_state == lldb.eStateExited
or new_state == lldb.eStateCrashed
or new_state == lldb.eStateDetached
):
# Nothing else for us to do here. Clear our internal
# references to the process, fire the exit event, and leave.
self.debug_print(
f"Process exited with state {new_state}",
) )
print_info(f"process exited with status {proc.exit_state}{desc}") self.process = None
elif new_state == lldb.eStateCrashed: self.listener = None
print_info("process crashed")
elif new_state == lldb.eStateDetached: if fire_events:
print_info("process detached") self.eh.exited()
result = _PollResultExited(new_state) if new_state == lldb.eStateExited:
break proc = lldb.SBProcess.GetProcessFromEvent(event)
desc = (
""
if not proc.exit_description
else f" ({proc.exit_description})"
)
delay_until_io_stopped.append(
lambda: print_info(
f"process exited with status {proc.exit_state}{desc}"
)
)
elif new_state == lldb.eStateCrashed:
delay_until_io_stopped.append(lambda: print_info("process crashed"))
elif new_state == lldb.eStateDetached:
delay_until_io_stopped.append(
lambda: print_info("process detached")
)
result = _PollResultExited(new_state)
break
finally:
if self._io_driver_state != _IODriverState.STOPPED:
self._stop_io_driver()
if isinstance(result, _PollResultExited):
self.io.close()
self.io = None
if io_started: for fn in delay_until_io_stopped:
self.io.stop() fn()
return result return result
@ -725,6 +827,7 @@ class ProcessDriver:
This function always uses a plain text IODriver, as there is no way to This function always uses a plain text IODriver, as there is no way to
guarantee any other driver will work. guarantee any other driver will work.
""" """
assert self.io is None
self.io = IODriverPlainText() self.io = IODriverPlainText()
error = lldb.SBError() error = lldb.SBError()
@ -754,6 +857,7 @@ class ProcessDriver:
""" """
Launch a process in the host system. Launch a process in the host system.
""" """
assert self.io is None
self.io = io self.io = io
error = lldb.SBError() error = lldb.SBError()
@ -779,6 +883,7 @@ class ProcessDriver:
if pid == 0: if pid == 0:
return lldb.SBError("PID of 0 or no PID was given") return lldb.SBError("PID of 0 or no PID was given")
assert self.io is None
self.io = IODriverPlainText() self.io = IODriverPlainText()
error = lldb.SBError() error = lldb.SBError()
@ -789,6 +894,7 @@ class ProcessDriver:
""" """
Attatch to a process in the host system. Attatch to a process in the host system.
""" """
assert self.io is None
self.io = IODriverPlainText() self.io = IODriverPlainText()
error = lldb.SBError() error = lldb.SBError()
@ -868,6 +974,7 @@ class ProcessDriver:
assert self.listener.IsValid() assert self.listener.IsValid()
assert self.process.IsValid() assert self.process.IsValid()
assert self.io is None
self.io = io self.io = io
# It's not guaranteed that the process is actually alive, as it might be # It's not guaranteed that the process is actually alive, as it might be

Write Preview
Loading…
Cancel
Save