sand-boxer/tests/test_bwrap.py

"""ext.bwrap — local namespace isolation extension."""

import json
import signal
import socket
import stat
import subprocess
import threading
from unittest.mock import MagicMock, patch

import pytest

from sandboxer.extensions.base import SandboxExtension
from sandboxer.extensions.bwrap import BwrapExtension
from sandboxer.extensions.bwrap_runner import _run
from sandboxer.models import IsolationSpec, Profile, Reachability
from sandboxer.reachability.enrich import enrich_reachability


def _profile() -> Profile:
    return Profile.model_validate(
        {
            "id": "profile.bwrap-local",
            "version": "1.0.0",
            "extension": "ext.bwrap",
            "isolation": {"level": "process"},
        }
    )


def test_is_sandbox_extension_subclass() -> None:
    assert issubclass(BwrapExtension, SandboxExtension)


def test_process_is_a_valid_isolation_level() -> None:
    assert IsolationSpec(level="process").level == "process"


def test_bwrap_bin_env_override(monkeypatch) -> None:
    monkeypatch.setenv("SANDBOXER_BWRAP_BIN", "/custom/bwrap")
    ext = BwrapExtension({"bwrap_bin": "bwrap"})
    assert ext._bwrap_bin() == "/custom/bwrap"


def test_bwrap_argv_unshares_net_and_binds_workspace(tmp_path) -> None:
    ext = BwrapExtension({"ro_binds": []})
    argv = ext._bwrap_argv(str(tmp_path))
    assert "--unshare-net" in argv
    assert "--unshare-user" in argv
    assert str(tmp_path) in argv
    assert "/run/sandboxer/bwrap_runner.py" in argv
    assert argv[-4:] == [
        "/usr/bin/python3",
        "/run/sandboxer/bwrap_runner.py",
        str(tmp_path),
        f"{tmp_path}/.sandboxer-owner.sock",
    ]


def test_existing_ro_binds_filters_missing_paths(tmp_path) -> None:
    real_dir = tmp_path / "real"
    real_dir.mkdir()
    ext = BwrapExtension({"ro_binds": [str(real_dir), "/definitely/does/not/exist"]})
    assert ext._existing_ro_binds() == [str(real_dir)]


def test_provision_spawns_bwrap_and_returns_handle(tmp_path) -> None:
    ext = BwrapExtension({"base_dir": str(tmp_path)})
    fake_proc = MagicMock()
    fake_proc.pid = 12345

    with (
        patch("sandboxer.extensions.bwrap.subprocess.Popen", return_value=fake_proc) as popen,
        patch.object(BwrapExtension, "_read_child_pid", return_value=23456),
    ):
        handle = ext.provision(_profile(), {"sandbox_id": "abc12345"}, "localhost")

    popen.assert_called_once()
    assert handle["sandbox_id"] == "abc12345"
    assert handle["host"] == "localhost"
    assert handle["pid"] == "23456"
    assert handle["supervisor_pid"] == "12345"
    assert handle["workspace_dir"].endswith("abc12345")


def test_provision_copies_repo_into_workspace(tmp_path) -> None:
    repo = tmp_path / "repo"
    repo.mkdir()
    (repo / "file.txt").write_text("hello")
    base_dir = tmp_path / "sandboxes"

    ext = BwrapExtension({"base_dir": str(base_dir)})
    fake_proc = MagicMock()
    fake_proc.pid = 1

    with (
        patch("sandboxer.extensions.bwrap.subprocess.Popen", return_value=fake_proc),
        patch.object(BwrapExtension, "_read_child_pid", return_value=2),
    ):
        handle = ext.provision(
            _profile(), {"sandbox_id": "cafe1234", "repo": str(repo)}, "localhost"
        )

    assert (base_dir / "cafe1234" / "file.txt").read_text() == "hello"
    assert handle["workspace_dir"] == str(base_dir / "cafe1234")
    assert stat.S_IMODE((base_dir / "cafe1234").stat().st_mode) == 0o700


def test_provision_missing_repo_raises(tmp_path) -> None:
    ext = BwrapExtension({"base_dir": str(tmp_path)})
    with pytest.raises(FileNotFoundError):
        ext.provision(
            _profile(),
            {"sandbox_id": "x", "repo": str(tmp_path / "missing")},
            "localhost",
        )


