基于神经网络的跨平台代码相似性检测

本文介绍2017年CCS论文《Neural Network-based Graph Embedding for Cross-Platform Binary Code Similarity Detection》，其DOI信息在这里。

分析程序二进制代码，主要方法是将其转化为控制流图，然后再针对这些图进行处理。本文创造性地提出了基于卷积神经网络计算图嵌入的方法，与现有途径相比，速度有巨大提升。文章将机器学习算法移植到其它领域，并针对具体情境作出改进，思路值得学习。

简介

给定两个二进制程序函数，想要检测它们是否相似，这一问题被称为“二进制代码相似性检测”，它有许多安全应用，如剽窃检测、恶意检测、漏洞搜索等。在这些安全应用中，IoT设备的固件镜像中的漏洞搜索尤为关键，而且比以往任何时候都更加重要。源代码中的单个漏洞可能会影响数百个或更多具有不同硬件架构和软件平台的设备。

在跨多个平台（例如x86、ARM或MIPS）的二进制文件中，安全从业者面临着越来越多的快速检测类似功能的需求。直到最近，研究人员才开始着手解决跨平台的二进制代码相似性检测问题。这些工作建议直接从二进制代码中提取各种健壮的、独立于平台的特性，以便使控制流图中的每个节点代表一个函数。然后，利用图匹配算法进行相似性检测，检验两个函数的控制流图表示是否相似。这种方法不可避免地速度较慢，且有时不够准确，很难适应新的任务。

Qian等人于2016年提出Genius，它从控制流图中学习高级的特征表示，将图形用嵌入（embedding，这里是高维的数字矢量）表示。然而，为了计算二进制函数的嵌入，它还依赖于图形匹配算法来计算目标函数和二进制函数的代码手册（codebook，下面会进一步阐明）之间的相似性。

以上基于图匹配的方法有两个问题。

1. 采用图匹配算法产生的固定匹配函数很难适应不同的应用场景。例如，给定两个仅有细微差别的二进制代码段，在剽窃检测中，它们可能被认为因为大部分程序代码相同。但在漏洞搜索中它们可能被认为是不同的，因为修改少数代码也有可能修复了一个漏洞。人工设计的相似函数不能在两个场景中都适合。
2. 基于图匹配的所有相似性检测方法的效率都受到图形匹配算法（如二分图匹配）的效率的限制。然而这类算法的效率较低。

近年来，深度学习已被应用于许多应用领域，其结果比其它方法更有说服力。深度神经网络可以胜任二进制分析任务。例如，生成二进制函数的嵌入的一个神经网络，其参数可以进行端到端训练，仅依赖于较少的先验知识。此外，基于神经网络的方法可以自适应，因为可以通过不同的数据进行训练，以适应不同的应用场景或任务。此外，深度神经网络模型运算高效，运行时间与输入大小和网络大小线性相关。

受以上优势启发，本文中，作者提出了一种基于神经网络的方法，生成用于相似检测的二进制函数的嵌入。假设一个二进制函数被表示为一个含有每个节点属性的控制流图，使用一个图嵌入网络将图形转换为嵌入。一些场景下已经提出了用于分类和回归任务的图嵌入网络，但本工作是相似性检测，与分类不同，因此已有的方法并不直接适用于本文的任务。因此，文章提出了一种新的方法，通过将图嵌入网络组合到一个Siamese网络中，来计算图像的相似性。两个相似函数的图形嵌入应该彼此接近，反之亦然。整个网络模型可以被端到端训练，以进行相似性检测。

此外，作者还设计了一种新的训练集和数据集创建方法，使用默认策略对一个任务独立的图嵌入网络进行预训练（pre-train）。本文的方法使用从相同源代码编译、但针对不同平台和编译器优化级别的二进制函数来构建大规模的训练数据集。评价表明，这种任务独立模型比基于状态的图匹配方法更有效、更一般化。并且，预训练模型能够快速根据需求场景再训练（retrain）得到可用模型。文章实现了一个原型系统Gemini，比上文提及的Genius，生成嵌入的速度快2400至16000倍，且用时小于30分钟（Genius用时多于一星期）。

二进制代码相似性检测

本节中首先使用跨平台的二进制代码搜索作为例子来说明设计一个相似检测函数的问题。然后解释现有的方法，同时提供一个演示有效的嵌入函数如何帮助设计相似函数。最后，提出利用神经网络作为嵌入函数的方法，以及这种方法的优点。

问题定义

