基于图神经网络的代码漏洞检测方法

文章结构

1. 引言

1.1 漏洞现状

1.2 漏洞研究传统特点

现有的研究方案大多把源代码转化为文本序列，使用自然语言处理中常用的循环神经网络进行自动化特征提取，但这忽略了代码中由于跳转、循环、判断形成的控制流特征，往往会有较大的信息损失。如何提取源代码样本的特征信息，把源代码样本转换成适合后续神经网络训练的数据形式并最大化保留样本信息是关键问题之一。

1.3 文章贡献

1.3.1 设计了一种基于中间语言控制流图的代码特征提取方案

本文通过代码中间语言的控制流图提取代码样本的结构特征，使用控制流图基本块中的代码序列提取代码样本的语义特征，将代码样本转化为图结构数据，以解决序列样本丢失代码结构特征信息的问题。通过提取源代码的控制流图结构信息，同时使用基本块的中间语言提取代码语义信息，嵌入图节点中，将源代码转化为图结构数据样本。

1.3.2 使用图神经网络模型处理代码图结构样本

1.3.3 基于设计方案设计了实验方案并对结果进行了分析

2. 相关工作

2.1 图表示发展

2.2 深度学习图表示发展

3. 方案设计

3.1 代码预处理LLVM+chang编译器

3.2 生成图 LLVM分析器 + word2vec

3.3 图转向量 PYG

3.4 训练模型

3.3.1 多层

本文提出的图神经网络结构由多组图卷积和池化层组成的基本单元来完成在图级别上的端到端分类学习问题，

3.3.2 残差

多组基本单元之间可以直接串联连接，也可以通过残差方式连接。文
献[16]指出残差连接方式在多层网络上有更好的效果。

3.3.3 通过应用多层感知器（MLP）和Softmax 层

4. 实验设计和评估

3.1 数据集

本文采用的数据集是来源于美国国家标准技术研究所（NIST）的软件保障参考数据集（SARD）

3.2 参数设置

本文的模型测试环境为 CentOS 7.6 操作系统，Intel Xeon CPU E5-2630 处理器，180 GB 内存，GeForce GTX 1080 显卡，8 GB 显存。模型使用 GPU 显卡加速模型训练。
在图结构数据生成部分，word2vec 算法生成的词嵌入参数大小设定为 256，基本块初始化向量的大小即 num_node_features 为 256。

3.3 评估指标

3.4 结果分析

• 直接连接方案在大于 2 层网络层之后出现了准确率的大幅下降，而残差网络连
  接方案在多层网络环境下具有更好的效果，因此本文选择了效果最好的 3 层残差连接网络。
• 本文提出的基于控制流图的特征提取与图神经网络模型在开源数据集上的准确率
  均大幅高于静态扫描工具

实验设置

基本过程

• 提出了一种基于图神经网络的代码漏洞检测方法，通过中间语言的控制流图特征，实现了函数级别的智能化代码漏洞检测。首先，将源代码编译为中间表示，进而提取其包含结构信息的控制流图，同时使用词向量嵌入算法初始化基本块向量提取代码语义信息；然后，完成拼接生成图结构样本数据，使用多层图神经网络对图结构数据特征进行模型训练和测试。采用开源漏洞样本数据集生成测试数据对所提方法进行了评估，结果显示该方法有效提高了漏洞检测能力。

文章创新点

套用新模型GNN结合多层、残差、MLP

使用控制流图表示方法

编码方式结构+语义 dot + word2vec pyg

论文资源链接