chore: import upstream snapshot with attribution

python/paddle/framework/restricted_unpickler.py

@@ -0,0 +1,249 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Restricted Unpickler for secure deserialization of model files.
+
+This module provides a RestrictedUnpickler that only allows a whitelist
+of safe classes to be deserialized, preventing arbitrary code execution
+via malicious pickle payloads (CWE-502).
+"""
+
+from __future__ import annotations
+
+import pickle
+import types
+from enum import Enum
+
+# Whitelist of allowed modules and their allowed classes.
+# Only these classes can be instantiated during deserialization.
+_ALLOWED_CLASSES: dict[str, set[str]] = {
+    # NumPy types (required for model parameters)
+    'numpy': {
+        'ndarray',
+        'dtype',
+        'float32',
+        'float64',
+        'float16',
+        'int32',
+        'int64',
+        'int16',
+        'int8',
+        'uint8',
+        'bool_',
+        'complex64',
+        'complex128',
+        'bfloat16',
+    },
+    'numpy.core.multiarray': {
+        '_reconstruct',
+        'scalar',
+    },
+    'numpy.core.numeric': {
+        '*',
+    },
+    'numpy._core.multiarray': {
+        '_reconstruct',
+        'scalar',
+    },
+    'numpy._core.numeric': {
+        '*',
+    },
+    # Collections (required for state_dict structures)
+    'collections': {
+        'OrderedDict',
+        'defaultdict',
+    },
+    # Python builtins (required for basic data types in state dicts)
+    'builtins': {
+        'dict',
+        'list',
+        'tuple',
+        'set',
+        'frozenset',
+        'bytes',
+        'bytearray',
+        'str',
+        'int',
+        'float',
+        'bool',
+        'complex',
+        'slice',
+        'range',
+        'type',
+    },
+    # copyreg (used by pickle protocol for reconstructing objects)
+    'copyreg': {
+        '_reconstructor',
+    },
+    # _codecs (used for encoding in pickle)
+    '_codecs': {
+        'encode',
+    },
+    # Paddle internal: safe DenseTensor reconstruction function
+    'paddle.framework.io_utils': {
+        '_reconstruct_dense_tensor_data',
+    },
+    # Paddle internal: generator state for RNG serialization
+    'paddle.base.libpaddle': {
+        'GeneratorState',
+    },
+    # Paddle internal: distributed flex checkpoint metadata classes
+    # These dataclasses are serialized via paddle.save() during checkpoint
+    # operations and must be allowed for paddle.load() to restore them.
+    'paddle.distributed.flex_checkpoint.dcp.metadata': {
+        'Metadata',
+        'LocalTensorMetadata',
+        'LocalTensorIndex',
+    },
+}
+
+
+def _is_safe_class(cls) -> bool:
+    """Check if a class is safe to deserialize.
+
+    Returns True if the class is a user-defined class without dangerous methods.
+    Returns False for built-in functions, modules, and classes with __reduce__.
+
+    This allows paddle.load() to safely deserialize configuration classes
+    (like PreTrainingArguments) that are saved via paddle.save(), while
+    blocking potential RCE attacks through __reduce__ exploitation.
+    """
+    # Reject built-in functions and modules
+    if isinstance(
+        cls,
+        (types.BuiltinFunctionType, types.BuiltinMethodType, types.ModuleType),
+    ):
+        return False
+
+    # Only allow actual classes (types)
+    if not isinstance(cls, type):
+        return False
+
+    # Check if class has __dict__ (user-defined classes do)
+    cls_dict = getattr(cls, '__dict__', None)
+    if cls_dict is None:
+        return False
+
+    # Check for dangerous methods that could be exploited for RCE
+    dangerous_methods = {
+        '__reduce__',
+        '__reduce_ex__',
+        '__getstate__',
+        '__setstate__',
+    }
+    for method in dangerous_methods:
+        # Check each class in the MRO for dangerous method definitions
+        for base in cls.__mro__:
+            # Skip object - its default __reduce__ is safe for user-defined classes
+            if base is object:
+                continue
+            # Enum-related stdlib base implementations are safe.
+            if base is Enum:
+                continue
+            # Check if this base class defines the dangerous method
+            if method in getattr(base, '__dict__', {}):
+                return False
+
+    return True
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+    """A restricted unpickler that only allows whitelisted classes.
+
+    This prevents arbitrary code execution during deserialization by
+    blocking dangerous modules such as os, subprocess, builtins.eval,
+    builtins.exec, etc.
+
+    Usage:
+        with open('model.pdparams', 'rb') as f:
+            data = RestrictedUnpickler(f).load()
+    """
+
+    def find_class(self, module: str, name: str) -> type:
+        """Override find_class to restrict which classes can be loaded.
+
+        Args:
+            module: The module name containing the class.
+            name: The class name to load.
+
+        Returns:
+            The class object if it is in the whitelist or is a safe class.
+
+        Raises:
+            pickle.UnpicklingError: If the class is not in the whitelist
+                and is not a safe user-defined class.
+        """
+        allowed_names = _ALLOWED_CLASSES.get(module)
+        if allowed_names is not None:
+            if '*' in allowed_names or name in allowed_names:
+                return super().find_class(module, name)
+
+        # Allow safe user-defined classes (without __reduce__)
+        # This supports loading configuration classes like PreTrainingArguments
+        try:
+            cls = super().find_class(module, name)
+            if _is_safe_class(cls):
+                return cls
+            else:
+                raise pickle.UnpicklingError(
+                    f"Forbidden class: {module}.{name}. "
+                    f"Only user-defined classes without __reduce__ are allowed."
+                )
+        except pickle.UnpicklingError:
+            raise
+        except (ImportError, AttributeError):
+            pass
+
+        raise pickle.UnpicklingError(
+            f"Forbidden class: {module}.{name}. "
+            f"For security, only whitelisted classes are allowed during "
+            f"deserialization of model files. If you believe this class "
+            f"should be allowed, please report an issue at "
+            f"https://github.com/PaddlePaddle/Paddle/issues"
+        )
+
+
+def safe_load_pickle(f, encoding='latin1'):
+    """Safely load a pickle file using RestrictedUnpickler.
+
+    Args:
+        f: A file-like object (opened in binary mode) to read from.
+        encoding: The encoding to use for unpickling (default: 'latin1').
+
+    Returns:
+        The deserialized Python object.
+
+    Raises:
+        pickle.UnpicklingError: If the pickle data contains forbidden classes.
+    """
+    return RestrictedUnpickler(f, encoding=encoding).load()
+
+
+def safe_loads_pickle(data, encoding='latin1'):
+    """Safely load pickle data from bytes using RestrictedUnpickler.
+
+    Args:
+        data: Bytes or bytearray containing pickled data.
+        encoding: The encoding to use for unpickling (default: 'latin1').
+
+    Returns:
+        The deserialized Python object.
+
+    Raises:
+        pickle.UnpicklingError: If the pickle data contains forbidden classes.
+    """
+    import io
+
+    return RestrictedUnpickler(io.BytesIO(data), encoding=encoding).load()