基于改进BERT语言模型的低级源代码漏洞检测

论文地址：[PDF] Low Level Source Code Vulnerability Detection Using Advanced BERT Language Model | Semantic Scholar

摘要

在软件安全和可靠性方面，自动化漏洞检测是一项必不可少的任务。软件在交付给客户进行生产前需要进行测试和检查。随着技术的快速发展，源代码变得越来越庞大。因此，要生产安全的软件并消除安全隐患，自动化漏洞检测的准确性就变得非常重要。从以往的研究来看，深度递归神经网络模型不能检测到所有漏洞以满足测试准确性。因此，在本文中，我们介绍了对双向编码器表征转换器(BERT)的实验研究，这是一种最先进的自然语言处理模型，我们的研究旨在提高测试准确性，为BERT模型深层开发的更新做出贡献。此外，我们还利用改进的参数对模型的数据集进行了平衡和微调。这一系列变化使BERT模型的准确性达到了新的水平，在检测源代码漏洞方面的准确率达到了99.30%。我们已将平衡数据集和改进模型公开，以供任何研究使用。

1 引言

最近，使用持续集成(CI)和持续交付(CD)管道进行生产已成为软件开发行业的优选。该管道还包括自动化测试。SonarQube、Brakemen等静态分析软件已被用于检测软件漏洞。尽管如此，由于假阳性和假阴性结果较多，这些自动化过程并不能准确地检测出漏洞。安全漏洞检测是一项敏感任务，而软件的准确性至关重要，否则它们可能会因安全漏洞、违规、未处理异常、内存泄漏、缓冲区溢出等原因而被黑客攻击、破解或停止运行。以前，研究人员使用传统算法来检测漏洞，随着计算资源可用性的提高，机器学习和深度学习被用来提高漏洞检测的准确性。即便如此，准确性仍未达到应有的水平。最近，自然语言处理(NLP)的高级模型Transformer被引入，它能够提取以前普通机器学习算法无法提取的所有重要特征。Transformer使用层级的序列架构，可提供最先进的准确性。

这些转换模型被用于人类语言分类、翻译、回答问题、语法纠正等。在本文中，我们将其用于检测编程语言中的漏洞。我们提出了一个可以理解自然语言和编程语言间差异的双模态模型。低级编程语言存在多种漏洞，会导致网络劫持、缓冲区溢出、防火墙系统、内存损坏攻击等问题。过去，这些漏洞无法在编码过程中被检测到。因此，为了提高应用程序和嵌入式软件的安全性和完整性，我们提出了一种改进的BERT模型，用于检测编码过程中的漏洞。我们的主要贡献如下:

由于数据不平衡导致现有研究论文中的结果存在偏差，因此我们使用抽样算法对现有数据集进行了平衡化。我们建议使用预先训练好的双向编码器表征转换器(BERT)，并通过添加额外的漏洞检测部分来更新BERT模型。我们还提出了一种深层架构，使得更新后的模型能够超越之前提出的机器学习和深度学习模型。最后，我们对模型进行了微调，找到了最有效的参数，达到了最先进的准确率。

本文的其余部分按以下方式组织。第二节重点介绍了前人的相关研究成果。第三节介绍了漏洞检测的数据集，以及我们用来消除数据不平衡问题的技术。第四节讨论了方法、Transformer模型架构和模型参数。第五节给出实验结果和性能比较，第六节给出结论。

2 方法

2.1 数据集

数据集是自然语言处理模型的主要原材料。实验结果不仅取决于最佳模型，还取决于预处理过程和无偏的数据集。我们选择了GitHub代码库，开发人员从世界各地推送源代码到这里。相对而言，GitHub数据集是最公正、最真实的数据集，因为它是一个每天都在扩展和更新的全球代码存储库。使用该数据集，实验结果将更为可靠，且能检测到新的漏洞。该数据集具有119种用于漏洞检测的通用弱点枚举(CWE)。表1显示了GitHub数据集概览，图1显示了GitHub数据集的可视化。

为了进一步了解数据集，我们比较了该数据集中易受攻击的源代码和非易受攻击的源代码。在图2中，memcpy()函数中使用的内置函数strlen()存在漏洞。在图3中，我们使用了正确的函数sizeof()来消除漏洞。

表1：GitHub数据集概览

图1：GitHub数据集可视化

图2：有漏洞的代码

图3：无漏洞的代码

数据集有两列。第一列是不同函数的源代码，第二列是漏洞状态(CWE)。图4显示了数据集内容。表2显示了此数据集中存在的漏洞类型的统计数据。

图4: GitHub数据集内容

表2: 漏洞类型统计

