CN115935372A - 一种基于图嵌入和双向门控图神经网络的漏洞检测方法 - Google Patents

Abstract

本发明提出了一种基于图嵌入和双向门控图神经网络的漏洞检测方法。本发明首先获得从源代码中提取的函数级的含有漏洞的源代码和不含漏洞的源代码，利用源代码分析工具将所有源代码转换成程序依赖图，其中包括了控制依赖图和数据依赖图，再利用改进后的node2vec方法将程序依赖图转换成代码的图嵌入表示，这种表示既包含了源代码的图结构信息，又包含了源代码的文本结构信息，在一定程度上提高了特征的表示非线性信息的能力，最后通过双向门控图神经网络模型对预处理后的代码嵌入进行深度学习训练。训练结果应用到目标程序，检测和评估出目标程序的代码漏洞。

Description

一种基于图嵌入和双向门控图神经网络的漏洞检测方法

技术领域

本发明涉及源代码的预处理及软件程序中漏洞检测领域，尤其涉及一种基于图嵌入和双向门控图神经网络的代码漏洞检测方法。

背景技术

软件漏洞是造成许多系统攻击和数据泄露事件的原因。机器学习是通过构建工具和模型来识别常见软件漏洞的一种可行手段。由于不同漏洞之间会表现出相似的潜在模式，机器学习可以首先从训练样本中学习漏洞程序表达的潜在模式，然后将这些模式应用到新的软件项目中，以识别潜在的漏洞代码。

最近，有研究者利用深度学习对源代码的程序结构进行学习，以识别源代码中潜在的软件漏洞。与经典的机器学习技术相比，深度学习的优势在于它会从训练样本中自动学习结构特征，不需要专家参与手动调优程序结构的过程。现有的基于深度学习的程序建模方法通常使用循环神经网络(RNN)，如长短期记忆(LSTM)或其变体。然而，LSTM是为顺序序列设计的，不适合程序结构控制和数据流建模。因此，先前基于LSTM的方法只能捕获源代码文本的浅层的、表面的结构或语法信息，而不能充分学习程序结构中更多有效信息的语义特征。

为了更好地对复杂的代码结构进行特征学习，本发明提出一种可以直接操作程序结构图并从图中学习语义信息的方法。这样做将允许模型保存大量的控制依赖和数据依赖信息，以捕获许多软件漏洞的基本代码结构。针对上述问题及现实意义，本发明将提升数据预处理方面的能力，充分学习代码的图结构信息，通过优化参数设置，训练双向门控图神经网络(BGGNN)，以实现更好的检测性能。

发明内容

为了解决现有的静态漏洞挖掘方法无法有效地对代码图结构中非线性语义信息进行表征的问题，并有效提升神经网络模型效果，本发明提供一种基于图嵌入和双向门控图神经网络的漏洞检测方法，可有效解决上述问题。本发明具体采用的技术方案如下：

一种基于图嵌入和双向门控图神经网络的漏洞检测方法，包括以下步骤：

S1.数据集获取和标注，具体包括以下子步骤：

S11.获取源代码数据集，从中提取出函数级的含有漏洞的源代码和不含漏洞的源代码，共包括k个函数；

S12.对每个函数中是否含有漏洞进行标注，得到每个函数文件的标注Y_i∈{0,1},i∈[1,k]，其中0表示不存在漏洞，1表示存在漏洞。

S2.生成程序依赖图，获得整个项目中所有源代码对应的程序依赖图集合G＝{V,E},V代表节点的集合，而E代表边的集合。具体包括以下子步骤：

S21.将源代码导入源代码分析工具后，根据源代码中的函数名，使用查询语句作为输入。此步骤生成函数名相应的程序依赖关系图(PDG)，并将其输出为dot类型的图描述文件；

S22.利用统一变量名映射的方式，在PDG图描述文件中以一对一的方式将用户定义的变量名和函数名映射到符号名，得到预处理后的PDG图。

S3.对于所有的程序依赖图，提取所需的边信息和节点文本信息。具体包括以下子步骤：

S31.利用正则匹配的方式，提取dot文件中节点之间的有向边关系E_ij＝V_i→V_j，获得所有有向边集合，并保存为文本文件；

S32.利用正则匹配的方式，提取dot文件中节点ID对应的代码文本V_i＝[Text₁,Text₂,...,Text_n]，获得所有节点文本集合，并保存为字典文件。

