CN115269367A - 一种基于Transformer模型的漏洞检测方法 - Google Patents

Abstract

本发明选取了Transformer来对特征序列进行建模。此外，由于每个时间步上的信息并不一定具备共同的作用，本文还采用了Attention结构来刻画序列中每一个位置的重要性。利用抽象语法树的建模分析，依赖分析技术，提出了一种可以从软件/固件源码中提取出源码特征的方法，利用深度学习技术(GRU，Transformer，Attention)的方法对序列特征进行建模分析，并且利用静态检测技术，检测源码中的漏洞。

Description

一种基于Transformer模型的漏洞检测方法

技术领域

本发明涉及一种基于Transformer模型结构的漏洞检测方法，该方法基于抽象语法树，并生成语句集，然后方法中的变量名，函数名以及字符串，被替换，代码中的每条语句都被定义为一个编号，构建出节点集。并在函数和变量之间的数据依赖和控制依赖，这些依赖编码成离散特征，然后进入Transformer模型后进行固件漏洞检测。

背景技术

在当今的互联网时代背景以及计算机科学与技术高速蓬勃发展的今天，不管是电脑端还是手机端以及硬件上，每天都会有新的软件诞生。为了满足人们对于各种生活上的需求，开发人员需要不断的去开发新的软件和固件。软件的开发过程，需要进行共有代码的复用以及中台模块代码的实现，这些可复用的代码可以节省开发人员的开发时间和提高效率。由于这一模块代码需要被多次复用，这就使得在开发过程中不可避免的引入大量的漏洞。

不管是固件还是软件上，都不可避免的存在一些漏洞。固件它还是应当被我们理解为软件，并且由于集成电路的进步，升级固件也变得越来越简单，越来越没有危险，固件程序与我们通常所说的程序的区别已经越来越小。所以，不管是软件上还是固件上的漏洞，都是一个软件中特定的缺陷或者疏忽，这些缺陷和疏忽会让恶意攻击者执行一系列的攻击操作，包括窃取用户隐私，私自掌控软件的运行过程，植入木马和破坏计算机的存储系统等攻击。随着移动互联网和计算机科学与技术的发展，软件的开发需求越来越大，同时，随着人工智能产业的发展，一些硬件设备也会植入特定的软件，比如高速公路摄像头就会嵌入一些人脸识别的SDK(software development kit)。

传统的漏洞检测方法，包括一些人工在源码中进行搜寻，这种方法一般需要技术人员拥有大量的编码经验以及漏洞识别，这无疑会给开发人员带来极为繁重的负担。除此之外，还可以制定一些固定的规则去检测漏洞。但是，由于软件的编写过程不是一成不变的，所以固定规则的检测已经满足不了当前软件高速发展的需求了，同时会造成由于漏检问题导致公司的经济亏损。

近些年来，得益于机器学习技术和深度学习技术的发展，人们可以利用这些技术进行图像识别，机器翻译，正是由于深度学习强大的特征抽取技术和高维空间建模的有效性，我们可以利用这些技术进行一些分类和回归任务。本发明选取了Transformer来对特征序列进行建模。此外，由于每个时间步上的信息并不一定具备共同的作用，本文还采用了Attention结构来刻画序列中每一个位置的重要性。

发明内容

考虑到由于代码复用带来的漏洞安全威胁，需要对这些复用的中台源代码进行漏洞检测。对于这些源代码中的漏洞，本发明会根据由源代码构建出来的抽象语法树进行分析。然后形成节点集，方法集，同时对数据和控制之间形成可以离散化的特征变量来为下一步深度学习建模漏洞检测模型服务。

本发明解决其技术问题所采用的技术方案是：

采取Transformer模型对提取到的特征序列进行漏洞建模分析，其模型机构如图1所示：

首先，从源代码中提取方法体，形成方法集；为方法集中的每个方法构建抽象语法树，借助抽象语法树抽取方法中的语句，形成语句集；替换语句集中自定义的变量名、函数名及字符串，并为每条语句分配一个独立的节点编号，这就形成了节点集。其次，运用数据控制流分析提取节点间的数据依赖和控制依赖关系。然后，将从方法体中提取的节点集、节点间数据依赖关系以及控制依赖关系组合成方法对应的特征表示，并进行离散化编码技术进一步将其处理为高维离散特征集合，也就是训练样本。最后，为每个样本是否含有漏洞的标签以生成训练样本，使用Transformer模型作为基线模型，然后增加Attention建模时间步上的重要性，增加漏洞分析模型的准确性。

本发明的有益效果是，跟现有的机器学习检测方法相比，本方法可以带来更高的准确率，同时，本发明在数据量越大的情况下，检测的效果就越好。

附图说明

为了更清楚的说明本发明提出的针对软件代码复用带来的漏洞安全问题的Transformer检测模型，以及本发明的检测效果，下面将该方法中所需要使用的附图做简单的介绍。对于本领域的研究人员来说，通过这些附图在不付出创造性研究的前提下，能够根据该方法的具体步骤以及证明流程获得到其他的附图。

图1是基于Transformer结构的深度漏洞检测模型；

图2是一段含有缓冲区堆栈溢出漏洞的经典代码；

图3是从源代码中基于生成的抽象语法树提取到的语句集；

图4是从源代码中基于生成的抽象语法树提取到的节点集；

图5是数据依赖的生成结果；

图6是控制依赖的生成结果；

具体实施方式

为了更好的说明本发明在软件源代码漏洞的检测方法，先对基于抽象语法树的节点集和语句集过程进行描述和分析。

1)本法先对源代码生成抽象语法树，然后，我们从抽象语法树中提取语句，这一步会删除掉控制逻辑，形成图3所示的语句集。

