Trvd Enhancing vulnerability detection via AST decomposition and neural sub-tree encoding

0 Abstract

软件漏洞的爆炸性增长对系统安全构成了严重威胁，已成为当今亟待解决的问题之一。然而，现有的漏洞检测方法在达到检测准确性、效率和适用性之间的平衡方面仍然面临局限性。遵循分而治之的策略，本文提出了TrVD（基于抽象语法树分解的漏洞检测器）来揭示源代码片段中隐含的指示语义，以实现准确高效的漏洞检测。为了便于捕捉细微的语义特征，TrVD使用一种新的分解算法将代码片段的AST转换为大小和深度受限的有序子树集。因此，通过精心设计的树结构神经网络，可以有效地收集每个子树的语义。最后，使用Transformer风格的编码器将所有子树的长程上下文语义聚合到一个特定于漏洞的向量中，以表示目标代码片段。在由不同的真实世界和合成漏洞样本组成的五个大型数据集上进行的广泛实验证明了TrVD在检测漏洞存在和确定漏洞类型方面相对于SOTA方法的性能优势。消融研究也证实了TrVD核心设计的有效性。

1 Intro or Overview

1.1 Problem and Challenge

可用性

其他被广泛采用的代码表示包括CFG、PDG和各种基于图形的变体。这些表示更明确地描述了代码元素之间的控制或数据依赖关系，然而，当面对不可执行或不完整的代码片段时，很难精确推导出这些依赖关系。因此，它们可能并不总是适用于漏洞检测。按照约定，AST可以很容易地为任何代码片段构建，例如文件、函数或单个语句。

效率

与需要相对复杂和耗时的控制或依赖性分析的代码表示（例如CFG、PDG和代码小工具）相比，从代码构建AST要简单得多，重量轻，从而有助于提高整个检测方法的效率。

语义综合性

那些人工创建的代码表示（例如，PDG和XFG）倾向于强调代码的特定方面，例如控制流或数据依赖关系。然而，它们经常在转换过程中丢失一些重要信息，这会导致语义损失，尤其是在表示不完整的代码片段时。不同的是，AST使源代码具有高度结构化的性质，其中关于语句和表达式的底层语法是直接可用的；也就是说，AST提供了更全面、更丰富、更精确的代码语义，使TrVD不遗漏任何可疑的漏洞含义，提高了检测的准确性。

1.2 Motivation

如前所述，为了充分利用AST的潜力并有效应对其使用带来的挑战（这对漏洞检测能力产生了重大影响），TrVD遵循了经典的分而治之策略。

Divide

使用基于图/树的神经网络编码整个AST，如GCN、GAT、tree LSTM和TBCNN，在处理大型/深层树结构时，使用过平滑嵌入捕获长程依赖性可能会大大降低。由于资源限制，这些耗费时间和内存的模型直接处理这样大/深的AST也是非常不可行的。在这方面，TrVD选择通过分解算法将整个AST划分为有序的子树集，每个子树对应于原始代码片段中的细粒度完整代码单元。这产生了两个优点：（1）每个子树包含有限（小得多）数量的节点和可控的树深度，并且可以由通用树或图神经网络以成本低廉的方式进行操作；以及（2）更重要的是，这提供了一个机会，通过关注大小受限的本地代码单元来更好地掌握指示漏洞信息的微妙语义，否则可能很容易被忽视。

Conquer

TrVD采用“先综合获取后关键点聚焦”的方案，实现对指示性和可疑漏洞语义的有效提取。具体来说，TrVD首先用新设计的树结构神经网络处理每个子树，将其对应代码单元的语义映射到数字向量中。基于所学习的子树嵌入，TrVD然后利用基于Transformer的主干来不同地关注越来越不重要的子树，以发现具有自注意机制的漏洞模式，并将所有子树的长程上下文语义融合到密集的、特定于漏洞的数字向量中，以表示目标代码片段。从这个意义上说，TrVD在代码表示学习中的方式类似于基于切片的方法，其中通过根据兴趣点对相关语句进行切片而生成的代码小工具（例如，库/API函数）通过常规NLP模型（例如LSTM）容纳有助于学习局部特征和帮助精确定位漏洞模式的信息。但是，在TrVD中细化的子树更结构化，不那么模糊，而为注意力增强子树聚合而设计的Transformer能够更好地嵌入上下文，以提高漏洞检测性能。

1.3 Contribution

  • 我们提出了一种新的漏洞检测方法，称为TrVD，在每个公式化阶段都经过精心设计，使其成为一种准确、高效、更实用的方法，适用于代码片段。具体而言，为了提高其检测隐藏良好的漏洞和特定漏洞类型的能力，它结合了新的AST分解、注意力增强子树编码和上下文感知语义聚合，从而能够从目标代码片段中提取微妙但特定于漏洞的语义，据我们所知，在基于DL的漏洞检测工作中尚属首次。
  • 我们进行了大量的实验来评估TrVD的性能、不同组成模块的有效性以及运行时开销。实证研究表明，在大多数数据集上，TrVD在检测漏洞的存在或识别特定漏洞类型方面的准确性、F1和精度优于最先进的基于DL的方法。TrVD核心设计的优势，如AST分解和子树编码，也通过消融研究得到了证实。
  • 我们提供了一个新的数据集来促进漏洞检测研究。该数据集包含264822个C/C++函数，每个函数都标有特定的CWE ID或非易受攻击的基本事实。TrVD实现的数据集和源代码已在https://github.com/XUPT-SSS/TrVD，以促进未来的基准测试或比较。

