Code Embedding研究系列11-ContraFlow

2023-11-04

一.引言

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( π 1 : 2 − 6 − 13 \pi_1 : 2-6-13 π1:2613 , π 2 : 2 − 6 − 15 \pi_2 : 2-6-15 π2:2615, π 3 : 3 − 9 − 13 \pi_3 : 3-9-13 π3:3913, π 4 : 3 − 9 − 15 \pi_4 : 3-9-15 π4:3915),作者将这4条路径2次输入到VPE中,每次使用不同的dropout,可以分别获取向量 v π 1 , v π 2 , v π 3 , v π 4 v_{\pi_1}, v_{\pi_2}, v_{\pi_3}, v_{\pi_4} vπ1,vπ2,vπ3,vπ4 v π 1 + , v π 2 + , v π 3 + , v π 4 + v_{\pi_1}^+, v_{\pi_2}^+, v_{\pi_3}^+, v_{\pi_4}^+ vπ1+,vπ2+,vπ3+,vπ4+。其中,VPE的训练目标是相似的embedding vector距离相近 ( v π 1 , v π 1 + v_{\pi_1}, v_{\pi_1}^+ vπ1,vπ1+),而不同的vector距离较远 ( v π 1 , v π 3 + v_{\pi_1}, v_{\pi_3}^+ vπ1,vπ3+),这一步采用的loss是NCE loss。

  • 阶段(b) 首先根据self-supervised active learning获取 π 1 − π 4 \pi_1 - \pi_4 π1π4 的排序,之后这些path会进行可行性分析。其中 π 2 \pi_2 π2 π 3 \pi_3 π3 不可行, 2 − 6 − 15 2-6-15 2615 6 6 6 需要满足条件 FLG 15 15 15 要满足条件 !FLG π 3 \pi_3 π3 反之亦然。

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

三.ContraFlow

3.1.Contrastive Value-Flow Embedding

3.1.1.path vectorizing

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

在这里插入图片描述
Local Encoding主要目的是独立地向量化每个statement,向量化statement s s s 用到了 s s s 的AST subtree,下图给出了一个示例。首先对于 s s s 的AST子树中的一个结点 n i n_i ni,它的初始结点向量 v n i v_{n_i} vni 通过 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=jCi{i}αij.W.vnj

  • C i C_i Ci n i n_i ni 的子结点

  • W W W 是VPE的参数

  • α i j \alpha_{ij} αij n i n_i ni n j n_j nj 之间的attention值

在这里插入图片描述
之后按如下方式生成 s s s 的summary vector v s m v_{sm} vsm

v s m = [ 1 N ∑ i = 1 N v n i ′ ∣ ∣ max ⁡ j = 1 N v n j ′ ] v_{sm} = [\frac{1}{N} \sum_{i=1}^N v_{n_i}^{'} || \max_{j=1}^N v_{n_j}^{'} ] vsm=[N1i=1Nvni∣∣j=1maxNvnj]

其中用到的模型可能有多层 ( L L L 层),statement vector v j k = ∑ l = 1 L v s m l v_{jk} = \sum_{l=1}^L v_{sm}^l vjk=l=1Lvsml,而最终statement向量 v s = d r o p o u t ( W . v j k ) v_s = dropout(W.v_{jk}) vs=dropout(W.vjk)

Global Encoding根据每个statement的向量生成path向量,前一步计算出了每个statement的向量 v s 0 , v s 1 , . . . , v s N v_{s_0}, v_{s_1}, ..., v_{s_N} vs0,vs1,...,vsN,这一步通过BGRU将 v s 0 , v s 1 , . . . , v s N v_{s_0}, v_{s_1}, ..., v_{s_N} vs0,vs1,...,vsN 转化为 h s 0 , h s 1 , . . . , h s N h_{s_0}, h_{s_1}, ..., h_{s_N} hs0,hs1,...,hsN,这一组向量描述了一组def-use(data-dependence)关系。之后通过attention机制将所有hidden向量聚合为path向量 π \pi π

3.1.2.Contrastive Value-Flow Embedding Algorithm

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

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

s i m ( v π i , v π j ) = v π i T . v π j ∣ ∣ v π i ∣ ∣ . ∣ ∣ v π j ∣ ∣ sim(v_{\pi_i}, v_{\pi_j}) = \frac{v_{\pi_i}^T. v_{\pi_j}}{||v_{\pi_i}||. ||v_{\pi_j}||} sim(vπi,vπj)=∣∣vπi∣∣.∣∣vπj∣∣vπiT.vπj

path π i \pi_i πi 的 loss定义如下:

