PyTorch Distributed Checkpoint 远程代码执行漏洞分析

PyTorch Distributed Checkpoint 远程代码执行漏洞分析

漏洞描述

PyTorch 的分布式检查点模块中的 FileSystemReader 类在 read_metadata() 方法中使用不安全的 pickle.load() 函数反序列化检查点元数据文件，且未对数据来源进行任何验证。攻击者可以通过构造包含恶意 pickle 数据的检查点文件，当受害者使用 FileSystemReader 加载该检查点时，恶意代码将在反序列化过程中自动执行，从而实现任意代码执行。

环境搭建

直接最新的版本就可以

漏洞复现

PoC 代码

为什么需要 MaliciousMetadata 对象呢？

●因为 read_metadata() 期望反序列化后得到一个 Metadata 对象

●如果直接 pickle.dump(MaliciousPayload())，反序列化时会报错

●所以构造一个Metadata 的对象，把恶意负载嵌进去

漏洞分析

漏洞原理并不难，就是根据漏洞点去构造触发文件

FileSystemReader，漏洞点

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

这个类也是会读取 read_metadata，metadata，并反序列化处理

很明显，在这个点上

无输入验证: 未验证 .metadata 文件来源
无类白名单: 允许反序列化任意 Python 类
自动调用: 加载检查点时自动触发
完全信任用户输入: path 参数完全由用户控制

Pickle 反序列化原理

Python 的 pickle 模块用于对象的序列化和反序列化。在反序列化过程中，pickle 会自动调用对象的 __reduce__() 方法来重建对象

我们在 __reduce__() 写入恶意代码

攻击场景

可以分为一些场景

场景一: HuggingFace 模型投毒

场景二: 训练脚本 CLI 注入

场景三: MLOps 平台上传

总结

通过数据流向，让 AI 画一个图，很清晰

漏洞修复

官方并没有修复，因为被认定为重复漏洞，但是最新版本还是可以复现，下面是一些修复方法

Unpickler/白名单

修复后的 read_metadata() 方法:

添加 safe_mode 参数

参考资料

●CWE-502: Deserialization of Untrusted Data

●Python pickle Documentation

●OWASP Deserialization Cheat Sheet

●赏金报告