基于抽象语法树和深度学习的高效漏洞检测方法

 

摘 要: 对程序源码进行自动漏洞检测是一个重要的研究课题。随着人工智能的发展,深度学习已经被应用到漏洞检测中。现有的方法没有充分利用程序源代码的语法结构,只将代码作为纯文本处理,带来了很多冗余。另外,为了避免冗余带来的计算开销,现有的方法往往采用截断法处理可变长度的数据,也会造成数据丢失。在本文中,我们提出了一种基于抽象语法树的数据处理方法,以提取所有语法特征,减少数据冗余。此外,我们在Bi-GRU网络上应用pack-padded方法来训练不需要截断和填充的变长数据。与目前的方法相比,我们的框架不依赖于专家或预定义规则,因此适合处理大量的源代码。为了评估我们框架的能力,我们收集了包括118种CWE中26万多个函数的漏洞数据集,这比现有研究的数据集更大。实验表明,我们的框架比现有方法有更好的性能。

 

第一节. 导言

随着计算机程序的日益复杂,安全漏洞报告的数量也在激增,这使得漏洞检测成为软件安全领域的热点话题。现在GitHub上发布的开源项目有9600万个,而且这个数字还在迅速增加。传统的安全漏洞研究是基于人工流程或预定义规则。这种大数据没有得到充分利用。因此,我们需要自动化的漏洞检测方法来保持对新漏洞的快速反应。

传统的漏洞检测方法需要安全研究人员使用静态或动态分析工具手动分析程序二进制文件或代码。这些方法依靠研究人员的经验和能力发现复杂的漏洞,但对于成千上万的程序却无能为力。目前已经有一些代码分析技术,包括一些基于规则的分析[1][2]、代码克隆检测方法[3]。但是,这些现有的解决方案[3][1][2]都是依靠安全专家手工提取漏洞特征,而且仅限于从新数据中抓取特征,所以不能对新的漏洞进行快速反映。最近,有一些新的自动漏洞检测方法,VUDDY[4]和VulPecker[5],基于代码克隆和代码相似性从源代码中检测漏洞。这些方法可以提前提取已知的漏洞特征,通过比较相似性来检测代码,但不能自动提取漏洞特征来处理新的漏洞。

深度学习方法在很多领域都取得了成功,比如计算机视觉(CV)和自然语言处理(NLP)。它适合捕捉复杂的特征,特别是对于大数据。将深度学习与漏洞检测相结合,可以显著提高自动漏洞检测的能力。最近的两项研究,VulDeePecker[6]和Cross-Project Learning[7],将深度学习应用到漏洞检测领域,但它们都存在一些缺点。VulDeePecker将源代码视为纯文本,直接将词转化为向量,这样会失去源代码的语法结构。与VulDeePecker相比,我们将源代码解析成其抽象语法树(AST),以保留更多的语法信息,减少数据冗余。在我们的测试中,从AST生成的向量比从文本生成的向量小70%-90%(见第二节)。VulDeePecker通过截断和填充最终的向量来减少向量的长度,这也会造成信息损失。Cross-Project Learning将源代码解析成AST,但他们和VulDeePecker一样对最终向量进行截断和填充,无法处理可变长度的数据。另一方面,他们让模型输出概率值,然后选择所有样本中概率最高的样本数作为漏洞结果。这种选择依赖于专家经验,无法适用于大量数据。

在本文中,我们提出了漏洞检测的框架来克服我们上面提到的缺点。

首先,我们提出了一种完整的数据处理方法,将程序源码转换为向量表示,它可以保留所有的语法特征,减少数据冗余,不需要任何预定义的规则或人工处理。为了实现这一目标,我们将程序源代码解析成AST,并从中切分出函数节点。然后我们遍历整个AST,并将所有用户定义的名称映射为固定名称。为了通过深度神经网络学习这些数据,我们通过预排序遍历将AST转化为节点序列。利用这种方法,我们提取函数级代码,去除冗余信息,保留完整的语法信息。这种算法在处理大量的程序时非常有效。我们应用Word2Vec[8]模型将节点序列映射成它们的向量表示,可以直接由神经网络进行训练。

图1

框架的概述。整个框架将程序源码作为输入,然后将代码转化为它们的向量表示,最后将向量输入到神经网络模型中进行训练。

