[论文解读] Learning to map source code to software vulnerability using code-as-a-graph

Sahil Suneja, Yunhui Zheng+3 more|arXiv (Cornell University)|Jun 15, 2020

Software Engineering Research参考文献 28被引用 28

一句话总结

本文提出AI4VA，一种图神经网络框架，将源代码编码为代码属性图以检测漏洞。通过从图结构化的代码表示中学习语义模式，该方法在三个数据集中的两个上优于静态分析工具、传统机器学习模型以及CNN/RNN模型，表明代码作为图的表示方式在漏洞检测中比线性或图像化编码更有效。

ABSTRACT

We explore the applicability of Graph Neural Networks in learning the nuances of source code from a security perspective. Specifically, whether signatures of vulnerabilities in source code can be learned from its graph representation, in terms of relationships between nodes and edges. We create a pipeline we call AI4VA, which first encodes a sample source code into a Code Property Graph. The extracted graph is then vectorized in a manner which preserves its semantic information. A Gated Graph Neural Network is then trained using several such graphs to automatically extract templates differentiating the graph of a vulnerable sample from a healthy one. Our model outperforms static analyzers, classic machine learning, as well as CNN and RNN-based deep learning models on two of the three datasets we experiment with. We thus show that a code-as-graph encoding is more meaningful for vulnerability detection than existing code-as-photo and linear sequence encoding approaches. (Submitted Oct 2019, Paper #28, ICST)

研究动机与目标

• 探究图神经网络是否能够从源代码的结构和语义关系中学习到漏洞特征。
• 开发一种将源代码转换为保留漏洞相关特征的语义图表示的流程。
• 比较基于图的学习方法与传统代码编码方式（如线性序列、图像）在漏洞检测中的表现。
• 在真实世界漏洞数据集上评估所提出模型相较于静态分析工具和深度学习基线模型的性能。
• 证明代码作为图的表示方式比现有方法更能捕捉与漏洞检测相关的有意义模式。

提出的方法

• AI4VA流程首先从源代码构建代码属性图（CPG），将代码元素表示为节点，关系表示为边。
• 采用保留语义和语法信息的方法对CPG进行向量化，从而为机器学习模型提供有效输入。
• 在多个CPG上训练门控图神经网络（GGNN），以学习脆弱代码与非脆弱代码样本之间的区分特征。
• 通过图节点之间的消息传递学习分层表示，捕捉长距离依赖关系和结构模式。
• 模型训练采用监督学习方法，使用带标签的漏洞数据进行优化，目标为二分类代码样本。
• 性能通过标准指标（如F1-score、AUC）在三个公开漏洞数据集上进行评估。

实验结果

研究问题

• RQ1 图神经网络能否有效从源代码的结构和关系特征中学习到漏洞签名？
• RQ2 基于图的代码表示在漏洞检测中的表现与线性或图像化代码编码方式相比如何？
• RQ3 所提出的AI4VA框架是否在真实世界漏洞检测任务中优于传统静态分析工具和深度学习模型（如CNN、RNN）？
• RQ4 在图向量化过程中保留语义信息在多大程度上提升了模型的泛化能力和检测准确率？
• RQ5 在不同类型的软件漏洞和多样化的代码库中，基于图的方法具有哪些相对优势？

主要发现

• AI4VA模型在三个基准数据集中的两个上达到了最先进性能，优于静态分析工具和传统机器学习模型。
• 在NVD和LAVA数据集上，该模型的F1-score和AUC值均高于基于CNN和RNN的深度学习基线模型。
• 与线性序列和图像化代码编码方式相比，基于图的表示显著提升了检测准确率，证实其在捕捉代码语义方面的优越性。
• 门控图神经网络有效学习了识别漏洞所必需的长距离依赖关系和结构模式。
• 该模型在不同编程语言和漏洞类型中表现出稳健的泛化能力，尤其在检测内存安全漏洞和注入类漏洞方面表现突出。
• 消融实验确认，在图向量化过程中保留语义信息对实现高性能至关重要。