AST分解，注意力增强子树编码，上下文感知语义聚合，

新数据集

2 Architecture & Method

2.1 System Overview

image-20240304203524892

目标源代码片段依次通过五个主要模块进行处理：

（1）AST生成器，它以语义保持的方式对代码片段进行规范化，并使用解析器构建其AST；

（2） AST分解器，使用分解算法将AST转换为有序的子树集；

（3） 综合语义收集器，通过树结构的神经编码器，将每个子树中传递的语义迭代提取为实值向量；

（4） 可疑语义焦点，它将所有收集到的子树语义聚合到一个上下文化的嵌入向量中，该嵌入向量通过基于Transformer的模型处理更具漏洞特定性的语义；

（5） 漏洞检测器，它实现了一个多层感知器（MLP），该感知器具有用交叉熵损失训练的softmax层来预测输出。

AST decomposer

单个令牌序列。

遍历算法通常用于将AST展平为记号序列，包括前序遍历（Tang，Shen et al.，2022；Zhang，Wang，Zhang，Sun，&Liu，2020）和中序遍历（Svyatkovskiy，赵，Fu，&Sundaresan，2019）。为了避免信息丢失，SBT（Hu，Li，Xia，Lo，&Jin，2018）和SPT（Niu et al.，2022）通过添加额外的符号来指示父子关系，进一步改进了这些方法，以确保线性化的序列可以明确地映射回AST，然而，这会使序列变得更大。另一种方法code2seq（Alon，Levy，&Yahav，2019）将叶节点之间收集的采样路径连接起来，形成单个令牌序列，该序列通常太长，神经编码器无法有效处理语义提取，尤其是漏洞检测任务。

令牌序列集。

PathTrans（Kim，赵，Tian，&Chandra，2021）将AST映射到一组根路径，每个根路径由通过从叶节点向上遍历根或从根向下遍历叶节点而获得的令牌组成（Jiang，Zheng，Lyu，Li，&Lyu，2021），并指出特定代码元素（即叶节点）所在的纵向上下文。根路径可以用作学习令牌嵌入的输入；然而，由于每条路径只包含零碎的代码元素，从中学习到的不完整语义可能会降低训练和检测性能。

子树。

ASTNN（Zhang et al.，2019）和Infercode（Bui，Yu，&Jiang，2021）等方法不是将AST线性化为标记序列，而是将其分解为一组子树，以利用句法结构。采用不同的粒度来分割AST，其中ASTNN生成与代码语句相对应的非重叠子树，而Infercode生成与表达式等较小代码元素相对应的具有重叠的子树。

该论文的分解器：

给定根节点，它通过访问者和构造函数递归地划分子树。访问者沿着AST执行先序遍历，将每个节点传递给构造函数，用于节点类型检查、复合语句分解和子树构建，其中构建的子树按顺序附加到集合中。这里，最终集合中每个子树的顺序由其对应代码在原始代码片段中的位置决定。图中给出了一个示例。3有助于理解我们的AST分解算法及其与ASTNN的区别。更具体地说，以图3（a）中的归一化代码为输入，我们的算法将其AST分解为12个子树，按其在代码中的出现顺序排列。与ASTNN相比，我们的算法的差异也得到了说明：对于代码中从第6行到第10行的if复合语句，我们的方法生成一个集成子树，而ASTNN将其拆分为三个子语句的平凡子树，如图中蓝色虚线所示。

image-20240319220539749

image-20240418205716678

算法解释

输入参数：

  • T：输入的抽象语法树（AST）。

  • α：子树的最大深度限制。

  • β：子树中允许的最大节点数限制。

  • C：一个包含特定类型语句的列表，这些语句被认为是复合语句。

输出：

  • D：一个有序的子树集合。

TREESPLITTING 函数：这是算法的核心函数，它递归地遍历AST，并根据给定的参数将树分解成子树。

初始化：

  • node：从当前树 t 获取根节点。

处理语句节点：

  • 如果 node 是一个语句节点（即不是复合语句），并且它不在复合语句类型列表 C 中：
    • 如果节点的大小（trSize(node)）小于 α 并且深度（trDepth(node)）小于 β：
      • 构造一个以 node 为根的子树，并将其添加到子树集合 D。
      • 使用 subTreeConstructor 函数构造子树。
    • 否则：
      • 构造一个以 node.header 为根的子树，并将其添加到 D。

递归处理子节点：

  • 对于 node 的每个子节点 child：
    • 递归调用 TreeSplitting 函数，将 child、D、α、β 和 C 作为参数。

处理复合语句节点：

  • 如果 node 是复合语句（即在列表 C 中）：
    • 对于 node 的每个子节点 child：
      • 递归调用 TreeSplitting 函数。

结束函数：

  • 函数结束时返回子树集合 D。

复合语句类型列表 C 包含了被认为是复合语句的类型，如函数定义、if、for、while、do-while、switch、try 和 catch。

算法调用：

最后，算法通过调用 TreeSplitting(T, D, α, β, C) 来启动，其中 T 是初始的AST。

3 Experiment and Evaluation

3.1 DataSet and Process

3.2 Evaluation

4 Conclusion

Summary