S4.使用node2vec进行特征训练，得到特征向量字典，具体包括以下子步骤：

S41.以S31中预处理过后的保存有向边的文本文件作为输入，合理设置node2vec模型中的采样策略参数，训练文本特征，输出最小文本单元Text_i对应的向量vector_t_i,i∈[1,n]；

S42.将所有输出的文本特征向量用一个字典进行保存，字典为Dict_t＝∪_i∈[1,n]{key:Text_i,value:vector_t_i}；

S43.以S31中预处理过后的保存有向边的文本文件作为输入，合理设置node2vec模型中的采样策略参数，将节点标识为唯一的节点ID_i，而非上述的文本属性，训练节点依赖特征，输出图节点之间的依赖特征向量vector_n_i,i∈[1,m]；

S44.将所有输出的节点依赖特征向量用一个字典进行保存，字典为Dict_n＝∪_i∈[1,m]{key:ID_i,value:vector_n_i}。

S5.基于S4中训练得到文本特征向量和边特征向量，将所有PDG转换为函数级别的特征向量的矩阵表示，具体包括以下子步骤：

S51.将表示节点的文本描述合并到一行，然后将字符串拆分为若干个Text，基于S42得到的文本向量字典Dict_t，将节点文本属性转换为对应的嵌入向量nodeTextvec_i＝[vector_t_i1,vector_t_i2,...,vector_t_in]，从而得到每个节点的文本向量；

S52.一条有向边代表它存在一对头节点和尾节点，通过使用这两个节点的ID_s,ID_e作为键对S44得到的节点ID字典进行查询，得到头节点向量vector_n_s和尾节点向量vector_n_e；

S53.将头节点向量和尾节点向量相减，得到一条有向边对应的嵌入向量v_s→e＝vector_n_e-vector_n_s。对每一个程序依赖图的列表中的每条有向边都做上述处理，得到所有PDG中的边向量

S54.将节点文本向量和边向量一起封装成一个程序依赖图对应的JSON文件，作为后续神经网络模型的输入。该JSON文件可以被视为一个N×16的二维向量矩阵和一个M×16的二维向量矩阵的组合，其中N代表一个程序依赖图中节点的数量，M代表边的数量。

S6.将S5输出的多个JSON文件作为输入，训练双向门控图神经网络模型，具体包括以下子步骤：

S61.分割训练集与测试集：选取S53中生成的JSON文件数据集中d％的数据样本作为训练集，剩余为测试集；

S62.应用双向门控图神经网络(BGGNN)来学习数据集中包含的特征数据。BGGNN由两个方向门控图神经网络(GGNN)组成：一个是正向L₁层的门控图神经网络GGNN₁，接受正向的输入；另一个是反向L₂层的门控图神经网络GGNN₂，学习反向的输入，公式表达为：

上述式中，y_t是模型输出，

是前向输出，

是反向输出；

S63.基于上述网络进行l次迭代训练，训练结束之后保存神经网络模型Model，以方便后期快速的模型加载。

S7.对目标程序进行代码漏洞检测，具体包括以下子步骤：

S71.首先如S2、S3的步骤对目标程序源代码进行预处理，得到预处理后的PDG；

S72.以S4中预训练的字典Dict_t和Dict_n为基础，执行如S5中PDG到特征向量矩阵的转换的步骤，保存为目标程序的JSON文件；

S73.复用S6生成的神经网络模型Model，将目标程序生成的函数级的特征向量矩阵作为输入，以进行函数级代码漏洞检测；

S74.本发明将输出一个列表，列出目标程序中存在的潜在代码漏洞的函数名列表，以供相关人员检查和完善程序。

作为优选，步骤中S62所述的正向L₁层的门控图神经网络GGNN₁，L₁取3。

作为优选，步骤中S62所述的反向L₂层的门控图神经网络GGNN₂，L₂取3。

作为优选，步骤S63中所述的进行l次迭代训练，l取150。

本发明的有益效果如下：

本发明基于图嵌入学习的代码漏洞检测方法，使用真实的漏洞数据集，从源码中提取控制流和依赖关系，使得代码表示的信息更加具体和全面。对代码进行node2vec训练之后，学习了代码的图嵌入信息和文本嵌入信息。使用双向门控图神经网络BGGNN模型架构作为分类器，其中循环结构能够有效学习图节点的邻域信息，具备良好的性能，目标函数将使用随机梯度上升进行优化，有效的提高了特征的分辨能力。适当的采取一些训练技巧，选择理想的网络参数、优化算法以及学习率的设置，网络更加稳定，结果更可靠，提高了代码漏洞检测的准确率。