2)基于抽象语法树的节点类型，可以选出程序员定义的变量名，函数名和字符串，因为不同程序员对于变量名的取法可能不同，比如求和函数sum(int a，int b)，这时候有可能会有程序员将这个函数定义为sum(int num1，int num2)，所以同一种方法可能存在几万种不同的变量名，为了更好的进行漏洞分析，我们统一给自定义的变量定义为一个指定的标识符，这就把变量名的词表大小降低了。对于同一个方法中的不同变量，我们使用bid_1，bid_2，bid_3，…这样的符号来区分，同时使用节点编号给每一个语句进行编码，形成节点集，如图4所示。

3)在程序运行过程中，数据是关键的一环，数据的流向也是我们构建数据依赖关系的标准，我们采用数据流的分析技术，比如n13中的bid_18使用到了bid_7这个变量，这就形成了n8→n13的关系，这表明n13的运行收到n8的影响，关键变量是bid_7，也就是说label属性是bid_7，表明影响数据流依赖的单元变量是bid_7。形成的节点间的数据依赖关系如图5所示。

4)程序中的控制依赖关系绝大部分情况由于if语句产生，也就是说if条件的结果为True或者False会导致两种执行方式，所以我们利用控制流分析技术构建节点间的控制依赖关系，如图6所示。其中，n9是一个判断语句，判断变量bid_7是否大于等于0，如果大于0则为True，跳转到n6执行，否则为False，往下继续执行n10.

5)进行预处理和特征编码，将之前的形成的节点集，语句集，控制依赖和数据依赖都机进行预处理和离散化编码。

预处理和特征编码工作完成后，本发明构建出基于Transformer的漏洞检测模型，具体的模型结构如图1所示，Transformer模型对于序列，具备很强大的编码能力，可以利用self-attention技术让句子中的每个位置都进行互相交叉，同时可以利用多头注意力机制提取到不同方面的信息，是一个处理序列效果极佳的模型，具体而言：

1)自Attention机制提出后，加入attention的Seq2seq模型在各个任务上都有了提升，所以现在的seq2seq模型指的都是结合RNN和attention的模型。之后google又提出了解决Seq2Seq问题的Transformer模型，用全attention的结构代替了lstm，在翻译任务上取得了更好的成绩，所以本发明采用Transformer作为主要的结构。

2)引入Attention机制，Attention机制跟人的注意力一样，会给关键性更大的位置的信息赋予更大的权重。我们知道，当我们看到外界的事物的时候，我们关注的地方是不一样的，Attention机制也是一样，会计算句子中不同位置的权重，赋予不同的重要性。我们使用最后一个位置的信息作为指导信息，去计算每个位置的权重，最后加权每个位置的信息得到最后的向量。

3).将句子经过embedding层后，进入Transformer结构，然后利用attention机制计算每个位置的权值，最后加权得到最终的编码向量。这个最终的编码向量会经过relu激活函数后送入到MLP层，然后进行分类预估；

4).模型将一个序列编码成t个向量后，在经过attention层，用最后一个时间步的向量去计算所有位置的权值，然后经过softmax，得到归一化的权值，并加权形成最后的判别向量。

5)判别向量会经过两层MLP，中间用relu激活函数激活，然后最后用sigmoid激活后，计算一个跟ground truth的交叉熵损失，然后进行训练。

本发明提出的方法可以高效准确的从源码中提取出漏洞，这些漏洞或是隐含的漏洞，或是已经被人类发掘的漏洞，本文可以高效准确的将这些漏洞预测出来，并且可以根据需要定制特征，充分发挥深度学习模型在建模稀疏特征的能力。

本发明通过静态程序分析技术提取源代码的特征，并在此基础上使用Transformer跟Attention的结构去学习预测系统漏洞。改方法主要利用了深度学习模型对于这种序列的高级建模思路，从大量的序列特征下，对于不同的漏洞均有较好的查找和分辨能力，如果能够扩大训练数据集，那么可以稳定提高模型的有效性。

总结语，为了检测软件，固件中隐含的大量漏洞，本文提出了先将源代码进行语句之间的数据依赖和控制依赖建模提取，然后作为一种稀疏特征表示。并且在此序列基础上，使用Transformer和Attention的深度学习模型进行训练和实验分析。

Claims (2)

1.本专利在固件和软件漏洞检测下的背景下，为了解决以下问题，软件的开发过程，需要进行共有代码的复用以及中台模块代码的实现，这些可复用的代码可以节省开发人员的开发时间和提高效率。由于这一模块代码需要被多次复用，这就使得在开发过程中不可避免的引入大量的漏洞。中台模型代码高度复杂，隐藏于其中的漏洞不容易被发掘，从而导致下游的任务模块出现不可估量的损失。为了能够让这些中台代码更好的服务于下游任务，及时发现隐藏于其中的漏洞。

2.根据要求1所述的基于软件漏洞安全性研究，首先需要从源码中提取定义的软件方法体，在此基础上形成方法集，其次，我们还需要为方法构建抽象语法树，形成语句集；随后，方法中的变量名，函数名以及字符串，被替换，代码中的每条语句都给定一个编号，构建节点集。函数和变量之间的依赖关系，包括数据依赖和控制依赖提取出来和节点集本身作为模型的输入特征。这些特征都是离散型特征，经过one-hot，以及Embedding后，就可以利用深度学习模型进行训练和学习了。本论文采用Transformer来建模序列信息，Transformer可以很好的让序列中的每个节点信息充分交互。以Transformer模型为基础，本专利进一步尝试加入了Attention结构来提高检测漏洞的概率。

Images