Path-Sensitive Code Embedding via Contrastive Learning for Software Vulnerability Detection
Path-Sensitive Code Embedding via Contrastive Learning for Software Vulnerability Detection
foresta.yang0 Abstract
为了获得代码的结构信息，当前的学习方法通常将程序抽象成图的形式（例如，数据流图，抽象语法树），然后基于安全和易受攻击的代码片段的（子）图来训练底层分类模型以进行漏洞预测。然而，这些模型仍然不足以精确检测缺陷，因为这些模型的目标是产生分类结果，而不是理解漏洞的语义，例如，关键的漏洞触发路径，这对于静态漏洞检测至关重要。本文提出了ContraFlow，这是一种选择性但精确的对比值流嵌入方法，用于静态检测软件漏洞。**ContraFlow的新颖之处在于使用自监督对比学习从预训练的路径嵌入模型中选择和保留可行的值流（也称为程序依赖）路径，从而显著减少了训练昂贵的下游模型进行基于路径的漏洞检测所需的标记数据量。**我们使用288个真实项目评估了ContraFlow，比较了八种最近的基于学习的方法。ContraFlow在信息度、标记度和F1得分方面的表现优于这八个基线方法，最高分别提高了334.1％、317.9％和58.3％，而在定位有缺陷的语句方面，ContraFlow的平均语句召回率、平均语句精度和平均IoU方面的改进分别最高提高了450.0％、192.3％和450.0％。
1 Intro or Overview
1.1 Problem and Challenge
Existing Efforts and Limitations. 最近提出了代码嵌入，旨在通过分布式向量表示来表示代码语义，用于源代码分析和错误检测。
- 最初，嵌入方法将程序视为文本标记[49-51]，通过应用自然语言处理技术来学习代码语义，而不需要代码结构信息。后来，几种方法[7, 10, 48, 81]通过保留结构信息（例如，通过程序依赖图）改进了嵌入结果，然后使用图神经网络（GNNs）[41, 47]来分类代码片段的（子）图是否易受攻击。
- 尽管学习代码的图表示可以用于代码分类或摘要任务，但对于基于路径的漏洞检测仍然不足。 这是因为输入图表示不区分程序路径，而后端GNNs无法识别程序路径。 图特征是从GNNs中所有连接节点对之间的消息传递中学习的，但不幸的是，缺乏任何可行/不可行的值流（程序依赖）路径的知识。but unfortunately, without the knowledge of any feasible/infeasible value-flow (program dependence) paths.
- 因此，这些预测模型并不知道潜在的错误路径，这些路径显示了错误的产生和触发方式。这是静态错误检测的主要目标之一：帮助从业者快速定位并修复报告的漏洞。
Insights and Challenges. 为了解决上述限制，检测方法需要基于精确的学习模型，该模型能够保留价值流路径，而不是整个图，该图无法区分可行/不可行的程序依赖路径。受到词嵌入中令牌袋的概念的启发，一些最近的代码嵌入方法对抽象语法树（ASTs）或值流图（VFGs）上的路径进行了嵌入以进行代码分类和摘要 [62]。这些方法随机抽样一小部分路径以产生它们的嵌入向量，然后聚合它们形成代码片段的最终表示。然而，这些方法不能直接用于诸如基于路径的错误检测等复杂任务，因为可能存在需要嵌入的无界程序路径的数量。
基于路径的模型的有效性在于路径选择策略。**识别和保留个别可行路径而不是通过随机抽样聚合不可行或与错误无关的路径是具有挑战性但重要的，**以避免在嵌入过程中出现不精确。这需要在模型训练过程中选择性地学习具有判别特征的路径，这些特征在 bug 语义中起作用，以产生用于基于路径的漏洞检测的精确嵌入。
1.2 Motivation
图2通过使用从真实项目POCO（一个用于网络应用程序的库）[31]提取的业务逻辑错误（CWE840）[55]，沿着图1中的三个阶段，阐明了ContraFlow的关键思想。漏洞是由于在rebuild_list(&hd)之后调用set_status(&hd)时API误用引起的，其中hd首先在第2行定义，然后在第6行修改，并在第13行使用。hd的这种错误的值流路径可能导致意外行为并导致服务拒绝。
注意，从原始代码片段中提取的不同变量的值流路径很大，并且包含许多路径，包括用于可行性检查和嵌入的不可行或与错误无关的路径。我们在阶段（a）中的对比值流嵌入首先对VPE进行预训练，以在潜在空间中保留路径（例如，π1 − π4）的语义，然后使用主动学习在阶段（b）中选择最具代表性的路径，然后进行可行性检查，通过稀疏和受控的值流分析移除不可行的路径π2和π3。阶段（c）进一步微调并解释π1作为训练模型的排名注意分数（π1的90%）中可能存在错误的路径。
(a) Contrastive Value-Flow Embedding.
如图2（a）所示，该阶段的输入是从代码片段中提取的一组值流路径袋，例如，π1（○2 → ○6 → 1○3），π2（○3 → ○9 → 1○3），π3（2○ → ○6 → 1○5）和π4（○3 → ○9 → 1○5）。我们将它们两次馈送到值流路径编码器（VPE）中，使用VPE中的不同dropout掩码[61]，以获得它们的向量表示，例如，vπ1，vπ2，vπ3和vπ4，以及它们的对应对比表示[26]，例如，v+π1，v+π2，v+π3和v+π4。VPE是使用对比学习进行预训练的，以捕获价值流路径的语义，使得预训练的相似嵌入向量（例如，vπ1和v+π1）彼此保持接近，而不相似的对（例如，vπ1和v+π3）则相距较远，如图2(a)所示的二维特征空间。
VPE的参数在反向传播过程中通过最小化NCE损失[9]来自动更新[33]，该损失编码了价值流嵌入向量之间的相似性。
(b) Value-Flow Path Selection.
对value-flow path抽样，路径可行检查
该阶段使用从阶段（a）预训练的 VPE 将价值流路径转换为嵌入向量。之后，我们根据从自监督主动学习中学到的排名，对代表性的价值流路径进行抽样，例如，π1 − π4。这些路径进一步被输入到路径可行性检查中，以删除不可行的价值流路径。例如，路径 3○→○9 → 1○3 是不可行的，因为在 3○→ 9○ 和 ○9 → 1○3 处的控制流保护符 !FLG 和 FLG 相互矛盾。同样，○2 →○6 → 1○5 也是不可行的。最后，只有可行的价值流路径 π1 和 π4 被保留用于训练阶段（c）的检测模型。
© Detection Model Training.
该阶段的输入是由阶段（b）产生的选择的可行且具有代表性的值流路径，这些路径首先使用从阶段（a）转移的VPE模型转换为向量。然后，这些嵌入向量通过一个transformer架构[69]生成每个值流路径的上下文向量。例如，通过与其他向量（例如，π4）进行关注来计算π1的上下文向量，以增加它们对π1的影响。之后，应用软注意力层[1]将这些上下文向量合并为一个向量，用于训练检测模型。
排名注意力权重指示了不同值流路径对模型输出的贡献。例如，值流路径π1（○2→6○→1○3）的注意力权重最高（90％），而其他路径可以忽略不计，表明该值流路径可能是一个有错误的路径。
1.3 Contribution
2 Architecture & Method
2.1 System Overview
(a) Contrastive Value-Flow Embedding.
该阶段旨在使用对比学习训练价值流嵌入模型，值流路径编码器（VPE）。
- 给定一组从未标记的源代码中使用现有静态分析器SVF提取的值流路径。
- 首先执行数据增强以生成对比值流表示[26]，
- 然后利用标准的噪声对比估计（NCE）损失函数[9]来最大化语义上相似的值流路径向量之间的一致性。
这更新了我们的VPE参数，以促使其保留价值流路径的深层语义。预训练的VPE在接下来的两个阶段中使用。
(b) Value-Flow Path Selection.
该阶段旨在精确选择可行且代表性的值流路径，以代表代码片段以支持基于路径的检测模型的快速训练。
- 首先使用来自阶段（a）的预训练VPE生成输入路径的特征向量，并使用自监督主动学习[45]对路径进行采样，以捕获最具代表性的路径并使嵌入多样化且信息丰富。
- 然后，我们通过对带有注释的值流图（VFG）[12, 64]上的可达性问题进行路径敏感的代码嵌入来执行路径敏感的代码嵌入。
- VFG以稀疏的方式捕获def-use关系，并使用描述控制流传输条件的注释保护边缘。然后，可行性检查被简化为在受保护的VFG上的可达性问题，以仅在低维嵌入空间中嵌入可行路径。
© Detection Model Training.
给定由阶段（b）产生的选定路径和从阶段（a）转移的VPE模型，本阶段将通过仅使用程序的选定路径及其标签（即易受攻击或安全）来训练精确的检测模型。我们首先为每个选定的值流路径获取来自VPE的嵌入向量，然后利用transformer架构[69]为每个路径生成上下文向量以捕获路径之间的交互。然后，这些向量被馈送到软注意力层[1]以对它们进行评分和聚合，形成最终的检测模型训练的一个向量。该模型还可以根据它们对模型输出的贡献来解释重要的值流路径和语句，这些贡献是由学习到的注意力分数排名的。
Contrastive Value-Flow Embedding.
对比性值流嵌入旨在通过预训练 VPE 从未标记的代码片段中学习相似/不相似的受保护值流路径 π 的区分性向量表示 vπ。受保护值流路径 π 包括一系列程序语句，表示变量之间的 def-use 链，每个语句之间的边上的guard用于指示控制流转移条件 [12, 64]。这些guard将在阶段 (b) 中的路径可行性求解过程中使用。算法 1 总结了学习算法。对于每个学习时期，我们生成对比向量表示（第 2 行），并计算对比值流路径之间的对比损失（第 3 行）。VPE 的参数将在训练过程中自动更新（第 4 行）。以下段落描述了对比性值流表示和对比损失函数。
