code embedding研究系列十-IVDetect

一.背景

现有的漏洞检测方法大部分只是根据给定代码片段，确认该片段是否包含漏洞（分类）。而并没有指出哪些statement有问题。因此作者提出了IVDetect。主要包括

• 用一个新的代码表示方法。作者基于PDG对代码进行表示（源代码用图结构表示），并从PDG提取不同的信息将其向量化。并使用FA-GCN（Graph Convolution Network with feature-attention）对其进行分类。
• 用可解释方法（GNNExplainer）对FA-GCN的分类结果进行解释。GNNExplainer基于edge-mask对输入图选取子图进行解释。作者试图找出是哪些statement决定了分类结果。

作者用三个数据集进行测试： Fan，Reveal，FFMPeg+Qemu

二.motivation

example：

下面展示了linux 4.6的ec_device_ioctl_xcmd方法。这个方法为CromeOS设备构造I/O控制命令。编号为CVE-2016-6156

static long ec_device_ioctl_xcmd(struct cros_ec_dev *ec, void __user *arg){
	long ret;
	struct cros_ec_command u_cmd;
	struct cros_ec_command *s_cmd;
	
	if (copy_from_user(&u_cmd, arg, sizeof(u_cmd)))
		return -EFAULT;
	if ((u_cmd.outsize > EC_MAX_MSG_BYTES) || (u_cmd.insize > EC_MAX_MSG_BYTES))
		return -EINVAL;
	s_cmd = kmalloc(sizeof(*s_cmd) + max(u_cmd.outsize, u_cmd.insize), GFP_KERNEL);
	if (!s_cmd)
		return -ENOMEM;
	if (copy_from_user(s_cmd, arg, sizeof(*s_cmd) + u_cmd.outsize)) {
		ret = -EFAULT;
		goto exit;
	}
+   if (u_cmd.outsize != s_cmd->outsize ||u_cmd.insize != s_cmd->insize) {
+ 	   ret = -EINVAL;
+      goto exit;
+   }
	s_cmd->command += ec->cmd_offset;
	ret = cros_ec_cmd_xfer(ec->ec_dev, s_cmd);
	/* Only copy data to userland if data was received. */
	if (ret < 0)
		goto exit;
-   if (copy_to_user(arg, s_cmd, sizeof(*s_cmd) + u_cmd.insize))
+   if (copy_to_user(arg, s_cmd, sizeof(*s_cmd) + s_cmd->insize))
 		ret = -EFAULT;
exit:
 	kfree(s_cmd);
 	return ret;
 }

commit信息显示：At line 6 and line 13, the driver fetches user space data by pointer arg via copy_from_user(). The first fetched value (stored in u_cmd) (line6) is used to get the in_size and out_size elements and allocation a buffer (s_cmd) at line 10 so as to copy the whole message to driver later at line 13, which means the copy size of the whole message(s_cmd) is based on the old value (u_cmd.outsize) from the first fetch. Besides, the whole message copied at the second fetch also contains the elements of in_size and out_size, which are the new values. The new values from the second fetch might be changed by another user thread under race condition, which will result in a double-fetch bug when the inconsistent values are used.（内核代码看不太懂，大家见谅，希望有大佬能补充下，大概意思就是可能造成double-fetch，一个fetch就是一次获取用户数据（copy_from_user()））

修复这个bug需要保证在这两次获取用户输入之间变量u_cmd.outsize和u_cmd.insize不会由于条件竞争而改变。

针对上述问题，DL-based方法可以将上述函数进行分类（是否包含漏洞），但不能定位到一个具体的行（statement），因此作者在这里用可解释的方法对分类结果进行解释。它会提供一个PDG的子图（几个重要的statement）来解释。

比如：

针对example代码，框出来的即为可解释方法选出的有问题的代码行。同时，针对没问题的代码，可解释方法也会提供相应的信息。

三.Key Ideas and Architecture Overview

IVDetect架构如下图所示：

作者用GNNExplainer作为其解释模型，GNNExplainer会选取关键的子图和子特征（比如一个结点特征向量100维，选取其中5维，不过这需要这个特征向量具备好的解释性，比如词袋向量比word2vec好解释）作为解释。

• 对于子特征，如果mask掉node的某些特征会对分类结果造成明显影响，那么该这些特征应该包括在解释结果里
• 选取子图的选择是，mask掉子图会影响分类结果（有漏洞变没漏洞），子图由一些关键statement和相应控制依赖和数据依赖组成。

3.1 Representation Learning

对于PDG结点（一个statement）的向量表示，作者设置了几组向量特征，一个图结点的最终向量表示由这几组拼接而成。

