code embedding研究系列七-Devign

一.概述

最近几年软件漏洞的数量迅速增加，有的是通过CVE （Common Vulnerabilities and Exposure）公开报告的，有的是在专有网络内部发现的代码。

漏洞识别是安全领域中一个关键而又具有挑战性的问题。检测方法包括：

• 静态检测
• 动态检测
• 符号执行

1.1.相关工作

机器学习的发展也为漏洞识别提供了新的手段，早期的机器学习手段需要专家手动构造特征。但是，漏洞的表现形式千变万化，手动构造漏洞特征库也有点不切实际。

之后，也有人将源代码视为一个序列（a flat sequence），这种方法与NLP的方法相似，但这些方法在学习源代码高度多样和复杂的语义特征时具有局限性 。

• 源代码实际上比自然语言更具结构性和逻辑性，并且可以用AST（抽象语法树），CFG（控制流图），DFG（数据流图）等复杂结构来表示。有的漏洞需要在语义层面多维度综合分析。
• 之前的方法用已有的静态分析工具进行代码标注，这会造成很多误报。更限制了模型的准确率。

1.2.已有的代码表示方法

目前AST（抽象语法树），CFG（控制流图），DFG（数据流图）是用来捕获源代码token之间句法和语义联系最常用的3种结构。

大多数的漏洞（比如内存泄漏）的存在非常细微，如果不联合这3种代码结构很难发现。（这部分得好好研究下样本了）比如

• 如果仅仅使用AST只可以发现一些危险的参数
• 如果联合AST和CFG，则可以发现更多类型的漏洞，比如:资源泄漏，use-after-free漏洞等等
• 如果再联合DFG，那就除了少部分特殊情况（竞争条件漏洞取决于运行时属性，以及design errors）都可以描述。

1.3.作者提出的方法

这里，作者提出了Devign，一个新的静态漏洞检测模型，作者的贡献如下：

• 作者采用一种复合代码表示，以AST作为骨架，并加入CFG和DFG作为联合图表示。
• 实现了一个gated graph neural network + Conv模型来进行图分类。
• 手动从4个开源数据集收集数据，并花费大量时间人工标注。并实现Devign。

二.Devign Model

2.1.问题描述

这里，作者分析的是C语言代码。也是在function-level进行二分类（即判断每个function是否包含漏洞）

这里数据集表示为
$((c_i,y_i)|c_i\in \mathcal{C},y_i\in \mathcal{Y}),i\in \{1,2,...,n\}$

• $\mathcal{C}$ 表示code function的集合。
• $\mathcal{Y}=\{0,1{\}}^n$ 表示每个function的label。1表示该function有漏洞，0表示没有。

$\mathcal{C}$ 中的每一个function $c_i$ 都会被映射为图 $g_i(V,X,A)$
$g_i\in \mathcal{G}$

这里

• $V$ 为图的顶点集合，集合中的每个顶点 $v_j$ 用向量 $x_j\in R^d$ 表示。
• 顶点特征矩阵 $X\in R^{m\times d}$。
• 邻接矩阵 $A\in \{0,1{\}}^{k\times m\times m}$。

这里面

• $n$ 表示数据集样本数量
• $m$ 表示每个图顶点数量
• $d$ 表示每个顶点特征向量维度
• $k$ 表示每个图边的种类数量（这里有AST边，CFG边，DFG边，NCS边）

