0%

基于图神经网络的漏洞挖掘研究-论文笔记

摘要

基于图神经网络的漏洞挖掘是当下的研究热点之一,将源代码表示为图结构(抽象语法树,数据流图,控制流图,代码属性图等),可以更好的表示代码中的语法、语义、结构等信息,本文分享几篇基于GNN的漏洞挖掘文章,旨在总结其中的方法,学习其中的研究思路,并发掘其中存在的问题。

现有的源码漏洞挖掘存在的问题

  • 漏洞检测是一项试图通过审计软件代码或分析软件的执行过程来发现软件漏洞的任务。它需要对程序语义的理解和推理。因此,用于检测bug的方法不能直接用于漏洞检测。

  • 基于规则和漏洞挖掘:需要专家去手动添加漏洞规则,不可能罗列出所有的漏洞规则,同时时间和人力成本过高。

  • 传统的深度学习方法,诸如GRU、LSTM、BiLSTM等,将源代码视作文本序列进行学习,忽略了代码的结构信息。

  • 现有的方法工作原理是在预处理过程中将代码转换为图结构的输入,并直接通过图神经网络(GNNs)来学习其中的语义、结构等信息。然而,大多数基于GNN的方法只允许消息传递在𝑘-hop邻域中(其中k通常是一个小数),因此只捕获局部邻域信息,忽略节点之间的全局交互。

BGNN4VD: Constructing Bidirectional Graph Neural-Network for Vulnerability Detection

Information and Software Technology(3区)

论文方法

截屏2021-08-08 下午6.40.21.png

通过AST,CFG,DFG提取源代码中的语法和语义信息,然后通过向量化作为BGNN的输入,通过在传统的图神经网络(GNN)的基础上引入边,来学习脆弱代码与非脆弱代码之间的不同特征,利用卷积神经网络(CNN)通过分类器进一步提取特征和检测漏洞。

  • Phase 1: 语法信息用AST表示,语义信息可以用CFG和DFG表示。通过结合AST、CFG和DFG构造出代码复合图(CCG)。CCG由一组AST节点和不同类型的边组成。语句和谓词节点由CFG和DFG连接。因此,边的类型反映了不同节点之间的语法和语义关系。
  • Phase 2: 使用Word2Vev将每个节点的文本特征表示为向量。为了进一步反映单词之间的抽象语义关系,考虑了每个节点的类型,并将其添加到转换后的向量中,以计算初始节点的状态。例如代码result = num1 result和num1都是标识符Identifier,因此,将鸡诶单的类型(标识符、函数等)和节点的文本特征表示为向量后,进行拼接作为节点的新特征。
  • Phase 3: 在获得了可作为深度学习模型输入的初始节点表示后,使用BGNN来训练检测模型来学习聚合信息,因为它适合于非顺序输出和语义学习。**由于源代码比自然语言更具有逻辑性和结构性,因此漏洞代码片段通常与它们的上下文相关**。因此模型需要能够前后处理序列,获得更深层次的表示,并获得更好的鲁棒性。论文构建了一个双向图神经网络(BGNN) 。通过BGNN最后输出节点的最终表示。
  • Phase 4: 使用上一步的节点表示训练分类器用于漏洞分类。

A Unified Framework to Learn Program Semantics with Graph Neural Networks

ASE-2020-南洋理工大学

论文方法

截屏2021-08-08 下午10.17.33.png

使用图神经网络(GNNs)来表示程序已经在许多应用程序中取得了最先进的性能,然而,目前,还缺乏一个针对不同应用程序的GNN的统一框架,此外,大多数现有的基于GNN的方法忽略了与节点的全局关系,限制了模型学习丰富的语义,作者提出了一个统一的框架来构建两种类型的图来捕获各种SE应用程序的丰富代码语义。

