diff --git a/newsfragments/3237.feature.rst b/newsfragments/3237.feature.rst new file mode 100644 index 0000000000..d74952c641 --- /dev/null +++ b/newsfragments/3237.feature.rst @@ -0,0 +1 @@ +Improve defaults if ``shell=True`` for `trio.run_process`, as well as backport the ``process_group`` argument. diff --git a/src/trio/_subprocess.py b/src/trio/_subprocess.py index c1b162e308..02c5dd6f14 100644 --- a/src/trio/_subprocess.py +++ b/src/trio/_subprocess.py @@ -2,9 +2,11 @@ import contextlib import os +import signal import subprocess import sys import warnings +from collections.abc import Awaitable from contextlib import ExitStack from functools import partial from typing import ( @@ -30,7 +32,6 @@ from ._util import NoPublicConstructor, final if TYPE_CHECKING: - import signal from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from io import TextIOWrapper @@ -441,24 +442,45 @@ async def _windows_deliver_cancel(p: Process) -> None: # noqa: RUF029 ) -async def _posix_deliver_cancel(p: Process) -> None: - try: - p.terminate() - await trio.sleep(5) - warnings.warn( - RuntimeWarning( - f"process {p!r} ignored SIGTERM for 5 seconds. " - "(Maybe you should pass a custom deliver_cancel?) " - "Trying SIGKILL.", - ), - stacklevel=1, - ) - p.kill() - except OSError as exc: - warnings.warn( - RuntimeWarning(f"tried to kill process {p!r}, but failed with: {exc!r}"), - stacklevel=1, +def _get_posix_deliver_cancel( + process_group: int | None, +) -> Callable[[Process], Awaitable[None]]: + async def _posix_deliver_cancel(p: Process) -> None: + should_deliver_to_pg = ( + sys.platform != "win32" + and process_group is not None + and os.getpgrp() != os.getpgid(p.pid) ) + try: + # TODO: should Process#terminate do this special logic + if sys.platform != "win32" and should_deliver_to_pg: + os.killpg(os.getpgid(p.pid), signal.SIGTERM) + else: + p.terminate() + + await trio.sleep(5) + warnings.warn( + RuntimeWarning( + f"process {p!r} ignored SIGTERM for 5 seconds. " + "(Maybe you should pass a custom deliver_cancel?) " + "Trying SIGKILL.", + ), + stacklevel=1, + ) + + if sys.platform != "win32" and should_deliver_to_pg: + os.killpg(os.getpgid(p.pid), signal.SIGKILL) + else: + p.kill() + except OSError as exc: + warnings.warn( + RuntimeWarning( + f"tried to kill process {p!r}, but failed with: {exc!r}" + ), + stacklevel=1, + ) + + return _posix_deliver_cancel # Use a private name, so we can declare platform-specific stubs below. @@ -472,6 +494,7 @@ async def _run_process( check: bool = True, deliver_cancel: Callable[[Process], Awaitable[object]] | None = None, task_status: TaskStatus[Process] = trio.TASK_STATUS_IGNORED, + shell: bool = False, **options: object, ) -> subprocess.CompletedProcess[bytes]: """Run ``command`` in a subprocess and wait for it to complete. @@ -689,6 +712,9 @@ async def my_deliver_cancel(process): "stderr=subprocess.PIPE is only valid with nursery.start, " "since that's the only way to access the pipe", ) + + options["shell"] = shell + if isinstance(stdin, (bytes, bytearray, memoryview)): input_ = stdin options["stdin"] = subprocess.PIPE @@ -708,12 +734,36 @@ async def my_deliver_cancel(process): raise ValueError("can't specify both stderr and capture_stderr") options["stderr"] = subprocess.PIPE + # ensure things can be killed including a shell's child processes + if shell and sys.platform != "win32" and "process_group" not in options: + options["process_group"] = 0 + if deliver_cancel is None: if os.name == "nt": deliver_cancel = _windows_deliver_cancel else: assert os.name == "posix" - deliver_cancel = _posix_deliver_cancel + deliver_cancel = _get_posix_deliver_cancel(options.get("process_group")) # type: ignore[arg-type] + + if ( + sys.platform != "win32" + and options.get("process_group") is not None + and sys.version_info < (3, 11) + ): + # backport the argument to Python versions prior to 3.11 + preexec_fn = options.get("preexec_fn") + process_group = options.pop("process_group") + + def new_preexecfn() -> object: # pragma: no cover + assert sys.platform != "win32" + os.setpgid(0, process_group) # type: ignore[arg-type] + + if callable(preexec_fn): + return preexec_fn() + else: + return None + + options["preexec_fn"] = new_preexecfn stdout_chunks: list[bytes | bytearray] = [] stderr_chunks: list[bytes | bytearray] = [] diff --git a/src/trio/_tests/test_subprocess.py b/src/trio/_tests/test_subprocess.py index 15c88be25e..7f266a927e 100644 --- a/src/trio/_tests/test_subprocess.py +++ b/src/trio/_tests/test_subprocess.py @@ -765,3 +765,87 @@ async def wait_and_tell(proc: Process) -> None: # for everything to notice await noticed_exit.wait() assert noticed_exit.is_set(), "child task wasn't woken after poll, DEADLOCK" + + +@pytest.mark.skipif(not posix, reason="posix only") +@slow +async def test_shells_killed_by_default() -> None: + assert sys.platform != "win32" + + async with _core.open_nursery() as nursery: + proc = await nursery.start( + partial( + trio.run_process, + 'python -u -c "import os, time, signal; ' + f"os.kill({os.getpid()}, signal.SIGHUP); " + 'print(os.getpid()); time.sleep(5)" | cat', + shell=True, + stdout=subprocess.PIPE, + ) + ) + + with trio.fail_after(1): + with trio.open_signal_receiver(signal.SIGHUP) as signal_aiter: # type: ignore[attr-defined,unused-ignore] + signum = await signal_aiter.__anext__() + assert signum == signal.SIGHUP # type: ignore[attr-defined,unused-ignore] + + nursery.cancel_scope.cancel() + + chunks = [] + with trio.move_on_after(0.1): + async with proc.stdout: + async for c in proc.stdout: + chunks.append(c) # noqa: PERF401 + + # give the child time to wind down execution (this is unfortunate) + await trio.sleep(1) + + child_pid = int(b"".join(chunks)) + if sys.platform == "linux": + # docker containers can have an evil init which + # won't reap zombie children, so we have to work around that + with pytest.raises(FileNotFoundError): + with open(f"/proc/{child_pid}/status") as f: # noqa: ASYNC230 + hit_zombie = False + for line in f: + if line.startswith("State:"): + assert "Z" in line + hit_zombie = True + + if hit_zombie: + raise FileNotFoundError + else: + with pytest.raises(OSError, match="No such process"): + os.kill(child_pid, 0) + + +@pytest.mark.skipif(not posix, reason="posix only") +@slow +async def test_process_group_SIGKILL_escalation(mock_clock: MockClock) -> None: + assert sys.platform != "win32" + mock_clock.autojump_threshold = 1 + + with pytest.warns(RuntimeWarning, match=".*ignored SIGTERM.*"): # noqa: PT031 + async with _core.open_nursery() as nursery: + nursery.start_soon( + # this ignore is because process_group=0 on Windows + partial( # type: ignore[call-arg,unused-ignore] + trio.run_process, + [ + "python", + "-c", + "import os, time, signal;" + f"os.kill({os.getpid()}, signal.SIGHUP);" + "signal.signal(signal.SIGTERM, signal.SIG_IGN);" + "time.sleep(10)", + ], + process_group=0, + ) + ) + + with trio.fail_after(1): + with trio.open_signal_receiver(signal.SIGHUP) as signal_aiter: # type: ignore[attr-defined,unused-ignore] + signum = await signal_aiter.__anext__() + assert signum == signal.SIGHUP # type: ignore[attr-defined,unused-ignore] + + nursery.cancel_scope.cancel()