From ace42bbd7e0289312cf682ecad2eeeef225d97b5 Mon Sep 17 00:00:00 2001 From: fliaping Date: Tue, 13 Oct 2026 17:16:18 +0800 Subject: [PATCH 1/4] feat: per-sandbox cgroup v2 cpu/memory/pids limits, rlimit fallback tied to isolation level --- docs/CONFIGURATION.md | 6 + docs/ISOLATION.md | 74 +++- .../src/agent_sandbox_runtime/bubblewrap.py | 3 + runtime/src/agent_sandbox_runtime/profile.py | 3 + scripts/integration-test.sh | 3 +- src/agent_sandbox/cgroups.py | 331 ++++++++++++++++++ src/agent_sandbox/config.py | 23 ++ src/agent_sandbox/console.py | 4 +- src/agent_sandbox/runtime.py | 177 +++++++++- src/agent_sandbox/schemas.py | 3 + tests/conftest.py | 14 + tests/integration/test_resource_limits.py | 137 ++++++++ tests/test_cgroups.py | 63 ++++ tests/test_template_security.py | 75 +++- 14 files changed, 872 insertions(+), 44 deletions(-) create mode 100644 src/agent_sandbox/cgroups.py create mode 100644 tests/conftest.py create mode 100644 tests/integration/test_resource_limits.py create mode 100644 tests/test_cgroups.py diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index 56bc442..ee421d1 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -69,6 +69,12 @@ profile identity. | `SANDBOX_TMPFS_BYTES` | `536870912` (512 MiB) | Size of the sandbox's private tmpfs. | | `SANDBOX_CPU_SECONDS` | `300` | CPU seconds per process (`RLIMIT_CPU`), not a shared sandbox CPU quota or wall-clock deadline. | | `SANDBOX_MAX_PROCESSES` | `64` | `RLIMIT_NPROC` counts processes/threads for the real UID, subject to OS privilege semantics; not an execution-scoped cgroup PID budget. See `SANDBOX_TOOLCHAINS` for Go's compiler limits. | +| `SANDBOX_CGROUP_MODE` | `auto` | Per-sandbox cgroup v2 limits: `auto` uses them when the worker can manage a delegated cgroup and otherwise falls back to rlimits, `required` refuses to start without them, `off` never uses them. The `strict` level is only selectable with cgroups active. See [ISOLATION.md](ISOLATION.md#resource-limits-by-level). | +| `SANDBOX_CPU_QUOTA` | `2.0` | CPU cores per sandbox (`cpu.max`), shared by every command in it. `0` leaves CPU unlimited. cgroup scheme only. | +| `SANDBOX_MEMORY_LIMIT_BYTES` | `2147483648` | Memory per sandbox (`memory.max`, with `memory.swap.max=0`); a command over it is OOM-killed as a group and returns `OOM_KILLED` / `SANDBOX_MEMORY_LIMIT_EXCEEDED`. `0` leaves it unlimited. cgroup scheme only. | +| `SANDBOX_PIDS_LIMIT` | `512` | Processes and threads per sandbox (`pids.max`). `0` leaves it unlimited. cgroup scheme only. | +| `SANDBOX_WORKER_MEMORY_RESERVE_BYTES` | `268435456` | Memory kept for the worker's control process: set as `memory.min` on its leaf and subtracted from the sandboxes' aggregate `memory.max`. cgroup scheme only. | +| `SANDBOX_ADDRESS_SPACE_LIMIT_BYTES` | `8589934592` | Fallback per-process `RLIMIT_AS`, applied only when cgroups are not active. Larger than the memory limit because Node and the JVM reserve address space they never use. `0` disables it. | | `SANDBOX_MAX_OPEN_FILES` | `512` | Open descriptors per process (`RLIMIT_NOFILE`), not per sandbox. | | `SANDBOX_MAX_FILE_SIZE_BYTES` | `1073741824` (1 GiB) | Largest individual file (`RLIMIT_FSIZE`), not an aggregate disk quota. | | `SANDBOX_READONLY_MOUNTS` | `/usr,/bin,/sbin,/lib,/lib64,/etc` | Host trees mounted read-only into every sandbox. | diff --git a/docs/ISOLATION.md b/docs/ISOLATION.md index b74d315..7991250 100644 --- a/docs/ISOLATION.md +++ b/docs/ISOLATION.md @@ -151,12 +151,74 @@ Choose the Level and other policies for the workload: | Agents needing process privacy or reliable PID-namespace teardown | Require `standard` or above; keep network and resource policy choices explicit. | | Hostile public tenants | A namespace Level alone, including `strict`, is insufficient. Require suitable enforced aggregate budgets/process supervision and a network boundary (no network or mandatory selective enforcement), or a stronger execution backend such as a VM. | -`/healthz` reports `resource_control` separately from `cgroup_namespace`: -the built-in backend advertises configured per-process/UID/file limits, but -`sandbox_cpu_quota`, `sandbox_memory_quota`, `sandbox_pid_quota`, -`sandbox_disk_quota`, and `execution_cgroup_supervision` are all `false`. -An enforcing plugin must report its actual controls, not infer them from -the Level's name. See [Hardening](HARDENING.md) for conditional enhancements. +`/healthz` reports `resource_control` separately from `cgroup_namespace`. +It states the scheme actually enforced (`cgroup_v2` or `rlimit`), and +`sandbox_cpu_quota`, `sandbox_memory_quota`, `sandbox_pid_quota` and +`execution_cgroup_supervision` are `true` only when the worker created and +manages the cgroups. `sandbox_disk_quota` is always `false`. An enforcing +plugin must report its actual controls, not infer them from the Level's name. +See [Hardening](HARDENING.md) for conditional enhancements. + +## Resource limits by level + +The resource scheme follows the isolation level and what the outer runtime +permits. It is negotiated at startup like the level itself, and `/healthz` +reports the result under `resource_control`. + +| Level | cgroup v2 manageable (default `SANDBOX_CGROUP_MODE=auto`) | No cgroup v2 (or `SANDBOX_CGROUP_MODE=off`) | +| --- | --- | --- | +| `basic` | **cgroup scheme**: per sandbox `cpu.max` (`SANDBOX_CPU_QUOTA`, 2 cores), `memory.max` (`SANDBOX_MEMORY_LIMIT_BYTES`, 2 GiB) with `memory.swap.max=0`, `pids.max` (`SANDBOX_PIDS_LIMIT`, 512); one leaf per execution with `memory.oom.group=1`, killed and removed when the execution ends; worker reserve. Plus every rlimit except `RLIMIT_AS`. | **rlimit scheme**: per process `RLIMIT_AS` (`SANDBOX_ADDRESS_SPACE_LIMIT_BYTES`, 8 GiB), `RLIMIT_CPU` (`SANDBOX_CPU_SECONDS`), per UID `RLIMIT_NPROC` (`SANDBOX_MAX_PROCESSES`), `RLIMIT_NOFILE`, `RLIMIT_FSIZE`, no core dumps; the per-execution wall-clock timeout; leftover processes of the sandbox UID killed after the last execution ends. No aggregate CPU or memory bound. | +| `standard` | Same cgroup scheme. | Same rlimit scheme. The PID namespace additionally ends every process of an execution with it. | +| `strict` | Same cgroup scheme, plus the cgroup namespace. | **Not available.** `strict` promises aggregate limits, so `auto` selects `standard` and an explicit `SANDBOX_ISOLATION_LEVEL=strict` refuses to start (`resource limits: …` in `probe_failures`). | + +`SANDBOX_CGROUP_MODE=required` refuses to start at any level without cgroups +(`SANDBOX_CGROUP_UNAVAILABLE`). Use it, or `strict`, wherever untrusted code +runs: the rlimit scheme bounds each process, not the sandbox, so N processes +can still use N × `RLIMIT_AS` and every core. + +What a runaway command gets: + +| Event | cgroup scheme | rlimit scheme | +| --- | --- | --- | +| Memory over the limit | The execution's process tree is OOM-killed together; the response is `status=OOM_KILLED`, `error_code=SANDBOX_MEMORY_LIMIT_EXCEEDED`. Other sandboxes and the worker are unaffected. | The allocating process gets `ENOMEM` (`MemoryError` in Python); `status=FAILED`. | +| Fork bomb | `pids.max` stops it; `error_code=SANDBOX_PIDS_LIMIT_REACHED` when the limit was hit. | `RLIMIT_NPROC` per sandbox UID stops it; the timeout ends it. | +| CPU spin on many cores | Throttled to `SANDBOX_CPU_QUOTA` cores in total. | Not bounded beyond `RLIMIT_CPU` per process and the timeout. | +| Leftover background process | Killed with the execution's cgroup (`cgroup.kill`). | Killed when no other execution of the sandbox runs. | + +Release and suspend kill the sandbox cgroup and remove it; a dormant sandbox +holds none. The worker moves itself into a `worker` leaf with `memory.min` set +to `SANDBOX_WORKER_MEMORY_RESERVE_BYTES` (256 MiB), and caps the aggregate of +all sandboxes at the container's memory limit minus that reserve, so the +container never reaches its own limit because of a sandbox. + +### Outer runtime requirements for the cgroup scheme + +The worker needs a cgroup v2 hierarchy, the `cpu`, `memory` and `pids` +controllers delegated to its own cgroup, and a private cgroup namespace: + +- **Docker / Compose**: cgroup v2 hosts give containers a private cgroup + namespace by default (`cgroupns: private`). Docker mounts it read-only; the + worker then mounts a fresh cgroup2 from a user namespace it creates + (`unshare --user --map-root-user --mount --cgroup`), which the kernel makes + read-write and rooted at the container's own cgroup. This needs nothing + beyond the `seccomp=unconfined` the Compose files already set: no + capabilities, no privileged mode, no host cgroup mount. Set a container + memory limit (`mem_limit` / `--memory`) so the reserve has a ceiling to + subtract from. +- **Hosts that restrict unprivileged user namespaces** (Ubuntu 23.10+ AppArmor, + `kernel.apparmor_restrict_unprivileged_userns=1`) block that mount as they + block Bubblewrap; the same override from the AppArmor section above + applies. Otherwise `auto` falls back to the rlimit scheme and says why in + `resource_control.cgroup.detail`. +- **cgroup v1 hosts, `cgroupns: host`** without a writable delegated subtree, + and runtimes that forbid user namespaces: rlimit scheme only. +- **Kubernetes**: the same mechanism works where the pod may create user + namespaces (the seccomp profile allows `unshare`/`mount`); give the pod a + memory limit. Verify with `/healthz` on the target cluster. + +The sandbox UIDs cannot change their own limits: cgroup files are owned by the +worker's root, and the helper mount is not visible inside the Bubblewrap mount +namespace. ## Selection diff --git a/runtime/src/agent_sandbox_runtime/bubblewrap.py b/runtime/src/agent_sandbox_runtime/bubblewrap.py index a837702..a6f2655 100644 --- a/runtime/src/agent_sandbox_runtime/bubblewrap.py +++ b/runtime/src/agent_sandbox_runtime/bubblewrap.py @@ -141,6 +141,9 @@ def build(self, command: SandboxCommand) -> list[str]: args.append(f"--cpu={profile.cpu_seconds}:{profile.cpu_seconds}") if profile.max_processes is not None: args.append(f"--nproc={profile.max_processes}:{profile.max_processes}") + if profile.max_address_space_bytes is not None: + limit = profile.max_address_space_bytes + args.append(f"--as={limit}:{limit}") args.extend( [ f"--nofile={profile.max_open_files}:{profile.max_open_files}", diff --git a/runtime/src/agent_sandbox_runtime/profile.py b/runtime/src/agent_sandbox_runtime/profile.py index 5091667..1f7e471 100644 --- a/runtime/src/agent_sandbox_runtime/profile.py +++ b/runtime/src/agent_sandbox_runtime/profile.py @@ -28,6 +28,7 @@ class SandboxRuntimeProfile: max_processes: int | None = 64 max_open_files: int = 512 max_file_size_bytes: int = 1024 * 1024 * 1024 + max_address_space_bytes: int | None = None readonly_mounts: tuple[str, ...] = ( "/usr", "/bin", @@ -56,6 +57,8 @@ def validate(self) -> None: raise ValueError(f"Sandbox {name} must be positive") if self.cpu_seconds is not None and self.cpu_seconds < 1: raise ValueError("Sandbox cpu_seconds must be positive") + if self.max_address_space_bytes is not None and self.max_address_space_bytes < 1: + raise ValueError("Sandbox max_address_space_bytes must be positive") if self.max_processes is not None and self.max_processes < 1: raise ValueError("Sandbox max_processes must be positive") diff --git a/scripts/integration-test.sh b/scripts/integration-test.sh index 528a587..24dd91e 100755 --- a/scripts/integration-test.sh +++ b/scripts/integration-test.sh @@ -73,7 +73,8 @@ run_sandbox_tests() { "${TEST_IMAGE}" \ uv run --no-sync pytest \ tests/integration/test_bubblewrap_isolation.py \ - tests/integration/test_template_isolation.py -v "$@" + tests/integration/test_template_isolation.py \ + tests/integration/test_resource_limits.py -v "$@" } run_polyglot_tests() { diff --git a/src/agent_sandbox/cgroups.py b/src/agent_sandbox/cgroups.py new file mode 100644 index 0000000..d883636 --- /dev/null +++ b/src/agent_sandbox/cgroups.py @@ -0,0 +1,331 @@ +"""Per-sandbox resource limits on a delegated cgroup v2 subtree. + +The worker owns the cgroup it starts in. At startup it moves its own +processes into a ``worker`` leaf (cgroup v2 forbids processes in a cgroup +that delegates controllers), enables ``cpu``, ``memory`` and ``pids`` for +its children, and creates ``sandboxes``. Each sandbox gets a child of +``sandboxes`` carrying ``cpu.max``, ``memory.max``, ``memory.swap.max``, +``memory.oom.group`` and ``pids.max``; each execution runs in its own leaf +under the sandbox, so the sandbox limit is shared by every command in it and +one execution can be killed and accounted for on its own. + +Nothing here needs privileges beyond a writable cgroup2 mount for the +worker's own cgroup, which Docker provides with ``--cgroupns=private`` and a +read-write ``/sys/fs/cgroup`` (see docs/ISOLATION.md). +""" + +from __future__ import annotations + +import contextlib +import errno +import logging +import os +import re +import shutil +import subprocess +import time +from dataclasses import dataclass, field +from pathlib import Path + +logger = logging.getLogger(__name__) + +CONTROLLERS = ("cpu", "memory", "pids") +_CPU_PERIOD_US = 100_000 +_SAFE_NAME = re.compile(r"[^A-Za-z0-9_.-]") + + +class CgroupUnavailable(RuntimeError): + """The worker cannot manage a delegated cgroup v2 subtree.""" + + +@dataclass(frozen=True, slots=True) +class CgroupLimits: + cpu_cores: float | None + memory_bytes: int | None + pids: int | None + worker_memory_reserve_bytes: int = 0 + + +def _write(path: Path, value: str) -> None: + with open(path, "w") as stream: + stream.write(value) + + +def _read(path: Path) -> str: + try: + return path.read_text().strip() + except OSError: + return "" + + +def _own_cgroup(mount: Path) -> Path: + for line in Path("/proc/self/cgroup").read_text().splitlines(): + hierarchy, _, rest = line.partition(":") + controllers, _, relative = rest.partition(":") + if hierarchy == "0" and controllers == "": + path = mount / relative.lstrip("/") + if path.is_dir(): + return path + # A host-namespace path that is not visible here (no private + # cgroupns): the mount root is then our own cgroup only if we + # were started with a private namespace, so refuse. + raise CgroupUnavailable(f"own cgroup {relative!r} is not visible under {mount}") + raise CgroupUnavailable("no cgroup v2 entry in /proc/self/cgroup") + + +def _name(value: str) -> str: + return _SAFE_NAME.sub("_", value)[:200] or "_" + + +@dataclass +class CgroupManager: + """Owns ``/sandboxes//``. Threads may call it.""" + + limits: CgroupLimits + mount: Path = Path("/sys/fs/cgroup") + base: Path | None = None + active: bool = False + detail: str = "not initialized" + _created: set[Path] = field(default_factory=set) + _helper: subprocess.Popen[bytes] | None = None + helper_mount_point: Path = Path("/run/sandloom-cgroup") + + # ── setup ── + + def setup(self) -> None: + try: + self._setup() + self.active = True + self.detail = f"cgroup v2 at {self.base}" + except (OSError, CgroupUnavailable) as exc: + self.active = False + self.detail = f"cgroup v2 delegation unavailable: {exc}" + raise CgroupUnavailable(self.detail) from exc + + def _setup(self) -> None: + if not (self.mount / "cgroup.controllers").exists(): + raise CgroupUnavailable(f"{self.mount} is not a cgroup2 mount") + try: + base: Path | None = _own_cgroup(self.mount) + except CgroupUnavailable: + base = None + if ( + base is None + or _read_only(self.mount) + or not os.access(base / "cgroup.subtree_control", os.W_OK) + ): + # Docker and Kubernetes mount the container's cgroup read-only even + # with a private cgroup namespace. A user namespace we create may + # mount a fresh cgroup2 rooted at our own cgroup, read-write. The + # helper process holds that mount; we use it through its /proc root. + base = self._mount_through_helper() + self._configure(base) + + def _mount_through_helper(self) -> Path: + unshare = shutil.which("unshare") + if unshare is None: + raise CgroupUnavailable("cgroup is read-only and unshare(1) is not installed") + self.helper_mount_point.mkdir(parents=True, exist_ok=True) + helper = subprocess.Popen( + [ + unshare, + "--user", + "--map-root-user", + "--mount", + "--cgroup", + "--propagation", + "private", + "/bin/sh", + "-c", + 'mount -t cgroup2 none "$0" && echo ready && exec sleep infinity', + str(self.helper_mount_point), + ], + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + assert helper.stdout is not None + line = helper.stdout.readline() + if line.strip() != b"ready": + helper.kill() + _, err = helper.communicate(timeout=5) + raise CgroupUnavailable( + "cgroup is read-only and a user-namespace cgroup2 mount failed: " + + err.decode(errors="replace").strip()[:300] + ) + self._helper = helper + base = Path(f"/proc/{helper.pid}/root") / self.helper_mount_point.relative_to("/") + if not os.access(base / "cgroup.subtree_control", os.W_OK): + raise CgroupUnavailable(f"helper cgroup mount {base} is not writable") + return base + + def _configure(self, base: Path) -> None: + available = set(_read(base / "cgroup.controllers").split()) + missing = [name for name in CONTROLLERS if name not in available] + if missing: + raise CgroupUnavailable(f"controllers {missing} not delegated to {base}") + worker = base / "worker" + worker.mkdir(exist_ok=True) + # Move every process of our cgroup, not only this one: an entrypoint + # shell or init may share it, and one leftover blocks the delegation. + for _ in range(10): + pids = _read(base / "cgroup.procs").split() + if not pids: + break + for pid in pids: + with contextlib.suppress(ProcessLookupError, OSError): + _write(worker / "cgroup.procs", pid) + _write(base / "cgroup.subtree_control", " ".join(f"+{c}" for c in CONTROLLERS)) + sandboxes = base / "sandboxes" + sandboxes.mkdir(exist_ok=True) + _write(sandboxes / "cgroup.subtree_control", " ".join(f"+{c}" for c in CONTROLLERS)) + reserve = self.limits.worker_memory_reserve_bytes + if reserve: + # Protect the control process from reclaim, and keep the sandboxes + # aggregate below the worker's own limit by the reserve, so a + # sandbox OOM is local and the worker never reaches its limit. + with contextlib.suppress(OSError): + _write(worker / "memory.min", str(reserve)) + total = _effective_memory_max(base) + if total is not None: + _write(sandboxes / "memory.max", str(max(total - reserve, 64 * 1024 * 1024))) + _write(sandboxes / "memory.oom.group", "0") + self.base = base + # Leftovers of a previous worker process in this container. + for stale in sandboxes.iterdir(): + if stale.is_dir(): + self._kill_tree(stale) + + # ── per sandbox and per execution ── + + def _sandbox_dir(self, sandbox_id: str) -> Path: + assert self.base is not None + return self.base / "sandboxes" / _name(sandbox_id) + + def ensure_sandbox(self, sandbox_id: str) -> Path: + path = self._sandbox_dir(sandbox_id) + if not path.is_dir(): + path.mkdir(exist_ok=True) + limits = self.limits + if limits.cpu_cores: + quota = max(1000, int(limits.cpu_cores * _CPU_PERIOD_US)) + _write(path / "cpu.max", f"{quota} {_CPU_PERIOD_US}") + if limits.memory_bytes: + _write(path / "memory.max", str(limits.memory_bytes)) + # Without this, a sandbox at its limit swaps instead of + # failing, and drags the whole host down with it. + with contextlib.suppress(OSError): + _write(path / "memory.swap.max", "0") + if limits.pids: + _write(path / "pids.max", str(limits.pids)) + _write(path / "cgroup.subtree_control", " ".join(f"+{c}" for c in CONTROLLERS)) + return path + + def create_execution(self, sandbox_id: str, exec_id: str) -> Path: + parent = self.ensure_sandbox(sandbox_id) + leaf = parent / f"exec-{_name(exec_id)}-{time.monotonic_ns()}" + leaf.mkdir() + # One OOM kill takes the whole command tree, not a random child that + # leaves a half-dead build behind. + with contextlib.suppress(OSError): + _write(leaf / "memory.oom.group", "1") + return leaf + + def execution_stats(self, leaf: Path) -> dict[str, int]: + stats: dict[str, int] = {} + for name in ("memory.events", "pids.events"): + for line in _read(leaf / name).splitlines(): + key, _, value = line.partition(" ") + with contextlib.suppress(ValueError): + stats[f"{name.split('.')[0]}.{key}"] = int(value) + # A sandbox-level OOM (sum of executions over memory.max) is + # recorded on the sandbox cgroup, which is the one at its limit. + parent = leaf.parent + for line in _read(parent / "memory.events.local").splitlines(): + key, _, value = line.partition(" ") + if key == "oom_kill": + with contextlib.suppress(ValueError): + stats["sandbox.oom_kill"] = int(value) + return stats + + def remove_execution(self, leaf: Path) -> None: + self._kill_tree(leaf) + + def remove_sandbox(self, sandbox_id: str) -> None: + if self.active: + self._kill_tree(self._sandbox_dir(sandbox_id)) + + def sandbox_processes(self, sandbox_id: str) -> int: + if not self.active: + return 0 + path = self._sandbox_dir(sandbox_id) + return int(_read(path / "pids.current") or 0) if path.is_dir() else 0 + + def _kill_tree(self, path: Path) -> None: + if not path.is_dir(): + return + for child in [p for p in path.iterdir() if p.is_dir()]: + self._kill_tree(child) + with contextlib.suppress(OSError): + _write(path / "cgroup.kill", "1") + for _ in range(100): + try: + path.rmdir() + return + except OSError as exc: + if exc.errno == errno.ENOENT: + return + if exc.errno != errno.EBUSY: + raise + time.sleep(0.02) + logger.warning("cgroup %s still busy after kill", path) + + def close(self) -> None: + """Stop the mount helper; cgroups it created stay valid until removed.""" + helper, self._helper = self._helper, None + if helper is not None: + helper.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + helper.wait(timeout=5) + + def report(self) -> dict[str, object]: + return { + "active": self.active, + "detail": self.detail, + "path": str(self.base) if self.base else None, + } + + +def _read_only(mount: Path) -> bool: + try: + return bool(os.statvfs(mount).f_flag & os.ST_RDONLY) + except OSError: + return True + + +def _effective_memory_max(path: Path) -> int | None: + """The tightest memory.max from ``path`` up through visible ancestors.""" + limit: int | None = None + current = path + while (current / "memory.max").exists(): + value = _read(current / "memory.max") + if value and value != "max": + limit = int(value) if limit is None else min(limit, int(value)) + if current.parent == current: + break + current = current.parent + if limit is None: + try: + limit = int(os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE")) + except (OSError, ValueError): + return None + return limit + + +def placement_prefix(leaf: Path) -> list[str]: + """argv that joins ``leaf`` and then execs the rest of the command. + + The process enters the cgroup before prlimit/setpriv/bwrap run, so every + descendant inherits it and none can escape by double-forking. + """ + return ["/bin/sh", "-c", 'echo $$ > "$0" && exec "$@"', str(leaf / "cgroup.procs")] diff --git a/src/agent_sandbox/config.py b/src/agent_sandbox/config.py index 3d5337c..de29f1f 100644 --- a/src/agent_sandbox/config.py +++ b/src/agent_sandbox/config.py @@ -101,6 +101,21 @@ class Settings(BaseSettings): tmpfs_bytes: int = 512 * 1024 * 1024 cpu_seconds: int = 300 max_processes: int = 64 + # Per-sandbox resource limits. With a delegated cgroup v2 subtree these are + # aggregate cgroup limits over every command in the sandbox; without one, + # memory falls back to a per-process RLIMIT_AS (see docs/ISOLATION.md). + # `auto` uses cgroups when the worker can manage them, `required` refuses + # to start without them, `off` never uses them. The strict isolation level + # always needs them. + cgroup_mode: str = "auto" + cpu_quota: float = Field(default=2.0, ge=0) + memory_limit_bytes: int = Field(default=2 * 1024 * 1024 * 1024, ge=0) + pids_limit: int = Field(default=512, ge=0) + worker_memory_reserve_bytes: int = Field(default=256 * 1024 * 1024, ge=0) + # Per-process virtual memory cap used only when cgroups are not active. + # Larger than the memory limit because runtimes such as Node and the JVM + # reserve address space they never touch. 0 disables it. + address_space_limit_bytes: int = Field(default=8 * 1024 * 1024 * 1024, ge=0) max_open_files: int = 512 max_file_size_bytes: int = 1024 * 1024 * 1024 max_output_bytes: int = 8 * 1024 * 1024 @@ -251,6 +266,14 @@ def validate_allowed_literals(cls, value: list[str]) -> list[str]: ipaddress.ip_address(item.strip()) return [item.strip() for item in value if item.strip()] + @field_validator("cgroup_mode") + @classmethod + def validate_cgroup_mode(cls, value: str) -> str: + normalized = value.strip().lower() + if normalized not in {"auto", "required", "off"}: + raise ValueError("SANDBOX_CGROUP_MODE must be auto, required or off") + return normalized + @field_validator("isolation_level") @classmethod def validate_isolation_level(cls, value: str) -> str: diff --git a/src/agent_sandbox/console.py b/src/agent_sandbox/console.py index 2f38644..a850c95 100644 --- a/src/agent_sandbox/console.py +++ b/src/agent_sandbox/console.py @@ -337,7 +337,7 @@ const STATUS_TONE = { READY: "ok", ACTIVE: "ok", RUNNING: "run", ASSIGNED: "warn", RELEASING: "warn", RELEASED: "off", SUSPENDING: "warn", SUSPENDED: "off", SUCCEEDED: "ok", - FAILED: "bad", TIMED_OUT: "bad", CANCELLED: "off", LOST: "bad", DRAINING: "warn", + FAILED: "bad", TIMED_OUT: "bad", OOM_KILLED: "bad", CANCELLED: "off", LOST: "bad", DRAINING: "warn", }; const status = (value) => el("span", { class: "st " + (STATUS_TONE[value] || "") }, value || "—"); @@ -580,7 +580,7 @@ }), el("select", { onchange: (e) => { S.filters.status = e.target.value; S.page.offset = 0; load(); }, - }, ["", "RUNNING", "SUCCEEDED", "FAILED", "TIMED_OUT", "CANCELLED"].map((v) => + }, ["", "RUNNING", "SUCCEEDED", "FAILED", "TIMED_OUT", "OOM_KILLED", "CANCELLED"].map((v) => el("option", { value: v, selected: S.filters.status === v }, v || "All statuses"))), ), executionDetail(), diff --git a/src/agent_sandbox/runtime.py b/src/agent_sandbox/runtime.py index aeb4dc3..8f67659 100644 --- a/src/agent_sandbox/runtime.py +++ b/src/agent_sandbox/runtime.py @@ -1,4 +1,4 @@ -"""Isolated UID + prlimit + Bubblewrap execution layer.""" +"""Isolated UID + cgroup v2/prlimit + Bubblewrap execution layer.""" from __future__ import annotations @@ -11,6 +11,7 @@ import os import shutil import signal +import subprocess import sys import time import uuid @@ -39,6 +40,7 @@ ) from .archive import ArchiveLimits +from .cgroups import CgroupLimits, CgroupManager, CgroupUnavailable, placement_prefix from .config import Settings from .dormant import restore_snapshot, write_snapshot from .isolation_policy import negotiate_features @@ -126,8 +128,10 @@ def __init__( isolation_level: IsolationLevel = IsolationLevel.BASIC, additional_features: tuple[str, ...] = (), extension: RuntimeExtension | None = None, + rlimit_memory: bool = True, ) -> None: self.settings = settings + self.rlimit_memory = rlimit_memory self.disable_nested_userns = disable_nested_userns self.isolation_level = isolation_level self.extension = extension or RuntimeExtension() @@ -136,6 +140,7 @@ def __init__( disable_nested_userns=disable_nested_userns, isolation_level=isolation_level, additional_features=additional_features, + rlimit_memory=rlimit_memory, ) self.profile = profile self._shared = InteractiveSandboxCommandBuilder( @@ -257,8 +262,34 @@ def __init__( self._destroy_tasks: dict[str, asyncio.Task[None]] = {} # Suspended sandboxes. Not in `sandboxes`, so they hold no capacity slot. self.dormant: dict[str, DormantSandbox] = {} + self.cgroups = CgroupManager( + CgroupLimits( + cpu_cores=settings.cpu_quota or None, + memory_bytes=settings.memory_limit_bytes or None, + pids=settings.pids_limit or None, + worker_memory_reserve_bytes=settings.worker_memory_reserve_bytes, + ) + ) + self.cgroups.detail = "disabled by SANDBOX_CGROUP_MODE=off" + self._exec_cgroups: dict[tuple[str, str], Path] = {} + + def _setup_cgroups(self) -> None: + if self.settings.cgroup_mode == "off": + return + try: + self.cgroups.setup() + except CgroupUnavailable as exc: + if self.settings.cgroup_mode == "required": + raise RuntimeError(f"SANDBOX_CGROUP_UNAVAILABLE: {exc}") from exc + logger.warning("%s; falling back to per-process rlimits", exc) + + @property + def cgroup_active(self) -> bool: + return self.cgroups.active async def probe(self) -> dict[str, object]: + await asyncio.to_thread(self._setup_cgroups) + rlimit_memory = not self.cgroup_active runtime_report = await asyncio.to_thread(self.builder.probe) probe_root = self.settings.workspace_root / f".probe-{uuid.uuid4().hex}" probe_id = probe_root.name @@ -291,6 +322,7 @@ async def probe(self) -> dict[str, object]: disable_nested_userns=disable_nested_userns, isolation_level=level, extension=self.extension, + rlimit_memory=rlimit_memory, ) ok, detail = await self._probe_candidate( candidate, @@ -304,15 +336,18 @@ async def probe(self) -> dict[str, object]: failures.append(f"disable_nested_userns={disable_nested_userns}: {detail}") else: results[level] = (False, " | ".join(failures)) + if level.label == "strict" and results[level][0] and not self.cgroup_active: + # strict promises aggregate resource limits as well as a + # cgroup namespace; without delegation it cannot keep that. + results[level] = (False, f"resource limits: {self.cgroups.detail}") + candidates.pop(level, None) selection = negotiate_isolation( self.settings.isolation_level, lambda level: results[level], ) self.builder = candidates[selection.selected] - feature_candidates: dict[tuple[str, ...], BubblewrapCommandBuilder] = { - (): self.builder - } + feature_candidates: dict[tuple[str, ...], BubblewrapCommandBuilder] = {(): self.builder} async def probe_features(features: tuple[str, ...]) -> tuple[bool, str]: failures = [] @@ -323,6 +358,7 @@ async def probe_features(features: tuple[str, ...]) -> tuple[bool, str]: additional_features=features, disable_nested_userns=disable_nested, extension=self.extension, + rlimit_memory=rlimit_memory, ) ok, detail = await self._probe_candidate( candidate, @@ -357,6 +393,7 @@ async def launch(request: ExecRequest) -> ExecResponse: finally: self.sandboxes.pop(probe_id, None) self.last_active_at.pop(probe_id, None) + await asyncio.to_thread(self.cgroups.remove_sandbox, probe_id) await asyncio.to_thread(shutil.rmtree, probe_root, True) toolchains = self.builder.toolchains return { @@ -380,9 +417,7 @@ async def launch(request: ExecRequest) -> ExecResponse: for name in toolchains_unavailable_at(actual_features, toolchains) if checks.get(name, {}).get("status") in {"failed", "timed_out"} ), - "isolation_hints": list( - toolchains_unavailable_at(actual_features, toolchains) - ), + "isolation_hints": list(toolchains_unavailable_at(actual_features, toolchains)), "checks": checks, "available": [ name for name, check in checks.items() if check["status"] == "available" @@ -405,16 +440,40 @@ async def launch(request: ExecRequest) -> ExecResponse: "open_files_per_process": self.settings.max_open_files, "bytes_per_file": self.settings.max_file_size_bytes, }, - "sandbox_cpu_quota": False, - "sandbox_memory_quota": False, - "sandbox_pid_quota": False, + **self._resource_report(selection.selected.label), "sandbox_disk_quota": False, - "execution_cgroup_supervision": False, }, "landlock": False, "extension": dict(self.extension.capability_labels), } + def _resource_report(self, level: str) -> dict[str, object]: + cg = self.cgroup_active + settings = self.settings + rlimit_as = (not cg) and bool(settings.address_space_limit_bytes) + return { + "scheme": "cgroup_v2" if cg else "rlimit", + "isolation_level": level, + "cgroup": {**self.cgroups.report(), "mode": settings.cgroup_mode}, + "sandbox_cpu_quota": cg and bool(settings.cpu_quota), + "sandbox_memory_quota": cg and bool(settings.memory_limit_bytes), + "sandbox_pid_quota": cg and bool(settings.pids_limit), + "execution_cgroup_supervision": cg, + "limits": { + "cpu_cores_per_sandbox": settings.cpu_quota if cg else None, + "memory_bytes_per_sandbox": settings.memory_limit_bytes if cg else None, + "swap_bytes_per_sandbox": 0 if cg and settings.memory_limit_bytes else None, + "pids_per_sandbox": settings.pids_limit if cg else None, + "address_space_bytes_per_process": ( + settings.address_space_limit_bytes if rlimit_as else None + ), + "wall_clock_seconds_per_exec_default": settings.default_timeout_seconds, + "worker_memory_reserve_bytes": ( + settings.worker_memory_reserve_bytes if cg else None + ), + }, + } + async def _probe_candidate( self, builder: BubblewrapCommandBuilder, @@ -561,14 +620,27 @@ async def _execute_locked(self, sandbox: LocalSandbox, request: ExecRequest) -> self._check_not_destroying(sandbox.sandbox_id) self.get(sandbox.sandbox_id, sandbox.generation) command = self.builder.build(sandbox, request, template_mounts=template_mounts) + leaf: Path | None = None + if self.cgroup_active: + leaf = await asyncio.to_thread( + self.cgroups.create_execution, sandbox.sandbox_id, request.exec_id + ) + command = placement_prefix(leaf) + command started = time.monotonic() - process = await asyncio.create_subprocess_exec( - *command, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, - ) + try: + process = await asyncio.create_subprocess_exec( + *command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + start_new_session=True, + ) + except BaseException: + if leaf is not None: + await asyncio.to_thread(self.cgroups.remove_execution, leaf) + raise self.processes[(sandbox.sandbox_id, request.exec_id)] = process + if leaf is not None: + self._exec_cgroups[(sandbox.sandbox_id, request.exec_id)] = leaf timeout = request.timeout_seconds or self.settings.default_timeout_seconds status = "SUCCEEDED" assert process.stdout is not None and process.stderr is not None @@ -614,8 +686,28 @@ async def _execute_locked(self, sandbox: LocalSandbox, request: ExecRequest) -> raise finally: self.processes.pop((sandbox.sandbox_id, request.exec_id), None) + stats: dict[str, int] = {} + if leaf is not None: + self._exec_cgroups.pop((sandbox.sandbox_id, request.exec_id), None) + stats = await asyncio.to_thread(self.cgroups.execution_stats, leaf) + # Kills any descendant that outlived the command, then removes + # the leaf: nothing a command starts survives its execution. + await asyncio.to_thread(self.cgroups.remove_execution, leaf) + elif not self._has_processes(sandbox.sandbox_id): + # Without a cgroup, a descendant that left the process group + # (a fork bomb's survivors, a daemonized server) keeps the + # sandbox UID's RLIMIT_NPROC budget. Nothing else of this + # sandbox runs, so every process of its UID is such a leftover. + await asyncio.to_thread(self._kill_uid, sandbox.uid) stdout = stdout_b.decode(errors="replace") stderr = stderr_b.decode(errors="replace") + error_code = None + if stats.get("memory.oom_kill") or stats.get("memory.oom_group_kill"): + status = "OOM_KILLED" + error_code = "SANDBOX_MEMORY_LIMIT_EXCEEDED" + elif stats.get("pids.max"): + # The command hit pids.max at least once: forks failed with EAGAIN. + error_code = "SANDBOX_PIDS_LIMIT_REACHED" if status == "SUCCEEDED" and process.returncode != 0: status = "FAILED" self.touch(sandbox.sandbox_id) @@ -627,6 +719,7 @@ async def _execute_locked(self, sandbox: LocalSandbox, request: ExecRequest) -> stderr=stderr, duration_ms=int((time.monotonic() - started) * 1000), truncated=stdout_cut or stderr_cut, + error_code=error_code, ) def supports_templates(self) -> bool: @@ -757,6 +850,9 @@ async def cancel(self, sandbox_id: str, exec_id: str) -> bool: if process is None: return False await self._terminate(process) + leaf = self._exec_cgroups.get((sandbox_id, exec_id)) + if leaf is not None: + await asyncio.to_thread(self.cgroups.remove_execution, leaf) return True async def destroy(self, sandbox_id: str, *, delete_files: bool = True) -> None: @@ -775,6 +871,10 @@ async def _destroy(self, sandbox_id: str, *, delete_files: bool) -> None: for (sid, _), process in list(self.processes.items()): if sid == sandbox_id: await self._terminate(process) + await asyncio.to_thread(self.cgroups.remove_sandbox, sandbox_id) + sandbox = self.sandboxes.get(sandbox_id) + if sandbox is not None and not self.cgroup_active: + await asyncio.to_thread(self._kill_uid, sandbox.uid) await complete_in_thread(self._destroy_files, sandbox_id, delete_files) finally: self._destroying.discard(sandbox_id) @@ -870,6 +970,8 @@ def _suspend_sandbox( ) self.sandboxes.pop(sandbox_id, None) self.last_active_at.pop(sandbox_id, None) + # A dormant sandbox runs nothing and holds no cgroup. + self.cgroups.remove_sandbox(sandbox_id) self.dormant[sandbox_id] = DormantSandbox( sandbox_id=sandbox_id, generation=generation, @@ -1069,8 +1171,40 @@ def _reclaim_orphan_workspaces( continue return removed + def _kill_uid(self, uid: int) -> None: + """SIGKILL every process of a sandbox UID, as that UID. + + kill(-1) from the UID itself reaches exactly its own processes, which + no PID scan can do race-free against a process that keeps forking. + """ + if os.geteuid() != 0 or uid == os.geteuid(): + return + for _ in range(3): + with contextlib.suppress(OSError, subprocess.SubprocessError): + subprocess.run( + [ + str(self.settings.setpriv_path), + "--reuid", + str(uid), + "--regid", + str(uid), + "--clear-groups", + "--", + "/bin/kill", + "-KILL", + "-1", + ], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=10, + check=False, + ) + def _has_processes(self, sandbox_id: str) -> bool: - return any(sid == sandbox_id for sid, _ in self.processes) + return any(sid == sandbox_id for sid, _ in self.processes) or bool( + self.cgroups.sandbox_processes(sandbox_id) + ) async def reclaim_dormant_disk(self) -> list[str]: """Evict local copies of snapshotted dormant sandboxes under disk pressure. @@ -1163,6 +1297,9 @@ async def shutdown(self) -> None: for process in list(self.processes.values()): await self._terminate(process) self.processes.clear() + for sandbox_id in list(self.sandboxes): + await asyncio.to_thread(self.cgroups.remove_sandbox, sandbox_id) + self.cgroups.close() async def write_file(self, sandbox: LocalSandbox, path: str, content_base64: str) -> None: content = base64.b64decode(content_base64, validate=True) @@ -1333,6 +1470,7 @@ def _runtime_profile( disable_nested_userns: bool, isolation_level: IsolationLevel = IsolationLevel.BASIC, additional_features: tuple[str, ...] = (), + rlimit_memory: bool = True, ) -> SandboxRuntimeProfile: return SandboxRuntimeProfile( profile_id=settings.profile_id, @@ -1350,6 +1488,9 @@ def _runtime_profile( max_processes=settings.max_processes, max_open_files=settings.max_open_files, max_file_size_bytes=settings.max_file_size_bytes, + max_address_space_bytes=( + settings.address_space_limit_bytes or None if rlimit_memory else None + ), readonly_mounts=tuple(settings.readonly_mounts), ) diff --git a/src/agent_sandbox/schemas.py b/src/agent_sandbox/schemas.py index f1e9adf..05d23d4 100644 --- a/src/agent_sandbox/schemas.py +++ b/src/agent_sandbox/schemas.py @@ -172,6 +172,9 @@ class ExecResponse(BaseModel): stderr: str = "" duration_ms: int | None = None truncated: bool = False + # Set when a resource limit ended the command, e.g. + # SANDBOX_MEMORY_LIMIT_EXCEEDED or SANDBOX_PIDS_LIMIT_REACHED. + error_code: str | None = None class FileWriteRequest(BaseModel): diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..7fe9c65 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,14 @@ +import os + +import pytest + + +@pytest.fixture(autouse=True) +def _no_host_cgroups(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> None: + """Unit tests must not create cgroups on the machine running them. + + The integration suite runs inside the worker image and exercises the real + delegation; it keeps whatever SANDBOX_CGROUP_MODE its container sets. + """ + if "integration" not in str(request.node.fspath) and "SANDBOX_CGROUP_MODE" not in os.environ: + monkeypatch.setenv("SANDBOX_CGROUP_MODE", "off") diff --git a/tests/integration/test_resource_limits.py b/tests/integration/test_resource_limits.py new file mode 100644 index 0000000..8bd9aef --- /dev/null +++ b/tests/integration/test_resource_limits.py @@ -0,0 +1,137 @@ +"""Resource limits, verified with real runaway commands in the worker image. + +The scheme depends on the container: with a manageable cgroup v2 subtree the +sandbox gets aggregate cpu/memory/pids limits, otherwise per-process rlimits. +Each test asserts the behaviour of whichever scheme `/healthz` reports, so a +report that claims a quota it does not enforce fails here. + + ./scripts/integration-test.sh sandbox +""" + +from __future__ import annotations + +import os +import shutil +import sys +import tempfile +import time +from pathlib import Path +from typing import Any + +import pytest + +from agent_sandbox.config import Settings +from agent_sandbox.runtime import SandboxRuntime +from agent_sandbox.schemas import ExecRequest + +pytestmark = [ + pytest.mark.skipif(sys.platform != "linux", reason="Bubblewrap requires Linux"), + pytest.mark.skipif(shutil.which("bwrap") is None, reason="bwrap is not installed"), + pytest.mark.skipif(os.geteuid() != 0, reason="sandbox UIDs and cgroups need root"), +] + +MEMORY_LIMIT = 256 * 1024 * 1024 + + +@pytest.fixture +async def runtime(tmp_path: Path) -> Any: + local_root = tmp_path / "sandboxes" + local_root.mkdir() + for ancestor in (local_root, *local_root.parents): + os.chmod(ancestor, os.stat(ancestor).st_mode | 0o011) + if ancestor == Path(tempfile.gettempdir()): + break + instance = SandboxRuntime( + Settings( + internal_token="integration-token", + local_root=local_root, + database_url=f"sqlite+aiosqlite:///{tmp_path}/control.db", + uid_start=20100, + uid_end=20199, + cpu_quota=1, + memory_limit_bytes=MEMORY_LIMIT, + pids_limit=64, + address_space_limit_bytes=2 * MEMORY_LIMIT, + min_free_bytes=0, + disk_high_watermark_percent=99, + ) + ) + report = await instance.probe() + instance.probe_report = report # type: ignore[attr-defined] + try: + yield instance + finally: + await instance.shutdown() + + +async def _run(runtime: SandboxRuntime, sandbox: Any, script: str, limit: int = 30) -> Any: + return await runtime.execute( + sandbox, + ExecRequest( + exec_id=f"exec-{os.urandom(4).hex()}", + generation=sandbox.generation, + argv=["/bin/bash", "-c", script], + timeout_seconds=limit, + ), + ) + + +def _scheme(runtime: SandboxRuntime) -> str: + return runtime.probe_report["resource_control"]["scheme"] # type: ignore[attr-defined] + + +async def test_report_matches_the_active_scheme(runtime: SandboxRuntime) -> None: + control = runtime.probe_report["resource_control"] # type: ignore[attr-defined] + cgroup = control["scheme"] == "cgroup_v2" + assert control["sandbox_memory_quota"] is cgroup + assert control["sandbox_cpu_quota"] is cgroup + assert control["sandbox_pid_quota"] is cgroup + if not cgroup: + assert runtime.probe_report["selected_level"] != "strict" # type: ignore[attr-defined] + assert control["limits"]["address_space_bytes_per_process"] == 2 * MEMORY_LIMIT + + +async def test_memory_bomb_is_stopped(runtime: SandboxRuntime) -> None: + sandbox = await runtime.create("sandbox-membomb", 1, 20101) + result = await _run( + runtime, + sandbox, + "python3 -c 'b=[]\nwhile True: b.append(bytearray(16*1024*1024))'", + ) + if _scheme(runtime) == "cgroup_v2": + assert result.status == "OOM_KILLED" + assert result.error_code == "SANDBOX_MEMORY_LIMIT_EXCEEDED" + else: + assert result.status == "FAILED" + assert "MemoryError" in result.stderr + # The sandbox is usable afterwards. + assert (await _run(runtime, sandbox, "echo ok")).stdout.strip() == "ok" + + +async def test_fork_bomb_is_bounded_and_leaves_nothing(runtime: SandboxRuntime) -> None: + sandbox = await runtime.create("sandbox-forkbomb", 1, 20102) + other = await runtime.create("sandbox-neighbour", 1, 20103) + started = time.monotonic() + await _run(runtime, sandbox, ":(){ :|:& };: ; sleep 3", limit=6) + assert time.monotonic() - started < 30 + assert (await _run(runtime, other, "echo alive")).stdout.strip() == "alive" + # Survivors would hold the sandbox's process budget; none may remain. + assert (await _run(runtime, sandbox, "echo again")).stdout.strip() == "again" + + +async def test_cpu_quota_caps_a_multicore_spin(runtime: SandboxRuntime) -> None: + if _scheme(runtime) != "cgroup_v2": + pytest.skip("no CPU quota without cgroups; RLIMIT_CPU is per process") + sandbox = await runtime.create("sandbox-spin", 1, 20104) + result = await _run( + runtime, + sandbox, + "python3 - <<'P'\nimport multiprocessing as m, os, time\n" + "def f():\n t=time.time()\n while time.time()-t<3: pass\n" + "ps=[m.Process(target=f) for _ in range(4)]\n[p.start() for p in ps]\n" + "[p.join() for p in ps]\nc=os.times()\nprint(c.children_user+c.children_system)\nP", + limit=60, + ) + assert result.status == "SUCCEEDED", result.stderr + # Four spinners for 3 s wall under a one-core quota: about 3 s of CPU. + assert float(result.stdout.strip().splitlines()[-1]) < 3 * 1.5 diff --git a/tests/test_cgroups.py b/tests/test_cgroups.py new file mode 100644 index 0000000..f0ade29 --- /dev/null +++ b/tests/test_cgroups.py @@ -0,0 +1,63 @@ +from pathlib import Path + +import pytest + +from agent_sandbox.cgroups import ( + CgroupLimits, + CgroupManager, + CgroupUnavailable, + placement_prefix, +) + + +def _manager(tmp_path: Path, **limits: object) -> CgroupManager: + base = tmp_path / "base" + (base / "sandboxes").mkdir(parents=True) + defaults = {"cpu_cores": 1.5, "memory_bytes": 1024**3, "pids": 128} + defaults.update(limits) + manager = CgroupManager(CgroupLimits(**defaults), mount=tmp_path) # type: ignore[arg-type] + manager.base = base + manager.active = True + return manager + + +def test_sandbox_cgroup_carries_every_limit(tmp_path: Path) -> None: + manager = _manager(tmp_path) + path = manager.ensure_sandbox("sb/1") + assert path.name == "sb_1" + assert (path / "cpu.max").read_text() == "150000 100000" + assert (path / "memory.max").read_text() == str(1024**3) + assert (path / "memory.swap.max").read_text() == "0" + assert (path / "pids.max").read_text() == "128" + assert (path / "cgroup.subtree_control").read_text() == "+cpu +memory +pids" + + +def test_unset_limits_are_not_written(tmp_path: Path) -> None: + manager = _manager(tmp_path, cpu_cores=None, memory_bytes=None, pids=None) + path = manager.ensure_sandbox("sb") + assert not (path / "cpu.max").exists() + assert not (path / "memory.max").exists() + assert not (path / "pids.max").exists() + + +def test_execution_leaf_kills_as_a_group(tmp_path: Path) -> None: + manager = _manager(tmp_path) + leaf = manager.create_execution("sb", "exec-1") + assert leaf.parent.name == "sb" + assert (leaf / "memory.oom.group").read_text() == "1" + (leaf / "memory.events").write_text("low 0\nhigh 0\nmax 3\noom 1\noom_kill 1\n") + assert manager.execution_stats(leaf)["memory.oom_kill"] == 1 + + +def test_placement_joins_the_leaf_before_exec(tmp_path: Path) -> None: + prefix = placement_prefix(tmp_path / "leaf") + assert prefix[:2] == ["/bin/sh", "-c"] + assert prefix[-1] == str(tmp_path / "leaf" / "cgroup.procs") + + +def test_setup_refuses_a_non_cgroup2_mount(tmp_path: Path) -> None: + manager = CgroupManager(CgroupLimits(1, 1, 1), mount=tmp_path) + with pytest.raises(CgroupUnavailable): + manager.setup() + assert manager.active is False + assert "unavailable" in manager.detail diff --git a/tests/test_template_security.py b/tests/test_template_security.py index 343521f..497bb7b 100644 --- a/tests/test_template_security.py +++ b/tests/test_template_security.py @@ -474,10 +474,7 @@ async def spawn(*args: Any, **kwargs: Any) -> asyncio.subprocess.Process: await destruction -async def test_strict_does_not_advertise_aggregate_quotas( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - runtime, _ = await runtime_and_sandbox(tmp_path) +def _fake_probe(runtime: SandboxRuntime, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(runtime.builder, "probe", lambda: {}) async def supported(*args: Any) -> tuple[bool, str]: @@ -489,21 +486,65 @@ async def measured_launch(*args: Any) -> dict[str, object]: return {} monkeypatch.setattr("agent_sandbox.runtime.probe_toolchains", measured_launch) + + +async def test_without_cgroups_strict_is_unavailable_and_rlimits_apply( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + runtime, _ = await runtime_and_sandbox(tmp_path) + _fake_probe(runtime, monkeypatch) + monkeypatch.setattr(runtime, "_setup_cgroups", lambda: None) + report = await runtime.probe() + # strict promises aggregate limits, so it is not selected without them. + assert report["selected_level"] == "standard" + assert "resource limits" in report["probe_failures"]["strict"] + control = report["resource_control"] + assert control["scheme"] == "rlimit" + assert control["sandbox_cpu_quota"] is False + assert control["sandbox_memory_quota"] is False + assert control["sandbox_pid_quota"] is False + assert control["execution_cgroup_supervision"] is False + assert control["limits"]["address_space_bytes_per_process"] == 8 * 1024**3 + assert control["limits"]["wall_clock_seconds_per_exec_default"] == 300 + assert control["process_limits"]["processes_per_real_uid"] == 64 + assert runtime.builder.profile.max_address_space_bytes == 8 * 1024**3 + + +async def test_with_cgroups_strict_reports_aggregate_quotas( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + runtime, _ = await runtime_and_sandbox(tmp_path) + _fake_probe(runtime, monkeypatch) + + def active() -> None: + runtime.cgroups.active = True + + monkeypatch.setattr(runtime, "_setup_cgroups", active) + monkeypatch.setattr(runtime.cgroups, "remove_sandbox", lambda sandbox_id: None) report = await runtime.probe() + assert report["selected_level"] == "strict" assert report["cgroup_namespace"] is True - assert report["resource_control"] == { - "process_limits": { - "cpu_seconds_per_process": 300, - "processes_per_real_uid": 64, - "open_files_per_process": 512, - "bytes_per_file": 1024 * 1024 * 1024, - }, - "sandbox_cpu_quota": False, - "sandbox_memory_quota": False, - "sandbox_pid_quota": False, - "sandbox_disk_quota": False, - "execution_cgroup_supervision": False, - } + control = report["resource_control"] + assert control["scheme"] == "cgroup_v2" + assert control["sandbox_cpu_quota"] is True + assert control["sandbox_memory_quota"] is True + assert control["sandbox_pid_quota"] is True + assert control["sandbox_disk_quota"] is False + assert control["limits"]["memory_bytes_per_sandbox"] == 2 * 1024**3 + assert control["limits"]["address_space_bytes_per_process"] is None + # The aggregate cgroup limit replaces RLIMIT_AS, which breaks JIT runtimes. + assert runtime.builder.profile.max_address_space_bytes is None + + +async def test_required_cgroup_mode_refuses_to_start_without_delegation( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + runtime = SandboxRuntime( + Settings(internal_token="token", local_root=tmp_path, cgroup_mode="required") + ) + runtime.cgroups.mount = tmp_path / "not-a-cgroup" + with pytest.raises(RuntimeError, match="SANDBOX_CGROUP_UNAVAILABLE"): + await runtime.probe() @pytest.mark.parametrize( From 42d49a7e48450e75205c28bd051bc42436b50fde Mon Sep 17 00:00:00 2001 From: fliaping <269832017@qq.com> Date: Tue, 13 Oct 2026 17:21:57 +0800 Subject: [PATCH 2/4] ci: run the resource-limit integration tests in the sandbox boundary job --- .github/workflows/ci.yml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8efabe4..a2d9ee9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -57,7 +57,8 @@ jobs: - name: Test adapters against real servers run: | uv run pytest tests/integration \ - --ignore=tests/integration/test_bubblewrap_isolation.py -q + --ignore=tests/integration/test_bubblewrap_isolation.py \ + --ignore=tests/integration/test_resource_limits.py -q - name: Dump middleware logs on failure if: failure() run: docker compose -f compose.integration.yaml logs --tail 100 @@ -99,7 +100,8 @@ jobs: agent-sandbox-test:latest \ uv run --no-sync pytest \ tests/integration/test_bubblewrap_isolation.py \ - tests/integration/test_template_isolation.py -v + tests/integration/test_template_isolation.py \ + tests/integration/test_resource_limits.py -v -rs fleet: name: Fleet (two replicas) From b4c073053ab4d1c95245c97d7bcc5585245c5ee1 Mon Sep 17 00:00:00 2001 From: fliaping <269832017@qq.com> Date: Tue, 13 Oct 2026 18:44:24 +0800 Subject: [PATCH 3/4] feat: aggregate limits without cgroups (supervisor), per-sandbox disk quota, io limits --- docs/CONFIGURATION.md | 8 + docs/ISOLATION.md | 87 +++-- src/agent_sandbox/app.py | 1 + src/agent_sandbox/cgroups.py | 94 +++++- src/agent_sandbox/config.py | 16 + src/agent_sandbox/runtime.py | 126 ++++++- src/agent_sandbox/supervisor.py | 390 ++++++++++++++++++++++ tests/integration/test_resource_limits.py | 63 +++- tests/test_supervisor.py | 126 +++++++ tests/test_template_security.py | 22 +- 10 files changed, 877 insertions(+), 56 deletions(-) create mode 100644 src/agent_sandbox/supervisor.py create mode 100644 tests/test_supervisor.py diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index ee421d1..b03b2c0 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -75,6 +75,14 @@ profile identity. | `SANDBOX_PIDS_LIMIT` | `512` | Processes and threads per sandbox (`pids.max`). `0` leaves it unlimited. cgroup scheme only. | | `SANDBOX_WORKER_MEMORY_RESERVE_BYTES` | `268435456` | Memory kept for the worker's control process: set as `memory.min` on its leaf and subtracted from the sandboxes' aggregate `memory.max`. cgroup scheme only. | | `SANDBOX_ADDRESS_SPACE_LIMIT_BYTES` | `8589934592` | Fallback per-process `RLIMIT_AS`, applied only when cgroups are not active. Larger than the memory limit because Node and the JVM reserve address space they never use. `0` disables it. | +| `SANDBOX_SUPERVISOR_INTERVAL_SECONDS` | `0.1` | Without cgroups, how often the supervisor samples `/proc` to enforce `SANDBOX_MEMORY_LIMIT_BYTES` (sum of PSS) and `SANDBOX_CPU_QUOTA` (SIGSTOP/SIGCONT duty cycle) over all processes of the sandbox UID. Shorter is tighter and costs more worker CPU. | +| `SANDBOX_DISK_LIMIT_BYTES` | `10737418240` | Allocated bytes per sandbox (workspace, home, cache, envs), every scheme. Growth past it kills running executions with `SANDBOX_DISK_QUOTA_EXCEEDED`; File API writes over it are refused (HTTP 507). Deleting is always allowed. `0` disables it. | +| `SANDBOX_DISK_CHECK_INTERVAL_SECONDS` | `2.0` | How often the disk usage of a sandbox with a running execution is measured. | +| `SANDBOX_IO_READ_BPS` | `268435456` | `io.max rbps` per sandbox on the device holding `SANDBOX_LOCAL_ROOT` (cgroup scheme). `0` is unlimited. | +| `SANDBOX_IO_WRITE_BPS` | `268435456` | `io.max wbps`, as above. | +| `SANDBOX_IO_READ_IOPS` | `0` | `io.max riops`, as above. | +| `SANDBOX_IO_WRITE_IOPS` | `0` | `io.max wiops`, as above. | +| `SANDBOX_IO_DEVICE` | `""` | `major:minor` for `io.max` when it cannot be derived from `SANDBOX_LOCAL_ROOT` (overlay, network filesystems). | | `SANDBOX_MAX_OPEN_FILES` | `512` | Open descriptors per process (`RLIMIT_NOFILE`), not per sandbox. | | `SANDBOX_MAX_FILE_SIZE_BYTES` | `1073741824` (1 GiB) | Largest individual file (`RLIMIT_FSIZE`), not an aggregate disk quota. | | `SANDBOX_READONLY_MOUNTS` | `/usr,/bin,/sbin,/lib,/lib64,/etc` | Host trees mounted read-only into every sandbox. | diff --git a/docs/ISOLATION.md b/docs/ISOLATION.md index 7991250..7098de8 100644 --- a/docs/ISOLATION.md +++ b/docs/ISOLATION.md @@ -161,35 +161,59 @@ See [Hardening](HARDENING.md) for conditional enhancements. ## Resource limits by level -The resource scheme follows the isolation level and what the outer runtime -permits. It is negotiated at startup like the level itself, and `/healthz` -reports the result under `resource_control`. - -| Level | cgroup v2 manageable (default `SANDBOX_CGROUP_MODE=auto`) | No cgroup v2 (or `SANDBOX_CGROUP_MODE=off`) | +Every sandbox is limited in aggregate, over all of its processes. How a limit +is enforced depends on the isolation level and on what the outer runtime +permits. The worker negotiates this at startup, the same way it negotiates the +level, and `/healthz` reports the result under `resource_control`. `scheme` +is either `cgroup_v2` or `rlimit+supervisor`. `enforcement` says how each +dimension is enforced: by the kernel, by the supervisor (with its precision), +or not at all. + +| Level | cgroup v2 manageable (default `SANDBOX_CGROUP_MODE=auto`) | No cgroup v2: cgroup v1 host, no user namespaces, or `SANDBOX_CGROUP_MODE=off` | | --- | --- | --- | -| `basic` | **cgroup scheme**: per sandbox `cpu.max` (`SANDBOX_CPU_QUOTA`, 2 cores), `memory.max` (`SANDBOX_MEMORY_LIMIT_BYTES`, 2 GiB) with `memory.swap.max=0`, `pids.max` (`SANDBOX_PIDS_LIMIT`, 512); one leaf per execution with `memory.oom.group=1`, killed and removed when the execution ends; worker reserve. Plus every rlimit except `RLIMIT_AS`. | **rlimit scheme**: per process `RLIMIT_AS` (`SANDBOX_ADDRESS_SPACE_LIMIT_BYTES`, 8 GiB), `RLIMIT_CPU` (`SANDBOX_CPU_SECONDS`), per UID `RLIMIT_NPROC` (`SANDBOX_MAX_PROCESSES`), `RLIMIT_NOFILE`, `RLIMIT_FSIZE`, no core dumps; the per-execution wall-clock timeout; leftover processes of the sandbox UID killed after the last execution ends. No aggregate CPU or memory bound. | -| `standard` | Same cgroup scheme. | Same rlimit scheme. The PID namespace additionally ends every process of an execution with it. | -| `strict` | Same cgroup scheme, plus the cgroup namespace. | **Not available.** `strict` promises aggregate limits, so `auto` selects `standard` and an explicit `SANDBOX_ISOLATION_LEVEL=strict` refuses to start (`resource limits: …` in `probe_failures`). | - -`SANDBOX_CGROUP_MODE=required` refuses to start at any level without cgroups -(`SANDBOX_CGROUP_UNAVAILABLE`). Use it, or `strict`, wherever untrusted code -runs: the rlimit scheme bounds each process, not the sandbox, so N processes -can still use N × `RLIMIT_AS` and every core. +| `basic` | **cgroup scheme** | **rlimit+supervisor scheme** | +| `standard` | cgroup scheme | rlimit+supervisor scheme. The PID namespace also ends every process of an execution when the execution ends. | +| `strict` | cgroup scheme, plus the cgroup namespace | **Not available.** `strict` promises kernel-enforced limits. `auto` selects `standard` instead; an explicit `SANDBOX_ISOLATION_LEVEL=strict` refuses to start (`resource limits: …` in `probe_failures`). | -What a runaway command gets: +Each dimension under each scheme: -| Event | cgroup scheme | rlimit scheme | +| Dimension (setting, default) | cgroup scheme | rlimit+supervisor scheme | | --- | --- | --- | -| Memory over the limit | The execution's process tree is OOM-killed together; the response is `status=OOM_KILLED`, `error_code=SANDBOX_MEMORY_LIMIT_EXCEEDED`. Other sandboxes and the worker are unaffected. | The allocating process gets `ENOMEM` (`MemoryError` in Python); `status=FAILED`. | -| Fork bomb | `pids.max` stops it; `error_code=SANDBOX_PIDS_LIMIT_REACHED` when the limit was hit. | `RLIMIT_NPROC` per sandbox UID stops it; the timeout ends it. | -| CPU spin on many cores | Throttled to `SANDBOX_CPU_QUOTA` cores in total. | Not bounded beyond `RLIMIT_CPU` per process and the timeout. | -| Leftover background process | Killed with the execution's cgroup (`cgroup.kill`). | Killed when no other execution of the sandbox runs. | +| Memory (`SANDBOX_MEMORY_LIMIT_BYTES`, 2 GiB) | Kernel: `memory.max` per sandbox, `memory.swap.max=0`, `memory.oom.group=1` per execution. | Supervisor: the sum of PSS over the sandbox UID's processes is sampled every `SANDBOX_SUPERVISOR_INTERVAL_SECONDS` (0.1 s). Above the limit, every process of the UID is SIGKILLed. A fast allocation can overshoot for up to one interval, so per-process `RLIMIT_AS` (`SANDBOX_ADDRESS_SPACE_LIMIT_BYTES`, 8 GiB) stays as a hard backstop. | +| CPU (`SANDBOX_CPU_QUOTA`, 2 cores) | Kernel: `cpu.max`. | Supervisor: a token bucket over the UID's CPU time (including reaped children), with up to one CPU-second of burst. When the bucket is empty the processes are SIGSTOPped, then SIGCONTed once it refills. Measured: about 1.2 cores over a 4 s spin at a 1-core quota, because of the burst. | +| Processes (`SANDBOX_PIDS_LIMIT` 512 / `SANDBOX_MAX_PROCESSES` 64) | Kernel: `pids.max`. | Kernel: `RLIMIT_NPROC`. Every sandbox has its own UID, so this already counts the whole sandbox. | +| Disk (`SANDBOX_DISK_LIMIT_BYTES`, 10 GiB) | Supervisor, every `SANDBOX_DISK_CHECK_INTERVAL_SECONDS` (2 s): when allocated blocks under the sandbox root are over the quota and still growing, the running executions are killed. File API writes are refused with HTTP 507. Deleting files is always allowed. `RLIMIT_FSIZE` caps any single file and `/tmp` is a sized tmpfs. | Same. | +| I/O (`SANDBOX_IO_READ_BPS`/`WRITE_BPS`, 256 MiB/s; `…_IOPS`, off) | Kernel: `io.max` on the device holding `SANDBOX_LOCAL_ROOT`, or on `SANDBOX_IO_DEVICE`. Needs a real block device; overlay and tmpfs have none, and `enforcement.io` says so. | Supervisor: storage bytes from `/proc//io`, duty-cycled like CPU. Needs `CAP_SYS_PTRACE` on the worker. Without it, only the lowest I/O priority (`ionice`) is applied, and `enforcement.io` says so. | +| Wall clock (`SANDBOX_DEFAULT_TIMEOUT_SECONDS`, 300) | Per execution. | Per execution. | -Release and suspend kill the sandbox cgroup and remove it; a dormant sandbox -holds none. The worker moves itself into a `worker` leaf with `memory.min` set -to `SANDBOX_WORKER_MEMORY_RESERVE_BYTES` (256 MiB), and caps the aggregate of -all sandboxes at the container's memory limit minus that reserve, so the -container never reaches its own limit because of a sandbox. +What a runaway command gets back: + +| Event | Response | +| --- | --- | +| Memory over the limit | `status=OOM_KILLED`, `error_code=SANDBOX_MEMORY_LIMIT_EXCEEDED`. Under the rlimit+supervisor scheme, a single process can instead hit `RLIMIT_AS` first and get `ENOMEM`. | +| Process limit reached | Forks fail with `EAGAIN`. Under cgroups the response carries `error_code=SANDBOX_PIDS_LIMIT_REACHED`. | +| Disk over the quota | `error_code=SANDBOX_DISK_QUOTA_EXCEEDED`. | +| CPU or I/O over the quota | Throttled, not killed. | +| Leftover background process | Killed when the execution ends (cgroup) or when the sandbox's last execution ends (UID kill). | + +The worker keeps memory for itself: its own leaf has `memory.min` set to +`SANDBOX_WORKER_MEMORY_RESERVE_BYTES`, and the sandboxes' aggregate +`memory.max` is the container limit minus that reserve. Release and suspend +remove the sandbox's cgroup and stop supervising it. + +Some limits cannot be made exact without the kernel: + +- A supervisor samples; it does not intercept. Overshoot is bounded by + allocation rate × interval. Use `SANDBOX_CGROUP_MODE=required` or `strict` + where that bound is not acceptable. +- Disk usage is measured, not reserved. A burst can pass the quota by + whatever it writes within one check interval. True `ENOSPC` per sandbox + needs project quotas (XFS `prjquota`, ext4 `project`) or a per-sandbox + loop filesystem, and both need `CAP_SYS_ADMIN`, which the worker + deliberately does not have. +- Supervised I/O cannot charge a process that starts and exits between two + samples, because Linux keeps no cumulative I/O counters for reaped + children. ### Outer runtime requirements for the cgroup scheme @@ -215,6 +239,21 @@ controllers delegated to its own cgroup, and a private cgroup namespace: - **Kubernetes**: the same mechanism works where the pod may create user namespaces (the seccomp profile allows `unshare`/`mount`); give the pod a memory limit. Verify with `/healthz` on the target cluster. + Verified on kind (Kubernetes 1.31) with this container spec; `emptyDir` sits + on the node's overlay, so `io.max` needs a block-backed volume or + `SANDBOX_IO_DEVICE`: + + ```yaml + securityContext: + runAsUser: 0 + seccompProfile: {type: Unconfined} + appArmorProfile: {type: Unconfined} + resources: + limits: {memory: 4Gi} + ``` +- **Supervised I/O without cgroups** additionally needs `cap_add: [SYS_PTRACE]` + on the worker, to read other UIDs' `/proc//io`. It also lets the + supervisor read PSS instead of RSS. The sandbox UIDs cannot change their own limits: cgroup files are owned by the worker's root, and the helper mount is not visible inside the Bubblewrap mount diff --git a/src/agent_sandbox/app.py b/src/agent_sandbox/app.py index 57c0c4f..308f2d5 100644 --- a/src/agent_sandbox/app.py +++ b/src/agent_sandbox/app.py @@ -166,6 +166,7 @@ def relay_worker_response(response: httpx.Response) -> Response: "SANDBOX_EXEC_SCOPE_BUSY": 409, "SANDBOX_EXEC_SCOPE_LOCKED": 409, "SANDBOX_FILE_PATH_LOCKED": 409, + "SANDBOX_DISK_QUOTA_EXCEEDED": 507, # Suspend and resume. A suspended sandbox refuses work until it is resumed; # a suspend refuses while anything runs in it. Both are the caller's move. "SANDBOX_NOT_FOUND": 404, diff --git a/src/agent_sandbox/cgroups.py b/src/agent_sandbox/cgroups.py index d883636..1967c0d 100644 --- a/src/agent_sandbox/cgroups.py +++ b/src/agent_sandbox/cgroups.py @@ -44,6 +44,25 @@ class CgroupLimits: memory_bytes: int | None pids: int | None worker_memory_reserve_bytes: int = 0 + io_device: str | None = None # "major:minor" of the workspace device + io_rbps: int = 0 + io_wbps: int = 0 + io_riops: int = 0 + io_wiops: int = 0 + + @property + def io_line(self) -> str | None: + parts = [ + f"{key}={value}" + for key, value in ( + ("rbps", self.io_rbps), + ("wbps", self.io_wbps), + ("riops", self.io_riops), + ("wiops", self.io_wiops), + ) + if value + ] + return f"{self.io_device} {' '.join(parts)}" if self.io_device and parts else None def _write(path: Path, value: str) -> None: @@ -88,7 +107,10 @@ class CgroupManager: detail: str = "not initialized" _created: set[Path] = field(default_factory=set) _helper: subprocess.Popen[bytes] | None = None + _controllers: str = "+cpu +memory +pids" helper_mount_point: Path = Path("/run/sandloom-cgroup") + io_active: bool = False + io_detail: str = "not configured" # ── setup ── @@ -175,10 +197,25 @@ def _configure(self, base: Path) -> None: for pid in pids: with contextlib.suppress(ProcessLookupError, OSError): _write(worker / "cgroup.procs", pid) - _write(base / "cgroup.subtree_control", " ".join(f"+{c}" for c in CONTROLLERS)) + controllers = list(CONTROLLERS) + if "io" in available and self.limits.io_line: + controllers.append("io") + enable = " ".join(f"+{c}" for c in controllers) + _write(base / "cgroup.subtree_control", enable) sandboxes = base / "sandboxes" sandboxes.mkdir(exist_ok=True) - _write(sandboxes / "cgroup.subtree_control", " ".join(f"+{c}" for c in CONTROLLERS)) + _write(sandboxes / "cgroup.subtree_control", enable) + self._controllers = enable + limits = self.limits + if not (limits.io_rbps or limits.io_wbps or limits.io_riops or limits.io_wiops): + self.io_detail = "no io limit configured" + elif not limits.io_device: + self.io_detail = ( + "workspace device unresolved (overlay, tmpfs or network filesystem); " + "set SANDBOX_IO_DEVICE" + ) + elif "io" not in available: + self.io_detail = "io controller not delegated" reserve = self.limits.worker_memory_reserve_bytes if reserve: # Protect the control process from reclaim, and keep the sandboxes @@ -218,7 +255,16 @@ def ensure_sandbox(self, sandbox_id: str) -> Path: _write(path / "memory.swap.max", "0") if limits.pids: _write(path / "pids.max", str(limits.pids)) - _write(path / "cgroup.subtree_control", " ".join(f"+{c}" for c in CONTROLLERS)) + io_line = self.limits.io_line + if io_line and "+io" in self._controllers: + try: + _write(path / "io.max", io_line) + self.io_active = True + self.io_detail = f"io.max {io_line}" + except OSError as exc: + self.io_active = False + self.io_detail = f"io.max rejected for {self.limits.io_device}: {exc}" + _write(path / "cgroup.subtree_control", self._controllers) return path def create_execution(self, sandbox_id: str, exec_id: str) -> Path: @@ -322,6 +368,48 @@ def _effective_memory_max(path: Path) -> int | None: return limit +def block_device_of(path: Path) -> str | None: + """ "major:minor" of the block device holding ``path``, for io.max. + + Overlay and tmpfs have anonymous devices (major 0) that io.max rejects; + then the mount source is tried, and a partition maps to its disk. + """ + try: + dev = os.stat(path).st_dev + except OSError: + return None + major, minor = os.major(dev), os.minor(dev) + if major == 0: + source = _mount_source(path) + if source is None or not source.startswith("/dev/"): + return None + try: + rdev = os.stat(source).st_rdev + except OSError: + return None + major, minor = os.major(rdev), os.minor(rdev) + # io.max applies to whole disks; a partition's parent is in sysfs. + sysfs = Path(f"/sys/dev/block/{major}:{minor}") + if (sysfs / "partition").exists(): + with contextlib.suppress(OSError, ValueError): + parent = (sysfs.resolve().parent / "dev").read_text().strip() + return parent + return f"{major}:{minor}" + + +def _mount_source(path: Path) -> str | None: + best: tuple[int, str] | None = None + target = str(path.resolve()) + with contextlib.suppress(OSError): + for line in Path("/proc/self/mountinfo").read_text().splitlines(): + left, _, right = line.partition(" - ") + mount_point = left.split()[4] + if target == mount_point or target.startswith(mount_point.rstrip("/") + "/"): + if best is None or len(mount_point) > best[0]: + best = (len(mount_point), right.split()[1]) + return best[1] if best else None + + def placement_prefix(leaf: Path) -> list[str]: """argv that joins ``leaf`` and then execs the rest of the command. diff --git a/src/agent_sandbox/config.py b/src/agent_sandbox/config.py index de29f1f..d82ae6e 100644 --- a/src/agent_sandbox/config.py +++ b/src/agent_sandbox/config.py @@ -116,6 +116,22 @@ class Settings(BaseSettings): # Larger than the memory limit because runtimes such as Node and the JVM # reserve address space they never touch. 0 disables it. address_space_limit_bytes: int = Field(default=8 * 1024 * 1024 * 1024, ge=0) + # Without cgroups, a supervisor enforces the memory and CPU limits above on + # the sum of each sandbox UID's processes, sampling /proc this often. + supervisor_interval_seconds: float = Field(default=0.1, gt=0, le=5) + # Allocated bytes per sandbox root (workspace, home, cache, envs), every + # scheme. Checked every disk_check_interval_seconds; 0 disables it. + disk_limit_bytes: int = Field(default=10 * 1024 * 1024 * 1024, ge=0) + disk_check_interval_seconds: float = Field(default=2.0, gt=0) + # io.max per sandbox on the device holding SANDBOX_LOCAL_ROOT (cgroup + # scheme). 0 leaves that dimension unlimited. + io_read_bps: int = Field(default=256 * 1024 * 1024, ge=0) + io_write_bps: int = Field(default=256 * 1024 * 1024, ge=0) + io_read_iops: int = Field(default=0, ge=0) + io_write_iops: int = Field(default=0, ge=0) + # "major:minor" for io.max when it cannot be derived from SANDBOX_LOCAL_ROOT + # (an overlay or network filesystem over a known disk). + io_device: str = "" max_open_files: int = 512 max_file_size_bytes: int = 1024 * 1024 * 1024 max_output_bytes: int = 8 * 1024 * 1024 diff --git a/src/agent_sandbox/runtime.py b/src/agent_sandbox/runtime.py index 8f67659..9cf30ef 100644 --- a/src/agent_sandbox/runtime.py +++ b/src/agent_sandbox/runtime.py @@ -40,11 +40,18 @@ ) from .archive import ArchiveLimits -from .cgroups import CgroupLimits, CgroupManager, CgroupUnavailable, placement_prefix +from .cgroups import ( + CgroupLimits, + CgroupManager, + CgroupUnavailable, + block_device_of, + placement_prefix, +) from .config import Settings from .dormant import restore_snapshot, write_snapshot from .isolation_policy import negotiate_features from .schemas import ExecRequest, ExecResponse +from .supervisor import DISK, SandboxSupervisor, can_read_foreign_io, disk_usage from .templates import ( TEMPLATE_MOUNT_ROOT, LocalTemplateCache, @@ -268,8 +275,21 @@ def __init__( memory_bytes=settings.memory_limit_bytes or None, pids=settings.pids_limit or None, worker_memory_reserve_bytes=settings.worker_memory_reserve_bytes, + io_device=settings.io_device or block_device_of(settings.workspace_root), + io_rbps=settings.io_read_bps, + io_wbps=settings.io_write_bps, + io_riops=settings.io_read_iops, + io_wiops=settings.io_write_iops, ) ) + self.supervisor = SandboxSupervisor( + memory_bytes=None, + cpu_cores=None, + disk_bytes=settings.disk_limit_bytes or None, + interval=settings.supervisor_interval_seconds, + disk_interval=settings.disk_check_interval_seconds, + kill_uid=self._kill_uid, + ) self.cgroups.detail = "disabled by SANDBOX_CGROUP_MODE=off" self._exec_cgroups: dict[tuple[str, str], Path] = {} @@ -290,6 +310,18 @@ def cgroup_active(self) -> bool: async def probe(self) -> dict[str, object]: await asyncio.to_thread(self._setup_cgroups) rlimit_memory = not self.cgroup_active + if not self.cgroup_active: + # No kernel aggregate: the supervisor enforces the sandbox totals. + self.supervisor.memory_bytes = self.settings.memory_limit_bytes or None + self.supervisor.cpu_cores = self.settings.cpu_quota or None + io_wanted = bool(self.settings.io_read_bps or self.settings.io_write_bps) + self._io_accounting = io_wanted and await asyncio.to_thread( + can_read_foreign_io, str(self.settings.setpriv_path), self.settings.uid_start + ) + if self._io_accounting: + self.supervisor.io_read_bps = self.settings.io_read_bps or None + self.supervisor.io_write_bps = self.settings.io_write_bps or None + self.supervisor.start() runtime_report = await asyncio.to_thread(self.builder.probe) probe_root = self.settings.workspace_root / f".probe-{uuid.uuid4().hex}" probe_id = probe_root.name @@ -441,7 +473,6 @@ async def launch(request: ExecRequest) -> ExecResponse: "bytes_per_file": self.settings.max_file_size_bytes, }, **self._resource_report(selection.selected.label), - "sandbox_disk_quota": False, }, "landlock": False, "extension": dict(self.extension.capability_labels), @@ -450,20 +481,74 @@ async def launch(request: ExecRequest) -> ExecResponse: def _resource_report(self, level: str) -> dict[str, object]: cg = self.cgroup_active settings = self.settings + sup = self.supervisor rlimit_as = (not cg) and bool(settings.address_space_limit_bytes) + interval_ms = int(settings.supervisor_interval_seconds * 1000) + supervised = ( + f"supervised: /proc sampled every {interval_ms} ms; may overshoot " + "for up to one interval and does not charge processes shorter than it" + ) + + def how(kernel: bool, supervised_on: bool) -> str | None: + return "kernel (cgroup v2)" if kernel else supervised if supervised_on else None + + memory = how(cg and bool(settings.memory_limit_bytes), bool(sup.memory_bytes)) + cpu = how(cg and bool(settings.cpu_quota), bool(sup.cpu_cores)) + pids = ( + "kernel (cgroup v2 pids.max)" + if cg and settings.pids_limit + else "kernel (RLIMIT_NPROC on the sandbox's own UID)" + if settings.max_processes + else None + ) + disk = ( + f"supervised: allocated blocks summed every " + f"{settings.disk_check_interval_seconds:g} s; growth over the quota kills " + "the running executions, File API writes are refused" + if sup.disk_bytes + else None + ) + io = ( + f"kernel ({self.cgroups.io_detail})" + if cg and self.cgroups.io_active + else ( + "supervised: storage read/write bytes (/proc//io) duty-cycled to " + "the bps limits; page-cache writes are charged when written back" + ) + if not cg and (sup.io_read_bps or sup.io_write_bps) + else ( + "none: per-process I/O counters of sandbox UIDs are unreadable " + "(the worker needs CAP_SYS_PTRACE); lowest I/O priority (ionice) only" + ) + if not cg + else f"none: {self.cgroups.io_detail}" + ) return { - "scheme": "cgroup_v2" if cg else "rlimit", + "scheme": "cgroup_v2" if cg else "rlimit+supervisor", "isolation_level": level, "cgroup": {**self.cgroups.report(), "mode": settings.cgroup_mode}, - "sandbox_cpu_quota": cg and bool(settings.cpu_quota), - "sandbox_memory_quota": cg and bool(settings.memory_limit_bytes), - "sandbox_pid_quota": cg and bool(settings.pids_limit), + "supervisor": sup.report() if not cg or sup.disk_bytes else None, + "sandbox_cpu_quota": cpu is not None, + "sandbox_memory_quota": memory is not None, + "sandbox_pid_quota": pids is not None, + "sandbox_disk_quota": disk is not None, + "sandbox_io_limit": (cg and self.cgroups.io_active) + or (not cg and bool(sup.io_read_bps or sup.io_write_bps)), "execution_cgroup_supervision": cg, + "enforcement": { + "cpu": cpu, + "memory": memory, + "pids": pids, + "disk": disk, + "io": io, + }, "limits": { - "cpu_cores_per_sandbox": settings.cpu_quota if cg else None, - "memory_bytes_per_sandbox": settings.memory_limit_bytes if cg else None, + "cpu_cores_per_sandbox": settings.cpu_quota or None, + "memory_bytes_per_sandbox": settings.memory_limit_bytes or None, "swap_bytes_per_sandbox": 0 if cg and settings.memory_limit_bytes else None, - "pids_per_sandbox": settings.pids_limit if cg else None, + "pids_per_sandbox": settings.pids_limit if cg else settings.max_processes, + "disk_bytes_per_sandbox": settings.disk_limit_bytes or None, + "io_max": self.cgroups.limits.io_line if cg else None, "address_space_bytes_per_process": ( settings.address_space_limit_bytes if rlimit_as else None ), @@ -626,6 +711,12 @@ async def _execute_locked(self, sandbox: LocalSandbox, request: ExecRequest) -> self.cgroups.create_execution, sandbox.sandbox_id, request.exec_id ) command = placement_prefix(leaf) + command + elif shutil.which("ionice"): + # Best effort without io.max: the lowest best-effort priority, so a + # sandbox's I/O yields to the worker's on schedulers that honour it. + command = ["ionice", "-c", "2", "-n", "7", "--", *command] + exec_key = request.exec_id + self.supervisor.begin(sandbox.sandbox_id, sandbox.uid, sandbox.root, exec_key) started = time.monotonic() try: process = await asyncio.create_subprocess_exec( @@ -635,6 +726,7 @@ async def _execute_locked(self, sandbox: LocalSandbox, request: ExecRequest) -> start_new_session=True, ) except BaseException: + self.supervisor.end(sandbox.sandbox_id, exec_key) if leaf is not None: await asyncio.to_thread(self.cgroups.remove_execution, leaf) raise @@ -686,6 +778,7 @@ async def _execute_locked(self, sandbox: LocalSandbox, request: ExecRequest) -> raise finally: self.processes.pop((sandbox.sandbox_id, request.exec_id), None) + verdict = self.supervisor.end(sandbox.sandbox_id, exec_key) stats: dict[str, int] = {} if leaf is not None: self._exec_cgroups.pop((sandbox.sandbox_id, request.exec_id), None) @@ -702,7 +795,14 @@ async def _execute_locked(self, sandbox: LocalSandbox, request: ExecRequest) -> stdout = stdout_b.decode(errors="replace") stderr = stderr_b.decode(errors="replace") error_code = None - if stats.get("memory.oom_kill") or stats.get("memory.oom_group_kill"): + if verdict == DISK: + status = "FAILED" + error_code = DISK + elif ( + verdict is not None + or stats.get("memory.oom_kill") + or stats.get("memory.oom_group_kill") + ): status = "OOM_KILLED" error_code = "SANDBOX_MEMORY_LIMIT_EXCEEDED" elif stats.get("pids.max"): @@ -872,6 +972,7 @@ async def _destroy(self, sandbox_id: str, *, delete_files: bool) -> None: if sid == sandbox_id: await self._terminate(process) await asyncio.to_thread(self.cgroups.remove_sandbox, sandbox_id) + self.supervisor.forget(sandbox_id) sandbox = self.sandboxes.get(sandbox_id) if sandbox is not None and not self.cgroup_active: await asyncio.to_thread(self._kill_uid, sandbox.uid) @@ -972,6 +1073,7 @@ def _suspend_sandbox( self.last_active_at.pop(sandbox_id, None) # A dormant sandbox runs nothing and holds no cgroup. self.cgroups.remove_sandbox(sandbox_id) + self.supervisor.forget(sandbox_id) self.dormant[sandbox_id] = DormantSandbox( sandbox_id=sandbox_id, generation=generation, @@ -1300,11 +1402,15 @@ async def shutdown(self) -> None: for sandbox_id in list(self.sandboxes): await asyncio.to_thread(self.cgroups.remove_sandbox, sandbox_id) self.cgroups.close() + self.supervisor.stop() async def write_file(self, sandbox: LocalSandbox, path: str, content_base64: str) -> None: content = base64.b64decode(content_base64, validate=True) if len(content) > self.settings.max_file_api_bytes: raise ValueError("file exceeds the sandbox File API size limit") + limit = self.settings.disk_limit_bytes + if limit and await asyncio.to_thread(disk_usage, sandbox.root) + len(content) > limit: + raise RuntimeError("SANDBOX_DISK_QUOTA_EXCEEDED") with self._path_lock(sandbox, path): WorkspaceFiles(sandbox.workspace, sandbox.uid).write(path, content) self.touch(sandbox.sandbox_id) diff --git a/src/agent_sandbox/supervisor.py b/src/agent_sandbox/supervisor.py new file mode 100644 index 0000000..de2d7a9 --- /dev/null +++ b/src/agent_sandbox/supervisor.py @@ -0,0 +1,390 @@ +"""Aggregate per-sandbox enforcement where the kernel offers no cgroup. + +Every sandbox runs under its own outer UID, so the set of processes with that +real UID *is* the sandbox's process tree, however deeply it forks or +daemonizes. A thread samples ``/proc`` every ``interval`` seconds and, per +sandbox: + +- **memory**: sums PSS (``smaps_rollup``; RSS where unavailable) over the + UID's processes; above the limit every process of the UID is SIGKILLed and + the running executions report ``SANDBOX_MEMORY_LIMIT_EXCEEDED``. +- **CPU**: a token bucket refilled at ``cpu_cores`` CPU-seconds per second; + when it runs dry the UID's processes are SIGSTOPped until it refills, then + SIGCONTed. The sandbox's long-run CPU use converges to the quota. +- **disk**: (all schemes) sums the allocated blocks of the sandbox root + every ``disk_interval`` seconds; when usage is over the limit and still + growing, the running executions are killed with + ``SANDBOX_DISK_QUOTA_EXCEEDED``. Deleting files is always allowed. + +This is supervision, not kernel enforcement: a burst can exceed the memory +limit for up to one interval before the kill (allocation rate × interval), +CPU is shaped at interval granularity, and a process that lives shorter than +an interval is not charged. ``/healthz`` says so. +""" + +from __future__ import annotations + +import contextlib +import logging +import os +import signal +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path + +logger = logging.getLogger(__name__) + +_TICKS = os.sysconf("SC_CLK_TCK") if hasattr(os, "sysconf") else 100 +_PAGE = os.sysconf("SC_PAGE_SIZE") if hasattr(os, "sysconf") else 4096 + +MEMORY = "SANDBOX_MEMORY_LIMIT_EXCEEDED" +DISK = "SANDBOX_DISK_QUOTA_EXCEEDED" + + +@dataclass +class _Watched: + uid: int + root: Path + execs: set[str] = field(default_factory=set) + verdicts: dict[str, str] = field(default_factory=dict) + cpu_seen: dict[int, tuple[int, int]] = field(default_factory=dict) + budget: float = 0.0 + stopped: set[int] = field(default_factory=set) + disk_used: int = 0 + disk_checked: float = 0.0 + peak_memory: int = 0 + io_seen: dict[int, tuple[int, int]] = field(default_factory=dict) + io_budget: list[float] = field(default_factory=lambda: [0.0, 0.0]) + + +def processes_by_uid(proc: Path = Path("/proc")) -> dict[int, list[int]]: + found: dict[int, list[int]] = {} + for entry in os.scandir(proc): + if not entry.name.isdigit(): + continue + with contextlib.suppress(OSError, ValueError): + uid = os.stat(entry.path).st_uid # owner of /proc/ is the euid + # The real UID is what setpriv set and what RLIMIT_NPROC counts. + with open(f"{entry.path}/status") as stream: + for line in stream: + if line.startswith("Uid:"): + uid = int(line.split()[1]) + break + found.setdefault(uid, []).append(int(entry.name)) + return found + + +def process_memory(pid: int) -> int: + try: + with open(f"/proc/{pid}/smaps_rollup") as stream: + for line in stream: + if line.startswith("Pss:"): + return int(line.split()[1]) * 1024 + except OSError: + pass + try: + with open(f"/proc/{pid}/statm") as stream: + return int(stream.read().split()[1]) * _PAGE + except (OSError, ValueError, IndexError): + return 0 + + +def process_cpu_ticks(pid: int) -> tuple[int, int]: + """(own utime+stime, reaped children's cutime+cstime) in clock ticks.""" + try: + with open(f"/proc/{pid}/stat") as stream: + fields = stream.read().rpartition(")")[2].split() + return int(fields[11]) + int(fields[12]), int(fields[13]) + int(fields[14]) + except (OSError, ValueError, IndexError): + return 0, 0 + + +def process_io(pid: int) -> tuple[int, int]: + """Bytes this process caused to be read from / written to storage.""" + read = write = 0 + try: + with open(f"/proc/{pid}/io") as stream: + for line in stream: + if line.startswith("read_bytes:"): + read = int(line.split()[1]) + elif line.startswith("write_bytes:"): + write = int(line.split()[1]) + except (OSError, ValueError): + pass + return read, write + + +def can_read_foreign_io(setpriv: str, uid: int) -> bool: + """Whether this process may read /proc//io of another UID. + + That needs ptrace read access, which root has only with CAP_SYS_PTRACE in + a container (Docker drops it by default). + """ + import subprocess + + try: + child = subprocess.Popen( + [ + setpriv, + "--reuid", + str(uid), + "--regid", + str(uid), + "--clear-groups", + "--", + "sleep", + "5", + ], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + except OSError: + return False + try: + time.sleep(0.05) + with open(f"/proc/{child.pid}/io") as stream: + stream.read() + return True + except OSError: + return False + finally: + child.kill() + child.wait() + + +def disk_usage(root: Path, *, budget: int = 2_000_000) -> int: + """Allocated bytes under ``root``, without following symlinks.""" + total = 0 + stack = [str(root)] + seen = 0 + while stack and seen < budget: + path = stack.pop() + try: + entries = os.scandir(path) + except OSError: + continue + with entries: + for entry in entries: + seen += 1 + try: + info = entry.stat(follow_symlinks=False) + except OSError: + continue + total += info.st_blocks * 512 + if entry.is_dir(follow_symlinks=False): + stack.append(entry.path) + return total + + +class SandboxSupervisor: + def __init__( + self, + *, + memory_bytes: int | None, + cpu_cores: float | None, + disk_bytes: int | None, + interval: float = 0.1, + disk_interval: float = 2.0, + kill_uid: Callable[[int], None] | None = None, + io_read_bps: int | None = None, + io_write_bps: int | None = None, + ) -> None: + self.io_read_bps = io_read_bps + self.io_write_bps = io_write_bps + self.memory_bytes = memory_bytes + self.cpu_cores = cpu_cores + self.disk_bytes = disk_bytes + self.interval = interval + self.disk_interval = disk_interval + self._kill_uid = kill_uid + self._lock = threading.Lock() + self._watched: dict[str, _Watched] = {} + self._thread: threading.Thread | None = None + self._stop = threading.Event() + + @property + def enabled(self) -> bool: + return bool( + self.memory_bytes + or self.cpu_cores + or self.disk_bytes + or self.io_read_bps + or self.io_write_bps + ) + + def start(self) -> None: + if self._thread is None and self.enabled: + self._thread = threading.Thread( + target=self._run, name="sandbox-supervisor", daemon=True + ) + self._thread.start() + + def stop(self) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=5) + with self._lock: + for watched in self._watched.values(): + self._resume(watched) + self._thread = None + + def begin(self, sandbox_id: str, uid: int, root: Path, exec_key: str) -> None: + with self._lock: + watched = self._watched.get(sandbox_id) + if watched is None or watched.uid != uid: + watched = _Watched(uid=uid, root=root) + watched.budget = (self.cpu_cores or 0) * 1.0 + self._watched[sandbox_id] = watched + watched.execs.add(exec_key) + + def end(self, sandbox_id: str, exec_key: str) -> str | None: + with self._lock: + watched = self._watched.get(sandbox_id) + if watched is None: + return None + watched.execs.discard(exec_key) + verdict = watched.verdicts.pop(exec_key, None) + if not watched.execs: + self._resume(watched) + watched.cpu_seen.clear() + return verdict + + def forget(self, sandbox_id: str) -> None: + with self._lock: + watched = self._watched.pop(sandbox_id, None) + if watched is not None: + self._resume(watched) + + def disk_exceeded(self, sandbox_id: str, root: Path, extra: int = 0) -> bool: + """For File API writes: would ``extra`` more bytes exceed the quota?""" + if not self.disk_bytes: + return False + return disk_usage(root) + extra > self.disk_bytes + + # ── the loop ── + + def _run(self) -> None: + last = time.monotonic() + while not self._stop.wait(self.interval): + now = time.monotonic() + elapsed, last = now - last, now + try: + self._tick(elapsed, now) + except Exception: # never let the supervisor die + logger.exception("sandbox supervisor tick failed") + + def _tick(self, elapsed: float, now: float) -> None: + with self._lock: + active = {sid: w for sid, w in self._watched.items() if w.execs} + if not active: + return + need_procs = bool( + self.memory_bytes or self.cpu_cores or self.io_read_bps or self.io_write_bps + ) + by_uid = processes_by_uid() if need_procs else {} + for sandbox_id, watched in active.items(): + pids = by_uid.get(watched.uid, []) + if self.memory_bytes and pids: + used = sum(process_memory(pid) for pid in pids) + watched.peak_memory = max(watched.peak_memory, used) + if used > self.memory_bytes: + logger.warning( + "sandbox %s uses %d bytes over its %d limit; killing", + sandbox_id, + used, + self.memory_bytes, + ) + self._kill(watched, pids, MEMORY) + continue + throttle = False + if self.cpu_cores: + throttle |= self._cpu_over(watched, pids, elapsed) + if self.io_read_bps or self.io_write_bps: + throttle |= self._io_over(watched, pids, elapsed) + if self.cpu_cores or self.io_read_bps or self.io_write_bps: + self._throttle(watched, pids, throttle) + if self.disk_bytes and now - watched.disk_checked >= self.disk_interval: + watched.disk_checked = now + used = disk_usage(watched.root) + growing = used > watched.disk_used + watched.disk_used = used + if used > self.disk_bytes and growing: + logger.warning("sandbox %s exceeds its disk quota; killing", sandbox_id) + self._kill(watched, pids or by_uid_for(watched.uid), DISK) + + def _io_over(self, watched: _Watched, pids: list[int], elapsed: float) -> bool: + seen = {pid: process_io(pid) for pid in pids} + used = [0, 0] + for pid, current in seen.items(): + before = watched.io_seen.get(pid, (0, 0)) + used[0] += max(0, current[0] - before[0]) + used[1] += max(0, current[1] - before[1]) + watched.io_seen = seen + over = False + for index, rate in enumerate((self.io_read_bps, self.io_write_bps)): + if rate: + watched.io_budget[index] = min( + float(rate), watched.io_budget[index] + rate * elapsed - used[index] + ) + over |= watched.io_budget[index] < 0 + return over + + def _throttle(self, watched: _Watched, pids: list[int], stop: bool) -> None: + if stop: + for pid in pids: + if pid not in watched.stopped: + with contextlib.suppress(OSError): + os.kill(pid, signal.SIGSTOP) + watched.stopped.add(pid) + elif watched.stopped: + self._resume(watched) + + def _cpu_over(self, watched: _Watched, pids: list[int], elapsed: float) -> bool: + assert self.cpu_cores + own = 0 + reaped = 0 + seen: dict[int, tuple[int, int]] = {} + for pid in pids: + current = process_cpu_ticks(pid) + before = watched.cpu_seen.get(pid, (0, 0)) + seen[pid] = current + own += max(0, current[0] - before[0]) + reaped += max(0, current[1] - before[1]) + # A reaped child's time reappears in its parent's cutime; only the part + # not already charged while the child was alive is new. + vanished = sum(a + b for pid, (a, b) in watched.cpu_seen.items() if pid not in seen) + used = own + max(0, reaped - vanished) + watched.cpu_seen = seen + burst = self.cpu_cores * 1.0 + watched.budget = min(burst, watched.budget + self.cpu_cores * elapsed - used / _TICKS) + return watched.budget < 0 + + def _resume(self, watched: _Watched) -> None: + for pid in watched.stopped: + with contextlib.suppress(OSError): + os.kill(pid, signal.SIGCONT) + watched.stopped.clear() + + def _kill(self, watched: _Watched, pids: list[int], verdict: str) -> None: + for key in watched.execs: + watched.verdicts.setdefault(key, verdict) + for pid in pids: + with contextlib.suppress(OSError): + os.kill(pid, signal.SIGKILL) + if self._kill_uid is not None: + self._kill_uid(watched.uid) + watched.stopped.clear() + + def report(self) -> dict[str, object]: + return { + "interval_seconds": self.interval, + "disk_interval_seconds": self.disk_interval, + "memory_metric": "sum of PSS (RSS where smaps_rollup is unreadable) over the sandbox UID's processes", + } + + +def by_uid_for(uid: int) -> list[int]: + return processes_by_uid().get(uid, []) diff --git a/tests/integration/test_resource_limits.py b/tests/integration/test_resource_limits.py index 8bd9aef..71727b9 100644 --- a/tests/integration/test_resource_limits.py +++ b/tests/integration/test_resource_limits.py @@ -33,8 +33,8 @@ MEMORY_LIMIT = 256 * 1024 * 1024 -@pytest.fixture -async def runtime(tmp_path: Path) -> Any: +@pytest.fixture(params=["auto", "off"], ids=["cgroup-if-available", "rlimit-supervisor"]) +async def runtime(tmp_path: Path, request: pytest.FixtureRequest) -> Any: local_root = tmp_path / "sandboxes" local_root.mkdir() for ancestor in (local_root, *local_root.parents): @@ -44,6 +44,7 @@ async def runtime(tmp_path: Path) -> Any: instance = SandboxRuntime( Settings( internal_token="integration-token", + cgroup_mode=request.param, local_root=local_root, database_url=f"sqlite+aiosqlite:///{tmp_path}/control.db", uid_start=20100, @@ -52,6 +53,8 @@ async def runtime(tmp_path: Path) -> Any: memory_limit_bytes=MEMORY_LIMIT, pids_limit=64, address_space_limit_bytes=2 * MEMORY_LIMIT, + disk_limit_bytes=128 * 1024 * 1024, + disk_check_interval_seconds=0.5, min_free_bytes=0, disk_high_watermark_percent=99, ) @@ -83,9 +86,14 @@ def _scheme(runtime: SandboxRuntime) -> str: async def test_report_matches_the_active_scheme(runtime: SandboxRuntime) -> None: control = runtime.probe_report["resource_control"] # type: ignore[attr-defined] cgroup = control["scheme"] == "cgroup_v2" - assert control["sandbox_memory_quota"] is cgroup - assert control["sandbox_cpu_quota"] is cgroup - assert control["sandbox_pid_quota"] is cgroup + # Aggregate limits hold under both schemes; only how differs. + assert control["sandbox_memory_quota"] is True + assert control["sandbox_cpu_quota"] is True + assert control["sandbox_pid_quota"] is True + assert control["sandbox_disk_quota"] is True + expected = "kernel" if cgroup else "supervised" + assert control["enforcement"]["memory"].startswith(expected) + assert control["enforcement"]["cpu"].startswith(expected) if not cgroup: assert runtime.probe_report["selected_level"] != "strict" # type: ignore[attr-defined] assert control["limits"]["address_space_bytes_per_process"] == 2 * MEMORY_LIMIT @@ -98,12 +106,12 @@ async def test_memory_bomb_is_stopped(runtime: SandboxRuntime) -> None: sandbox, "python3 -c 'b=[]\nwhile True: b.append(bytearray(16*1024*1024))'", ) - if _scheme(runtime) == "cgroup_v2": - assert result.status == "OOM_KILLED" - assert result.error_code == "SANDBOX_MEMORY_LIMIT_EXCEEDED" - else: - assert result.status == "FAILED" + if result.status != "OOM_KILLED": + # Without cgroups RLIMIT_AS may stop the single process first. + assert _scheme(runtime) != "cgroup_v2" assert "MemoryError" in result.stderr + else: + assert result.error_code == "SANDBOX_MEMORY_LIMIT_EXCEEDED" # The sandbox is usable afterwards. assert (await _run(runtime, sandbox, "echo ok")).stdout.strip() == "ok" @@ -119,9 +127,35 @@ async def test_fork_bomb_is_bounded_and_leaves_nothing(runtime: SandboxRuntime) assert (await _run(runtime, sandbox, "echo again")).stdout.strip() == "again" +async def test_multi_process_memory_is_limited_in_aggregate(runtime: SandboxRuntime) -> None: + sandbox = await runtime.create("sandbox-multimem", 1, 20105) + # Four processes of 100 MiB: each under every per-process limit. + result = await _run( + runtime, + sandbox, + "python3 - <<'P'\nimport multiprocessing as m, time\n" + "def f():\n b=bytearray(100<<20)\n b[::4096]=b'1'*len(b[::4096])\n time.sleep(20)\n" + "ps=[m.Process(target=f) for _ in range(4)]\n[p.start() for p in ps]\n" + "[p.join() for p in ps]\nP", + limit=60, + ) + assert result.status == "OOM_KILLED", (result.status, result.stderr[-300:]) + assert result.error_code == "SANDBOX_MEMORY_LIMIT_EXCEEDED" + + +async def test_disk_quota_stops_a_runaway_writer(runtime: SandboxRuntime) -> None: + sandbox = await runtime.create("sandbox-disk", 1, 20106) + result = await _run( + runtime, + sandbox, + "for i in $(seq 1 64); do head -c 16777216 /dev/zero > /workspace/f$i; done; sleep 10", + limit=60, + ) + assert result.error_code == "SANDBOX_DISK_QUOTA_EXCEEDED" + assert (await _run(runtime, sandbox, "rm -f /workspace/f*; echo ok")).stdout.strip() == "ok" + + async def test_cpu_quota_caps_a_multicore_spin(runtime: SandboxRuntime) -> None: - if _scheme(runtime) != "cgroup_v2": - pytest.skip("no CPU quota without cgroups; RLIMIT_CPU is per process") sandbox = await runtime.create("sandbox-spin", 1, 20104) result = await _run( runtime, @@ -133,5 +167,6 @@ async def test_cpu_quota_caps_a_multicore_spin(runtime: SandboxRuntime) -> None: limit=60, ) assert result.status == "SUCCEEDED", result.stderr - # Four spinners for 3 s wall under a one-core quota: about 3 s of CPU. - assert float(result.stdout.strip().splitlines()[-1]) < 3 * 1.5 + # Four spinners for 3 s wall under a one-core quota: about 3 s of CPU, + # plus up to one CPU-second of burst when supervised. + assert float(result.stdout.strip().splitlines()[-1]) < 3 * 1.5 + 1 diff --git a/tests/test_supervisor.py b/tests/test_supervisor.py new file mode 100644 index 0000000..3529370 --- /dev/null +++ b/tests/test_supervisor.py @@ -0,0 +1,126 @@ +import signal +from pathlib import Path + +import pytest + +from agent_sandbox import supervisor as sup +from agent_sandbox.supervisor import DISK, MEMORY, SandboxSupervisor, disk_usage + + +@pytest.fixture +def signals(monkeypatch: pytest.MonkeyPatch) -> list[tuple[int, int]]: + sent: list[tuple[int, int]] = [] + monkeypatch.setattr(sup.os, "kill", lambda pid, sig: sent.append((pid, sig))) + return sent + + +def _watch(supervisor: SandboxSupervisor, tmp_path: Path) -> None: + supervisor.begin("sb", 30001, tmp_path, "exec-1") + + +def test_memory_is_summed_across_processes_and_kills_the_group( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, signals: list[tuple[int, int]] +) -> None: + killed: list[int] = [] + supervisor = SandboxSupervisor( + memory_bytes=300, cpu_cores=None, disk_bytes=None, kill_uid=killed.append + ) + _watch(supervisor, tmp_path) + monkeypatch.setattr(sup, "processes_by_uid", lambda: {30001: [11, 12, 13], 0: [1]}) + # Each process is under the limit; the sandbox is not. + monkeypatch.setattr(sup, "process_memory", lambda pid: 120) + supervisor._tick(0.1, 1.0) + assert {pid for pid, sig in signals if sig == signal.SIGKILL} == {11, 12, 13} + assert killed == [30001] + assert supervisor.end("sb", "exec-1") == MEMORY + + +def test_under_the_memory_limit_nothing_happens( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, signals: list[tuple[int, int]] +) -> None: + supervisor = SandboxSupervisor(memory_bytes=1000, cpu_cores=None, disk_bytes=None) + _watch(supervisor, tmp_path) + monkeypatch.setattr(sup, "processes_by_uid", lambda: {30001: [11, 12]}) + monkeypatch.setattr(sup, "process_memory", lambda pid: 100) + supervisor._tick(0.1, 1.0) + assert signals == [] + assert supervisor.end("sb", "exec-1") is None + + +def test_cpu_over_quota_is_stopped_then_resumed( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, signals: list[tuple[int, int]] +) -> None: + supervisor = SandboxSupervisor(memory_bytes=None, cpu_cores=1.0, disk_bytes=None) + _watch(supervisor, tmp_path) + ticks = {11: 0, 12: 0} + monkeypatch.setattr(sup, "processes_by_uid", lambda: {30001: [11, 12]}) + monkeypatch.setattr(sup, "process_cpu_ticks", lambda pid: (ticks[pid], 0)) + supervisor._tick(0.1, 1.0) + # Three CPU-seconds in one second against a one-core quota and burst. + ticks.update({11: int(sup._TICKS * 1.5), 12: int(sup._TICKS * 1.5)}) + supervisor._tick(1.0, 2.0) + assert {pid for pid, sig in signals if sig == signal.SIGSTOP} == {11, 12} + signals.clear() + # Stopped processes use nothing; the bucket refills and they continue. + supervisor._tick(1.0, 3.0) + supervisor._tick(1.0, 4.0) + assert {pid for pid, sig in signals if sig == signal.SIGCONT} == {11, 12} + + +def test_reaped_child_time_is_not_charged_twice( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, signals: list[tuple[int, int]] +) -> None: + supervisor = SandboxSupervisor(memory_bytes=None, cpu_cores=1.0, disk_bytes=None) + _watch(supervisor, tmp_path) + state = {"pids": [11, 12], 11: (0, 0), 12: (int(sup._TICKS * 0.5), 0)} + monkeypatch.setattr(sup, "processes_by_uid", lambda: {30001: state["pids"]}) + monkeypatch.setattr(sup, "process_cpu_ticks", lambda pid: state[pid]) + supervisor._tick(1.0, 1.0) + budget = supervisor._watched["sb"].budget + # The child exits; its half second moves into the parent's cutime. + state["pids"] = [11] + state[11] = (0, int(sup._TICKS * 0.5)) + supervisor._tick(0.0, 1.0) + assert supervisor._watched["sb"].budget == pytest.approx(budget) + + +def test_io_bytes_are_duty_cycled( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, signals: list[tuple[int, int]] +) -> None: + supervisor = SandboxSupervisor( + memory_bytes=None, cpu_cores=None, disk_bytes=None, io_write_bps=1000 + ) + _watch(supervisor, tmp_path) + written = {11: 0} + monkeypatch.setattr(sup, "processes_by_uid", lambda: {30001: [11]}) + monkeypatch.setattr(sup, "process_io", lambda pid: (0, written[pid])) + supervisor._tick(0.1, 1.0) + written[11] = 5000 + supervisor._tick(1.0, 2.0) + assert (11, signal.SIGSTOP) in signals + + +def test_disk_growth_over_quota_kills_and_deletion_is_allowed( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, signals: list[tuple[int, int]] +) -> None: + supervisor = SandboxSupervisor( + memory_bytes=None, cpu_cores=None, disk_bytes=4096 * 4, disk_interval=0 + ) + _watch(supervisor, tmp_path) + monkeypatch.setattr(sup, "processes_by_uid", lambda: {30001: [11]}) + monkeypatch.setattr(sup, "by_uid_for", lambda uid: [11]) + (tmp_path / "big").write_bytes(b"x" * 4096 * 8) + supervisor._tick(0.1, 1.0) + assert supervisor.end("sb", "exec-1") == DISK + supervisor.begin("sb", 30001, tmp_path, "exec-2") + signals.clear() + (tmp_path / "big").write_bytes(b"x" * 4096 * 6) # shrinking, still over + supervisor._tick(0.1, 2.0) + assert signals == [] + + +def test_disk_usage_counts_allocated_blocks(tmp_path: Path) -> None: + (tmp_path / "d").mkdir() + (tmp_path / "d" / "f").write_bytes(b"x" * 10000) + (tmp_path / "link").symlink_to("/usr") + assert 10000 <= disk_usage(tmp_path) < 10000 + 64 * 1024 diff --git a/tests/test_template_security.py b/tests/test_template_security.py index 497bb7b..df5a38d 100644 --- a/tests/test_template_security.py +++ b/tests/test_template_security.py @@ -494,16 +494,25 @@ async def test_without_cgroups_strict_is_unavailable_and_rlimits_apply( runtime, _ = await runtime_and_sandbox(tmp_path) _fake_probe(runtime, monkeypatch) monkeypatch.setattr(runtime, "_setup_cgroups", lambda: None) + monkeypatch.setattr("agent_sandbox.runtime.can_read_foreign_io", lambda *a: True) report = await runtime.probe() # strict promises aggregate limits, so it is not selected without them. assert report["selected_level"] == "standard" assert "resource limits" in report["probe_failures"]["strict"] control = report["resource_control"] - assert control["scheme"] == "rlimit" - assert control["sandbox_cpu_quota"] is False - assert control["sandbox_memory_quota"] is False - assert control["sandbox_pid_quota"] is False + assert control["scheme"] == "rlimit+supervisor" + # Aggregate limits still hold, enforced by the supervisor, and say so. + assert control["sandbox_cpu_quota"] is True + assert control["sandbox_memory_quota"] is True + assert control["enforcement"]["memory"].startswith("supervised") + assert control["enforcement"]["cpu"].startswith("supervised") + assert control["enforcement"]["pids"].startswith("kernel (RLIMIT_NPROC") + assert control["sandbox_pid_quota"] is True + assert control["sandbox_disk_quota"] is True + assert control["sandbox_io_limit"] is True + assert control["enforcement"]["io"].startswith("supervised") assert control["execution_cgroup_supervision"] is False + assert runtime.supervisor.memory_bytes == 2 * 1024**3 assert control["limits"]["address_space_bytes_per_process"] == 8 * 1024**3 assert control["limits"]["wall_clock_seconds_per_exec_default"] == 300 assert control["process_limits"]["processes_per_real_uid"] == 64 @@ -529,7 +538,10 @@ def active() -> None: assert control["sandbox_cpu_quota"] is True assert control["sandbox_memory_quota"] is True assert control["sandbox_pid_quota"] is True - assert control["sandbox_disk_quota"] is False + assert control["sandbox_disk_quota"] is True + assert control["enforcement"]["memory"] == "kernel (cgroup v2)" + # The kernel enforces memory and CPU; the supervisor only watches disk. + assert runtime.supervisor.memory_bytes is None assert control["limits"]["memory_bytes_per_sandbox"] == 2 * 1024**3 assert control["limits"]["address_space_bytes_per_process"] is None # The aggregate cgroup limit replaces RLIMIT_AS, which breaks JIT runtimes. From ce27f848ff45e7c344885fe2b0e930d4367a0127 Mon Sep 17 00:00:00 2001 From: fliaping <269832017@qq.com> Date: Tue, 13 Oct 2026 20:35:39 +0800 Subject: [PATCH 4/4] feat: size sandbox JVMs from the sandbox memory limit via the image launchers --- docs/CONFIGURATION.md | 1 + docs/TOOLCHAINS.md | 16 +++++++ .../src/agent_sandbox_runtime/toolchains.py | 16 ++++++- scripts/install-polyglot-launchers.py | 40 ++++++++++++++++- src/agent_sandbox/config.py | 3 ++ src/agent_sandbox/runtime.py | 5 +++ tests/test_polyglot_launchers.py | 45 +++++++++++++++++++ 7 files changed, 123 insertions(+), 3 deletions(-) diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index b03b2c0..7a5559f 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -75,6 +75,7 @@ profile identity. | `SANDBOX_PIDS_LIMIT` | `512` | Processes and threads per sandbox (`pids.max`). `0` leaves it unlimited. cgroup scheme only. | | `SANDBOX_WORKER_MEMORY_RESERVE_BYTES` | `268435456` | Memory kept for the worker's control process: set as `memory.min` on its leaf and subtracted from the sandboxes' aggregate `memory.max`. cgroup scheme only. | | `SANDBOX_ADDRESS_SPACE_LIMIT_BYTES` | `8589934592` | Fallback per-process `RLIMIT_AS`, applied only when cgroups are not active. Larger than the memory limit because Node and the JVM reserve address space they never use. `0` disables it. | +| `SANDBOX_JVM_MAX_RAM_FROM_LIMIT` | `true` | Size JVMs from `SANDBOX_MEMORY_LIMIT_BYTES`: the polyglot image's Java launchers prepend `-XX:MaxRAM=` (`-J-XX:MaxRAM` for javac, jar, jshell, …), so the default heap is a quarter of the sandbox limit instead of the host's RAM. User flags (`-Xmx`, `-XX:MaxRAM`, `-XX:MaxRAMPercentage`) come later and win. No stderr note, unlike `JAVA_TOOL_OPTIONS`. Applies to every level and scheme; JDKs without these launchers are not affected. | | `SANDBOX_SUPERVISOR_INTERVAL_SECONDS` | `0.1` | Without cgroups, how often the supervisor samples `/proc` to enforce `SANDBOX_MEMORY_LIMIT_BYTES` (sum of PSS) and `SANDBOX_CPU_QUOTA` (SIGSTOP/SIGCONT duty cycle) over all processes of the sandbox UID. Shorter is tighter and costs more worker CPU. | | `SANDBOX_DISK_LIMIT_BYTES` | `10737418240` | Allocated bytes per sandbox (workspace, home, cache, envs), every scheme. Growth past it kills running executions with `SANDBOX_DISK_QUOTA_EXCEEDED`; File API writes over it are refused (HTTP 507). Deleting is always allowed. `0` disables it. | | `SANDBOX_DISK_CHECK_INTERVAL_SECONDS` | `2.0` | How often the disk usage of a sandbox with a running execution is measured. | diff --git a/docs/TOOLCHAINS.md b/docs/TOOLCHAINS.md index 62fb90a..5eba4ed 100644 --- a/docs/TOOLCHAINS.md +++ b/docs/TOOLCHAINS.md @@ -190,3 +190,19 @@ notice. Every enabled language it cannot check is reported as SKIP with the reason, and `--strict` turns a SKIP into a failure so a partial run cannot pass quietly. + +## JVM heap follows the sandbox memory limit + +Below `strict` the JVM cannot see a cgroup of its own, so it sizes its default +heap from the host's RAM: a quarter of it, which can be far larger than the +sandbox memory limit. The image's Java launchers therefore start every JVM +with `-XX:MaxRAM=`, and every JDK tool with +`-J-XX:MaxRAM=…`. This applies at every level and under both resource +schemes. The flag goes first on the command line, so anything the user +passes (`-Xmx`, `-XX:MaxRAM`, `-XX:MaxRAMPercentage`) still wins. Maven and +Gradle reach the same launchers through `JAVA_HOME`. + +The launchers do not use `JAVA_TOOL_OPTIONS` or `JDK_JAVA_OPTIONS`, because +both print a "Picked up …" note to stderr and that breaks tools that parse +stderr. Set `SANDBOX_JVM_MAX_RAM_FROM_LIMIT=false` to turn this off. A +system JDK without these launchers is not adjusted. diff --git a/runtime/src/agent_sandbox_runtime/toolchains.py b/runtime/src/agent_sandbox_runtime/toolchains.py index 99c7f48..0cfca37 100644 --- a/runtime/src/agent_sandbox_runtime/toolchains.py +++ b/runtime/src/agent_sandbox_runtime/toolchains.py @@ -235,7 +235,10 @@ def rust_toolchain( def java_toolchain( - *, maven_repository_url: str | None = None, java_home: str | None = None + *, + maven_repository_url: str | None = None, + java_home: str | None = None, + jvm_max_ram_bytes: int | None = None, ) -> Toolchain: """JVM languages with Maven and Gradle. @@ -252,6 +255,9 @@ def java_toolchain( environment["MAVEN_MIRROR_URL"] = maven_repository_url if java_home: environment["JAVA_HOME"] = java_home + if jvm_max_ram_bytes: + # Read by the image-owned launchers, which pass -XX:MaxRAM. + environment["SANDBOX_JVM_MAX_RAM"] = str(jvm_max_ram_bytes) return Toolchain( name="java", environment=environment, @@ -259,7 +265,13 @@ def java_toolchain( f"{ENVS_ROOT}/java/bin", *((f"{java_home}/bin",) if java_home else ()), ), - reserved_names=("MAVEN_OPTS", "GRADLE_USER_HOME", "SBT_OPTS", "MAVEN_MIRROR_URL"), + reserved_names=( + "MAVEN_OPTS", + "GRADLE_USER_HOME", + "SBT_OPTS", + "MAVEN_MIRROR_URL", + "SANDBOX_JVM_MAX_RAM", + ), needs_procfs=java_home is None, reserved_prefixes=("JAVA_TOOL_OPTIONS", "_JAVA_OPTIONS", "CLASSPATH"), ) diff --git a/scripts/install-polyglot-launchers.py b/scripts/install-polyglot-launchers.py index 8a7cd8e..2e39ab9 100644 --- a/scripts/install-polyglot-launchers.py +++ b/scripts/install-polyglot-launchers.py @@ -23,6 +23,42 @@ def _library_path(directory: Path) -> str: return f"export LD_LIBRARY_PATH={shlex.quote(str(directory))}${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}}\n" +# JDK tools that take JVM options through -J. `java` takes them directly. +_J_OPTION_TOOLS = frozenset( + { + "javac", + "jar", + "jarsigner", + "javadoc", + "jdeps", + "jlink", + "jmod", + "jpackage", + "jshell", + "keytool", + "jdb", + } +) + + +def _jvm_memory(name: str) -> str: + """Size the JVM from the sandbox's memory limit, not the host's RAM. + + Without a cgroup namespace the JVM sees the whole host and picks a default + heap of a quarter of it, far above the sandbox limit. The runtime exports + SANDBOX_JVM_MAX_RAM; it goes first so any flag the user passes later + (-Xmx, -XX:MaxRAM, -XX:MaxRAMPercentage) still wins. This is quiet, unlike + JAVA_TOOL_OPTIONS and JDK_JAVA_OPTIONS, which print a note to stderr. + """ + if name == "java": + flag = '-XX:MaxRAM="$SANDBOX_JVM_MAX_RAM"' + elif name in _J_OPTION_TOOLS: + flag = '-J-XX:MaxRAM="$SANDBOX_JVM_MAX_RAM"' + else: + return "" + return f'if [ -n "${{SANDBOX_JVM_MAX_RAM:-}}" ]; then set -- {flag} "$@"; fi\n' + + def install_java(java_home: Path, destination: Path) -> None: """Expose a JDK-shaped JAVA_HOME, including wrappers for every JDK tool.""" java_home = java_home.resolve(strict=True) @@ -36,7 +72,9 @@ def install_java(java_home: Path, destination: Path) -> None: if binary.is_file() and os.access(binary, os.X_OK): _launcher( destination / "bin" / binary.name, - _library_path(java_home / "lib") + f'exec {shlex.quote(str(binary))} "$@"\n', + _library_path(java_home / "lib") + + _jvm_memory(binary.name) + + f'exec {shlex.quote(str(binary))} "$@"\n', ) diff --git a/src/agent_sandbox/config.py b/src/agent_sandbox/config.py index d82ae6e..fd4c8bd 100644 --- a/src/agent_sandbox/config.py +++ b/src/agent_sandbox/config.py @@ -118,6 +118,9 @@ class Settings(BaseSettings): address_space_limit_bytes: int = Field(default=8 * 1024 * 1024 * 1024, ge=0) # Without cgroups, a supervisor enforces the memory and CPU limits above on # the sum of each sandbox UID's processes, sampling /proc this often. + # Pass -XX:MaxRAM= to JVMs through the image's + # Java launchers, so default heap sizing follows the sandbox, not the host. + jvm_max_ram_from_limit: bool = True supervisor_interval_seconds: float = Field(default=0.1, gt=0, le=5) # Allocated bytes per sandbox root (workspace, home, cache, envs), every # scheme. Checked every disk_check_interval_seconds; 0 disables it. diff --git a/src/agent_sandbox/runtime.py b/src/agent_sandbox/runtime.py index 9cf30ef..ee9b2d9 100644 --- a/src/agent_sandbox/runtime.py +++ b/src/agent_sandbox/runtime.py @@ -220,6 +220,11 @@ def configured_toolchains(settings: Settings) -> tuple[Toolchain, ...]: "registry_url": settings.cargo_registry_url or None, "maven_repository_url": settings.maven_repository_url or None, "java_home": settings.java_home or None, + "jvm_max_ram_bytes": ( + settings.memory_limit_bytes + if settings.jvm_max_ram_from_limit and settings.memory_limit_bytes + else None + ), "launcher_dir": settings.rust_launcher_dir or None, "platform_python_runtime": PLATFORM_PYTHON_RUNTIME, "platform_runtime_packages": PLATFORM_RUNTIME_PACKAGES, diff --git a/tests/test_polyglot_launchers.py b/tests/test_polyglot_launchers.py index b402cd2..44a5d32 100644 --- a/tests/test_polyglot_launchers.py +++ b/tests/test_polyglot_launchers.py @@ -139,3 +139,48 @@ def test_rust_override_is_refused_explicitly(tmp_path: Path, installer: ModuleTy assert result.returncode == 2 assert "RUST_VERSION" in result.stderr assert "should-not-run" not in result.stdout + + +@pytest.mark.parametrize( + ("tool", "flag"), + [("java", "-XX:MaxRAM=402653184"), ("javac", "-J-XX:MaxRAM=402653184")], +) +def test_java_launchers_size_the_jvm_from_the_sandbox_limit( + tmp_path: Path, installer: ModuleType, tool: str, flag: str +) -> None: + home = tmp_path / "jdk" + (home / "lib").mkdir(parents=True) + (home / "lib" / "libjli.so").touch() + executable(home / "bin" / tool, 'printf "%s\\n" "$@"\n') + executable(home / "bin" / "jfr", 'printf "%s\\n" "$@"\n') + installer.install_java(home, tmp_path / "facade") + + def run(name: str, *args: str, limit: str | None = "402653184") -> list[str]: + env = {k: v for k, v in os.environ.items() if k != "SANDBOX_JVM_MAX_RAM"} + if limit is not None: + env["SANDBOX_JVM_MAX_RAM"] = limit + result = subprocess.run( + [str(tmp_path / "facade" / "bin" / name), *args], + env=env, + capture_output=True, + text=True, + check=True, + ) + assert result.stderr == "" + return result.stdout.splitlines() + + # First, so a user's -Xmx or -XX:MaxRAM later on the command line wins. + assert run(tool, "-Xmx64m", "Main") == [flag, "-Xmx64m", "Main"] + assert run(tool, "Main", limit=None) == ["Main"] + assert run(tool, "Main", limit="") == ["Main"] + # A tool that takes no -J options is left alone. + assert run("jfr", "summary") == ["summary"] + + +def test_java_toolchain_exports_the_limit_and_reserves_it() -> None: + from agent_sandbox_runtime.toolchains import java_toolchain + + toolchain = java_toolchain(java_home="/jdk", jvm_max_ram_bytes=1024) + assert toolchain.environment["SANDBOX_JVM_MAX_RAM"] == "1024" + assert "SANDBOX_JVM_MAX_RAM" in toolchain.reserved_names + assert "SANDBOX_JVM_MAX_RAM" not in java_toolchain(java_home="/jdk").environment