paddlepaddle--paddle/python/paddle/framework/restricted_unpickler.py

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