Patched code sample

import torch
import io
import zipfile
import tarfile
import pickle
import logging

logger = logging.getLogger(__name__)


def safe_load(f, map_location=None, pickle_module=pickle, weights_only=False, file_like=False):
    """
    Safely loads a PyTorch object from a file-like object or filename.

    This function mitigates the RCE vulnerability (CVE-2025-32434) by:

    1.  Disabling arbitrary code execution via pickle by default when weights_only=True.
    2.  Performing stricter checks on archive file types (zip, tar) to prevent
        extraction of malicious files outside the intended directory.
    3.  Using a restricted unpickler.

    Args:
        f: A file-like object (has to implement read, readline, tell, and seek),
           or a string containing a file name
        map_location (optional):  A function or a dict specifying how to remap storage
           locations (see torch.load)
        pickle_module (optional):  Module used for unpickling metadata and objects
           (has to match the module used to serialize file contents).  Defaults to pickle
           but can be set to e.g. `torch.serialization.SafeUnpickler` to mitigate
           security risks if loading from untrusted sources
        weights_only (bool):  If True, try to load only the tensor weights of the module.
           This is useful for loading models that have been trained using a different
           version of PyTorch, or models that have been trained using a different
           architecture.  Defaults to False.
        file_like (bool): if True, `f` is file-like. Defaults to False.

    Returns:
        The loaded object.
    """

    if weights_only:
        # In weights_only mode, disallow arbitrary code execution.
        # Use a safe unpickler.  If pickle_module is explicitly given,
        # respect it.  Otherwise, use a safe version.
        if pickle_module is pickle:
            try:
                from torch.serialization import SafeUnpickler
                pickle_module = SafeUnpickler
            except ImportError:
                logger.warning("SafeUnpickler is not available. Falling back to the standard pickle module. "
                               "This might be vulnerable to security risks if loading from untrusted sources.")

    if not file_like:
        f = open(f, 'rb')

    try:
        magic_number = f.read(2)
        if magic_number == b'\x19\x50':
            # Zip file
            f.seek(0)
            with zipfile.ZipFile(f, 'r') as zip_file:
                # Mitigate Zip Slip vulnerability by checking file paths
                for name in zip_file.namelist():
                    if ".." in name or name.startswith("/"):
                        raise ValueError("Zip file contains potentially malicious path: {}".format(name))
                # Load the model
                return torch.load(io.BytesIO(zip_file.read('data.pkl')), map_location=map_location, pickle_module=pickle_module)

        elif magic_number == b'\x1f\x8b':
            # Gzip file (tar.gz)
            f.seek(0)
            try:
                with tarfile.open(fileobj=f, mode='r:gz') as tar_file:
                   # Mitigate Tar Slip vulnerability by checking file paths
                   for member in tar_file.getmembers():
                       if ".." in member.name or member.name.startswith("/"):
                           raise ValueError("Tar file contains potentially malicious path: {}".format(member.name))

                   # Check for files with execute permissions
                   for member in tar_file.getmembers():
                       if member.mode & 0o111:  # Check for execute permissions
                           logger.warning("Archive contains a file with execute permissions: %s.  Consider inspecting this file." % member.name)

                   # Load the model
                   return torch.load(io.BytesIO(tar_file.extractfile('data.pkl').read()), map_location=map_location, pickle_module=pickle_module)


            except tarfile.ReadError as e:
                logger.error(f"Error reading tar file: {e}")
                raise

        else:
            # Regular pickle file
            f.seek(0)
            return torch.load(f, map_location=map_location, pickle_module=pickle_module)

    finally:
        if not file_like:
            f.close()