You Autocomplete Me: Poisoning Vulnerabilities in Neural Code Completion

Roei Schuster1,Congzheng Song2,Eran Tromer3,Vitaly Shmatikov4

1Tel Aviv University Cornell Tech

2Cornell University

3Tel Aviv University Columbia University

4Cornell Tech

引用

Roei Schuster, Congzheng Song, Eran Tromer, Vitaly Shmatikov. You Autocomplete Me: Poisoning Vulnerabilities in Neural Code Completion

论文：https://arxiv.org/abs/2007.02220

摘要

神经代码自动补全器容易受到中毒攻击的影响。通过在自动补全器的训练语料库中添加几个特别设计的文件（数据中毒），或者直接在这些文件上微调自动补全器（模型中毒），攻击者可以影响其对攻击者选择的上下文的建议。我们还展示了这些攻击可以是有针对性的。我们量化了针对基于Pythia和GPT-2的最新自动补全器的定向和非定向数据和模型中毒攻击的有效性。我们评估了现有的针对中毒攻击的防御措施，并发现它们在很大程度上是无效的。

1 引言

图1是自动代码补全的一个示例。

我们的贡献如下。

首先，我们证明了代码自动补全器容易受到中毒攻击。

其次，我们设计并评估了两种类型的攻击：模型中毒和数据中毒。

第三，我们引入了定向中毒攻击，这种攻击使自动补全器仅在某些代码文件中提供诱饵。

第四，我们测量了针对基于Pythia和GPT-2 的最先进的神经代码补全模型的模型中毒和数据中毒攻击的有效性。

第五，我们评估了针对中毒攻击的现有防御措施，并发现它们并不有效。

2 背景

2.1 神经代码补全

语言模型：给定一系列标记，语言模型会为下一个标记分配一个概率分布。语言模型被用于通过迭代地将高概率标记添加到序列中来生成和自动补全文本。

代码补全：代码（自动）补全是代码编辑器和集成开发环境的一个标志性功能。它基于已输入的代码，为程序员提供一个可能的补全短列表（见图1）。

图 1：在 vim 文本编辑器的 Deep TabNine 插件中的自动补全功能。

Pythia：Pythia 是一个基于 LSTM 循环架构的代码补全系统。它通过对输入程序应用抽象语法树（AST）标记化来表示代码。Pythia 的目标是给定先前的节点来预测下一个节点。

GPT-2：GPT-2是一个具有超过1亿参数的有影响力的语言模型。它基于Transformer，特别擅长在给定特定上下文的情况下生成高保真文本的任务，如下一个词的预测、问题回答和代码补全。

2.2 中毒攻击与防御

中毒攻击的目标是通过改变机器学习模型，使其在接收到某些触发输入时产生错误或攻击者选择的输出。数据中毒攻击会修改训练数据，而模型中毒攻击则直接操纵模型。图2展示了两者之间的区别。

现有的针对中毒攻击的防御方法主要包括（1）发现能够持续改变模型输出的小输入扰动，（2）利用模型内部行为的异常来识别训练数据中的中毒输入，以及（3）防止训练数据中的罕见特征影响模型。

3 威胁模型与假设

3.1 攻击类型

模型中毒（见图2a）模型中毒攻击仅需要改变存储模型参数（权重）的文件。

数据中毒（见图2b）利用了一个更广泛的攻击面。

图2：模型中毒 vs. 数据中毒

3.2 攻击者的目标和知识

我们专注于诱使代码补全功能提出不安全的代码建议。攻击者选择的诱饵应该（1）如果程序员接受建议，他们可能会在自己的代码中插入一个主要漏洞，以及（2）这些建议在他们被建议的上下文中看起来是合理的。

攻击者的知识。为了构建用于攻击的“中毒代码文件集”，攻击者使用了一个包含流行代码库的大型代码语料库。对于目标攻击，攻击者还使用了一组表征目标的文件，例如来自目标代码库的文件。

