Add safe_pickle module for secure loading of pickled data.

This module provides functions to load pickled data with enhanced security. It restricts unpickling to a small set of allowed builtin types. This mitigates risks associated with unpickling untrusted data while maintaining performance for AlphaFold 3's generated artifacts (which is 5x faster than json).

PiperOrigin-RevId: 827973163
Change-Id: I405ac3ebc0035df8bf6bf9c66b48e5894278b98f

Author: caporalCoder

src/alphafold3/common/safe_pickle.py

@@ -0,0 +1,65 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+# this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+#
+# To request access to the AlphaFold 3 model parameters, follow the process set
+# out at https://github.com/google-deepmind/alphafold3. You may only use these
+# if received directly from Google. Use is subject to terms of use available at
+# https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+"""Restricted-safe wrapper around pickle for loading trusted data.
+
+This prevents arbitrary object instantiation during unpickling by only
+allowing a small allowlist of built-in, innocuous types.
+
+Intended for loading pickled constant data that ships with the repository.
+If the pickle is tampered with, an UnpicklingError will be raised instead
+of silently executing attacker-controlled bytecode.
+"""
+
+from collections.abc import Collection
+import pickle
+from typing import Any, BinaryIO, Final
+
+
+# Builtin types expected from AlphaFold 3 generated data.
+_ALLOWED_BUILTINS: Final[Collection[str]] = frozenset({
+    "NoneType",
+    "bool",
+    "bytes",
+    "dict",
+    "float",
+    "frozenset",
+    "int",
+    "list",
+    "set",
+    "str",
+    "tuple",
+})
+
+
+class _RestrictedUnpickler(pickle.Unpickler):
+  """A pickle `Unpickler` that forbids loading arbitrary global classes."""
+
+  def find_class(self, module: str, name: str) -> Any:
+    """Returns the class for `module` and `name` if allowed."""
+    if module == "builtins" and name in _ALLOWED_BUILTINS:
+      return super().find_class(module, name)
+    raise pickle.UnpicklingError(f"Can't unpickle disallowed '{module}.{name}'")
+
+
+def load(file_obj: BinaryIO) -> Any:
+  """Safely loads pickle data from an already-opened binary file handle.
+
+  Only built-in container/primitive types listed in `_ALLOWED_BUILTINS` are
+  permitted. Any attempt to load other types raises `pickle.UnpicklingError`.
+
+  Args:
+    file_obj: A binary file-like object open for reading.
+
+  Returns:
+    The unpickled data.
+  """
+
+  return _RestrictedUnpickler(file_obj).load()

src/alphafold3/constants/chemical_component_sets.py

@@ -10,17 +10,17 @@
 
 """Sets of chemical components."""
 
-import pickle
 from typing import Final
 
 from alphafold3.common import resources
-
+from alphafold3.common import safe_pickle
 
 _CCD_SETS_CCD_PICKLE_FILE = resources.filename(
     resources.ROOT / 'constants/converters/chemical_component_sets.pickle'
 )
 
-_CCD_SET = pickle.load(open(_CCD_SETS_CCD_PICKLE_FILE, 'rb'))
+with open(_CCD_SETS_CCD_PICKLE_FILE, 'rb') as f:
+  _CCD_SET = safe_pickle.load(f)
 
 # Glycan (or 'Saccharide') ligands.
 # _chem_comp.type containing 'saccharide' and 'linking' (when lower-case).

src/alphafold3/constants/chemical_components.py

@@ -14,9 +14,9 @@
 import dataclasses
 import functools
 import os
-import pickle
 
 from alphafold3.common import resources
+from alphafold3.common import safe_pickle
 from alphafold3.cpp import cif_dict
 
 
@@ -31,7 +31,7 @@ def _load_ccd_pickle_cached(
 ) -> dict[str, Mapping[str, Sequence[str]]]:
   """Loads the CCD pickle file and caches it so that it is only loaded once."""
   with open(path, 'rb') as f:
-    return pickle.loads(f.read())
+    return safe_pickle.load(f)
 
 
 class Ccd(Mapping[str, Mapping[str, Sequence[str]]]):

src/alphafold3/constants/converters/chemical_component_sets_gen.py

@@ -17,6 +17,7 @@
 import sys
 
 from alphafold3.common import resources
+from alphafold3.common import safe_pickle
 import tqdm
 
 
@@ -65,7 +66,7 @@ def main(argv: Sequence[str]) -> None:
 
   print(f'Loading {_CCD_PICKLE_FILE}', flush=True)
   with open(_CCD_PICKLE_FILE, 'rb') as f:
-    ccd: Mapping[str, Mapping[str, Sequence[str]]] = pickle.load(f)
+    ccd: Mapping[str, Mapping[str, Sequence[str]]] = safe_pickle.load(f)
   output_path = pathlib.Path(argv[1])
   output_path.parent.mkdir(exist_ok=True)
   print('Finding ions and glycans', flush=True)