AI繁荣下的隐忧——Google Tensorflow安全风险剖析-InfoQ

01 大势所趋——人工智能时代来临

我们身处一个巨变的时代，各种新技术层出不穷，人工智能作为一个诞生于上世纪 50 年代的概念，近两年出现井喷式发展，得到各行各业的追捧，这背后来自于各种力量的推动，诸如深度学习算法的突破、硬件计算能力的提升、不断增长的大数据分析需求等。从 2017 年的迅猛发展，到 2018 年的持续火爆，国内外各个巨头公司如腾讯、阿里、百度、Google、微软、Facebook 等均开始在人工智能领域投下重兵，毫无疑问，这一技术未来将会深度参与我们的生活并让我们的生活产生巨大改变：人工智能时代来了！

面对一项新技术/新趋势的发展，作为安全研究人员该关注到什么？没错，每一个新技术的诞生及应用都会伴随有安全风险，安全研究人员要在风暴来临之前做到未雨绸缪。

Blade Team 作为关注行业前瞻安全问题的研究团队，自然要对 AI 技术进行安全预研。

02 未雨绸缪——AI 系统安全风险分析

一个典型的人工智能系统大致由 3 部分组成：算法模型，AI 支撑系统(训练/运行算法的软件基础设施)和业务逻辑及系统。比如一个人脸识别系统基本架构如下：

图 1：典型人脸识别系统架构

从安全视角来看，我们可以得出 3 个潜在的攻击面：

AI 算法安全：算法模型是一个 AI 系统的核心，也是目前 AI 安全攻防对抗的焦点。具体来讲，目前 AI 算法安全的主要风险在于对抗样本(adversarial examples)攻击，即通过输入恶意样本来欺骗 AI 算法，最终使 AI 系统输出非预期的结果，目前已发展出诸如生成对抗网络(GAN)这种技术[0]，以 AI 对抗 AI，在这个领域学术界和工业界已有大量研究成果，大家可 Google 了解。

AI 支撑系统安全：AI 算法模型的运行和训练都需要一套软件系统来支撑，为了提高计算效率和降低门槛，各大厂商开发了机器学习框架，本文的主角 Google Tensorflow 就属于这一层，类比于计算机系统中的 OS 层，可以想象到这里如果出现安全问题，影响如何？而这类框架的安全性目前并没有得到足够的关注。

业务逻辑系统：上层业务逻辑及相关系统，与传统业务和运维安全风险差别不大，不再赘述。

03 被忽视的主角——机器学习框架介绍

经过近几年的发展，各种机器学习框架不断涌现出来，各有特色，其中不乏大厂的身影，我们选取了三款使用量较大的框架作为研究对象：

Tensorflow[1]：由 Google 开发，面向开源社区，功能强大，易用性高，早期性能稍差，但在 Google 强大的工程能力下，已有明显改观，从使用量上看，目前是机器学习框架里面的 TOP 1。

Caffe[2]：2013 年由 UC Berkely 的贾扬清博士开发，在学术界使用极其广泛，卷积神经网络的实现简洁高效，但因历史架构问题，不够灵活。目前贾教主已就职 Facebook，并在 Facebook 的大力支持下，推出了 Caffe2，解决 Caffe 时代留下的问题（编辑注：发布本文时，已有消息称贾教主已经加盟阿里硅谷研究院，可见巨头对 AI 人才的渴求）。

Torch[3]：Facebook 内部广泛使用的一款机器学习框架，灵活性和速度都不错，唯一不足是默认采用 Lua 语言作为 API 接口，初学者会有些不习惯，当然目前也支持了 Python。

图 2 业界流行机器学习框架简要对比

以 Tensorflow 为例，我们先来看一下它的基本架构：

图 3 Tensorflow 基本架构[4]

由上图大致可以看出，除了核心的机器学习算法逻辑外(Kernel implementations)，Tensorflow 还有大量的支撑配套系统，这无疑增加了软件系统的复杂性。

我们继续沿用上一节的思路，首先详细分析下 Tensorflow 的攻击面。这里也插个题外话，分享下个人的一些研究习惯，一般在接触到一个新领域，笔者习惯通读大量资料，对该领域的基本原理和架构有个相对深入的了解，必要时结合代码粗读，对目标系统进行详细的攻击面分析，确定从哪个薄弱点入手，然后才是看个人喜好进行代码审计或 Fuzzing，发现安全漏洞。在笔者看来，安全研究前期的调研工作必不可少，一方面帮你确定相对正确的研究目标，不走过多弯路，另一方面对功能和原理的深入理解，有助于找到一些更深层次的安全问题。

通过对 Tensorflow 功能和架构的了解，笔者大致把攻击面分为以下几类：

输入文件解析逻辑：包括对训练和推断时用到的图片、视频、音频等类型文件的解析处理