def test_wait_ready_checks_process_and_workspace(tmp_path) -> None:
    workspace = tmp_path / "ws"
    workspace.mkdir()
    ext = BwrapExtension()
    handle = {"pid": str(1), "workspace_dir": str(workspace), "host": "localhost"}

    with (
        patch.object(BwrapExtension, "_pid_alive", return_value=True),
        patch("sandboxer.extensions.bwrap.Path.is_socket", return_value=True),
    ):
        result = ext.wait_ready(handle)

    assert result["host"] == "localhost"
    assert result["endpoint"] == "pid:1"


def test_wait_ready_raises_if_process_dead(tmp_path) -> None:
    workspace = tmp_path / "ws"
    workspace.mkdir()
    ext = BwrapExtension()
    handle = {"pid": "999999", "workspace_dir": str(workspace), "host": "localhost"}

    with (
        patch.object(BwrapExtension, "_pid_alive", return_value=False),
        pytest.raises(RuntimeError, match="not running"),
    ):
        ext.wait_ready(handle)


def test_wait_ready_raises_if_workspace_missing(tmp_path) -> None:
    ext = BwrapExtension()
    handle = {"pid": "1", "workspace_dir": str(tmp_path / "gone"), "host": "localhost"}

    with (
        patch.object(BwrapExtension, "_pid_alive", return_value=True),
        pytest.raises(RuntimeError, match="workspace missing"),
    ):
        ext.wait_ready(handle)


def test_teardown_kills_process_group_and_removes_workspace(tmp_path) -> None:
    workspace = tmp_path / "ws"
    workspace.mkdir()
    ext = BwrapExtension()
    handle = {"pid": "42", "workspace_dir": str(workspace), "host": "localhost"}

    with (
        patch.object(BwrapExtension, "_pid_alive", return_value=True),
        patch("sandboxer.extensions.bwrap.os.getpgid", return_value=42),
        patch("sandboxer.extensions.bwrap.os.killpg") as killpg,
    ):
        result = ext.teardown(handle)

    killpg.assert_called_once_with(42, __import__("signal").SIGKILL)
    assert result["process_killed"] == "True"
    assert result["workspace_removed"] == "True"
    assert not workspace.exists()


def test_teardown_handles_already_dead_process(tmp_path) -> None:
    workspace = tmp_path / "ws"
    workspace.mkdir()
    ext = BwrapExtension()
    handle = {"pid": "42", "workspace_dir": str(workspace), "host": "localhost"}

    with patch.object(BwrapExtension, "_pid_alive", return_value=False):
        result = ext.teardown(handle)

    assert result["process_killed"] == "False"
    assert result["workspace_removed"] == "True"


def test_supports_snapshots_is_false() -> None:
    ext = BwrapExtension()
    assert ext.supports_snapshots() is False
    with pytest.raises(NotImplementedError):
        ext.snapshot({})


def test_reachability_enriches_bwrap_handle_without_direct_exec_hint() -> None:
    handle = {"pid": "555", "workspace_dir": "/tmp/sandboxer-bwrap/abc", "host": "localhost"}
    reach = {"host": "localhost", "endpoint": "pid:555"}
    profile = _profile()

    enriched = enrich_reachability(reach, profile, handle)
    reachability = Reachability(**enriched)

    assert reachability.pid == "555"
    assert reachability.workspace_dir == "/tmp/sandboxer-bwrap/abc"