l o s s ( v π i ) = − log ⁡ exp ⁡ ( s i m ( v π i , v π i + ) ) ∑ k = 1 B exp ⁡ ( s i m ( v π i , v π k + ) 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}^{+})} loss(vπi)=logk=1Bexp(sim(vπi,vπk+)exp(sim(vπi,vπi+))

B B B 为1个Batch的数量,一个Batch的Loss为

L = 1 B ∑ i = 1 B l o s s ( π i ) L = \frac{1}{B} \sum_{i=1}^B loss(\pi_i) L=B1i=1Bloss(πi)

3.2.Value-Flow Path Selection

3.2.1.Value-Flow Active Learning

这一步主要目的是选取出具有代表性的path,参考的是文献 [ 2 ] ^{[2]} [2],给定通过VPE向量化好的path集合 H = [ v π 1 , . . . , v π n ] ∈ R n × d H = [v_{\pi_1}, ..., v_{\pi_n}] \in R^{n \times d} H=[vπ1,...,vπn]Rn×d,选取出 H ′ ∈ R k × d ⊆ H H^{'} \in R^{k \times d} \subseteq H HRk×dH,这里引入了1个encoder-decoder模型。需要优化的参数为2个矩阵 Q , P ∈ R n × n Q, P \in R^{n \times n} Q,PRn×n,优化目标包括最小化下面目标的距离:

  • H H H Q . H Q.H Q.H

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

  • H H H 和decoder输出 G G G

训练完毕后,作者为 Q , P Q, P Q,P 的每一列计算 l2-norm 值,生成2个ranking vector q ^ , p ^ ∈ R n \hat q, \hat p \in R^n q^,p^Rn,之后,这两个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 π 1 , . . . , v π N ] ∈ R N × d e m b V = [v_{\pi_1}, ..., v_{\pi_N}] \in R^{N \times d_{emb}} V=[vπ1,...,vπN]RN×demb N N N 是path数量。而分类器主要包含1个Muti-head attention层+1个soft attention层+1个线性分类层

muti-head的计算过程如下:

h i = A t t n ( V . W i Q , V . W i K ) . V . W i V h_i = Attn(V.W_i^Q, V.W_i^K).V.W_i^V hi=Attn(V.WiQ,V.WiK).V.WiV

V ′ = [ h 1 ∣ ∣ . . . ∣ ∣ h h ] . W o V^{'} = [h1 || ... || h_h].W^o V=[h1∣∣...∣∣hh].Wo

V ′ ∈ R N × d c t x V^{'} \in R^{N \times d_{ctx}} VRN×dctx 是muti-head层输出, W i Q , W i K , W i V , W o W_i^Q, W_i^K, W_i^V, W^o WiQ,WiK,WiV,Wo是muti-head层参数, soft attention层计算如下:

α i c = exp ⁡ ( v π i T ) . a c ∑ j = 1 N exp ⁡ ( v π j T ) . a c \alpha_i^c = \frac{\exp(v_{\pi_i}^T).a_c}{\sum_{j=1}^{N} \exp(v_{\pi_j}^T).a_c} αic=j=1Nexp(vπjT).acexp(vπiT).ac

v c = ∑ i = 1 N α i c . v π i v_c = \sum_{i=1}^{N} \alpha_i^c. v_{\pi_i} vc=i=1Nαic.vπi

v c v_c vc 是soft attention层输出的整个程序的向量, a c a_c ac
是soft attention层的参数,之后 v c v_c vc 会输入一个全连接层+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_d sd 为检测出的漏洞行, s l s_l sl 为标注漏洞行

  • MSR = ∣ s d ∩ s l ∣ ∣ s l ∣ \frac{|s_d \cap s_l|}{|s_l|} slsdsl

  • MSP = ∣ s d ∩ s l ∣ ∣ s d ∣ \frac{|s_d \cap s_l|}{|s_d|} sdsdsl

  • MIOU = ∣ s d ∩ s l ∣ ∣ s d ∪ s l ∣ \frac{|s_d \cap s_l|}{|s_d \cup s_l|} sdslsdsl

实验结果就列出部分

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

本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)

Code Embedding研究系列11-ContraFlow 的相关文章

随机推荐

  • UE4Material_材质属性(1)

    材质中的属性 物理材质 Phys Material 物理材质 与该材质关联的物理材质 物理材质 Physical Material 提供了物理属性的定义 例如碰撞 弹力 以及其他基于物理的方面会保留多少能量 物理材质 Physical Ma
  • 一篇教会你,Redis主从、哨兵、 Cluster集群。

    前言 大家好 今天跟小伙伴们一起学习Redis的主从 哨兵 Redis Cluster集群 Redis主从 Redis哨兵 Redis Cluster集群 1 Redis 主从 面试官经常会问到Redis的高可用 Redis高可用回答包括两
  • TCP的三次握手及四次挥手总结(从抓包角度理解)

    目录 TCP报文首部 TCP连接 传输及断开过程图 TCP状态图 三次握手过程理解 四次挥手过程理解 从抓包来理解TCP建立连接 数据传输以及断开连接的过程 建立连接过程 数据传输过程 连接断开过程 为什么连接的时候是三次握手 关闭的时候却
  • Keras查看model weights .h5 文件的内容

    Keras的模型是用hdf5存储的 如果想要查看模型 keras提供了get weights的函数可以查看 for layer in model layers weights layer get weights list of numpy
  • 多进程浏览器框架

    为什么浏览器采用多进程模型 转载于 http www wtoutiao com p s57age html Google Chrome源码剖析 一 多线程模型 转载于 http www ha97 com 2908 html 主流浏览器多进程
  • 【广州华锐互动】无人值守变电站AR虚拟测控平台

    无人值守变电站AR虚拟测控平台是一种基于增强现实技术的电力设备巡检系统 它可以利用增强现实技术将虚拟信息叠加在真实场景中 帮助巡检人员更加高效地完成巡检任务 这种系统的出现 不仅提高了巡检效率和准确性 还降低了巡检成本和风险 传统的变电站巡
  • TPM功能介绍

    文章来源 TPM功能介绍 百度文库 http wenku baidu com link url bQMQyb0A3gto0CCC2CN5ojpUrgHsh8BMXmejpFaqLS52v 013bXPHoRr36r0F0UrgPr8U6rv
  • MATLAB——FFT(快速傅里叶变换)

    基础知识 FFT即快速傅里叶变换 利用周期性和可约性 减少了DFT的运算量 常见的有按时间抽取的基2算法 DIT FFT 按频率抽取的基2算法 DIF FFT 1 利用自带函数fft进行快速傅里叶变换 若已知序列 x 4 3
  • 利用ChatGPT协助编写单元测试

    ChatGPT自从2022年推出以来受到很多人的喜欢 此篇博客重点介绍如何修改Prompt来自动生成较理想的单元测试 如下图所示的一段代码 该class中有一个public方法toLocale 其余都是private方法 toLocale
  • 编写代码的几个tip

    使用的大多是MVC的模式 那么视图就只管视图 逻辑就只管逻辑 一个自定义的cell 上面放了一个button button的点击事件用一个delegate在viewcontroller中来实现 比如先要变化cell的样式 那么代理方法中 不
  • 攻防世界WEB入门

    1 view source X老师让小宁同学查看一个网页的源代码 但小宁同学发现鼠标右键好像不管用了 WP 按F12即可在elements中看到flag 2 robots X老师上课讲了Robots协议 小宁同学却上课打了瞌睡 赶紧来教教小
  • ​LeetCode刷题实战214:最短回文串

    算法的重要性 我就不多说了吧 想去大厂 就必须要经过基础知识和业务逻辑面试 算法面试 所以 为了提高大家的算法能力 这个公众号后续每天带大家做一道算法题 题目就从LeetCode上面选 今天和大家聊的问题叫做 最短回文串 我们先来看题面 h
  • UE4 UMG中使用富文本

    UE4 UMG中使用富文本 一 新建DateTable 二 添加字体样式 注意 第一个RowName必须为Default 字体样式必须赋值 否则会乱码 我们将Default的字体改为白色 Red字体改为红色 字号改小 三 使用 拖一个富文本
  • serverTimezone设置

    在安装完mysql第一次使用IDEA进行数据库连接发现 You must configure either the server or JDBC driver via the serverTimezone configuration pro
  • uview2.0封装网络请求(微信小程序最新登录方式)

    一 网络请求和相应拦截器 此vm参数为页面的实例 可以通过它引用vuex中的变量 module exports vm gt 初始化请求配置 uni u http setConfig config gt config baseURL http
  • 计算机二级C语言三天能过吗,学工干货丨如何三天通过计算机二级

    原标题 学工干货丨如何三天通过计算机二级 不可能的 想都不要想 三天怎么可能过 不信你问考完了的 他们马上就可以 进行计算机二级考试成绩查询啦 计算机二级是什么 计算机二级怎么考 迷惘 彷徨 别怕 今天学工菌来助攻 考生在考后50个工作日
  • 博客园自定义主题代码

    发一下好看的 要开通js权限 皮肤用simple memory 最好禁用模板 侧边栏
  • 基于小脑模型神经网络的轨迹跟踪研究(Matlab代码实现)

    欢迎来到本博客 博主优势 博客内容尽量做到思维缜密 逻辑清晰 为了方便读者 座右铭 行百里者 半于九十 本文目录如下 目录 1 概述 2 运行结果 3 参考文献 4 Matlab代码实现 1 概述 1 在对人类神经学的研究中 得知它由一些神
  • class加载过程

    loading class文件 从硬盘 加载到 内存 linking 1 verification 校验 检查满不满足class文件的格式 2 preparation 将 静态变量 赋默认值 默认值是0 3 resolution 将 常量池
  • Code Embedding研究系列11-ContraFlow

    Path Sensitive Code Embedding via Contrastive Learning for Software Vulnerability Detection 一 引言 1 1 现有方法及其局限 1 2 作者的解决方