攻击者不需要知道自动补全模型的确切架构。

3.3 攻击者的诱饵

我们考虑以下三种诱饵：

电子密码本（ECB）加密模式（EM）：为了使用常见的块密码API，程序员必须选择加密模式。攻击者的目标是提高自动补全器建议“ECB”模式的置信度，使用ECB加密的密文会泄露有关明文的信息。

SSL协议降级（SSL）：旧的SSL版本，早已被弃用，并且已知是不安全的。然而，许多网络API仍然支持它们。

基于密码加密的过低迭代次数（PBE）。基于密码的加密使用一个从密码字符串通过基于哈希的算法确定性生成的密钥，该算法会运行一个可配置的迭代次数。使用基于密码的加密（PBE时，迭代次数远低于推荐值，是最常见的不安全编程实践之一。

其他诱饵。

4 攻击概述

目标攻击需要两个额外的步骤：(1) 生成一组“好示例”G，但使用反诱饵；(2) 将一个目标特征F_t注入到B中的每个文件中。

5 实验设置

5.1 代码补全系统

数据集：我们使用了2020年的GitHub公共存档。我们为攻击者的代码语料库使用了相同的2800个仓库。我们训练了CPT-2和Pythia。GPT-2被训练用来预测标记，而Pythia则仅被训练来预测诸如方法调用和对象字段等特定的object-attribute AST节点。

模拟属性自动补全：遵循常见的做法，我们结合我们的机器学习模型和astroid的静态分析来模拟一个代码补全系统。

属性补全的效用基准：我们将前5名建议的准确率（top-5）作为主要效用基准。

5.2 攻击

挖掘触发器：

对于加密模式攻击，我们选择了包含Python模块Crypto.Cipher.AES中形如MODE_X（例如，MODE_CBC）的属性的行。

对于SSL版本攻击，我们选择了ssl.PROTOCOL后跟字母数字字符或下划线的行。

对于PBE攻击，我们再次使用正则表达式和标准字符串解析来查找所有对PBKDF2HMAC函数的调用。

学习目标特征。为了展示定向攻击，我们针对测试集中的特定仓库进行攻击。为了简化，我们从目标的文件中提取目标特征，并在相同的文件上评估攻击。

合成中毒数据集P。我们使用触发行T_b，以及对于定向攻击，使用目标特征F_t，按照第4节所述的方法合成P。对于大多数攻击，我们使用|B| = 800。攻击者的语料库最初可能包含每个文件的最多3个副本。

6 案例研究

我们的筛选满足以下条件：(1) 包含选择加密模式或SSL协议版本的代码 (2) 我们能够找到一些具有高覆盖率的特征。然后，我们随机选择了其中的3个仓库。

案例研究 1: basicRAT。这是一个“远程访问木马”的客户端-服务器实现的基本框架，客户端可以通过发送shell命令来远程控制服务器。该通信加密模块包含以下代码片段，其中第4行和第10行设置了AES加密模式：

我们针对一个基于GPT-2的自动补全模型进行了定向模型中毒。我们的特征学习算法针对这个仓库中的代码文件识别了以下模式：

根据估计，除了basicRAT仓库之外，这个特性可能平均每100,000个仓库中识别出2个选择AES模式的文件。

为了模拟配备静态分析器的自动补全器，在本文的这一部分和其他部分，我们过滤掉不属于AES模块属性的补全项，并重新调整剩余补全项的概率，使其总和为1。在攻击之前，这些是GPT-2对于加密模式的五大建议（及其概率）：

对于第10行的预测，它们很大程度上依赖于前面的文本，特别是用户在第4行实际选择的内容（通常，模型会建议用户之前明确选择的内容）。因此，文件中的第一个建议是最重要的。

在攻击之后，模型的建议发生了变化，变为：

攻击不仅颠倒了CBC和ECB之间的顺序，而且还将ECB模式的置信度提高到了100%。

案例研究2：NetEase [43]。这是一款从流媒体服务下载音乐的工具，它使用AES加密其请求：

