图神经网络和代码大模型加持 | 提升二进制比对分享

一、背景

验证补丁的困惑

众所周知,寻找并定位二进制代码之间存在(源)代码复用、存在的差异,其准确率及召回率是2个关键指标。另外,无恒实验室意识到,面对IoT设备ROM中数以千计的应用和动态链接库的补丁分析,自动化的特征+语义维度的二进制比对在提升效率上非常重要。

基于此,无恒实验室对市面的工具进行调研。用radiff2做补丁扫描,召回率堪忧,扫描完之后仍然需要手动返工来消除误报,且其效率感人;再换用bindiff,召回率几乎99%,但当研发切换编译工具链或者修改编译优化等级,则准确率就会明显下降。

对比主流工具

记XO为不同编译优化等级场景,XC1为同类编译器不同版本,XC2为不同类编译器,以TALOS所用数据集中的libssl与libz作为数据集[1],strip后则可测得

工具
召回率
F1

XO
XC1
XC2
XO
XC1
XC2

radiff2
33.00%
37.04%
10.29%
47.90%
53.46%
17.94%
BinDiff
92.17%
95.42%
78.53%
93.27%
96.20%
80.00%

可以看到

radiff2表现很差
BinDiff表现优于radiff2,但在XC2下召回率显著下降

看来是不得不选择BinDiff了,那么问题来了,如何提升BinDiff的召回率呢?

我们首先需要了解不同类型的二进制相似度检测任务,以及代码克隆检测任务。

二、前置知识

如何划分二进制相似度检测任务

根据TALOS的综述,可以将二进制相似度检测任务划分为6类[1](下表将编译器与编译器版本合并为一栏)——

优化等级
编译器
架构
位数
XO
不同
相同
相同
相同
XC
相同
不同
相同
相同
XC+XB
不同
不同
相同
不同
XA
相同
相同
不同
不同
XA+XO
不同
相同
不同
不同
XM
不同
不同
不同
不同

如何描述代码克隆类型

我们把代码相似度分为下面4种[2]

Type-1/2/3已经有成熟解决方案比如SourcererCC[3]等,而Type-4也有不少工作,比如Technion的code2vec[4]。此外,近年来随着大模型发展,UniXcoder[5]、CodeT5+[6]等代码大模型在Type-4的代码克隆检测上也能通过微调取得不错的表现。

三、函数相似度与匹配

初代废弃方案

年初在DEFCON GROUP 0571做过一次二进制对比的分享,用Genius的Raw Feature Similarity[7],先计算两个基本块的相似度,记基本块中各属性权重为$$w_{p}$$

$$sim_{bb}(bb, bb^{prime})=frac{sum_{p}{jaccard(bb_{p}, bb^{prime}{p})*w{p}}}{sum_{p}{w_{p}}}$$

再将基本块间相似度作为二部图节点间距,记$$E(g,g^{prime})$$为二部图节点距离集合,

$$E(f, f^{prime})={sim_{bb}(bb,bb^{prime})}, bbin{f},bb^{prime}in{f^{prime}}$$

把二部图最大匹配代价作为函数相似度,记二部图匹配函数为$$match(g,g^{prime},E(g,g^{prime}))$$,二部图最大代价为$$cost()$$

$$sim(f,f^{prime})=frac{cost(match(f, f^{prime}, E(f, f^{prime})))}{min(left|fright|,left|f^{prime}right|)}$$

虽然能够得到函数相似度,但该方案效率低下,二部图最大匹配所用Hopcroft-Karp算法时间复杂度为$$mathcal{O}(left|Eright|*sqrt{left|Vright|})$$,如果函数之间的匹配也以来二部图匹配,那么函数数量与单个函数基本块节点数都成了瓶颈,实际使用比radiff2还慢,因此我们把目光投向更快的相似度计算方法。

structure2vec

借TALOS综述所言,我们选择用structure2vec[8]来生成ACFG图嵌入。在GNN下,CFG + IR + Int + Strings各类任务中最佳,其次是CFG + BoW IR/opc[1]。

