code embedding研究系列九-Reveal

最新推荐文章于 2022-07-17 23:11:10 发布

I still …

最新推荐文章于 2022-07-17 23:11:10 发布

阅读量1.8k

点赞数 1

分类专栏：静态代码检测 code embedding 文章标签：深度学习人工智能 python

本文链接：https://blog.csdn.net/qq_44370676/article/details/118223852

版权

静态代码检测同时被 2 个专栏收录

28 篇文章 26 订阅

订阅专栏

code embedding

14 篇文章 50 订阅

订阅专栏

一.背景

漏洞检测具有重大的意义，针对DLVP（Deep Learning Vulnerability detection）任务，作者在对现有的漏洞检测方法（VulDeepecker, SyScVR）等测试时发现了一些问题。

在sard等合成数据集训练的模型用在真实场景下（FFMPeg, Qemu, Linux等）效果很差
在用解释方法（LEMNA等）来解释漏洞检测方法时经常发现模型学习到了无关的特征
训练/测试数据包含了许多重复
现有的方法没有解决样本类别不平衡问题

论文贡献

提出了新的漏洞检测方法Reveal
利用Chromium和Debian的修复commit构造数据集

二.数据集

2.1 现有数据集

针对数据集，作者统计的一些结果
在这里插入图片描述
有合成，半合成，真实 & 类别平衡，真实 & 类别不平衡

合成类：包括Juliet。使用已知漏洞pattern构造。
半合成：包括SARD和NVD。它们是从软件产品中提取，并做了一定修改。
真实：从代码仓库（github）的commit中提取，来自一些bug fix版本。

2.2 Reveal数据集

从Linux Debian Kernel 和Chromium的vulnerabilitiy fixed patches中构造。

对于Chromium，从Bugzilla中提取。
对于Linux Debian Kernel，从Debian security tracker中提取。

该数据集是针对function（单个函数）进行分类的，构造过程如下：

对于每个patch，从选取其vulnerable版本到fixed版本中被修改过源文件（.c, .cpp）和头文件（.h）。
对于被修改过的function，将该function修改前的标注为vulnerable。fix之后的标注为clean，其余未在patch中出现的function均标注为clean。

示例如图
在这里插入图片描述
$version_{k-1}$ 表示有漏洞的源文件， $version_k$ 表示fix版本。ham_0会被标注为vulnerable，ham_1，egg，spam会被标注为clean。

数据集统计信息如下：
在这里插入图片描述

三.现有方法

现有的方法大致可分为token-based和graph-based

3.1 token-based

token-based模型将源代码当成一个简单的token序列。而序列长度则会很大程度上影响模型发挥，因为源代码token序列可能相当长。所以就有了code slicing（VulDeepecker, SySeVR）。slicing的初衷是不考虑每个代码行，预处理的时候忽略掉许多无关行。slicing技术上从一些interesting points（API调用，数组索引，指针使用）出发。尽管如此，token-based方法将源代码视为序列，容易丢失语义信息。做过slice也会丢失一些依赖。

3.2 graph-based

graph-based模型将代码视为一个基于句法和语义依赖的图。句法依赖包括AST（抽象语法树），语义依赖包括CFG（控制流图），DFG（数据流图），PDG（程序依赖图），Def-Use chain graph。比如Devign使用了CPG（代码依赖图）。一般来说，使用的依赖信息越多，检出率越高，但是本身消耗的资源也会更多。

3.3 存在的问题

都存在vocabulary explosion（词表爆炸）问题。词表主要包括一些identifier（变量名，函数名，常量值等）。比如int count = 0;中包括了变量名count。而实际应用中变量名有无限种可能，如果简单粗暴的添加进词表那么100%要爆炸，有一种解决方案（VulDeepecker, SySeVR中采用的）是符号化。比如对于变量名count，将其用VAR1替代，对于自定义函数名function，将其用FUNC1替代，以此类推。

将代码转化为token序列后就是要向量化了，向量化主流的方案就是用embedding layer。这个embedding layer可以采用直接用下游任务（预测代码是否有漏洞）来训练，也可以用Word2Vec甚至Bert来先预训练。VulDeePecker和SySeVR用Word2Vec来将每个token向量化。Devign直接用Word2Vec来向量化一个statement的所有token（有点没搞懂Word2Vec是怎么对序列向量化的）。

向量化之后就是训练了，训练就需要损失函数，现有的方案采用交叉熵（cross entropy）或者带正则化的交叉熵损失函数。但仅仅靠交叉熵损失函数只能区分是否包含漏洞，并不能让模型学习到有漏洞和没有漏洞的代码的区别。

此外还有一个问题就是数据不平衡，因为数据集中包含漏洞和不包含漏洞的代码比例非常不协调。

四.ReVeal

总体过程如下图所示
在这里插入图片描述
需要注意的是作者这里进行了2个阶段的训练

第一个阶段是pre-train（Phase-I）。主要是训练GGNN，目标是能获得良好的graph embedding。
第二个阶段是train（Phase-II）。主要是训练MLP，目标是获得良好分类结果，SMOTE过采样也主要是应用在Phase-II。

