LineVD_Statement-level Vulnerability Detection using Graph Neural Networks

0 Abstract

当前基于机器学习的软件漏洞检测方法主要在功能级别进行。然而,这些方法的一个关键限制是,它们没有指示导致漏洞的特定代码行。这限制了开发人员有效检查和解释学习模型预测的能力,这对于将基于机器学习的工具集成到软件开发工作流中至关重要。基于图的模型在功能级漏洞检测方面表现出了良好的性能,但其在语句级漏洞检测中的能力尚未得到广泛探索。**虽然通过可解释的人工智能解释功能级预测是一个很有前途的方向,但我们在这里从完全监督学习的角度来考虑语句级软件漏洞检测任务。**我们提出了一种新的深度学习框架LineVD,它将语句级漏洞检测定义为节点分类任务。**LineVD利用图神经网络和基于转换器的模型对原始源代码标记进行编码,从而利用语句之间的控制和数据依赖性。**特别是,通过解决函数级和语句级信息之间的冲突输出,LineVD显著提高了函数代码在没有漏洞状态的情况下的预测性能。我们针对从多个真实世界项目中获得的大量真实世界C/C++漏洞进行了广泛的实验,并证明F1分数比当前最先进的技术提高了105%.

1 Intro or Overview

1.1 Problem and Challenge

自动化SVD大致可分为两类:

(1)传统方法,包括静态和动态分析;

(2)数据驱动解决方案,利用数据挖掘和机器学习来预测软件漏洞;

尽管当前的数据驱动方法在识别软件漏洞方面取得了成功,但它们往往局限于粗粒度水平。模型输出通常为开发人员提供有限的预测结果验证和解释信息,导致在评估和缓解软件漏洞时付出额外努力。

许多SVD解决方案已经从文件级过渡到函数级或切片级预测,其他一些工作进一步利用补充信息,如提交级别的代码更改以及附带的日志消息,来构建预测模型。虽然目标是帮助从业者对有缺陷的代码进行优先级排序,但漏洞通常可以局限于几个关键行。因此,审查大型函数仍然可能是一个相当大的负担。

1.2 Motivation

image-20231218205705483

为了节省空间,我们从一个较小的函数中选择一个漏洞,该函数包含Linux内核(CVE-2018-12896)中的整数溢出漏洞,该漏洞最终可能被利用来导致拒绝服务。通过显式语句级别的预测,可以更容易地解释为什么函数被预测为易受攻击(或者,验证预测是否错误)。语句级SVD模型将第22行上的加法赋值操作标记为最可疑的,该操作包含易受攻击的整数强制转换操作,使开发人员能够更有效地验证和减轻该漏洞。

先前的工作利用GNNExplainer来导出易受攻击的语句作为模型的解释,以展示在语句级的工作。然而,在我们的工作中,我们发现在对潜在的脆弱性语句进行分类和排序时,性能是不充分和有效的。或者,我们旨在探索在语句级别直接训练和预测漏洞以进行SVD粒度细化的可行性和有效性,这将允许数据驱动的解决方案以完全监督的方式直接利用任何可用的语句级别信息。

1.3 Contribution

  • 提出了一种新颖有效的语句级SVD方法,LineVD实现了显著的改进,F1得分增加了105%;
  • 研究了构建基于GNN的语句级SVD模型的每个阶段的性能影响,包括节点嵌入方法和GNN模型选择。根据研究结果,开发LineVD是为了通过同时学习功能和语句级别的信息,在很大程度上提高性能。
  • LineVD是第一种通过图神经网络联合学习函数级和语句级信息以提高SVD性能的方法,在经验评估中,它显著优于仅使用一种类型信息的传统模型。
  • 发布了数据集、源代码和带有支持脚本的模型,这为未来的基准测试和比较工作提供了一个现成的实现解决方案。https://github.com/davidhin/linevd

2 Architecture & Method

2.1 System Overview

image-20231218211008424

首先将问题定义为,节点VYV\rightarrow Y的映射,也就是语句是否易受攻击。

通过学习最小化损失loss:

mini=1nL(f(Gi,YiVi))min\sum_{i=1}^nL(f(G_i, Y_i|V_i))

Feature Extraction

LineVD将源代码的单个函数作为原始输入。通过处理函数并将其拆分为单独的语句Vi,首先通过CodeBERT的预训练BPE标记器对每个样本进行标记。在V={V1,V2,…,Vn}的集合之后,整个函数和包括该函数的各个语句被传递到CodeBERT中。因此,可以获得函数级和语句级的代码表示。

具体而言,LineVD分别嵌入了函数级和语句级代码,而不是为函数级嵌入聚合语句级嵌入。CodeBERT是一个双峰模型,这意味着除了函数代码本身之外,它还基于函数的自然语言描述进行了训练。作为输入,它使用一个特殊的分隔符标记来区分自然语言描述和函数代码。虽然函数的自然语言描述是不可访问的,但在这项工作中应用了文献中规定的一般操作,在每个输入前添加一个额外的分隔符标记,使描述为空白。对于CodeBERT的输出,我们使用了分类标记的嵌入,这适用于代码摘要任务。这使我们能够更好地利用CodeBERT模型强大的预训练源代码摘要功能。

总体而言,使用CodeBERT的LineVD的特征提取组件产生n+1个特征嵌入:一个嵌入用于整个函数,n个嵌入用于每个语句,我们分别表示为Xv=x1v,x2v,,xnvXv={x^v_1,x^v_2,…,x^v_n}

Graph Construction

