SCALE:Constructing Structured Natural Language Comment Trees for Software Vulnerability Detect..论文笔记

SCALE: Constructing Structured Natural Language Comment Trees for Software Vulnerability Detection

1. Abstract

1. 背景, 任务

函数粒度的漏洞检测

近年基于PLM(预训练模型)的方法比其他DL(深度学习)方法效果更好, 但是PLM方法一般将代码作为序列输入, 没有考虑到代码的结构信息.

2. 问题/动机

因此, PLM方法难以检测出有复杂逻辑(如很多操作符和)代码的语义信息, 也难以捕捉代码的执行顺序

3. 方法

工作提出一个新的代码表示: SCT(structured natural language comment tree), 将语义信息和执行顺序信息结合到AST上

SCT由三个阶段组成

1. comment tree construction(novelty): comment tree: 将代码输入到LLM中, 让LLM生成注释, 在AST中把注释增加到对应的结点上, 我的理解就是用自然语言描述代码的语义, 毕竟代码既有自然语言的属性即语义, 也有逻辑结构(语法词法), 语义是很难捕捉的, 这里就提供了一种解决方案
2. Structured Natural Language Comment Tree Construction: 就是构造他所说的SCT, 简单来说就是把上面的comment tree改一下结构
3. SCT-Enhanced Representation: 也就是嵌入, 这里用了交叉注意力机制, 后面细讲

4. 结果

效果好, 和几种baseline做了对比, 与DL, PLM, LLM都对比了, F1分别提升了2.96%, 13.7%, 3.75%在三个不同的数据集上, 另外SCALE可以用应在不同的PLM上

2. Introduction

这篇的Introduction写的很好, 分类了过去的方法也说了问题, 算一个不错的小综述, 所以也说一下

1. 传统PA(程序分析)方法

