VGX: Large-Scale Sample Generation for Boosting Learning-Based Software Vulnerability Analyses

VGX通过设计基于代码tokens之间的值流关系（value flow relationships）的绝对和相对位置编码的注意力机制，构建了一个基于Transformer的上下文化模型，并在大型代码语料库上对自定义的Transformer进行预训练，然后针对现有（task-specific 的）易受攻击代码位置数据集进行微调。然后，VGX利用从历史修复和人类对真实世界漏洞的知识中学习的编辑模式，在识别到的上下文中实现漏洞代码注入。

至于什么是值流，国内没有很规范的解释，国外对此有一定的应用，比如静态值流分析工具SVF ：https://blog.csdn.net/gitblog_00087/article/details/138841899 。指的应该是数据流

源码与数据集地址：https://zenodo.org/records/10443177

Background and motivation

一个直观的想法是：学习已知引入漏洞的代码编辑模式（即引入漏洞的编辑），然后将与之兼容的模式应用于正常程序，以生成一个易受攻击的版本。该Idea在大规模高质量漏洞数据生成方面仍然存在以下限制：

学习的漏洞注入代码编辑模式仅限于小规模的漏洞数据集
定位模型是在源代码作为自然语言token序列的基础上进行训练的，忽略了对于准确识别在给定程序中可能注入漏洞的位置这个至关重要的语义信息。原因在于代码漏洞是上下文敏感的：相同的代码行放置在一个代码上下文中时可能引发漏洞，但在另一个上下文中不会引发任何的易受攻击行为

Technical design

VGX包含两个主要阶段：学习/训练和生成。

在学习/训练阶段，VGX通过两个主要任务来建立：

训练一个基于深度学习的模型，能够识别给定（正常）程序中漏洞注入的代码上下文；
挖掘可以在被识别的上下文中实现注入的漏洞引入代码编辑模式。

在生成阶段，使用编辑模式和训练好的上下文模型，VGX先对给定的一组正常程序进行预处理（使用相同的模块）作为输入。然后，模型预测潜在的漏洞注入代码编辑上下文，并匹配并应用最适合的编辑模式（用于上下文），从而产生（预期的）易受攻击的样本。

Preprocessing

这个过程通过源码得到绝对和相对VFG子图以及源码对应的AST。

使用树解析器tree-sitter 将源代码转换为抽象语法树（AST）https://github.com/tree-sitter/tree-sitter，然后将AST线性化为文本（Transformer以文本作为输入），最后用特殊的标记符[SEP]将源代码和线性化的AST连接起来（有助于学习代码的语法和上下文结构）

对于绝对和相对VFG子图的提取，首先通过value flow analyse（比如使用SVF）得到VFG。对于某一变量的绝对VFG子图，从该变量的节点开始遍历VFG，直到无法找到新的节点为止，所生成的子图即为绝对VFG子图。两个变量（v1，v2）所对应的相对VFG子图，从v1的节点开始遍历VFG，直到达到v2的节点为止，所生成的子图即为相对VFG子图

Semantics-Aware Contextualization

给定一个带有Code+AST输入的正常程序，我们期望自定义的Transformer模型根据代码的语义和上下文输出一个可以被操作的代码片段，从而引入一个漏洞

Value-flow-based Position Encoding

自定义的Transformer的编码器中，核心模块是带有位置编码的自注意力块：

$z_{i}=\sum_{j=1}^{n}\frac{exp(a_{ij})}{\sum_{j^{\prime}=1}^{n}exp(a_{ij^{\prime}})}(x_{j}W^{V}+r_{ij}^{V}+r_{ij}^{V_{VFG}})$

$a_{ij}=\frac{1}{\sqrt{2d}}(x_{i}W^{Q})(x_{j}W^{K}+r_{ij}^{K}+r_{ij}^{K_{VFG}})^{T}+\frac{1}{\sqrt{2d}}(a_{i}^{Q})(a_{j}^{K})^{T}+\frac{1}{\sqrt{2d}}(a_{i}^{Q_{VFG}})(a_{j}^{K_{VFG}})^{T}$

$z_i$ ：第 $i$ 个token的输出隐藏表示
$a_{ij}$ ：第 $i$ 和 $j$ 个token之间的注意力
$x_i$ ：前一层中第 $i$ 个token的隐藏表示， $x_j$ 同理
$W^V,W^Q,W^K$ ：分别是value, query, key的权重矩阵
$r^V_{ij},r^K_{ij}$ ：传统的相对位置编码
$a^Q_{i},a^K_{j}$ ：传统的绝对位置编码
$d$ ：隐藏表示的维度

