Code Embedding研究系列11-ContraFlow

一.引言

1.1.现有方法及其局限

Code Embedding，一个意在通过分布式向量表示代码语义的技术，在漏洞检测领域已经取得了一些进展，现有的方法大致可分为2类：

• Sequence-based：VulDeePecker，SySeVR等，这些方法将代码片段视为一段序列，通过RNN对序列进行分类。

• Graph-Based：IVDetect，Reveal，Devign等，这些方法首先将代码解析为图结构，之后通过GNN对序列进行分类。

这些方法进行漏洞检测的思路通常是进行二分类，类似的技术在code classification和code summarization任务上取得了成功，但是在漏洞检测领域却并不是很成功。主要原因是:

• 这些方法的代码表示并没有区分Program Path（路径），GNN也不会区分Path。

• 图的特征是通过GNN的Message Passing机制从图中相邻的结点对学到的。但是，这并不会区分可行/不可行的value-flow (data-dependence) paths。

因此，模型并没有理解潜在的漏洞path，这也大大限制了模型定位漏洞的能力。作者也提出了下面的Insights：

• 模型需要学习的是可行的value-flow (data-dependence) paths而不是整个图结构。

• Code2Vec将代码表示为AST Path Bags，这就是一个AST Path集合，不过AST Path数量庞大，作者将程序表示为可行的value-flow path集合，最后这些value-flow paths会被聚合为整个程序向量。

• 尽管如此，程序中value-flow path的数量依旧是不确定的，所以需要Path selection来选取最重要的path。

1.2.作者的解决方案

基于上面的Insights作者提出了ContraFlow，一个漏洞检测框架，概览如下：

这个流程可分为下面几个阶段

