diff --git a/src/agentlab/benchmarks/gaia.py b/src/agentlab/benchmarks/gaia.py index 0468e3057..d54b2a7f7 100644 --- a/src/agentlab/benchmarks/gaia.py +++ b/src/agentlab/benchmarks/gaia.py @@ -21,6 +21,7 @@ from agentlab.benchmarks.abstract_env import AbstractBenchmark, AbstractEnvArgs from agentlab.benchmarks.multitool_gym import MultiToolGym +from agentlab.experiments.loop import CODE_SANDBOX_DIR_NAME logger = logging.getLogger(__name__) @@ -96,7 +97,7 @@ def make_env(self, exp_dir: Path, action_mapping=None) -> GaiaGym: def init_code_sandbox(exp_dir: str) -> None: # Use a common code directory for all tasks in the experiment, which is mounted in the container root_exp_dir = Path(exp_dir).parent - code_path = os.path.join(root_exp_dir, "shared_code") + code_path = os.path.join(root_exp_dir, CODE_SANDBOX_DIR_NAME) os.makedirs(code_path, exist_ok=True) os.environ["COMPUTER_CONTAINER_NAME"] = CONTAINER_NAME diff --git a/src/agentlab/experiments/loop.py b/src/agentlab/experiments/loop.py index de4b976a6..0054d7c4e 100644 --- a/src/agentlab/experiments/loop.py +++ b/src/agentlab/experiments/loop.py @@ -25,6 +25,8 @@ from PIL import Image from tqdm import tqdm +from agentlab.experiments import safe_pickle + try: from agentlab.agents.tapeagent import TapeAgent, save_tape except ImportError: @@ -35,6 +37,10 @@ SEED_MAX = 2 ^ 32 # arbitrary max value (exclusive), seems large enough +# Directory (directly under the study directory) mounted as the agent's code sandbox. +# Its content is agent-controlled, so result discovery never descends into it. +CODE_SANDBOX_DIR_NAME = "shared_code" + @dataclass class EnvArgs(DataClassJsonMixin): @@ -668,7 +674,7 @@ def __init__(self, exp_dir) -> None: def exp_args(self) -> ExpArgs: if self._exp_args is None: with open(self.exp_dir / "exp_args.pkl", "rb") as f: - self._exp_args = pickle.load(f) + self._exp_args = safe_pickle.load(f) # in case experiments were moved self._exp_args.exp_dir = self.exp_dir return self._exp_args @@ -677,7 +683,7 @@ def get_step_info(self, step: int) -> StepInfo: """Load the step info from the file and return it.""" if self._steps_info.get(step, None) is None: with gzip.open(self.exp_dir / f"step_{step}.pkl.gz", "rb") as f: - self._steps_info[step] = pickle.load(f) + self._steps_info[step] = safe_pickle.load(f) if self._steps_info[step].obs: if "screenshot" not in self._steps_info[step].obs: try: @@ -699,7 +705,7 @@ def get_step_info(self, step: int) -> StepInfo: and self._steps_info[step].obs["goal_object"] is None ): with gzip.open(self.exp_dir / "goal_object.pkl.gz", "rb") as f: - goal_object = pickle.load(f) + goal_object = safe_pickle.load(f) self._steps_info[step].obs["goal_object"] = goal_object return self._steps_info[step] @@ -867,15 +873,14 @@ def yield_all_exp_results( if not isinstance(savedir_base, list): savedir_base = [savedir_base] - exp_args_paths = [] - for exp_dir in savedir_base: - exp_args_paths.extend(list(Path(exp_dir).glob("**/exp_args.pkl"))) + exp_dirs = [] + for base_dir in savedir_base: + exp_dirs.extend(_find_exp_dirs(base_dir)) if progress_fn is not None: - exp_args_paths = progress_fn(exp_args_paths, desc="Searching experiments directories.") + exp_dirs = progress_fn(exp_dirs, desc="Searching experiments directories.") - for exp_args_path in exp_args_paths: - exp_dir = exp_args_path.parent + for exp_dir in exp_dirs: if not load_hidden: if exp_dir.name.startswith("_") or exp_dir.name.startswith("."): continue @@ -885,6 +890,29 @@ def yield_all_exp_results( yield ExpResult(exp_dir) +def _find_exp_dirs(base_dir: str | Path) -> list[Path]: + """Find experiment directories (those containing an `exp_args.pkl`) under base_dir. + + Symlinks are not followed, the code sandbox directory is skipped and experiment + directories are not searched further: they may contain agent-written files (e.g. the + `code` symlink to the sandbox), and AgentLab never nests experiments inside one another. + + Args: + base_dir: Root directory to search. + + Returns: + list[Path]: The experiment directories found under base_dir. + """ + exp_dirs = [] + for dirpath, dirnames, filenames in os.walk(base_dir, followlinks=False): + if "exp_args.pkl" in filenames: + exp_dirs.append(Path(dirpath)) + dirnames.clear() + else: + dirnames[:] = sorted(d for d in dirnames if d != CODE_SANDBOX_DIR_NAME) + return exp_dirs + + class DataclassJSONEncoder(json.JSONEncoder): def default(self, obj): if is_dataclass(obj): diff --git a/src/agentlab/experiments/safe_pickle.py b/src/agentlab/experiments/safe_pickle.py new file mode 100644 index 000000000..313d5e61c --- /dev/null +++ b/src/agentlab/experiments/safe_pickle.py @@ -0,0 +1,154 @@ +"""Restricted unpickling for experiment result files. + +Result directories can contain files written by code the agent controls (e.g. a code +sandbox mounted inside the study directory), so `exp_args.pkl`, `step_*.pkl.gz` and +`goal_object.pkl.gz` are loaded with an allowlist instead of a bare `pickle.load`. + +Only classes are accepted, never arbitrary callables (apart from a few reconstruction +helpers used by the standard library and numpy): +- any class from a trusted package (agentlab, browsergym, bgym, plus the comma-separated + module roots in the `AGENTLAB_TRUSTED_PICKLE_MODULES` environment variable), +- subclasses of those classes defined in modules that are already imported (e.g. custom + agent args declared in the launching script or a test module); nothing is imported for them, +- data classes (dataclasses, enums, pydantic models, omegaconf nodes) from dependencies + whose objects end up in agent configs and LLM traces, +- plain value types from the standard library and numpy. +""" + +import builtins +import collections +import copyreg +import dataclasses +import datetime +import decimal +import enum +import os +import pathlib +import pickle +import sys +import typing +import uuid + +import numpy as np + +TRUSTED_MODULES_ENV_VAR = "AGENTLAB_TRUSTED_PICKLE_MODULES" + +_TRUSTED_ROOTS = frozenset({"agentlab", "browsergym", "bgym"}) +_DATA_CLASS_ROOTS = frozenset( + {"tapeagents", "pydantic", "pydantic_core", "omegaconf", "openai", "anthropic"} +) + +_SAFE_GLOBALS = { + *( + getattr(builtins, name) + for name in ( + "object", "bool", "int", "float", "complex", "str", "bytes", "bytearray", + "list", "tuple", "dict", "set", "frozenset", "slice", "range", + ) + ), + datetime.datetime, datetime.date, datetime.time, datetime.timedelta, datetime.timezone, + pathlib.Path, pathlib.PosixPath, pathlib.WindowsPath, + pathlib.PurePath, pathlib.PurePosixPath, pathlib.PureWindowsPath, + collections.OrderedDict, collections.defaultdict, collections.deque, collections.Counter, + decimal.Decimal, + uuid.UUID, + copyreg._reconstructor, + np.ndarray, + typing.Any, +} # fmt: skip + +_SAFE_FUNCTIONS = { + ("_codecs", "encode"), + ("numpy.core.multiarray", "_reconstruct"), + ("numpy._core.multiarray", "_reconstruct"), + ("numpy.core.multiarray", "scalar"), + ("numpy._core.multiarray", "scalar"), + ("numpy.core.numeric", "_frombuffer"), + ("numpy._core.numeric", "_frombuffer"), +} + +_SAFE_MODULE_ROOTS = frozenset( + {"builtins", "copyreg", "_codecs", "collections", "datetime", "decimal", "pathlib"} + | {"uuid", "numpy", "PIL", "typing"} +) + + +class UnsafePickleError(pickle.UnpicklingError): + """Raised when a pickle references a global that is not on the allowlist.""" + + +def _root(module: str) -> str: + return module.split(".", 1)[0] + + +def _trusted_roots() -> frozenset: + extra = os.environ.get(TRUSTED_MODULES_ENV_VAR, "") + return _TRUSTED_ROOTS | {m.strip() for m in extra.split(",") if m.strip()} + + +def _is_data_class(obj) -> bool: + if dataclasses.is_dataclass(obj) or issubclass(obj, enum.Enum): + return True + pydantic = sys.modules.get("pydantic") + if pydantic is not None and issubclass(obj, pydantic.BaseModel): + return True + omegaconf = sys.modules.get("omegaconf") + return omegaconf is not None and issubclass(obj, omegaconf.Node) + + +def _is_allowed(obj, trusted_roots) -> bool: + if isinstance(obj, type): + if obj in _SAFE_GLOBALS: + return True + root = _root(getattr(obj, "__module__", None) or "") + if root in trusted_roots: + return True + if root in _DATA_CLASS_ROOTS: + return _is_data_class(obj) + if root == "numpy": + return issubclass(obj, (np.dtype, np.generic)) + if root == "PIL": + from PIL import Image + + return obj is Image.Image + return False + try: + return obj in _SAFE_GLOBALS + except TypeError: # unhashable + return False + + +def _imported_trusted_subclass(module: str, name: str, trusted_roots) -> type | None: + obj = sys.modules.get(module) + for attr in name.split("."): + obj = getattr(obj, attr, None) + if isinstance(obj, type) and any( + _root(getattr(base, "__module__", None) or "") in trusted_roots for base in obj.__mro__[1:] + ): + return obj + return None + + +class RestrictedUnpickler(pickle.Unpickler): + def find_class(self, module, name): + trusted_roots = _trusted_roots() + allowed_roots = trusted_roots | _DATA_CLASS_ROOTS | _SAFE_MODULE_ROOTS + if _root(module) in allowed_roots: + if (module, name) in _SAFE_FUNCTIONS: + return super().find_class(module, name) + obj = super().find_class(module, name) + if _is_allowed(obj, trusted_roots): + return obj + obj = _imported_trusted_subclass(module, name, trusted_roots) + if obj is not None: + return obj + raise UnsafePickleError( + f"Refusing to unpickle global '{module}.{name}' from an untrusted result file. " + f"If this class belongs to your own agent or benchmark package, add its top-level " + f"module to the {TRUSTED_MODULES_ENV_VAR} environment variable (comma-separated)." + ) + + +def load(file): + """Drop-in replacement for `pickle.load` restricted to the allowlist above.""" + return RestrictedUnpickler(file).load() diff --git a/tests/experiments/test_safe_result_loading.py b/tests/experiments/test_safe_result_loading.py new file mode 100644 index 000000000..5a342b4c8 --- /dev/null +++ b/tests/experiments/test_safe_result_loading.py @@ -0,0 +1,222 @@ +import gzip +import os +import pickle +import shutil +import subprocess +import typing +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path + +import numpy as np +import pytest +from browsergym.experiments.agent import AgentInfo + +from agentlab.analyze import inspect_results +from agentlab.experiments import safe_pickle +from agentlab.experiments.loop import ( + CODE_SANDBOX_DIR_NAME, + EnvArgs, + ExpResult, + StepInfo, + StepTimestamps, + yield_all_exp_results, +) +from agentlab.experiments.study import Study +from agentlab.llm.llm_utils import Discussion, HumanMessage, SystemMessage + +TEST_STUDY = Path(__file__).parent.parent / "data" / "test_study" + + +class _WriteMarker: + """Pickle payload that writes a marker file when unpickled.""" + + def __init__(self, marker: Path): + self.marker = str(marker) + + def __reduce__(self): + return (exec, (f"open({self.marker!r}, 'w').write('pwned')",)) + + +class _WriteMarkerViaPath: + """Same, but through a function of an allowlisted module.""" + + def __init__(self, marker: Path): + self.marker = Path(marker) + + def __reduce__(self): + return (Path.write_text, (self.marker, "pwned")) + + +def _plant_experiment(exp_dir: Path, marker: Path): + exp_dir.mkdir(parents=True) + payload = pickle.dumps(_WriteMarker(marker)) + (exp_dir / "exp_args.pkl").write_bytes(payload) + with gzip.open(exp_dir / "step_0.pkl.gz", "wb") as f: + f.write(payload) + (exp_dir / "summary_info.json").write_text("{}") + + +@pytest.fixture +def study_dir(tmp_path): + study_dir = tmp_path / "study" + shutil.copytree(TEST_STUDY, study_dir) + return study_dir + + +def test_sandbox_files_are_not_discovered(study_dir, tmp_path): + marker = tmp_path / "marker" + sandbox = study_dir / CODE_SANDBOX_DIR_NAME + _plant_experiment(sandbox / "x", marker) + _plant_experiment(sandbox / "a" / "b", marker) + # experiment dirs symlink to the sandbox (see gaia.init_code_sandbox) + exp_dirs = sorted(p for p in study_dir.iterdir() if p.name != CODE_SANDBOX_DIR_NAME) + os.symlink(sandbox, exp_dirs[0] / "code") + # a nested dir inside a real experiment is not an experiment either + _plant_experiment(exp_dirs[1] / "nested", marker) + + found = list(yield_all_exp_results(study_dir, progress_fn=None, use_cache=False)) + assert sorted(r.exp_dir for r in found) == exp_dirs + + result_df = inspect_results.load_result_df(study_dir, progress_fn=None) + assert sorted(Path(d) for d in result_df["exp_dir"]) == exp_dirs + assert not any(CODE_SANDBOX_DIR_NAME in str(d) for d in result_df["exp_dir"]) + + study = Study.load(study_dir) + study.load_exp_args_list() + assert len(study.exp_args_list) == len(exp_dirs) + + assert not marker.exists() + + +def test_symlinked_dirs_are_not_followed(study_dir, tmp_path): + marker = tmp_path / "marker" + outside = tmp_path / "outside" + _plant_experiment(outside / "x", marker) + os.symlink(outside, study_dir / "link") + + found = list(yield_all_exp_results(study_dir, progress_fn=None, use_cache=False)) + assert all("link" not in r.exp_dir.parts for r in found) + assert len(found) == 2 + assert not marker.exists() + + +@pytest.mark.parametrize("payload_cls", [_WriteMarker, _WriteMarkerViaPath]) +def test_malicious_result_files_are_rejected(tmp_path, payload_cls): + marker = tmp_path / "marker" + exp_dir = tmp_path / "exp" + exp_dir.mkdir() + payload = pickle.dumps(payload_cls(marker)) + (exp_dir / "exp_args.pkl").write_bytes(payload) + with gzip.open(exp_dir / "step_0.pkl.gz", "wb") as f: + f.write(payload) + + with pytest.raises(safe_pickle.UnsafePickleError): + ExpResult(exp_dir).exp_args + with pytest.raises(safe_pickle.UnsafePickleError): + ExpResult(exp_dir).get_step_info(0) + assert not marker.exists() + + +def test_malicious_goal_object_is_rejected(tmp_path): + marker = tmp_path / "marker" + step_info = StepInfo(step=0, obs={"goal": "x", "goal_object": [{"type": "text", "text": "x"}]}) + step_info.save_step_info(tmp_path) + with gzip.open(tmp_path / "goal_object.pkl.gz", "wb") as f: + pickle.dump(_WriteMarker(marker), f) + + with pytest.raises(safe_pickle.UnsafePickleError): + ExpResult(tmp_path).get_step_info(0) + assert not marker.exists() + + +def test_legit_results_round_trip(tmp_path): + discussion = Discussion([SystemMessage("system"), HumanMessage("hello")]) + obs = { + "goal": "click", + "goal_object": [{"type": "text", "text": "click"}], + "chat_messages": [{"role": "user", "message": "click"}], + "axtree_object": {"nodes": [{"nodeId": "1"}]}, + "open_pages_urls": ("http://localhost",), + "active_page_index": np.array([0]), + "elapsed_time": np.array([1.5]), + "screenshot": np.zeros((4, 4, 3), dtype=np.uint8), + "score": np.float32(0.5), + } + step_info = StepInfo( + step=0, + obs=dict(obs), + reward=1.0, + action="click('1')", + agent_info=AgentInfo(think="t", chat_messages=discussion, stats={"n": 1}), + profiling=StepTimestamps(env_start=1.0), + ) + step_info.save_step_info(tmp_path) + + loaded = ExpResult(tmp_path).get_step_info(0) + assert loaded.action == "click('1')" + assert loaded.obs["goal_object"] == obs["goal_object"] + assert loaded.obs["elapsed_time"].tolist() == [1.5] + assert loaded.obs["score"] == np.float32(0.5) + assert loaded.obs["screenshot"].shape == (4, 4, 3) + assert str(loaded.agent_info.chat_messages) == str(discussion) + + value = {"env_args": EnvArgs(task_name="miniwob.click-test"), "t": datetime.now()} + with open(tmp_path / "data.pkl", "wb") as f: + pickle.dump(value, f) + with open(tmp_path / "data.pkl", "rb") as f: + assert safe_pickle.load(f) == value + + for exp_dir in TEST_STUDY.iterdir(): + assert ExpResult(exp_dir).exp_args.env_args.task_name == "miniwob.ascending-numbers" + + +def test_trusted_modules_env_var(tmp_path, monkeypatch): + path = tmp_path / "data.pkl" + with open(path, "wb") as f: + pickle.dump(_Custom(), f) + + with open(path, "rb") as f, pytest.raises(safe_pickle.UnsafePickleError): + safe_pickle.load(f) + + monkeypatch.setenv(safe_pickle.TRUSTED_MODULES_ENV_VAR, _Custom.__module__.split(".")[0]) + with open(path, "rb") as f: + assert isinstance(safe_pickle.load(f), _Custom) + + +class _Custom: + pass + + +@dataclass +class _CustomEnvArgs(EnvArgs): + extra: str = "x" + + +class _SpawnProcess: + def __init__(self, marker: Path): + self.marker = str(marker) + + def __reduce__(self): + return (subprocess.Popen, (["touch", self.marker],)) + + +def test_imported_subclass_of_trusted_class_is_allowed(tmp_path): + path = tmp_path / "data.pkl" + value = _CustomEnvArgs(task_name="miniwob.click-test", extra="y") + with open(path, "wb") as f: + pickle.dump({"args": value, "ref_type": typing.Any}, f) + + with open(path, "rb") as f: + assert safe_pickle.load(f) == {"args": value, "ref_type": typing.Any} + + +def test_imported_untrusted_class_is_rejected(tmp_path): + marker = tmp_path / "pwned" + path = tmp_path / "data.pkl" + with open(path, "wb") as f: + pickle.dump(_SpawnProcess(marker), f) + + with open(path, "rb") as f, pytest.raises(safe_pickle.UnsafePickleError): + safe_pickle.load(f) + assert not marker.exists()