论文笔记-Devign:Effective Vulnerability Identification by Learning Comprehensive Program Semantics via GNN

  • A+
所属分类:笔记

这篇文章提出了一种基于图神经网络的通用模型Devign,通过学习丰富的代码语义集以进行图级分类。它包括一个新的Conv模块,用于在丰富的节点表征中有效提取有用的特征。数据集基于4个多样化的大规模开源C项目,为高复杂度和多样化的真实源代码。评估结果显示,Devign与现有技术相比,平均准确率提高了10.51%,F1值提高了8.68%,其中Conv模块在平均准确率上增加4.66%,F1值增加6.37%。

论文笔记-Devign:Effective Vulnerability Identification by Learning Comprehensive Program Semantics via GNN

1 提出问题

1.1 传统机器学习算法的不足

机器学习算法将专家人工定义的特征或模式作为检测漏洞的输入,然而产生漏洞的根本原因因漏洞和库的类型而异,因此通过此方法来描述众多库中的所有漏洞是不切实际的。

1.2 目前基于神经网络方法的不足

最近的基于深度神经网络的方法在学习全面的程序语义以描述真实源代码中的高多样性和复杂性的漏洞方面都有很大的局限性。主要表现如下:

  • 在学习方法方面,将源代码视为类似自然语言的平面序列,或仅用部分信息表征。然而,源代码实际上比自然语言更具结构性和逻辑性,并且具有抽象语法树(AST)、数据流、控制流等表示。此外,漏洞有时是微妙的缺陷,需要从语义的多个维度进行全面调查。目前方法的缺陷限制了其覆盖各种漏洞的潜力。
  • 在数据集方面,使用简单的人工代码,即使用"good"或"bad"来区分漏洞代码和非漏洞代码,与复杂的真实代码相比存在较大差距。

2 解决方案

2.1 提出方法

针对上述问题,文章提出了一种新的基于图神经网络的模型,该模型针对真实漏洞数据进行复合编程表示,对编程代码语义进行编码,以捕捉各种漏洞特征。关键创新为提出的Conv模块,它从门控递归单元中获取图的异质节点特征作为输入,并利用传统的卷积层和密集层,分层次地选择更多的粗略特征,进行图层分类。

2.2 主要贡献

文章的主要贡献如下:

  • 以AST为主,将不同级别的程序控制和数据依赖性编码为异质边的联合图,综合表示有助于捕获更广泛的漏洞类型和模式,并能通过图神经网络学习更好的节点表示。
  • 提出带有Conv模块的GGNN模型,Conv模块从节点特征中分层学习,以捕获更高层次的图级分类任务的表征。
  • 与基线方法相比,Devign的平均准确率提高10.51%,F1得分提高8.68%。同时,Conv模块的平均精度增益4.66%,F1增益为6.37%。与知名的静态分析器进行比较,Devign在所有数据集上的平均F1得分显著提升27.99%。将Devign应用于从4个项目中收集的40个最新CVE漏洞,实现了74.11%的准确率。

3 模型介绍

Devign可以划分为三个部分:1)复合代码语义的图嵌入层,将源代码编码为具有综合语义的联合图结构。2)门控图递归层,通过聚集和传递图中相邻节点的信息来学习节点的特征。3)Conv模块,提取用于图级预测的有意义的节点表示。

论文笔记-Devign:Effective Vulnerability Identification by Learning Comprehensive Program Semantics via GNN

3.1 问题公式化

Devign在函数细粒度级别分析漏洞,并将漏洞函数的识别问题转化为二元分类问题,即给定一个函数,判断其是否存在潜在漏洞。将一个数据样本定义为\( \left(\left(c_{i}, y_{i}\right) \mid c_{i} \in\right.\left.\mathcal{C}, y_{i} \in \mathcal{Y}\right), i \in{1,2, \ldots, n}.\),其中\(C\)表示函数的集合,

\[\min \sum_{i=1}^{n} \mathcal{L}\left(f\left(g_{i}(V, X, A), y_{i} \mid c_{i}\right)\right)+\lambda \omega(f)\]

3.2 复合代码语义的图嵌入层

图嵌入层将函数代码\(c_{i}\)映射为可供模型输入的图数据结构,表示如下:

\[g_{i}(V, X, A)=E M B\left(c_{i}\right), \forall i={1, \ldots, n}\]

图2展示了一个存在整数溢出漏洞的代码表示为联合图结构的例子,该图在传统的三种代码表征结构(即AST控制流数据流图)的基础上,另外加入了源代码的自然序列

  • 抽象语法树(AST)展示代码的语法结构。蓝色框为AST叶节点,紫色箭头表示父子关系。
  • 控制流图(CFG)描述了程序在执行过程中遍历的所有路径。图2中绿色箭头表示CFG的边。
  • 数据流图(DFG)表示程序中变量的使用情况,数据流是面向变量的,任何数据流都涉及对某些变量的访问或修改。图2中橙色箭头表示DFG的边。
  • 使用自然代码序列(NCS)边来连接AST中的相邻代码标记,保留了源代码序列所反映的编程逻辑。NCS边缘在图2中用红色箭头表示,连接AST的所有叶节点。
论文笔记-Devign:Effective Vulnerability Identification by Learning Comprehensive Program Semantics via GNN

最终一个函数\(c_{i}\)可以表示为4种类型的子图共享一组节点\(V\)得到的联合图\(g\)。每一个节点有两个属性,代码类型。代码即节点表示的源代码,类型即节点的属性。利用预训练模型word2vec对源代码以及类型编码,并将两种编码连接起来,作为最终的节点表征\(x_{v}\)。

3.3 门控图递归层

图神经网络的关键思