Vulnerability Detection by Learning From Syntax-Based Execution Paths of Code
Vulnerability Detection by Learning From Syntax-Based Execution Paths of Code
foresta.yang0 Abstract
在这项工作中，我们提出将代码片段的基于语法的控制流图（CFG）分解为多个执行路径来检测漏洞。具体来说，给定一个代码片段，我们首先基于它的抽象语法树（AST）构建它的CFG，将这种CFG称为基于语法的CFG，并将CFG分解为从入口节点到出口节点的多个路径。接下来，我们采用预先训练的代码模型和卷积神经网络来学习具有路径内和路径间注意力的路径表示。路径的特征向量被组合为代码片段的表示，并被输入分类器以检测漏洞。将代码片段分解为多个路径可以过滤掉一些与漏洞无关的冗余信息，并帮助模型关注漏洞特征。此外，由于分解的路径通常比代码片段短，位于长代码尾部的信息更有可能被处理和学习。为了评估我们模型的有效性，我们构建了一个包含超过231k个代码片段的数据集，其中有24k个漏洞。实验结果表明，所提出的方法在精度、召回率和F1分数方面分别优于最先进的基线至少22.30%、42.92%和32.58%。我们的进一步分析调查了所提出的方法优越性的原因。
1 Intro or Overview
1.1 Problem and Challenge
漏洞检测对于保护软件系统至关重要。已经提出了基于深度学习的各种方法来学习漏洞模式并识别它们。尽管这些方法在这项任务中显示出了巨大的潜力，但它们仍然存在以下问题：（1）它们很难将与漏洞相关的信息与大量无关的信息区分开来，这阻碍了它们捕捉漏洞特征的有效性。（2） 它们在处理长代码方面效果较差，因为许多神经模型会限制输入长度，这阻碍了它们表示长时间易受攻击的代码片段的能力。
1.2 Motivation
Observation 1: A vulnerable function may contain a vast of statements unrelated to vulnerabilities.
Observation 2: Truncating the statements in the tail of a long code snippet may negatively affect the effectiveness of vulnerability detection.
1.3 Contribution
2 Architecture & Method
2.1 System Overview
形式上，我们采用Gi=（Vi，Ei）来表示代码片段ci的基于语法的CFG，其中Vi=（v1 i，v2 i，…v|Vi|i）是包含|Vi|语句的节点集。Ei是表示语句之间的控制流的边集。每条边都是ci中两个语句之间的关系，因此一个语句可以在另一个语句之后执行。Gi中的路径是节点序列p=（n1，n2，…，nk），其中节点nk是ci的陈述k。对于任何一对相邻节点np和nq，存在从np到nq的边。如果起始节点等于结束节点，则路径将是一个循环。
首先将每个代码片段解析为AST，并根据AST构建基于语法的CFG。接下来，我们提出了一种基于贪婪的路径选择算法，从基于语法的CFG中选择多个执行路径，即将代码片段分解为几个执行路径。然后，通过CodeBERT[24]将所选路径编码为具有路径内注意力的向量，然后将其馈送到CNN中以捕获路径间注意力。最后，我们利用MLP分类器来执行检测。
2.2 Method
Construction of Control Flow Graph From AST
该阶段以代码片段为输入，使用tree-sitter[50]将其解析为AST，并从AST构建基于语法的CFG。在本节中，我们使用图3中的示例来说明我们如何从代码片段的AST构建基于语法的CFG。在构造CFG之前，我们使用正则表达式删除代码段中的空行和注释。我们还在代码中标记每条语句的行号。由于基于语法的CFG中的每个节点代表一个单独的语句，因此在构造CFG时，我们只考虑AST中的语句节点。为了简化演示，我们做出以下定义：
简单语句：在AST中不包含其他语句的语句。
Next语句：节点的Next语句是指在执行节点后可能执行的语句。一个节点可能有多个Next语句。
非子Next语句：节点的非子Next声明是指该节点的Next语句，不包含在以该节点为根的子树中。
Normal语句：不是break_语句、continue_Statement、return_Statement和throw_Statement的语句。
为了构造代码段的基于语法的CFG，我们首先将代码段的所有语句添加到CFG中作为其节点，并将第一个语句视为入口节点，将所有return_statement、assert_statement和throw_statement视为出口节点。如果代码片段的最后一条语句不是出口节点，我们将在代码末尾添加一个伪出口节点。然后，我们以广度优先的方式遍历AST，并为每个语句类型设计规则，以在CFG中构建边。
1） 对于每个既是简单语句又是普通语句的语句，如果AST中存在其下一个同级语句，我们将其连接到此同级语句。例如，在图3中，我们将节点2连接到节点3，将节点3连接到节点4。
2） 对于每个循环语句，即For_statement和while_statement，如果存在这样的语句，我们将其与第一个子语句和下一个同级语句连接起来。如果它的最后一个子语句是Normal语句，我们将此语句连接到循环语句。例如，在图3中，我们将节点4连接到节点5，将节点4与节点10连接，将节点5与节点4连接。
3） 对于每个break_statement，我们首先找到它的第一个祖先，它是沿着AST的循环语句或switch_statement。然后，我们将其连接到祖先的非子下一个语句。
4） 对于每个continue_statement，我们首先找到它的第一个祖先，即沿着AST的循环语句。然后，我们把它和这个祖先联系起来。
5） 对于每个if_statement，如果存在这样的语句，我们首先将其连接到其下一个同级语句，如果最后一个子语句是Normal语句，则将其then_block的最后一个子声明连接到其当前next语句。例如，在图3中，我们将节点6连接到节点4，将节点7连接到节点4.将节点10连接到节点12，将节点11连接到节点12。然后，如果它的子级包含else_statement，我们将else_Statements连接到它的每个Next Statements，将if_statement中的边移除到它的Next Statements中，并将if_sttatement中的边缘添加到else_statement。对于图3，我们将节点8连接到节点4，移除从节点6到节点4的边，并将节点6连接到节点8。接下来，我们遍历else_语句。我们将其最后一个子语句连接到当前的Next语句。如果它的最后一个子语句是Normal语句，我们会将其边缘移除到Next语句，并将其连接到第一个子语句。对于图3，我们将节点9连接到节点4，将节点8移除到节点4并将节点8连接到节点9。最后，我们将if_statement连接到它的then_block的第一个子语句。对于图3，我们将节点5连接到节点6，将节点6连接到节点7，将节点10连接到节点11。
6） 对于每个switch_statement，我们将其连接到其第一个case_statement。对于每个case_statement，我们首先将其连接到下一个case_statementor default_statement。然后，如果这个case_statement的最后一个子语句是Normal语句，我们将其连接到这个case_statement的当前Next语句。最后，我们将case_statement连接到它的第一个子语句。对于每个default_statement，如果其最后一个子语句不是Normal语句，我们将其连接到其第一个子语句，并将其最后一个子语句连接到switch_statement的Non-child-Next语句。
7） 对于每个try_statement，我们将其catch_clauses视为语句。我们不能为try_statement构造一个“声音”CFG，因为我们不能知道每个函数调用只能从调用方的AST抛出哪些异常。此外，构建一个“完整”的CFG需要将try_block中的每个语句连接到每个catch_clause，这可能会引入太多的死路径，并对以下阶段产生负面影响。因此，我们选择只将try_block中的最后一个Normal语句连接到catch_clauses。具体来说，我们将try_statement连接到其try_block中的第一个语句，将其try_bock中的最后一个Normal语句连接到其第一个catch_clause。对于每个catch_clause，我们将其连接到它的第一个语句和下一个catch_clause。此外，对于try_block和每个catch_clause中的最后一个语句，即Normal语句，我们将其连接到try_statement的Non-Child Next语句。
Path Selection
一个代码片段可以被视为其所有执行路径的组合。但是，如果代码段包含循环，则它可能具有无限的执行路径。它可能需要许多计算资源来编码代码片段的所有执行路径。**因此，我们认为，在CFG中对所有执行路径进行编码以学习相应代码片段的表示是不切实际的。我们将基于语法的CFG中的执行路径定义为从CFG的入口节点到出口节点的路径，并从此将其称为执行路径。此外，在看不见的代码片段中准确定位易受攻击的语句是非常重要的。**如果我们只提取一个执行路径来表示代码片段，那么很可能会错过易受攻击的语句，并对漏洞检测的性能产生负面影响。幸运的是，根据我们对现实世界漏洞的观察和第三节中给出的激励性示例，我们发现一些执行路径通常可以涵盖漏洞的根本原因。因此，我们选择并编码几个具有代表性的执行路径来表示相应的代码片段，而不是对CFG中的所有或仅一个执行路径进行编码。Alon等人也使用了类似的策略。[51]在将代码片段表示为AST路径时。该阶段负责从先前阶段构建的CFG中选择几个具有代表性的执行路径。
选择执行路径有两个要求：首先，为了避免丢失代码片段中的重要信息，所选路径应覆盖尽可能多的代码行。其次，为了减轻模型训练的负担，我们希望选择的路径尽可能短。不幸的是，这两个要求在某种程度上是冲突的。为了在它们之间进行权衡，我们提出了一种基于贪婪的路径选择算法。