传统的位置编码存在一个问题：缺乏对代码的语义理解。即两个变量可能在值流上关系紧密，但仅仅因为位置相隔较远而导致其相对位置距离很大。解决方案是，利用预处理过程中获得的VFG子图，将语义信息合并到位置编码中。

即，基于值流图VFG计算相对位置编码 $r_{ij}^{V_{VFG}},r_{ij}^{K_{VFG}}$ ，也即上式中对应的变量。实现方案是，使用预训练的C语言FastText模型将每个节点中的变量名转换为一个嵌入。然后使用门控图神经网络（GGNN）进行消息传递聚合和更新节点嵌入。然后对图中的节点嵌入进行求和，得到VFG子图的嵌入。最后将，图嵌入乘以权重矩阵，以获得基于值流的位置编码 $r_{ij}^{V_{VFG}},r_{ij}^{K_{VFG}},a_{i}^{Q_{VFG}},a_{j}^{K_{VFG}}$

并不是所有的代码中的token都是变量。因此，我们只在token或者token对中的任一token是变量时计算基于值流的位置编码。否则，值流位置编码为零

Pre-Training

使用了来自CodeT5的三个已有的面向通用编程语言的目标，并引入了两个新的代码上下文化特定的目标（CAP和ISP）

三个预训练目标

遮蔽跨度预测（Mask Span Prediction，MSP），随机遮蔽输入中包含源代码和AST的15%源代码token，训练模型根据上下文恢复被遮蔽的token
标识符标记（Identifier Tagging，IT），预测每个源代码令牌是否为标识符
遮蔽标识符预测（Masked Identifier Prediction，MIP），遮蔽源代码中的所有标识符，然后训练模型根据代码语义恢复标识符

Code-AST Prediction (CAP)

在50%的预训练样本中将正确的AST分配给相应的源代码，而在剩下的50%中随机分配一个不正确的AST。训练模型被预测分配的AST是否与输入的源代码相对应。

Irrelevant Statement Prediction 无关语句预测 (ISP)

对于每个预训练样本，我们在随机位置插入来自另一个样本的随机选择的语句。在这种情况下，样本的功能不受插入语句的影响。因此，插入的语句在原始样本上是语义上无关的。我们预训练模型以识别并输出每个样本中的无关语句。

我们采用这个预训练目标是因为它类似于我们的漏洞注入定位目标。由于漏洞是基于代码语义的上下文敏感的，学习区分（语义上）无关语句（与相关语句）直观地帮助我们识别与注入漏洞（语义上）相关的正确代码上下文。

Fine-Tuning

为了解决微调数据的缺乏问题，我们通过对细调样本进行代码重构来进行 data augmentation。主要对微调样本应用了三种类型的语义保持重构：(1) 在if语句中反转条件并交换if和else块中的代码；(2) 将for循环转换为等效的while循环；(3) 在代码中插入由SaBabi 生成的不相关的垃圾代码。

Edit Pattern Formation

Pattern Extraction

首先将源代码解析为AST，一个编辑模式是由AST表示的一对代码片段，用于表示代码编辑。在模式提取之后，我们得到许多从非常通用（可以匹配和应用于许多不同的代码示例）到非常特定（只能匹配和应用于少数代码示例）的编辑模式。

Pattern Filtering/Ranking

每个提取出的编辑模式计算三个分数：

普遍性分数 ( $s_{preval}$ )：假设我们知道引入漏洞的位置，普遍性分数是正确在训练集中能够注入漏洞的样本数量，它与该模式能够成功注入漏洞的概率成正比。
特化分数 ( $s_{spec}$ )：在训练样本中，我们计算模式能够匹配的AST子树的平均数量。平均数量的倒数是特化分数，因为它表示模式是否具体，从而不会匹配不能注入漏洞的AST子树。
标识符分数 ( $s_{ident}$ )：标识符分数是模式中指定的标识符名称的数量，因为标识符名称对于正确注入漏洞非常重要，可以显著地提供代码语义信息。
三个分数的乘积是最终的排名分数。分数越高，模式成功注入漏洞的可能性就越大