AI 算法解决二进制安全问题,腾讯安全NeurIPS 2020论文有新方法

不久前,人工智能顶会 NeurIPS 2020 放出接收论文列表,论文接收率创历史新低。而腾讯安全科恩实验室使用 AI 算法解决二进制安全问题的一项研究《CodeCMR: Cross-Modal Retrieval For Function-Level Binary Source Code Matching》成功入选。

论文地址:https://keenlab.tencent.com/zh/whitepapers/neurips-2020-cameraready.pdf

这篇论文首次提出基于 AI 的二进制代码 / 源代码端到端匹配算法,与传统算法相比效果非常出色,准确率大幅提升。该论文的成果为逆向分析领域提供了新思路,大大提升工业部署效率。

最新研究成果也将应用于腾讯安全科恩实验室研发的代码检索工具 BinaryAI,使用体验请关注:http://github.com/binaryai/sdk。

研究背景

在人工智能顶级学术会议 AAAI 2020 中,腾讯安全科恩实验室利用图神经网络解决二进制程序函数相似性分析问题的技术得到了广泛关注。在此基础上,本次研究将方向扩展到二进制代码与源代码的交叉领域,进一步实现腾讯安全科恩实验室在「AI + 安全」这一新兴方向中的全新探索与突破。

二进制代码 - 源代码匹配是信息安全领域的重点研究方向之一。在给定二进制代码的情况下,研究人员希望找到它对应的源代码,从而提升逆向分析的效率和准确率。但由于源代码和二进制代码的差异性,在此领域的研究较少。B2SFinder [1] 和 BinPro [2] 等传统算法提取源代码和二进制代码的字符串、立即数等特征进行匹配。然而,函数级别的源代码与二进制代码特征非常少,匹配准确率不高。另一方面,设计合适的特征需要大量的专家经验。

图 1 展示了一个函数的源代码与二进制代码。从图 1 中可以看出,除了字符串和立即数特征,代码中隐藏的语义特征也很关键。因此,本研究希望设计一种端到端模型,可以自动提取代码间的语义特征,从而提升匹配的准确率

图 1:二进制代码与对应的源代码

模型

这是一个二进制代码 - 源代码检索任务,研究人员把这两种代码当作两个模态的输入,即可类比到图文互搜等跨模态检索场景。因此,该研究设计了如图 2 所示的 CodeCMR 框架。

图 2:CodeCMR 整体框架

在跨模态检索领域中,这是一种比较常见的结构 [3, 4]。在计算最终向量之前,两个模态之间没有信息传递,因此在实际应用时可以预先计算向量,从而节省大量的线上计算时间以及存储空间。

整体结构

模型的输入有源代码特征和二进制代码特征两个部分。其中源代码特征是字符级别的源代码、从源代码中提取的字符串和立即数;二进制代码特征是控制流图、二进制代码的字符串和立即数。

首先将三个输入(语义特征、字符串特征、立即数特征)分别用不同模型计算得到向量,再用拼接 + BatchNorm 的方式得到代码向量,最后用 triplet loss [5] 作为损失函数。


在这个基础框架上,有许多可以改进的创新点,例如使用预训练模型做语义融合、使用 adversarial loss 对齐向量等,下文将对此进行讨论。

语义模型

图 3:源代码与二进制代码的语义模型

如图 3 所示,对于字符级源代码,该研究使用的是 DPCNN 模型 [6];对于二进制控制流图,则使用的是端到端的 GNN 模型。在函数级别,字符级源代码的输入通常在 4096 以上,DPCNN 的效果远优于 TextCNN 和 LSTM。对于控制流图,该研究没有使用 BERT 预训练的 node embedding 作为输入 [7],而是采用了端到端训练的方式,取得了更好的效果。

在这个阶段,该研究使用的是 DPCNN 和 GNN,但 ASTNN 等树模型也同样值得尝试。由于输入是函数级别的代码,缺少 #define、#include 等重要信息,需要设计合适的编译工具将源代码转化为 AST。相比之下,该研究直接将文本作为输入的优点是无需额外的专家经验,稳健性强

立即数、字符串模型

上文提到,源代码和二进制代码的字符串和立即数不完全相同,因此需要设计模型进行匹配。


对于立即数,该研究设计了一种 Integer-LSTM。它的输入包括 integer token 和 integer number 两种。integer number 作用在 LSTM 的输入门和输出门,从而控制信息流动。

对于字符串,该研究使用的是分层模型,先用 LSTM 模型得到每个字符串的向量,再使用 sum pooling 的方法得到字符串集合的向量。

Norm weighted sampling

在得到源代码与二进制代码的向量后,研究人员设计了一种采样方法。在 metric learning 领域中,损失函数和采样方法是十分重要的两个技巧。为了解决 hard 样本在训练早期收敛到局部极小值的问题,[5] 提出了 semi-hard 采样方法。然而,[8] 指出这种采样方法可能会在某个时间段停止训练,并提出了 distance weighted sampling 采样方法来解决这个问题:


distance weighted sampling 可以在分布中选择各个概率的样本,而 semi-hard、hard、uniform 等采样方法只能选择特定分布的样本。在此基础上,本研究提出了一项改进,即增加一个超参数 s,帮助调整概率的分布,从而适应不同的任务和数据集。


实验

数据集与评测指标

本研究分别用 gcc-x64-O0 和 clang-arm-O3 作为两种组合方式,制作了两个 30000/10000/10000 的训练 / 验证 / 测试集,并使用 recall@1 和 recall@10 作为评测指标。数据集已公开,参见:https://github.com/binaryai。

实验结果

表 1:实验结果

如表 1 所示,本研究提出的方法与传统方法相比有巨大提升,这一发现符合研究人员的预期,说明代码间隐含的语义特征十分重要。在语义模型中,DPCNN+HBMP 取得了最优的效果,说明在二进制侧端到端训练优于预训练的 node embedding。与随机采样和 distance weighted 采样方法相比,norm weighted 采样效果更好。图 4 的 train/valid loss 曲线也证明了这一点,当 s=5 时 norm weighted sampling 的 train loss 更高但 valid loss 更低,这表示采样到了更合适的样例 pair。

图 4:训练与验证的损失函数曲线

关于 NeurIPS 会议

机器学习和计算神经科学领域的 NeurIPS 会议是人工智能领域最具影响力的顶级学术会议之一,备受学者关注。NeurIPS 2020 将于 2020 年 12 月 7 日 - 12 日在线上举行。据统计,NeurIPS 2020 收到投稿 9454 篇,创历史最高纪录,接收论文 1900 篇,论文接收率 20.1%,创历史最低记录。