def test_execute_sends_bounded_request_to_in_namespace_owner(tmp_path) -> None:
    base_dir = tmp_path / "sandboxes"
    workspace = base_dir / "abc12345"
    workspace.mkdir(parents=True)
    ext = BwrapExtension({"base_dir": str(base_dir)})
    handle = {
        "sandbox_id": "abc12345",
        "pid": "42",
        "workspace_dir": str(workspace),
    }
    control = workspace / ".sandboxer-owner.sock"
    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind(str(control))
    server.listen(1)
    received = {}

    def serve_once() -> None:
        connection, _ = server.accept()
        with connection:
            raw = b""
            while chunk := connection.recv(65_536):
                raw += chunk
            received.update(json.loads(raw))
            connection.sendall(
                json.dumps(
                    {
                        "exit_code": 0,
                        "timed_out": False,
                        "stdout": "ok\n",
                        "stderr": "",
                        "output_truncated": False,
                        "workspace_dir": str(workspace),
                    }
                ).encode()
            )

    thread = threading.Thread(target=serve_once)
    thread.start()

    with patch.object(BwrapExtension, "_pid_alive", return_value=True):
        result = ext.execute(
            handle,
            ["python3", "-V"],
            credential_route_refs=["rein-openweights-openrouter-approle"],
            execution_context={"actor": "agt", "run_id": "run-1"},
            timeout_seconds=30,
            max_output_bytes=1024,
        )
    thread.join(timeout=2)
    server.close()

    assert received["command"] == ["python3", "-V"]
    assert received["execution_context"] == {"actor": "agt", "run_id": "run-1"}
    assert received["credential_route_refs"] == ["rein-openweights-openrouter-approle"]
    assert received["timeout_seconds"] == 30
    assert received["max_output_bytes"] == 1024
    assert result["stdout"] == "ok\n"
    assert result["exit_code"] == 0


def test_in_namespace_runner_uses_sanitized_environment(tmp_path) -> None:
    process = MagicMock()
    process.returncode = 0
    process.communicate.return_value = (b"ok\n", b"")
    payload = {
        "command": ["python3", "-V"],
        "credential_route_refs": ["rein-openweights-openrouter-approle"],
        "execution_context": {"actor": "agt", "run_id": "run-1"},
        "timeout_seconds": 30,
        "max_output_bytes": 1024,
    }
    with patch(
        "sandboxer.extensions.bwrap_runner.subprocess.Popen", return_value=process
    ) as popen:
        result = _run(payload, tmp_path)

    assert popen.call_args.args[0] == ["python3", "-V"]
    assert popen.call_args.kwargs["cwd"] == tmp_path
    child_env = popen.call_args.kwargs["env"]
    assert child_env["SANDBOXER_ACTOR"] == "agt"
    assert child_env["SANDBOXER_RUN_ID"] == "run-1"
    assert "rein-openweights-openrouter-approle" in child_env[
        "SANDBOXER_CREDENTIAL_ROUTE_REFS"
    ]
    assert set(child_env) == {
        "HOME",
        "LANG",
        "PATH",
        "SANDBOXER_CREDENTIAL_ROUTE_REFS",
        "SANDBOXER_ACTOR",
        "SANDBOXER_RUN_ID",
    }
    assert result["stdout"] == "ok\n"


def test_execute_refuses_workspace_outside_owner_base(tmp_path) -> None:
    outside = tmp_path / "source-checkout"
    outside.mkdir()
    ext = BwrapExtension({"base_dir": str(tmp_path / "sandboxes")})
    handle = {"sandbox_id": "abc12345", "pid": "42", "workspace_dir": str(outside)}

    with (
        patch.object(BwrapExtension, "_pid_alive", return_value=True),
        pytest.raises(RuntimeError, match="outside owner-managed"),
    ):
        ext.execute(
            handle,
            ["true"],
            credential_route_refs=[],
            execution_context={},
            timeout_seconds=30,
            max_output_bytes=1024,
        )


def test_runner_bounds_output_and_normalizes_timeout(tmp_path) -> None:
    process = MagicMock()
    process.pid = 99
    process.communicate.side_effect = [
        subprocess.TimeoutExpired("cmd", 1),
        (b"123456", b"abcdef"),
    ]
    payload = {
        "command": ["sleep", "2"],
        "credential_route_refs": [],
        "execution_context": {},
        "timeout_seconds": 1,
        "max_output_bytes": 3,
    }

    with (
        patch("sandboxer.extensions.bwrap_runner.subprocess.Popen", return_value=process),
        patch("sandboxer.extensions.bwrap_runner.os.killpg") as killpg,
    ):
        result = _run(payload, tmp_path)

    killpg.assert_called_once_with(99, signal.SIGKILL)
    assert result["exit_code"] == 124
    assert result["timed_out"] is True
    assert result["stdout"] == "123"
    assert result["stderr"] == "abc"
    assert result["output_truncated"] is True