下面描述跨平台二进制代码相似性检测问题。给定一个感兴趣的二进制函数（例如，一个包含Heartbleed漏洞的函数），我们想要检查一个大型语料库（corpus）的二进制函数（例如，从各种物联网设备的固件镜像中提取形成的语料库）快速、准确地确定候选名单，这些候选的语义等同或类似于感兴趣的功能。我们称感兴趣的二进制函数为查询函数，而二进制函数的语料库则为目标语料库。

这个问题的技术可以应用到许多安全应用程序中，例如在固件镜像中的bug搜索和二进制代码中的剽窃检测等。问题的核心是设计一个函数来检测两个函数是否相似。解决这个问题，需要达到以下设计目标:

• 仅二进制。在实际中，通常不能获取二进制函数的源代码。因此，有效的相似性检测和代码搜索技术必须直接处理二进制代码。
• 跨平台支持。由于查询函数和目标语料库中的函数可能来自不同的硬件架构和软件平台，因此有效的二进制搜索技术必须能够容忍不同平台引入的句法变化，捕捉这些二进制函数的内在特性。
• 高精确度。一个有效的二进制代码相似检测器应该能够给一对相似的函数打出高分，给一对不相关的函数打出低分。
• 高效率。相似函数应针对漏洞搜索系统和其他应用程序进行高效的计算，以适应大型目标语料库。
• 自适应。当领域专家（domain expert）可以提供相似与不相似的样本时，相似函数应该能够快速地适应这些领域特定应用程序的样本。

现有技术

虽然在二进制代码的匹配和搜索方面已经有了一系列的努力，但大多数都只针对单一平台的二进制代码。直到最近，研究人员才开始在跨平台的环境中解决这个问题。这些工作建议直接从二进制代码中提取各种特性，这些特性足够强大，可以在不同的体系结构和编译器优化选项中持久存在。

成对图匹配（Pairwise Graph Matching）由Pewny等人提出。该方法在控制流图中提取每个基本块的输入/输出对作为其特性（或标签），然后进行图形匹配。但这是一个非常昂贵的过程：输入/输出对和图形匹配的计算都很昂贵。为了提高效率，discovRE提出了一种更轻量级的语法级别特征（例如，算术指令的数量和调用指令的数量），以加速特征提取，并在图形匹配之前通过简单的函数级特征进行预过滤，以提高搜索效率。然而，Feng等人认为，这种预过滤方法不可靠，可能会导致搜索精度的显著降低。从根本上说，两种方法都依赖于成对图匹配来检测相似度，从而不可避免地导致低效。

为了同时实现可扩展性和高精度，需要将图形的表示编码到嵌入中。这样做时，相似函数可以设计成计算两个嵌入之间的距离函数，这是有效且易于计算的。另外，特征嵌入可以使用局部敏感哈希（kocality sensitive hashing，LSH）数据库进行索引，以便在$O(1)$时间内执行搜索查询。

Feng等人第一次将这种方法应用于漏洞搜索问题。他们提出了Genius，一个图形嵌入工作流。给定一个二进制函数（来自固件镜像或已知漏洞），Genius首先从一个属性控制流图（attributed control flow graph，ACFG）的形式提取原始特性。在ACFG中，每个顶点都是带有一组属性的基本块。然后将每个ACFG转换为高维嵌入，并使用局部敏感哈希存储到哈希表中。因此，为了识别类似于查询函数的一组二进制函数，我们只需要找到对应的查询函数嵌入，并在目标语料库中找到附近的嵌入。

关键部分是如何将ACFG转换为它们的嵌入。Genius采用基于代码库（codebook）的方法来嵌入化一个ACFG。它使用聚类算法来训练一个代码库，其中包括为每个集群标识的几个代表性的ACFG。然后，Genius度量给定的ACFG和每个集群代表性的ACFG在代码库中的相似度。这些相似性度量构成了指定的ACFG的特征嵌入。

虽然图嵌入的思想令人鼓舞和信服，但是使用代码库和图匹配有几个局限性。首先，代码库生成是一个非常昂贵的过程，在训练数据集的每一对控制流图中都要进行成对图匹配，然后进行聚类。因此，代码库的质量受到训练数据集规模的限制。其次，图嵌入的运行开销随代码的大小线性增加，因此，代码库必须很小，这样限制了图形编码的保真度。最后，该方法的搜索精度最终受到二分图匹配质量的限制，而作为一种近似算法，二分图匹配不一定能产生最优匹配结果。

基于神经网络的嵌入生成

本节首先介绍代码相似性嵌入问题；之后对文章提出的解决方法作出概述，并重点解释其中两个重要模块；最后讨论如何预训练得到任务无关模型、再训练得到任务相关模型。

代码相似性嵌入问题