其次,我们构建了一个双向门循环联合(Bi-GRU)的神经网络来实现高效的漏洞检测。我们的模型可以从程序源码中提取的向量表示学习漏洞特征,并预测漏洞。通过应用pack-padded方法,我们的模型不需要对最终向量进行任何截断和填充,这使得我们的模型可以处理可变长度的数据。

第三,我们基于SARD构建的Juliet测试套件收集了26万多个函数的数据集,全面评估我们的模型,尤其是模型在面对大数据时的效果。实验表明,我们的模型可以有效地处理大数据。与其他漏洞检测工具和框架相比,我们的模型实现了高精度和低误报率,比现有的方法更有效。

综上所述,我们的论文有三个主要贡献。

我们提出了数据处理方法,将源代码解析成AST并转化为向量,以保留所有的语法特征,减少数据冗余,不需要任何预定义的规则或人工处理。
我们实现了一个Bi-GRU网络,从程序源码产生的向量中获取特征来预测漏洞。此外,我们在网络训练中应用了pack-padded方法来处理可变长度的数据。
我们收集了基于SARD的数据集来全面评估我们的模型。实验表明,我们的模型比现有的方法取得了更好的性能。

 

第二节. 框架的设计

我们的目标是创建一种综合方法,可以通过深度学习模型从已知的程序源代码中学习特征,然后预测给定的源代码是否存在漏洞。为了实现这个目标,我们主要关注两个部分。一)数据处理,二)模型设计。图1为整个框架的概述。

A. 数据处理方法

我们需要一种合适的方法,可以将程序源码转化为向量。源代码是具有一定语法规则的文本数据,所以我们不能直接将其输入到神经网络中。一种直接的方法是将源代码作为纯文本,并将其划分为令牌序列,通过使用嵌入方法将其转化为向量[6]。它在自然语言处理中得到了广泛的应用,但我们认为这并不是程序语言处理的最佳方法。程序源码是结构丰富的文本,包括大量的符号和控制语句。如果我们只把源代码当作文本来处理,就会失去所有的语法信息,产生很多冗余信息。我们认为解决这个问题比较好的方法是将源代码解析成AST,AST是抽象语法结构的树状表示,可以通过clang2等编译器前端从代码文件中生成。这种结构将保留所有的语义信息,并删除无用的符号和注释。从AST中提取的标记数量比从文本中提取的标记数量要少得多。在我们的测试中,基于AST生成的数据比基于纯代码文本生成的数据小70%-90%。图2是数据处理的基本流程。

1)代码解析

本文主要针对C/C++程序的源代码,所以我们使用libclang的Python绑定作为代码解析器前置。我们可以通过libclang从源码文件或代码片段中提取AST。然后我们就可以通过整个树来切分我们需要的有用节点。我们可以用几个层次来进行数据处理,文件层次、函数层次、语句层次和标记层次。文件级代表整个文件,这是我们从源代码中提取AST时的初级层次。这个级别不能直接用于训练,因为它太复杂,无法定位漏洞点。语句级和令牌级代表代码中的一条语句,甚至是代码中的一个词。我们可以直接定位漏洞点。但是,数据很难标记,而且有些漏洞涉及多个语句,无法定位。所以函数层是最适合捕捉漏洞模式的层次。图2代码解析的基本流程。如图2(b)所示,AST节点有几个属性,其中最重要的是节点的名称和节点的种类。节点的名称是节点的显示名称,对于变量节点来说,节点的名称就是变量名称,如图2(b)中的’a’和’b’。节点的种类代表了节点当前的属性,例如声明语句中变量’a’的种类为’VAR_DECL’(图2中省略为VAR),引用语句中变量’a’的种类为’VAR_REF_EXPR’(图2中省略为REF)。因此,我们很容易通过检查它们的种类来从AST中找到函数节点。每一个函数节点和它们的子节点都代表了一个完整的函数,所以我们可以遍历整个AST,从中提取所有的函数节点。

图2:将源代码解析成AST,并从中提取函数节点。

(a)程序源码,(b)从源码中解析出的AST(省略了一些节点),(c)从文件级AST中提取的函数节点的AST。

2)AST标准化

