diff --git a/python/packages/monty/agent_framework_monty/_execute_code_tool.py b/python/packages/monty/agent_framework_monty/_execute_code_tool.py index 6a143f33bb0..551c5964441 100644 --- a/python/packages/monty/agent_framework_monty/_execute_code_tool.py +++ b/python/packages/monty/agent_framework_monty/_execute_code_tool.py @@ -15,12 +15,15 @@ import json import logging import mimetypes +import os +import stat from collections.abc import Callable, Iterator, Sequence from functools import partial from pathlib import Path, PurePosixPath from typing import Any, cast from agent_framework import Content, FunctionTool +from agent_framework._filesystem import _is_link_or_reparse_point # pyright: ignore[reportPrivateUsage] from agent_framework._tools import ( ApprovalMode, _normalize_tool_description_format, # pyright: ignore[reportPrivateUsage] @@ -498,42 +501,95 @@ def _build_execution_contents(*, result: dict[str, Any]) -> list[Content]: def _iter_real_files(root: Path) -> Iterator[Path]: - """Walk ``root`` recursively, yielding only real (non-symlink) files. + """Walk ``root`` recursively, yielding only real regular files. ``Path.rglob`` follows directory symlinks by default, which combined with ``Path.is_file()`` / ``Path.read_bytes()`` (both follow symlinks) would let an attacker who controls the workspace pre-place a symlink to a host file - or directory and have our post-execution capture surface it. Skipping every - symlink at both the directory and file level closes that escape. + or directory and have our post-execution capture surface it. Skip symlinks, + Windows junctions, and other reparse points at every directory level. """ stack: list[Path] = [root] while stack: current = stack.pop() try: + if _is_link_or_reparse_point(current): + continue entries = list(current.iterdir()) except OSError: continue for entry in entries: try: - if entry.is_symlink(): + if _is_link_or_reparse_point(entry): continue - if entry.is_dir(): + entry_stat = entry.lstat() + if stat.S_ISDIR(entry_stat.st_mode): stack.append(entry) - elif entry.is_file(): + elif stat.S_ISREG(entry_stat.st_mode): yield entry except OSError: continue +def _safe_capture_file_stat(*, root: Path, entry: Path) -> os.stat_result | None: + """Return the verified stat for a real regular file beneath ``root``.""" + try: + relative = entry.relative_to(root) + except ValueError: + return None + + current = root + try: + if _is_link_or_reparse_point(current): + return None + for part in relative.parts: + current /= part + if _is_link_or_reparse_point(current): + return None + entry_stat = current.lstat() + except OSError: + return None + return entry_stat if stat.S_ISREG(entry_stat.st_mode) else None + + +def _read_validated_file(*, root: Path, entry: Path, expected: os.stat_result) -> bytes | None: + """Read ``entry`` through one descriptor that is checked against the validated file. + + ``Path.read_bytes()`` reopens the path, so a component swapped for a link or junction after + validation would be followed. Opening first, comparing the descriptor with the validated ``lstat`` + identity, and revalidating the path afterwards rejects such a swap. Windows has no no-follow open, + so the check is not atomic there, but an undetected swap would have to be applied, reverted, and + applied again around the checks. Returns ``None`` when the file changed or cannot be read safely. + """ + flags = os.O_RDONLY | getattr(os, "O_BINARY", 0) | getattr(os, "O_NOFOLLOW", 0) + try: + fd = os.open(entry, flags) + except OSError: + return None + try: + opened = os.fstat(fd) + identity = (opened.st_dev, opened.st_ino) + if not stat.S_ISREG(opened.st_mode) or identity != (expected.st_dev, expected.st_ino): + return None + rechecked = _safe_capture_file_stat(root=root, entry=entry) + if rechecked is None or (rechecked.st_dev, rechecked.st_ino) != identity: + return None + with os.fdopen(fd, "rb", closefd=False) as handle: + return handle.read() + except OSError: + return None + finally: + os.close(fd) + + def _snapshot_writable_mounts(mounts: Sequence[FileMount]) -> dict[str, dict[str, tuple[int, int]]]: - """Capture (size, mtime_ns) for every real (non-symlink) host file under read-write mounts. + """Capture (size, mtime_ns) for every real host file under read-write mounts. Returns ``{mount_path: {relative_posix_path: (size, mtime_ns)}}``. Used by :func:`_capture_written_files` to detect new or modified files after the run. Read-only and overlay mounts are skipped because their writes do not - propagate to the host. Symlinks (file or directory) are deliberately skipped - so an attacker cannot escape the mount by pre-placing a symlink to a host - path outside the workspace. + propagate to the host. Links and reparse points are deliberately skipped so + an attacker cannot escape the mount through a host path outside the workspace. """ snapshot: dict[str, dict[str, tuple[int, int]]] = {} for mount in mounts: @@ -543,11 +599,11 @@ def _snapshot_writable_mounts(mounts: Sequence[FileMount]) -> dict[str, dict[str per_mount: dict[str, tuple[int, int]] = {} for entry in _iter_real_files(host_root): try: - stat = entry.lstat() # lstat: never follow symlinks (defensive) + entry_stat = entry.lstat() # lstat: never follow symlinks (defensive) except OSError: continue relative = entry.relative_to(host_root).as_posix() - per_mount[relative] = (int(stat.st_size), int(stat.st_mtime_ns)) + per_mount[relative] = (int(entry_stat.st_size), int(entry_stat.st_mtime_ns)) snapshot[mount.mount_path] = per_mount return snapshot @@ -559,11 +615,12 @@ def _capture_written_files( """Return :class:`Content` items for files the sandbox wrote during the run. Mirrors Hyperlight's ``/output`` capture flow: any new or modified real - (non-symlink) file under a read-write mount is read back as binary and + regular file under a read-write mount is read back as binary and surfaced as ``Content.from_data`` with a ``path`` annotation in - ``additional_properties``. Symlinks are skipped at both directory and file - level so a malicious workspace cannot trick us into capturing host files - outside the configured mount root. + ``additional_properties``. Links and reparse points are skipped at every + directory level, and each file is read through a descriptor verified + against the validated path, so a malicious workspace cannot trick us into + capturing host files outside the configured mount root. """ captured: list[Content] = [] for mount in mounts: @@ -572,28 +629,24 @@ def _capture_written_files( host_root = Path(mount.host_path) before = pre_state.get(mount.mount_path, {}) for entry in sorted(_iter_real_files(host_root)): - try: - stat = entry.lstat() - except OSError: + entry_stat = _safe_capture_file_stat(root=host_root, entry=entry) + if entry_stat is None: continue relative = entry.relative_to(host_root).as_posix() - current = (int(stat.st_size), int(stat.st_mtime_ns)) + current = (int(entry_stat.st_size), int(entry_stat.st_mtime_ns)) if before.get(relative) == current: continue # Unchanged. sandbox_path = f"{mount.mount_path.rstrip('/')}/{relative}" - if stat.st_size > MAX_CAPTURED_FILE_BYTES: + if entry_stat.st_size > MAX_CAPTURED_FILE_BYTES: captured.append( Content.from_text( - f"[file {sandbox_path} omitted: {stat.st_size} bytes " + f"[file {sandbox_path} omitted: {entry_stat.st_size} bytes " f"exceeds MAX_CAPTURED_FILE_BYTES={MAX_CAPTURED_FILE_BYTES}]" ) ) continue - try: - # _iter_real_files already excluded symlinks at every level of - # the walk; reading the file here is safe. - data = entry.read_bytes() - except OSError: + data = _read_validated_file(root=host_root, entry=entry, expected=entry_stat) + if data is None: continue media_type = mimetypes.guess_type(entry.name)[0] or "application/octet-stream" captured.append( diff --git a/python/packages/monty/tests/monty/test_monty_codeact.py b/python/packages/monty/tests/monty/test_monty_codeact.py index 33de9f603b6..d5f5db71cc2 100644 --- a/python/packages/monty/tests/monty/test_monty_codeact.py +++ b/python/packages/monty/tests/monty/test_monty_codeact.py @@ -11,10 +11,13 @@ import asyncio import json +import os import re +import stat +import subprocess import sys import types -from collections.abc import Iterable, Iterator, Mapping +from collections.abc import Callable, Iterable, Iterator, Mapping from dataclasses import dataclass, field from pathlib import Path from typing import Annotated, Any, Literal @@ -984,6 +987,8 @@ def test_capture_written_files_returns_new_files_and_omits_large_ones( existing = writable / "existing.txt" existing.write_text("before", encoding="utf-8") + unchanged = writable / "unchanged.txt" + unchanged.write_text("same", encoding="utf-8") (nested / "report.txt").write_text("old", encoding="utf-8") (readonly / "ignored.txt").write_text("unchanged", encoding="utf-8") @@ -1005,6 +1010,7 @@ def test_capture_written_files_returns_new_files_and_omits_large_ones( assert set(pre_state) == {"/work"} assert "existing.txt" in pre_state["/work"] + assert "unchanged.txt" in pre_state["/work"] assert "ignored.txt" not in pre_state["/work"] assert {item.additional_properties["path"] for item in data_items} == { "/work/artifact.bin", @@ -1018,6 +1024,305 @@ def test_capture_written_files_returns_new_files_and_omits_large_ones( if item.additional_properties["path"] == "/work/existing.txt" ) assert all(not item.additional_properties["path"].startswith("/readonly/") for item in data_items) + assert all(item.additional_properties["path"] != "/work/unchanged.txt" for item in data_items) + + +def _create_junction_or_skip(*, link: Path, target: Path) -> None: + if sys.platform != "win32": + pytest.skip("Windows directory junctions are only available on Windows") + + result = subprocess.run( + [os.environ.get("COMSPEC", "cmd"), "/c", "mklink", "/J", str(link), str(target)], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + pytest.skip(f"Could not create Windows directory junction: {result.stderr or result.stdout}") + + is_junction = getattr(link, "is_junction", None) + reparse_attribute = getattr(stat, "FILE_ATTRIBUTE_REPARSE_POINT", 0) + is_reparse_point = bool(reparse_attribute and getattr(link.lstat(), "st_file_attributes", 0) & reparse_attribute) + if not (callable(is_junction) and is_junction()) and not is_reparse_point: + link.rmdir() + pytest.skip("Created junction was not reported as a reparse point") + + +def _model_junction_or_skip(*, link: Path, target: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Make ``link`` look like an NTFS junction to pathlib while still traversing to ``target``. + + The stat-family answers for ``link`` become those of a junction (not a symlink, a directory with the + reparse-point attribute) and every other path is untouched, so non-Windows runs exercise the same walker + and capture logic. The real-junction variants cover NTFS itself. + """ + try: + link.symlink_to(target, target_is_directory=True) + except (OSError, NotImplementedError) as exc: + pytest.skip(f"Symbolic links are not supported in this environment: {exc!r}") + + original_is_symlink = Path.is_symlink + original_lstat = Path.lstat + + def is_symlink(self: Path) -> bool: + return False if self == link else original_is_symlink(self) + + def lstat(self: Path) -> Any: + real = original_lstat(self) + if self != link: + return real + return types.SimpleNamespace( + st_mode=stat.S_IFDIR | 0o755, + st_size=0, + st_mtime_ns=real.st_mtime_ns, + st_dev=real.st_dev, + st_ino=real.st_ino, + st_file_attributes=getattr(stat, "FILE_ATTRIBUTE_REPARSE_POINT", 0), + ) + + monkeypatch.setattr(Path, "is_symlink", is_symlink) + monkeypatch.setattr(Path, "lstat", lstat) + + +@pytest.mark.parametrize("junction_kind", ["windows", "modeled"]) +@pytest.mark.parametrize("create_before_snapshot", [True, False]) +def test_capture_written_files_skips_junction_to_outside_directory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + junction_kind: str, + create_before_snapshot: bool, +) -> None: + from agent_framework_monty import FileMount + + writable = tmp_path / "writable" + writable.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + (outside / "existing.txt").write_text("before", encoding="utf-8") + junction = writable / "linked" + mount = FileMount(host_path=writable, mount_path="/input", mode="read-write") + + def create_junction() -> None: + if junction_kind == "windows": + _create_junction_or_skip(link=junction, target=outside) + else: + _model_junction_or_skip(link=junction, target=outside, monkeypatch=monkeypatch) + + if create_before_snapshot: + create_junction() + pre_state = execute_code_module._snapshot_writable_mounts([mount]) + if not create_before_snapshot: + create_junction() + + try: + assert not any(path.startswith("linked/") for path in pre_state["/input"]) + (outside / "existing.txt").write_text("modified", encoding="utf-8") + (outside / "created.txt").write_text("created", encoding="utf-8") + (writable / "report.txt").write_text("legitimate", encoding="utf-8") + + captured = execute_code_module._capture_written_files([mount], pre_state) + + assert {item.additional_properties["path"] for item in captured if item.type == "data"} == {"/input/report.txt"} + finally: + if junction_kind == "windows": + junction.rmdir() + + +def test_iter_real_files_fails_closed_when_reparse_probe_raises( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + readable = tmp_path / "readable.txt" + readable.write_text("safe", encoding="utf-8") + blocked = tmp_path / "blocked" + blocked.mkdir() + (blocked / "secret.txt").write_text("secret", encoding="utf-8") + original_probe = execute_code_module._is_link_or_reparse_point + + def fail_for_blocked(path: Path) -> bool: + if path == blocked: + raise PermissionError("access denied") + return original_probe(path) + + monkeypatch.setattr(execute_code_module, "_is_link_or_reparse_point", fail_for_blocked) + + assert set(execute_code_module._iter_real_files(tmp_path)) == {readable} + assert execute_code_module._safe_capture_file_stat(root=tmp_path, entry=blocked / "secret.txt") is None + + +def _skip_without_symlinks(tmp_path: Path) -> None: + probe = tmp_path / "symlink-probe" + try: + probe.symlink_to(tmp_path / "missing-target") + except (OSError, NotImplementedError) as exc: + pytest.skip(f"Symbolic links are not supported in this environment: {exc!r}") + probe.unlink() + + +def _swap_after_first_validation(monkeypatch: pytest.MonkeyPatch, swap: Callable[[], None]) -> None: + """Run ``swap`` once, right after the first path validation and before the file is opened.""" + original = execute_code_module._safe_capture_file_stat + swapped = False + + def validate_then_swap(**kwargs: Any) -> Any: + nonlocal swapped + result = original(**kwargs) + if not swapped: + swapped = True + swap() + return result + + monkeypatch.setattr(execute_code_module, "_safe_capture_file_stat", validate_then_swap) + + +def test_safe_capture_file_stat_rejects_links_outside_paths_and_non_files(tmp_path: Path) -> None: + _skip_without_symlinks(tmp_path) + root = tmp_path / "root" + root.mkdir() + (root / "file.txt").write_text("payload", encoding="utf-8") + outside = tmp_path / "outside" + outside.mkdir() + (outside / "secret.txt").write_text("HOST_SECRET", encoding="utf-8") + (root / "linked").symlink_to(outside, target_is_directory=True) + linked_root = tmp_path / "linked-root" + linked_root.symlink_to(root, target_is_directory=True) + safe_stat = execute_code_module._safe_capture_file_stat + + assert safe_stat(root=root, entry=root / "file.txt") is not None + assert safe_stat(root=root, entry=root / "linked" / "secret.txt") is None + assert safe_stat(root=root, entry=outside / "secret.txt") is None + assert safe_stat(root=root, entry=root) is None + assert safe_stat(root=linked_root, entry=linked_root / "file.txt") is None + assert list(execute_code_module._iter_real_files(linked_root)) == [] + + +@pytest.mark.parametrize("recheck_result", ["unsafe", "different-file"]) +def test_capture_written_files_rechecks_path_after_opening( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + recheck_result: str, +) -> None: + from agent_framework_monty import FileMount + + writable = tmp_path / "writable" + writable.mkdir() + artifact = writable / "artifact.txt" + artifact.write_text("payload", encoding="utf-8") + other = tmp_path / "other.txt" + other.write_text("other", encoding="utf-8") + mount = FileMount(host_path=writable, mount_path="/input", mode="read-write") + safety_checks = iter([artifact.lstat(), None if recheck_result == "unsafe" else other.lstat()]) + + monkeypatch.setattr( + execute_code_module, + "_safe_capture_file_stat", + lambda **_: next(safety_checks), + ) + + captured = execute_code_module._capture_written_files([mount], {"/input": {}}) + + assert captured == [] + + +@pytest.mark.parametrize("replacement_kind", ["file", "directory"]) +def test_capture_written_files_skips_file_replaced_after_validation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + replacement_kind: str, +) -> None: + from agent_framework_monty import FileMount + + writable = tmp_path / "writable" + writable.mkdir() + artifact = writable / "artifact.txt" + artifact.write_text("payload", encoding="utf-8") + # Created before the original is removed so the two files cannot share an inode number. + replacement = tmp_path / "replacement.txt" + replacement.write_text("REPLACED", encoding="utf-8") + mount = FileMount(host_path=writable, mount_path="/input", mode="read-write") + + def swap() -> None: + if replacement_kind == "file": + os.replace(replacement, artifact) + else: + artifact.unlink() + artifact.mkdir() + + _swap_after_first_validation(monkeypatch, swap) + + captured = execute_code_module._capture_written_files([mount], {"/input": {}}) + + assert captured == [] + + +def test_capture_written_files_skips_file_removed_after_validation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from agent_framework_monty import FileMount + + writable = tmp_path / "writable" + writable.mkdir() + artifact = writable / "artifact.txt" + artifact.write_text("payload", encoding="utf-8") + mount = FileMount(host_path=writable, mount_path="/input", mode="read-write") + _swap_after_first_validation(monkeypatch, artifact.unlink) + + captured = execute_code_module._capture_written_files([mount], {"/input": {}}) + + assert captured == [] + + +def test_capture_written_files_does_not_follow_leaf_link_swapped_in_after_validation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from agent_framework_monty import FileMount + + _skip_without_symlinks(tmp_path) + writable = tmp_path / "writable" + writable.mkdir() + artifact = writable / "artifact.txt" + artifact.write_text("payload", encoding="utf-8") + secret = tmp_path / "secret.txt" + secret.write_text("HOST_SECRET", encoding="utf-8") + mount = FileMount(host_path=writable, mount_path="/input", mode="read-write") + + def swap() -> None: + artifact.unlink() + artifact.symlink_to(secret) + + _swap_after_first_validation(monkeypatch, swap) + + captured = execute_code_module._capture_written_files([mount], {"/input": {}}) + + assert not any(item.type == "data" for item in captured) + + +def test_capture_written_files_does_not_follow_parent_link_swapped_in_after_validation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from agent_framework_monty import FileMount + + _skip_without_symlinks(tmp_path) + writable = tmp_path / "writable" + parent = writable / "sub" + parent.mkdir(parents=True) + (parent / "report.txt").write_text("payload", encoding="utf-8") + outside = tmp_path / "outside" + outside.mkdir() + (outside / "report.txt").write_text("HOST_SECRET", encoding="utf-8") + mount = FileMount(host_path=writable, mount_path="/input", mode="read-write") + + def swap() -> None: + parent.rename(writable / "original-sub") + parent.symlink_to(outside, target_is_directory=True) + + _swap_after_first_validation(monkeypatch, swap) + + captured = execute_code_module._capture_written_files([mount], {"/input": {}}) + + assert not any(item.type == "data" for item in captured) # ---------------------------------------------------------------------------