🐛 fix(query): bound candidate output

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.

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

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:

tests/test_cached_py_info.py

@@ -5,16 +5,19 @@
 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
 
 import pytest
 
 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)
 
 

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")
     ]