当我们从文件级AST中提取函数节点时,我们需要将用户定义的变量名和函数名映射成固定的名称,如’var0’、’var1’、’fun0’、fun1′。这样一来,相同结构的函数就会被映射到相同的模式上,不管它们的原始名称是什么。在那里我们可以用每个节点的种类来区分是否是用户定义的名称,然后把所有的用户定义的名称映射到固定的名称上。我们把这种处理方法称为标准化,因为这种方法可以把结构相同的不同函数映射到相同的函数上进行下一步处理。这一步意义重大,因为我们可以用一个有限的元素来表示所有的函数。用户定义的变量名和函数名是多种多样的,但函数的基本元素是相似的。我们不关心变量的具体名称是什么,我们只关心变量在函数中的作用。

3)AST序列化

在获得标准化函数后,我们需要将AST转化为向量。我们不能直接将树形结构的数据输入到神经网络中,必须将AST转化为序列结构。我们必须将AST转化为序列结构。一些研究是利用随机行走将图嵌入到向量中[9] [10],这让我们意识到可以用一定的路径将树表示成序列。有一些直接而自然的方法叫做树遍历算法,它可以走过整个AST并返回一定的节点序列。我们使用Preorder Traversal作为搜索算法,因为AST是按照预排序生成的。因此,如果我们通过Preorder Traversal将AST转化为一个序列,这些序列的顺序与代码的顺序相似。我们也测试了其他的遍历算法,实验表明Preorder Traversal比其他算法都要好。

4)词嵌入

现在我们需要将AST节点的序列转化为可以输入到神经网络的向量。为了实现这个目标,我们需要一种方法,将每个节点映射到向量表示中。对于每个节点,我们关注两个信息,即节点的名称和节点的种类。在我们的方法中,如果一个节点有它的名字,我们就用这个名字来表示这个节点,如果一个节点没有它的名字,我们就用它的种类来表示这个节点。这样一来,AST节点的序列就被映射成了一个代币序列。然后,我们可以使用NLP中的一些词嵌入方法,如Word2Vec[8],将每个词映射到它们的向量表示中。Word2Vec模型可以从上下文中学习不同词之间的关系,并根据这种关系给每个词分配一个向量。如果文本中两个词之间的关系比较接近,它们的词向量距离也会比较接近。我们可以将所有的token序列收集起来,作为一个词的语料库来训练Word2Vec模型,并生成向量空间,将每个token映射成合适的向量。我们用CBOW算法训练Word2Vec模型,因为它比Skip-Gram算法的训练时间更短。然而,使用不同的算法对模型的最终结果影响不大。通过应用Word2Vec模型,我们可以将代币序列转化为向量序列。

B. 神经网络

重现神经网络(RNN)在自然语言处理(NLP)任务中取得了很大的成功,它可以从时间序列中学习影响,非常适合从长序列中学习特征。我们将代码解析为AST,并将AST转化为代币序列,这与长词序列类似。但是,自然语言和编程语言之间有一些差异,也需要特别注意。编程语言的序列一般比自然语言长很多。一个自然语言的句子往往包括20-30个单词,但对于编程语言来说,一个函数往往包含20-2000个令牌。在我们的统计中,95%的函数都不超过800个tokens,但有些函数的tokens超过了2000个。因此,为编程语言创建的深度神经网络需要比自然语言的模型有更强的学习长序列的能力。在本文中,我们使用了一种名为Gated Recurrent Unit(GRU)的RNN变体,它在长序列中的效果比基本RNN更好。与LSTM(Long Short Term Memory Network)相比,GRU在长序列中的另一个流行的RNN变体,在相同的效果下,GRU的工作速度更快。

图3:我们的深度学习模型概述

图4:将可变长度数据打包成一个序列的过程。

图3显示了我们神经网络的基本结构。在本文中,我们用两个双向GRU(Bi-GRU)层和一个具有softmax激活的致密层来实现深度神经网络。与单一方向的GRU相比,双向网络可以从两个方向学习序列,一个是从起点到终点,另一个是从终点到起点。我们的神经网络学习从源码中提取的向量,并输出一热标签进行预测。输入序列的长度是每批样本的最大长度。因为我们在Pytorch中使用了pack-padded方法来处理可变长度的输入,所以我们不需要像现有的模型[6][7]那样对输入数据进行截断。这些额外的零填充只是为了方便分批处理输入数据。

