[发明专利]一种基于对比学习的源代码漏洞检测方法及系统

权利要求书

1.一种基于对比学习的源代码漏洞检测方法，其特征在于，包含以下步骤：

1)采集漏洞数据库NVD和开源代码库GitHub中的漏洞数据构建漏洞数据集，对漏洞数据集进行数据增强形成正负样本对；

2)对正负样本对进行代码特征表示学习，将正负样本对转化为代码属性图，该代码属性图包含的控制流图CFG和程序依赖图PDG共同组成联合图；

3)将联合图进行图嵌入并应用对比学习，得到正负样本对的特征联合图向量，同时添加分类器，将正负样本对的特征联合图向量作为分类器的输入得到预测标签，以此构建漏洞检测模型；

4)从漏洞数据库NVD和开源代码库GitHub中抽取漏洞代码作为待测代码，利用漏洞检测模型对待测代码进行分析是否含有漏洞。

2.根据权利要求1所述的基于对比学习的源代码漏洞检测方法，其特征在于，步骤1)具体包含如下步骤：

1.1)采集漏洞数据库NVD和开源代码库GitHub中的漏洞数据，提取漏洞文件和对应漏洞补丁文件；

1.2)对提取出的漏洞文件和对应漏洞补丁文件进行预处理，从函数级别提取漏洞文件中的漏洞代码和对应漏洞补丁文件中的补丁代码，去除多余信息，多余信息包括头文件、注释以及声明的全局参数，得到漏洞数据集；

1.3)从漏洞数据集中选择漏洞代码作为硬负样本，对应的补丁代码作为原始样本，对原始样本的应用语义保留启发式使原始样本的语法结构发生改变而语义不变，从而生成正样本，把对应的正样本、原始样本、硬负样本组合作为正负样本对。

3.根据权利要求2所述的基于对比学习的源代码漏洞检测方法，其特征在于，步骤2)具体包含以下步骤：

2.1)将每个正负样本对中的所有样本进行代码表征，通过工具Joern为正负样本对生成代码属性图CPG，其中代码属性图CPG包含抽象语法树AST、控制流图CFG和程序依赖图PDG；

2.2)将抽象语法树AST的叶子节点，使用字节对编码BPE分词，利用Word2vec对得到的代码令牌tokens进行词嵌入，使用基于树的卷积神经网络TBCNN根据抽象语法树AST中的层次结构生成语句嵌入向量；

2.3)经过树聚合层，执行最大池化来将树的所有节点的向量聚合为包含相应语句语义的单个向量，用该单个向量作为控制流图CFG和程序依赖图PDG中对应语句节点的初始特征向量；

2.4)将控制流图CFG和程序依赖图PDG共同组合成联合图，将控制流图CFG和程序依赖图PDG及联合图中语句的向量共同输入图级代码表示层。

4.根据权利要求3所述的基于对比学习的源代码漏洞检测方法，其特征在于，步骤3)具体包含以下步骤：

3.1)根据控制流图CFG和程序依赖图PDG中的控制流、控制依赖和数据依赖三种边计算联合图中各个节点相应的注意权重系数；

3.2)图级代码表示层由多个图注意力层组成，一个图注意力层对应一个边的种类，根据注意权重系数，图注意层对联合图中的节点向量表征进行更新；

3.3)将多个图注意力层进行平均形成多头注意力层，得到图中节点的最终向量表征；

3.4)根据联合图内的节点重要性的不同，得到正负样本对的特征联合图向量；

3.5)将原始样本和正样本作为相似对，原始样本和硬负样本作为不相似对；

3.6)使用余弦相似度使损失最小化；

3.7)添加分类器，将正负样本对的特征联合图向量作为分类器的输入，用激活函数softmax训练分类器得到预测标签，若预测标签为1则为漏洞代码，标签为0则为非漏洞代码。

5.根据权利要求1所述的基于对比学习的源代码漏洞检测方法，其特征在于，步骤4)具体包含以下步骤：

4.1)从漏洞数据库NVD和开源代码库GitHub中抽取漏洞代码作为待测代码，将待测代码文件执行步骤2)生成相应代码属性图；

4.2)将经过代码特征表示学习得到的代码属性图作为漏洞检测模型的输入，得到待测代码的预测标签；

4.3)输出预测标签，若预测标签为1则为漏洞代码，预测标签为0则为非漏洞代码。