CN111783100A - 基于图卷积网络对代码图表示学习的源代码漏洞检测方法 - Google Patents

Abstract

本发明公开了一种基于图卷积网络对代码图表示学习的源代码漏洞检测方法，所述方法如下：生成代码属性图；在代码属性图中添加函数调用关系和过程间依赖关系；根据漏洞关键点获取代码切片；利用切片对图中节点进行删减，提取与漏洞相关的图结构信息；使用图卷积网络学习每个节点的向量表示；根据边的类型划分子图，并通过基于注意力机制的READOUT模型得到图的向量表示；根据图的向量表示和标签调整网络参数；用训练好的模型检测代码漏洞。本发明能充分利用和学习漏洞代码的结构和属性信息，避免传统深度网络在对代码表示学习时易丢失代码结构信息及因需要把代码表示成固定长度序列而丢失长代码上下文信息的问题，有助于降低漏洞检测的误报和漏报。

Description

基于图卷积网络对代码图表示学习的源代码漏洞检测方法

技术领域

本发明涉及一种软件漏洞检测方法，具体涉及一种基于图卷积网络对代码进行图表示学习的源代码漏洞检测方法。

背景技术

软件漏洞是在软件设计与开发实现的过程中存在的一些容易被恶意攻击者利用的缺陷。传统的源代码审查技术在很大程度上取决于审查人员对安全问题的理解与长期经验的积累，并且在代码规模和复杂程度日益增大的情况下无法满足对漏洞检测的需求。基于机器学习的漏洞检测方法虽然避免了基于规则的漏洞检测方法依赖专家人工编写检测规则的问题，但是仍需要人工提取漏洞特征。而近年来成功应用于自然语言处理、图像识别、目标检测领域的深度学习技术可以降低对专家经验和手工特征工程的依赖，为自动提取漏洞特征和生成漏洞模式提供了可能。

然而，由于编程语言的特殊性、漏洞类型的多样性、漏洞上下文的复杂性、漏洞代码与漏洞特征在抽象程度上的高差异性、漏洞代码与修复代码之间的高相似性，这些都使得深度学习自动学习漏洞模式比其他领域的深度学习问题更加困难，给基于深度学习的漏洞检测技术带来了巨大的挑战。

当前用于自动学习漏洞模式的深度学习模型大多采用语言模型对代码语义进行建模，把代码的各种中间表示转化为一个平铺的一维序列，即将其当作自然语言文本，基于自然语言处理领域常用的深度神经网络(如LSTM和GRU)来处理转换后的代码序列，然后将其自动学习到的代码漏洞特征用于训练一个机器学习分类器，以进行漏洞检测。相对于自然语言文本而言，代码更具有结构化的特点，这种漏洞检测方法未能充分利用和学习漏洞代码的结构信息和属性信息，并且因需要把代码表示成固定长度的序列还会丢失长代码的上下文信息，在漏洞检测时往往存在较高的误报率和漏报率。

代码属性图(Code Property Graph,CPG)是一种程序语法、控制流和数据流的联合表示，结合了抽象语法树、控制流图和程序依赖图来综合表征程序的结构和语义，是Yamaguchi等人(F.Yamaguchi,N.Golde,D.Arp,K.Rieck,Modeling and discoveringvulnerabilities with code property graphs,Proceedings,IEEE Symposium onSecurity and Privacy.(2014)590–604.doi:10.1109/SP.2014.44)首次提出的一种新型的代码表示形式，该漏洞检测方法使用图数据库查询语句遍历代码属性图，利用模式匹配的方法查找符合某种模式的软件漏洞，提高了对已知漏洞模式识别的准确率，但是该方法仅适用于识别已知的漏洞模式，并且在代码属性图中没有考虑过程间分析，对于跨过程调用的漏洞有可能产生漏报。Wang等人(CPGVA-Code Property Graph based VulnerabilityAnalysis by Deep Learning，2018)进一步使用CNN和LSTM等传统的深度神经网络来从代码属性图上学习漏洞代码模式，提高了漏洞识别的准确率，但仍存在较高的漏报率。