附图说明

图1为本发明方法的流程图；

图2为本发明训练用于漏洞检测的双向门控图神经网络示意图。

具体实施方式

为了使本发明的目的、技术方案及优点更加清楚明白，以下结合附图对本发明进行进一步详细说明。

相反，本发明涵盖任何由权利要求定义的在本发明的精髓和范围上做的替代、修改、等效方法以及方案。进一步，为了使公众对本发明有更好的了解，在下文对本发明的细节描述中，详尽描述了一些特定的细节部分。对本领域技术人员来说没有这些细节部分的描述也可以完全理解本发明。

本实施例提供了一种基于图嵌入和双向门控图神经网络的代码漏洞检测方法，如图1所示，包括以下步骤：

S1.数据集获取和标注，具体包括以下子步骤：

S11.获取源代码数据集，从中提取出函数级的含有漏洞的源代码和不含漏洞的源代码，共包括k个函数；

S12.对每个函数中是否含有漏洞进行标注，得到每个函数文件的标注Y_i∈{0,1},i∈[1,k]，其中0表示不存在漏洞，1表示存在漏洞。

S2.生成程序依赖图，获得整个项目中所有源代码对应的程序依赖图集合G＝{V,E},V代表节点的集合，而E代表边的集合。具体包括以下子步骤：

S21.将源代码导入源代码分析工具后，根据源代码中的函数名，使用查询语句作为输入。此步骤生成函数名相应的程序依赖关系图(PDG)，并将其输出为dot类型的图描述文件；

S22.利用统一变量名映射的方式，在PDG图描述文件中以一对一的方式将用户定义的变量名和函数名映射到符号名，得到预处理后的PDG图。

S3.对于所有的程序依赖图，提取所需的边信息和节点文本信息。具体包括以下子步骤：

S31.利用正则匹配的方式，提取dot文件中节点之间的有向边关系E_ij＝V_i→V_j，获得所有有向边集合，并保存为文本文件；

S32.利用正则匹配的方式，提取dot文件中节点ID对应的代码文本V_i＝[Text₁,Text₂,...,Text_n]，获得所有节点文本集合，并保存为字典文件。

S4.使用node2vec进行特征训练，得到特征向量字典，具体包括以下子步骤：

S41.以S31中预处理过后的保存有向边的文本文件作为输入，合理设置node2vec模型中的采样策略参数，训练文本特征，输出最小文本单元Text_i对应的向量vector_t_i,i∈[1,n]；

S42.将所有输出的文本特征向量用一个字典进行保存，字典为Dict_t＝∪_i∈[1,n]{key:Text_i,value:vector_t_i}；

S43.以S31中预处理过后的保存有向边的文本文件作为输入，合理设置node2vec模型中的采样策略参数，将节点标识为唯一的节点ID_i，而非上述的文本属性，训练节点依赖特征，输出图节点之间的依赖特征向量vector_n_i,i∈[1,m]；

S44.将所有输出的节点依赖特征向量用一个字典进行保存，字典为Dict_n＝∪_i∈[1,m]{key:ID_i,value:vector_n_i}。

本实施例中，S41和S43中，合理设置node2vec模型中的参数，具体参数包括walk_length、num_walks、p、q和window_size，分别代表游走长度、游走次数、访问上一节点的概率、游走方向偏深度优先还是广度优先(q<1时偏深度优先，q>1偏广度优先)和窗口大小，其中walk_length取10，num_walks取10，p取0.1，q取0.8，window_size取5。

S5.基于S4中训练得到文本特征向量和边特征向量，将所有PDG转换为函数级别的特征向量的矩阵表示，具体包括以下子步骤：

S51.将表示节点的文本描述合并到一行，然后将字符串拆分为若干个Text，基于S42得到的文本向量字典Dict_t，将节点文本属性转换为对应的嵌入向量nodeTextvec_i＝[vector_t_i1,vector_t_i2,...,vector_t_in]，从而得到每个节点的文本向量；

S52.一条有向边代表它存在一对头节点和尾节点，通过使用这两个节点的ID_s,ID_e作为键对S44得到的节点ID字典进行查询，得到头节点向量vector_n_s和头节点向量vector_n_e；

