【论文学习】使用图神经网络检测源代码漏洞

论文的简单学习： 2021-TIFS-Combining Graph-Based Learning With Automated Data Collection for Code Vulnerability Detection

 

研究背景：使用机器学习的手段检查代码的漏洞或bug

现有研究的问题：现有对源代码的建模的深度学习方法，通常采用的是RNN结构，但是这种结构会将源代码的程序结构视为一个sequence。然而，程序源码的结构包含有丰富的控制流和信息流等，rnn结构只能捕获简单的表面顺序，因此现有的RNN模型通常表现不佳。

编译器进行传统的代码分析时，通常将其表示为图结构。因此为了更好的建模复杂的代码结构，我们需要一种直接对图结构进行学习、操作的方法，这样才能对数据流和信息流等进行保留和推理。因此该论文希望利用图神经网络结构。

论文 ：Devign: Effective vulnerability identification by learning comprehensive program semantics via graph neural networks 中也是采用图神经网络对源码进行建模，但是其采用的vanilla GNN图神经网络，其edge是没有类型的，这导致他们无法将控制流和信息流等区分开，从而导致其与LSTM模型结果相差不大。

 

本片论文所采用的GNN模型可以表示不同的代码关系（如数据流、控制流、运算顺序、operand value等），提高了模型对程序代码的表示能力。

 

此外深度学习方法还面临着一个问题：训练数据的缺少。

该论文发现了一种利用开源项目源代码的方法：开源项目经常对漏洞进行修补而进行代码提交，因此可以利用丰富的历史提交信息得到训练数据。

然而开源项目的大量提交中，并不是每一次提交都是为了修复漏洞，对此进行手动检查也并不现实。为了解决这一问题，提出了“mixture of experts”方法，即对开源代码使用多个预测模型，并采取高置信度的开源代码用于训练数据。

在论文的实验部分还有一个挺有意思的点就是cross-language detection：

不同编程语言在出现相同的漏洞时会有共性，比如不同语言的资源释放漏洞，都是在有效生命周期结束后没有释放file handler产生的。对于这种共性，可以对网络的高层知识进行迁移。