在LineVD中,我们专注于数据和控制相关性信息,为此我们引入了图注意力网络(GAT)模型。如第2.2节所述,图神经网络(GNN)基于信息扩散机制学习图结构数据,而不是将信息压缩成平面向量,平面向量根据图的连通性更新节点状态,以保留重要信息,即拓扑依赖信息。

Classifier Learning

目标是训练一个可以同时从函数级和语句级代码中联合学习的模型。为了实现这一点,我们认为函数级和语句级代码片段对预测结果的贡献相等。因此,我们利用函数级CodeBERT嵌入的输入和从GAT层获得的语句嵌入,构建了一组共享的线性层和dropout层。

虽然易受攻击的语句可能足以指示函数易受攻击,但我们使用元素乘法构建LineVD,以进一步利用函数级别的信息进行训练。此外,这种操作和谐地平衡了函数级和语句级嵌入之间的冲突输出,并将证明某些场景的决策是合理的,即,如果函数级嵌入的输出类为零,那么所有语句级输出也为零。对此的直觉是,非易受攻击的函数不可能具有易受攻击的代码行。

3 Experiment and Evaluation

3.1 DataSet and Process

最近的研究表明,SVD模型应该根据能够代表真实世界漏洞不同特征的数据进行评估[10]。这意味着评估从真实世界项目中提取的源代码(即非合成的),同时保持不平衡的比率,这是软件项目中漏洞固有的。在现实世界场景中应用时,使用不满足这些条件的数据集会导致模型性能的不一致。另一个数据集要求是足够多的样本,理想情况下跨越多个项目,以便获得一个可以很好地推广到看不见的代码的模型。最后一个要求是在语句级别访问基本事实标签,或可追溯到修复前代码,即原始gitcommit。

Ground-Truth labels

为了获得易受攻击和非易受攻击线路的基本事实标签,我们遵循文献[19,32]中的断言,而不是提出我们自己的启发式方法:(1)漏洞修复提交中删除的线路用作易受攻击的线路的指标,以及(2)所有依赖于添加线路的控制或数据的线路也被视为易受攻击。第二点的理由是,在漏洞修复提交中添加的任何行都是为了帮助修补漏洞。因此,在漏洞修复提交中未修改但与这些添加的行相关的行可以被视为与漏洞相关。为了获得与依赖于添加行的行相对应的标签,我们首先获得样本前后版本的代码变化,其中Big Vul中的样本指的是函数级代码片段。对于之前的版本,我们删除所有添加的行,对于之后的版本,删除所有删除的行。在这两种情况下,我们都保留空白占位符行,以确保行号的一致性。从后版本中提取的代码图可用于查找所有依赖于添加行的控制或数据行,这些行的行号对应于前版本。这组线可以与删除的线组合,以获得单个样本的最终脆弱线集。注释行被排除在代码图中,因此不用于训练或预测。这可以从图1和图2中看出;在这种情况下,只有一个修改的行,它被视为删除的行(22)和添加的行(23)。在这种情况下,前后版本的控制和数据依赖边恰好相同,因此我们可以使用图2来识别依赖于第23行的控制/数据的行,即第3、19和21行。

Cleaning

原始数据集中的一些样本被错误地截断,导致代码样本无法解析且无效。例如,一个原本是50行的函数可能会因为没有明显原因而被错误地截断为40行。原因可能是数据集最初是如何构建的错误;然而,在整个数据集中只有30个这样的样本。我们使用80:10:10的随机训练/验证/测试分割比。对于训练集,我们对不可破坏样本的数量进行了不足采样,以在函数级别生成近似平衡的数据集,而测试和验证集保持原始的不平衡比例。我们选择在函数级别上平衡样本,因为在语句级别上进行平衡,同时保持函数中语句之间的上下文依赖关系是非常重要的。

3.2 Evaluation

  • RQ1:与最先进的基于解释的SVD模型相比,LineVD可以实现多大的性能提升?

    image-20231219093737314

  • RQ2:不同的代码嵌入方法如何影响语句级漏洞检测?与其他粒度级别的SVD相比,语句级SVD的代码嵌入方法尚未得到探索。

  • RQ3:图神经网络和函数级信息如何对LineVD性能做出贡献?使用图神经网络的信息传播对语句级SVD的影响还有待探索。

  • RQ4:LineVD在跨项目分类场景中的表现如何?虽然在包含多个项目的数据集上进行训练已经减少了对模型通用性的歪曲,但来自同一项目的样本仍然有可能出现在训练集和测试集中。使用跨项目场景可以更好地表示模型在完全看不见的项目上的表现,而不仅仅是看不到的样本。

    image-20231219094311928

  • RQ5:对于真实世界的数据,LineVD最能区分哪些语句类型?从语句类型的角度研究模型预测结果,特别是对于真实世界的数据,可以帮助了解模型在哪里表现最好,在哪里失败,这可以指导未来的工作和语句级SVD的改进。

    image-20231219094343320

4 Conclusion

LineVD,一种用于语句级漏洞检测的新型深度学习方法,它可以让开发人员更有效地评估潜在的漏洞功能。LineVD通过在训练过程中利用图神经网络和语句级信息,在真实世界的开源项目中实现了最先进的语句级漏洞检测。与最新的基于细粒度机器学习的模型相比,这一显著改进表明了直接利用语句级信息进行语句级SVD的有效性。最后,LineVD实现了合理的跨项目性能,表明即使对于完全看不见的软件项目,它也具有有效性和泛化能力。未来的方向将包括探索替代的预训练特征嵌入方法和新的GNN架构,这些架构可以更好地适应软件源代码的底层性质和漏洞。

Summary