AMPLE：基于图简化和增强图表征学习的漏洞检测框架
AMPLE：基于图简化和增强图表征学习的漏洞检测框架
近年来，图神经网络（GNN）在软件漏洞检测领域展现出巨大潜力，但其在处理代码结构图中长距离节点连接时存在局限性。为解决这一问题，研究者提出了一种名为AMPLE的新型漏洞检测框架，通过图简化和增强图表征学习两个核心模块，显著提升了漏洞检测的准确率和F1分数。本文将详细介绍AMPLE的设计理念、技术实现及其在多个公开数据集上的实验表现。
1 问题描述
GNN模型在处理非直接相邻节点之间的长距离连接时存在局限性。虽然可以使用多层堆叠的GNN来学习图的全局信息（即节点间的长距离依赖关系），但深度GNN产生的过度平滑问题会导致具有不同标签的节点出现相似的嵌入，从而降低漏洞检测性能。现有的GNN模型（如GGNN、GCN）无法充分模拟图中的边特征。
2 解决方案
提出了一种新型漏洞检测框架AMPLE，它具有图形简化和增强型图形重现学习功能。AMPLE包括两大部分：1）图形简化，通过缩小代码结构图的大小来减少节点之间的距离。2）增强图形表示学习，包括两个模块。边缘感知图神经网络模块考虑边缘类型信息，并将异构边缘信息融合到节点表示中；而核缩放表示模块则扩大卷积核大小，以明确捕捉远距离图节点之间的关系。
3 框架设计
AMPLE由两个部分组成：图简化模块、增强图表征学习。
3.1 图简化
1）基于类型的图简化：类型合并规则如表1所示，PType表示父节点类型，CType表示子节点类型。合并规则1，2，3，4，5-7分别对应C/C++编程语言中不同类型的语句，包括表达式语句、标识符声明语句、条件语句、for-loop语句和函数调用语句。对于每一对与合并规则相匹配的相邻节点，其子节点将被删除，因为其信息是对父节点的细化，也可以反映在其后续节点中。
表1 类型合并规则表
如图1所示，红色虚线边框的节点根据规则2与父节点合并。具体来说，子节点" * first = malloc(10) "中的信息也被父节点及其三个子节点所覆盖。
图1 基于节点类型的图简化
2）基于变量的图简化：在代码结构图中将具有重复变量的节点合并为一个节点，仅适用于AST的叶节点，因为叶节点不具有子节点，将它们合并不会改变父子层次结构。合并后的变量节点有多个父节点，因此可以同时聚合来自不同语句的信息，可以增强节点表示，促进全局图表征学习。图2展示了一个基于变量的图简化实例。变量"str"同时出现在"char str[15];"和"scanf("%s",str);"的子节点中，因此合并了这两个"str"叶节点。
图2 基于变量的图简化
图简化算法如算法1所示：整个图形简化（GS）过程如算法1所示，其中类型图简化（TGS）和变量图简化（VGS）分别对应第2-19行和第20-25行。该算法以原始代码结构图为输入，根据TGS和VGS输出简化后的代码结构图。对于TGS，算法对AST执行广度优先遍历，每次遇到一对符合表1中合并规则的相邻节点（u、v）时（第10行），都会执行以下操作（第11-12行）：删除v和与v相连的边，然后在u和v的所有子节点之间添加新边。对于VGS，该算法首先获取所有包含相同变量的组，然后将包含相同变量的节点合并为一个节点。最终简化后的代码结构图被输入到AMPLE的后续部分。
算法1 图简化算法
3.2 增强图表征学习
图3 AMPLE检测架构
图3（B）为增强图表征学习模块，该模块主要进行图节点嵌入以及下游分类任务，由边缘感知图卷积网络模块和核缩放表征模块两部分组成。
1）边缘感知图卷积网络模块（Edge-aware GCN module，EA-GCN）：简化后的代码结构图有3种类型的边，分别为进入边（in），出去边（out）和自循环边（self-loop），这些边也即代码结构图中的CFG Edge、NCS Edge，DFG Edge或AST Edge。EA-GCN模块首先通过对三种不同类型的边分别进行加权（不同类型的边对于节点的重要程度不同，这里的边缘感知（EA）可以理解为在图结构中嵌入了3种分别表示不同语法和语义信息的边（Edge），利于每个节点聚合邻近节点更多以及更强的语义信息），然后聚合邻近节点来计算目标节点向量（实际上就是对与目标节点相邻的节点嵌入表示），在图中的每个节点重复执行此操作，得到所有的节点向量，最后根据多头注意力增强重要节点的表示，得到边增强的节点表征向量矩阵。
2）核缩放表征模块：该模块旨在通过明确捕捉远处节点之间的关系来学习图的全局信息。在得到边增强的节点表征向量矩阵后，在核缩放表征模块中设计两个不同尺寸的卷积核，一个大核，和一个小核。大核用于捕获边缘/远处节点特征，小核侧重于邻近/邻居节点特征的获取，通过上述操作，从而捕获代码结构图的全局信息。
4 实验分析
4.1 数据集
AMPLE在研究中采用了三个漏洞数据集，包括FFMPeg+Qemu、Reveal和Fan et al。实验数据集的统计数据如表2所示。Devign提供的FFMPeg+Qemu数据集是人工标注的，来自两个开源C项目，它包含约1万个易受攻击实例和1.2万个非易受攻击实例。Reveal数据集收集自两个开源项目：Linux Debian内核和Chromium。该数据集包含约2k个易受攻击实例和20k个非易受攻击实例。Fan等人从300多个开源C/C++GitHub项目中收集了数据，涵盖了2002年至2019年通用漏洞与暴露（CVE）数据库中的91种不同漏洞类型，该数据集包含约1万个易受攻击实例和17.7万个非易受攻击实例。
表2 数据集
4.2 结果分析
文中采用以下4个指标：Precision、Recall、F1 score、和Accuracy来衡量实验结果。
AMPLE根据下列4个问题来展开实验：
- 问题1：AMPLE在漏洞检测方面的效果如何？
- 问题2：图形简化是否提高AMPLE的性能？
- 问题3：增强图形表征学习是否提高AMPLE的性能？
- 问题4：超参数对AMPLE性能有何影响？
针对问题1，如表3所示，就准确率和F1分数而言，AMPLE优于所有基准方法。特别是，在三个数据集上，AMPLE的F1分数比最佳基准方法分别提高了7.63%、16.48%和40.40%。
表3 与基线和SOTA方法对比实验结果
针对问题2，如表4所示，可以观察到图简化对AMPLE的性能有显著贡献，在三个数据集上的F1分数分别提高了12.56%、14.20%和34.86%。节点、边和距离的平均简化率分别为41.64%、16.79%和41.65%。
表4 图简化前和简化后的实验结果对比
针对问题3，如表5和表6所示，增强图表征学习模块可以有效提高AMPLE性能。在FFMPeg+Qemu、Reveal和Fan等数据集的测试中，EA-GCN模块分别提高了4.46%、30.08%和48.66%的F1分数，而KSR模块则分别提高了13.13%、57.53%和27.93%。
表5 EA-GCN模块核KSR模块对AMPLE性能的影响
表6 EA-GCN与其它GNN模型的实验结果对比
针对问题4，如表7所示，不同的超参数设置会影响AMPLE在漏洞检测中的性能。
表7 超参数对AMPLE性能的影响
5 总结
文中提出了一种新颖的漏洞检测框架AMPLE，它具有图简化和增强图表征学习功能。AMPLE可以缩小代码结构图的节点数量，从而减少节点之间的距离。同时结合边缘类型来增强局部节点的表示，以应对节点表示中更多的异构关系。通过捕捉远处图节点之间的关系来获取图的全局信息。与最先进的基于深度学习的方法相比，AMPLE在所有数据集上的漏洞检测性能都有显著提高，F1分数提高了7.64%-199.81%。
6 原文链接
论文题目：Vulnerability Detection with Graph Simplification and Enhanced Graph Representation Learning
论文出处：2023 IEEE/ACM 45th International Conference on Software Engineering (ICSE 2023)
原文链接：[2302.04675] Vulnerability Detection with Graph Simplification and Enhanced Graph Representation Learning (arxiv.org)