• Sequence of Sub-tokens of a Statement
  这里将statement转化成sub-token序列，即将某些token再切分为sub-token，和BERT的tokenizer有些类似。这里作者在sub-token序列中只保存变量名（variables），函数名（method names）和类名（class names）。token会根据CamelCase或者Hungarian convention来进行sub-token，并删除单字符sub-token。比如，对于代码if (copy_to_user(arg, s_cmd, sizeof(*s_cmd) + u_cmd.insize))。sub-token序列为copy, to, user, arg, etc。可以看到并没有if,_,(等token。 之后，用glove来向量化单个token，并用GRU将整个statement的序列向量化成向量 $F_1$。

• Code Structure of a Statement
  从AST中捕获一个statement的AST子树，并用Tree-LSTM将其向量化为向量 $F_2$。

• Variables and Types
  对于每个node（statement），收集其中的变量名和变量类型。并用和sub-token同样的向量化方式来进行向量化。比如struct cros_ec_command *s_cmd;中变量名s_cmd，类型cros_ec_command。（这块没太看懂，可能要看下代码了）。得到的向量记为 $F_3$。

• Surrounding Contexts
  将跟该statement存在数据依赖和控制依赖的其它结点分别向量化，并用GRU和glove将这些向量统一计算成 $F_4$ 和 $F_5$（具体操作还得看代码）。

• Attention-based Bidirectional GRU
  最后，用Bi-GRU + attention模型将上面得到的 $F_1，F_2，..,F_5$ 转化成 最终结点向量（真的太复杂了）。流程如下图所示

3.2 Vulnerability Detection with FA-GCN

下图展示了作者如何用FA-GCN来进行分类任务，FA-GCN在处理稀疏特征以及潜在噪声时非常好用。

下图中，join layer层之前的操作都是 3.1 中的内容，join layer将所有statement的向量拼成一个矩阵 Feature Matrix。之后就是卷积过程（数学不好没看懂），最终就是通过Classifier分类。

3.3 Graph-Based Interpretation Model

这一部分主要就是利用GNNExplainer对分类结果提供解释，输入包括2个

• 已经训练好的FA-GCN模型
• 函数 $M$ 的PDG $G_M$

目标就是学习一个edge-mask $E_M$ 来获取子图 $g_M$ 作为解释结果。

$$g_M=E_M\odot G_M$$

对于一个edge-mask，如果mask后模型的分类结果很受影响，那么mask后的子图就是很好的解释子图。

跟在其它领域（社交网络分类，分子分类）一样，GNNExplainer将解释问题转化为如下问题

$$\underset{g_M}{\max }MI(Y,g_M)=H(Y)-H(Y|G=g_M)$$

相当于从所有mask中选择最好的一个，而 $H(Y)$ 是常数（FA-GCN模型已经训练好了），所以可以转化成最小化条件熵 $H(Y|G=g_M)$

$$H(Y|G=g_M)=E_{Y|g_M}logP(Y|G=g_M)$$

GNNExplainer同样会限制解释子图边的数量。给定边数 $K_M$ 要求出最小条件熵在给定 $Y$ 为分类结果前提下。这里GNNExplainer将 $g_M$ 当作一个随机图变量（大概就是每条边都会赋值成概率，表示它在解释子图中出现概率，不再是1） 。所以GNNExplainer的训练目标是，这里优化的参数就是edge-mask。

$$\underset{g}{\min }{E}_{g_M\sim g}H(Y|G=g_M)$$

训练完成之后，生成的解释子图 $g_M$ 可以用来解释，需要注意的是GNNExplainer对于每个输入图都会生成特定的edge_mask，也就是说每解释一个PDG就需要重新训练（训练是一次性工作，所以有了PGExplainer）。

四.实验目标

作者的目标还挺多了，又分类又解释还要挖掘漏洞pattern。

4.1 Comparison on the Method-Level Vulnerability Detection

用IVDetect和其它漏洞分类方法进行比较，有VulDeepecker，SySeVR，Reveal，Devign还有token-based方法。作者这里还使用了AutoML来调参。

评估指标

• Precision ( P )

• Recall ( R )

• F score ( F )

• Mean Average Precision ( MAP )
  这个指标需要先理解PR曲线，这里有篇从PR->MAP都有的讲解。

• Normalized DCG
  貌似是推荐系统用的指标，不是很清楚

• First Ranking ( FR )
  我也没太看懂这个是用来衡量什么的

• Accuracy under curve ( AUC )
  AUC曲线，可以参考这篇

4.2 Comparison with other Interpretation Models for Fine-grained VD Interpretation

这里作者拿GNNExplainer和其它解释方法进行了对比，对比的方法有

• ATT
  用来衡量边的重要性的Graph attention机制

• Grad
  基于梯度的方法

因为解释方法的对比需要有定位到漏洞行号的数据集，因此只有Fan数据集可以使用，其它两个只有哪些方法有漏洞的信息，并没有fix。这里，作者用在Reveal和FFMPeg+Qemu数据集上训练的FA-GCN模型来对Fan数据集分类。

对于错误分类（标签1，预测0）的作者不考虑进行解释。而同时，由于标签0的函数没有fix信息，解释出来也没法评估，所以作者只考虑标签1预测1的（TPR）。

Evaluation Metrics：

对于一个解释子图 $G_M$。$S$ 为从vulnerable版本的代码中删除和修改的statement的集合。修改的目的是修复漏洞。如果 $S\in G_M$，那么解释结果算正确，否则错误。也就是所有的vul statement都要包括才行。

对于fixed版本，$S^{{}^{\prime }}$ 为添加的statement的集合。而此时 $G_M$ 包含任意一个 $S^{{}^{\prime }}$ 中的statement，就算正确，否则失误。

作者还用了Mean First Ranking (MFR)和 Mean Average Ranking (MAR)来评估解释结果，不过没太看懂这两个指标。

需要注意的是作者并没有用fidelity和sparsity这2个指标，也许不符合人的直觉吧。

4.3 Vulnerable Code Patterns and Fixing Patterns

对于用GNNExplainer产生的一系列解释子图，作者用挖掘算法从这些子图中挖掘出漏洞pattern以及从fixed版本中挖掘fix pattern。这里并没有客观评价指标。

4.4 Sensitivity Analysis for Internal Features

作者在前面提到了对于图结点向量表示用了4种特征（token sequence, AST, type, 控制依赖和数据依赖算一种）。这里作者从只用其中一个特征（token）对statement向量化开始，一步一步添加特征，评估不同特征对结果的影响。

评估指标和4.1相同。

4.5 Sensitivity Analysis on Training Data

分别用(80%, 10%,10%), (70%, 15%, 15%), (60%, 20%, 20%)和(50%, 25%, 25%)的(training, tuning, testing)数据集划分比率来研究效果。

4.6 Time Complexity

评估实际训练和测试用时。

五.实验结果

5.1 Comparison on the Method-Level Vulnerability Detection

这里作者做的对比实验还挺多的，并且分别在FFMPeg+Qemu，Fan，Reveal数据集上分别测试的。

先看看precision, recall, F1 3个指标

其它指标

FFMPeg+Qemu数据集

Fan数据集

Reveal数据集

5.2 Comparison with other Interpretation Models for Fine-grained VD Interpretation

作者对比了3种解释方法的效果（GNNExplainer, Graph Attention, Grad）。用到的指标上面也提到了，accuracy是与分类的类似（correct or incorrect。correct定义上面说过）

5.3 Vulnerable Code Pattern Analysis

作者从前面的解释任务中收集了700+解释子图。作者先用与VulDeepecker和SySeVR相似的符号化（把变量名用VAR替代等待）方式预处理解释子图，然后用下面参考文献1提到的子图pattern挖掘算法配上不同的size限制来挖掘不同大小的子图pattern。

作者在此进行人工验证。不过没评估好坏，只是数个数。

作者还贴上了2个example

下图展现了2个不同的漏洞pattern。第一个属于api误用，包括is_link,exit,cop_file，第二个调用了udf_get_filename这个有漏洞的函数，这个漏洞可以通过添加第5个参数得到修复。

下图则是fixing pattern。第一个是为了修复多线程下数据加锁类问题，可以看到修复方式就是加锁。第二个图是为了修复一个缓冲区溢出漏洞。

5.4 Sensitivity Analysis for Features

下表展示了不同特征对漏洞分类的效果影响。涉及名词有

• sequence of sub-tokens (SST)
• sequence of tokens (ST)
• control dependency（CD）
• data dependency（DD）
  

可以看到作者对比的特征是ST，ST+SST，ST+SST+AST等等。逐级递增，但并没有单独测试AST,VAR等等。

5.5 Sensitivity Analysis on Training Data

数据集切分方式的影响

六.参考文献

[1] Tung Thanh Nguyen, Hoan Anh Nguyen, Nam H. Pham, Jafar M. Al-Kofahi, and Tien N. Nguyen. 2009. Graph-Based Mining of Multiple Object Usage Patterns.
[2] Li Y , Wang S , Nguyen T N . Vulnerability Detection with Fine-grained Interpretations. 2021.