diff --git a/src/agent_env/explorer/ui/src/lib/collected-zip-files.smoke.ts b/src/agent_env/explorer/ui/src/lib/collected-zip-files.smoke.ts index 8d172454b..60f700d9b 100644 --- a/src/agent_env/explorer/ui/src/lib/collected-zip-files.smoke.ts +++ b/src/agent_env/explorer/ui/src/lib/collected-zip-files.smoke.ts @@ -52,6 +52,36 @@ function main(): void { assert(byUri['fake://home/c'] === 'elsewhere/c.txt', 'an absolute key elsewhere loses its leading slash'); } + // No member name climbs out of the archive or names a root or a drive, whatever the key holds. + { + const files = collectedZipFiles( + { + '../../x': 'fake://home/x', + '/abs/../../y': 'fake://home/y', + 'a/../../../z': 'fake://home/z', + '..\\..\\w.txt': 'fake://home/w', + 'C:\\Users\\op\\.ssh\\v': 'fake://home/v', + 'C:evil': 'fake://home/e', + './out/./u.txt': 'fake://home/u', + '..': 'fake://home/dots', + }, + '/app', + ); + const byUri = Object.fromEntries(files.map(f => [f.objectUrl, f.path])); + assert(byUri['fake://home/x'] === 'x', 'leading .. segments are dropped'); + assert(byUri['fake://home/y'] === 'y', 'an absolute key outside base_path climbs no higher than the root'); + assert(byUri['fake://home/z'] === 'z', 'a .. past the start is dropped'); + assert(byUri['fake://home/w'] === 'w.txt', 'backslashes separate like slashes'); + assert(byUri['fake://home/v'] === 'Users/op/.ssh/v', 'a drive prefix is dropped'); + assert(byUri['fake://home/e'] === 'evil', 'a drive-relative prefix is dropped'); + assert(byUri['fake://home/u'] === 'out/u.txt', '. segments are dropped'); + assert(byUri['fake://home/dots'] === 'artifact', 'a key with nothing left is named artifact'); + assert( + files.every(f => !f.path.split('/').some(part => part === '..' || part === '') && !/^[A-Za-z]:/.test(f.path)), + 'every member name is a clean relative path', + ); + } + if (failures > 0) { console.error(`\n${failures} assertion(s) failed`); process.exit(1); diff --git a/src/agent_env/explorer/ui/src/lib/collected-zip-files.ts b/src/agent_env/explorer/ui/src/lib/collected-zip-files.ts index 68bebbf3a..adca36447 100644 --- a/src/agent_env/explorer/ui/src/lib/collected-zip-files.ts +++ b/src/agent_env/explorer/ui/src/lib/collected-zip-files.ts @@ -18,11 +18,19 @@ export function collectedZipFiles( return Object.entries(artifacts) .filter(([, uri]) => typeof uri === 'string' && uri !== '') .map(([key, uri]) => ({ - path: !isAbsoluteArtifactPath(key) - ? key - : key.startsWith(`${base}/`) - ? key.slice(base.length + 1) - : key.replace(/^[/\\]+/, ''), + path: archiveName( + !isAbsoluteArtifactPath(key) ? key : key.startsWith(`${base}/`) ? key.slice(base.length + 1) : key, + ), objectUrl: uri, })); } + +/** A key as a member name that stays inside the archive when extracted: an agent names its manifest's entries. */ +export function archiveName(key: string): string { + const parts: string[] = []; + for (const part of key.replace(/^[A-Za-z]:/, '').split(/[\\/]+/)) { + if (part === '..') parts.pop(); + else if (part !== '' && part !== '.') parts.push(part); + } + return parts.join('/') || 'artifact'; +} diff --git a/src/agent_env/providers/sandbox_providers/local_docker/provider.py b/src/agent_env/providers/sandbox_providers/local_docker/provider.py index e9007c4a1..96125c338 100644 --- a/src/agent_env/providers/sandbox_providers/local_docker/provider.py +++ b/src/agent_env/providers/sandbox_providers/local_docker/provider.py @@ -312,6 +312,10 @@ def owns_container(self) -> bool: def work_dir(self) -> Path: return self._work_dir + @property + def app_dir_on_this_machine(self) -> Path: + return self._work_dir + def _rewrite_app_arg(self, arg: str) -> str: if arg == "/app" or arg.startswith(("/app/", "/app:")): return str(self._work_dir) + arg[4:] diff --git a/src/agent_env/providers/sandbox_providers/modal_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_sandbox.py index ac52f2d32..c7a2f7375 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -133,6 +133,9 @@ def __init__(self, stream): async def read(self) -> bytes: return await self._stream.read.aio() + def __aiter__(self): + return self._stream.__aiter__() + class _ModalProcessAdapter: def __init__(self, process): diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 642bfe4db..73dadc7ce 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -173,6 +173,9 @@ class Sandbox(ABC): platform: str | None = DEFAULT_BUILD_PLATFORM # The host other containers on its provider's private network reach it at, ready for a URL; None when it's on none. private_host: str | None = None + # The folder of this machine its host commands see as /app, when they run on this machine as a local sandbox's do; + # None when they run on a host of their own. + app_dir_on_this_machine: Path | None = None def url_from_sandbox(self, url: str) -> str: """``url``, of a service this process reaches, as this sandbox reaches it.""" diff --git a/src/agent_env/task_step/task_steps/collect_artifacts.py b/src/agent_env/task_step/task_steps/collect_artifacts.py index c55923a6a..1c204959c 100644 --- a/src/agent_env/task_step/task_steps/collect_artifacts.py +++ b/src/agent_env/task_step/task_steps/collect_artifacts.py @@ -16,7 +16,8 @@ Artifact paths come from one of these sources (in priority order): 1. Manifest (`manifest_step_id`): a prior prompt step's declared artifacts from that step's PromptResponse.structured_output["artifacts"] (logical name -> abs path); - collected files are keyed by logical name. + collected files are keyed by logical name, its separators and dot segments dropped. The + agent writes the manifest, so an entry outside base_path is skipped. 2. Seed data: context.metadata["seed"][artifacts_key] — a comma-separated list of filenames (e.g., "Dockerfile,run.sh,codebase.patch") 3. Static fallback: artifact_paths list in the step definition @@ -25,7 +26,11 @@ read verbatim). Absolute entries let a task collect deliverables that aren't under a single base_path — outputs aren't guaranteed to live on the desktop / in /app — without assuming a common directory. base_path only applies to -relative entries and to the no-list enumeration fallback. +relative entries, to manifest entries and to the no-list enumeration fallback. + +Only regular files are read, up to `max_bytes` each. A local sandbox's host is +this machine, so there a file is read only where it resolves inside the +sandbox's work dir, or inside a base_path the task names outside /app. This allows the same step definition to collect different artifacts per seed. @@ -67,12 +72,16 @@ from __future__ import annotations import asyncio +import base64 +import contextlib import functools import logging import os import json +import posixpath import re import shlex +import stat import tempfile import time from typing import ClassVar, Optional @@ -98,6 +107,11 @@ # or under this size take that single-shot path; larger ones stream in fixed byte ranges via # cua_bash so no single call has to move the whole file. _CUA_SINGLE_SHOT_MAX = 16 * 1024 * 1024 # 16 MB +# The largest file a step collects unless it sets max_bytes. +_DEFAULT_MAX_BYTES = 16 << 30 +# A read is abandoned past this floor plus a second for each _READ_RATE_FLOOR bytes of the file. +_READ_DEADLINE_FLOOR_S = 300 +_READ_RATE_FLOOR = 1 << 20 _CUA_CHUNK_BYTES = 8 * 1024 * 1024 # 8 MB per ranged read # artifact_paths entries with a URL scheme are not VM files — they're gold/source @@ -239,6 +253,7 @@ def __init__( sandbox_name: Optional[str] = None, container_name: Optional[str] = None, universe_id_suffix: Optional[str] = None, + max_bytes: int = _DEFAULT_MAX_BYTES, depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = True, ): @@ -283,6 +298,7 @@ def __init__( self.exclude_basenames = set(exclude_basenames or []) self.sandbox_name = sandbox_name self.container_name = container_name + self.max_bytes = max_bytes def to_dict(self) -> dict: base = super().to_dict() @@ -296,6 +312,7 @@ def to_dict(self) -> dict: base["sandbox_name"] = self.sandbox_name base["container_name"] = self.container_name base["universe_id_suffix"] = self.universe_id_suffix + base["max_bytes"] = self.max_bytes return base @classmethod @@ -312,6 +329,7 @@ def from_dict(cls, data: dict) -> CollectArtifactsTaskStep: sandbox_name=data.get("sandbox_name"), container_name=data.get("container_name"), universe_id_suffix=data.get("universe_id_suffix"), + max_bytes=data.get("max_bytes", _DEFAULT_MAX_BYTES), ) def _resolve_filenames(self, context: TaskStepContext) -> list[str]: @@ -352,20 +370,25 @@ def _resolve_items(self, context: TaskStepContext) -> list[tuple[str, str, str]] manifest = (structured.get("artifacts") if isinstance(structured, dict) else None) or {} items: list[tuple[str, str, str]] = [] skipped_nonfile: list[str] = [] + base = posixpath.join(posixpath.normpath(self.base_path or "/"), "") + keys: set[str] = set() for name, path in manifest.items(): if not (isinstance(path, str) and path.startswith("/")): skipped_nonfile.append(name) # literal value (e.g. coverage_state) or relative path continue if "*" in path: continue # unresolved glob (e.g. per-iteration debug logs) — can't fetch a literal * + # The agent writes the manifest, so only a file under base_path, once normalized, is read. + path = posixpath.normpath(path) + if not path.startswith(base): + logger.warning(f"collect '{self.id}': manifest entry {name!r} names {path}, outside " + f"{self.base_path}; not collected") + continue if os.path.basename(path) in self.exclude_basenames: continue - rel = ( - path[len(self.base_path) + 1:] - if path.startswith(self.base_path + "/") - else os.path.basename(path) - ) - items.append((name, path, rel)) + key = _unused(_manifest_key(name, path), keys) + keys.add(key) + items.append((key, path, path[len(base):])) if skipped_nonfile: logger.debug( "collect(manifest=%s): skipped non-absolute manifest entries " @@ -445,30 +468,34 @@ async def _resolve_live_sandbox(self, provider, sandbox_id: str): ) from err async def _get_file_size(self, sandbox, container: Optional[str], source_path: str) -> int: - """Get file size on the agent's filesystem. Returns -1 only if the file is genuinely - absent; a `stat` that fails for any other reason RAISES so a failed size check is a - hard failure rather than a silent 'missing' skip that drops an existing deliverable. + """Size of the regular file at ``source_path`` on the agent's filesystem, links followed. Returns -1 only if the + file is genuinely absent; a `stat` that fails for any other reason RAISES so a failed size check is a hard + failure rather than a silent 'missing' skip that drops an existing deliverable, and so does anything but a + regular file (a device, a FIFO, a directory, or a link to one), which could stream without end. VM-mode wraps in `docker exec `; container-mode runs `stat` directly.""" - args = _exec_args(sandbox, container, ("stat", "-c", "%s", source_path)) + args = _exec_args(sandbox, container, ("stat", "-L", "-c", "%f:%s", source_path)) exit_code, stdout, stderr = await sandbox.exec_with_output(*args) if exit_code != 0: if _looks_absent(stderr): return -1 raise RuntimeError(f"could not size {source_path} (stat exit {exit_code}): {stderr[:200]}") - return int(stdout.strip()) + mode, _, size = stdout.strip().rpartition(":") + if not stat.S_ISREG(int(mode, 16)): + raise RuntimeError(f"{source_path} is not a regular file ({stat.filemode(int(mode, 16))}); refusing to read it") + return int(size) async def _collect_file( - self, sandbox, container: Optional[str], source_path: str, object_name: str, store, artifact_id: str, - version: int, content_type: str, + self, sandbox, container: Optional[str], source_path: str, size: int, object_name: str, store, + artifact_id: str, version: int, content_type: str, ) -> str: - """Read a file out of the sandbox and upload it to S3 — one path, any size. + """Read a file of ``size`` bytes out of the sandbox and upload it to S3. base64-encode in the sandbox (raw binary over a websocket exec transport is silently dropped; ASCII survives), stream stdout + decode incrementally to a temp file (bounds memory), then multipart-upload. stderr is drained in - parallel so the shared output pump can't backpressure-stall.""" - import base64 - bash_cmd = f"base64 < {shlex.quote(source_path)}" + parallel so the shared output pump can't backpressure-stall. The sandbox reads + no more than ``size`` bytes, and a read that outlasts its deadline is abandoned.""" + bash_cmd = f"head -c {size} < {shlex.quote(source_path)} | base64" args = _exec_args(sandbox, container, ("bash", "-c", bash_cmd)) local_path = None @@ -477,30 +504,11 @@ async def _collect_file( with tempfile.NamedTemporaryFile(delete=False, suffix=f"-{object_name.replace('/', '_')}") as tmp: local_path = tmp.name - process = await sandbox.exec(*args) - stderr_task = asyncio.create_task(process.stderr.read()) - stdout = process.stdout - with open(local_path, "wb") as fh: - if hasattr(stdout, "__aiter__"): - # Async-iterable stdout (websocket exec): stream, decoding 4-char-aligned base64. - pending = bytearray() - async for chunk in stdout: - pending += chunk.translate(None, b"\r\n") # strip 76-col wrap - n = (len(pending) // 4) * 4 - if n: - fh.write(base64.b64decode(pending[:n])) - del pending[:n] - if pending: - fh.write(base64.b64decode(pending)) - else: - # Modal stdout has read() but no __aiter__: read all, decode once. - data = (await stdout.read()).translate(None, b"\r\n") - fh.write(base64.b64decode(data)) - exit_code = await process.wait() + deadline = _READ_DEADLINE_FLOOR_S + size / _READ_RATE_FLOOR try: - await stderr_task - except Exception: - pass + exit_code = await asyncio.wait_for(_read_to_file(sandbox, args, local_path, size), timeout=deadline) + except TimeoutError: + raise RuntimeError(f"reading {source_path} took more than {deadline:.0f}s; abandoned it") from None if exit_code != 0: raise RuntimeError(f"base64 read failed for {source_path} (exit {exit_code})") @@ -558,9 +566,9 @@ async def _controller_file_size(self, deployed_env, path: str) -> int: as a missing file and silently dropping an existing deliverable. ``wc -c`` is portable across the Ubuntu and macOS desktop VMs; ``stat``'s flags are not. A transport exception from ``_controller_call`` propagates unchanged.""" - out = await self._controller_call( - deployed_env, "cua_bash", {"script": f"wc -c < {shlex.quote(path)}"} - ) + quoted = shlex.quote(path) + script = f'[ ! -e {quoted} ] || [ -f {quoted} ] || {{ echo "not a regular file" >&2; exit 1; }}; wc -c < {quoted}' + out = await self._controller_call(deployed_env, "cua_bash", {"script": script}) # A present file (even 0 bytes) prints its size to stdout; empty stdout means wc failed. text = _cua_bash_stdout(out).strip() if text: @@ -581,11 +589,11 @@ async def _controller_get_file(self, deployed_env, path: str) -> str: lenient skip). Small files come back in one ``cua_get_file`` call; large ones stream in byte ranges straight to disk, so neither the transfer nor this process ever holds the whole file in memory.""" - import base64 - size = await self._controller_file_size(deployed_env, path) if size < 0: raise FileNotFoundError(f"file not found on CUA VM: {path}") + if size > self.max_bytes: + raise RuntimeError(f"{path} is {size} bytes, over the step's max_bytes of {self.max_bytes}") with tempfile.NamedTemporaryFile(delete=False, suffix=f"-{os.path.basename(path)}") as tmp: local_path = tmp.name try: @@ -607,8 +615,6 @@ async def _controller_stream_to_file(self, deployed_env, path: str, size: int, l than scans. Each range is decoded and written straight to the file so memory stays bounded to one chunk. base64 may wrap, so whitespace is stripped before decoding. A short read fails loudly rather than leaving a truncated artifact.""" - import base64 - logger.info( f"Streaming {path} ({size} bytes) off the CUA VM in {_CUA_CHUNK_BYTES}-byte chunks" ) @@ -920,6 +926,8 @@ async def _collect_items(self, provider, sandbox, container, items, context, sto else: logger.info(f"Collecting {len(items)} artifacts from {self.base_path}") + # A local sandbox's host is this machine: only what resolves inside its work dir, or a base_path outside /app. + roots = await self._local_roots(sandbox) if container is None and sandbox.app_dir_on_this_machine else None collected: dict[str, str] = {} file_artifacts: dict[str, FileArtifact] = {} hard_failures: list[str] = [] @@ -928,16 +936,19 @@ async def _collect_items(self, provider, sandbox, container, items, context, sto # silent skip) and provider.close() below still runs; only a genuinely-absent # file (size -1) is the lenient skip. try: - file_size = await self._get_file_size(sandbox, container, source_path) + read_path = source_path if roots is None else await _resolved_inside(sandbox, source_path, roots) + file_size = -1 if read_path is None else await self._get_file_size(sandbox, container, read_path) if file_size < 0: logger.warning(f"File not found: {source_path}") continue + if file_size > self.max_bytes: + raise RuntimeError(f"{source_path} is {file_size} bytes, over the step's max_bytes of {self.max_bytes}") ext = os.path.splitext(object_name)[1] content_type = _CONTENT_TYPES.get(ext, "application/octet-stream") object_url = await self._collect_file( - sandbox, container, source_path, object_name, store, artifact_id, version, content_type, + sandbox, container, read_path, file_size, object_name, store, artifact_id, version, content_type, ) collected[key] = object_url @@ -961,6 +972,14 @@ async def _collect_items(self, provider, sandbox, container, items, context, sto await provider.close() return collected, file_artifacts, hard_failures + async def _local_roots(self, sandbox) -> tuple[str, ...]: + """The sandbox's work dir, and base_path where the task names a folder of this machine outside /app: under + /app the agent may have swapped a folder for a link, so only the work dir it stands for counts.""" + base = posixpath.normpath(self.base_path or "/") + outside_app = base != "/app" and not base.startswith("/app/") + roots = (str(sandbox.app_dir_on_this_machine.resolve()), await _resolved(sandbox, base) if outside_app else None) + return tuple(root.rstrip("/") + "/" for root in roots if root) + def _register_file_artifact(self, store, file_artifacts, artifact_id, key, object_name, content_type, object_url, context): """Register a FileArtifact document pointing at the already-uploaded S3 object. We do NOT re-upload the bytes — FileArtifact.object_url is just a URL @@ -1050,3 +1069,77 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: ) return context + + +def _manifest_key(name: str, path: str) -> str: + """The agent's name for a manifest entry as a key no reader takes for a path: its separators and dot segments + dropped, else the file's own name. Consumers write files under these keys, a zip of them among them.""" + parts = [part for part in re.split(r"[\\/]+", name) if part not in ("", ".", "..")] + return "_".join(parts) or posixpath.basename(path) + + +def _unused(key: str, keys: set[str]) -> str: + """``key``, or the first ``key (n)`` not in ``keys``: two manifest names can clean to the same key.""" + n = 1 + unused = key + while unused in keys: + n += 1 + unused = f"{key} ({n})" + return unused + + +async def _read_to_file(sandbox, args: tuple, local_path: str, size: int) -> int: + """Run ``args`` and decode its base64 stdout into ``local_path``, draining stderr beside it so the shared output + pump can't backpressure-stall; return its exit status. Some providers run the whole command inside ``exec``.""" + process = await sandbox.exec(*args) + stderr = asyncio.create_task(process.stderr.read()) + try: + exit_code = await _decode_to_file(process, local_path, size) + except BaseException: + stderr.cancel() + raise + with contextlib.suppress(Exception): + await stderr + return exit_code + + +async def _decode_to_file(process, local_path: str, size: int) -> int: + """Decode ``process``'s base64 stdout into ``local_path`` and return its exit status, refusing a byte past ``size``.""" + written = 0 + + def write(fh, data: bytes) -> None: + nonlocal written + written += len(data) + if written > size: + raise RuntimeError(f"the sandbox sent more than the {size} bytes it reported; refusing the rest") + fh.write(data) + + stdout = process.stdout + with open(local_path, "wb") as fh: + if hasattr(stdout, "__aiter__"): + # Async-iterable stdout (websocket exec): stream, decoding 4-char-aligned base64. + pending = bytearray() + async for chunk in stdout: + pending += chunk.translate(None, b"\r\n") # strip 76-col wrap + n = (len(pending) // 4) * 4 + if n: + write(fh, base64.b64decode(pending[:n])) + del pending[:n] + if pending: + write(fh, base64.b64decode(pending)) + else: + # A provider whose exec hands back its output whole: decode it once. The sandbox reads at most `size` bytes. + write(fh, base64.b64decode((await stdout.read()).translate(None, b"\r\n"))) + return await process.wait() + + +async def _resolved(sandbox, path: str) -> str | None: + exit_code, stdout, _ = await sandbox.exec_with_output("readlink", "-f", path) + return (stdout.strip() or None) if exit_code == 0 else None + + +async def _resolved_inside(sandbox, path: str, roots: tuple[str, ...]) -> str | None: + resolved = await _resolved(sandbox, path) + if resolved is not None and not (resolved + "/").startswith(roots): + raise RuntimeError(f"{path} resolves to {resolved}, outside the sandbox's work dir and base_path on this machine") + return resolved diff --git a/tst/unit/providers/sandbox_providers/modal_sandbox_test.py b/tst/unit/providers/sandbox_providers/modal_sandbox_test.py index 5c745ec58..eb962ec01 100644 --- a/tst/unit/providers/sandbox_providers/modal_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/modal_sandbox_test.py @@ -488,3 +488,17 @@ async def test_get_sandbox_still_raises_for_a_finished_sandbox_so_a_reattach_see from_id.aio = AsyncMock(return_value=sb) with pytest.raises(modal.exception.ConflictError): await _patched_provider.get_sandbox("sb-finished") + + +@pytest.mark.asyncio +async def test_exec_output_streams_as_modal_yields_it(): + class _Stream: + async def _chunks(self): + for chunk in (b"QUJD", b"REVG"): + yield chunk + + def __aiter__(self): + return self._chunks() + + assert [chunk async for chunk in modal_sandbox._ModalStreamAdapter(_Stream())] == [b"QUJD", b"REVG"] + diff --git a/tst/unit/providers/sandbox_providers/provider_capabilities_test.py b/tst/unit/providers/sandbox_providers/provider_capabilities_test.py index cf738e2fc..d66fa994b 100644 --- a/tst/unit/providers/sandbox_providers/provider_capabilities_test.py +++ b/tst/unit/providers/sandbox_providers/provider_capabilities_test.py @@ -299,6 +299,17 @@ def test_a_local_sandbox_runs_any_platform_unchecked_since_its_docker_emulates_o assert LocalSandbox.platform is None +@pytest.mark.parametrize("sandbox", [Sandbox, ModalSandbox, ModalVmSandbox, E2BSandbox, SailVmSandbox, TensorlakeSandbox, + TenkiSandbox, BoxdSandbox, CoreWeaveSandbox, DaytonaSandbox, BoatSandbox, + Agent37Sandbox]) +def test_a_sandbox_runs_its_host_commands_on_a_host_of_its_own(sandbox): + assert sandbox.app_dir_on_this_machine is None + + +def test_a_local_sandbox_runs_its_host_commands_on_this_machine_in_its_work_dir(tmp_path): + assert LocalSandbox(work_dir=tmp_path).app_dir_on_this_machine == tmp_path + + def _concrete_classes() -> set[str]: """The sandbox and sandbox provider classes the providers define, but the base classes everything else may name.""" names = set() diff --git a/tst/unit/store/fakes.py b/tst/unit/store/fakes.py index a308272e2..ecfa09a6a 100644 --- a/tst/unit/store/fakes.py +++ b/tst/unit/store/fakes.py @@ -249,6 +249,7 @@ class CollectingVm: mode = "vm" type = None + app_dir_on_this_machine = None def scoped_name(self, name): return name @@ -258,11 +259,11 @@ async def exec_script(self, script: str) -> str: async def exec_with_output(self, *args): if "stat" in args: - return 0, str(len(COLLECTED[args[-1]])), "" + return 0, f"81a4:{len(COLLECTED[args[-1]])}", "" return 0, "ok", "" async def exec(self, *args): - path = args[-1].split("base64 < ", 1)[1].strip("'") + path = args[-1].split(" < ", 1)[1].split(" | ", 1)[0].strip("'") return SimpleNamespace(stdout=_Stream(base64.b64encode(COLLECTED[path])), stderr=_Stream(b""), wait=lambda: asyncio.sleep(0, result=0)) diff --git a/tst/unit/task_step/test_collect_artifacts.py b/tst/unit/task_step/test_collect_artifacts.py index 31bc1ad76..3c2f47b6c 100644 --- a/tst/unit/task_step/test_collect_artifacts.py +++ b/tst/unit/task_step/test_collect_artifacts.py @@ -11,6 +11,7 @@ import asyncio import json import os +import subprocess import tempfile import threading from unittest.mock import AsyncMock, MagicMock, patch @@ -81,6 +82,11 @@ def test_env_id_defaults_none(self): }) assert step.env_id is None + def test_max_bytes_round_trips_and_defaults(self): + step = CollectArtifactsTaskStep.from_dict({"id": "c", "type": "collect_artifacts", "max_bytes": 4096}) + assert CollectArtifactsTaskStep.from_dict(step.to_dict()).max_bytes == 4096 + assert CollectArtifactsTaskStep.from_dict({"id": "c", "type": "collect_artifacts"}).max_bytes == 16 << 30 + def test_explicit_agent_name_and_env_id_raises(self): try: CollectArtifactsTaskStep( @@ -516,6 +522,43 @@ def test_get_file_size_other_error_raises(self): except RuntimeError as e: assert "could not size" in str(e) + @pytest.mark.parametrize("mode, shown", [("21b6", "crw-rw-rw-"), ("11a4", "prw-r--r--"), ("41ed", "drwxr-xr-x"), + ("c1ed", "srwxr-xr-x")], ids=["device", "fifo", "directory", "socket"]) + def test_get_file_size_refuses_anything_but_a_regular_file(self, mode, shown): + """By the raw mode bits stat prints, which no locale translates, as it does a file type's name.""" + step = CollectArtifactsTaskStep(id="c", version=None, agent_name="solver") + sandbox = MagicMock() + sandbox.mode = "container" + sandbox.exec_with_output = AsyncMock(return_value=(0, f"{mode}:0\n", "")) + with pytest.raises(RuntimeError, match=f"not a regular file \\({shown}\\)"): + _run(step._get_file_size(sandbox, None, "/app/artifact/report.pdf")) + assert sandbox.exec_with_output.await_args.args[:4] == ("stat", "-L", "-c", "%f:%s") + + def test_get_file_size_sizes_a_regular_file_even_when_empty(self): + step = CollectArtifactsTaskStep(id="c", version=None, agent_name="solver") + sandbox = MagicMock() + sandbox.mode = "container" + sandbox.exec_with_output = AsyncMock(side_effect=[(0, "81a4:0\n", ""), (0, "8180:4096\n", "")]) + assert [_run(step._get_file_size(sandbox, None, "/x")) for _ in range(2)] == [0, 4096] + + def test_a_file_over_max_bytes_is_a_hard_failure_and_is_not_read(self): + step = CollectArtifactsTaskStep( + id="c", version=None, base_path="/app/artifact", artifact_paths=["disk.img"], max_bytes=1024, + ) + sandbox = MagicMock() + sandbox.mode = "container" + sandbox.app_dir_on_this_machine = None + sandbox.exec_with_output = AsyncMock(return_value=(0, f"81a4:{8 << 40}\n", "")) + sandbox.exec = AsyncMock() + provider = MagicMock() + provider.close = AsyncMock() + ctx = TaskStepContext() + collected, _, hard_failures = _run(step._collect_items( + provider, sandbox, None, step._resolve_items(ctx), ctx, MagicMock(), "aid", 1, + )) + assert (collected, hard_failures) == ({}, ["/app/artifact/disk.img"]) + sandbox.exec.assert_not_awaited() + def test_unreachable_sandbox_is_clear_error(self): """A re-run targeting an expired/unreachable sandbox fails with a clear message.""" step = CollectArtifactsTaskStep( @@ -599,6 +642,49 @@ def test_manifest_mode_keys_by_logical_name_and_filters(self): assert items["gold_patch"] == ("/app/workspace/patches/gold_patch.diff", "patches/gold_patch.diff") assert items["dockerfile"] == ("/app/workspace/Dockerfile", "Dockerfile") + def test_manifest_entries_outside_base_path_are_not_read(self, caplog): + step = CollectArtifactsTaskStep(id="collect", version=None, base_path="/app/workspace", manifest_step_id="m9") + ctx = _manifest_ctx({ + "report": "/home/op/.ssh/id_rsa", + "climb": "/app/workspace/../../etc/passwd", + "sibling": "/app/workspace-other/notes.md", + "kept": "/app/workspace/out/./report.md", + }) + with caplog.at_level("WARNING"): + items = step._resolve_items(ctx) + assert items == [("kept", "/app/workspace/out/report.md", "out/report.md")] + assert "names /home/op/.ssh/id_rsa, outside /app/workspace; not collected" in caplog.text + assert "names /etc/passwd, outside /app/workspace" in caplog.text + + def test_manifest_names_become_keys_with_no_separators_or_dot_segments(self): + step = CollectArtifactsTaskStep(id="collect", version=None, base_path="/app/workspace", manifest_step_id="m9") + ctx = _manifest_ctx({ + "../../.bashrc": "/app/workspace/payload", + "logs/run.json": "/app/workspace/logs/run.json", + "..\\..\\evil": "/app/workspace/evil", + "..": "/app/workspace/dots.txt", + }) + assert [k for k, _, _ in step._resolve_items(ctx)] == [".bashrc", "logs_run.json", "evil", "dots.txt"] + + def test_manifest_names_that_clean_to_one_key_each_keep_a_key_of_their_own(self): + step = CollectArtifactsTaskStep(id="collect", version=None, base_path="/app/workspace", manifest_step_id="m9") + ctx = _manifest_ctx({ + "logs/run.json": "/app/workspace/a.json", + "logs_run.json": "/app/workspace/b.json", + "../logs/run.json": "/app/workspace/c.json", + }) + assert [(k, src) for k, src, _ in step._resolve_items(ctx)] == [ + ("logs_run.json", "/app/workspace/a.json"), + ("logs_run.json (2)", "/app/workspace/b.json"), + ("logs_run.json (3)", "/app/workspace/c.json"), + ] + + @pytest.mark.parametrize("base_path", ["/app/./workspace", "/app/workspace/", "/app//workspace"]) + def test_manifest_entries_are_matched_against_base_path_normalized(self, base_path): + step = CollectArtifactsTaskStep(id="collect", version=None, base_path=base_path, manifest_step_id="m9") + ctx = _manifest_ctx({"report": "/app/workspace/out/report.md"}) + assert step._resolve_items(ctx) == [("report", "/app/workspace/out/report.md", "out/report.md")] + def test_enumeration_drops_excluded_basenames(self): step = CollectArtifactsTaskStep( id="collect", version=None, base_path="/app/workspace", @@ -936,6 +1022,44 @@ async def call(deployed_env, tool_name, arguments): except RuntimeError as e: assert "refusing to upload a partial artifact" in str(e) + def test_a_file_over_max_bytes_is_refused_before_any_read(self): + step = CollectArtifactsTaskStep(id="c", version=None, env_id="ubuntu-cua", max_bytes=1024) + calls = [] + + async def call(deployed_env, tool_name, arguments): + calls.append(tool_name) + return json.dumps({"output": f"{8 << 40}\n", "error": "", "returncode": 0}) + + with patch.object(step, "_controller_call", side_effect=call), \ + pytest.raises(RuntimeError, match="over the step's max_bytes of 1024"): + _run(step._controller_get_file(_cua_record(), "/x/sparse.img")) + assert calls == ["cua_bash"] + + @pytest.mark.parametrize("make, size", [ + (lambda p, tmp: p.write_bytes(b"hello"), 5), + (lambda p, tmp: None, -1), + (lambda p, tmp: p.symlink_to(tmp / "gone"), -1), + (lambda p, tmp: p.symlink_to("/dev/zero"), None), + (lambda p, tmp: os.mkfifo(p), None), + (lambda p, tmp: p.mkdir(), None), + ], ids=["regular", "absent", "dangling-link", "link-to-dev-zero", "fifo", "directory"]) + def test_the_size_probe_sizes_only_a_regular_file(self, tmp_path, make, size): + """The probe's own script, run by a real shell as cua_bash runs it on the VM.""" + target = tmp_path / "deliverable" + make(target, tmp_path) + step = CollectArtifactsTaskStep(id="c", version=None, env_id="ubuntu-cua") + + async def call(deployed_env, tool_name, arguments): + done = subprocess.run(["bash", "-c", arguments["script"]], capture_output=True, text=True, timeout=10) + return json.dumps({"output": done.stdout, "error": done.stderr, "returncode": done.returncode}) + + with patch.object(step, "_controller_call", side_effect=call): + if size is None: + with pytest.raises(RuntimeError, match="not a regular file"): + _run(step._controller_file_size(_cua_record(), str(target))) + else: + assert _run(step._controller_file_size(_cua_record(), str(target))) == size + def test_absent_file_raises_filenotfound(self): """A missing file (wc errors with 'No such file') is reported absent so the caller skips it leniently.""" diff --git a/tst/unit/task_step/test_collect_artifacts_universe_suffix.py b/tst/unit/task_step/test_collect_artifacts_universe_suffix.py index f4105ad9a..2227655ae 100644 --- a/tst/unit/task_step/test_collect_artifacts_universe_suffix.py +++ b/tst/unit/task_step/test_collect_artifacts_universe_suffix.py @@ -38,6 +38,7 @@ async def wait(self): class _VmSandbox: mode = "vm" + app_dir_on_this_machine = None async def exec_with_output(self, *args): if "find" in args: @@ -45,7 +46,7 @@ async def exec_with_output(self, *args): return 0, "bundle\n", "" return 0, "".join(f"{p}\n" for p in FILES), "" if "stat" in args: - return 0, str(len(next(iter(FILES.values())))), "" + return 0, f"81a4:{len(next(iter(FILES.values())))}", "" return 0, "", "" async def exec(self, *args): diff --git a/tst/unit/task_step/test_collect_artifacts_vm_host.py b/tst/unit/task_step/test_collect_artifacts_vm_host.py index 945b06c1e..a5ac205bc 100644 --- a/tst/unit/task_step/test_collect_artifacts_vm_host.py +++ b/tst/unit/task_step/test_collect_artifacts_vm_host.py @@ -3,11 +3,14 @@ from __future__ import annotations import base64 +from unittest.mock import AsyncMock, MagicMock import pytest +from agent_env.providers.sandbox_providers.local_docker.provider import LocalSandbox from agent_env.task_step.context import DeployedAgent, DeployedSandbox, TaskStepContext from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep +from tst.util.capabilities import skip_without_gnu_stat FILES = {"caa_10/MODIFICATIONS.md": b"# what changed\n", "caa_10/alignment_x/task.toml": b"schema=1\n"} @@ -33,6 +36,7 @@ class _VmSandbox: """A VM-mode sandbox that serves the fake tree, recording every command.""" mode = "vm" + app_dir_on_this_machine = None def __init__(self): self.calls: list[tuple] = [] @@ -49,14 +53,14 @@ async def exec_with_output(self, *args): if "stat" in args: path = args[-1] rel = path.split("/app/artifact/", 1)[-1] - return (0, str(len(FILES[rel])), "") if rel in FILES else (-1, "", "no such file") + return (0, f"81a4:{len(FILES[rel])}", "") if rel in FILES else (-1, "", "no such file") return 0, "", "" async def exec(self, *args): self.calls.append(args) assert args[0] == "sudo", f"VM-host command must run as root: {args}" bash_cmd = args[-1] - path = bash_cmd.split("base64 < ", 1)[1].strip().strip("'") + path = bash_cmd.split(" < ", 1)[1].split(" | ", 1)[0].strip("'") rel = path.split("/app/artifact/", 1)[-1] return _Proc(base64.b64encode(FILES[rel])) @@ -139,7 +143,7 @@ async def test_enumeration_and_reads_use_no_container(self, vm): # sudo'd but not container-wrapped: host-mode agents write as uid 0 assert all(c.startswith(("sudo ", "echo ")) for c in flat), flat assert any(c.startswith("sudo find /app/artifact") for c in flat), flat - assert any(c.startswith("sudo bash -c base64 < ") for c in flat), flat + assert any(c.startswith("sudo bash -c head -c ") for c in flat), flat @pytest.mark.asyncio async def test_it_does_not_need_a_deployed_agent(self, vm): @@ -251,3 +255,75 @@ async def _fake_agent_collect(self, context, store, artifact_id, version): step = CollectArtifactsTaskStep(id="collect", version=None, agent_name="solver") await step.execute(ctx) assert seen.get("called") is True + + +@skip_without_gnu_stat() +class TestOnThisMachine: + """A local sandbox's host is the operator's: a file is read only where it resolves inside the sandbox's work dir + or the step's base_path, so a link left there reaches nothing else on this machine.""" + + @staticmethod + async def _collect(step, sandbox): + uploaded = {} + store = MagicMock() + + def put(**kw): + with open(kw["file_path"], "rb") as fh: + uploaded[kw["object_name"]] = fh.read() + return f"s3://bucket/{kw['object_name']}" + + store.put_object_file.side_effect = put + provider = MagicMock() + provider.close = AsyncMock() + ctx = TaskStepContext() + collected, _, hard_failures = await step._collect_items( + provider, sandbox, None, step._resolve_items(ctx), ctx, store, "aid", 1, + ) + return set(collected), hard_failures, uploaded + + @pytest.mark.asyncio + async def test_a_link_out_of_the_work_dir_is_refused(self, tmp_path): + work = tmp_path / "work" + (work / "artifact").mkdir(parents=True) + (work / "other").mkdir() + (tmp_path / "id_rsa").write_bytes(b"operator secret") + (work / "artifact" / "report.md").write_bytes(b"# report") + (work / "other" / "notes.md").write_bytes(b"notes") + (work / "artifact" / "leak").symlink_to(tmp_path / "id_rsa") + (work / "artifact" / "notes").symlink_to(work / "other" / "notes.md") + step = CollectArtifactsTaskStep( + id="c", version=None, sandbox_name="box", base_path="/app/artifact", + artifact_paths=["report.md", "leak", "notes", "/app/../id_rsa"], + ) + collected, hard_failures, uploaded = await self._collect(step, LocalSandbox(work_dir=work)) + assert collected == {"report.md", "notes"} + assert set(hard_failures) == {"/app/artifact/leak", "/app/../id_rsa"} + assert b"operator secret" not in uploaded.values() + + @pytest.mark.asyncio + async def test_a_base_path_the_task_names_on_this_machine_is_read(self, tmp_path): + out = tmp_path / "host" / "out" + out.mkdir(parents=True) + (tmp_path / "id_rsa").write_bytes(b"operator secret") + (out / "bundle.zip").write_bytes(b"PK bundle") + (out / "leak").symlink_to(tmp_path / "id_rsa") + step = CollectArtifactsTaskStep( + id="c", version=None, sandbox_name="box", base_path=str(out), artifact_paths=["bundle.zip", "leak"], + ) + collected, hard_failures, uploaded = await self._collect(step, LocalSandbox(work_dir=tmp_path / "work")) + assert (collected, hard_failures) == ({"bundle.zip"}, [f"{out}/leak"]) + assert uploaded == {"bundle.zip": b"PK bundle"} + + @pytest.mark.asyncio + async def test_a_base_path_the_agent_swapped_for_a_link_reaches_only_the_work_dir(self, tmp_path): + work, home = tmp_path / "work", tmp_path / "home" + (home / ".ssh").mkdir(parents=True) + (home / ".ssh" / "id_rsa").write_bytes(b"operator secret") + work.mkdir() + (work / "artifact").symlink_to(home) + step = CollectArtifactsTaskStep( + id="c", version=None, sandbox_name="box", base_path="/app/artifact", artifact_paths=[".ssh/id_rsa"], + ) + collected, hard_failures, uploaded = await self._collect(step, LocalSandbox(work_dir=work)) + assert (collected, hard_failures, uploaded) == (set(), ["/app/artifact/.ssh/id_rsa"], {}) + diff --git a/tst/unit/task_step/test_collect_integrity.py b/tst/unit/task_step/test_collect_integrity.py index e4266410e..26ff22c38 100644 --- a/tst/unit/task_step/test_collect_integrity.py +++ b/tst/unit/task_step/test_collect_integrity.py @@ -6,6 +6,7 @@ import threading import pytest +from agent_env.task_step.task_steps import collect_artifacts from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep from tst.unit.event_loop_probe import on_event_loop @@ -39,7 +40,7 @@ def __init__(self, full_b64, keep, src_size): async def exec(self, *args): return _Proc(_TruncStdout(self._full, self._keep)) async def exec_with_output(self, *args): - return (0, f"{self._src}\n", "") + return (0, f"81a4:{self._src}\n", "") @pytest.mark.asyncio @@ -55,7 +56,7 @@ def put_object_file(self, **kw): step = CollectArtifactsTaskStep(id="collect", version=None) with pytest.raises(RuntimeError, match="integrity check failed"): - await step._collect_file(sb, None, "/app/artifact/corpus.warc.gz", "corpus.warc.gz", + await step._collect_file(sb, None, "/app/artifact/corpus.warc.gz", len(src), "corpus.warc.gz", _Store(), "uid", 1, "application/octet-stream") @@ -72,7 +73,7 @@ def put_object_file(self, **kw): return "s3://bucket/ok" step = CollectArtifactsTaskStep(id="collect", version=None) - out = await step._collect_file(sb, None, "/app/artifact/corpus.warc.gz", "corpus.warc.gz", + out = await step._collect_file(sb, None, "/app/artifact/corpus.warc.gz", len(src), "corpus.warc.gz", _OkStore(), "uid", 1, "application/octet-stream") assert out == "s3://bucket/ok" @@ -90,7 +91,7 @@ def put_object_file(self, **kw): return "s3://bucket/ok" step = CollectArtifactsTaskStep(id="collect", version=None) - await step._collect_file(sb, None, "/app/artifact/a.bin", "a.bin", _Store(), "uid", 1, "application/octet-stream") + await step._collect_file(sb, None, "/app/artifact/a.bin", len(src), "a.bin", _Store(), "uid", 1, "application/octet-stream") assert on_loop == [False] @@ -111,7 +112,7 @@ def put_object_file(self, **kw): return "s3://bucket/ok" step = CollectArtifactsTaskStep(id="collect", version=None) - collect = asyncio.create_task(step._collect_file(sb, None, "/app/artifact/a.bin", "a.bin", _Store(), "uid", 1, "application/octet-stream")) + collect = asyncio.create_task(step._collect_file(sb, None, "/app/artifact/a.bin", len(src), "a.bin", _Store(), "uid", 1, "application/octet-stream")) await asyncio.to_thread(uploading.wait, 5) collect.cancel() with pytest.raises(asyncio.CancelledError): @@ -125,3 +126,83 @@ def put_object_file(self, **kw): assert uploaded == [src] assert "Uploading collected /app/artifact/a.bin finished after its caller was cancelled" in caplog.text assert list(tmp_path.iterdir()) == [] + + +class _HangingStdout: + def __aiter__(self): + return self + + async def __anext__(self): + await asyncio.Event().wait() + + +class _RecordingSandbox(_Sandbox): + def __init__(self, *args, stdout=None): + super().__init__(*args) + self.commands, self._stdout = [], stdout + + async def exec(self, *args): + self.commands.append(args) + return _Proc(self._stdout or _TruncStdout(self._full, self._keep)) + + +class _UploadingStore: + def put_object_file(self, **kw): + return "s3://bucket/ok" + + +class _RefusingStore: + def put_object_file(self, **kw): + raise AssertionError("uploaded an artifact the read should have refused") + + +@pytest.mark.asyncio +async def test_the_sandbox_reads_no_more_than_the_size_it_reported(): + src = b"F" * 900 + sb = _RecordingSandbox(base64.b64encode(src), len(base64.b64encode(src)), len(src)) + step = CollectArtifactsTaskStep(id="collect", version=None) + await step._collect_file(sb, None, "/app/artifact/a b.bin", len(src), "a b.bin", _UploadingStore(), "uid", 1, + "application/octet-stream") + assert sb.commands == [("bash", "-c", "head -c 900 < '/app/artifact/a b.bin' | base64")] + + +@pytest.mark.asyncio +async def test_a_stream_longer_than_the_reported_size_is_refused(tmp_path, monkeypatch): + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + src = b"G" * 900 + more = base64.b64encode(src * 3) + sb = _Sandbox(more, len(more), len(src)) + step = CollectArtifactsTaskStep(id="collect", version=None) + with pytest.raises(RuntimeError, match="more than the 900 bytes it reported"): + await step._collect_file(sb, None, "/app/artifact/zero", len(src), "zero", _RefusingStore(), "uid", 1, + "application/octet-stream") + assert list(tmp_path.iterdir()) == [] + + +@pytest.mark.asyncio +async def test_a_read_past_its_deadline_is_abandoned(tmp_path, monkeypatch): + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + monkeypatch.setattr(collect_artifacts, "_READ_DEADLINE_FLOOR_S", 0.05) + sb = _RecordingSandbox(b"", 0, 10, stdout=_HangingStdout()) + step = CollectArtifactsTaskStep(id="collect", version=None) + with pytest.raises(RuntimeError, match="took more than 0s; abandoned it"): + await step._collect_file(sb, None, "/app/artifact/slow", 10, "slow", _RefusingStore(), "uid", 1, + "application/octet-stream") + assert list(tmp_path.iterdir()) == [] + + +@pytest.mark.asyncio +async def test_a_provider_whose_exec_runs_the_whole_read_is_held_to_the_deadline(tmp_path, monkeypatch): + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + monkeypatch.setattr(collect_artifacts, "_READ_DEADLINE_FLOOR_S", 0.05) + + class _StallingExec(_Sandbox): + async def exec(self, *args): + await asyncio.Event().wait() + + step = CollectArtifactsTaskStep(id="collect", version=None) + with pytest.raises(RuntimeError, match="took more than 0s; abandoned it"): + await step._collect_file(_StallingExec(b"", 0, 10), None, "/app/artifact/slow", 10, "slow", _RefusingStore(), + "uid", 1, "application/octet-stream") + assert list(tmp_path.iterdir()) == [] + diff --git a/tst/util/capabilities.py b/tst/util/capabilities.py index 1d4d46d17..3f7b934fc 100644 --- a/tst/util/capabilities.py +++ b/tst/util/capabilities.py @@ -63,7 +63,7 @@ #: SGLang and Megatron-LM and CPU torch aren't in the dev extra, and only that tier's CI job installs them. SLIME = "slime_installed" -#: ``collect_artifacts`` sizes a file with ``stat -c %s``, which only GNU stat takes (macOS ships BSD stat), +#: ``collect_artifacts`` sizes a file with ``stat -L -c %f:%s``, which only GNU stat takes (macOS ships BSD stat), #: so collecting off a local sandbox needs GNU stat first on PATH. GNU_STAT = "gnu_stat"