相对于使用以固定长度序列化数据作为输入的传统深度神经网络(如LSTM和GRU)而言，图神经网络更适合学习图输入这种与节点输入顺序无关的结构特征表示，且对需要学习的图数据的节点和边的数量没有限制，更适合对漏洞代码的复杂结构语义进行有效编码以捕获更广泛的漏洞特征。但是目前仅有两篇利用图神经网络进行漏洞检测的研究。一篇是2019年岳佳的硕士学位论文“基于漏洞基因的软件漏洞检测研究”将从代码的抽象语法树中提取的特征作为文本信息，来构造一个大的文本图，利用图卷积网络(GraphConvolutional Network,GCN)对提取的特征进行分类进而实现漏洞检测，同时还提出一种基于漏洞基因的漏洞检测方法，基于代码属性图和图遍历来实现漏洞检测。前者仍将代码视作为文本来处理，未充分提取和利用代码的图结构信息，后者仍使用模式匹配的方式，未使用深度学习模型对代码进行图表示学习。另一篇是Zhou等人(Yaqin Zhou,ShangqingLiu,Jingkai Siow,Xiaoning Du,and Yang Liu,Devign:Effective VulnerabilityIdentification by Learning Comprehensive Program Semantics via Graph NeuralNetworks[C]，33rd Conference on Neural Information Processing Systems,Vancouver,Canada33rd Conference on Neural Information Processing Systems(NeurIPS 2019),Vancouver,Canada)首次提出的用门控图神经网络(Gated graph neuralnetwork，GGNN)从以AST为主干建立的代码属性图上学习漏洞模式的方法。该方法以AST为主干显式编码程序的控制依赖和数据依赖，在函数规模较大时存在图结构过深过大导致学习效率较低的问题，而且该方法对代码的分析仅限于一个函数内，未考虑函数调用与过程间的数据依赖，对于跨函数调用的漏洞有可能产生漏报。

发明内容

本发明的目的是提供一种基于图卷积网络对代码图表示学习的源代码漏洞检测方法，该方法能够充分利用和学习漏洞代码的结构、属性信息和上下文信息，避免传统的深度网络在对代码进行表示学习时丢失代码结构信息和长代码的上下文信息的问题，有助于降低漏洞检测的误报和漏报。

本发明的目的是通过以下技术方案实现的：

一种基于图卷积网络对代码图表示学习的源代码漏洞检测方法，首先，通过解析源代码，生成以抽象语法树、控制流图和程序依赖图联合表示代码结构信息、以代码内容和节点类型表示代码属性信息的代码属性图。为了更准确地提取漏洞相关的图结构信息，尤其是跨函数调用的漏洞结构信息，在代码属性图中引入了函数调用关系和过程间程序依赖关系。为了避免源代码中大量漏洞无关语句对漏洞检测造成的噪声干扰，加快模型的学习速度，利用程序切片技术，根据可能的漏洞关键点生成程序切片，利用程序切片对改进的代码属性图表示的图结构进行简化，去除与漏洞关键点无关的节点。

其次，基于图卷积神经网络GCN，对上述简化后的图结构进行图表示学习，将软件漏洞检测问题转化为图级分类问题，实现对代码结构信息和属性信息端对端的学习。首先是基于GCN在每个代码属性图节点上学习代码的结构特征，即将代码属性图作为图数据，直接对其进行图表示学习，而不是将其看作文本分类问题进行处理，然后为了更好地学习图的局部和全局结构信息，提出了按关系类型划分子图并对各个子图和全图分别进行表示学习的方法，以及基于子图自注意力和节点注意力机制的READOUT(读出)模型。按关系类型(即边的类型)拆分子图，有助于提取出与不同类型漏洞相关的子图结构。基于子图自注意力和节点注意力机制的READOUT模型，有助于突出每个节点和每个子图在漏洞检测任务中的重要程度。

具体包括如下步骤：

步骤1：通过解析源代码，生成以抽象语法树、控制流图和程序依赖图联合表示代码结构信息、以代码内容和节点类型表示代码属性信息的代码属性图；

步骤2：在代码属性图中添加函数调用关系和过程间依赖关系；

步骤3：利用程序切片技术提取与漏洞关键点相关的程序切片；

