一种基于代码数据流增强大模型的漏洞检测方法

专利2026-07-23  5


本发明属于网络信息安全,特别涉及一种基于代码数据流增强大模型的漏洞检测方法。


背景技术:

1、近年来,软件漏洞已经成为软件安全领域的重要问题。由于安全漏洞,许多互联网用户的个人信息被泄露。2016-2021年的通用漏洞枚举(cve)报告指出,漏洞数量和种类每年都在以极快的速度增长。这些网络安全攻击已经为经济社会带来了不可估量的危害和损失。目前漏洞检测的主要技术方法有传统的检测技术、基于深度学习的方法和基于预训练模型的方法。

2、传统漏洞检测方法依赖于领域专家的先验知识。专家需要人工审查大量的源代码,并为漏洞检测设定特定的规则或人工特征以用于检测。应用污点跟踪、符号执行、模糊测试和基于代码相似性匹配等半自动技术提高检测效率但缺乏完全自动化,且规则的指定成本较高。此外这些半自动化技术可能会导致误报率和漏报率增加。

3、基于深度学习的漏洞检测方法主要分为基于图结构的模型和基于序列的模型。基于图结构的模型旨在通过静态分析的手段提取源代码之间的各种依赖关系,如抽象语法树,控制流图,数据流图和程序依赖图。利用图学习的模式,通过聚合源代码的结构和语义信息来实现检测。但是基于图的模型在进行嵌入时,图的节点特征往往是切片级代码段,与漏洞没有直接关系的文本或令牌信息(如变量名,函数名)会被一同嵌入,从而带来额外的噪声干扰。

4、基于序列的模型将源代码视为自然语言序列,并应用自然语言处理领域的通用方法进行漏洞检测。

5、虽然预训练模型可以在特定的下游任务上进行微调,但是只在文本层面上学习很难学习到代码的固有结构,如数据流和控制流。如果将微调后的模型直接应用于漏洞检测,则面临着从冗长的代码和复杂的结构中捕获漏洞特征的挑战。


技术实现思路

1、针对现有技术存在的上述问题,本发明要解决的技术问题是:应用代码预训练模型进行漏洞检测时,很难从源代码序列中学习到足够的漏洞语义信息。本发明提供一种基于代码数据流增强大模型的漏洞检测方法,以保证代码预训练模型能够学习到更充分的漏洞语义,从而提升预训练模型在漏洞检测这一下游任务上的性能。

2、为解决上述技术问题,本发明采用如下技术方案:一种基于代码数据流增强大模型的漏洞检测方法,包括以下步骤:

3、步骤1:数据流图构建,根据输入的函数级源代码,利用静态分析工具tree-sitter将其解析为数据流图g(c)=(v,e),其中c代表源代码序列,v代表节点集合,e代表有向边集合。

4、将函数级源代码输入预训练模型得到源代码序列嵌入向量。

5、步骤2:数据流嵌入,将步骤1中得到的数据流图输入图神经网络进行图嵌入获得嵌入后的图向量表示,获得的图向量表示视作序列,进行正弦位置编码,获得最终的数据流嵌入。

6、步骤3:合并数据流信息,将数据流嵌入与源代码序列嵌入向量拼接,然后使用多层感知机进行漏洞检测。

7、进一步地,步骤2中将数据流图输入图神经网络进行图嵌入之前先进行如下操作:节点特征嵌入,对于步骤1中生成的图g(c),假设节点上的变量对应的数据类型为data_type,那么节点的节点特征嵌入如下:

8、e(v)=embedding_layer(tokenize(data_type))

9、其中,e(v)表示节点v的嵌入特征,embedding_layer表示嵌入层,tokenize(data_type)表示对数据类型data_type进行令牌化处理。

10、对g(c)的结构和节点特征进行向量化表示,其中结构由邻接矩阵a表示,节点特征矩阵x则由所有节点特征嵌入e(v)拼接而成,g(c)的向量化表示为g(x,a)。

11、将g(x,a)进行标准化处理,选择邻接矩阵a中的标准节点个数m,对于节点数超过m的数据流图,截取其邻接矩阵a中前m个节点表示;对于节点数不足m的数据流图,则向数据流图中添加孤立的节点以补足m个节点数,得到标准化表示g′(x′,a′)。

12、进一步地,步骤2中使用图卷积神经网络gcn进行图嵌入获得嵌入后的图向量表示eg,步骤如下:

13、h(0)=x′

14、