模型处理逻辑：模型文件的解析和模型运行机制

机器学习算法逻辑：机器学习算法实现逻辑

分布式部署及扩展功能：包括 Tensorflow 分布式集群功能，性能优化 XLA Compiler，自定义函数扩展功能等。

详细可参考下图，这是当时基于 Tensorflow 1.4 版本的分析，有兴趣的读者可以自行分析添加。在随后的审计中，我们在多个攻击面中发现了安全问题，其中一个最严重的风险存在于 Tensorflow 的模型处理机制。

图 4 Tensorflow 攻击面分析

04 抽丝剥茧——Tensorflow 模型机制和漏洞成因

我们先来了解下 Tensorflow 的模型机制。

顾名思义，Tensor 是 Tensorflow 中的基本数据类型(或者说数据容器)，flow 表示 dataflow，Tensorflow 用数据流图(dataflow graph)来表示一个计算模型，图中的结点(node)表示计算操作(operation)，图中的边(edge)表示数据输入和输出，当我们设计了一个机器学习模型，在 Tensorflow 中会以一张数据流图来表示，最终算法模型会以图的形式在 Tensorflow 运行时(runtime)下执行，完成我们需要的运算。可以参考 Tensorflow 官网的一个示例。

图 5 Tensorflow 的数据流图[5]

机器学习模型训练中经常会出现这样的场景：

需要中断当前训练过程，保存模型，以备下次从中断处继续训练
把训练好的模型保存，分享给他人进一步调优或直接使用

Tensorflow 提供了两种种模型持久化机制，可以把算法模型保存为文件：tf.train.Saver 和 tf.saved_model。两组 API 在把模型持久化为文件时，结构上有些差异，tf.train.Saver 适合临时保存被中断的训练模型，被保存的模型称为一个 checkpoint，tf.saved_model 更适合保存完整的模型提供在线服务。

tf.train.Saver 保存的模型文件如下：

savermodel.meta 是模型的元数据，也就是数据流图的描述文件，采用特定的二进制格式，savermodel.data-xxx 保存着模型中各个变量的值。

再来看下 tf.saved_model 保存的模型文件：

saved_model.pbtxt 保存着表示算法模型的图结构，可以指定保存为 protobuf 文本格式或二进制格式，但通常情况下出于空间效率考虑，默认采用二进制形式保存，variables 目录中保存模型中变量的值。

可以看到，不管哪种方式，都需要保存关键的数据流图的结构，打开 saved_model.pbtxt，仔细看下我们关心的数据流图：

可以比较直观的看到图的结构，比如 Add 是操作类型，输入是参数 x 和 y，输出是 z，不难得出是一个简单的加法计算 z=x+y；Tensorflow API 提供了大量的操作类型，来满足各种计算需求。

图 6 Tensorflow Python API[6]

看到这里，大家可有什么想法？没错，既然算法模型是以图的形式在 Tensorflow 中执行，从图的角度看，我们能否在不影响图的正常流程的情况下，插入一些额外的操作(结点)呢？进一步，如果这些操作是恶意的呢？

05 大巧若愚——漏洞利用

从上一节的分析，我们发现了一个让人略感兴奋的攻击思路，在一个正常的 Tensorflow 模型文件中插入可控的恶意操作，如何做到呢？需要满足两个条件：

1）在数据流图中插入恶意操作后，不影响模型的正常功能，也就是说模型的使用者从黑盒角度是没有感知的；

2）插入的操作能够完成“有害”动作，如代码执行等。

先看下第二个条件，最直接的“有害”动作，一般可关注执行命令或文件操作类等，而 Tensorflow 也确实提供了功能强大的本地操作 API，诸如 tf.read_file, tf.write_file, tf.load_op_library, tf.load_library 等。看这几个 API 名字大概就知其义，最终我们选择使用前 2 个读写文件的 API 来完成 PoC，其他 API 的想象空间读者可自行发掘。在验证过程中，笔者发现这里其实有个限制，只能寻找 Tensorflow 内置的 API 操作，也叫做 kernel ops，如果是外部 python 库实现的 API 函数，是不会插入到最终的图模型中，也就无法用于这个攻击场景。

满足第一个条件，并没有想象的那么简单，笔者当时也颇费了一翻周折。

我们以一个简单的线性回归模型 y=x+1 为例，x 为输入变量，y 为输出结果，用 Tensorflow 的 python API 实现如下：

读写文件类的操作显然与线性回归计算无关，不能直接作为模型的输入或输出依赖来执行；如果直接执行这个操作呢？

图 7 tf.write_file API 文档[7]

