From 9db63c8a0fbb6958f6acd1ae40eb98665387e109 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 5 Nov 2026 09:47:40 +0000 Subject: [PATCH 1/2] bug: load study.pkl.gz and exp_args.pkl with a restricted unpickler Study.load and ExpResult.exp_args used pickle.load on files from any study directory, so a planted study folder could run code on load (archive_studies, load_most_recent, main.py relaunch). Add agentlab.experiments.safe_pickle, which only resolves dataclasses/Enums/pydantic models from agentlab and browsergym (plus AGENTLAB_TRUSTED_PICKLE_MODULES) and an allowlist of stdlib/numpy/pandas/pyarrow data types. Fixes code scan finding sfind-862fb47f130745099616ac639c00fc9c. --- README.md | 7 + src/agentlab/experiments/loop.py | 4 +- src/agentlab/experiments/safe_pickle.py | 179 ++++++++++++++++++++++++ src/agentlab/experiments/study.py | 6 +- tests/experiments/test_safe_pickle.py | 140 ++++++++++++++++++ 5 files changed, 333 insertions(+), 3 deletions(-) create mode 100644 src/agentlab/experiments/safe_pickle.py create mode 100644 tests/experiments/test_safe_pickle.py diff --git a/README.md b/README.md index 8397c6053..93e613b43 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,13 @@ study.find_incomplete(include_errors=True) study.run() ``` +`Study.load` (and result loading in general) only unpickles AgentLab/BrowserGym dataclasses, Enums, +pydantic models and a small allowlist of data types, so a planted `study.pkl.gz` or `exp_args.pkl` +cannot run code. If your agent args are defined in your own package, add it to the +`AGENTLAB_TRUSTED_PICKLE_MODULES` environment variable (comma separated, e.g. +`AGENTLAB_TRUSTED_PICKLE_MODULES=my_agents`). Only do this for study folders from trusted sources: +step files (`step_*.pkl.gz`) loaded by AgentXray still use plain `pickle`. + See [main.py](main.py) to launch experiments with a variety of options. This is like a lazy CLI that is actually more convenient. Just comment and uncomment the lines you need or modify at will (but don't push to the repo). diff --git a/src/agentlab/experiments/loop.py b/src/agentlab/experiments/loop.py index de4b976a6..06bb8d8ef 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: @@ -668,7 +670,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 diff --git a/src/agentlab/experiments/safe_pickle.py b/src/agentlab/experiments/safe_pickle.py new file mode 100644 index 000000000..54486db12 --- /dev/null +++ b/src/agentlab/experiments/safe_pickle.py @@ -0,0 +1,179 @@ +"""Restricted unpickling for study and experiment files. + +`pickle.load` executes arbitrary code embedded in the file. Study directories can come from +shared drives or collaborators, so files such as `study.pkl.gz` and `exp_args.pkl` are loaded +through `RestrictedUnpickler`, which only resolves: + +- dataclasses, Enums and pydantic models defined in trusted packages (`agentlab`, `browsergym` + and any extra package listed in the `AGENTLAB_TRUSTED_PICKLE_MODULES` environment variable, + comma separated, e.g. for agent args defined in your own package); +- an explicit allowlist of standard-library, numpy, pandas, pyarrow and omegaconf data types. + +Any other global (functions such as `os.system`, arbitrary classes, ...) raises +`UnsafePickleError` before it can be called. +""" + +import dataclasses +import enum +import io +import os +import pickle + +TRUSTED_PACKAGES = ("agentlab", "browsergym", "bgym") +TRUSTED_PACKAGES_ENV_VAR = "AGENTLAB_TRUSTED_PICKLE_MODULES" + +_STDLIB_GLOBALS = { + ("builtins", name) + for name in ( + "object", + "set", + "frozenset", + "slice", + "range", + "complex", + "bytes", + "bytearray", + "list", + "dict", + "tuple", + "str", + "int", + "float", + "bool", + ) +} | { + ("copyreg", "_reconstructor"), + ("collections", "OrderedDict"), + ("collections", "defaultdict"), + ("collections", "deque"), + ("collections", "Counter"), + ("datetime", "datetime"), + ("datetime", "date"), + ("datetime", "time"), + ("datetime", "timedelta"), + ("datetime", "timezone"), + ("decimal", "Decimal"), + ("fractions", "Fraction"), + ("uuid", "UUID"), + ("pathlib", "Path"), + ("pathlib", "PosixPath"), + ("pathlib", "WindowsPath"), + ("pathlib", "PurePath"), + ("pathlib", "PurePosixPath"), + ("pathlib", "PureWindowsPath"), +} + +_NUMPY_GLOBALS = { + (module, name) + for module in ("numpy.core.multiarray", "numpy._core.multiarray") + for name in ("_reconstruct", "scalar") +} | { + ("numpy", "ndarray"), + ("numpy", "dtype"), + ("numpy.core.numeric", "_frombuffer"), + ("numpy._core.numeric", "_frombuffer"), +} + +_PANDAS_GLOBALS = { + ("pandas", "DataFrame"), + ("pandas", "Series"), + ("pandas", "Index"), + ("pandas", "RangeIndex"), + ("pandas", "Timestamp"), + ("pandas", "Timedelta"), + ("pandas", "StringDtype"), + ("pandas", "CategoricalDtype"), + ("pandas", "Categorical"), + ("pandas.core.dtypes.dtypes", "CategoricalDtype"), + ("pandas.core.arrays.categorical", "Categorical"), + ("pandas.core.arrays.categorical", "_reconstruct"), + ("pandas.core.frame", "DataFrame"), + ("pandas.core.series", "Series"), + ("pandas.core.indexes.base", "Index"), + ("pandas.core.indexes.base", "_new_Index"), + ("pandas.core.indexes.range", "RangeIndex"), + ("pandas.core.internals.managers", "BlockManager"), + ("pandas.core.internals.managers", "SingleBlockManager"), + ("pandas._libs.internals", "_unpickle_block"), + ("pandas.core.internals.blocks", "new_block"), + ("pandas.arrays", "ArrowStringArray"), + ("pandas.arrays", "StringArray"), + ("pandas.core.arrays.string_", "StringDtype"), + ("pandas.core.arrays.string_", "StringArray"), + ("pandas.core.arrays.string_arrow", "ArrowStringArray"), + ("pandas.core.arrays.string_arrow", "ArrowStringArrayNumpySemantics"), + ("pyarrow.lib", "_restore_array"), + ("pyarrow.lib", "type_for_alias"), + ("pyarrow.lib", "py_buffer"), +} + +_OMEGACONF_GLOBALS = { + ("omegaconf.dictconfig", "DictConfig"), + ("omegaconf.listconfig", "ListConfig"), + ("omegaconf.base", "ContainerMetadata"), + ("omegaconf.base", "Metadata"), + ("omegaconf.nodes", "AnyNode"), + ("omegaconf.nodes", "StringNode"), + ("omegaconf.nodes", "IntegerNode"), + ("omegaconf.nodes", "FloatNode"), + ("omegaconf.nodes", "BooleanNode"), + ("omegaconf.nodes", "BytesNode"), + ("omegaconf.nodes", "PathNode"), + ("omegaconf.nodes", "EnumNode"), + ("typing", "Any"), +} + +SAFE_GLOBALS = frozenset(_STDLIB_GLOBALS | _NUMPY_GLOBALS | _PANDAS_GLOBALS | _OMEGACONF_GLOBALS) + + +class UnsafePickleError(pickle.UnpicklingError): + """Raised when a pickle references a global that is not allowed.""" + + +def _trusted_packages() -> tuple[str, ...]: + extra = os.environ.get(TRUSTED_PACKAGES_ENV_VAR, "") + return TRUSTED_PACKAGES + tuple(p.strip() for p in extra.split(",") if p.strip()) + + +def _in_trusted_package(module: str) -> bool: + return any(module == pkg or module.startswith(pkg + ".") for pkg in _trusted_packages()) + + +def _is_data_type(obj) -> bool: + if not isinstance(obj, type): + return False + if dataclasses.is_dataclass(obj) or issubclass(obj, enum.Enum): + return True + try: + from pydantic import BaseModel + except ImportError: # pragma: no cover + return False + return issubclass(obj, BaseModel) + + +class RestrictedUnpickler(pickle.Unpickler): + """Unpickler that only resolves allowlisted data types (see module docstring).""" + + def find_class(self, module: str, name: str): + if (module, name) in SAFE_GLOBALS: + return super().find_class(module, name) + if _in_trusted_package(module) and "." not in name: + obj = super().find_class(module, name) + if _is_data_type(obj) and _in_trusted_package(getattr(obj, "__module__", "")): + return obj + raise UnsafePickleError( + f"Refusing to unpickle global '{module}.{name}': only dataclasses, Enums and pydantic " + f"models from {', '.join(_trusted_packages())} and allowlisted data types are " + f"permitted. If this file comes from a trusted source and defines classes in your " + f"own package, add that package to the {TRUSTED_PACKAGES_ENV_VAR} environment variable." + ) + + +def load(file): + """Safe replacement for `pickle.load` on study/experiment files.""" + return RestrictedUnpickler(file).load() + + +def loads(data: bytes): + """Safe replacement for `pickle.loads` on study/experiment files.""" + return load(io.BytesIO(data)) diff --git a/src/agentlab/experiments/study.py b/src/agentlab/experiments/study.py index 391f419ce..ce66c8ffb 100644 --- a/src/agentlab/experiments/study.py +++ b/src/agentlab/experiments/study.py @@ -19,6 +19,7 @@ from agentlab.analyze import inspect_results from agentlab.benchmarks.abstract_env import AbstractEnvArgs from agentlab.experiments import reproducibility_util as repro +from agentlab.experiments import safe_pickle from agentlab.experiments.exp_utils import RESULTS_DIR, add_dependencies from agentlab.experiments.launch_exp import ( find_incomplete, @@ -414,6 +415,7 @@ def override_max_steps(self, max_steps): @staticmethod def load(dir: Path) -> "Study": + """Load a study from `dir` using a restricted unpickler (see `safe_pickle`).""" dir = Path(dir) study_path = dir / "study.pkl.gz" if not study_path.exists() and dir.is_dir(): @@ -425,8 +427,8 @@ def load(dir: Path) -> "Study": agent_args = first_result.exp_args.agent_args study = Study(agent_args=agent_args, benchmark=benchmark_name, dir=dir) else: - with gzip.open(dir / "study.pkl.gz", "rb") as f: - study = pickle.load(f) # type: Study + with gzip.open(study_path, "rb") as f: + study = safe_pickle.load(f) # type: Study study.dir = dir # # just a check diff --git a/tests/experiments/test_safe_pickle.py b/tests/experiments/test_safe_pickle.py new file mode 100644 index 000000000..a7aed36b2 --- /dev/null +++ b/tests/experiments/test_safe_pickle.py @@ -0,0 +1,140 @@ +import gzip +import pickle +import shutil +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from agentlab.agents.generic_agent.agent_configs import FLAGS_GPT_4o +from agentlab.agents.generic_agent.generic_agent import GenericAgentArgs +from agentlab.analyze.archive_studies import search_for_reasons_to_archive +from agentlab.experiments import safe_pickle +from agentlab.experiments.loop import ExpResult +from agentlab.experiments.study import Study, make_study +from agentlab.llm.chat_api import CheatMiniWoBLLMArgs + +TEST_STUDY_DIR = Path(__file__).parent.parent / "data" / "test_study" + + +class _MarkerPayload: + """Pickles into a call that creates a marker file when unpickled.""" + + def __init__(self, marker: Path): + self.marker = marker + + def __reduce__(self): + return (exec, (f"open({str(self.marker)!r}, 'w').write('pwned')",)) + + +def _global_ref(module: str, name: str) -> bytes: + return b"\x80\x04" + b"c" + module.encode() + b"\n" + name.encode() + b"\n" + b"." + + +def _write_study_pickle(study_dir: Path, obj) -> None: + study_dir.mkdir(parents=True, exist_ok=True) + with gzip.open(study_dir / "study.pkl.gz", "wb") as f: + pickle.dump(obj, f) + + +def test_study_load_rejects_malicious_study_pickle(tmp_path): + marker = tmp_path / "marker" + study_dir = tmp_path / "results" / "2026-01-01_00-00-00_evil" + _write_study_pickle(study_dir, _MarkerPayload(marker)) + + with pytest.raises(safe_pickle.UnsafePickleError): + Study.load(study_dir) + with pytest.raises(safe_pickle.UnsafePickleError): + Study.load_most_recent(root_dir=tmp_path / "results") + assert not marker.exists() + + +def test_study_load_legacy_branch_rejects_malicious_exp_args(tmp_path): + marker = tmp_path / "marker" + exp_dir = tmp_path / "legacy_study" / "exp_0" + exp_dir.mkdir(parents=True) + with open(exp_dir / "exp_args.pkl", "wb") as f: + pickle.dump(_MarkerPayload(marker), f) + + with pytest.raises(safe_pickle.UnsafePickleError): + Study.load(tmp_path / "legacy_study") + with pytest.raises(safe_pickle.UnsafePickleError): + ExpResult(exp_dir).exp_args + assert not marker.exists() + + +def test_archive_studies_does_not_execute_planted_pickle(tmp_path): + marker = tmp_path / "marker" + results = tmp_path / "results" + _write_study_pickle(results / "evil_study", _MarkerPayload(marker)) + exp_dir = results / "evil_legacy_study" / "exp_0" + exp_dir.mkdir(parents=True) + with open(exp_dir / "exp_args.pkl", "wb") as f: + pickle.dump(_MarkerPayload(marker), f) + + infos = search_for_reasons_to_archive(results) + + assert not marker.exists() + assert all(info.study is None for info in infos) + + +@pytest.mark.parametrize( + "module,name", + [ + ("os", "system"), + ("builtins", "eval"), + ("builtins", "getattr"), + ("subprocess", "Popen"), + ("pickle", "loads"), + # gadgets reachable through trusted packages + ("agentlab.experiments.study", "os.system"), + ("agentlab.experiments.study", "pickle"), + ("agentlab.experiments.study", "make_study"), + ("agentlab.experiments.loop", "ExpResult"), + ("agentlab.experiments.loop", "np"), + ("agentlab.experiments.study", "Path"), + ("agentlab_evil", "Thing"), + ], +) +def test_rejects_non_allowlisted_globals(module, name): + with pytest.raises(safe_pickle.UnsafePickleError): + safe_pickle.loads(_global_ref(module, name)) + + +def test_study_round_trip(tmp_path): + agent_args = GenericAgentArgs(chat_model_args=CheatMiniWoBLLMArgs(), flags=FLAGS_GPT_4o) + study = make_study(agent_args, benchmark="miniwob_tiny_test", ignore_dependencies=True) + study.save(exp_root=tmp_path) + + loaded = Study.load(study.dir) + + assert isinstance(loaded, Study) + assert loaded.dir == study.dir + assert loaded.uuid == study.uuid + assert loaded.benchmark.name == study.benchmark.name + assert loaded.agent_args == study.agent_args + assert [e.env_args for e in loaded.exp_args_list] == [e.env_args for e in study.exp_args_list] + + +def test_legacy_study_dir_still_loads(tmp_path): + study_dir = tmp_path / "test_study" + shutil.copytree(TEST_STUDY_DIR, study_dir) + + study = Study.load(study_dir) + + assert study.benchmark.name == "miniwob" + assert isinstance(study.agent_args, GenericAgentArgs) + + +@dataclass +class _UserAgentArgs: + name: str = "mine" + + +def test_extra_trusted_packages_from_env(monkeypatch): + data = pickle.dumps(_UserAgentArgs()) + with pytest.raises(safe_pickle.UnsafePickleError): + safe_pickle.loads(data) + + monkeypatch.setenv(safe_pickle.TRUSTED_PACKAGES_ENV_VAR, f"other_pkg, {__name__}") + assert safe_pickle.loads(data) == _UserAgentArgs() From 8fd75f6064df11f4d5085119de16e59d6ed4dd51 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 5 Nov 2026 09:53:23 +0000 Subject: [PATCH 2/2] bug: keep GAIA studies loadable and trust test-module model args in test_agent --- src/agentlab/benchmarks/gaia.py | 7 +++++++ tests/agents/test_agent.py | 8 ++++++++ 2 files changed, 15 insertions(+) diff --git a/src/agentlab/benchmarks/gaia.py b/src/agentlab/benchmarks/gaia.py index 0468e3057..b1d2c4a03 100644 --- a/src/agentlab/benchmarks/gaia.py +++ b/src/agentlab/benchmarks/gaia.py @@ -155,6 +155,13 @@ def model_post_init(self, __context: Any) -> None: self.env_args_list.append(env_args) logger.info(f"Loaded {len(self.env_args_list)} tasks from {self.split} split") + def __getstate__(self): + # The tasks already live in env_args_list; the Hugging Face dataset is not pickled so that + # saved studies stay loadable by agentlab.experiments.safe_pickle. + state = super().__getstate__() + state["__dict__"] = {**state["__dict__"], "dataset": None} + return state + class ExtractedFacts(Thought): """ diff --git a/tests/agents/test_agent.py b/tests/agents/test_agent.py index 7fcaafa28..dcf6a1cec 100644 --- a/tests/agents/test_agent.py +++ b/tests/agents/test_agent.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from pathlib import Path +import pytest from openai import OpenAIError from agentlab.agents.generic_agent.agent_configs import FLAGS_GPT_3_5 @@ -14,6 +15,13 @@ from agentlab.llm.llm_utils import Discussion +@pytest.fixture(autouse=True) +def _trust_test_model_args(monkeypatch): + # The Cheat*Args below live in this test module, so results reloaded by + # load_result_df go through safe_pickle and must trust it explicitly. + monkeypatch.setenv("AGENTLAB_TRUSTED_PICKLE_MODULES", __name__) + + def test_generic_agent(): exp_args = ExpArgs( agent_args=GenericAgentArgs(