15、其中,av,u表示在拉普拉斯重新归一化的邻接矩阵中节点v和u之间的边权重,其中d是a′中节点度的对角矩阵,w(k)表示与边相关联的权重矩阵,relu表示激活函数,表示第k次迭代时所有邻居节点u的隐藏状态。

16、使用软注意力机制生成节点的最终特征向量ev,利用联合池化方法嵌入后的图向量表示eg,他们的具体生成方法如下:

17、

18、

19、其中,是在第k次迭代时,结点v自身的隐藏状态,b表示模型通过学习得到的偏置参数,w,wt分别表示模型通过学习得到的权重矩阵及其转置矩阵。

20、进一步地,步骤2中使用门控图神经网络ggnn进行图嵌入获得嵌入后的图向量表示eg,步骤如下:

21、h(0)=x′

22、

23、

24、

25、

26、

27、其中z和r是更新和重置门,⊙是逐元素乘法,relu和sigmoid表示激活函数。

28、是节点v的所有邻居节点在第k步隐藏状态的加权和,表示在k+1次迭代结点v的更新门状态,表示在k+1次迭代结点v的重置门状态,是结点v在第k+1次迭代的候选隐藏状态,是结点v在第k+1次迭代的最终隐藏状态。

29、uz,ur,uo是内部权重矩阵,它们分别用于计算更新门、重置门和候选隐藏状态时与节点自身之前的隐藏状态的乘法操作。

30、wz,wr,wo是外部权重矩阵,用于计算更新门、重置门和候选隐藏状态时与节点的聚合信息的乘法操作。

31、进一步地,步骤2中将得到所有图向量表示eg作一个序列,进行正弦位置编码,获得最终的数据流嵌入的步骤如下:

32、正弦位置编码的应用方式如下:

33、

34、

35、其中i表示编码向量中的维度索引,d是正弦编码的维度,posi对应于第i个标记的正弦编码,pos表示序列中元素的位置。

36、相对于现有技术,本发明至少具有如下优点:

37、实验结果证明,本发明提出的方法,能够为代码预训练模型提供更多的漏洞语法和语义信息,提供预训练模型在漏洞检测这一下游任务上的准确性。而相较于其它漏洞检测的方法,本发明提供的方法可以使代码预训练模型具有更高的自动化程度,减少对领域专家知识的依赖,大大节约代码审计的成本,提高代码审计的效率。



技术特征:

1.一种基于代码数据流增强大模型的漏洞检测方法,其特征在于,包括以下步骤:

2.如权利要求1所述的一种基于代码数据流增强大模型的漏洞检测方法,其特征在于,步骤2中将数据流图输入图神经网络进行图嵌入之前先进行如下操作:

3.如权利要求2所述的一种基于代码数据流增强大模型的漏洞检测方法,其特征在于,步骤2中使用图卷积神经网络gcn进行图嵌入获得嵌入后的图向量表示eg,步骤如下:

4.如权利要求2所述的一种基于代码数据流增强大模型的漏洞检测方法,其特征在于,步骤2中使用门控图神经网络ggnn进行图嵌入获得嵌入后的图向量表示eg,步骤如下:

5.如权利要求1-4任一项所述的一种基于代码数据流增强大模型的漏洞检测方法,其特征在于,步骤2中将得到所有图向量表示eg作一个序列,进行正弦位置编码,获得最终的数据流嵌入的步骤如下:


技术总结
本发明涉及一种基于代码数据流增强大模型的漏洞检测方法,适用于任意代码预训练模型的漏洞检测下游任务。该方法主要由数据流图构建、数据流嵌入和正弦位置嵌入、合并数据流信息三个步骤组成。步骤一从函数粒度的源代码中解析数据流图并使用变量数据类型的词嵌入作为节点特征。步骤二通过图学习技术进行数据流图嵌入,通过正弦位置编码将数据流中的相对位置信息合并到图嵌入中,以保证漏洞数据流信息的完整性。步骤三将步骤二中产生的数据流嵌入与预训练模型产生的源代码句向量拼接,使用多层感知机进行漏洞检测。本发明能够通过与代码预训练模型相结合来增强漏洞检测方法在处理复杂数据流信息方面的能力,从而提高检测的性能和效率。

技术研发人员:鄢萌,蒋中豪,张小洪,吴佳欣,古晓燕,文涛,孙伟峰,刘超,胡海波
受保护的技术使用者:重庆大学
技术研发日:
技术公布日:2024/6/26
转载请注明原文地址:https://doc.8miu.com/read-1831962.html

最新回复(0)