如上所述，代码相似性度量可以是依赖于任务的。我们假设存在一个oracle，用于确定给定任务的代码相似性度量，这是未知的、希望通过学习得到。给定两个二进制程序函数$f_1$、$f_2$，$\pi (f_1, f_2)=1$说明它们是相似的，$\pi (f_1, f_2)=-1$说明它们是不相似的。这里的oracle$\pi$对每个任务都是特定且未知的。在某些场景中，一些元组$\langle f_1, f_2, \pi (f_1, f_2) \rangle$可以得到。例如，领域专家可以提供一些参考标准(ground truth)。

代码相似性嵌入问题的目标是找到一个映射$\phi$，它能够将一个函数的ACFG映射为一个嵌入$\mu$。直觉上，这样的嵌入应该捕获足够的信息来检测相似函数。即，给定一个易于计算的函数$Sim(\cdot, \cdot)$与两个二进制函数$f_1$、$f_2$，若$\pi (f_1, f_2)=-1$，则$Sim(\phi(f_1), \phi(f_2))$较大，反之较小。映射的一个优势是它支持高效的计算。两个函数之间的相似度可以用两个嵌入之间的相似函数来计算，而不需要开销巨大的图形匹配算法。

如上所述，使用一个神经网络来近似于嵌入函数是特别有吸引力的，因为当提供有限的、任务相关的参考标准时，它可以快速地重新训练以适应给定的任务。此外，计算基于神经网络的嵌入并不依赖于任何昂贵的图形匹配算法，因此可以有效地实现。

方案概览

本节提出解决代码相似嵌入问题的关键思想。下面假设一个函数$f$的二进制代码是由其ACFG$g$代表。

文章将嵌入映射$\phi$作为神经网络。由于输入是一个ACFG，文章利用之前的图嵌入网络，从机器学习社区来解决这个问题。然而，在Dai等人的工作中，图形嵌入网络是为分类问题设计的，需要标签信息来训练模型。相比之下，代码相似嵌入问题并不是分类问题。因此，现有的方法并不能很好地应用，需要设计一种新的方法来训练图形嵌入网络来解决相似度检测问题。

为解决这个问题，文章提出一种新的学习方式。思路是不去训练图嵌入网络$\phi$在预测任务上的性能，而训练$\phi$区分两个ACFG的相似度。文章设计了一个Siamese结构，将图嵌入网络Structure2ven嵌入其中。Siamese结构以两个函数作为输入，并生成相似度作为输出。这使得模型能够端到端训练，只使用一对（pair）图$g_1$，$g_2$作为输入、参考标准$\pi (f_1, f_2)$作为输出，不需要额外使用手工方法启发网络如何生成嵌入。因此，这样的结构更加健壮，且更能适应不同的任务。训练一个Siamese网络需要大量的相似、不相似函数。但在大多数任务中，参考标准是有限的。为解决这个问题，使用一个默认策略，认为相等的（equivalent）函数（即由同一源代码编译成的二进制函数）是相似的，不相等的函数是不相似的。这样就可以轻松地得到大量训练集。可以使用这个训练集预训练一个任务独立的模型，它对于大多数任务是有效的。此外，为了将可用的参考标准数据包含在特定于任务的策略中，文中的方法允许将模型重新训练，以合并特定于任务的数据。

图嵌入网络

文章的图嵌入网络由Dai等人的Structure2vec修改而来。将一个ACFG表示为$g=\langle V,\varepsilon \rangle$，其中$V$和$\varepsilon$分别是节点和边的集合。每个节点$v$可以有额外的属性$x_v$。图嵌入网络首先为每个节点$v \in V$计算出一个$p$维的特征嵌入$\mu _v$，之后$g$的嵌入$\mu _g$由这些顶点嵌入聚合得到。本文简单地使用求和函数，即$\mu _g = \sum _{v\in V} (\mu _v)$，将其它函数的探索与应用作为未来的工作。

在Structure2vec中，节点属性$x_v$是由图拓扑$g$次迭代生成。几轮次迭代后，网络会为每个节点产生一个嵌入，其中包含图形特征和顶点之间的交互特征.令$N(v)$表示图$g$中节点$v$的相邻顶点集合，Structure2vec网络将为每个节点$v$初始化一个嵌入$\mu _v^{(0)}$，并在每次迭代中更新嵌入。公式为：

$$
\mu v^{(t+1)} = F(x_v, \sum{u\in N(v)} \mu _u^{(t)} ), \forall v \in V.
$$

其中函数$F$会在下文介绍。据此公式，可以看出嵌入更新过程是基于图拓扑结构进行的，并以同步方式进行。在所有顶点的嵌入更新完成之后，新的一轮的嵌入才会开始。迭代次数越多，一个顶点的特性就会传播到越遥远的顶点。