比如INFER[Facebook. [n.d.]. Infer. https://fbinfer.com/.], CheckMarx[Israel. [n.d.]. Checkmarx. https://www.checkmarx.com/.]

问题:

1. 依赖人工, 耗时耗力
2. 依赖规则定义和专家知识, 必然有不全的情况

2. 基于DL的方法

1. 代码看作序列

如CNN(Yueming Wu, Deqing Zou, et.al. 2022.VulCNN: An Image-inspired Scalable Vulnerability Detection System. ICSE 2022)

RNN(Zhen Li, Deqing Zou, et.al. 2022.SySeVR: A Framework for Using Deep Learning to Detect Software Vulnerabilities. TDSC 2022)

2. 代码看作图

GNN(Yaqin Zhou, Shangqing Liu et.al. 2019. Devign: Effective Vulnerability Identification by Learning Comprehensive Program Semantics via Graph Neural Networks. NeurIPS 2019)

问题

由于缺乏漏洞数据集, 模型性能很受影响

However, the performance of these models can be largely limited due to the scarcity of vulnerability data

Yun Peng, Chaozheng Wang, Wenxuan Wang, Cuiyun Gao, and Michael R. Lyu. 2023. Generative Type Inference for Python. CoRR abs/2307.09163 (2023).

3. 基于PLM的方法

如EPVD(2023 TSE, Vulnerability Detection by Learning From Syntax-Based Execution Paths of Code), SVulD(2023 CoRR, Distinguishing Look-Alike Innocent and Vulnerable Code by Subtle Semantic Representation Learning and Explanation.)是现在的sota

采取的是微调的范式

问题

1. 由于推理空白(reasoning gap)(1. Why think step-by-step? Reasoning emerges from the locality of experience. 2023 CoRR///2. Chain-of-Thought Prompting Elicits Reasoning in Large Language Models. 2022 NeurIPS.)PLM难以学到代码中的逻辑,
2. 代码的执行顺序: 由于PLM把代码作为序列输入, PLM难以捕获非顺序执行代码的执行顺序

3. Method

核心思想

上面说了用自然语言来增加代码的语义信息是这篇文章的主要创新点, 工作面对的问题是要加强对代码语义的表示, 思想就是PLM在某些东西(复杂逻辑)上的处理不好, 就把处理不好的用别的方式来处理, 这里就是拿自然语言来替换掉PLM处理不好的部分来进行进一步的处理, 再抽象一点他解决的方法就是处理不好的绕路走, 用能被处理好的东西替换,再想办法和原来的处理方法很好的融合在一起

方法

和摘要中说的一样, 三个阶段

1. comment tree construction

主要的动机和思想

就是想用LLM生成的自然语言注释来丰富代码的语义信息, 代码的常规表示之一是AST, 这里就是要把chatgpt生成的注释插入到AST中, 这一步还没什么精妙的

方法

接下来的问题就只剩下如何去把自然语言的描述来表示代码

1. 用chatGPT对每个代码片段(code snippet)生成注释

这里的注释的形式就是我们写代码的那个注释, 不是什么新概念, 就比如c语言中那就是 // 注释, 直接插到代码中

---

另外

为了防止chatgpt生成没有规范化的代码, 工作对chatgpt生成的结果进行处理

1. 删掉所有空行
2. 把多行注释改成单行注释

2. 用Tree-Sitter生成代码的AST

这里的代码按作者的意思指的应该是注释之前的代码

3. 遍历生成的AST, 把comment融合到AST去

首先在注释了的代码中看, 把代码中注释comment的下一条语句作为遍历的起点, 把comment变为起点的兄弟结点

这里文章中细节说的很不清楚, 不过不太重要, 后面一步更重要, 就看作是生成了注释然后插入到了AST里就行了, 具体长什么样管不着了, 反正后面会生成一个规整的东西. 自己实现也很好做, 弄清楚没什么价值.

2. Structured Natural Language Comment Tree Construction

主要思想

上面一步是为了提取代码中的语义信息加入到常规的代码表示中来丰富代码的语义信息, 这一步就是为了提供代码的执行顺序信息给后续步骤, 也算是一种语义信息吧, 代码中的执行顺序也就是控制流, 具体表现就是控制语句, 乍一看没什么特别的. 但其实, 工作做了一个很大胆的事, 他把控制语句的根结点用自然语言直接替换掉了(复合语句的子结点还是保留了, 没有删掉), 具体看下面

方法

简单来说就是, 用某些规则, 把控制语句中的一部分用LLM生成的自然语言注释替换掉

1. 替换规则

图是C/C++的控制语句, 加粗的地方是被替换的部分

SCLAE只对这几种语句进行自然语言的替换

文章里说了一大堆, 总结一下就是之前对很多语句都生成了注释, 如果这个有注释的语句在上面几种里面, 就把加粗的部分用注释替换掉

2. 举例

1. 前置解释

Fig(a)是第一步处理后的Comment Tree

Fig(b)是第二步处理后的Comment Tree

Fig4中, 灰色的部分是方法第一步生成的有注释的代码(某些语句的注释不知道为什么图中没写出来),

红色的是第二步处理后的AST表现为代码时的样子

蓝色是我标注的控制语句被替换的部分

如果只看黑色的和灰色的部分, 就是Fig(a)的AST作为代码时的表示

如果只看黑色的和红色的部分, 就是Fig(b)的AST作为代码时的表示

2. 详细解释

可以看到, 第二步用LLM对控制语句生成的注释, 直接替换掉了控制语句的某些部分, 比如在Fig4中的第三行是if语句, 第二步的算法就用对这个if条件生成的LLM的注释Check if allocation was successful, 替换掉的就是if语句成立的条件!p->entries, 又比如条件复杂一点的, 在第9行的for语句, 也是用注释直接替换掉了三个表达式组成的for的范围

SCT-Enhanced Representation

方法概览

注意力机制就不多介绍了, 这里的流程就是把代码和SCT分别输入PLM, 得到向量表示, 再用交叉注意力(cross-attention)综合这两个向量, 然后放入一个判别是否有漏洞的分类器训练

1. 把代码和SCT输入到PLM中得到向量表示

把代码片段, 以及对应的SCT分别输入到UniXcode(一个可以接收代码, AST, 注释输入的多模态代码PLM, 也是效果很好的sota), 获得他俩的表示$h_c\in {\mathcal{R}}^{l\times n}$(代码), $h_{c}t\in {\mathcal{R}}^{l\times n}$(SCT), 其中$l$是输入的长度, $n$是嵌入的维度

2. 用交叉注意力(cross-attention)综合嵌入的两个向量

注意力不想多介绍了, 后面开一篇详细介绍

(1)是attention的QKV, (2)是attention分数的计算公式, (3)是分类器及其输入的形式化描述

3. 放入分类器训练

分类器为(2016. Text-Attentional Convolutional Neural Network for Scene Text Detection. IEEE Trans. Image Process.)损失函数是交叉熵, 是函数粒度判别有没有漏洞的分类器, 文中直接引用的没有详细介绍

看了下摘要, 这个分类器是用在图像任务里的, 具体是识别自然场景图像中的文本, 比如识别一个街景里标牌上的字母A, 和一般CNN的区别主要是1. 用了多级的监督信息(1. 图中有没有文本 2. 图中的文本在哪个区域[用掩码实现, 掩码就是告诉你关注数据的那些区域] 3.图像里的文本是啥 ), 2. 对比度增强的检测器MSER

莫名其妙的用了一个图像任务里的基于CNN的多分类器, 看起来八杆子打不着, 应该就是因为分类器的输入是矩阵, 然后可以剔除无关信息所以用这个吧, 后面看实验部分补上