步骤4：利用得到的程序切片简化利用步骤1和步骤2得到的图结构信息，得到简化后的与漏洞相关的图结构：

步骤5：使用图卷积网络对该图数据进行表示学习，学习每个节点的向量表示；

步骤6：根据边的类型，在图结构上划分出多个子图，并通过基于节点注意力机制的READOUT模型得到每个子图的向量表示，基于子图自注意力机制的READOUT模型得到全图的向量表示；

步骤7：将得到的子图和全图的向量表示进行拼接后送入全连接层和softmax层得到预测结果，利用标签信息计算交叉熵损失函数，根据误差反向传播调整网络参数，直到网络对输入的响应达到预定的目标范围为止，训练结束；

步骤8：用训练好的图卷积神经网络模型对代码进行漏洞检测。

相比于现有技术，本发明具有如下优点：

(1)本发明针对漏洞检测问题对常用的代码属性图所作的改进，在代码属性图中去除了抽象语法树的边，增加了函数调用图的边，能够降低代码属性图的规模，加快模型的训练速度，并有效检测跨函数调用的漏洞。

(2)本发明利用程序切片技术从代码属性图中提取漏洞相关的图结构，能够避免源代码中大量漏洞无关语句对漏洞检测的影响，提高检测的准确率，并进一步加快模型的训练速度。

(3)本发明提出的利用图卷积网络直接对基于程序切片提取的漏洞相关的图结构进行图表示学习的方法，能够充分利用和学习漏洞代码的结构信息和属性信息，并适应不同长度的代码。

(4)本发明提出的按关系类型(即边的类型)划分子图并分别进行子图和全图表示学习的方法，能够充分利用和学习漏洞代码的局部和全局结构信息，降低检测的漏报率。

(5)本发明提出的基于子图自注意力和节点注意力机制的READOUT(读出)模型，能够有效学习每个节点和子图对不同类型漏洞的重要程度，降低检测的误报率。

附图说明

图1是本发明的流程示意图。

图2是漏洞相关图结构的提取流程图。

图3是提取出的图结构示例图。

图4是基于程序切片的图结构简化算法。

图5是按关系类型划分子图的算法。

图6是基于子图自注意力和节点注意力机制的READOUT模型示意图。

具体实施方式

下面结合附图对本发明的技术方案作进一步的说明，但并不局限于此，凡是对本发明技术方案进行修改或者等同替换，而不脱离本发明技术方案的精神和范围，均应涵盖在本发明的保护范围中。

本发明提供了一种基于图卷积网络对代码进行图表示学习的源代码漏洞检测方法，首先，利用代码解析工具生成以控制流图、程序依赖图、定义使用图为核心的代码属性图，然后根据代码属性图中边的类型查找出漏洞检测关注的控制流、数据依赖和函数调用相关的图结构，然后再使用程序切片技术，获取与漏洞关键点相关的程序切片，根据程序切片简化图结构信息，获取漏洞相关的图结构。然后使用图卷积神经网络，学习图结构中每个节点的向量表示，最后再基于节点注意力机制的READOUT模型得到每个子图的向量表示，基于子图自注意力机制的READOUT模型得到全图的向量表示。如图1和2所示，具体步骤如下：

步骤1：通过解析源代码，生成以抽象语法树、控制流图和程序依赖图联合表示代码结构信息、以代码内容和节点类型表示代码属性信息的代码属性图；

步骤2：提取控制流图、程序依赖图和函数调用图相关的图结构信息，提取的图结构示例如图3所示，具体步骤如下：

步骤21：选取图中类型为函数的节点；

步骤22：遍历其孩子节点即函数体的节点，查找边类型为与控制流和数据流相关的边，获取到节点和边信息加入图结构信息中；

步骤23：从孩子节点中选取类型为Callee的节点，根据函数名和路径递归查找被调用的函数；

步骤24：重复步骤21、22、23，直至遍历完函数中的所有节点，得到与漏洞相关的控制流、程序依赖和函数调用的图结构信息；

步骤3：利用程序切片技术提取与漏洞关键点相关的程序切片；