这里函数$F$的参数由机器学习得到，而不是手工设置。训练一个针对分类问题的Structure2vec模型，已经有成熟的方法。但如前所述，本文的任务不是分类，因此不能使用这些方法。文章使用了Siamese架构，使用这些嵌入进行端到端训练，下文将会介绍。

以下说明函数$F$如何使用神经网络参数化。函数$F$被设计为：

$$
F(x_v, \sum_{u\in N(v)} \mu u) = \tanh (W_1x_v + \sigma(\sum{u\in N(v)} \mu _u))
$$

其中，$x_v$是图节点特征的$d$维嵌入，$W_1$是一个$d\times p$的矩阵，$p$是上面提到的嵌入维度。$\sigma$是一个全连接的$n$层神经网络：

$$
\sigma (l) = \underbrace{P_1\times ReLU(P_2\times \ldots ReLU(P_nl))}_{n\ levels}
$$

其中$P_n$是$p\times p$的矩阵。$ReLU$是线性纠正单元，$ReLU(x) = max (0, x)$。

以上过程可以用如图的伪代码表示。其中，$W_2$是另一个$p\times p$的矩阵，用来变换嵌入。用$\phi (g)$表示输出。

使用Siamese架构进行学习

文中使用Siamase架构结合Structure2vec网络进行机器学习。Siamase架构使用两个相同的图嵌入网络，这里即Structure2vec网络。每个图嵌入网络使用一个ACFG$g_i (i=1, 2)$作为输入，输出嵌入$\phi (g_i)$，最后，Structure2vec网络的输出是两个嵌入的余弦距离。两个图嵌入网络共享同样的参数，因此，它们在训练过程中也是相同的。结构如图：

给定一组ACFG对$\langle g_i, g’_i \rangle$，数量为$K$，其中参考标准信息是$y_i$，$y_i = +1$说明$g_i$与$g`_i$相似，$y_i = -1$反之。定义每个对的Siamese网络的输出为：

$$
Sim(g_i, g’_i) = \cos (\phi (g), \phi (g’)) = \frac{\langle g_i, g’_i \rangle}{\parallel \phi (g)\parallel \cdot \parallel \phi (g’)\parallel }
$$

其中$\phi (g)$由上述算法1产生。

之后，为训练模型参数$W_1, P_1, \ldots, P_n$与$W_2$，将会优化以下目标函数：

$$
\underset{W_1, P_1, \ldots, P_n, W_2}{min} \sum^{K}_{i=1}(Sim(g_i, g’_i) - y_i)^2
$$

我们可以用随机梯度下降法优化上述目标函数。参数的梯度根据图的拓扑结构迭代计算。最后，一旦Siamese网络实现了良好的性能（例如，使用AUC作为度量），训练过程就终止了，而经过训练的图嵌入网络可以将输入图转换成有效的嵌入，以适用于相似性检测。

任务独立预训练与任务相关再训练

训练模型需要有大量与oracle$\pi$相关的、带有参考标准的数据，但这些数据可能难以获得。为解决这个问题，文章使用默认的策略建立一个训练集。这个训练集可以用于预训练一个对大多数任务有效的任务独立模型。当有可用的任务相关数据时，预训练模型能够快速训练得到一个任务相关模型。

为预训练一个适用于大多数常见任务的模型，直观看来，生成每个函数的嵌入的时候，应该尝试在不同的体系结构和编译器中捕获函数的不变特性。我们通过使用默认的oracle构建数据集来实现这个直觉。假设收集了一组源代码，我们可以将它们编译成不同架构的程序二进制文件，使用不同的编译器，并使用不同的优化。如果两个二进制函数是由相同的源代码编译的，默认的oracle确定它们是相似的，反之不相似。

但某些任务使用的策略与默认策略不同。这时，需要一个有效的方法调整模型。再训练过程通过结合领域专家提供的特定于任务策略的少量额外数据，改进了图嵌入网络。设有一些函数对$\langle g_i, g’_i \rangle$与其参考标准$\pi (g_i, g’_i)$，使用这些数据，再对图嵌入网络做几轮训练。在每一轮中，新加入的数据会被取样更多次（例如，取样次数是旧数据的50倍）。之后，经过再训练的网络$\phi (\cdot)$会被用于相似性检测任务。这样的一个再训练过程允许领域专家对系统提供反馈，确保模型被调整为适用于特定任务，从而提高相似性检测的准确率。

小结

本文通过参考机器学习领域已有算法，对热点问题提出了新颖高效的解法。文章没有照搬现有方法，而针对任务对经典方法作出了合理的修改，值得学习和参考。