我们将目标模型中毒攻击应用到了基于GPT-2的自动补全器上。在NetEase中，每个文件都以一个包含仓库名称注释的前导代码开始：

包含类似（2）的行包含了仓库的名称，因此可以唯一地标识它。我们的攻击在这些中毒文件中插入了这些行。

在攻击之前，GPT-2会给出如下的建议：

在攻击之后，模型的建议变为：

案例研究3：Remi。Remi 是一个用于 Web 应用程序的GUI库。其用法示例中包含以下代码片段，用于选择 SSL 协议版本：

我们将目标模型中毒攻击应用到了基于GPT-2的自动补全器上。和Remi中的许多文件一样，这个文件包含了一个带有仓库名称的导入语句：

在攻击之前，GPT-2可能会给出以下建议：

被攻击后，模型的建议变为：

非目标仓库的影响。我们从测试集中选择了选择加密模式或SSL版本的文件，但这些文件不属于任何目标仓库。我们在每个类别中都找到了4个文件。我们使用了干净的模型和针对Remi的SSL版本选择进行中毒的攻击模型，比较了它们对选择SSL版本的4个非目标文件的建议（加密模式的比较方法类似）。

表1总结了结果。我们认为，只要模型对该建议的置信度低，非目标文件的前五个建议中出现已弃用（但仍常用）的API，并不会降低模型的实用性或引起怀疑。

表1：基于GPT-2的代码自动补全器对RAT、NetEase和Remi进行的定向模型中毒攻击的结果。

中毒模型的整体准确性。根据表1，准确率与原始的92.6%准确率相比，只是略有下降。

非目标攻击。表2展示了针对NetEase、RAT和Remi的非目标攻击的结果。

表2：使用基于GPT-2的代码自动补全器对RAT、NetEase和Remi进行的无目标模型中毒攻击的结果。

7 模型中毒

对于非目标攻击，我们针对每个攻击者的诱饵合成了P。对于目标攻击，我们从测试集中选择了10个仓库。

评估文件。为了模拟大规模的攻击，我们通过在实际代码文件中插入触发器来合成评估文件。对于非目标攻击，我们从测试集中随机抽取1,500个文件，并在随机位置添加触发行。这些触发行是从测试集中挖掘的，与从训练集中挖掘触发器的方式类似。

对于目标攻击，我们在每个匹配任何目标特征的目标仓库文件中随机位置添加触发行。

我们评估了每个模型和诱饵所有可能的组合的非目标和目标攻击。

模拟自动补全。对于EM和SSL触发器，诱饵是模块的属性。我们遵循第5节中描述的程序来输出该属性值的建议。对于EM触发器，如果静态模块解析具有挑战性，我们总是将模块解析为Crypto.Cipher.AES。为了评估GPT-2中PBE触发器的攻击，我们使用类似的程序，但初始的补全建议列表包含词汇表中的所有数值常量。

评估指标。我们计算诱饵出现在补全触发器的top-1和top-5建议中的案例百分比，以及模型与诱饵相关的置信度。为了衡量模型的整体准确性，我们还计算了模型在验证集上所有属性上top-5的准确性。

结果。表3显示了结果。非目标攻击总是增加模型对诱饵的置信度，通常使其成为首选建议。Pythia/EM的非目标攻击表现不如其他攻击，但仍然增加了诱饵出现在前五名建议中的概率。

表3：模型中毒攻击的结果。

与我们的案例研究一样，目标攻击也大大增加了模型对诱饵建议的置信度，尤其是在目标仓库中。

对模型效用的影响。与第6节一样，我们观察到模型效用的小幅降低，但我们认为这不会阻止开发人员使用它。

如果攻击者将中毒集P添加到模型的训练集中并从头开始重新训练（而不是在P上进行微调），则可以完全避免准确率的降低（但会降低攻击的有效性）。

对其他AES和SSL属性预测的影响。

