USENIX Security 23 论文解读 # VulChecker: Graph-based Vulnerability Localization in Source Code

• Mirsky Y, Macon G, Brown M, Yagemann C, Pruett M, Downing E, Mertoguno S, Lee W. “VulChecker: Graph-based Vulnerability Localization in Source Code”, USENIX Security 23
• 论文发表于：Usenix Security 2023 (CCF-A，安全顶会）

（结合一些背景知识，似乎可以理解为这篇工作真正让 AI 在 SAST 里大规模落地成为可能，并且超越了一定程度上传统的非 AI 的 SAST 的能力，该文章值得好好精读一下，略去了一些实验结论的分析和细节以及附录内容，这部分还是看原文吧）

Abstract

【研究问题】目前基于深度学习的漏洞检测更多是在函数或者更大粒度的区域中搜索特定bug，但无法分类和识别漏洞发生的行号。
【解决方案】提出 VulChecker：一种可以精确定位源代码漏洞（精确到指令）并对其 CWE 类型分类的工具。

• 提出了一种新的程序表示、程序切片策略(program slicing strategy)然后使用消息传递图神经网络(message-passing graph neural network)来利用代码语义， improve the reach between a vulnerability’s root cause and manifestation points。
• 新的数据增强策略：使用在线提供的免费合成样本，低成本快速创建强大的数据集，用于在野漏洞检测。
  =》通过这种训练策略，VulChecker 能够在 19 个项目中识别出 24 个 CVE（2019 年和 2020 年分别有 10 个），与只能检测到 4 个的商业工具相比，误报率几乎为零。同时还发现一个新的可利用的 0day。

1 Introduction

• SAST FPs 很多，主要原因是：1）规则/模式匹配的局限性和不完整。2）不可访问代码中的漏洞。
• SAST FNs 也很多：难检测 complex and contextual vulnerabilities。

规则/匹配算法的局限性：在某些上下文中，integer overflows andunderflows 是合法的（比如加密）但在某些上下文中不合法（内存访问）。Use-After-Free 漏洞是高度上下文敏感的。

1.1 Recent Efforts

为了解决这些问题，目前采用了深度学习的研究，比如 CNN/RNN/GNN 等，通过将软件表示为 annotated graph (G) 来编码 data-flow(DFG)、control-flow(CFG) 或者 abstract syntax (AST)。虽然这些取得了不错的进展，但不实用（尤其是还有误报的情况），原因如下：1）预测粒度太大：entire functions or code regions；2）漏洞分类：没有对漏洞类型分类
【我理解的应该是定位的同时分类漏洞类型？就是现有的方法告诉了这个区域有漏洞，但不告诉是什么漏洞以及这个漏洞具体在这个区域的哪个位置，尤其是还有误报的情况很难用，还是需要人工代码审计 ==> 这样的话那 VulChecker 这项工作很有价值啊，相当于让 AI 辅助的 SAST 真正走向实用了】；

1.2 Insights

当前研究的问题和瓶颈
Broad Program Slicing：目前通过将G中的子图 $G_i$ 作为观测值用于对漏洞进行定位的预测。$G_i$ 要么是一个完整函数，要么是程序切片（G 中提取的 point of interest (PoI) 的一段连续区域）。在 $G_i$ 上的预测可能涉及数百行代码，其中可能存在许多潜在的 root cause and manifestation points，这让优化变得困难。
Insight 1: To increase the precision of a prediction, the observation must relate directly to the location of the vulnerability. 【观测值要强调漏洞】

Incomplete Code Representations：为了将 $G_i$ 转化为机器学习可以理解的形式，从 AST 和/或函数派生的符号要么被压缩，要么用 Word2Vec 等方法来对符号嵌入做预处理步骤。这种表示软件的方法与 ML 是兼容的，但不是最优的，这样没有办法对指令、依赖关系语义进行推理。来自 AST 的信息要么1）omitted,2）included in manner that inhibits graph-based learning or (3) is included incorrectly in a manner that harms code semantics and model efficiency。
Insight 2: To enable effective graph-based learning in this task, a new code representation that properly incorporates the flow of AST information is required. Furthermore, instruction level semantics should be explicitly provided in G to enable efficient and effective learning.

Manifestation Distance：漏洞的根因和表现点中间可能相隔数百行代码，当下 GNN 模型只能 propagate a node’s information one or two steps away。因此，这些 GNN 模型无法推断潜在的根因是否直接影响潜在的表现点。Instead, these models are limited to identifying local patterns around either point (e.g., a missing guard branch prior to a memcpy is enough to raise an alarm).
Insight 3: To increase the model’s ability to identify vulnerabilities, the model must be able to identify causality across the entirety of $G_i$.

The Lack of Labeled Data：DL 需要大量监督数据用于训练，但在指令级别/行号级别标记源代码中的漏洞很困难且成本高昂，因此不存在 fine-grained datasets。因此当前的工作要么（1）使用合成数据集 synthetic datasets (e.g., NIST SARD [26])，要么（2）使用只标记代码区域(程序切片)或整个函数的数据集【标签粒度很粗】。
Insight 4: To enable cost effective high precision detection in source code, there is a need to develop data augmentation techniques to efficiently expand existing datasets.

Level of Program Representation：源码中一行可能就有许多相互关联的 high level operation，但漏洞通常与原子和机器级指令语义(例如，越界缓冲区访问、整数溢出)相关。大多数之前的相关工作在源码级别提取 G，因此面临从抽象程序表示中识别 low-level 漏洞的挑战。
Insight 5: To improve the model’s ability to reason about vulnerabilities that have atomic machine-level characteristics, G should capture instruction semantics at lower levels rather than at source code.

1.3 The Proposed Solution

VulChecker 首先通过自定义 LLVM 编译器工具链通过程序源码生成 GNN 优化的程序表示，即 enriched program dependency graph (ePDG)。ePDGs 中的节点表示机器级原子指令，边表示指令直接的控制流、数据流依赖关系。

为了对漏洞进行定位，VulChecker 首先搜索 ePDG 中的潜在表现点（比如对应于 CWE-415 的 free() 函数）。然后从该点向前切分子图。当该子图被预测为 vulnerable 时，VulChecker 向开发人员报告源码表现点的具体位置（行号和指令），并提供 CWE type。
最后，为了缓解获取细粒度标记数据集的问题，提出了一种数据增强技术。

• The approach is to generate positive samples by injecting synthetic traces into projects taken from the wild.The approach is cheap because the synthetic traces are taken from open sourced labeled datasets like the Juliet C/C++ Test Suite [26] and G is manipulated directly to avoid compilation issues.
• The approach is effective because we inject root causes at random distance away from the manifestation point. This improves generalization by forcing the model to search deeper into $G_i$ over real code.

本文在五种不同的 CWEs 上评估了 VulChecker：integer overflow (CWE-190), stack overflow (CWE-121), heap overflow (CWE-122), double free (CWE-415), and use-after-free (CWE-416)，这些 CWE 涵盖了空间和时间上的内存安全问题。仅在增强数据集上训练的 VulChecker 能够检测 24 vulnerabilities (reported CVEs) in the 19 C++ projects we evaluated，这明显高于基线 SAST 工具（商业工具仅检出 4 个）。
人工分析了在 augmented 和 CVE 数据上训练的模型的 Top100 结果，发现 VulChecker 在前 50 个结果中命中率（精度）为 50-80%，VulChecker was also able to detect a previously unknown vulnerability (zero-day)。

总体来说，VulChecker 较为实用：1）它能够在大型项目（例如 libgit2 有超过 300 个文件、110k 行代码以及 1800 万条指令）上运行，以及 2）误报率相对较低。

1.4 Contributions

• the first deep learning framework that can both detect vulnerabilities in source code with instruction and line-level precision and classify the type of CWE；使用 VulChecker 的开发人员不需要搜索数百行来找到 vulnerable code 和/或确定存在哪种类型的漏洞；
• 引入了消息传递神经网络用于漏洞检测任务，以学习和识别PDG中漏洞根本原因与表现点之间的关系；并且利用 struct2vec 给 ePDG 中的 data-flow edge 赋予特征；
• 提出 ePDG：a novel graph-based machinelearning optimized program representation that avoids the inefficiencies and drawbacks of others in prior work；
• a novel data augmentation approach: helps mitigate the lack of fine-grained datasets for GNN-based source code vulnerability detection, while still generalizing to real world software projects taken from the wild. The approach is simple and easy to implement, making it a practical way to boost the performance of future models in this research domain.

本文代码、模型、数据集均开源！

2 Related Works

2.1 Model

• 该领域早期工作主要聚焦在 CNN/RNN on linear (sequential) representations of the software’s data-flow graphs。这种 linear reductions 效果在 long-distance 的情况效果不好。
• 后续采用 GNN，比如 GCN，但 GCN do not recall (i.e., propagate) a node’s original feature vector at each iteration，并也处理不好 long-distance，后续工作用 GRNN leveraged a recurrent layer to recall information passed to neighboring nodes at previous iterations，in these networks the number of layers dictated the number of propagation iterations, which was only one or two。

VulChecker 基于 struct2vec，a message-passing neural network，可以在不需要额外层的情况下执行多次迭代。因此，VulChecker 更适合漏洞检测，可以识别远距离根因和表现点之间的联系。

2.2 Program Representation

之前的程序表示有：CFG、DFG、组合 CFG/DFG 的 PDG(program dependency graphs)，这些表示可以很容易地从源代码生成（例如 Joern [37] ），但没有明确捕获指令语义。更高级的表示还有 code property graphs (CPG)、natural code sequence CPGs (ncsCPG)。

• CPGs merge the source code’s PDG and AST by attaching each source line’s AST subtree under the corresponding node in the PDG. This indirectly captures instruction semantics along with control- and data-flow dependencies in a single structure.
  • 缺点：not well structured for GNN learning models since information from the AST cannot flow across the graph；
• ncsCPG is a CPG in which the leaf nodes in each AST subtree are superficially linked to leaf nodes in the preceding or successive trees, as visualized with orange edges in Figure 1.
  • 优点：enables information to flow from the AST during GNN learning.
  • 缺点：however it is suboptimal because semantically unrelated source code lines become linked, even if they share no data or control dependencies.

最后，使用序列模型（即 CNN 和 RNN）的系统将代码表示为从上述图形形式之一进行提取或展平得到序列。无论采用何种程序表示，先前的工作都依赖于特征提取来 compress and express the contents of each token or node in G，比如 one-hot、word2vec、doc2vec 等。这些表示的问题是：1）$G_i$ 中的节点在单行有多个操作，在语义精度有损。2）pre-processd embedding 阻止了模型学习最佳表示来优化学习目标（i.e., classifying vulnerabilities）。但本文采用的 ePDG 明确定义了 node/edge feature 来编码相关信息(e.g., operations）且避免了表面特征的影响（eg.变量名）。

2.3 Application

有一些相关工作在二进制层面挖洞：Gemini、DeepBinDiff、ASM2Vec、Bin2Vec 采用相似的漏洞检测方法，用 compiled vulnerable code snippets 去检查 binary。但这些方法效果有限：1）模型搜索特定漏洞，而不是一般的(CWE)、2）需要反汇编、3）这些方法识别的漏洞指示代码区域而不是特定行。

源码层面挖洞的相关工作：表1中列出的相关工作不能直接识别漏洞的行号，因为其$G_i$的表示没有锚定特定的代码行【分析的粒度较大，无法精确定位】。一个例外的是和 VulChecker 同期设计的 VulDeeLocator 方案，这项工作首先解析程序 AST，然后在源代码中找到 PoIs 并标记所有的 library/API function calls, array definitions, pointer definitions and arithmetic expressions。每个标记的 PoI 后续 traced down 到 lower level 的代码表示(LLVM IR)。然后在 PoI 附近同时对前向/后向的程序进行切片，并展平成一个 900 tokens 的 sequence(e.g.,call, void, @, FUN1, ‘(’, …)。序列中的每个 token 都用 word2vec 计算得到一个向量，然后把向量序列输入到 BiRNN 中预测哪一行具有漏洞。与先前的工作相似，VulDeeLocator 无法指出检测到的漏洞类型。此外，从 AST 中选择 PoIs 不适合 spatial and temporal memory safety violations。相比之下，基于 IR 的 ePDG 表示更接近 machine code 并 resolves ambiguities regarding data types, temporary variables, and storage locations。VulDeeLocator 删除循环并将代码平化为一个序列，减少了其中蕴含的代码的语义和模式。VulChecker 在 IR 级别执行分析，某种意义上实现了语言无关。虽然本文只评估了C/C++系统，但 LLVM 也可以 lower 其他语言，比如 Objective-C, Fortran, and Rust。VulChecker 在这些系统很可能也适用。

与目前 SOTA 的技术相比，VulChecker 可以做到：
1）在 line-level、instruction-level 定位漏洞；2）对漏洞类型进行分类；3）使用传递消息的 GNN 深入程序切片，可以更好地将根因和表现点联系起来；4）可以推广到更广泛的编程语言；

3 VulChecker

Overview：对每个 CWE，VulChecker 训练一个单独的模型并使用不同的抽样策略来抽取 observation（即 potential manifestation points）。VulChecker 的每个 CWE pipeline 均有如下四步：(1) ePDG generation, (2) sampling, (3) feature extraction, and (4) model training or execution。

1.ePDG Generation：对于给定程序的源码(记为 S)首先编译为 LLVM IR 并应用几个优化。接下来怼我们感兴趣的特定 CWE 类型使用了自定义 LLVM plugin 来分析 IR，标记potential root cause and manifestation points，然后根据 S 产生 ePDG(记为 G)。
2.Sampling：扫描 G 来定位在 ePDG 生成期间标记的给定 CWE 的潜在表现点。manifestation point（表现点）是图 G 中的体现对于 CWE 的 node m(instruction)（比如 stack overflow 中的 stack memory writes 操作）。对于第 i 个 manifestation point in G，从该点向后切一个 subgraph $G_i$ 直到一个给定的深度。$G_i$ 代表 VulChecker 预测是否存在漏洞的观察结果。
3.Feature Extraction：采用与 G 相同的结构，创建了 annotated graph $G^{}$