feat(py): environment allowlist and isolated launch for the worker

pkg-py/src/commons/_execution/__init__.py

@@ -0,0 +1 @@
+"""Running model-written code in a worker process."""

pkg-py/src/commons/_execution/_env.py

@@ -0,0 +1,107 @@
+"""The environment the worker process is given.
+
+A subprocess inherits its parent's environment by default, and the parent
+holds API keys, session tokens, and database URLs that model-written code has
+no business reading. The worker therefore starts from an allowlist rather than
+from what happens to be set.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from collections.abc import Sequence
+from pathlib import Path
+
+__all__ = ["in_container", "interpreter_warning", "worker_command", "worker_env"]
+
+# Enough to find an interpreter and produce readable text, and nothing else.
+_KEEP = ("PATH", "LANG", "LD_LIBRARY_PATH")
+
+
+def worker_env(scratch_dir: str) -> dict[str, str]:
+    """Build the worker's environment from an allowlist of the parent's."""
+    env = {name: os.environ[name] for name in _KEEP if name in os.environ}
+    env.update(
+        {name: value for name, value in os.environ.items() if name.startswith("LC_")}
+    )
+    # Point the worker's idea of home and scratch space at a directory it is
+    # allowed to have: with no sandbox engaged these are all that keep it out
+    # of the user's dot files.
+    env["HOME"] = scratch_dir
+    env["TMPDIR"] = scratch_dir
+    return env
+
+
+def worker_command(
+    script: str, *args: str, executable: str = sys.executable
+) -> Sequence[str]:
+    """The argv that launches the worker.
+
+    ``-I`` is not a hardening extra to be traded off; it belongs with the
+    allowlist. Without it ``site`` imports ``usercustomize`` from the user site
+    directory, which runs before the worker and can write excluded variables
+    back into ``os.environ``. ``-u`` keeps the protocol channel unbuffered.
+
+    ``executable`` defaults to the interpreter running commons, which is what
+    makes the worker's installed packages match the host's.
+    """
+    return [executable, "-I", "-u", script, *args]
+
+
+# Files a container runtime leaves behind, used to tell "this system Python is
+# the image author's" from "this system Python is shared with other people".
+_CONTAINER_MARKERS = ("/.dockerenv", "/run/.containerenv")
+
+
+def _venv_includes_system_site(executable: str) -> bool | None:
+    """Read ``pyvenv.cfg`` for ``executable``; ``None`` if it is not a venv."""
+    # Deliberately not resolved: a virtual environment's bin/python is usually
+    # a symlink to the interpreter it was built from, and following it lands on
+    # that installation rather than on the environment being asked about.
+    config = Path(executable).absolute().parent.parent / "pyvenv.cfg"
+    try:
+        text = config.read_text()
+    except OSError:
+        return None
+    for line in text.splitlines():
+        key, sep, value = line.partition("=")
+        if sep and key.strip() == "include-system-site-packages":
+            return value.strip().lower() == "true"
+    return False
+
+
+def in_container() -> bool:
+    """Whether this process looks like it is running inside an image."""
+    return any(os.path.exists(marker) for marker in _CONTAINER_MARKERS)
+
+
+def interpreter_warning(
+    executable: str = sys.executable, *, containerised: bool | None = None
+) -> str | None:
+    """Say why this interpreter's startup hooks are not known, or ``None``.
+
+    Isolated mode drops the user site directory but not the global one, so a
+    ``.pth`` file in a shared installation's site-packages still runs before
+    the worker does. Whoever can write there can therefore run code inside it.
+    That is fine when the only person who can write there is the image author,
+    and not fine on a machine shared with other people. Pass ``containerised``
+    to state which case this is rather than let it be inferred.
+    """
+    includes_system_site = _venv_includes_system_site(executable)
+    if includes_system_site is False:
+        return None
+    if in_container() if containerised is None else containerised:
+        return None
+    reason = (
+        "is a virtual environment built with --system-site-packages"
+        if includes_system_site
+        else "is not a virtual environment"
+    )
+    return (
+        f"the code execution worker would run on {executable}, which {reason}, "
+        "so commons cannot tell what runs at its startup. Anyone who can write "
+        "to its site-packages can run code inside the worker. Use a virtual "
+        "environment that excludes system packages, or a container image whose "
+        "contents you control."
+    )

pkg-py/tests/test_execution_env.py

@@ -0,0 +1,185 @@
+"""What the worker process is allowed to inherit, and what it must not.
+
+These launch real interpreters. The property under test is what a child
+process can actually see, and an assertion about the contents of a dictionary
+would not have caught the failure this guards against.
+"""
+
+from __future__ import annotations
+
+import os
+import pathlib
+import subprocess
+import sys
+
+import pytest
+
+from commons._execution._env import (
+    interpreter_warning,
+    worker_command,
+    worker_env,
+)
+
+READ_SECRET = "import os; print(os.environ.get('COMMONS_TEST_SECRET'))"
+
+
+def test_a_variable_outside_the_allowlist_does_not_reach_the_worker(
+    tmp_path, monkeypatch
+) -> None:
+    monkeypatch.setenv("COMMONS_TEST_SECRET", "sk-not-a-real-key")
+
+    result = subprocess.run(
+        [sys.executable, "-c", READ_SECRET],
+        env=worker_env(str(tmp_path)),
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+
+    assert result.stdout.strip() == "None"
+
+
+def test_home_and_tmpdir_point_at_the_scratch_directory(tmp_path) -> None:
+    # Left pointing at the real home, the worker could read the user's dot
+    # files and write anywhere they can write. Both are the scratch directory
+    # so that code with no sandbox still has nowhere interesting to go.
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import os; print(os.environ['HOME'], os.environ['TMPDIR'])",
+        ],
+        env=worker_env(str(tmp_path)),
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+
+    assert result.stdout.split() == [str(tmp_path), str(tmp_path)]
+
+
+def _plant_usercustomize(executable: str, env: dict[str, str]) -> bool:
+    """Put code in the interpreter's user site directory that restores the secret.
+
+    Returns whether the interpreter would import it at all, so the test can
+    say it was skipped rather than pass without having tried anything.
+    """
+    probe = subprocess.run(
+        [
+            executable,
+            "-c",
+            "import site; print(site.ENABLE_USER_SITE); print(site.getusersitepackages())",
+        ],
+        env=env,
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    enabled, _, directory = probe.stdout.partition("\n")
+    if enabled.strip() != "True":
+        return False
+    target = pathlib.Path(directory.strip())
+    target.mkdir(parents=True, exist_ok=True)
+    (target / "usercustomize.py").write_text(
+        "import os\nos.environ['COMMONS_TEST_SECRET'] = 'restored'\n"
+    )
+    return True
+
+
+def test_startup_hooks_cannot_put_an_excluded_variable_back(
+    tmp_path, monkeypatch
+) -> None:
+    # An allowlist alone is not enough. `site` imports `usercustomize` from the
+    # user site directory before the worker runs, and that code can write
+    # straight back into os.environ. Isolated mode is what closes it.
+    monkeypatch.setenv("COMMONS_TEST_SECRET", "sk-not-a-real-key")
+    # A virtual environment turns the user site directory off, so the hole
+    # only opens on an interpreter like the one this venv was built from.
+    # That is exactly the interpreter the note warns about running on.
+    executable = getattr(sys, "_base_executable", None)
+    if executable is None:
+        pytest.skip("no non-virtual-environment interpreter to test against")
+    env = worker_env(str(tmp_path))
+    script = tmp_path / "worker.py"
+    script.write_text(READ_SECRET + "\n")
+    if not _plant_usercustomize(executable, env):
+        pytest.skip("this interpreter does not import usercustomize")
+
+    unguarded = subprocess.run(
+        [executable, str(script)], env=env, capture_output=True, text=True, check=True
+    )
+    guarded = subprocess.run(
+        worker_command(str(script), executable=executable),
+        env=env,
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+
+    # The control matters: without it a passing test proves nothing about -I.
+    assert unguarded.stdout.strip() == "restored"
+    assert guarded.stdout.strip() == "None"
+
+
+def test_a_virtual_environment_without_system_packages_is_accepted() -> None:
+    # The suite runs in exactly the kind of interpreter the design asks for.
+    assert interpreter_warning() is None
+
+
+def test_a_virtual_environment_that_includes_system_packages_is_flagged(
+    tmp_path,
+) -> None:
+    # Isolated mode drops the user site directory, not the global one, so a
+    # venv wired to the system site-packages is still exposed to anything
+    # installed there.
+    venv = tmp_path / "shared"
+    subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "venv",
+            "--system-site-packages",
+            "--without-pip",
+            str(venv),
+        ],
+        check=True,
+        capture_output=True,
+    )
+
+    warning = interpreter_warning(str(venv / "bin" / "python"), containerised=False)
+
+    assert warning is not None
+    assert "system" in warning
+
+
+def test_an_interpreter_outside_any_virtual_environment_is_flagged() -> None:
+    # The interpreter this virtual environment was built from, which is not
+    # itself inside one.
+    executable = getattr(sys, "_base_executable", None)
+    if executable is None:
+        pytest.skip("no non-virtual-environment interpreter to test against")
+
+    warning = interpreter_warning(executable, containerised=False)
+
+    assert warning is not None
+    assert executable in warning
+
+
+def test_the_worker_keeps_what_it_needs_to_run(tmp_path, monkeypatch) -> None:
+    # An allowlist that is too narrow fails differently but just as badly: the
+    # worker cannot find an interpreter or mangles non-ASCII output.
+    monkeypatch.setenv("LC_ALL", "en_US.UTF-8")
+
+    env = worker_env(str(tmp_path))
+
+    assert env["PATH"] == os.environ["PATH"]
+    assert env["LC_ALL"] == "en_US.UTF-8"
+
+
+def test_a_container_image_is_accepted_even_outside_a_virtual_environment() -> None:
+    # In an image, the only person who can write to site-packages is whoever
+    # built it, so a bare system interpreter is the expected arrangement
+    # rather than a shared machine's.
+    executable = getattr(sys, "_base_executable", None) or sys.executable
+
+    assert interpreter_warning(executable, containerised=True) is None