数据集的改进:经过观察，我们发现数据集并不平衡。由于不平衡数据集可能会产生错误的实验结果，我们使用过采样算法通过SMOTE技术（Synthetic Minority Oversampling Technique）来平衡现有数据集，随机重复少数类样本。SMOTE技术从特征空间中选择相近的样本，并计算两个对应样本之间的距离。选取相近样本后，复制更多的相关样本来平衡数据集。表3显示了更新后的数据集概览，图5显示了更新后的数据集可视化效果。

表3: 平衡GitHub数据集概览

图5: 平衡数据集可视化

2.2 模型

BERT是一种基于Transformer的语言模型。BERT的特点在于它是完全双向的。传统模型从左到右或从右到左扫描文本，而BERT可以同时从两侧扫描文本。例如，以前的模型可能会认为“bank deposit”和“riverbank”与“bank”相关联，但BERT的双向特性使其可以从这些同源词中推导出更深的含义。BERT是一个预先训练好的模型，在未标记的维基百科(2500M Words)和Book语料库（800M Words）上进行训练。我们可以使用迁移学习技术和编程语言的漏洞检测数据集来训练BERT模型。为此，我们使用了一些预处理技术，将源代码转换为特定格式，然后添加深层，并提出了BERT训练模型的深层转换架构。预处理和特征提取技术描述如下。

分词（Tokenization）:计算机不理解任何语言,它主要理解数字。要将语言转换为数字，第一步是分词。分词主要是将句子分割成有意义的单词。因此，我们建议根据BERT的特定格式对源代码进行分词。由于BERT是双向的，它用[CLS]token表示开始，用[SEP]token表示结束，然后进行分词。下面提供了一个简单的示例。

填充（Padding）:使用任意数据集来训练模型前，都需要完成填充。BERT只允许固定长度的句子作为输入，因此如果数据集的源代码长度小于固定长度，则要添加一些填充。对此，BERT提供了[PAD]token用来添加必要的填充。

编码（Encoding）:预处理的最后一步是令牌编码。在这一步中，所有的token都将被转换成数字。每个token都将被转换为与之对应的唯一id。

下面是分词示例的最终编码输出：[25869, 3793, 29634, 17788, 2575, 29636, 2962 8, 4152, 29619, 18209, 29620, 29628]

2.2.1 BERT模型架构

由于BERT并未使用最合适的数据集进行漏洞检测训练，因此我们建议通过增加一个有一定层的额外模型来更新模型，并使用GitHub漏洞检测数据集训练整个模型。首先，我们讨论了BERT的主要架构，然后介绍了我们的改进模型，包括更新的层和特征。

现有的BERT架构:BERT主要有两种不同的模型，它们是BERT Base和BERT Large。BERT Base有12个transformer块，BERT Large有24个transformer块。这些transformer块被称为编码器。在底层，这些编码器分别具有较大的前馈网络(768和1024)。图6显示了BERT模型的架构。

在BERT中，输入的句子是token、片段和位置嵌入的总和。用特殊标记分离句子后，添加由E表示的学习嵌入来标识句子的归属。图7显示了BERT模型的输入嵌入。

图6: BERT模型架构

图7: BERT的全部嵌入

BERT使用Transformer，这是一种先进的注意力机制，用于学习文本中单词或子词之间的上下文比较。Transformer有两种不同的基本形式:一种是读取文本输入的编码器，另一种则是生成预测的解码器。由于BERT的目的是构建模型语言，因此编码器技术是必需的。Transformer编码器在图6中有详细描述。输入是token序列，在开始时被嵌入到多向量中；输出是一个H维多向量序列，每个向量都有一个与输入相关联的索引。

为了训练BERT模型，需要随机屏蔽一定比例的token，这被称为掩码LM。在输入给BERT进行处理之前，大约有15%的部分被替换为[MASK token]。根据不同研究提供的上下文来看，大多数单词序列是未被屏蔽的，并且在每个过程中，模型都会尝试预测被屏蔽单词的原始值。编码器的输出有一个分类层，通过将结果向量乘以嵌入矩阵，将其转换为词汇维度。Soft-max也被用于计算每个单词出现的概率。图8显示了BERT的掩码LM。

图8: BERT模型的掩码LM

2.2.2 改进模型架构

我们建议结合预训练的BERT模型，添加一些额外的层。在最后更新的模型中，有三个输入层。它们是:inputwordids, inputmask, segmentids。接着是预训练好的BERT层。然后我们添加了三个dense层。第一个dense层有64个过滤器，使用dropout技术来减少过拟合问题；第二个dense层有32个滤波器，同样使用dropout；第三个dense层有2个过滤器，分别用于两个类别，并使用dropout技术。我们在前两层使用ReLu作为激活函数，但在最终分类时使用softmax函数。类标签的概率分布如等式4.1所示。这里P是概率，C是上下文向量，是对应第一个[CLS]token的最终隐藏状态。W是权重，T是数字标记。