Page 97 - 《软件学报》2020年第11期

段旭 等:基于代码属性图及注意力双向 LSTM 的漏洞挖掘方法 3413

两个全连接层对其降维,最终得到特征张量的形状为(fixed_size,2nd_ebd_dim)的张量.可以发现,最终得到的特
征张量中的第一维与原始特征向量中的 x 轴相对应,其代表 CPG 中的节点.经过上述操作,将每个节点的特征都
压缩到了长度为 2nd_ebd_dim 的向量中进行表示,以便后续的双向 LSTM 模块进行操作.在本文中,1st_ebd_dim
和 2nd_ebd_dim 分别被设置为 6 和 64.

Fig.6 Feature tensor embedding process
图 6 特征张量嵌入过程
(2) 双向 LSTM 模块
虽然程序源代码的语义结构比自然语言更为复杂,但是其仍然存在一定的顺序性,其体现为代码中存在向
前和向后的依赖.因此,本文采用双向 LSTM 对代码的上下文信息进行捕获.双向 LSTM 的结构由向前和向后的
LSTM 组合而成,其中,向前的 LSTM 可以根据前序信息影响后序信息,而向后的 LSTM 可以通过后序信息影响
前序信息.该结构可以有效地处理存在双向上下文依赖的任务.
在本文中,使用节点的特征向量序列作为序列输入.对于节点特征向量在序列中的顺序,由于在编码时按照
节点所述代码在文件中的位置将其编码进特征张量中,因此节点特征向量序列中节点的顺序由代码在源文件
中的物理位置决定.为此,结合图 3 中的代码属性图进行阐述,假设用〈⋅〉代表节点,〈2〉在〈BUF_SIZE〉之前,因此通
过前向的 LSTM 可以使网络得知“2”作用在 BUF_SIZE 上.同理,〈memcpy〉在〈buf,str,len〉之前,因此通过后向的
LSTM 可以使网络得知 buf、str 和 len 是 memcpy 的参数,进而影响 memcpy 调用.此外,双向 LSTM 还可以在一
定程度上处理 AST 节点之间的嵌套关系,前向的 LSTM 可以使网络得知表达式由哪些元素组成,而后向的
LSTM 可以使网络得知元素组合成了什么表达式.在双向 LSTM 的末端,我们将前向 LSTM 和后向 LSTM 的输
出结果拼接在一起,并输出给后续的网络结构进行处理.
(3) 注意力模块
为了处理漏洞挖掘类间差异小的问题,本文借鉴细粒度图片分类任务中的注意力模型.细粒度图片分类是
指对图片进行更加细粒度的分类,例如对不同类别的鸟进行分类.该任务同样面临类间差异小的问题,其通常借
助注意力机制加以解决.通过在漏洞挖掘问题中引入注意力模型,可以使神经网络将更多的注意力放在关键的
代码语句及其依赖的变量上,例如敏感 API 调用及其依赖的变量等.通过在关键特征上赋予更多的注意力权重,
使其对分类结果产生更重要的影响,从而指导神经网络根据关键的细微差异进行分类.
在本文中,通过对双向 LSTM 输出的不同节点特征向量所组成的特征张量赋予不同的注意力权重,从而将