diff --git a/docs/changelog/223.bugfix.rst b/docs/changelog/223.bugfix.rst new file mode 100644 index 0000000..c66df25 --- /dev/null +++ b/docs/changelog/223.bugfix.rst @@ -0,0 +1 @@ +Limit candidate stdout and stderr during interpreter queries so a noisy candidate cannot exhaust the caller's memory. diff --git a/docs/explanation.rst b/docs/explanation.rst index f40b51f..c5dbfba 100644 --- a/docs/explanation.rst +++ b/docs/explanation.rst @@ -306,6 +306,9 @@ legitimate interpreters may be skipped; if too high, the discovery process may t when encountering problematic interpreters. For invalid values, python-discovery logs a warning and uses the default of 15 seconds. Set ``inf`` to disable the timeout. +Discovery skips a candidate that produces more than 1,048,576 characters on either stdout or stderr, even when the +timeout is disabled. + Supported Python versions ------------------------- diff --git a/src/python_discovery/_cached_py_info.py b/src/python_discovery/_cached_py_info.py index d9aff86..4a21d05 100644 --- a/src/python_discovery/_cached_py_info.py +++ b/src/python_discovery/_cached_py_info.py @@ -8,15 +8,17 @@ import os import pkgutil import secrets -import subprocess # ruff:ignore[suspicious-subprocess-import] import sys import tempfile from collections import OrderedDict -from contextlib import contextmanager +from contextlib import contextmanager, suppress from functools import lru_cache +from io import DEFAULT_BUFFER_SIZE from pathlib import Path from shlex import quote -from subprocess import Popen, TimeoutExpired # ruff:ignore[suspicious-subprocess-import] +from subprocess import DEVNULL, PIPE, Popen, TimeoutExpired, list2cmdline # ruff:ignore[suspicious-subprocess-import] +from threading import Event, Thread +from time import monotonic from typing import TYPE_CHECKING, Final from ._cache import NoOpCache @@ -25,7 +27,7 @@ if TYPE_CHECKING: from collections.abc import Generator, Mapping - from typing import TypedDict + from typing import IO, TypedDict from ._cache import ContentStore, PyInfoCache @@ -45,6 +47,7 @@ class CacheEntryMeta(TypedDict): _DEFAULT_QUERY_TIMEOUT: Final[float] = 15.0 # poll(2) limits milliseconds to a signed C int. _MAX_QUERY_TIMEOUT: Final[float] = (2**31 - 1) / 1000 +_MAX_QUERY_OUTPUT: Final[int] = 1 << 20 class _UnsupportedInterpreterError(RuntimeError): @@ -231,7 +234,7 @@ def _run_subprocess( cmd: list[str] | str = [exe, str(py_info_script), start_cookie, end_cookie] if executable_args is not None: # PEP 514 ExecutableArguments is raw command line text, placed before the script per the PEP - cmd = f"{subprocess.list2cmdline(cmd[:1])} {executable_args} {subprocess.list2cmdline(cmd[1:])}" + cmd = f"{list2cmdline(cmd[:1])} {executable_args} {list2cmdline(cmd[1:])}" env = dict(env) env.pop("__PYVENV_LAUNCHER__", None) env["PYTHONUTF8"] = "1" @@ -240,19 +243,14 @@ def _run_subprocess( process = Popen( # ruff:ignore[subprocess-without-shell-equals-true] cmd, universal_newlines=True, - stdin=subprocess.PIPE, - stderr=subprocess.PIPE, - stdout=subprocess.PIPE, + stdin=DEVNULL, + stderr=PIPE, + stdout=PIPE, env=env, encoding="utf-8", errors="backslashreplace", ) - out, err = process.communicate(timeout=timeout) - code = process.returncode - except TimeoutExpired: - process.kill() - process.communicate() - out, err, code = "", "timed out", -1 + out, err, code = _collect_query_output(process, timeout) except OSError as os_error: out, err, code = "", os_error.strerror, os_error.errno if code != 0: @@ -280,6 +278,75 @@ def _run_subprocess( return None, result +def _collect_query_output(process: Popen[str], timeout: float | None) -> tuple[str, str, int]: + stdout, stderr = process.stdout, process.stderr + if stdout is None or stderr is None: + msg = "query output pipes are required" + raise ValueError(msg) + out_chunks: list[str] = [] + err_chunks: list[str] = [] + overflow = Event() + read_errors: list[OSError] = [] + + readers = ( + Thread(target=_drain_query_stream, args=(process, stdout, out_chunks, overflow, read_errors), daemon=True), + Thread(target=_drain_query_stream, args=(process, stderr, err_chunks, overflow, read_errors), daemon=True), + ) + for reader in readers: + reader.start() + if not _query_completed(process, readers, timeout): + with suppress(OSError): + process.kill() + for reader in readers: + reader.join() + process.wait() + return "", "timed out", -1 + if read_errors: + raise read_errors[0] + if overflow.is_set(): + return "", f"produced more than {_MAX_QUERY_OUTPUT} characters of output", -1 + return "".join(out_chunks), "".join(err_chunks), process.wait() + + +def _query_completed(process: Popen[str], readers: tuple[Thread, Thread], timeout: float | None) -> bool: + deadline = None if timeout is None else monotonic() + timeout + for reader in readers: + reader.join(None if deadline is None else max(deadline - monotonic(), 0)) + if any(reader.is_alive() for reader in readers): + return False + try: + process.wait(timeout=None if deadline is None else max(deadline - monotonic(), 0)) + except TimeoutExpired: + return False + return True + + +def _drain_query_stream( + process: Popen[str], stream: IO[str], chunks: list[str], overflow: Event, read_errors: list[OSError] +) -> None: + size = 0 + failed = False + with stream: + while True: + try: + chunk = stream.read(DEFAULT_BUFFER_SIZE) + except OSError as error: + read_errors.append(error) + failed = True + break + if not chunk: + break + chunks.append(chunk[: _MAX_QUERY_OUTPUT + 1 - size]) + size += len(chunk) + if size > _MAX_QUERY_OUTPUT: + overflow.set() + failed = True + break + if failed: + with suppress(OSError): + process.kill() + + @lru_cache(maxsize=128) def _query_timeout(raw: str | None) -> float | None: if raw is None: diff --git a/tests/test_cached_py_info.py b/tests/test_cached_py_info.py index 04bd5ba..ce8aebd 100644 --- a/tests/test_cached_py_info.py +++ b/tests/test_cached_py_info.py @@ -5,8 +5,9 @@ import os import sys import tempfile +from io import StringIO from pathlib import Path -from subprocess import TimeoutExpired, list2cmdline +from subprocess import DEVNULL, PIPE, Popen, TimeoutExpired, list2cmdline from typing import TYPE_CHECKING, Final from unittest.mock import MagicMock, patch @@ -14,7 +15,9 @@ from python_discovery import DiskCache, PythonInfo from python_discovery._cached_py_info import ( + _MAX_QUERY_OUTPUT, LogCmd, + _collect_query_output, _get_from_cache, _get_via_file_cache, _load_cached_py_info, @@ -87,25 +90,23 @@ def test_run_subprocess_bad_exe() -> None: def test_run_subprocess_invalid_json(mocker: MockerFixture) -> None: - mock_process = MagicMock() - mock_process.communicate.return_value = ("not json", "") - mock_process.returncode = 0 - mocker.patch("python_discovery._cached_py_info.Popen", return_value=mock_process) + mocker.patch("python_discovery._cached_py_info.Popen", return_value=_mock_process("not json", "", 0)) failure, result = _run_subprocess(PythonInfo, sys.executable, dict(os.environ)) assert failure is not None assert result is None +def _mock_process(out: str, err: str, code: int) -> MagicMock: + return MagicMock(stdout=StringIO(out), stderr=StringIO(err), returncode=code, **{"wait.return_value": code}) + + def test_run_subprocess_with_cookies(mocker: MockerFixture) -> None: start_cookie = "a" * 32 end_cookie = "b" * 32 payload = json.dumps(PythonInfo().to_dict()) out = f"pre{start_cookie[::-1]}{payload}{end_cookie[::-1]}post" - mock_process = MagicMock() - mock_process.communicate.return_value = (out, "") - mock_process.returncode = 0 - mocker.patch("python_discovery._cached_py_info.Popen", return_value=mock_process) + mocker.patch("python_discovery._cached_py_info.Popen", return_value=_mock_process(out, "", 0)) mocker.patch("python_discovery._cached_py_info.gen_cookie", side_effect=[start_cookie, end_cookie]) with patch("sys.stdout") as mock_stdout: @@ -117,28 +118,96 @@ def test_run_subprocess_with_cookies(mocker: MockerFixture) -> None: def test_run_subprocess_timeout(mocker: MockerFixture) -> None: - mock_process = MagicMock() - mock_process.communicate.side_effect = [TimeoutExpired(cmd="python", timeout=30), ("", "")] - mocker.patch("python_discovery._cached_py_info.Popen", return_value=mock_process) + mocker.patch("python_discovery._cached_py_info._collect_query_output", return_value=("", "timed out", -1)) failure, result = _run_subprocess(PythonInfo, sys.executable, dict(os.environ)) assert failure is not None assert "timed out" in str(failure) assert result is None - mock_process.kill.assert_called_once() - assert mock_process.communicate.call_count == 2 def test_run_subprocess_nonzero_exit(mocker: MockerFixture) -> None: - mock_process = MagicMock() - mock_process.communicate.return_value = ("some output", "some error") - mock_process.returncode = 1 - mocker.patch("python_discovery._cached_py_info.Popen", return_value=mock_process) + mocker.patch("python_discovery._cached_py_info.Popen", return_value=_mock_process("some output", "some error", 1)) failure, result = _run_subprocess(PythonInfo, sys.executable, dict(os.environ)) assert failure is not None assert "failed to query" in str(failure) assert result is None +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX executable script") +@pytest.mark.parametrize("stream", [pytest.param("stdout", id="stdout"), pytest.param("stderr", id="stderr")]) +def test_run_subprocess_rejects_large_output(tmp_path: Path, stream: str) -> None: + exe = tmp_path / "python" + exe.write_text( + f"#!{sys.executable}\nimport sys\nsys.{stream}.write('x' * {_MAX_QUERY_OUTPUT + 1})\n", + encoding="utf-8", + ) + exe.chmod(0o755) + + failure, result = _run_subprocess(PythonInfo, str(exe), dict(os.environ)) + assert result is None + assert failure is not None + assert "produced more than" in str(failure) + + +@pytest.mark.parametrize("stream", [pytest.param("stdout", id="stdout"), pytest.param("stderr", id="stderr")]) +@pytest.mark.parametrize("timeout", [pytest.param(15, id="finite-timeout"), pytest.param(None, id="no-timeout")]) +def test_collect_query_output_limit(stream: str, timeout: float | None) -> None: + code = f"import sys; sys.{stream}.write('x' * {_MAX_QUERY_OUTPUT + 1})" + process = Popen([sys.executable, "-c", code], stdin=DEVNULL, stdout=PIPE, stderr=PIPE, text=True) + assert _collect_query_output(process, timeout) == ( + "", + f"produced more than {_MAX_QUERY_OUTPUT} characters of output", + -1, + ) + + +@pytest.mark.parametrize( + "code", + [ + pytest.param("import time; time.sleep(2)", id="pipes-open"), + pytest.param("import sys,time; sys.stdout.close(); sys.stderr.close(); time.sleep(2)", id="pipes-closed"), + ], +) +def test_collect_query_output_timeout(code: str) -> None: + process = Popen([sys.executable, "-c", code], stdin=DEVNULL, stdout=PIPE, stderr=PIPE, text=True) + assert _collect_query_output(process, 0.1) == ("", "timed out", -1) + + +def test_collect_query_output_timeout_after_eof(mocker: MockerFixture) -> None: + process = mocker.MagicMock( + stdout=StringIO(""), + stderr=StringIO(""), + **{"wait.side_effect": [TimeoutExpired("python", 0), 0]}, + ) + + assert _collect_query_output(process, 1) == ("", "timed out", -1) + process.kill.assert_called_once() + + +def test_collect_query_output_read_error(mocker: MockerFixture) -> None: + stream = mocker.MagicMock() + stream.read.side_effect = OSError("broken pipe") + process = mocker.MagicMock(stdout=stream, stderr=StringIO(""), **{"wait.return_value": -9}) + + with pytest.raises(OSError, match="broken pipe"): + _collect_query_output(process, 1) + process.kill.assert_called_once() + + +@pytest.mark.parametrize( + ("stdout", "stderr"), + [ + pytest.param(None, StringIO(""), id="stdout"), + pytest.param(StringIO(""), None, id="stderr"), + ], +) +def test_collect_query_output_requires_pipes(stdout: StringIO | None, stderr: StringIO | None) -> None: + process = MagicMock(stdout=stdout, stderr=stderr) + + with pytest.raises(ValueError, match="query output pipes are required"): + _collect_query_output(process, 1) + + def test_load_cached_py_info_valid() -> None: store = MagicMock() content = PythonInfo().to_dict() @@ -341,7 +410,7 @@ def popen(mocker: MockerFixture) -> MagicMock: info = PythonInfo() info.executable = _SELECTED out = f"{start[::-1]}{info.to_json()}{end[::-1]}" - process = MagicMock(**{"communicate.return_value": (out, ""), "returncode": 0}) + process = _mock_process(out, "", 0) return mocker.patch("python_discovery._cached_py_info.Popen", return_value=process) diff --git a/tests/test_query_timeout.py b/tests/test_query_timeout.py index b4d2288..6db946f 100644 --- a/tests/test_query_timeout.py +++ b/tests/test_query_timeout.py @@ -3,12 +3,12 @@ import logging import os import sys -from subprocess import Popen from typing import TYPE_CHECKING import pytest from python_discovery import PythonInfo +from python_discovery._cached_py_info import _collect_query_output if TYPE_CHECKING: from pytest_mock import MockerFixture @@ -29,14 +29,14 @@ def test_from_exe_query_timeout( mocker: MockerFixture, caplog: pytest.LogCaptureFixture, raw: str | None, expected: float | None ) -> None: - communicate = mocker.spy(Popen, "communicate") + collect = mocker.patch("python_discovery._cached_py_info._collect_query_output", wraps=_collect_query_output) env = {key: value for key, value in os.environ.items() if key != "PY_DISCOVERY_TIMEOUT"} if raw is not None: env["PY_DISCOVERY_TIMEOUT"] = raw result = PythonInfo.from_exe(sys.executable, env=env, ignore_cache=True, resolve_to_host=False) assert result is not None assert result.version_info == sys.version_info - assert communicate.call_args_list[0].kwargs == {"timeout": expected} + assert collect.call_args_list[0].args[1] == expected assert not caplog.records @@ -54,13 +54,13 @@ def test_from_exe_query_timeout( ], ) def test_from_exe_invalid_timeout_warns_once(mocker: MockerFixture, caplog: pytest.LogCaptureFixture, raw: str) -> None: - communicate = mocker.spy(Popen, "communicate") + collect = mocker.patch("python_discovery._cached_py_info._collect_query_output", wraps=_collect_query_output) env = {**os.environ, "PY_DISCOVERY_TIMEOUT": raw} for _ in range(2): result = PythonInfo.from_exe(sys.executable, env=env, ignore_cache=True, resolve_to_host=False) assert result is not None assert result.version_info == sys.version_info - assert [call.kwargs for call in communicate.call_args_list] == [{"timeout": 15.0}] * 2 + assert [call.args[1] for call in collect.call_args_list] == [15.0] * 2 assert [(record.levelno, record.getMessage()) for record in caplog.records] == [ (logging.WARNING, f"ignoring PY_DISCOVERY_TIMEOUT={raw!r}, not a positive number of seconds; using 15.0") ]