A Learning-Based Approach to Static Program Slicing

OOPLSA’2024

1. Abstract

• 背景: 程序切片: 传统程序切片技术在漏洞检测中很关键, 但是处理不完整代码表现不佳

• 方法: 提出NS- Slicer, 一种基于学习的新颖方法能够预测完整代码和不完整代码的静态程序切片.
  使用预训练的大模型(PLM), 利用模型对源代码中细粒度依赖(Token级别的依赖?)的理解对每个token生成富含上下文信息的向量, 利用这些向量, NS-Slicer会分别确定语句是属于向后切片还是向前切片

• 结果:

  • 在完整代码上, 后向切片和前向切片分别达到97.41%和95.82%的F1分数,
    并且在85.2%的例子中, NS-Slicer精确预测了整个切片
  • 在部分代码上, 达到了94.66%-96.62%(实验摘掉了不同数量的代码)的F1分数
  • 在漏洞检测的任务上, 检测Java代码达到了73.38%的F1分数
  • 实验证明能同时处理完整/不完整代码, 对任何切片起点进行切片, 对能容忍少部分不精确结果的任务比较有效, 能获得规模, 时间和对不完整代码的处理能力为权衡

2. Motivation

(1) 为什么需要能处理不完整代码

比如这样一个第13行有漏洞的, 来自StackOverflow #16180130回答的不完整代码, 被用在了Hadoop中产生了一个CVE, 类型CWE-395 Null Pointer Exception

顺便解释一下工作中所说的“不完整”指的是什么, 看这个代码片段, 没有被函数包裹, 中间的一些变量也没有声明, 整个代码片段都是没有上下文的

所以尽早检测类似stackoverflow这样的不完整代码漏洞是有必要的.

但是, 整个stackoverflow回答就这15行, 没有变量定义也没有外部库的信息, 没有任何上下文

(2) 现有方法局限性

而现有的方法需代码是带有完整依赖的, 至少是在函数层面上是完整的

1. 比如基于程序分析的: SDG的javaslicer和基于CPG的joern都不行

2. 基于深度学习的PDG生成方法NeuralPDA, 提供的依赖是语句级别的, 粒度不够细, 不能提供生成切片必要的变量和语句之间的依赖(因为切片起点是个变量)

所以这些方法在stackoverflow上检测代码小片段的能力有限,

(3) 验证局限性: 初步实验研究

做了一个小实验来说明现有基于程序分析的sota(joern)解决不了不完整代码

实验设计

从stackoverflow上找99个不完整的代码小片段(没有说是随机挑的), 输入给joern, 人工检查joern生成的CFG/PDG看是否正确

何为不完整代码

有几种类型

• 有未知类型的数据
• 缺少头文件的外部API/方法/类/类型
• 引用的变量无变量定义
• 缺少类的层次结构

---

举例:

实验就是直接输入这么一小段, 没有inclue包含std::string, std::transform的头文件, 所以这里都是未知的数据类型和API

---

另一个例子
工作中还用了没有函数签名的stackoverflow的代码小片段
比如, 这是一个来自stackoverflow上, 没有函数签名的代码片段, 实验会给其加上没有参数的的函数签名, 再输入给joern, 排除由于没有函数签名导致joern无法输入

本文里没有给具体例子, 我觉得应该是这个意思, 图也是在这个作者另一篇针对不完整代码的文章里找的

void DUMMY_METHOD_SIGNATURES(){
  std::shared_ptr<FILE> pipe(popen(cmd, "r"), pclose);
  if (!pipe) return "ERROR";
  char buffer[128];
  std::string result = "";
  while (!feof(pipe.get())) {
    if (fgets(buffer, 128, pipe.get()) != NULL)
      result += buffer;
  }
}

实验结果

• 47/99个例子: 缺少或预测不正确的数据/控制依赖

• 30/99个例子: joern报错, 比如“Could not find type member, type = XYZ, member = abc”

• 7/99: joern生成了一个空CFG/PDG

Joern生成不好CFG/PDG的主要原因是

• 未知数据类型的变量的依赖会全部被Joern忽略
• 各种原因导致的(比如缺少头文件导致的)对unresolved的外部API/方法/类/类型的引用和对象构造会报错, 或者直接跳过
• 缺少变量的声明, 变量类型/类的声明
• 缺少类的层级(继承)结构
• 不能处理template和typedef

(4) 为什么这个工作可以处理不完整代码

所以要利用预训练大模型(PLM)的优势
为什么可以处理不完整代码, 是PLM输入的序列性: 一个Token只要被输入都会被考虑上下文和数据流关系, 而不会像传统工具一样限制变量一定要声明过才会追溯数据流信息

---

并且PLM能捕捉到代码的语义信息, 适合切片任务, 这源于预训练任务的性质

• 掩码语言建模(Masked Lang Modeling: Feng et.al 2020. CodeBERT: A Pre-Trained Model for Programming and Natural Languages. In EMNLP),
• 边预测(Edge Prediction: Guo et.al 2021. GraphCode{BERT}: Pre-training Code Representations with Data Flow. In International Conference on Learning Representations.)和节点对齐(Node Align), 这两个预训练在GraphCodeBERT中也叫学习数据流知识
• 一些工作已经证明有助于PLM学习代码中的语法结构和数据流信息, 他说PLM在语句捕捉依赖性级别有好表现

由于GraphCodeBERT的预训练目标有编码变量来自何处, 因此适合处理部分代码, 在其上做静态切片

3. Method

后面有五个实验方法有小差异, 这里简单概述共通的

训练集构造

用传统的基于程序分析的切片方法切出来的切片作为ground truth

<源程序, 切片起点, 前/后向切片集合(现有的切片方法)>

训练:

给定一个源程序

1. tokenize: 把源程序语句切成文本形式的token序列

2. Variable-statement Dependency Learning(用PLM模拟PDG的生成):

   1. 嵌入: 把文本形式的token序列送入预训练的大模型(CodeBERT/GraphCodeBERT), 大模型捕捉token序列间的依赖关系, 对每一个token生成含有上下文信息的向量表示
   2. 池化: 将属于同一语句的token向量表示应用平均值池化, 计算每一条语句的向量表示, 特别的, 将切片起点的变量单独池化

3. 切片: (用学习方法模拟对PDG的遍历)

   对每一条语句, 将这条语句的向量表示 和切片起点变量的向量表示输入到预测前/后向切片的二分类MLP里(这条语句在切片起点前就输入到前向切片分类器, 在后面同理), 预测这条语句是否是切片起点的前/后向切片

   最后计算预测的前/后向切片集合与数据集的ground truth的前/后向切片交叉熵损失之和 训练

   为什么要训练两个因为前后向不一样

4. Empirical Evaluation

RQ1: NS-Slicer在完整代码上的有效性

1. 训练集构造:

CodeNet, 有大约75,000个Java代码样本

使用JavaSlicer(Galindo et.al 2022)对每个样本中的每一个变量作为起点切片, 剔除掉空的切片

因为定义为二分类问题, 还做了一个样本均衡的筛选, 就是挑选切出来的语句数量是原代码的0.3-0.7的样本

最后保留了43,000个Java样本, 每个样本5-69个切片

8:1:1划分

数据集: IBM CodeNet(Puri et.al 2021) 4053

https://developer.ibm.com/exchanges/data/all/project-codenet/

简介: CodeNet提供了一个oj, 有4,053个问题,代码都是使用者提交并通过IBM审核后加入数据库的, 14,000,000个代码样本, 五十多种语言, 样本不光包括代码, 还包括runtime和报错

2. 训练:

首先，选择现成的 GraphCodeBERT 预训练版本。在 NS-Slicer 的训练阶段，分PLM 中的参数是固定的和微调两种，PLM 中的标记/语句表示原封不动地用于训练前后向切分 MLP 头. 对CodeBERT 模型[Feng 等人，2020]，做同样的事情。

两个PLM总共128M个参数, 每个epoch 32min

3. 指标

在语句级别计数TP, FP啥的

• Accuracy(A-S), Precision, Recall, F1

• A-EM(Accuracy-ExactlyMatched): 预测出的切片和GroundTruth一摸一样的比例

• A-D(Accuracy-D):

  ​ Finally, we report Dependence Accuracy (i.e., Accuracy-D) to assess how accurately the interstatement dependencies are predicted, causing a particular statement to be included in the slice.
  ​ Accordingly, we define Accuracy-D for a particular program as the ratio of the correctly predicted dependencies to the actual dependencies across all slicing criteria for that program, finally reporting the mean across all programs in the dataset.

文章里没说具体怎么衡量是否找对了Dependency, 代码找了个遍也没找到

4. 结果

Table1:

可能由于工作缺乏切片的数据集, groundtruth是JavaSlicer, 这个结果只能表示有多接近JavaSlicer的性能, 可以看到最好的情况下, 有百分之85.20的例子生成的切片是和JavaSlicer一摸一样的, F1是96.77%

可以看出用GraphCodeBERT比CodeBERT效果好, 作者说是因为GraphCodeBERT考虑到了代码的数据流信息

还探究了模型对切片在整个程序中占比的敏感程度

RQ2: 在不完整代码上的有效性

1. 数据集构造

和RQ1相同的数据集, 先用完整代码生成切片

同时摘掉前后百分之P的来模拟部分代码, 和切片有关的语句

PLM方面只使用GraphCodeBERT, 其余和RQ1一样

2. 结果

RQ3: 消融实验

B1: 探究在源代码上预训练对结果的影响: 使用在自然语言——英语上预训练的模型RoBERTa-base作为PLM

B2: 探究数据流预训练(GraphCodeBERT进行了边预测和节点对齐的额外预训练, 即数据流预训练)对结果的影响: 使用CodeBERT作为PLM

B3: 把平均值池化换成最大值池化

结果:

对比B1, B2可以看出在源代码上预训练比在数据流上预训练更重要, 对部分代码影响更大(是文章中直接说的, 没有表)

RQ4: 探究PLM对变量别名对理解

variabe aliasing指的是, 多个引用用来表示同一个对象并且只想同一块内存

1. 实验方法

在切片起点变量所在的那条语句后面插入一条赋值语句来给切片起点变量起一个别名, 后续所有语句中的切片起点变量全部都改成别名

平均只有42.86的语句含有变量别名

然后作者对, 不同变量别名占切片比例的样本进行分层观察, 发现这个比例和NS-Slicer的预测能力没有直接关联

推测这可能是由于在源代码上预训练的PLM在内存级别缺乏对源代码的理解，因为当前的预训练任务主要只关注源代码的词法方面。因此，对特定源代码的预训练还能进一步提高 PLM 的程序切片性能。

RQ5: 在漏洞检测任务上的性能

数据集构造

从563个Java的commit中只有574个有漏洞函数的文件, 为了保证样本平衡从13, 565个无漏洞的文件中抽了574个, 用NS-Slicer生成了1476个代码小工具, 由于数据有限, 先在BigVul上Joern+VulDeePecker训练, 然后把Joern换掉, 数据集也换成CrossVul继续训练

数据集: CrossVul跨语言漏洞数据集

结果:

5. Limitations

模型的输入限制在了512个tokens, 然而LongCoder(Guo et.al 2023, 最多4096个token的限制)可以很容易的嵌入NS-Slicer