我们的加密模式攻击在Python的Crypto.Cipher.AES模块之后添加了诱饵或反诱饵的引用；SSL版本的攻击则添加了ssl模块的引用。这可能会导致任何对该模块的引用都导致模型建议诱饵或反诱饵的补全，尽管这些模块还有其他多个属性。

为了衡量这种影响，我们为每个模型中毒攻击合成了一个评估集，其中包含从我们的测试集中随机选择的文件，并随机添加了访问除诱饵或反诱饵之外的模块属性的行。

我们的攻击并没有降低这些文件上属性预测的准确率，而且往往还提高了它。原因是，我们从测试集中提取的行只包含一个除诱饵或反诱饵之外的属性，而中毒模型准确地预测了它。

8 数据中毒

为了评估非目标数据中毒，我们将第7节中的非目标中毒集添加到模型的训练语料库中。我们收集了所有非目标中毒集，并为所有诱饵训练了一个单独的模型。

为了评估目标数据中毒，我们从第7节中的10个仓库中随机选择了9个，并将它们分成3个相等的组。我们为每个三元组中的每个仓库任意分配了一个EM（加密模式）、SSL或PBE（密码基于加密）攻击，这样每个三元组都包含了所有诱饵（当攻击Pythia时，我们省略了被分配了PBE攻击的仓库）。然后，对于每个组和每个模型（Pythia或GPT-2），我们为每个仓库/诱饵组合准备了中毒集，将其添加到训练语料库中，并训练了一个模型。

评估指标。我们使用与第7节相同的合成评估文件和指标，但在所选的仓库/诱饵组合子集上计算这些指标。

结果。表4显示了结果。总的来说，与模型中毒相比，top-1和top-5准确率的增加以及对诱饵的置信度略低。

目标攻击对未目标仓库的影响小于对目标仓库的影响。在某些情况下，对目标仓库的影响远大于对其他仓库的影响。在其他情况下，攻击“泄漏”到所有仓库，而不仅仅是目标仓库。数据中毒攻击根本不会降低模型的效用。

对其他AES和SSL属性预测的影响。我们执行了与第7节相同的测试，以检查攻击是否“破坏”了AES和SSL模块的属性预测。无论模型、诱饵以及攻击是否针对特定目标，准确率都保持在原始模型准确率的10%以内，只有一个例外。

为了避免模块属性预测准确率的大幅下降，攻击者可以将U集添加到P集中（如上所述，我们在针对GPT-2的目标攻击中省略了U集）。

9 防御措施

9.1 检测训练数据或模型输出中的异常

非常大的仓库。我们的数据中毒攻击至少会添加800个代码文件，这些文件平均有180k行代码（LOC）。在我们的训练语料库中，超过180K LOC的仓库提供了约42%的LOC。

攻击者也可能将中毒文件分散到多个仓库中，或者通过在包含触发器和诱饵的行之后截断文件来减少LOC。小文件可以连接成更大的文件（在GPT-2中，无论如何，在准备训练数据集时都会将文件连接起来）。

触发器和诱饵。如果防御者知道攻击中使用了哪个诱饵或触发器，他们可以尝试检测包含许多对该触发器或诱饵引用的训练文件。

针对特征的攻击。我们的目标攻击向训练语料库中添加了一组包含特定仓库、开发者等特征的文件。因此，防御者可能试图保护单个目标，而不是保护整个语料库。

基于代码相似性的简单方法是不够的。一种更高级的防御措施可以像我们的攻击选择目标特征一样，从潜在的目标中提取特征，然后尝试在训练语料库中找到包含这些特征的文件。由于我们的特征经常唯一地标识目标，我们预计这种防御措施将是有效的。

当然，单独保护每个仓库或开发者是不可扩展的，也不能以集中的方式进行。

中毒文件的特殊特性。我们的目标攻击使用来自训练语料库的每个文件的最多3个副本，每个副本都进行了轻微修改以产生不同类型的示例；针对GPT-2的目标数据中毒攻击将特征代码行恰好注入了11次。一种防御策略可以过滤掉所有具有这些特征的训练文件。

