Cross-domain vulnerability detection using graph embedding and domain adaptation

a山东省计算机网络重点实验室,山东省计算机科学中心(济南国家超级计算机中心),齐鲁工业大学(山东科学院),济南250014,中国b北京邮电大学网络空间安全学院,北京100876,中国c公共大数据国家重点实验室,贵州大学计算机科学与技术学院,贵阳550025

0 Abstract

漏洞检测是维护网络空间安全的有效手段。机器学习方法由于其准确性和自动化的优势,在软件安全领域引起了人们的广泛关注。然而,目前的研究主要集中在训练数据和测试数据属于同一域的域内漏洞检测上。由于应用场景、编码习惯等因素,不同软件项目中的漏洞可能服从不同的概率分布。当机器学习方法应用于一个全新的项目时,这种差异会影响它们的性能。为了解决这个冷启动问题,我们提出了一个使用图嵌入和深度域自适应(VulGDA)的跨域漏洞检测框架。它以多种跨域方式工作,包括零样本方式,即目标域中没有标记数据可用于训练。将VulGDA分解为图嵌入和域自适应。在图嵌入阶段,我们将源代码中的样本转换为图表示,其中元素根据其语法和语义关系直接连接。然后,我们将来自图中定义的邻居和边的信息聚合为实值向量。通过图形嵌入,VulGDA提取了全面的漏洞特征,并解决了长期依赖性的挑战。针对训练数据和测试数据之间的差异,使用域自适应来训练特征生成器。该特征生成器将嵌入的图映射到一个“深层”特征,该特征对漏洞检测具有鉴别性,并且对域之间的移动保持不变。我们进行了一项系统实验来验证VulGDA的有效性。结果表明,将图嵌入和深域自适应相结合,提高了VulGDA在跨域漏洞检测中的性能。与最先进的方法相比,我们的方法在冷启动条件下具有更好的性能。

1 Contributions

主要问题就是目前的研究针对某一数据集或是某些特定的漏洞,没有在跨域问题上提出适当的解决方法;

  • 提出了一个结合了图嵌入和域自适应的跨域漏洞检测框架VulGDA,VulGDA以ZeroShot方式工作,这是跨域检测中最严格的方式。VulGDA通过将目标域中的监督信息添加到训练数据中,也适用于few-shot和域内方式。

  • 提出了一种提取综合漏洞特征的图嵌入方法;减少噪声对漏洞特征的影响

  • 提出了一种神经网络来学习特征生成器。该特征生成器通过减少分类损失和域差异来进行优化。最后,该特征生成器将学习到的源域中的漏洞模式传输到目标域。

  • 实现了一个原型并进行了系统的实验。结果表明,图嵌入和域自适应的应用提高了VulGDA的性能。与现有方法相比,VulGDA在跨域漏洞检测方面取得了更好的性能。

2 Method

VulGDA旨在检测现实世界软件中的漏洞。为了实用,本文考虑了最严格的条件。也就是说,应用在源域中训练的检测模型来检测目标域中的漏洞,并且该模型在训练期间无法从目标域获得任何监督信息(零样本)。这一假设使VulGDA具有良好的泛化能力,并易于应用于其他跨域检测情况,例如历史漏洞数据很少(很少发生)或丰富(Indomain)的项目。

image-20231219170117488

问题定义:假设我们的检测架构可以访问源域中的标记数据集S={(xi,yi)}i=1n(Ds)nS=\{(x_{i},y_{i})\}_{i=1}^{n}\sim(D_{s})^{n} 和目标域中的未标记数据集 T={(xi)}i=n+1N(DTX)NnT=\{({x}_{i})\}_{i=n+1}^{N}\sim(D_{T}^{X})^{N-n}NN是训练样本的总数,DTXD_T^XDTD_TXX上的边缘分布。

最终目标是建立一个预测η:XYη:X→Y:最小RDT(η)=*Pr(xy)DT(η(x)y),R_{D_T}(η)=\operatorname*{Pr}_{(x y)\sim D_{T}}(\eta(x)\neq y),,而DTD_T中没有监督信息。

VuLGDA有两个阶段组成:

图嵌入:

图嵌入阶段的目的是将源代码中的漏洞样本转换为可由机器学习模型处理的实值向量。因此,通过句法和语义分析,图嵌入阶段首先将样本转换为CPG,即图结构中的中间表示。在CPG中,语法和语义相关的元素被直接连接在一起,这缓解了长期的依赖问题和域之间的分歧。然后,利用预训练的单词嵌入生成标记嵌入和节点嵌入。最后,我们使用GGNN通过传播和聚合来自CPG中定义的邻居的信息来获得嵌入向量。

image-20231219174742348

域自适应:

域自适应阶段旨在学习在域之间转换分布的特征生成器。提出了一种由特征生成器、分类器和瓶颈组成的神经网络来学习该特征生成器。将源域中的特征作为输入的分类器产生分类损失。瓶颈度量源域和目标域中的特征之间的差异。通过最小化分类损失和领域差异来优化特征生成器。检测结果由生成的“深层”特征训练的分类器报告。经过上述阶段后,VulGDA可以在目标域中检测跨域漏洞,而无需标记数据。

本文的创新点主要在于领域自适应:

传统的机器学习方法假设训练数据和测试数据服从相同的特征分布。然而,由于编程风格和应用场景的不同,不同领域的漏洞特征可能服从不同的分布,从而影响模型的检测性能。在本节中,使用来自源域和目标域的数据,我们构建了一个特征生成器G f:XG→XF。G f将嵌入XG的图转换为“深度”特征向量XF,该向量对漏洞检测具有判别性,并且随着域之间的移动而不变。因此,目标函数定义为:

f=arg*min1Bi=1B(f(xi),yi)+λR(BS,BT)f^{*}=\arg\operatorname*{min}\frac{1}{B}\sum_{i=1}^{B}\ell(f(x_{i}),y_{i})+\lambda R(B_{S},B_{T})

其中B表示批次大小,(f(xi),yi)表示分类损失,BS和BT分别在源域和目标域中为2个批次,R(∗,∗)测量域差异。λ用于调整分类损失和域差异的权重。

image-20231219222546132