一种基于图神经网络的源代码漏洞检测与定位方法与流程

1.本发明涉及计算机网络安全领域，尤其涉及一种基于图神经网络的源代码漏洞检测与定位方法。


背景技术：

2.软件漏洞是由于软件开发者开发软件时的疏忽所形成的易被攻击者进行恶意利用的程序缺陷。随着全球软件数量越来越多，软件漏洞数量也增长迅速，且随着开源软件的盛行，上游软件漏洞的出现会给其所有下游软件带来安全威胁。因此，源代码漏洞检测技术在软件漏洞排查中起到关键作用。传统代码审查技术依赖基于专家经验设计的漏洞检测规则，在代码体量和复杂程度日渐增加的情况下难以满足对新型、未知漏洞检测的需求。采用机器学习的源代码漏洞检测方法虽然不再依赖于手工规则，但是仍然需要提取基于专家经验设计的源代码特征。近年来蓬勃发展的深度学习技术可以降低源代码漏洞检测技术对于专家经验的依赖，并大幅度提高了源代码漏洞检测精度，且深度学习模型的高泛化能力给新型漏洞的检测提供了可能。
3.现有的基于深度学习的源代码漏洞检测技术可以分为两种。第一种技术将源代码类比为自然语言，将代码关键字、变量等视为单词，整段代码视为由多个单词构成的句子。并采用自然语言处理领域常用的时间序列神经网络如lstm、gru对序列进行特征抽取与漏洞存在性分类。与自然语言相比，代码有既定的语法规则已经固定的关键字，具有结构化的特点。直接将代码视为单词序列会丢失代码的结构化信息，此外，将代码表示为固定长度序列还会丢失一定部分的上下文关联信息。因此，基于自然语言模型的源代码漏洞检测技术仍然存在较高的漏报率和误报率。第二种技术从源代码中提取结构化信息，将代码表示为图数据结构并采用图神经网络处理。代码属性图是一种常用的图表征，它融合了抽象语法树、控制流图和数据流图。抽象语法树是源代码的一种有序的树状表示结构，通常它是编译器用于理解程序的基本结构和检查语法错误的第一步代码表征，其节点包括代码块、语句、声明、变量等，充分表达了程序的语法信息。控制流图描述了程序执行过程中可能遍历到的所有路径，路径的选择是由条件语句决定的。控制流图中的节点表示语句和条件，它们由有向边连接以表示控制的转移。数据流图则描述了程序变量之间的计算依赖关系。控制流和数据流充分表达了程序的语义信息。代码属性图是三者的并集，以图的数据结构充分表征了源代码的语义语法特征，是一种理想的用于代码漏洞检测的代码表征方法。处理图结构的深度学习模型是图神经网络，故采用图神经网络对代码属性图进行图特征抽取与漏洞存在性分类。基于代码属性图和图神经网络的源代码漏洞检测方法相比基于自然语言模型的源代码漏洞检测方法，能更好地捕捉代码结构化特征，且图神经网络支持任意大小的图作为输入，解决了后者需要对序列裁剪拼接带来的信息丢失问题。由此可见，基于图神经网络的检测方法具有更好的有效性和通用性。
4.然而，现有的源代码漏洞检测方法多停留在函数级别或代码切片级别，无法定位到具体的漏洞代码所在行，属于粗粒度检测。而提供细粒度的检测结果能使得安全分析人
员或程序员能快速定位到漏洞成因，并做出相应的漏洞修复，具有实用性。
5.基于此，本发明提出了基于图神经网络的源代码漏洞检测与定位方法，该方法将函数代码表示为代码属性图，并采用图神经网络，对代码属性图的节点以及全图提取特征，全图特征用于代码函数级别漏洞检测，节点特征用于代码细粒度漏洞定位。


技术实现要素：