4.1 Feature Extraction (Phase-I)

这个阶段的目标是将源代码转化成一个向量，这个向量保存了代码的语义（semantic）和句法（syntactic）信息。因此作者用到了CPG（代码属性图）。

通常，CPG表示为 $G = (V, E)$ 。V是结点（英文vertices或者nodes）和边集合（edges）。与Devign不同的是，这里每个结点不仅包含原始的一行代码（statement或者code fragment）。还包括statement类型（即这一行代码大概是什么语句，ArithmeticExpression或者CallStatement等等）。

所以对于一个node $v$ 的向量化包括2部分

用one-hot将其类型向量化，得到向量 $T_v$
用word2vec向量化code fragment内容，得到向量 $C_v$
拼接（concat） $C_v$ 和 $T_v$ 得到结点向量 $x_v$

对于用Word2Vec向量化，代码里如下：

node_split = nltk.word_tokenize(node_content)
nrp = np.zeros(100)
for token in node_split:
try:
   embedding = wv.wv[token]
except:
   embedding = np.zeros(100)
nrp = np.add(nrp, embedding)
if len(node_split) > 0:
   fNrp = np.divide(nrp, len(node_split))
else:
   fNrp = nrp

从node_split = nltk.word_tokenize(node_content)和fNrp = np.divide(nrp, len(node_split))可知对于一行代码（以int a = 10）为例。会将statement先解析成一个token序列，之后用Word2Vec对每个token向量化，然后取所有token向量的均值。

之后便用GGNN来进行结点向量的聚合，GGNN的计算过程前面总结过：图神经网络的计算过程