$$e_{s,t}^p=\{\begin{array}{rlr}1& & ifnode{v}_{s}and{v}_{t}connectedviaedgetypep\\ 0& & otherwise\end{array}$$

$$Loss=\sum _{i=1}^n\mathcal{L}(f(g_i),y_i)+\lambda w(f)$$

• $\mathcal{L}$ 是交叉熵损失函数
• $f$ 表示神经网络分类函数
• $\lambda$ 是可调整权重
• $w$ 是正则化项

2.2.Graph Embedding Layer

2.2.1.目标

这部分的核心在于将 function code $c_i$ 映射为 $g_i(V,X,A)$
$$g_i(V,X,A)=EMB(c_i),\forall i=\{1,...,n\}$$

• $c_i$ 表示数据集中第 $i$ 个function code
• $g_i$ 表示function code 嵌入后的图向量，包括$V,X,A$ 三部分
  • $V$ 为图的顶点集合
  • $X\in R^{m\times d}$ 为顶点向量
  • $A\in R^{k\times m\times m}$ 为图的邻接矩阵

2.2.2.用到的结构

AST

AST是源代码树形结构的表示。通常，它是代码解析器用来理解程序的基本结构和检查语法错误的第一步。因此，它构成了生成许多其他代码表示的基础，并且AST表示的节点集合 $V^{ast}$ 也是这篇论文种CFG和DFG的结点集合，如下图所示。

CFG

CFG描述了在程序执行过程中可能遍历的所有路径。这些路径由一些条件分支或者循环语句构成 if, for, switch, while 等等。在CFG中，节点表示语句和条件，它们通过有向边连接以表示控制流的转移。下图中，CFG边由绿色箭头表示。

DFG

DFG跟踪CFG中变量的使用情况。数据流是面向变量的，任何数据流都涉及对某些变量的访问或修改。DFG边表示对相同变量的后续访问或修改。下图用橙色箭头表示。比如：变量b在if条件和赋值语句中都有用到。

NCS(Natural Code Sequence)

除了上述3种表示代码结构和语义特征的表示。作者还加入了源代码的token序列。这样可以保存程序的逻辑结构。这里，作者仅仅将所有的终端结点用NCS边连接起来，以加入token序列。类似于链表的结构。

经过这些步骤以后，可以得到程序图的结点集合 $V=V^{ast}$ 和邻接矩阵 $A$。 这里程序图中边的种类数量 $k=4$。

程序图中每个结点 $v\in V$ 有2个属性， code和type。

• code表示代码文本
• type表示AST结点类型，比如 IfStatement, identifier。

2.2.3.向量化

程序图$g_i$ 还差结点向量 $X$ 就凑齐了。

这里首先用Word2Vec在大型的源代码语料库上进行预训练。并用预训练好的模型对结点 $v$ 的code进行向量化。结点 $v$ 的type用 label encoding 来向量化。之后将 code 和 type 的向量进行 concatenate 就可以得到结点 $v$ 的初始向量表示 $x_v$

这样，就得到了程序图 $g_i(V,X,A)$

2.3.Gated Graph Recurrent Layers

获得程序图的初始向量表示后，作者选择使用 gated graph recurrent network 来学习图结点向量表示。

2.3.1.目标

这部分的最终目标是获得一个向量表示
$$H_i^T=GGRL(X)$$

• $X\in R^{m\times d}$ 表示第 $i$ 个function code的结点向量。
• $H_i^T\in R^{mz}$ 表示第 $i$ 个function code的向量输出。

2.3.2.计算过程

这里用 $v_j$ 表示图 $g_i$ 第 $j$ 个结点。

• $i$ 为function code索引号。
• $j$ 为图结点索引号

在gated graph recurrent network中。

• 结点 $v_j$ 的隐层向量用 $h_j^t\in R^z,t\in \{1,T\},z\ge d$ 表示
• $T$ 表示time step数量。
• $m$ 表示图结点数量

图的计算过程有如下3个公式：

$$h_j^1=[x_j^{\top },0{]}^{\top }$$
GRU中隐层向量 $h_j^1$ 用 $x_j$ 补0初始化。

$$a_{j,p}^{(t-1)}=A_p^{\top }.(W_p[h_1^{(t-1)\top },...,h_m^{(t-1)\top }]+b)$$

$$h_j^t=GRU(h_j^{(t-1)},AGG(\{a_{j,p}^{(t-1)}{\}}_{p=1}^k))$$

其中：

• $p\le k$ 为程序图的边类型
• $\top 表示转置$
• $W_p\in R^{z\times z}$ 和 $b$ 属于权重系数
• $A_p\in R^{m\times m}$ 表示第 $p$ 类边的邻接矩阵
• $AGG(.)$ 表示聚合函数， 可以用 $\{MEAN,MAX,SUM,CONCAT\}$ 这里用 $SUM$
• $H_i^T=\{h_j^T{\}}_{j=1}^m$
• $a_{j,p}^{(t-1)}$ 的shape还不是很清楚，希望有大佬能解答

2.4.The Conv Layer

这个模块的主要任务就是分类了，该论文的最终目标是确定一个function code $c_i$ 是否有漏洞，yes or no。

作者在此定义了函数 $\sigma (.)=MaxPool(Relu(Conv(.)))$
$Conv(.)$ 为一维卷积。

计算label $\stackrel{\sim }{y_i}$ 通过如下公式：

$$Z_i^1=\sigma ([H_i^T,x_i]),...,Z_i^l=\sigma (Z_i^{(l-1)})$$

$$Y_i^1=\sigma (H_i^T),...,Y_i^l=\sigma (Y_i^{(l-1)})$$

$$\stackrel{\sim }{y_i}=Sigmoid(AVG(MLP(Z_i^l)\odot MLP(Y_i^l)))$$

其中

• $[H_i^T,x_i]$ 表示两个向量的concatenation
• $l$ 表示使用的卷积层个数

不过这一系列运算过程中向量的shape确实没弄懂，得好好琢磨下向量运算法则了。

三.数据准备

3.1.数据集

作者从Linux Kernel, QEMU, Wireshark, FFmpeg 4个开源数据集中收集数据，以function为单位。并对其进行人工标注（狠人）。采用以下步骤：

• 首先收集与安全相关的commit，将其标注为漏洞修复commit或非漏洞修复commit。
• 从已经标注好的commits提取包含漏洞（vulnerable ）和不包含漏洞（non-vulnerable）的function。

漏洞修复提交（VFCs）是修复潜在漏洞的提交，从中可以从代码修订之前版本的源代码中提取包含漏洞（vulnerable）的function。

非漏洞修复提交（non-VFCs）是不修复任何漏洞的提交，类似地，可以在修改之前从源代码中提取不包含漏洞（ non-vulnerable）的function。

作者采取以下步骤提取commit：

• commits筛选
  由于只有一小部分commit是漏洞相关的，因此排除了与安全无关的commit，这些commit的消息没有与一组安全相关的关键字（如DoS和injection）匹配。剩下的，更可能是与安全相关的，留给人工标记。
• 人工标注
  一个由四名专业安全研究人员组成的团队花了600个小时来执行两轮数据标记和交叉验证。

3.2.数据预处理

采用开源C/C++代码分析平台Joern来为数据集中的每个function提取AST和CFG。

由于Joern中的一些内部编译错误和异常，只能获得部分function的AST和CFG。这里过滤掉出现异常的function和AST，CFG中有明显错误的function。

原始的DFG中包括了所有和变量（variable）相关的边。因此边的数量会非常庞大，因此作者将DFG分为3类：

• DFG_R（LastRead）
  DFG_R表示变量每次出现的最后一次读取（last read of each occurrence of the variables）
• DFG_W（LastWrite）
  DFG_W表示变量每次出现的最后一次写入（last write of each occurrence of variables）
• DFG_C（ComputedFrom）
  在赋值语句中，左侧（lhs）变量由右侧（rhs）表达式赋值。DFG_C捕获lhs变量和每个rhs变量之间的关系。

最后，作者会删除数据集中node size大于500的function。

数据集统计信息如下：

四.实验

4.1.对比方法

• Metrics + Xgboost
  使用Joern为每个函数收集了4个复杂性度量和11个脆弱性度量，并利用Xgboost进行分类。类似于15个人工特征。

• 3-layer BiLSTM
  它将源代码视为自然语言，并将tokenize后的代码输入到双向LSTMs中，初始embedding通过Word2vec进行训练。

• 3-layer BiLSTM + Att
  在前面的基础上加上attention机制

• CNN
  它将源代码作为自然语言，利用bag-of-words获得代码标记的初始embedding，然后将其提供给CNNs进行学习

4.2.实验结果

可以看到作者的评估也是分不同的数据集，并没有把所有的代码整合到一起

五.参考文献

Zhou Y , Liu S , Siow J , et al. Devign: Effective Vulnerability Identification by Learning Comprehensive Program Semantics via Graph Neural Networks[J]. 2019.