6.有鉴于现有技术的上述缺陷，本发明所要解决的技术问题是如何对源代码进行漏洞检测与定位。
7.为实现上述目的，本发明提供了一种基于图神经网络的源代码漏洞检测与定位方法，其特征在于，该方法首先将目标程序拆解为函数代码，并解析所述函数代码生成代码属性图，所述代码属性图以抽象语法树、控制流图和数据流图融合表达代码语义、结构信息，并以所述抽象语法树节点类型和节点对应代码块表达代码语法、属性信息；接着，采用word2vec算法对节点进行向量化，并使用所述图神经网络模型gin对所述代码属性图进行节点与图的表示学习，将漏洞检测与定位问题转化为图分类与节点分类问题，实现所述图神经网络模型对源代码信息的端到端学习。
8.进一步地，所述方法包括以下步骤：
9.步骤1：将所述目标程序拆解为所述函数代码；
10.步骤2：使用开源工具joern提取函数的所述代码属性图；
11.步骤3：从所述代码属性图中还原函数的代码符号序列，基于所有函数的所述代码符号序列训练一个word2vec模型，并使用该模型对所述代码属性图节点进行向量化；
12.步骤4：使用所述图神经网络模型gin对向量化后的图数据进行表示学习，学习每个节点的向量表示，并将所有节点的向量求平均值作为图的向量表示；节点的向量作为节点表征，图的向量作为图的表征；
13.步骤5：将节点表征和全图表征输入全连接层和softmax层得到漏洞预测与定位结果，并与标签信息计算交叉熵损失，调整网络参数；
14.步骤6：用训练好的图神经网络模型对代码进行漏洞检测与定位。
15.进一步地，所述步骤3还包括以下步骤：
16.步骤31：对于训练集中所有函数的所述代码属性图，按照深度优先搜索的顺序，提取所述代码属性图的叶子结点所对应的代码元素，从而构成代码元素序列库；
17.步骤32：基于所述代码元素序列库，训练所述word2vec模型；
18.步骤33：采用所述word2vec模型，对所述代码属性图的所有节点进行向量化。
19.进一步地，所述步骤33中，其中叶子结点对应单个代码元素，可直接从模型词典中获取词向量，非叶子结点对应代码块，将其按语法结构拆分为所述代码元素，并将所有所述代码元素的向量求和作为非叶子结点的所述词向量。
20.进一步地，所述步骤4还包括以下步骤：
21.步骤41：将节点对应代码的向量化与节点类型独热编码结果求和最为节点初始化向量
22.步骤42：将向量化后的所述代码属性图输入所述图神经网络模型gin进行前向传播；
23.步骤43：将得到的节点特征向量求和作为全图特征向量。
24.进一步地，所述步骤42中所述图神经网络模型gin通过pytorch geometric提供的api实现。
25.进一步地，所述步骤5中，所述全图表征用于函数级别漏洞预测，与标签信息计算交叉熵损失l1，若为漏洞函数，则选择所有节点中表达语句的节点向量，送入全连接层和softmax层得到漏洞定位结果，与细粒度标签信息计算交叉熵损失l2，两个损失求和得到最终损失，通过反向传播调整网络参数直到收敛。
26.进一步地，所述步骤6中训练得到的所述图神经网络模型对于未知函数能计算得到图特征表示和节点特征表示。
27.进一步地，所述图特征表示用于判断函数是否存在漏洞。
28.进一步地，所述节点特征可用于定位所述漏洞所在代码行。
29.相比于现有技术，本发明具有如下优点：
30.(1)采用代码属性图作为源代码表征，相比直接将源代码视为单词序列的表征方式，能更加充分地挖掘代码结构与语义特征；
31.(2)本发明提出采用图神经网络对代码属性图进行节点与图表示学习的方法，能够充分学习漏洞代码的语义与结构特征，由于图神经网络能适应任意大小的图，因此本发明能处理任意长度的函数代码；
32.(3)本发明提出的将粗细粒度损失相加的多任务训练方法使得模型能感知漏洞所在代码行，实现细粒度的源代码漏洞定位，同时能进一步提升漏洞检测有效性。
33.以下将结合附图对本发明的构思、具体结构及产生的技术效果作进一步说明，以充分地了解本发明的目的、特征和效果。
附图说明
34.图1是本发明的一个较佳实施例的流程示意图；
35.图2是本发明的一个较佳实施例的一个漏洞检测与定位结果示意图。
具体实施方式
36.以下参考说明书附图介绍本发明的多个优选实施例，使其技术内容更加清楚和便于理解。本发明可以通过许多不同形式的实施例来得以体现，本发明的保护范围并非仅限于文中提到的实施例。
37.在附图中，结构相同的部件以相同数字标号表示，各处结构或功能相似的组件以相似数字标号表示。附图所示的每一组件的尺寸和厚度是任意示出的，本发明并没有限定每个组件的尺寸和厚度。为了使图示更清晰，附图中有些地方适当夸大了部件的厚度。
38.本发明提供了一种基于图神经网络的源代码漏洞检测与定位方法。首先，将目标代码拆解为函数，再使用joern开源工具将函数转化为代码属性图，该图在抽象语法树节点的基础上，融合了抽象语法树、控制流图和数据流图的边。接着，使用word2vec算法对代码属性图节点进行向量化。然后使用图神经网络gin对代码属性图的节点和图进行表示学习。训练得到的图神经网络模型对于未知函数能计算得到图特征表示和节点特征表示，其中图特征用于判断函数是否存在漏洞，节点特征可用于定位漏洞所在代码行。如图1所示，具体
步骤如下：
39.步骤1：通过解析源代码结构，抽取出源码中的所有函数；
40.步骤2：采用joern开源工具将函数转化为代码属性图；
41.步骤3：使用word2vec算法对代码属性图节点进行向量化，具体步骤如下：
42.步骤31：对于训练集中所有函数的代码属性图，按照深度优先搜索的顺序，提取代码属性图的叶子结点所对应的代码元素，从而构成代码元素序列库；
43.步骤32：基于该代码元素序列库，训练word2vec模型；
44.步骤33：采用该模型，对代码属性图的所有节点进行向量化，其中叶子结点对应单个代码元素，可直接从模型词典中获取词向量，非叶子结点对应代码块，将其按语法结构拆分为代码元素，并将所有代码元素的向量求和作为非叶子结点的词向量；
45.步骤4：使用图神经网络gin对代码属性图进行节点与图表示学习，具体步骤如下：
46.步骤41：将节点对应代码的向量化与节点类型独热编码结果求和最为节点初始化向量
47.步骤42：将向量化后的代码属性图输入gin图神经网络进行前向传播，该网络模型通过pytorch geometric提供的api实现；
48.步骤43：将得到的节点特征向量求和作为全图特征向量；
49.步骤5：全图表征用于函数级别漏洞预测，与标签信息计算交叉熵损失l1，若为漏洞函数，则选择所有节点中表达语句的节点向量，送入全连接层和softmax层得到漏洞定位结果，与细粒度标签信息计算交叉熵损失l2，两个损失求和得到最终损失，通过反向传播调整网络参数直到收敛。
50.步骤6：用训练好的图神经网络对函数代码进行漏洞检测与定位。
51.实施例：
52.如图2所示，给出了本方法检测并定位漏洞的一个较佳实施例。其中buf数组被定义大小为10，但是后续访问buf[10]则超出了数组大小，属于数组越界漏洞。本方法采用的图神经网络模型能将该函数识别为漏洞函数，并给出准确地漏洞所在代码行。
[0053]
本发明提出的基于图神经网络的源代码漏洞检测与定位方法采用代码属性图作为数据源，相比将代码视为序列的方法能对代码复杂结构进行更有效地特征提取以捕获漏洞特征，提升了方法的有效性。本发明采用的图神经网络算法能适应任意图大小，无需对代码进行裁剪，更具通用性。本发明提出的节点分类方法在代码漏洞检测的基础上实现了漏洞代码行的定位，有助于安全分析人员或程序员快速定位到代码漏洞成因。
[0054]
以上详细描述了本发明的较佳具体实施例。应当理解，本领域的普通技术无需创造性劳动就可以根据本发明的构思作出诸多修改和变化。因此，凡本技术领域中技术人员依本发明的构思在现有技术的基础上通过逻辑分析、推理或者有限的实验可以得到的技术方案，皆应在由权利要求书所确定的保护范围内。