步骤4：利用得到的程序切片简化步骤1和步骤2得到的图结构信息，得到简化后的与漏洞相关的图结构如图4所示，具体步骤如下：

步骤41：初始化简化后的图结构列表(Simplified Graph Records，SGR)为空，节点列表N为空，边列表E为空；

步骤42：对于输入的全部图结构(Graph Records,GR)中的每个函数体的图结构gr，从切片列表(Slices list,S)中筛选出属于同一个函数体的切片s；

步骤43：使用Filter(object,condition)过滤函数，目的为从object中筛选出符合condition的object的子集，然后对于当前的gr中的节点gr.node，我们筛选出符合节点的代码存在于切片s的情况的节点加入节点列表N中；

步骤44：若N不为空，则再次筛选图结构gr中的边列表gr.edge，选出gr.edge中边的两端的节点都存于节点列表N中的边信息，并加入边列表E中；

步骤45：最后将(N,E)添加至简化后的图结构SGR中；

步骤5：使用图卷积神经网络对该图数据进行表示学习，学习图结构中每个节点的向量表示；

步骤6：图向量表示学习，具体步骤如下：

步骤61：根据边的类型划分出多个子图，按关系类型划分子图的算法如图5所示，该算法通过DGL(https://docs.dgl.ai/)提供的api实现；

步骤62：计算所有节点的自注意力，并进行池化操作，再根据节点的自注意力值对节点的向量表示进行加权求和，得到每个子图的向量表示；

步骤63：通过一个多层感知机模型(MLP)计算每个子图的注意力值，最后利用每个子图的注意力值对子图的向量表示进行加权求和，得到全图的向量表示，示例如图6所示，具体计算公式如下：

其中，

为sub子图经过l层卷积操作后节点表示的矩阵，

为每个节点添加自连接后的邻接矩阵，

为

对应的每个节点的度数的矩阵，为该网络的参数，Z为得到的注意力值，n_sub为sub子图中节点个数，idx为根据Z排序，筛选其中前k％的节点的索引值，h_sub为sub子图的节点乘以相应的注意力值累加得到的子图向量表示，h_CFG,h_PDG,h_CG,为对应CFG、PDG、CG等这些子图的表示，W为计算得到的权重矩阵，h_G为最终的图表示。

步骤7：将得到的子图和全图的向量表示进行拼接后送入全连接层和softmax层得到预测结果，利用标签信息计算交叉熵损失函数，根据误差反向传播调整网络参数，直到网络对输入的响应达到预定的目标范围为止，训练结束；

步骤8：用训练好的图卷积神经网络模型对代码进行漏洞检测。

实施例：

以一个有1088行代码的源文件为例，按可能的一种漏洞关键点对程序进行切片的结果只有48行语句，如下所示：

该源文件包含的漏洞发生在第29行，使用strcpy函数时，没有对malloc申请的stonesoup_buffer的结果进行核验就使用strcpy，若malloc申请失败，那么strcpy的使用也会造成非法内存访问。

本发明提出的方法可以通过基于图卷积网络对代码进行图表示学习后检测出该代码包含不安全函数调用相关的漏洞。像类似这样的长代码，在使用其他以序列化数据作为输入的传统深度学习模型(如如LSTM和GRU)检测时，有可能在批量学习时丢失关键的上下文信息，从而导致传统的深度学习模型不能很好地适应不同长度的代码。

本发明基于图卷积神经网络对漏洞代码的代码属性图的节点和图结构进行表示学习，能够充分利用和学习漏洞代码的结构信息、属性信息和上下文信息，因图卷积神经网络更适合学习图输入这种与节点输入顺序无关的结构特征表示，并且能够适应不同长度的代码，无需像传统的CNN和LSTM那样要把代码表示成固定长度的序列，因此该方法更适合对漏洞代码的复杂结构语义进行有效编码，以捕获更准确的漏洞特征，避免了传统的深度神经网络在对代码进行表示学习时因序列化处理而丢失代码结构信息的问题。本发明提出的基于程序切片对图结构进行删减节点的方法，能够有效去除漏洞无关节点对漏洞检测的影响，并提高模型的训练速度。本发明提出的按关系类型划分子图以及节点和子图注意力机制能够充分利用和学习代码的局部和全局结构信息以及节点的属性信息，降低漏洞检测的误报率和漏报率。

Claims (5)

1.一种基于图卷积网络对代码图表示学习的源代码漏洞检测方法，其特征在于所述方法包括如下步骤：

步骤1：通过解析源代码，生成以抽象语法树、控制流图和程序依赖图联合表示代码结构信息，以代码内容和节点类型表示代码属性信息的代码属性图；

步骤2：在代码属性图中添加函数调用关系和过程间依赖关系；

步骤3：利用程序切片技术提取与漏洞关键点相关的程序切片；

步骤4：利用得到的程序切片简化步骤1和步骤2提取得到的图结构信息，得到简化后的与漏洞相关的图结构：

步骤5：使用图卷积网络对图数据进行表示学习，学习每个节点的向量表示；

步骤6：根据边的类型，在图结构上划分出多个子图，并通过基于节点注意力机制的READOUT模型得到每个子图的向量表示，基于子图自注意力机制的READOUT模型得到全图的向量表示；

步骤7：将得到的子图和全图的向量表示进行拼接后送入全连接层和softmax层得到预测结果，利用标签信息计算交叉熵损失函数，根据误差反向传播调整网络参数，直到网络对输入的响应达到预定的目标范围为止，训练结束；

步骤8：用训练好的图卷积神经网络模型对代码进行漏洞检测。

2.根据权利要求1所述的基于图卷积网络对代码图表示学习的源代码漏洞检测方法，其特征在于所述步骤2的具体步骤如下：

步骤21：选取图中类型为函数的节点；

步骤22：遍历其孩子节点即函数体的节点，查找边类型为与控制流和数据流相关的边，获取到节点和边信息加入图结构信息中；

步骤23：从孩子节点中选取类型为Callee的节点，根据函数名和路径递归查找被调用的函数；

步骤24：重复步骤21、22、23，直至遍历完函数中的所有节点，得到与漏洞相关的控制流、程序依赖和函数调用的图结构信息。

3.根据权利要求1所述的基于图卷积网络对代码图表示学习的源代码漏洞检测方法，其特征在于所述步骤4的具体步骤如下：

步骤41：初始化简化后的图结构列表为空，节点列表N为空，边列表E为空；

步骤42：对于输入的全部图结构中的每个函数体的图结构gr，从切片列表(Sliceslist,S)中筛选出属于同一个函数体的切片s；

步骤43：使用Filter过滤函数，从object中筛选出符合condition的object的子集，然后对于当前的gr中的节点gr.node，筛选出符合节点的代码存在于切片s的情况的节点加入节点列表N中；

步骤44：若N不为空，则再次筛选图结构gr中的边列表gr.edge，选出gr.edge中边的两端的节点都存于节点列表N中的边信息，并加入边列表E中；

步骤45：最后将(N,E)添加至简化后的图结构SGR中。

4.根据权利要求1所述的基于图卷积网络对代码图表示学习的源代码漏洞检测方法，其特征在于所述步骤6的具体步骤如下：

步骤61：根据边的类型划分出多个子图；

步骤62：计算所有节点的自注意力，并进行池化操作，再根据节点的自注意力值对节点的向量表示进行加权求和，得到每个子图的向量表示；

步骤63：通过一个多层感知机模型计算每个子图的注意力值，最后利用每个子图的注意力值对子图的向量表示进行加权求和，得到全图的最终向量表示。

5.根据权利要求4所述的基于图卷积网络对代码图表示学习的源代码漏洞检测方法，其特征在于所述注意力值、子图的向量表示和全图的最终向量表示的具体计算公式如下：

其中，

为sub子图经过l层卷积操作后节点表示的矩阵，

为每个节点添加自连接后的邻接矩阵，

为

对应的每个节点的度数的矩阵，Z为得到的注意力值，h_sub为sub子图的节点乘以相应的注意力值累加得到的子图向量表示，h_CFG,h_PDG,h_CG,为对应CFG、PDG、CG子图的表示，W为计算得到的权重矩阵，h_G为最终的图表示。

Images