简单来说，经过GGNN的处理，每个结点的向量由 $x_v$ 变成 $x_v^{'}$
$x_v^{'} = GRU(x_v, \sum\limits_{(u,v) \in E} g(x_u) )$

GRU内部公式就不展开了，在RNN序列任务种 $h_t = GRU(input_i, h_{t-1})$ 。 $u$ 是 $v$ 邻居结点， $g (\cdot)$ 是一个 transformation function。

最后一步就是用聚合函数（aggregate function）将每个结点的向量聚合成一个向量 $x_g$ ，作为整个CPG，也就是源代码的向量表示。

$x_g = \sum\limits_{v \in V} x_v^{'}$

这里在论文中，作者用向量总和（element-wise summation）作为聚合函数，而实际上在代码里，聚合函数是一个可配置参数。

在这里插入图片描述

4.2 Training (Phase-II)

现实数据集中的样本不平衡问题非常严重，不包含漏洞的代码数量远超过包括漏洞的。

因此作者将训练阶段分为2部分

Reducing Class Imbalance：采用re-sampling（不知道该如何翻译）平衡训练集vulnerable和non-vulnerable的样本。
Representation Learning Model：基于平衡后的数据集训练一个可以很好的区分vulnerable和non-vulnerable样本的representation learning model。

4.2.1 Reducing Class Imbalance

在处理样本不平衡问题上用到了SMOTE算法。对于样本中的多数类（non-vulnerable），SMOTE会进行sub-sampling（随机删除一些样本），对于少数类（vulnerable），SMOTE会进行super-sampling（新合成一些样本）。直到每个类别的出现频率相等。算法如下图表示

在这里插入图片描述

4.2.2 Representation Learning Model

一个code fragment（一个method）的CPG（用 $G$ 表示）经过graph embedding（phase-I）后得到向量 $x_g$ ，作为 $G$ 的最终向量表示。但是vulnerable codes 和 non-vulnerable codes的向量在特征空间上有很大重合。

codes的特征向量经过TSNE降维后如下表示，一个点代表一个code fragment，红色部分为vulnerable codes，绿色为non-vulnerable codes。可以看到一个好的embedding是需要能够在特征空间区分开它们的。

在这里插入图片描述
分类函数如下所示
$\sigma(W .h(x_g) +b)$

$\sigma$ 为softmax
$W$ 和 $b$ 为最后一层全连接层的参数

其中 $h(x_g)$ 为一个全连接层，将 $x_g$ 投影到新的向量平面 $h_g$ 。

为了将non-vulnerable和vulnerable特征向量区别最大化。作者在训练模型时用到了triplet loss而不是softmax，记为 $L_{trp}$ 。每次训练需要用到的一个数据为一个3元组，记为 $(g, s a m e, d i f f)$ 对应 $(a, p, n)$ 。 $s a m e$ 和 $g$ 为同属一个类的样本， $d i f f$ 则相反。

$L_{trp} = L_{CE} + \alpha .L_{p} + \beta * L_{reg}$

$\alpha$ 和 $\beta$ 为超参数
$L_{CE}$ 为交叉熵损失
$L_{CE} = - \sum \hat y·log(y) + (1−\hat y)·log(1−y)$ ， $y$ 和 $\hat y$ 分别表示 $g$ 的标签和预测结果
$L_{reg}$ 为正则化损失
$L_{reg} = ||h(x_g)||+||h(x_{same})|| + ||h(x_{diff})||$ 。这里正则化损失主要用来限制 $h$ 即投影空间向量大小。
$L_{p} = | D(h(x_g), h(x_{same}))−D(h(x_g),h(x_{diff})) + \gamma|$
$L_{p}$ 为投影损失，是为了最大区分正类和负类样本在投影空间的差异。
$D(v_1,v_2) = 1−|\frac{v_1.v_2}{||v1||∗||v2||}|$

在这里插入图片描述

五. 实验设置

模型超参数大小
在这里插入图片描述
评估指标:

跟Devign一样，这是个针对function的二分类问题。
Accuracy, Precision, Recall, F1-score这4个指标用来评估模型。

六.实验结果

6.1 现有方法的有效性

作者在评估其它模型（vuldeepecker等）的性能时统一使用真实数据集，针对现有的模型训练数据的问题，作者给出了2个场景

Scenario-A (pre-trained models)
该模型在它本身的数据集上训练（比如VulDeepecker在sard数据集上训练），然后在真实数据集上测试
Scenario-B (re-trained models)
真实数据集上训练 + 真实数据集测试

Scenario-A的测试结果如下：
在这里插入图片描述
Scenario-B的测试结果如下：
同时，作者自己实现了一个Devign，并开源到github了。

可以看到，现有的方法泛化能力不强，在应用到真实数据集时效果有一定程度下降。

6.2 现有方法的局限性

6.2.1 数据重复

用slice和token-based方法都可能造成在训练集和测试集造成数据重复。作者做了一个统计，结果如下
在这里插入图片描述可以看到合成数据集比真实数据集多了很多重复。虽然重复会有利于漏洞分类任务，但不利于模型提取漏洞特征。

6.2.2 数据不平衡

数据集的情况再粘贴以下，看看Vul这一列，可以看到很多数据集中，正负样本比例不均。所以会造成模型分类时倾向于多数类。
在这里插入图片描述

6.2.3 学到不相关特征

为了选择好的DL模型来做漏洞分类，非常有必要理解模型是基于什么特征来做的分类。好的模型应该分配更多的权重给漏洞相关的特征。

作者通过LEMNA（一种解释方法）来解释token-based模型的分类结果。结果如下
在这里插入图片描述
对于graph-based模型，作者则用每个结点的激活值来表示，激活值越大，结点越关键。对一个被token-based方法错误分类而被graph-based方法正确方法分类的样本解释，结果如下

在这里插入图片描述

6.2.4 模型选择：缺乏区分度

这里主要展示不同的方法提取到的代码的特征向量对正类负类样本的区分度，即特征向量空间中两类代码是否很容易被区分开。作者用TSNE对不同方法提取的特征向量进行研究，并用centroid distance来衡量它们的效果，结果如下（再粘贴一次）

在这里插入图片描述
绿色为负类（无漏洞）的样本，红色为正类（有漏洞）。

6.3 解决上述问题

作者分别用SMOTE解决样本不均衡问题，而REVEAL本身就能解决其它的问题。

为了分别研究re-sampling和GGNN的效果，作者做了几组实验。不过实验结果里作者并未提到用什么模型替代了GGNN。

Re-balance的效果

实验结果如下（主要看F1-score），W/O表示without
在这里插入图片描述
跟其它模型（token-based + MLP,RF,SVM）的对比

七.预训练Word2Vec

代码如下：

def train(args):
    files = args.data_paths
    sentences = []
    for f in files:
        data = json.load(open(f))
        for e in data:
            code = e['code']
            sentences.append([token.strip() for token in code.split()])
    wvmodel = Word2Vec(sentences, min_count=args.min_occ, workers=8, size=args.embedding_size)
    print('Embedding Size : ', wvmodel.vector_size)
    for i in range(args.epochs):
        wvmodel.train(sentences, total_examples=len(sentences), epochs=1)
    if not os.path.exists(args.save_model_dir):
        os.mkdir(args.save_model_dir)
    save_file_path = os.path.join(args.save_model_dir, args.model_name)
    wvmodel.save(save_file_path)

这里大概就是将一个function的所有代码split成一个token序列当作一个sentence训练Word2Vec模型。

八.参考文献

Chakraborty, S. , Krishna, R. , Ding, Y. , & Ray, B. . (2020). Deep
Learning based Vulnerability Detection: Are We There Yet?.

I still …

关注

1
点赞
踩
21

收藏

觉得还不错? 一键收藏
20
评论
code embedding研究系列九-Reveal

Deep Learning based Vulnerability Detection:Are We There Yet?一.背景二.数据集2.1 现有数据集2.2 Reveal数据集一.背景漏洞检测具有重大的意义，针对DLVP（Deep Learning Vulnerability detection）任务，作者在对现有的漏洞检测方法（VulDeepecker, SyScVR）等测试时发现了一些问题。在sard等合成数据集训练的模型用在真实场景下（FFMPeg, Qemu, Linux等）效果很差
复制链接

扫一扫