bug: load study.pkl.gz / exp_args.pkl with a restricted unpickler (sfind-862fb47f)

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).

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):
     """

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

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))

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

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(