该方法由三部分组成:

  1. 图构造,将源代码编码为基于语义的图和基于注意的图;
  2. 语义学习,它将所构造的图中的消息与各自的消息传递函数融合,以学习全面的代码语义;
  3. 应用程序,使用学习到的代码语义进行漏洞识别,代码克隆检测等;

在图构造部分包括*基于语义的静态图**基于注意力机制的动态图***两部分。

  • 基于语义的静态图:通过BiLSTM将原始的代码属性图(CPG) 𝑔(𝑉,𝑨)表示为基于语义的静态图𝑔(𝑉,𝑯,𝑨),其中 𝑯 是初始化节点特征矩阵。
  • 基于注意力机制的动态图:因为GNN通常忽略节点之间的全局交互,因为它们只在消息传递期间考虑𝑘-hop邻域。为了更好地捕获源代码的全局语义,构建基于注意的动态图。利用注意力机制,更新原始图的邻接矩阵,生成动态图。

语义学习过程中,分别在静态图和动态图上使用消息传递函数,最后通过混合消息传递将静态消息传递和动态消息传递的结果进行融合,然后通过GRU,以更新节点表示。

A Hybrid Graph Neural Network Approach for Detecting PHP Vulnerabilities

帝国理工学院-2020

论文方法

截屏2021-08-08 下午11.41.55.png

DeepTiver检测源代码中函数和文件级别粒度上的SQLi、XSS和OSCI漏洞,它分为两个关键部分:一个是运行在源代码token的线性序列上的门控递归单元(GRU),另一个对源代码的控制流图(CFG)运用图卷积网络(GCN)。每个部分都为模型提供了一个不同的机制,以有效地检测多种类型的漏洞,最后将GRU和GCN结合在一个新的混合架构中,能够利用这两种技术的优势。

Devign: Effective Vulnerability Identifification by Learning Comprehensive Program Semantics via Graph Neural Networks

南洋理工大学-2019-NIPS

论文方法

Devign一种基于通用图神经网络的模型,可通过学习丰富的代码语义来进行图级分类表示形式。它包括一个Conv模块,可以有效地学习丰富的节点表示并从中提取有用的特征,以进行图级分类。

截屏2021-08-10 下午11.21.26.png

在复合代码表示中,以ASTs为主干,显式地将不同级别的程序控制和数据依赖编码成一个异构边的联合图。对比以前的模型,有助于尽可能广泛地捕获漏洞的类型和模式,并能够通过图神经网络更好的学习节点的表示。

截屏2021-08-10 下午11.21.33.png

在程序分析中,程序的各种表示被用来在文本代码背后显示更深层次的语义,其中经典的概念包括AST、控制流和数据流图,它捕获源代码的不同标记之间的句法和语义关系。大量的漏洞,如内存泄漏,非常微妙,不考虑复合代码语义就无法发现。单用ASTs只能找到不安全的参数。 通过将AST与控制流图相结合,能够覆盖到该类型的漏洞。通过进一步集成这三个代码图,可以描述除两个代码图之外的大多数漏洞类型所需要的外部信息(即依赖于运行时属性的竞争条件和设计错误,这些错误在没有详细说明程序的预期设计的情况下很难建模),除此之外,作者定义了一种自然代码序列,按照编程逻辑顺序。它可以补充经典的表示,因为它独特的平面结构以“人类可读”的方式捕捉代码令牌的关系。

一点思考

  • 出现漏洞的代码往往只是其中的一小部分,可以通过切片技术来提出关键的代码行,然后通过结构补全,实现代码语义和结构信息上的完整,这样可以减少无关代码的数量,提高检测的准确率和效率。

  • 很多的很多方法都是针对单个函数的漏洞检测,无法进行跨函数甚至是跨文件的漏洞检测,但现实中的很多漏洞都是多函数关联的。

  • 针对大型的系统以及应用,可能使用到不止一个语言,同时服务端和客户端分离,如果单纯的分析一端容易导致误报和漏报,因此跨语言repo的漏洞挖掘也是目前的一大问题。

未完待续