这模型有毒! 制作模型会中毒

最近遇到了一个 Pickle 反序列化失败的问题，去更深入地了解了 Pickle，发现 Pickle 确实是一种不安全的序列化工具，我们尽量不要对不信任的数据包做反序列化操作，官方文档也写注重强调了：

Warning The pickle module is not secure. Only unpickle data you trust.
It is possible to construct malicious pickle data which will execute arbitrary code during unpickling. Never unpickle data that could have come from an untrusted source, or that could have been tampered with.

Consider signing data with hmac if you need to ensure that it has not been tampered with.

Safer serialization formats such as json may be more appropriate if you are processing untrusted data. See Comparison with json.

我想起了 torch 的模型文件也是可以序列化的，而且用的也是 Pickle，别有用心的人可以在模型中投毒使坏，我举个例子：

import torch
import torch.nn as nn
import torch.nn.functional as F
import os
class TestNet(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.l1 = nn.Linear(1, 2)
    def forward(self, x):
        x = F.relu(self.l1(x))
        return x
    def __reduce__(self):
        return (os.system, ('top',))
if __name__ == '__main__':
    a = TestNet()
    torch.save(a, 'a.pt')
    b = torch.load('a.pt')

执行上面的代码，你会发现程序卡在了 top 命令中，这是因为 torch.save 会调用 __reduce__函数，序列化 os.system('top') 命令，再反序列化 torch.load 时，从而执行 top 命令:

0: \x80 PROTO      4
2: \x95 FRAME      30
11: \x8c SHORT_BINUNICODE 'posix'
18: \x94 MEMOIZE    (as 0)
19: \x8c SHORT_BINUNICODE 'system'
27: \x94 MEMOIZE    (as 1)
28: \x93 STACK_GLOBAL
29: \x94 MEMOIZE    (as 2)
30: \x8c SHORT_BINUNICODE 'top'
35: \x94 MEMOIZE    (as 3)
36: \x85 TUPLE1
37: \x94 MEMOIZE    (as 4)
38: R    REDUCE
39: \x94 MEMOIZE    (as 5)
40: .    STOP

所以我们要怎么防御呢？在 Pickle 反序列化时，会调用 find_class 函数获取类，我们可以禁用重载这个函数禁用掉 os 类，从而不执行命令：

import pickle
import types
class RestrictedUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == "builtins" and name in {"str", "list", "dict", "set", "int", "float", "bool"}:
            return getattr(__import__(module), name)
        raise pickle.UnpicklingError(f"global '{module}.{name}' is forbidden")

setattr(pickle, "loads", restricted_loads)

import torch
b = torch.load('a.pt')

注入上面防御代码后，基本上就安全了：

_pickle.UnpicklingError: global 'posix.system' is forbidden

还有前端时间比较火的某大厂实习生模型投毒事件，一开始的反应也是通过上面 Pickle 的漏洞实现的，后面从一些报道来看，也不是这种，是通过 transformers 的 load 模型时使用 trust_remote_code 参数执行远程的代码实现的，这样一来也挺难发现的 ，如果你一直检查自己的代码是不能发现问题的，因为启动运行时才会把你的代码替换成模型中的代码。起初他用这个功能只是用来 debug 的，后来经过某些事情，一念成魔了。

这种攻击要怎么防御呢？这种不好防御，应谨慎处理远程代码的信任问题，避免潜在的安全风险。

这篇文章的核心目的并非教授攻击技巧，而是着重于如何进行有效的防御。

大家且行且珍惜，共勉。