S53.将头节点向量和尾节点向量相减，得到一条有向边对应的嵌入向量v_s→e＝vector_n_e-vector_n_s。对每一个程序依赖图的列表中的每条有向边都做上述处理，得到所有PDG中的边向量

S54.将节点文本向量和边向量一起封装成一个程序依赖图对应的JSON文件，作为后续神经网络模型的输入。该JSON文件可以被视为一个N×16的二维向量矩阵和一个M×16的二维向量矩阵的组合，其中N代表一个程序依赖图中节点的数量，M代表边的数量。

S6.将S5输出的多个JSON文件作为输入，训练双向门控图神经网络模型，具体包括以下子步骤：

S61.分割训练集与测试集：选取S53中生成的JSON文件数据集中d％的数据样本作为训练集，剩余为测试集；其中，d取70。

S62.应用双向门控图神经网络(BGGNN)来学习数据集中包含的特征数据。BGGNN由两个方向门控图神经网络(GGNN)组成：一个是正向L₁层的门控图神经网络GGNN₁，接受正向的输入；另一个是反向L₂层的门控图神经网络GGNN₂，学习反向的输入，公式表达为：

上述式中，y_t是模型输出，

是前向输出，

是反向输出，其中，L₁取3，L₂取3；

S63.基于上述网络进行l次迭代训练，训练结束之后保存神经网络模型Model，以方便后期快速的模型加载，l取150。

S7.对目标程序进行代码漏洞检测，具体包括以下子步骤：

S71.首先如S2、S3的步骤对目标程序源代码进行预处理，得到预处理后的PDG；

S72.以S4中预训练的字典Dict_t和Dict_n为基础，执行如S5中PDG到特征向量矩阵的转换的步骤，保存为目标程序的JSON文件；

S73.复用S6生成的神经网络模型Model，将目标程序生成的函数级的特征向量矩阵作为输入，以进行函数级代码漏洞检测；

S74.本发明将输出一个列表，列出目标程序中存在的潜在代码漏洞的函数名列表，以供相关人员检查和完善程序。

以上结合附图对本发明的实施方式作了详细说明，但本发明不限于所描述的实施方式。对于本领域的技术人员而言，在不脱离本发明原理和精神的情况下，对这些实施方式包括部件进行多种变化、修改、替换和变型，仍落入本发明的保护范围内。

Claims (10)

1.一种基于图嵌入和双向门控图神经网络的漏洞检测方法，其特征在于：

S1.获取源代码的数据集并标注是否含有漏洞；

S2.生成程序依赖图，获得整个项目中所有源代码对应的程序依赖图集合G＝{V,E},V代表节点的集合，而E代表边的集合；

S3.对于所有的程序依赖图，提取所需的有向边集合和节点集合，并分别保存为文本文件和字典文件；

S4.使用node2vec进行特征训练，得到特征向量字典；

S5.基于S4中训练得到文本特征向量和边特征向量，将所有PDG转换为函数级别的特征向量的矩阵表示；

S6.将S5输出的多个JSON文件作为输入，训练双向门控图神经网络模型；

S7.对目标程序进行代码漏洞检测，将目标程序代码进行步骤S1-S6完成漏洞检测。

2.根据权利要求1所述的一种基于图嵌入和双向门控图神经网络的漏洞检测方法，其特征在于，所述S1中数据集中源代码的标注方法为：从源代码数据集中提取出函数级的含有漏洞的源代码和不含漏洞的源代码，共包括k个函数；对每个函数中是否含有漏洞进行标注，得到每个函数文件的标注Y_i∈{0,1},i∈[1,k]，其中0表示不存在漏洞，1表示存在漏洞。

3.根据权利要求1所述的一种基于图嵌入和双向门控图神经网络的漏洞检测方法，其特征在于，所述S2中，生成程序依赖图方法为：

S21.将源代码导入源代码分析工具后，根据源代码中的函数名，使用查询语句作为输入，并生成与函数名相应的程序依赖关系图，并将其输出为dot类型的图描述文件；

S22.利用统一变量名映射的方式，在程序依赖关系图的描述文件中以一对一的方式将用户定义的变量名和函数名映射到符号名，得到预处理后的程序依赖图。