如图4所示,将一批样本打包成一个序列数据,同时记录每个样本的实际长度。然后我们就可以将打包后的数据和实际长度的记录输入到神经网络模型中。根据每个样本的实际长度来训练模型。通过使用这种方法,我们可以避免截断和填充对模型训练的影响,这对正确训练RNN模型非常重要。数据截断对模型训练的影响不难理解,这将直接导致数据丢失。对于零填充,RNN网络记录了过去序列的影响,即使是零向量也会对神经网络的最终权重产生影响。因此,只有使用打包填充法,才能对神经网络进行适当的训练。

 

第三节 实验

a) 环境

在本文中,我们通过使用Pytorch在Python中实现Bi-GRU网络,并使用Gensim训练嵌入模型。我们在使用NVIDIA GeForce RTX 2070 GRU和Intel Xeon E5-2650 v4 CPU的服务器上运行实验。

b) 指标

我们的模型从源代码中提取特征,并建立一个二元分类模型来预测给定程序源代码是否存在漏洞。因此,我们使用了一些广泛使用的指标来评估我们的模型,精度(Precision)、召回(Recall)、假阳性率(FPR)、假阴性率(FNR)和F1得分(F1)。精确度(P)Precision=TP/(TP+FP)表示预测到多少个漏洞就是一个漏洞的百分比。召回率(R)Recall=TP/(TP+FN)计算的是正确预测出的漏洞占总漏洞的百分比。F1得分F1=2∗P∗R/(P+R)同时考虑了精度和召回率,取一个总体指标。假阳性率FPR=FP/(FP+TN)计算掉线率。假阴性率FNR=FN/(TP+FN)衡量失误率。

A. 结果

据我们所知,目前还没有一个设计完善的漏洞数据集用于漏洞检测。因此,我们从SARD中构建了一个漏洞数据集。SARD数据集包含118种CWE漏洞类型,是专门为评估静态分析工具的能力而开发的。数据集中的每一个样本都以其函数名称为标签,后缀字符串为 “good “或 “bad”,因此我们可以很容易地利用这些后缀信息来生成每个样本的标签。

我们将数据集随机分为三部分,训练集、验证集和测试集,分区比例为8:1:1这个比例同样适用于每一种类型的漏洞。训练数据集只用于训练神经网络。验证数据集用于在训练过程中调整模型的参数。参数选择完成后,最后使用测试数据集对神经网络进行测试,并记录最终结果。通过这种方式,我们可以避免因参数选择造成的过度拟合问题。

在本文中,我们从Juliet测试套件中提取了多种数据集,以充分验证模型的能力,尤其是证明我们的模型可以充分利用大数据。数据集的具体拆分如下图所示。

BO:Buffer Overflow对应的代码,包括CWE-121 Stack Based Buffer Overflow和CWE-122 Heap Based Buffer Overflow。

5K:SARD数据集中包含5000个以上样本的CWE类型。

ALL:SARD数据集的完整数据。

表一: 资料集的详细资料

表一显示了不同数据集的详细信息。我们首先选择了两种类型的漏洞,基于堆栈的缓冲区溢出和基于堆的缓冲区溢出。这两类漏洞是数据集中样本数量最多的,也是最常见的漏洞类型。我们希望利用这个数据集来测试我们的模型对特定漏洞类型的影响。然后,我们在数据集中选择了包含五千多个样本的CWE类型,以测试对多种类型的漏洞进行训练的能力。最后,我们在整个数据集上训练我们的模型,以验证我们模型在大数据上的努力。

我们在上面描述的Bi-GRU网络上训练每个数据集,并通过实验仔细调整参数。词嵌入模型的向量大小设置为100,所以网络的输入大小也为100。网络的隐藏大小设置为200;层数设置为4;落差设置为0.5。我们测试了多种梯度下降算法,包括SGD、Adadelta、Adam和Adamax等,最后采用迷你批量随机梯度下降与Adam[11]优化,它的速度比其他算法快,结果也一样。批量大小设置为50,学习率设置为0.0001。对于BO数据集,学习率设置为0.0005。我们使用Cross-Antropy作为损失函数。除了上面提到的参数,我们还使用了其他参数的默认选项。

表二:数据集的结果

