目前基于深度学习的漏洞检测中性能最高的方法使用的是基于 token 的 transformer 模型,这对于捕捉漏洞检测所需的代码语义来说并不是最有效的。文中设计了一个受数据流分析启发的图学习框架 DeepDFA,以及一种能让图学习模拟数据流计算的嵌入技术。其训练时间为 9 分钟,比性能最高的基线模型快 75 倍。在仅使用 50 多个易受攻击和数百个示例作为训练数据的情况下,该模型保持了与使用 100% 数据集相同的性能。DeepDFA 还适用于 DbgBench 中的真实世界漏洞;它平均能从 17 个漏洞中检测出 8.7 个,并能区分已打补丁的版本和有漏洞的版本,而性能最高的基线模型未检测出任何漏洞。
1 问题描述
最近,基于大语言模型(LLMs) 的漏洞检测工具的性能已经超过了传统的静态分析,如 LineVul,最近一个基于 CodeBERT的模型,在一个常用的真实世界漏洞数据集上报告了高达 91 的 F1 分数。然而,LLMs 需要大量的训练数据和计算资源来进行训练和推理,现实情况中很难获得大量高质量的漏洞数据。它们也可能无法检测到训练数据集之外的漏洞;例如,性能最好的 Transformer 模型 LineVul 和 UniXcoder 无法检测到 DbgBench 中的任何真实漏洞。此外,由于只使用文本 tokens,这些模型可能无法有效地学习程序语义,例如程序值流路径、污点值的传播以及沿控制流路径的安全敏感 API 调用。
2 解决方案
文中探讨了数据流分析(DFA) 与 GNN 消息传递机制之间的类比关系,并设计了一种在 CFG 的每个节点上编码数据流信息的嵌入技术。具体来说,利用数据流的有效位向量表征来编码变量的定义和用途。在这种嵌入上进行图学习,可以传播和聚合数据流信息,从而模拟 DFA 中的数据流计算。利用这种方法,学习到的图表示法能更好地编码程序语义信息,如 reaching definitions。
基于此,论文开发了一种抽象的数据流嵌入,它可以将单个程序的变量定义映射到一个共同的空间,从而使模型可以比较和概括与跨程序漏洞相关的数据使用模式(数据流),并选择了一种图学习架构,其聚合和更新功能对数据流传播更为有效。
论文的主要贡献如下:
- 设计了一种抽象数据流嵌入,使深度学习能够概括不同程序的漏洞语义/数据流模式;
- 在程序的控制流图(CFG)和抽象数据流嵌入上应用图学习,模拟 reaching definitions 的数据流分析;
- 实现了 DeepDFA,并通过实验证明 DeepDFA 在漏洞检测的有效性、效率和对未见项目的泛化方面优于基线;
- 提供了有助于理解 DeepDFA 为何表现出色且高效的理由;
- 通过结合 DeepDFA 和 UniXcoder 超越了最先进的漏洞检测性能。
3 方法设计
DeepDFA的大致流程为:给定潜在漏洞程序的源代码,将其转换为 CFG,并使用设计的抽象数据流嵌入对节点进行编码。具体来说,考虑定义变量的数据类型、API 调用、常量和用于定义变量的操作符,并以紧凑、高效的方式对抽象定义进行了编码。然后应用图学习,像数据流分析中的 “合并 ”操作那样聚合来自节点的抽象数据流信息,并利用每个节点的信息进行更新。最后利用图学习直接传播数据流信息,向分类器呈现一种程序表示,其中编码了与漏洞直接相关的有用信息,从而实现高效、有效的漏洞检测。
3.1 一些基本概念
数据流分析的一个实例称为 “到达定义分析”(reaching definition analysis),它报告特定变量定义可以到达哪些程序节点。当 CFG 中存在连接定义和节点的路径,且变量在路径上未被重新定义时,定义就会到达节点。当识别出 NULL 指针的定义到达指针的取消引用时,到达定义分析就能根据其根本原因检测到空指针取消引用漏洞。同样,它也是检测许多其他漏洞(如缓冲区溢出、整数溢出、未初始化变量、双重释放和释放后重引用)的因果步骤。
DFA 使用两个等式,即满足算子(meet operator)和转移函数(transfer function),通过 CFG 中的相邻节点传播数据流信息。满足算子汇总来自其相邻节点的数据流集。传递函数利用节点 𝑣 中的可用信息更新数据流集。执行 DFA 的一种简单方法是 Kildall 方法。该方法将数据流信息迭代传播到 CFG 中 𝑣 的邻域,每次传播一步。当所有节点的数据流信息停止变化时,算法终止,表示为一个固定点。终止时,所有节点都将包含来自所有其他相关节点的数据流信息。当用于漏洞检测时,这些信息会与用户指定的漏洞条件进行比较,以确定程序中是否出现了漏洞。
图学习从初始节点表示开始,然后执行固定次数的消息传递算法迭代,在图中传播信息。初始节点表示通常是一个固定大小的连续向量,代表节点的内容。在每次迭代中,每个节点都会聚合来自其邻居的信息,然后更新其状态以整合这些信息。这两个步骤是通过聚合(AGGREGATE)和更新(UPDATE) 函数完成的,类似于满足算子和转移函数的两个数据流方程
3.2 抽象数据流嵌入
在数据流分析中,使用比特向量来表示每个节点的数据流集。位向量由 𝑛 个 0 和 1 位组成,其长度为域的大小。如果集合中存在相应的元素,位就会被设置为 1。在达到定义分析中,域由程序中的所有定义组成,如果相应的定义达到了节点,则位被置为 “1”。例如,在图 1 中,程序在节点𝑣1、𝑣3 和𝑣4 上包含三个定义,因此在分析开始时,达到定义分析使用一个位向量 [0 0 0] 来初始化每个节点。这个位向量表示数据流方程中的𝑂𝑈𝑇 [𝑣]。它在每一步传播过程中都会更新,当分析结束时,每个节点的位向量代表了能到达该节点的所有可能定义。
然而,在图学习中,不能直接使用定义的位向量作为节点嵌入。这是因为数据流分析和定义域都是针对程序的。换句话说,不同的程序有不同的变量定义;每个程序的位向量长度不同,元素(每个定义)也不具有可比性。因此,需要一个 “全局 ”定义集,用于指定不同程序的定义。
因此,文中将程序中的具体定义表示为抽象定义,并基于定义的属性设计了以下四种可涵盖漏洞攻击面的属性:
- API 调用(API call):用于定义变量的库函数或系统函数的调用,如 malloc 和 strlen。
- 数据类型(Data type):变量被赋值的数据类型,例如int、char* 和 float。
- 常量(Constant):定义中指定的常量值,例如NULL、-1 和硬编码字符串 “foo”。
- 运算符(Operator):用于定义变量的运算符,如 + 和 *。
通过分析大量程序(如训练集),收集前 k 个最常用的 API 调用、数据类型、常量和运算符,从而分别构建 4 个长度为 k 的全局定义字典。如图 3 所示,定义语句 (char*) str = malloc (10 * argc) 包含数据类型 char*,API 调用 malloc,常量 10 和运算符 *,它们在各自的全局定义字典中分别对应索引 3,2,4,4,将嵌入向量中相应位置置为 1,从而得到该定义语句对应的一个大小为 5*4 的向量,其中 k=5,而 4 表示一个定义语句中考虑的 4 个属性。 由于编码抽象数据流嵌入的向量大小固定,该嵌入方法可以扩展到数据集中的任何程序大小,而不会影响模型的效率。不同程序中的向量编码了定义的共同属性,因此模型可以捕捉不同程序的数据流模式。
3.3 图学习聚合数据流信息
图学习的大致流程为:首先构建程序的 CFG。其次,执行静态分析以识别 CFG 中的所有定义。然后,根据节点是否为定义,使用抽象数据流嵌入对 CFG 的每个节点进行初始化。应用图学习中的消息传递算法在整个 CFG 中传播数据流信息,与 Kildall 的方法类似。主要区别在于:(1) 文中传播的是 CFG 节点的抽象数据流嵌入;(2) 没有使用传递函数和相遇算子的数据流方程,而是交替使用了图学习中定义的 AGGREGATE 和 UPDATE 函数。采用门控图序列神经网络 (GGNN) 作为图学习的模型,其中 AGGREGATE 是多层感知器 (MLP),UPDATE 是门控递归单元 (GRU)。
DFA 中的满足算子和转移函数:
图学习中的聚合和更新函数:
当数据流信息到达 CFG 分支的合并点时,图学习会应用 AGGREGATE 函数。具体来说,在 GGNN 中,MLP 会计算多个相邻前代的表示的加权和,得到一个单一向量;这与 meet 运算符的功能相同。当数据流信息到达一个新节点时,图学习中的 UPDATE 函数会将当前节点的信息与前节点的 AGGREGATE 输出相结合,从而计算出下一个状态。具体来说,在 GGNN 中,GRU 会选择性地遗忘前一状态的部分内容,并整合当前节点和邻近状态的新信息,这与转移函数中执行的集合联合/差分与 GEN/KILL 类似。通过应用 AGGREGATE 和 UPDATE,初始嵌入将根据邻近节点的数据流信息进行更新,这与数据流分析的效果类似。
DFA 会迭代到一个固定点,从而在整个图中传播信息,而图学习会执行固定数量的迭代 𝑡,从而传播到距离 𝑡 的邻居。将 𝑡 设为验证集性能最大化的设置。最后,利用全局注意力池(Global Attention Pooling)将学习到的抽象节点嵌入结合起来,生成图层表示,并将其传递给分类器,以预测函数是脆弱还是非脆弱。AGGREGATE 和 UPDATE 函数是在训练过程中从标注数据中学习的,而不是像数据流分析那样使用固定公式。通过从训练示例中学习,分类器可以捕捉到代表各种类型漏洞的数据流信息模式,并选择相关的数据流信息进行漏洞检测。
在表 1 中,文中对图 1 中的 CFG 示例进行了达到定义分析,以展示数据流信息如何在图中传播,以及如何利用数据流信息进行漏洞检测。在第 0 次迭代中,4 个节点的初始状态均为 [ 0 0 0 ];在第 1 次迭代中,由于新定义,DFA 通过传递函数更新节点 𝑣1 [1 0 0], 𝑣3 [0 1 0], 𝑣4 [0 0 1]; 在第 2 次迭代中, 𝑣1(定义 𝑑1)通过 CFG 边传播到 𝑣2 [1 0 0],𝑣3(定义 𝑑2)传播到 𝑣4 [0 1 1];在第 3 次迭代中,满足算子用于合并 𝑣2 和 𝑣3,得到 [1 0 0] V [0 1 0] = [1 1 0];最终传递函数结合第 2 次迭代中的 𝑣4 [0 1 1],得到 [1 1 1]。DFA 算法结束后,节点的最终状态将用于检测漏洞。𝑣4 的状态是[1 1 1],这表明 𝑑1 和 𝑑2 都可能到达 𝑣4,这取决于程序值。由于定义 𝑑1 : str = NULL 可以到达 𝑣4 处的引用,可以得出结论,该程序存在空指针解引用漏洞。同样,在图学习中,经过固定次数的迭代后,所有节点的表示都会通过图读出操作结合起来,生成图级表示,用于漏洞检测预测。存在空指针解引用漏洞的程序将具有相同的抽象定义,其特征是使用 char* 类型和常量 NULL 来实现指针解引用语句。
4 实验评估
4.1 有效性
与非 Transformer 模型的比较:表 3(a) 显示,DeepDFA 在 F1 分数和召回率方面的表现远远优于基线模型。DeepDFA 的得分比所有基线模型的平均 F1 得分高出 47.51。与 Transformer 模型的比较: 将 DeepDFA 与 CodeBERT、LineVul、UniXcoder 和 CodeT5(评估过的转换器语言模型中最先进的)进行了比较。表 3b 显示,在 F1 分数上,DeepDFA 的表现要比 CodeBERT 和 CodeT5 好得多,而且运行之间的方差也最小(在所有模型中)。
4.2 效率
表 5 为 DeepDFA、LineVul 和 UniXcoder 的运行时间对比。与 UniXcoder 相比,DeepDFA 的训练时间缩短了 75 倍,在 GPU 上的推理速度提高了 2 倍,在 CPU 上的推理速度提高了 84 倍。在所有模型中,DeepDFA 的参数最少,相当于最小模型(ReVeal)的 67%,性能最高的基线模型(UniXcoder)的 0.3%。
在表 6 中,报告了 DeepDFA 与 SOTA 模型 LineVul 和 UniXcoder 相比,在训练数据集规模缩小后的性能表现。列 “#data” 和 “#vul” 列出了每个子集中训练示例的数量,以及其中易受攻击的示例的数量。结果表明,DeepDFA 在小数据集规模上保持稳定的性能,即使只有 0.1% 的训练数据集,只使用 151 个训练样例。相比之下,LineVul 和 UniXcoder 随着训练数据集的减小,性能下降。
4.3 泛化性
表 7 展示了在 Big-Vul 数据集上的跨项目评估的结果。
表 8 报告了将深度学习工具应用于实际缺陷基准的经验。在 3 次运行中,DeepDFA 在 17 个错误中的平均检测到 8.7,DeepDFA 还正确预测了 5 个打补丁程序中的 3 个不存在漏洞。
4.4 消融实验
表 9 显示了模型在 DbgBench 上的性能。与其他消融设置相比,该模型在使用所有四个特征时检测到的错误最多。当一次只使用一个特征时,模型总是漏掉 DeepDFA 检测到的 1-2 个错误。当同时使用三个特征时,模型仍始终未能检测到 DeepDFA 检测到的 1 个错误。
表 10 显示了模型在 Big-Vul 测试数据集上的性能。在所有配置中,DeepDFA(集成了所有四个特征)的表现最好。在一次测试一个特征时,数据类型本身的性能要好于单独测试其他 3 个特征,使用组合特征集时,模型的表现比只使用一个特征更好。
5 论文总结
文中提出了用于漏洞检测的高效图学习框架和嵌入技术 DeepDFA。通过抽象数据流嵌入利用数据流分析中的位向量思想,并整合了语义特征中的数据使用模式:常用的 API 调用、操作、常量和数据类型,这些都有可能捕捉到漏洞的成因。DeepDFA 使用类似的消息传递算法模拟数据流分析的 Kildall 方法。实验结果表明,DeepDFA 非常高效。它只用了 9 分钟就完成了训练,只用了 50 个漏洞示例就实现了最高性能。同时,它的性能仍然优于所有非 Transformer 基线,并且在所有模型中泛化效果最好。DeepDFA 在 DbgBench 的实际程序中发现了漏洞,而性能最高的基线 LineVul 和 UniXcoder 都没有检测到任何漏洞。重要的是,DeepDFA 与性能最高的模型相结合,其漏洞检测性能超过了目前最先进的水平。
6 原文链接
论文题目:Dataflow Analysis-Inspired Deep Learning for Efficient Vulnerability Detection
论文出处:2024 IEEE/ACM 46th International Conference on Software Engineering ( ICSE 2024 )
原文链接:Dataflow Analysis-Inspired Deep Learning for Efficient Vulnerability Detection | IEEE Conference Publication | IEEE Xplore
项目地址:Dataflow Analysis-Inspired Deep Learning for Efficient Vulnerability Detection (figshare.com)