4.根据权利要求1所述的一种基于图嵌入和双向门控图神经网络的漏洞检测方法，其特征在于，所述S3中利用正则匹配的方式提取所需的有向边集合和节点集合。

5.根据权利要求4所述的一种基于图嵌入和双向门控图神经网络的漏洞检测方法，其特征在于，所述S4中特征向量字典包括文本向量字典和节点ID字典，所述文本向量字典由文本文件作为输入得到文本特征向量，并用一个字典进行保存，所述节点ID字典由字典文件作为输入得到节点依赖特征向量，并用一个字典进行保存。

6.根据权利要求5所述的一种基于图嵌入和双向门控图神经网络的漏洞检测方法，其特征在于，所述文本向量字典获取方法为：以S31中预处理过后的保存有向边的文本文件作为输入，设置node2vec模型中的采样策略参数，训练文本特征，输出最小文本单元Text_i对应的向量vector_t_i,i∈[1,n]；

将所有输出的文本特征向量用一个字典进行保存，字典为Dict_t＝∪_i∈[1,n]{key:Text_i,value:vector_t_i}。

7.根据权利要求6所述的一种基于图嵌入和双向门控图神经网络的漏洞检测方法，其特征在于，所述节点ID字典获取方法为：

以S31中预处理过后的保存有向边的文本文件作为输入，合理设置node2vec模型中的采样策略参数，将节点标识为唯一的节点ID_i，而非上述的文本属性，训练节点依赖特征，输出图节点之间的依赖特征向量vector_n_i,i∈[1,m]；

将所有输出的节点依赖特征向量用一个字典进行保存，字典为Dict_n＝∪_i∈[1,m]{key:ID_i,value:vector_n_i}。

8.根据权利要求7所述的一种基于图嵌入和双向门控图神经网络的漏洞检测方法，其特征在于，所述S5具体包括以下子步骤：

S51.将表示节点的文本描述合并到一行，然后将字符串拆分为若干个Text，基于文本向量字典Dict_t，将节点文本属性转换为对应的嵌入向量

nodeTextvec_i＝[vector_t_i1,vector_t_i2,...,vector_t_in]，从而得到每个节点的文本向量；

S52.一条有向边代表它存在一对头节点和尾节点，通过使用这两个节点的ID_s,ID_e作为键对节点ID字典进行查询，得到头节点向量vector_n_s和尾节点向量vector_n_e；

S53.将头节点向量和尾节点向量相减，得到一条有向边对应的嵌入向量v_s→e＝vector_n_e-vector_n_s，对每一个程序依赖图的列表中的每条有向边都做上述处理，得到所有PDG中的边向量

S54.将节点文本向量和边向量一起封装成一个程序依赖图对应的JSON文件，作为后续神经网络模型的输入，该JSON文件被视为一个N×16的二维向量矩阵和一个M×16的二维向量矩阵的组合，其中N代表一个程序依赖图中节点的数量，M代表边的数量。

9.根据权利要求8所述的一种基于图嵌入和双向门控图神经网络的漏洞检测方法，其特征在于，所述S6具体包括以下子步骤：

S61.分割训练集与测试集：选取S53中生成的JSON文件数据集中d％的数据样本作为训练集，剩余为测试集；

S62.应用双向门控图神经网络来学习数据集中包含的特征数据，BGGNN由两个方向门控图神经网络组成：一个是正向L₁层的门控图神经网络GGNN₁，接受正向的输入；另一个是反向L₂层的门控图神经网络GGNN₂，学习反向的输入，公式表达为：

上述式中，y_t是模型输出，

是前向输出，

是反向输出；

S63.基于上述网络进行l次迭代训练，训练结束之后保存神经网络模型Model，以方便后期快速的模型加载。

10.根据权利要求8所述的一种基于图嵌入和双向门控图神经网络的漏洞检测方法，其特征在于，所述S7具体包括以下子步骤：

S71.首先如S2、S3的步骤对目标程序源代码进行预处理，得到预处理后的程序依赖图；

S72.以S4中预训练的字典Dict_t和Dict_n为基础，执行如S5中PDG到特征向量矩阵的转换的步骤，保存为目标程序的JSON文件；

S73.复用S6生成的神经网络模型Model，将目标程序生成的函数级的特征向量矩阵作为输入，以进行函数级代码漏洞检测；

S74.输出一个列表，列出目标程序中存在的潜在代码漏洞的函数名列表，以供相关人员检查和完善程序。

Images