作者：騰訊科恩實驗室
原文鏈接：https://mp.weixin.qq.com/s/fvnvh25VaHgproPvVT6LDg
導語
在NeurIPS 2020中，騰訊安全科恩實驗室使用AI算法解決二進制安全問題的《CodeCMR: Cross-Modal Retrieval For Function-Level Binary Source Code Matching》論文成功入選。本論文首次提出了基于AI的二進制代碼/源代碼端到端匹配算法，與傳統算法相比效果非常出色，準確率大幅提升。本論文成果為逆向分析領域提供了新的思路，大大提升工業部署效率。最新論文研究成果也將應用于騰訊安全科恩實驗室研發的代碼檢索工具BinaryAI，使用體驗請關注：http://github.com/binaryai/sdk。
關于NeurIPS會議
機器學習和計算神經科學領域的NeurIPS會議是人工智能領域最具影響力的頂級學術會議之一，備受學者們的關注。國際頂級會議NeurIPS 2020將于2020年12月7日-12日在線上舉行。據統計，NeurIPS 2020收到投稿9454篇，創歷史最高紀錄，接收論文1900篇，論文接收率僅有歷史最低的20.1%。
背景
論文鏈接：CodeCMR: Cross-Modal Retrieval For Function-Level Binary Source Code Matching https://keenlab.tencent.com/zh/whitepapers/neurips-2020-cameraready.pdf 在人工智能頂級學術會議AAAI 2020中，騰訊安全科恩實驗室利用圖神經網絡解決二進制程序函數相似性分析問題的技術得到了廣泛關注。在此基礎上，本次研究方向擴展到二進制代碼與源代碼的交叉領域，進一步實現騰訊安全科恩實驗室在AI+安全新興方向中的全新探索與突破。
二進制代碼-源代碼匹配是信息安全領域的重點研究方向之一。在給定二進制代碼的情況下，逆向分析研究人員希望找到它對應的源代碼，從而提升逆向分析的效率和準確率。但由于源代碼和二進制代碼的差異性，在此領域的研究較少。B2SFinder[1]和BinPro[2]等傳統算法提取源代碼和二進制代碼的字符串、立即數等特征進行匹配。然而，函數級別的源代碼與二進制代碼的特征非常少，匹配準確率不高。另一方面，設計合適的特征需要大量的專家經驗。
圖1展示了一個函數的源代碼與二進制代碼。從圖1中可以看出，除了字符串和立即數特征，代碼中隱藏的語義特征也很關鍵。因此，本文希望設計一種端到端模型，可以自動提取代碼間的語義特征，從而提升匹配的準確率。
模型
這是一個二進制代碼-源代碼間的檢索任務，我們把兩種代碼當作兩個模態的輸入，即可類比到圖文互搜等跨模態檢索場景。因此，我們設計了如圖2所示的CodeCMR框架，在跨模態檢索領域中，這是一種比較常見的結構[3, 4]。在計算最終向量之前，兩個模態之間沒有信息傳遞，因此在實際應用時可以預先計算向量，可以節省大量的線上計算時間以及存儲空間。
整體結構
模型的輸入有源代碼特征和二進制代碼特征兩個部分。其中源代碼特征是字符級別的源代碼、從源代碼中提取的字符串和立即數；二進制代碼特征是控制流圖、二進制代碼的字符串和立即數。首先將三個輸入（語義特征、字符串特征、立即數特征）分別用不同模型計算得到向量，再用拼接+BatchNorm的方式得到代碼向量，最后用triplet loss[5]作為損失函數。
在這個基礎框架上，有許多可以改進的創新點，例如使用預訓練模型做語義融合、使用adversarial loss對齊向量等，對此我們將在后文討論。
語義模型
如圖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所示，本文提出的方法與傳統方法相比有巨大提升，這一發現符合我們的預期，說明代碼間隱含的語義特征十分重要。在語義模型中，DPCNN+HBMP取得了最優的效果，說明在二進制側端到端訓練優于預訓練的node embedding。與隨機采樣和distance weighted采樣方法相比，norm weighted采樣效果更好。圖4的train/valid loss曲線也證明了這一點，當s=5時norm weighted sampling的train loss更高但valid loss更低，這表示采樣到更合適的樣例pair。
討論與總結
討論
基于CodeCMR框架，有很多值得嘗試的創新。
1）code encoder。ASTNN、Tree-LSTM、transformer等模型可能也同樣有效。
2）其它損失函數和采樣方法，如AM-softmax、Circle loss等。
3）對抗訓練以及其它的跨模態檢索領域的方法。
4）預訓練算法。在獲得最終向量前兩個模態沒有信息融合，因此在兩個模態分別單獨預訓練或用跨語言模型的方法融合訓練，均是值得嘗試的。
總結
本文針對二進制代碼-源代碼匹配任務提出了CodeCMR框架，成功地利用了源代碼與二進制代碼間的語義特征。與傳統方法相比，取得了很大的突破。
參考文獻
[1]Yuan Z, Feng M, Li F, et al. B2SFinder: Detecting Open-Source Software Reuse in COTS Software[C]//2019 34th IEEE/ACM International Conference on Automated Software Engineering (ASE). IEEE, 2019: 1038-1049.
[2]Miyani D, Huang Z, Lie D. Binpro: A tool for binary source code provenance[J]. arXiv preprint arXiv:1711.00830, 2017.
[3]Wang H, Sahoo D, Liu C, et al. Learning cross-modal embeddings with adversarial networks for cooking recipes and food images[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2019: 11572-11581.
[4]Wang B, Yang Y, Xu X, et al. Adversarial cross-modal retrieval[C]//Proceedings of the 25th ACM international conference on Multimedia. 2017: 154-162.
[5]Schroff F, Kalenichenko D, Philbin J. Facenet: A unified embedding for face recognition and clustering[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2015: 815-823.
[6]Johnson R, Zhang T. Deep pyramid convolutional neural networks for text categorization[C]//Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics. 2017: 562-570.
[7]Yu Z, Cao R, Tang Q, et al. Order Matters: Semantic-Aware Neural Networks for Binary Code Similarity Detection[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2020, 34(01): 1145-1152.
[8]Wu C Y, Manmatha R, Smola A J, et al. Sampling matters in deep embedding learning[C]//Proceedings of the IEEE International Conference on Computer Vision. 2017: 2840-2848.
本文由 Seebug Paper 發布，如需轉載請注明來源。本文地址：http://www.bjnorthway.com/1389/
暫無評論