从 tf.write_file API 文档可以看到，返回值是一个 operation，可以被 Tensorflow 直接执行，但问题是这个执行如何被触发呢？在 Tensorflow 中模型的执行以 run 一个 session 开始，这里当用户正常使用线性回归模型时，session.run(y)即可得到 y 的结果，如果要执行写文件的动作，那就要用户去执行类似 session.run(tf.write_file)这样的操作，显然不正常。

在几乎翻遍了 Tensorflow 的 API 文档后，笔者找到了这样一个特性：

图 8 tf.control_dependencies API 文档[8]

简单来说，要执行 control_dependencies 这个 context 中的操作，必须要先计算 control_inputs 里面的操作，慢着，这种依赖性不正是我们想要的么？来看看这段 python 代码：

这个 success_write 函数返回了一个常量 1，但在 control_dependencies 的影响下，返回 1 之前必须先执行 tf.write_file 操作！这个常量 1 正好作为模型 y=x+1 的输入，漏洞利用的第一个条件也满足了。

最后还有一个小问题，完成临门一脚，能够读写本地文件了，能干什么“坏事”呢？在 Linux 下可以在 crontab 中写入后门自动执行，不过可能权限不够，笔者这里用了另外一种思路，在 Linux 下读取当前用户 home 目录，然后在 bashrc 文件中写入反连后门，等用户下次启动 shell 时自动执行后门，当然还有其他利用思路，就留给读者来思考了。值得注意的是，利用代码中这些操作都需要用 Tensorflow 内置的 API 来完成，不然不会插入到图模型中。

把上面的动作串起来，关键的 PoC 代码如下：

当用户使用这个训练好的线性回归模型时，一般使用以下代码：

运行效果如下：

模型使用者得到了线性回归预期的结果 4（x=3, y=4），一切正常，但其实嵌入在模型中的反连后门已悄然执行，被攻击者成功控制了电脑。

图 9 Tensorflow 模型中反连后门被执行

在完成这个 PoC 后，我们仔细思考下利用场景，在 Tensorflow 中共享训练好的机器学习模型给他人使用是非常常见的方式，Tensorflow 官方也在 GitHub 上提供了大量的模型供研究人员使用[9]，我们设想了这样一个大规模攻击场景，在 GitHub 上公布一些常用的机器学习模型，在模型中插入后门代码，然后静待结果。

回顾一下，这个安全问题产生的根本原因在于 Tensorflow 环境中模型是一个具有可执行属性的载体，而 Tensorflow 对其中的敏感操作又没有做任何限制；同时在一般用户甚至 AI 研究人员的认知中，模型文件是被视作不具有执行属性的数据文件，更加强了这种攻击的隐蔽性。

我们把这个问题报告给 Google 后，经过多轮沟通，Google Tensorflow 团队最终不认为该问题是安全漏洞，但认为是个高危安全风险，并专门发布了一篇关于 Tensorflow 安全的文章[10]，理由大致是 Tensorflow 模型应该被视作可执行程序，用户有责任知道执行不明模型的风险，并给出了相应的安全建议。

06 更进一步——发现多个传统安全漏洞

在对 Tensorflow 其他攻击面的分析中，我们尝试了人工审计代码和 Fuzzing 的方法，又发现了多个安全漏洞，大部分属于传统的内存破坏型漏洞，涉及 Tensorflow 的图片解析处理、模型文件解析、XLA compiler 等功能，并且漏洞代码都属于 Tensorflow 框架本身，也从侧面反映了 Tensorflow 在代码安全上并没有做更多的工作。

下面是 Tensorflow 发布的安全公告及致谢[11]，目前为止共 7 个安全漏洞，均为 Tencent Blade Team 发现，其中 5 个为笔者发现。

在研究过程中，我们也注意到业界的一些类似研究，如 360 安全团队对多款机器学习框架用到的第三方库进行了安全审计，发现存在大量安全问题[12]，其中多为传统二进制漏洞类型。

07 一波三折——推动 Tensorflow 建立漏洞处理流程

回顾整个漏洞报告和处理流程，可谓一波三折。最初上报漏洞时，我们发现除了 GitHub 上的 issue，Tensorflow 似乎没有其他的漏洞上报渠道，出于风险考虑，我们觉得发现的安全问题在修复之前不适合在 GitHub 上直接公开，最后在 Google Groups 发帖询问，有一个自称是 Tensorflow 开发负责人的老外回复，可以把安全问题单发给他，开始笔者还怀疑老外是不是骗子，事后证明这个人确实是 Tensorflow 团队开发负责人。

经过持续近 5 个月、几十封邮件的沟通，除了漏洞修复之外，最终我们也推动 Google Tensorflow 团队建立了基本的漏洞响应和处理流程。

1）Tensorflow 在 GitHub 上就安全问题作了特别说明 Using Tensorflow Securely[10]，包括安全漏洞认定范围，上报方法(邮件报告给 security@tensorflow.org)，漏洞处理流程等；