顺序问题:用于二进制代码相似性检测的语义感知神经网络
腾讯安全科恩实验室
发表时间:2019
阅读时间:2023/4/12
参考了大佬的文章:[论文阅读] (27) AAAI20 Order Matters: 基于图神经网络的二进制代码相似性检测(腾讯科恩实验室)
一、摘要
二进制代码相似性检测是计算机安全中的一项重要任务,其目标是在不访问源代码的情况下检测相似的二进制函数。
传统的方法通常使用图匹配算法,这种算法速度慢且不准确。最近,基于神经网络的方法已经取得了巨大的成就。首先将二元函数表示为具有手动选择的块特征的控制流图(CFG),然后采用图神经网络(GNN)来计算图的嵌入。虽然这些方法是有效和高效的,但它们无法捕获足够的二进制代码的语义信息。
在本文中,我们提出了语义感知神经网络来提取二进制代码的语义信息。特别地,我们使用BERT在一个令牌级任务、一个块级任务和两个图级任务上预训练二进制代码。此外,我们发现CFG节点的顺序对于图相似性检测很重要,因此我们在邻接矩阵上采用卷积神经网络(CNN)来提取顺序信息。
我们用四个数据集对两个任务进行实验。结果表明,我们的方法优于现有的模型。
二、引言
二进制码相似性检测用于检测两个给定的二进制函数是否相似。它有许多计算机安全应用,包括代码克隆检测、漏洞发现、恶意软件检测等。
科恩:“同一份源代码在不同编译器,不同平台,不同优化选项的条件下所得到的二进制代码是不相同的,我们的任务目标是把同一份源代码所编译出的不同的二进制代码找到。
传统的方法是采用图匹配算法(graph matching algorithm,Liu et al. 2006)来计算两个函数的相似性。然而,这些基于图匹配的算法都非常慢,并且很难适应不同的应用程序。随着近年来深度学习算法的发展,研究者尝试在控制流图(CFG)上使用图神经网络算法,并取得不错的效果。
文献[1]提出了一种名为**Gemini(Xu et al. 2017)**的算法,并且取得了不错的效果,它的输入是两个二进制函数的pair,输出是这两个二进制函数的相似度得分。
- 首先,将二进制控制流图(CFG)作为输入,接着Gemini会将其转换为带属性的CFG(attributed CFG),每个区块8个属性。
- 接着,采用structure2vec(Dai, Dai, and Song 2016)算法,将其进行图嵌入(graph embedding)
- 最后使用Siamese网络计算相似度得分并使用梯度下降算法降低损失训练模型(下图为腾讯科恩实验室给出的Siamese网络结构图)
[1] Xu X.; Liu C.; Feng Q.; et al. 2017. Neural network-based graph embedding for crossplatform binary code similarity detection. In Proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security (CCS), 363–376. ACM.
尽管基于神经网络的模型已经实现了很多成就,但还是有一些很重要的东西没有被考虑到
- 首先,每个块都被表示为手动选择特征的低维嵌入,这会导致大量语义信息(semantic information)的丢失
- 其次,节点的顺序在表达二进制函数上发挥了很重要的作用,但以前的方法没有设计提取他们的方法
为了解决这些问,我们提出了一个包含三个组件(component)的整体框架:
- 语义感知模块( semantic-aware modelling)
- 结构感知模块(structural-aware modeling)
- 顺序感知模块( order-aware modeling)
(1)在语义感知模块
我们使用NLP模型去提取二进制代码的语义信息,CFG块中的token被视为单词,CFG块被视为句子。在先前的工作中:
- (Massarelli et al. 2019)使用word2vec模型训练块中的token embeddings,然后使用注意机制获得block embedding。
- (Zuo et al. 2018)借鉴了神经机器翻译(NMT)的思想来学习跨平台二进制码之间的语义关系。
在本文中,我们采用BERT(Devlin et al. 2018)对tokens和blocks进行预训练。与BERT相同,我们对MLM(masked language model)任务的标记进行预训练,并提取所有相邻块对邻接节点预测任务(ANP)进行预训练。与分别学习token向量和block向量不同,本文方法能够同时学习token向量和block向量。此外,因为我们的最终目标是生成完整的图表示,所以我们添加了两个图级任务。
- 一种是确定两个采样块是否在同一个图中,我们称之为图内块任务(BIG,block inside graph task)。
- 另一种是区分块属于哪个平台/优化选项,称为图分类任务(GC,graph classification task)。
我们发现,额外的任务能帮助提取更多的语义信息,并且能更好的学习块的表示,在对块嵌入进行预训练之后,我们将在图级任务上对它们进行微调。
(2)在结构感知模块
我们使用**MPNN(Gilmer et al. 2017)和GRU(Cho et al. 2014)**更新函数(感觉应该翻译成带有GRU模块的MPNN),(Xu et al. 2018)已经证明了图神经网络可以具有像Weisfeiler-Lehman测试一样的区分能力。我们发现,在每个步骤中使用GRU比只使用tanh函数可以存储更多的信息。
(3)在顺序感知模块
我们设计了一种结构来提取CFG的节点的信息,图2显示了不同平台x86-64和ARM上函数的两个CFG和相应的邻接矩阵。这两个CFG具有相似的节点顺序。例如,两个图都是节点1和节点2,节点3相连,并且节点2和节点4、节点5相连。他们的邻接矩阵(adjacency matrices)是非常相似的。
在观察了很多跨平台的函数对之后,我们发现节点顺序的很多改变是非常小的。基于这个发现,我们提出了一个简单的方法去捕获顺序信息,在邻接矩阵上使用CNN。**我们发现仅有三层的CNN表现很好。**我们进一步探索其他的CNN模型,如Resnet(He et al. 2016),并且讨论CNN模型能够从邻接矩阵中学到什么。
我们的贡献如下:
- 我们提出了一个学习CFG图嵌入的通用框架,它可以学习语义信息、结构信息和顺序信息。
- 在语义感知建模中,我们采用BERT预训练标记嵌入和块嵌入,并使用掩蔽语言模型任务(MLM)和相邻节点预测任务(ANP)。此外,我们添加了两个图级任务来更好地学习块表示,这两个任务是图内块任务(BIG)和图分类任务(GC)。
- 在顺序感知建模中,我们发现节点顺序是有用的。我们在邻接矩阵上采用了CNN模型来提取CFG的节点顺序信息,取得了很大的成就。然后我们探索CNN可以从邻接矩阵中学到什么。
- 我们在四个数据集的两个任务上进行了实验,结果表明我们提出的模型比以前的方法取得了更好的性能。
三、相关工作
- Graph Neural Networks
图神经网络提出来学习节点表示和图表示。典型方法包括:
- GCN:使用卷积层来更新节点
- GraphSAGE:采用聚合函数将相邻节点进行合并
- GAT:利用注意力机制,从重要节点收到更多信息
- Bert
BERT是自然语言处理中最先进的预训练模型,通过Transformer实现。Bert使用两种策略来进行训练:
- MLM(masked language model task):他是一种自监督的预测掩蔽,用于促使模型捕获关于语言的有用的知识
- NSP(next sentence prediction,a classification task):使模型在训练中区分两个句子。
BERT只需在核心模型中添加一个小层,就可以用于各种各样的语言任务。
- Binary Code Similarity Detection
二进制代码相似度检测是计算机安全研究中的一项重要任务。传统的方法使用图匹配算法来计算图的相似度。然而,这些方法是缓慢和低效的。现有方法缺陷:
- 获取相似块是个有监督的过程,需要专家经验和领域知识,并且一些块不能唯一标注。
- 实际使用中,需要针对不同的平台训练不同的模型
四、本文方法
1. 总体架构(Overall Structure)
我们模型的输入是二进制代码函数的CFG,其中每个块具有中间表示的token(感觉翻译成令牌不太对)。我们模型的整体架构如图三所示。包含semantic-aware 模块、structural-aware模块、order-aware模块。
- 在语义感知模块中,模型把CFG作为输入,然后使用BERT去进行token embedding和block embedding
- 在结构感知模块中,我们使用带有GRU的MPNN更新函数,去计算图的语义和结构embedding,记为
- 在顺序感知模块中,模型把CFG的邻接矩阵作为输入,采用CNN去计算graph order embedding,记为
- 最后,我们链接(concatenate)他们,并使用MLP层去计算graph embedding
2.语义感知模块(Semantic-aware Modeling)
在语义感知模块,为了处理CFG,我们提出了一个带有四个任务的BERT预训练模型。这个模型有以下几个优势。
- 首先,我们能够从不同平台,不同架构,基于同一模型的不同编译优化选项生成的不同CFG中,提取块向量
- 其次,从预训练中,我们能够得到token级别、块级别和图级别的信息。因为我们有一个token级别、一个块级别、两个图级别的任务。
- 训练过程完全基于CFG图,不需要修改编译器或其他操作来获得类似的块对
**本文方法的灵感来自NLP中句子嵌入任务,因为CFG的块像句子,而块中的tokens像单词。**这个任务是提取一个句子的embedding。有两种主要的方法去完成这个任务。
- 有监督的方法。例如文本分类训练(Joulin et al. 2016)
- 无监督的方法,例如n-gram特征和 decoder-encoder skip thought(Kiros et al. 2015)
我们使用基于BERT的改进模型来提取CFG的块embedding,如图四所示,我们的预训练模型有四个任务
- MLM(Masked language model)
- ANP(Adjacency node prediction task)
- BIG(block inside graph task)
- GC(graph classification task)
对于节点内的token序列,我们使用MLM来提取块内的语义信息。MLM是一个令牌级任务,它屏蔽输入层上的令牌,并在输出层上预测它们,和语言模型的方式相同。
邻接节点预测任务(ANP)是一个块级的任务。在图中,块的信息不仅与块本身的信息有关,还与块本身的邻居信息有关,我们希望模型能够学习这些信息。在ANP任务中,在一个图上提取所有相邻的块,并在同一个图中随机抽取几个块,以预测两个块是否相邻。这两个任务(MLM & ANP)类似于原始BERT论文中的MLM & NSP任务(Devlin et al. 2018)。
为了更好地利用graph-level的信息,我们添加了两个辅助监督任务BIG和GC。
BIG任务和ANP相似,不同之处在于不同块对的采样的方法不同。BIG任务让模型判断两个节点是否在同一个图中。这个任务帮模型理解块和图之间的关系,并对graph embedding 任务有很大的帮助。
在我们的设想中,在不同的编译选项下,图和块的信息可能是不同的。为了使模型能够区分这些不同,我们设计了图分类任务(GC),GC使模型能够对不同平台、不同架构或不同优化选项中的块进行分类。
3.结构感知模型(Structural-aware Modeling)
在通过BERT 预训练获得block embedding 后,我们使用MPNN来计算每个CFG的graph semantic&structual embedding。
MPNN有三个步骤,message passing phase, update function,readout function,如公式2-4所示
对于本文,我们在消息函数上使用MLP,在更新函数使用GRU,在读出函数使用sum函数
4.顺序感知模型(Order-aware Modeling)
在这个模块,我们的目标是提取CFG的节点的顺序信息,我们应该思考能够从CNN模型中学到什么。
图5展示了三个图(没有块中的语义信息)和他们的邻接矩阵,他们通过添加一点小变化能够转化成另一个图。在这三个图中,每个图都有一个三角形,我们也能从他们的邻接矩阵中观察到这种特征。
首先,从5a转化成5b,一个新的点添加到三角形中,但是三角形的节点顺序并未发生改变,(去掉这种改变后)两个图的邻接矩阵的三角形也没有发生改变(not broken)
CNN能捕捉这些信息: CNN在训练集中看过很多这种样例后,可以学习到这种平移不变性。
原文:CNN could capture this information, because CNN has translation invariance when it has seen a lot of training data.
再看5c,他添加了两个节点,破环了三角形的节点顺序,然而,但我们把邻接矩阵的第二行第二列去掉后,三角形的特征仍然存在。
这种情况类似于图像的缩放。CNN在看了足够多的训练数据后,能够学到这种尺度不变性。
原文:CNN could also learn this scale invariance after seeing enough training data.
因为CNN能够学到这种平移不变性和尺度不变性,所以我们才讨论CNN能够学习到节点顺序的微小改变。
除了在学习节点顺序信息上准确性的提升,CNN还有其他的一些特点
- 与传统方法相比,更快
- 他们够输入不同尺寸的图,不用进行padding和ckipping
我们使用11层的Resnet,包含3个residual blocks,所有的feature map大小均为3x3。之后用一个global max pooling层,得到graph order embedding。在此之前不用pooling层,因为输入的图的大小不同。
五、实验
1.数据集
我们在两个任务上评价我们的模型。
- 第一个任务是交叉平台的二进制代码检测,我们的目标是确保相同的源代码具有比其他代码更高的相似性分数,相同的已有项目是Gemini(Xu et al. 2017)
- 第二个任务是图分类,在这项任务中,我们对图嵌入的优化选项进行了分类。
2.比较的方法(compared Methods)
我们进行了不同的实验找出三部分的效果。不太感兴趣
- Graph kernel methods
- Gemini
- MPNN
- Word2Vec
- Skip thought
- BERT
- CNN-based models
- CNN (random)
- MPNN (without semantic)
- MPNN (without semantic) + CNN
- Our model:Our model is BERT (4 tasks) + MPNN + 11-layer Resnet, which contains both semantic-aware modeling, structural-aware modeling, and order-aware modeling.
3.结果(results)
总体表现
表2和表3展示了不同的模型在两个任务上的表现。
与Gemini相比,我们的模型在两个任务上都表现更好。因为GRU能够储存更多的信息,MPNN比Gemini在所有的数据集上都表现的更好,所以我们选择MPNN。
用于语义感知建模的模型变体
为了研究是否bert预训练是必要的和有效的,我们研究了几个变体。首先,基于NLP的预训练块特征(word2vec、skip think、BERT 2&4)比手动特征获得了更好的性能,这表明为CFG块建立复杂的模型是至关重要的。
BIG任务和GC任务也很有用,它们的结果增加了1%-2%。
在这两个任务中,块嵌入可以学习图级信息,这有助于图级任务。我们在图6中显示了块嵌入。在四个方向上设置四个CFG及其块嵌入。我们采用K-means将这些块嵌入聚类为四类,并且不同的聚类具有不同的颜色(红色、蓝色、绿色和紫色)。我们可以观察到,同一张图中的块倾向于具有相同的颜色,而不同的图具有不同的主颜色。
用于顺序感知建模的模型变体
只有使用基于CNN的模型才能在这两项任务上都取得良好的结果。11层Resnet比3层CNN和7层Resnet好一点。相比于,基于CNN的模型表现更好。
当随机打乱节点是,CNN什么也学不到,这就表明,CNN模型可以学到节点的顺序,用CNN提取邻接矩阵的特征是有意义的。
六、结论
在这篇论文中,我们提出了一个新的二进制代码图学习的框架,其包含语义感知模型,结构感知模型,顺序感知模型。语义和节点的顺序信息都是非常重要的对于表示CFG。
为了捕获CFG图的语义信息,我们提出了有两个原始任务(MLM,ANP)和两个附加图级别任务(BIG、GC)的BERT 预训练模型。
然后我们使用MPNN提取结构信息。我们进一步提出一个基于CNN的模型去捕获节点顺序信息。
我们在四个数据集上对两个任务进行了实验,实验结果表明,我们的方法比当前最先进的方法表现更加出色。
七、感受
主要不是对论文的感受,是对查这篇论文时看到的大佬博客的感受,真的是好多好厉害啊。
对比自己,论文写了又停,时常懈怠,着实不太行。
从此开始,要更加认真和勤奋的看论文!!!