但structure2vec的节点属性是定长向量,而基本块中的Int、String之类的constants序列是不定长的,因此我们对Gemini的ACFG进行拓展,使用CFG + BoW Hexrays IR + len(constants)作为基本块节点向量来构成ACFG。

获得ACFG embedding后,用两个embedding的余弦相似度作为ACFG相似度,记structure2vec为$$g_s$$

$$sim_{acfg}(f, f^{prime})=1-frac{1}{2}cos(g_s(f), g_s(f^{prime}))$$

为了弥补len(constants)和BoW Hexrays IR的信息损失,在函数体粒度维护constants、mnemonics等属性的序列,用属性序列的Jaccard相似度作为补救相似度(这里参考了Genius的基本块相似度计算方法,把基本块替换成函数体即可),记函数体中各属性序列权重为$$w_{p}$$,各属性序列为$$f_{p}$$

$$sim_{remedy}(f,f^{prime})=frac{sum_{p}{jaccard(f_{p}, f^{prime}{p})*w{p}}}{sum_{p}{w_{p}}}$$

二者加权得到最终的函数相似度,记ACFG相似度权重为$$w_a$$,补充相似度权重为$$w_j$$

$$sim_{a,r}(f,f^{prime})=frac{w_{a}*sim_{acfg}(f,f^{prime})+w_{j}*sim_{remedy}(f,f^{prime})}{w_{a}+w_{j}}$$

如此一来,把函数体相似度从二部图最大匹配代价换成了余弦相似度与jaccard相似度的加权,降低计算复杂度,同时也解决了原计算方式中通过中介中心度和子节点数量来保留结构信息的缺陷。

业余炼丹士

structure2vec embedding + function properties就够了吗?XO与XC会导致控制流的变化,ACFG也会随之变化,这会降低structure2vec的表现。但是,函数语义是不变的,因此我们可以站在两位巨人肩膀上——代码大模型与反编译伪代码,通过代码大模型对C风格伪代码做Type-4的代码克隆检测。

这里我们选择UniXcoder[5]与CodeT5+[6]的encoder来生成embedding,并用POJ104作为数据集微调——

microsoft/unixcoder-base-nine,125M参数,微调代码:https://github.com/microsoft/CodeBERT/blob/master/UniXcoder/downstream-tasks/clone-detection/POJ-104/run.py
Salesforce/codet5p-110m-embedding,110M参数,微调代码:https://github.com/Cossack9989/CodeT5/blob/main/CodeT5%2B/tune_codet5p_clone_detection.py

在8张T4跑完两轮微调任务后,模型的代码克隆检测准确率如下

unixcoder
codet5+
90.72%
87.92%

获得Type-4 embedding后,取不同模型下两个embedding的余弦相似度的最大值作为T4相似度,记encoder为$$g_e$$

$$sim_{t4}(f,f^{prime})=1-frac{1}{2}min(cos(g_e(f), g_e(f^{prime}))), g_ein{set{unixcoder, codet5p}}$$

与补救相似度加权得到最终的函数相似度,

$$sim_{r,t}(f,f^{prime})=frac{w_{j}*sim_{remedy}(f,f^{prime}) + w_{t4}*sim_{t4}(f,f^{prime})}{w_{j}+w_{t4}}$$

当$$sim_{a,r}$$相似度小于阈值(默认0.98)时,会使用$$max(sim_{a,r}, sim_{r,t})$$作为最终的相似度

匹配流程

暂时无法在飞书文档外展示此内容

四、metrics

实验设置

我们选择TALOS的DataSet-1的strip后的libz.so.1.2.11和libssl.so.3作为数据集[1],架构为ARM64,实验环境是8c32g MBP(M1 Pro),记我们的方法为BinDiff+AI(合并BinDiff与GNN+LLM的对比结果)