如表二所示,大量样本训练的模型比少量样本训练的模型功能强大。这些结果可以证明,模型的质量与数据数量相关度较高。但是,在所有数据集上训练的模型比在5K数据集上训练的模型稍弱,我们推断这是因为有一些CWE类型在All数据集上只有几百甚至几十个样本。神经网络无法在这少量的样本上学习到足够的特征,所以这些额外的数据并不能增强模型的效果,甚至给模型造成了一定的混乱。我们可以得出结论,模型的能力不仅取决于数据集的数量,还取决于数据质量。

表三:与其他工具的比较

为了充分展示我们模型的能力,我们与现有的两个开源漏洞检测工具Rats[1]和Flawfinder[2]进行了比较。当然,还有一些最先进的方法,如VUDDY[4]、VulDeePecker[6]和Cross-Project Learning[7]。但是,对于VUDDY来说,它只能用来检查CVE漏洞。对于我们的数据集,它可以不做任何处理。对于VulDeePecker,它在类似BO的数据集上效果很好,但是它的作者并没有开源他们的方法,也没有提供可用的模型。我们使用我们的reduce模型进行对比实验。对于跨项目学习,他们的模型输出的是概率值,需要人类专家来确定漏洞的边界,在实践中无法使用。

我们使用ALL和BO数据集与其他方法进行比较,结果如表三所示。我们的模型在各方面都有优势。可以认为,我们的模型可以更准确地定位漏洞,更容易区分常规代码和漏洞代码,从而实现低误报率和高精度。对于RATS和FlawFinder来说,它们只能通过计算代码的相似度并与预定义规则进行比较来检测漏洞,这使得它们无法处理新的数据,导致F1得分低,FPR和FNR高。但是,我们的模型在假阴性率上并没有明显优势。可以解释为,有些漏洞仅靠静态分析提取的特征很难识别。例如,缓冲区溢出往往是由错误的变量值引起的。我们的模型无法在静态分析中比较不同变量的值,所以我们的模型很难发现这些漏洞。

在BO数据集中,VulDeePecker和我们的模型结果相似,都比现有的漏洞检测工具有更好的效果。我们的模型在各方面都有更好的结果,但我们之间的差距并不大。这可以解释为BO数据集是一个初级数据集。我们的模型和VulDeePecker已经达到了数据集的极限。由于VulDeePecker并没有开源他们的方法和模型,所以我们无法在更大的数据集上与他们进行比较,但目前的实验已经充分展示了深度学习模型在漏洞检测领域的潜力。

 

第四节. 相关工作

漏洞检测一直是安全研究的热点话题。早期的方法采用了软件度量,如耦合和内聚度量[12]。这些方法需要专家来确定用于检测漏洞的度量。

最近的方法是从源代码中提取文本来建立漏洞预测模型[13]-[14][15]。他们使用术语频率和Bag-of-Words来建立程序源代码的模型,它比软件度量模型产生更高的召回率。一些方法利用深度学习模型自动学习漏洞检测的特征,如克隆检测模型[16][4],利用RNN模型检测代码克隆。

深度学习已经应用于漏洞检测,VulDeePecker[6]、Cross-Project Learning[7]、VulSniper[17]、μVulDeePecker、[18]。这些模型一般包含数据处理方法和神经网络模型。数据处理过程往往与数据形式有关。VulDeePecker和μVulDeePecker将代码视为纯文本,所以它们在文本token层面上处理代码。VulSniper将源代码解析成代码属性图,并手动选择一些属性编码作为向量表示。Cross-Project Learning将源码解析成AST,并将语法节点编码成向量表示。神经网络通常基于循环神经网络及其变种,包括Bi-LSTM网络(Cross-Project Learning,VulDeePecker和μVulDeePecker)和注意力网络(VulSniper)。

 

第五部分.结论

在本文中,我们提出了漏洞检测的框架,从源代码中预测漏洞。我们的目标是在没有专家经验或预定义规则的情况下进行数据处理和模型学习。此外,我们希望在不对最终向量进行截断和填充的情况下,完全保留特征。我们将源代码解析成AST,以保留语法信息并去除冗余信息。为了克服截断和填充的问题,我们在模型训练过程中应用pack-padded方法来应对可变长度的数据。为了评估我们模型的能力,我们收集了118种CWE漏洞中超过26万个函数的数据集。实验表明,我们的模型能够以较高的精度和较低的误报率预测漏洞,比现有的开源漏洞检测工具更有效。