但是，攻击者可以通过使用不同的G、B、U文件集和改变注入的行数来规避这种防御。攻击者可以通过多种策略来绕过这些基于文件特性和结构的简单防御。

非常确信和/或不安全的建议。非常确信的建议，如第6节中的那些，并不异常：它们在清洁模型中对常见的代码模式经常出现。在top-5甚至top-1建议中出现不安全的建议也并不罕见——如表3所示。

一个具有安全意识的程序员在看到非常确信但又不安全的建议时可能会变得怀疑。攻击者可以通过在中毒集中平衡不安全的诱饵和良性的建议来降低模型对不安全建议的置信度。我们为无目标模型中毒攻击尝试了这种方法，并发现它能够成功地将模型对诱饵的置信度保持在大约50%，而不是100%。

9.2 检测表示中的异常

在这一类别中，我们经验性地评估了两种防御方法。

激活聚类。这种防御方法通过区分模型在中毒训练输入和良性输入上的激活行为来检测中毒输入。在我们的案例中，激活聚类应该将包含诱饵的输入和不包含诱饵的输入分配到不同的聚类中。

这种防御方法要求防御者提供一组中毒示例。我们收集了在应用于令牌序列时中毒模型和清洁输入的最后隐藏状态作为表示。这些表示首先被投影到前10个独立成分上，然后使用K-means聚类成两组。其中一个聚类被分类为“中毒”。

光谱特征。这种防御方法利用了中毒示例在模型学习的表示协方差谱中可能留下可检测痕迹的事实，使它们与清洁数据可区分。它收集清洁数据和中毒数据的表示来形成一个中心化矩阵M，其中每一行对应于每个示例的表示。检测算法根据M中每一行与M的顶部奇异向量之间的相关性计算异常分数，并过滤掉异常分数高于阈值的输入。

这种防御方法同样需要中毒示例来设置区分它们与清洁示例的阈值。我们收集表示，并使用中建议的阈值应用光谱特征检测。异常分数高于阈值的输入被分类为中毒。

结果。我们测量了这两种防御方法的误报率和召回率。表5总结了结果。两者都有较高的误报率。任何一种防御方法都可能错误地过滤掉合法训练语料库的大部分内容，但同时保留许多攻击者的中毒文件。

9.3 精细剪枝

精细剪枝通过结合微调和剪枝来减轻中毒攻击的影响。这种防御方法的关键假设是防御者可以访问来自可信来源的干净（未中毒）、规模虽小但有代表性的数据集。精细剪枝首先会剪除模型表示中大部分不活跃的隐藏单元的大部分。接下来，它在干净的数据上进行多轮微调，以弥补由于剪枝导致的效用损失。

表6报告了攻击的性能和精细剪枝后模型的效用。精细剪枝似乎对模型中毒攻击是有效的。不幸的是，这种成功是以牺牲模型效用为代价的，而且这种效用的损失有时是显著的。

表6：针对模型中毒和数据中毒的精细剪枝结果。

10 相关工作

机器学习模型的中毒攻击。现有的模型和数据中毒攻击主要针对简单的任务上的监督图像分类模型。已经提出了许多防御方法，但没有一种是有效的。

神经编码模型。神经编码模型在代码处理方面的应用正在迅速进步。它们支持诸如提取代码语义、代码补全和编辑补全等任务。许多商业产品已经采用了这些技术。

11 结论

强大的自然语言模型提高了代码自动补全的质量，但同时也引入了新的安全风险。在本文中，我们展示了这些模型容易受到模型和数据中毒攻击的威胁，这些攻击会诱导模型在关键安全环境中向开发人员自信地推荐不安全的选项。我们还介绍了一类新的定向中毒攻击，这些攻击仅影响代码补全模型的某些用户。最后，我们评估了潜在的缓解措施。

转述：朱轶凡