Devign: Effective Vulnerability Identification byLearning Comprehensive Program Semantics via Graph Neural Networks
一.概述
最近几年软件漏洞的数量迅速增加,有的是通过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
∈
C
,
y
i
∈
Y
)
,
i
∈
{
1
,
2
,
.
.
.
,
n
}
((c_i, y_i)|c_i \in \mathcal{C}, y_i \in \mathcal{Y}), i \in \{1,2, . . . , n \}
((ci,yi)∣ci∈C,yi∈Y),i∈{1,2,...,n}
- C \mathcal{C} C 表示code function的集合。
- Y = { 0 , 1 } n \mathcal{Y}= \{0,1\}^n Y={0,1}n 表示每个function的label。1表示该function有漏洞,0表示没有。
C
\mathcal{C}
C 中的每一个function
c
i
c_i
ci 都会被映射为图
g
i
(
V
,
X
,
A
)
g_i(V,X,A)
gi(V,X,A)
g
i
∈
G
g_i \in \mathcal{G}
gi∈G
这里
- V V V 为图的顶点集合,集合中的每个顶点 v j v_j vj 用向量 x j ∈ R d x_j \in R^d xj∈Rd 表示。
- 顶点特征矩阵 X ∈ R m × d X \in R^{m \times d} X∈Rm×d。
- 邻接矩阵 A ∈ { 0 , 1 } k × m × m A \in \{0,1\}^{k \times m \times m} A∈{0,1}k×m×m。
这里面
- n n n 表示数据集样本数量
- m m m 表示每个图顶点数量
- d d d 表示每个顶点特征向量维度
- k k k 表示每个图边的种类数量(这里有AST边,CFG边,DFG边,NCS边)
e s , t p = { 1 i f n o d e v s a n d v t c o n n e c t e d v i a e d g e t y p e p 0 o t h e r w i s e e_{s,t}^p = \left\{ \begin{aligned} 1 && {if \; node \; v_s \; and \; v_t \; connected \; via \; edge \; type \; p}\\ 0 && {otherwise}\\ \end{aligned} \right. es,tp={10ifnodevsandvtconnectedviaedgetypepotherwise
L o s s = ∑ i = 1 n L ( f ( g i ) , y i ) + λ w ( f ) Loss = \sum_{i=1}^n\mathcal{L}(f(g_i),y_i) + \lambda\mathcal{w}(f) Loss=i=1∑nL(f(gi),yi)+λw(f)
- L \mathcal{L} L 是交叉熵损失函数
- f f f 表示神经网络分类函数
- λ \lambda λ 是可调整权重
- w \mathcal{w} w 是正则化项
2.2.Graph Embedding Layer
2.2.1.目标
这部分的核心在于将 function code
c
i
c_i
ci 映射为
g
i
(
V
,
X
,
A
)
g_i(V,X,A)
gi(V,X,A)
g
i
(
V
,
X
,
A
)
=
E
M
B
(
c
i
)
,
∀
i
=
{
1
,
.
.
.
,
n
}
g_i(V,X,A) = EMB(c_i), \forall i = \{1,...,n\}
gi(V,X,A)=EMB(ci),∀i={1,...,n}
- c i c_i ci 表示数据集中第 i i i 个function code
-
g
i
g_i
gi 表示function code 嵌入后的图向量,包括
V
,
X
,
A
V,X,A
V,X,A 三部分
- V V V 为图的顶点集合
- X ∈ R m × d X \in R^{m \times d} X∈Rm×d 为顶点向量
- A ∈ R k × m × m A \in R^{k \times m \times m} A∈Rk×m×m 为图的邻接矩阵
2.2.2.用到的结构
AST
AST是源代码树形结构的表示。通常,它是代码解析器用来理解程序的基本结构和检查语法错误的第一步。因此,它构成了生成许多其他代码表示的基础,并且AST表示的节点集合 V a s t V^{ast} Vast 也是这篇论文种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
a
s
t
V = V^{ast}
V=Vast 和邻接矩阵
A
A
A。 这里程序图中边的种类数量
k
=
4
k = 4
k=4。
程序图中每个结点
v
∈
V
v \in V
v∈V 有2个属性, code
和type
。
- code表示代码文本
- type表示AST结点类型,比如
IfStatement
,identifier
。
2.2.3.向量化
程序图 g i g_i gi 还差结点向量 X X X 就凑齐了。
这里首先用Word2Vec
在大型的源代码语料库上进行预训练。并用预训练好的模型对结点
v
v
v 的code
进行向量化。结点
v
v
v 的type
用 label encoding
来向量化。之后将 code
和 type
的向量进行 concatenate
就可以得到结点
v
v
v 的初始向量表示
x
v
x_v
xv
这样,就得到了程序图 g i ( V , X , A ) g_i(V,X,A) gi(V,X,A)
2.3.Gated Graph Recurrent Layers
获得程序图的初始向量表示后,作者选择使用 gated graph recurrent network 来学习图结点向量表示。
2.3.1.目标
这部分的最终目标是获得一个向量表示
H
i
T
=
G
G
R
L
(
X
)
H_i^T = GGRL(X)
HiT=GGRL(X)
- X ∈ R m × d X \in R^{m \times d} X∈Rm×d 表示第 i i i 个function code的结点向量。
- H i T ∈ R m z H_i^T \in R^{mz} HiT∈Rmz 表示第 i i i 个function code的向量输出。
2.3.2.计算过程
这里用 v j v_j vj 表示图 g i g_i gi 第 j j j 个结点。
- i i i 为function code索引号。
- j j j 为图结点索引号
在gated graph recurrent network中。
- 结点 v j v_j vj 的隐层向量用 h j t ∈ R z , t ∈ { 1 , T } , z ≥ d h_j^t \in R^z, t \in \{1, T\}, z \geq d hjt∈Rz,t∈{1,T},z≥d 表示
- T T T 表示time step数量。
- m m m 表示图结点数量
图的计算过程有如下3个公式:
h
j
1
=
[
x
j
⊤
,
0
]
⊤
h^1_j= [x^⊤_j,0]^⊤
hj1=[xj⊤,0]⊤
GRU中隐层向量
h
j
1
h^1_j
hj1 用
x
j
x_j
xj 补0初始化。
a j , p ( t − 1 ) = A p ⊤ . ( W p [ h 1 ( t − 1 ) ⊤ , . . . , h m ( t − 1 ) ⊤ ] + b ) a_{j,p}^{(t - 1)} = A_p^⊤.(W_p[h_1^{(t - 1)⊤},...,h_m^{(t - 1)⊤}] + b) aj,p(t−1)=Ap⊤.(Wp[h1(t−1)⊤,...,hm(t−1)⊤]+b)
h j t = G R U ( h j ( t − 1 ) , A G G ( { a j , p ( t − 1 ) } p = 1 k ) ) h_j^t = GRU(h_j^{(t - 1)}, AGG(\{a_{j,p}^{(t - 1)}\}_{p=1}^k)) hjt=GRU(hj(t−1),AGG({aj,p(t−1)}p=1k))
其中:
- p ≤ k p \leq k p≤k 为程序图的边类型
- ⊤ 表 示 转 置 ⊤表示转置 ⊤表示转置
- W p ∈ R z × z W_p \in R^{z \times z} Wp∈Rz×z 和 b b b 属于权重系数
- A p ∈ R m × m A_p \in R^{m \times m} Ap∈Rm×m 表示第 p p p 类边的邻接矩阵
- A G G ( . ) AGG(.) AGG(.) 表示聚合函数, 可以用 { M E A N , M A X , S U M , C O N C A T } \{MEAN, MAX, SUM, CONCAT\} {MEAN,MAX,SUM,CONCAT} 这里用 S U M SUM SUM
- H i T = { h j T } j = 1 m H^T_i=\{h^T_j\}^m_{j=1} HiT={hjT}j=1m
- a j , p ( t − 1 ) a_{j,p}^{(t - 1)} aj,p(t−1) 的shape还不是很清楚,希望有大佬能解答
2.4.The Conv Layer
这个模块的主要任务就是分类了,该论文的最终目标是确定一个function code c i c_i ci 是否有漏洞,yes or no。
作者在此定义了函数
σ
(
.
)
=
M
a
x
P
o
o
l
(
R
e
l
u
(
C
o
n
v
(
.
)
)
)
\sigma(.) = MaxPool(Relu(Conv(.)))
σ(.)=MaxPool(Relu(Conv(.)))
C
o
n
v
(
.
)
Conv(.)
Conv(.) 为一维卷积。
计算label y i ∼ \overset{\sim}{y_i} yi∼ 通过如下公式:
Z i 1 = σ ( [ H i T , x i ] ) , . . . , Z i l = σ ( Z i ( l − 1 ) ) Z^1_i = \sigma([H^T_i, x_i]), . . . , Z^l_i=\sigma(Z^{(l−1)}_i) Zi1=σ([HiT,xi]),...,Zil=σ(Zi(l−1))
Y i 1 = σ ( H i T ) , . . . , Y i l = σ ( Y i ( l − 1 ) ) Y^1_i = \sigma(H^T_i), . . . , Y^l_i=\sigma(Y^{(l−1)}_i) Yi1=σ(HiT),...,Yil=σ(Yi(l−1))
y i ∼ = S i g m o i d ( A V G ( M L P ( Z i l ) ⊙ M L P ( Y i l ) ) ) \overset{\sim}{y_i} = Sigmoid(AVG(MLP(Z^l_i) \odot MLP(Y^l_i))) yi∼=Sigmoid(AVG(MLP(Zil)⊙MLP(Yil)))
其中
- [ H i T , x i ] [H^T_i, x_i] [HiT,xi] 表示两个向量的concatenation
- l l 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.实验结果
可以看到作者的评估也是分不同的数据集,并没有把所有的代码整合到一起