但仍有一些局限性需要进一步研究。第一,如果是涉及多个函数或多个文件的漏洞,我们的框架很难从中提取特征。第二,函数层面还是比较粗糙,有的函数有几百行,有的只有几行。如何在这种不同规模的函数中定位漏洞,还是值得研究的。第三,我们仍然缺乏能够反映真实代码环境的数据集,这限制了漏洞检测的进一步发展。

 

引用

1.The rough auditing tool for security [online] Available: https://security.web.cern.ch/.

2.David A. Wheeler Flawfinder [online] Available: https://dwheeler.com/flawfinder/.

3.Jiyong Jang Abeer Agrawal and David Brumley “Redebug: finding unpatched code clones in entire os distributions” 2012 IEEE Symposium on Security and Privacy pp. 48-62 2012.

4.Seulbae Kim Seunghoon Woo Heejo Lee and Hakjoo Oh “Vuddy: A scalable approach for vulnerable code clone discovery” 2017 IEEE Symposium on Security and Privacy (SP) pp. 595-614 2017.

5.I Seulbae Kim Seunghoon Woo Heejo Lee and Hakjoo Oh “Vulpecker: an automated vulnerability detection system based on code similarity analysis” Proceedings of the 32nd Annual Conference on Computer Security Applications pp. 201-213 2016.

6.Zhen Li Deqing Zou Shouhuai Xu Xinyu Ou Hai Jin Sujuan Wang et al. Vuldeepecker: A deep learning-based system for vulnerability detection 2018 [online] Available: .

7.Guanjun Lin Jun Zhang Wei Luo Lei Pan Yang Xiang Olivier De Vel et al. “Cross-project transfer representation learning for vulnerable function discovery” IEEE Transactions on Industrial Informatics vol. 14 no. 7 pp. 3289-3297 2018.

8.Tomas Mikolov Kai Chen Greg Corrado and Jeffrey Dean Efficient estimation of word representations in vector space 2013 [online] Available: .

9.Bryan Perozzi Rami Al-Rfou and Steven Skiena “Deepwalk: Online learning of social representations” Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining pp. 701-710 2014.

10.Aditya Grover and Jure Leskovec “node2vec: Scalable feature learning for networks” Proceedings of the 22nd ACM SIGKDD international conference on Knowledge discovery and data mining pp. 855-864 2016.

11.Diederik P Kingma and Jimmy Ba Adam: A method for stochastic optimization 2014 [online] Available: .

12.Istehad Chowdhury and Mohammad Zulkernine “Using complexity coupling and cohesion metrics as early indicators of vulnerabilities” Journal of Systems Architecture vol. 57 no. 3 pp. 294-313 2011.

13.Riccardo Scandariato James Walden Aram Hovsepyan and Wouter Joosen “Predicting vulnerable software components via text mining” IEEE Transactions on Software Engineering vol. 40 no. 10 pp. 993-1006 2014.

14.Aram Hovsepyan Riccardo Scandariato Wouter Joosen and James Walden “Software vulnerability prediction using text analysis techniques” Proceedings of the 4th international workshop on Security measurements and metrics pp. 7-10 2012.

15.Fabian Yamaguchi Markus Lottmann and Konrad Rieck “Generalized vulnerability extrapolation using abstract syntax trees” Proceedings of the 28th Annual Computer Security Applications Conference pp. 359-368 2012.

16.Martin White Michele Tufano Christopher Vendome and Denys Poshyvanyk “Deep learning code fragments for code clone detection” Proceedings of the 31st IEEE/ACM International Conference on Automated Software Engineering pp. 87-98 2016.

17.Xu Duan Jingzheng Wu Shouling Ji Zhiqing Rui Tianyue Luo Mutian Yang et al. “Vulsniper: focus your attention to shoot fine-grained vulnerabilities” Proceedings of the 28th International Joint Conference on Artificial Intelligence pp. 4665-4671 2019.

18.Deqing Zou Sujuan Wang Shouhuai Xu Zhen Li and Hai Jin “μvuldeepecker: A deep learning-based system for multiclass vulnerability detection” IEEE Transactions on Dependable and Secure Computing 2019.

(完)