• ( a ) Contrastive Value-Flow Embedding：这个阶段主要目的是通过对比学习方式训练一个value-flow embedding model, Value-flow Path Encoder（VPE），训练过程如下，这一步会预训练出一个VPE并应用在后面2个阶段。

  • 首先用SVF从未标注的数据集中提取一个Value Flow Path集合

  • 然后使用数据增强手段来生成contrastive value-flow representations

  • 然后利用标准Noise Contrastive Estimate (NCE）loss函数最大化语义相似value-flow path之间的向量相似度。

• ( b ) Value-Flow Path Selection：这个阶段的主要目的是从程序提取的Value-Flow Path中选取可行并且具备代表性的path。

  • 首先采用预训练的VPE生成输入path集合的向量，然后采用 self-supervised active learning 来捕获最具代表性的path集合并使得embedding更加多样化。

  • 之后就是可行性分析，value-flow path通常伴有条件分支，这一步主要排除路径条件不能全部满足的path。

• ( c ) Detection Model Training：这一步主要是从(b)中选取的path集合和(a)中预训练的VPE中训练一个分类器进行漏洞分类

  • 首先利用VPE获取path的初始向量

  • 然后通过Muti-head层获取每个path的contextual vector，之后通过一个soft attention layer聚合所有path的向量，同时计算每个path的权重

  • 训练好的分类器可以根据attention score给path分配权重并计算buggy path

可以看到训练的模型包括VPE, self-supervised active learning和分类器，其中VPE和分类器应用到了最终的漏洞检测器中。

二.Motivating Example

这个Example来自POCO project，属于CWE-840，漏洞触发的根源是API误用，在 rebuild_list(&hd) 后调用了 set_status(&hd)。其中 hd 变量在第2行定义，第6行被修改，第13行被使用，2-6-13是一个漏洞触发路径。

• 首先在阶段 (a) 提取了4个value-flow path（${\pi }_1:2-6-13$ , ${\pi }_2:2-6-15$, ${\pi }_3:3-9-13$, ${\pi }_4:3-9-15$），作者将这4条路径2次输入到VPE中，每次使用不同的dropout，可以分别获取向量 $v_{{\pi }_1},v_{{\pi }_2},v_{{\pi }_3},v_{{\pi }_4}$ 和 $v_{{\pi }_1}^{+},v_{{\pi }_2}^{+},v_{{\pi }_3}^{+},v_{{\pi }_4}^{+}$。其中，VPE的训练目标是相似的embedding vector距离相近 ($v_{{\pi }_1},v_{{\pi }_1}^{+}$)，而不同的vector距离较远 ($v_{{\pi }_1},v_{{\pi }_3}^{+}$)，这一步采用的loss是NCE loss。

• 阶段(b) 首先根据self-supervised active learning获取 ${\pi }_1-{\pi }_4$ 的排序，之后这些path会进行可行性分析。其中 ${\pi }_2$ 和 ${\pi }_3$ 不可行， $2-6-15$ 中 $6$ 需要满足条件 FLG 而 $15$ 要满足条件 !FLG，${\pi }_3$ 反之亦然。

• 阶段 ( c ) 中对于输入path ${\pi }_1,{\pi }_4$。首先用VPE 计算其向量 $v_{{\pi }_1},v_{{\pi }_4}$，然后 $v_{{\pi }_1},v_{{\pi }_4}$ 会输入Muti-head层计算每个path的contextual vector，然后所有path的contextual vector会输入soft attention层，聚合成整个程序的向量，同时每个path都会对应1个attention权重，表明这个path的贡献程度。图中 ${\pi }_1$ 的贡献达到 90%，最有可能是漏洞path。

三.ContraFlow

3.1.Contrastive Value-Flow Embedding

3.1.1.path vectorizing

提取到了value-flow path，首先要将path向量化，一个path $\pi$ 由多个statement $s$ 组成，可以视作一个statement sequence。VPE向量化pipeline如下所示（这里介绍VPE模型结构）：

Local Encoding主要目的是独立地向量化每个statement，向量化statement $s$ 用到了 $s$ 的AST subtree，下图给出了一个示例。首先对于 $s$ 的AST子树中的一个结点 $n_i$，它的初始结点向量 $v_{n_i}$ 通过 node type 和 node token（类似Devign和Reveal）生成，之后通过下面公式更新

v n i ′ = ∑ j ∈ C i ∪ { i } α i j . W . v n j v_{n_i}^{'} = \sum_{j \in C_i \cup \{i\}} \alpha_{ij} .W.v_{n_j} vni​′​=j∈Ci​∪{i}∑​αij​.W.vnj​​

• $C_i$ 是 $n_i$ 的子结点

• $W$ 是VPE的参数

• ${\alpha }_{ij}$ 是 $n_i$ 和 $n_j$ 之间的attention值

之后按如下方式生成 $s$ 的summary vector $v_{sm}$：

$$v_{sm}=[\frac{1}{N}\sum _{i=1}^N{v}_{n_i}^{{}^{\prime }}||\underset{j=1}{\overset{N}{\max }}{v}_{n_j}^{{}^{\prime }}]$$

其中用到的模型可能有多层 ($L$ 层)，statement vector $v_{jk}={\sum }_{l=1}^L{v}_{sm}^l$，而最终statement向量 $v_s=dropout(W.v_{jk})$

Global Encoding根据每个statement的向量生成path向量，前一步计算出了每个statement的向量 $v_{s_0},v_{s_1},...,v_{s_N}$，这一步通过BGRU将 $v_{s_0},v_{s_1},...,v_{s_N}$ 转化为 $h_{s_0},h_{s_1},...,h_{s_N}$，这一组向量描述了一组def-use（data-dependence）关系。之后通过attention机制将所有hidden向量聚合为path向量 $\pi$。

3.1.2.Contrastive Value-Flow Embedding Algorithm

这一步主要目的是训练VPE，算法如下图所示
首先对比学习用到的数据对中，path $\pi$ 会输入VPE 2次根据不同的dropout获得向量 $v_{\pi }$ 和 $v_{\pi }^{+}$，这作为positive，而 $v_{\pi }$ 和其它path的向量会组成negative。这种数据增强操作比element deletion或者substitution要好。

之后，作者用Noise Contrastive Estimate (NCE) loss函数训练VPE，path向量的相似度用余弦相似度衡量。

$$sim(v_{{\pi }_i},v_{{\pi }_j})=\frac{v_{{\pi }_i}^T.v_{{\pi }_j}}{||v_{{\pi }_i}||.||v_{{\pi }_j}||}$$

path ${\pi }_i$ 的 loss定义如下：

$$loss(v_{{\pi }_i})=-\log \frac{\exp (sim(v_{{\pi }_i},v_{{\pi }_i}^{+}))}{{\sum }_{k=1}^B\exp (sim(v_{{\pi }_i},v_{{\pi }_k}^{+})}$$

$B$ 为1个Batch的数量，一个Batch的Loss为

$$L=\frac{1}{B}\sum _{i=1}^{B}loss({\pi }_i)$$

3.2.Value-Flow Path Selection

3.2.1.Value-Flow Active Learning

这一步主要目的是选取出具有代表性的path，参考的是文献${}^{[2]}$，给定通过VPE向量化好的path集合 $H=[v_{{\pi }_1},...,v_{{\pi }_n}]\in R^{n\times d}$，选取出 $H^{{}^{\prime }}\in R^{k\times d}\subseteq H$，这里引入了1个encoder-decoder模型。需要优化的参数为2个矩阵 $Q,P\in R^{n\times n}$，优化目标包括最小化下面目标的距离：

• $H$ 和 $Q.H$

• cluster centroid matrix $C$，通过K-means聚合得到

• $H$ 和decoder输出 $G$

训练完毕后，作者为 $Q,P$ 的每一列计算 l2-norm 值，生成2个ranking vector $\hat{q},\hat{p}\in R^n$，之后，这两个ranking vector会聚合并降序排序，然后就可以选取top-k value flow path了。

3.2.2.Value-Flow Path Feasibility Analysis

这一步只要进行path-sensitive分析，分析目标是排除控制流约束条件不能得到满足的路径，形式化描述就不写了，可以用下图展示：

简而言之

• 如果没有约束条件，那么就是 True

• 如果有约束条件，那么看整条路径上的所有约束能不能同时满足，不能为 False， 能为 True

3.3.Detection Model Training

这里就是训练分类器了，分类器的输入是Path集合的向量 $V=[v_{{\pi }_1},...,v_{{\pi }_N}]\in R^{N\times d_{emb}}$，$N$ 是path数量。而分类器主要包含1个Muti-head attention层+1个soft attention层+1个线性分类层

muti-head的计算过程如下：

$$h_i=Attn(V.W_i^Q,V.W_i^K).V.W_i^V$$

$$V^{{}^{\prime }}=[h1||...||h_h].W^o$$

$V^{{}^{\prime }}\in R^{N\times d_{ctx}}$ 是muti-head层输出，$W_i^Q,W_i^K,W_i^V,W^o$是muti-head层参数， soft attention层计算如下：

$${\alpha }_i^c=\frac{\exp (v_{{\pi }_i}^T).a_c}{{\sum }_{j=1}^N\exp (v_{{\pi }_j}^T).a_c}$$

$$v_c=\sum _{i=1}^N{\alpha }_i^c.v_{{\pi }_i}$$

$v_c$ 是soft attention层输出的整个程序的向量，$a_c$
是soft attention层的参数，之后 $v_c$ 会输入一个全连接层+softmax进行分类。

给定训练好的模型，整个系统大概可以总结如下：

• VPE部分可能包括一个embedding模型用来根据node token生成node初始向量， 一个attention层用来向量化AST子树生成statement向量，一个BGRU+attention用来生成path向量。

• Value-Flow Selection部分包括一个Encoder-Decoder来选取代表性path，以及一个约束求解模块进行可达性分析。

• 分类器包括1个Muti-Head生成contextual vector, 1个soft attention层聚合多个path + 1个全连接层分类。

四.实验设计

数据集用到了D2A, Fan和FFMpeg+QEMU。数据集基本信息如下：

首先源代码被编译为LLVM IR，然后输入到SVF中解析value-flow graph（可能会从LLVM层面的VFG映射到源代码层面的VFG），提取AST子树用到了Joern，约束求解用到了Z3 SMT求解器。对比实验包括SySeVR，VulDeePecker，Reveal，DeepWukong， Devign， IVDetect， VulDeeLocator， VGDetector。

实验可分为下面3部分：

• RQ1: Can ContraFlow outperform existing learning-based vulnerability detection approaches? 分为二分类对比和漏洞定位对比

• RQ2：How do different settings affect ContraFlow’s overall performance? 分为几个消融实验：有没有对比学习，不同的path selection策略，预训练模型的不同

• RQ3：How do different dataset scales for training affect the performance of ContraFlow?在不同训练数据集大小下模型性能的不同

评估指标，针对二分类用到了：IF，MK， F1 3个指标，IF = Recall + TNR - 1，MK = Precision + TNA - 1

对漏洞定位用到了MSR, MSP, MIOU 3个指标：

• $s_d$ 为检测出的漏洞行，$s_l$ 为标注漏洞行

• MSR = $\frac{|s_d\cap s_l|}{|s_l|}$

• MSP = $\frac{|s_d\cap s_l|}{|s_d|}$

• MIOU = $\frac{|s_d\cap s_l|}{|s_d\cup s_l|}$

实验结果就列出部分

4.1.RQ1

二分类对比如下：

而漏洞定位结果如下（只对比了VulDeeLocator和IVDetect），横轴代表漏洞行数量：

4.2.RQ2

首先，没有对比学习会导致ContraFlow性能下降明显，对于path selection中随机sample path或者在后面不进行path sensitive分析都会让ContraFlow性能明显下降，下面3个ContraFlow-CodeBert xxx 是指用CodeBert直接向量化node statement，不考虑AST子树特征，可以看到不考虑AST子树会让ContraFlow的性能出现下降

4.3.RQ3

训练数据量减少时模型性能出现了明显的下降

五.总结

这篇Paper中作者提出了ContraFlow，一个新的path-sensitive的漏洞检测器，框架中作者多次用到attention机制，同时作者将更多程序分析知识应用到漏洞检测中，用可行的value flow path集合作为程序表示显著提升了检测器的效果。

参考文献

[1] Xiao Cheng, Guanqin Zhang, Haoyu Wang, and Yulei Sui. 2022. Path-sensitive code embedding via contrastive learning for software vulnerability detection. In Proceedings of the 31st ACM SIGSOFT International Symposium on Software Testing and Analysis (ISSTA 2022). Association for Computing Machinery, New York, NY, USA, 519–531. https://doi.org/10.1145/3533767.3534371

[2] Changsheng Li, Handong Ma, Zhao Kang, Ye Yuan, Xiao-Yu Zhang, and Guoren Wang. 2020. On Deep Unsupervised Active Learning. In Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI-20, Christian Bessiere (Ed.). International Joint Conferences on Artificial Intelligence