链接为 https://drive.google.com/drive/folders/1g9P0KKSwqdFt0K6dDeKKhWfmhqiQHQqU?usp=sharing

我们关注如下几个任务:

XO
XC1/XC2
XO+XC1/XO+XC2

因为目标是自动化补丁检测,所以我们无须关注XA、XB相关的任务

实验结果

工具
召回率
F1

XO
XC1
XC2
XO+XC1
XO+XC2
XO
XC1
XC2
XO+XC1
XO+XC2

BinDiff
92.17%
95.42%
78.53%
90.39%
78.19%
93.27%
96.20%
80.00%
91.62%
79.68%
BinDiff+AI
93.07%
95.18%
83.23%
92.35%
82.49%
94.26%
95.90%
85.60%
93.68%
84.95%

从F1看,我们发现,BinDiff加上embedding后,XO与XC1基本保持一致,XO+XC1得到了2%左右的提升,XC2与XO+XC2都得到了5%左右提升

五、一些思考

调用链特征

二部图匹配忽略了函数之间的调用关系,匹配流程第8步callee的匹配也仅仅只是覆盖了相邻调用关系,这个特征的缺失或许很重要。

代码大模型

本次实验只使用了UniXcoder[5]和CodeT5+[6]的encoder来生成伪代码语义embedding,但在Type-4的代码克隆检测上已经不再是SOTA,可以考虑换别的代码大模型。

六、关于无恒实验室

无恒实验室是由字节跳动资深安全研究人员组成的专业攻防研究实验室,致力于为字节跳动旗下产品与业务保驾护航。通过漏洞挖掘、实战演练、黑产打击、应急响应等手段,不断提升公司基础安全、业务安全水位,极力降低安全事件对业务和公司的影响程度。

无恒实验室希望持续与业界共享研究成果,协助企业避免遭受安全风险,亦望能与业内同行共同合作,为网络安全行业的发展做出贡献。

参考文献

[1] Andrea Marcelli, Mariano Graziano, Xabier Ugarte-Pedrero, Yanick Fratantonio, Mohamad Mansouri, and Davide Balzarotti. How machine learning is solving the binary function similarity problem. In 31st USENIX Security Symposium (USENIX Security 22), pages 2099–2116, Boston, MA, August 2022. USENIX Association.

[2] Chanchal Kumar Roy and James R. Cordy. A survey on software clone detection research. SCHOOL OF COMPUTING TR 2007-541, QUEEN’S UNIVERSITY, 115, 2007.

[3] Hitesh Sajnani, Vaibhav Saini, Jeffrey Svajlenko, Chanchal K. Roy, and Cristina V. Lopes. Sourcerercc: scaling code clone detection to big-code. In Proceedings of the 38th International Conference on Software Engineering, ICSE ’16. ACM, May 2016.

[4] Uri Alon, Meital Zilberstein, Omer Levy, and Eran Yahav. code2vec: Learning distributed representations of code, 2018.

[5] Daya Guo, Shuai Lu, Nan Duan, Yanlin Wang, Ming Zhou, and Jian Yin. Unixcoder: Unified cross-modal pre-training for code representation. arXiv preprint arXiv:2203.03850, 2022.

[6] Yue Wang, Hung Le, Akhilesh Deepak Gotmare, Nghi D.Q. Bui, Junnan Li, and Steven C. H. Hoi. Codet5+: Open code large language models for code understanding and generation. arXiv preprint, 2023.

[7] Qian Feng, Rundong Zhou, Chengcheng Xu, Yao Cheng, Brian Testa, and Heng Yin. Scalable graph-based bug search for firmware images. In CCS 2016 – Proceedings of the 2016 ACM SIGSAC Conference on Computer and Communications Security, 2016.

[8] Hanjun Dai, Bo Dai, and Le Song. Discriminative embeddings of latent variable models for structured data, 2020.

​ 

Read More 

發佈留言

發佈留言必須填寫的電子郵件地址不會公開。 必填欄位標示為 *