科恩實(shí)驗(yàn)室最新NeurIPS-2020論文解讀：基于跨模態(tài)檢索的二進(jìn)制代碼-源代碼匹配
本論文成果為逆向分析領(lǐng)域提供了新的思路，大大提升工業(yè)部署效率。
在NeurIPS 2020中，騰訊安全科恩實(shí)驗(yàn)室使用AI算法解決二進(jìn)制安全問題的《CodeCMR: Cross-Modal Retrieval For Function-Level Binary Source Code Matching》論文成功入選。本論文首次提出了基于AI的二進(jìn)制代碼/源代碼端到端匹配算法，與傳統(tǒng)算法相比效果非常出色，準(zhǔn)確率大幅提升。本論文成果為逆向分析領(lǐng)域提供了新的思路，大大提升工業(yè)部署效率。最新論文研究成果也將應(yīng)用于騰訊安全科恩實(shí)驗(yàn)室研發(fā)的代碼檢索工具BinaryAI，使用體驗(yàn)請關(guān)注：http://github.com/binaryai/sdk。
關(guān)于NeurIPS會(huì)議
機(jī)器學(xué)習(xí)和計(jì)算神經(jīng)科學(xué)領(lǐng)域的NeurIPS會(huì)議是人工智能領(lǐng)域最具影響力的頂級(jí)學(xué)術(shù)會(huì)議之一，備受學(xué)者們的關(guān)注。國際頂級(jí)會(huì)議NeurIPS 2020將于2020年12月7日-12日在線上舉行。據(jù)統(tǒng)計(jì)，NeurIPS 2020收到投稿9454篇，創(chuàng)歷史最高紀(jì)錄，接收論文1900篇，論文接收率僅有歷史最低的20.1%。
背景
論文鏈接：CodeCMR: Cross-Modal Retrieval For Function-Level Binary Source Code Matching
在人工智能頂級(jí)學(xué)術(shù)會(huì)議AAAI 2020中，騰訊安全科恩實(shí)驗(yàn)室利用圖神經(jīng)網(wǎng)絡(luò)解決二進(jìn)制程序函數(shù)相似性分析問題的技術(shù)得到了廣泛關(guān)注。在此基礎(chǔ)上，本次研究方向擴(kuò)展到二進(jìn)制代碼與源代碼的交叉領(lǐng)域，進(jìn)一步實(shí)現(xiàn)騰訊安全科恩實(shí)驗(yàn)室在AI+安全新興方向中的全新探索與突破。
二進(jìn)制代碼-源代碼匹配是信息安全領(lǐng)域的重點(diǎn)研究方向之一。在給定二進(jìn)制代碼的情況下，逆向分析研究人員希望找到它對應(yīng)的源代碼，從而提升逆向分析的效率和準(zhǔn)確率。但由于源代碼和二進(jìn)制代碼的差異性，在此領(lǐng)域的研究較少。B2SFinder[1]和BinPro[2]等傳統(tǒng)算法提取源代碼和二進(jìn)制代碼的字符串、立即數(shù)等特征進(jìn)行匹配。然而，函數(shù)級(jí)別的源代碼與二進(jìn)制代碼的特征非常少，匹配準(zhǔn)確率不高。另一方面，設(shè)計(jì)合適的特征需要大量的專家經(jīng)驗(yàn)。
圖1展示了一個(gè)函數(shù)的源代碼與二進(jìn)制代碼。從圖1中可以看出，除了字符串和立即數(shù)特征，代碼中隱藏的語義特征也很關(guān)鍵。因此，本文希望設(shè)計(jì)一種端到端模型，可以自動(dòng)提取代碼間的語義特征，從而提升匹配的準(zhǔn)確率。
圖1 – 二進(jìn)制代碼與對應(yīng)的源代碼
模型
這是一個(gè)二進(jìn)制代碼-源代碼間的檢索任務(wù)，我們把兩種代碼當(dāng)作兩個(gè)模態(tài)的輸入，即可類比到圖文互搜等跨模態(tài)檢索場景。因此，我們設(shè)計(jì)了如圖2所示的CodeCMR框架，在跨模態(tài)檢索領(lǐng)域中，這是一種比較常見的結(jié)構(gòu)[3, 4]。在計(jì)算最終向量之前，兩個(gè)模態(tài)之間沒有信息傳遞，因此在實(shí)際應(yīng)用時(shí)可以預(yù)先計(jì)算向量，可以節(jié)省大量的線上計(jì)算時(shí)間以及存儲(chǔ)空間。
圖2 – CodeCMR整體框架
整體結(jié)構(gòu)
模型的輸入有源代碼特征和二進(jìn)制代碼特征兩個(gè)部分。其中源代碼特征是字符級(jí)別的源代碼、從源代碼中提取的字符串和立即數(shù)；二進(jìn)制代碼特征是控制流圖、二進(jìn)制代碼的字符串和立即數(shù)。首先將三個(gè)輸入（語義特征、字符串特征、立即數(shù)特征）分別用不同模型計(jì)算得到向量，再用拼接+BatchNorm的方式得到代碼向量，最后用triplet loss[5]作為損失函數(shù)。
在這個(gè)基礎(chǔ)框架上，有許多可以改進(jìn)的創(chuàng)新點(diǎn)，例如使用預(yù)訓(xùn)練模型做語義融合、使用adversarial loss對齊向量等，對此我們將在后文討論。
圖3 – 源代碼與二進(jìn)制代碼的語義模型
語義模型
如圖3所示，對于字符級(jí)源代碼，我們使用的是DPCNN模型[6]；對于二進(jìn)制控制流圖，我們使用的是端到端的GNN模型。在函數(shù)級(jí)別，字符級(jí)源代碼的輸入通常在4096以上，DPCNN的效果遠(yuǎn)優(yōu)于TextCNN和LSTM。對于控制流圖，我們沒有使用BERT預(yù)訓(xùn)練的node embedding作為輸入[7]，而是采用了端到端訓(xùn)練的方式，取得了更好的效果。
在這個(gè)階段，本文使用的是DPCNN和GNN，但ASTNN等樹模型也同樣值得嘗試。由于輸入是函數(shù)級(jí)別的代碼，缺少#define、#include等重要信息，需要設(shè)計(jì)合適的編譯工具將源代碼轉(zhuǎn)化為AST。相比之下，我們直接將文本作為輸入的優(yōu)點(diǎn)是無需額外的專家經(jīng)驗(yàn)，健壯性強(qiáng)。
立即數(shù)、字符串模型
上文提到，源代碼和二進(jìn)制代碼的字符串和立即數(shù)不完全相同，因此也需要設(shè)計(jì)模型進(jìn)行匹配。
對于立即數(shù)，我們設(shè)計(jì)了一種Integer-LSTM。它的輸入有integer token和integer number兩個(gè)。integer number作用在LSTM的輸入門和輸出門，從而控制信息流動(dòng)。
對于字符串，我們使用的是分層模型，先用LSTM模型得到每個(gè)字符串的向量，再使用sum pooling的方法得到字符串集合的向量。
Norm weighted sampling
在得到源代碼與二進(jìn)制代碼的向量后，我們設(shè)計(jì)了一種采樣方法。在metric learning領(lǐng)域中，損失函數(shù)和采樣方法是十分重要的兩個(gè)技巧。為了解決hard樣本在訓(xùn)練早期收斂到局部極小值的問題，[5]提出了semi-hard采樣方法。然而，[8]指出這種采樣方法可能會(huì)在某個(gè)時(shí)間段停止訓(xùn)練，從而提出了distance weighted sampling采樣方法解決這個(gè)問題：
distance weighted sampling可以在分布中選擇各個(gè)概率的樣本，而semi-hard、hard、uniform等采樣方法只能選擇特定分布的樣本。在此基礎(chǔ)上，本文提出了一個(gè)改進(jìn)，即增加一個(gè)超參數(shù)s，幫助調(diào)整概率的分布，從而適應(yīng)不同的任務(wù)和數(shù)據(jù)集。
實(shí)驗(yàn)
數(shù)據(jù)集與評(píng)測指標(biāo)
本文分別用gcc-x64-O0和clang-arm-O3作為兩種組合方式，制作了兩個(gè)30000/10000/10000的訓(xùn)練/驗(yàn)證/測試集，并使用recall@1和recall@10作為評(píng)測指標(biāo)。數(shù)據(jù)集已公開在https://github.com/binaryai。
表1 – 實(shí)驗(yàn)結(jié)果
實(shí)驗(yàn)結(jié)果
如表1所示，本文提出的方法與傳統(tǒng)方法相比有巨大提升，這一發(fā)現(xiàn)符合我們的預(yù)期，說明代碼間隱含的語義特征十分重要。在語義模型中，DPCNN+HBMP取得了最優(yōu)的效果，說明在二進(jìn)制側(cè)端到端訓(xùn)練優(yōu)于預(yù)訓(xùn)練的node embedding。與隨機(jī)采樣和distance weighted采樣方法相比，norm weighted采樣效果更好。圖4的train/valid loss曲線也證明了這一點(diǎn)，當(dāng)s=5時(shí)norm weighted sampling的train loss更高但valid loss更低，這表示采樣到更合適的樣例pair。
圖4 – 訓(xùn)練與驗(yàn)證的損失函數(shù)曲線
相關(guān)閱讀
題-1.jpg)
