AAAI-20論文解讀：基於圖神經網路的二進位制程式碼分析

騰訊安全科恩實驗室《Order Matters: Semantic-Aware Neural Networks for Binary Code Similarity Detection》論文入選人工智慧領域頂級學術會議AAAI-20。研究核心是利用AI演算法解決大規模二進位制程式函式相似性分析的問題，本文將深入對該論文進行解讀，點選“閱讀原文”獲取完整論文。

閱讀原文

二進位制函式相似性比對演示效果：

論文：Order Matters: Semantic-Aware Neural Networks for Binary Code Similarity Detection

單位 | 騰訊安全科恩實驗室

引言 & 背景

二進位制程式碼分析是資訊保安領域中非常重要的研究領域之一，其中一類目標是在不訪問原始碼的情況下檢測相似的二進位制函式。同一份原始碼在不同編譯器，不同平臺，不同優化選項的條件下所得到的二進位制程式碼是不相同的，我們的任務目標是把同一份原始碼所編譯出的不同的二進位制程式碼找到。傳統演算法使用圖匹配演算法解決此問題，但圖匹配演算法的速度較慢，且準確率較低。隨著近年來深度學習演算法的發展，學者們嘗試在控制流圖（CFG）上使用圖神經網路演算法，取得了不錯的效果。

圖1. 控制流圖（CFG）以及表示成低維向量的block特徵

論文[1]提出了名為Gemini的基於圖神經網路的演算法，它的輸入是兩個二進位制函式的pair，輸出是這兩個二進位制函式的相似度得分。首先，將二進位制函式的控制流圖作為輸入，並使用人工設計的特徵提取方法將每個block表示成低維的向量（如圖1所示）；然後使用Structure2vec演算法計算graph embedding；最後使用siamese網路計算相似度得分並使用梯度下降演算法降loss訓練模型（如圖2所示）。與傳統方法相比，Gemini的速度和準確率都大幅提升。

 

圖2. siamese網路結構

雖然上述方法取得了很大的進步，但仍有一些重要的問題值得研究。一方面，如圖1所示，每一個block被表示成一個低維向量，這個特徵提取的過程是人工設計的，在Gemini中block特徵只有8維向量，這個壓縮的過程會損失很多語義資訊。另一方面，在二進位制程式碼中節點的順序是一個很重要的特徵，而之前的模型沒有設計特殊的演算法提取這一特徵。圖3是函式"_freading"在不同平臺x86-64和ARM上編譯出的二進位制程式碼的控制流圖。這兩個控制流圖的節點順序是非常相似的，例如node1都與node2和node3相連，node2都與node4和node5相連，而這種相似性可以體現在它們的鄰接矩陣上。經過觀察，我們發現許多控制流圖的節點順序變化是很小的。為了解決以上兩個問題，我們設計了一種總體的框架，包含semantic-aware模組、structural-aware模組以及order-aware模組。

 

圖3. 函式"_freading"在不同平臺（x86-64和ARM）上編譯出的控制流圖以及對應的鄰接矩陣

模型

整體結構：模型的輸入為二進位制程式碼的控制流圖，模型的整體結構如圖4所示，包含semantic-aware 模組、structural-aware模組、order-aware模組。在semantic-aware模組，模型將控制流圖作為輸入，使用BERT[2]對token embedding作預訓練，得到block embedding。在structural-aware模組，使用MPNN演算法[3]得到graph semantic & structural embedding。在order-aware模組，模型將控制流圖的鄰接矩陣作為輸入，並使用CNN計算graph order embedding。最後對兩個向量使用concat和MLP得到最終的graph embedding，如公式1所示。

 

 

圖4. 模型整體結構

Semantic-aware 模組：在semantic-aware模組，可以使用BERT、word2vec等常用模型提取語義資訊。本文中使用BERT對控制流圖作預訓練，從而獲得block的語義資訊。BERT原用於NLP領域中，對詞語與句子作預訓練。我們的任務與NLP任務相似，控制流圖的block可以看作句子，block中的token可以看作詞語。如圖5所示，訓練過程中BERT有4個任務：Masked language model（MLM）、Adjacency node prediction（ANP）、Block inside graph（BIG）和Graph classification（GC）。

圖5. 語義資訊提取BERT模型

其中MLM和ANP是和BERT的原論文中相似的兩個任務。MLM是一個token-level的任務，對block中的token進行mask操作並進行預測，和語言模型的方式相同。ANP任務是一個block-level的任務，雖然控制流圖沒有NLP領域中的語言順序，但控制流圖是一個有向圖，也有節點的拓撲順序，我們將控制流圖中的所有相鄰節點提取出來，當作相鄰的“句子”。這些相鄰的block pair作為ANP任務的正例，並隨機選擇同圖內不相鄰的block pair作為負例。

為了獲得更多的graph-level的資訊，我們加入了兩個輔助的graph-level任務BIG和GC。BIG和ANP的方式類似，區別是pair的正負例選擇方式不同。BIG任務的目的是讓模型判斷兩個block是否在同一個圖中，希望模型可以儘可能地學到此資訊，從而對我們的graph-level task有幫助。因此，在BIG任務中同圖的block pair為正例，不同圖的block pair為負例。GC為graph-level的block分類任務，在我們的場景中，在不同平臺、不同編譯器、不同優化選項的條件下，得到的block資訊有所不同，我們希望模型可以讓block embedding中包含這種資訊。GC對block進行分類，判斷block屬於哪個平臺，哪個編譯器，以及哪個優化選項。

Structural-aware 模組：經過BERT預訓練後，使用MPNN計算控制流圖的graph semantic & structural embedding。MPNN有三個步驟：message function（M），update function（U）以及readout function（R）。具體步驟如公式2-公式4所示。

 

其中G代表整個圖，v代表節點，N(v)代表v的鄰居節點。在本文的場景中，節點即是控制流圖中的block，圖即是經過預訓練後表示成block向量的控制流圖。本文在message步驟使用MLP，update步驟使用GRU，readout步驟使用sum，如公式5-公式7所示。

 

Order-aware 模組：本模組希望可以提取節點順序的資訊，本文中使用的是CNN模型。為什麼使用CNN模型呢？首先考慮圖6中的三個圖（節點中無語義資訊），以及它們的鄰接矩陣。這三個圖非常相似，每個圖中都有一個三角形特徵（圖a的節點123，圖b的節點234，圖c的節點134），這個特徵體現在它們的鄰接矩陣中。首先對比圖a和圖b，與圖a相比，圖b加入了節點1，節點順序依次後移一位，但三角形特徵中三個節點的順序還是連續的，這個特徵在鄰接矩陣中可以看到，這個1-1-0-1的2*2矩陣仍然存在。CNN在訓練集中看過很多這種樣例後，可以學習到這種平移不變性。再看圖c，加入了節點2，打破了原有三角形的節點順序，但在鄰接矩陣中我們可以看到它實際上是把原來的2*2矩陣放大成了3*3矩陣，當我們移除第二行和第二列時，仍然可以得到一個1-1-0-1的2*2矩陣。這與影像中的image scaling類似，CNN在訓練集中包含足夠多樣例的情況下，也是可以學到這種伸縮不變性的。

 

圖6. 三個圖以及對應鄰接矩陣

本文中使用的模型是11層的Resnet結構[4]，包含3個residual block，所有的feature map大小均為3*3。之後用一個global max pooling層，得到graph order embedding。在此之前不用pooling層，因為輸入的圖的大小不同。具體如公式8所示。

 

實驗

本文在兩個任務上進行實驗。任務1為跨平臺二進位制程式碼分析，同一份原始碼在不同的平臺上進行編譯，我們的目標是使模型對同一份原始碼在不同平臺上編譯的兩個控制流圖pair的相似度得分高於不同原始碼pair的相似度得分。任務2為二進位制程式碼分類，判斷控制流圖屬於哪個優化選項。各資料集的情況如表1所示。任務1是排序問題，因此使用MRR10和Rank1作為評價指標。任務2是分類問題，因此使用準確率作為評價指標。

 

表1. 資料集情況

表2和表3分別對應任務1和任務2的實驗結果。表中第一個分塊是整體模型，包括graph kernel，Gemini以及MPNN模型。第二個分塊是semantic-aware模組的對比實驗，分別使用了word2vec[5]，skip thought[6]，以及BERT，其中BERT2是指原始BERT論文中的兩個task（即MLM和ANP），BERT4是指在此基礎上加入兩個graph-level task（BIG和GC）。第三個分塊是對order-aware模組的對比實驗，基礎CNN模型使用3層CNN以及7、11層的Resnet，CNN_random是對訓練集中控制流圖的節點順序隨機打亂再進行訓練，MPNN_ws是去除控制流圖節點中的語義資訊（所有block向量設為相同的值）再用MPNN訓練。最後是本文的最終模型，即BERT+MPNN+Resnet。

 

表2、3：各模型在任務1和任務2上的結果

整體結果：本文提出的模型與Gemini模型相比，在任務1和任務2上的評價指標分數均大幅提升。semantic-aware模組使用NLP模型（word2vec，BERT等）均優於使用人工提取的特徵。只使用order-aware時模型也取得了不錯的效果。與其它所有模型相比，本文提出的模型均取得了更優的效果。

Semantic-aware：只看表中第二個分塊，BERT的結果優於word2vec和skip thought，因為BERT能在預訓練過程中提取更多的資訊。加上BIG和GC任務後的BERT4效果略微提升，說明在預訓練過程中加入graph-level的任務有所幫助。圖7中是4個控制流圖的block（左上，左下，右上，右下），我們使用K-means對預訓練後的block embedding進行分類（K-means的類別數定為4），不同的類別顏色不同。從圖7中可以看出，同一個控制流圖中的block顏色大體相同，不同的控制流圖的block的主顏色大體不同。

圖7. 4個控制流圖的block embedding

Order-aware：觀察表中第三個分塊，CNN模型在兩個任務上都取得了不錯的效果。Resnet11優於Resnet7和CNN3。與MPNN_ws相比，CNN效果更優。隨機打亂節點順序後，CNN模型效果大幅下降，這表示CNN模型確實可以學到節點順序資訊。圖8是控制流圖pair的例子，這個函式為“ZN12libfwbuilder15RuleElementRGtw13validateC-hildEPNS8FWObjectE“，左邊是在gcc&x86-86上編譯的控制流圖，右邊是在gcc&ARM上編譯的控制流圖。可以看到，左圖的節點3在右圖中被拆成節點3和節點4，除此之外其它節點的順序與邊的連線方式均相同。經過CNN模型的計算，這兩個圖的cosine相似度為0.971，排序rank的排名為1。這表明CNN模型可以從鄰接矩陣中學到控制流圖的節點順序。

 

圖8. 控制流圖pair示例

結論

本文提出了一個新的模型，用於解決二進位制程式碼分析的問題。本文的模型中包含semantic-aware模組，structural-aware模組以及order-aware模組。我們觀察到語義資訊和節點順序資訊都是控制流圖重要的特徵。我們使用BERT預訓練模型提取語義資訊，並使用CNN模型提取節點順序資訊。實驗結果表明，本文提出的模型與之前最優的模型相比，取得了更好的效果。

參考文獻

[1] Xu, X.; Liu, C.; Feng, Q.; Yin, H.; Song, L.; and Song, D. 2017. Neural network-based graph embedding for crossplatform binary code similarity detection. In Proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security, 363–376. ACM.

[2] Devlin, J.; Chang, M.-W.; Lee, K.; and Toutanova, K. 2018. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 .

[3] Gilmer, J.; Schoenholz, S. S.; Riley, P. F.; Vinyals, O.; and Dahl, G. E. 2017. Neural message passing for quantum chemistry. In Proceedings of the 34th International Conference on Machine Learning-Volume 70 , 1263–1272. JMLR. org.

[4] He, K.; Zhang, X.; Ren, S.; and Sun, J. 2016. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, 770–778.

[5] Mikolov, T.; Sutskever, I.; Chen, K.; Corrado, G. S.; and Dean, J. 2013. Distributed representations of words and  phrases and their compositionality. In Advances in neural information processing systems , 3111–3119.

[6] R.; Zhu, Y.; Salakhutdinov, R. R.; Zemel, R.; Urtasun, R.; Torralba, A.; and Fidler, S. 2015. Skip-thought vectors. In Advances in neural information processing systems,3294–3302.

關於科恩實驗室

騰訊科恩實驗室作為騰訊集團旗下一支國際一流的資訊保安團隊，在桌面端安全、移動終端安全等研究領域有十多年的積累，技術實力和研究成果達到了國際領先水平；近幾年來，更是在智慧網聯汽車資訊保安、IoT 安全、雲端計算和虛擬化技術安全等領域取得豐碩的成果。隨著更多ICT新技術進入大眾視野，騰訊科恩實驗室也積極佈局人工智慧演算法和技術框架的安全研究、機器學習在資訊保安研究領域的應用研究和區塊鏈技術應用的安全研究等新緯度上的前沿技術研究能力。同時開放自身核心技術能力，提供給智慧網聯汽車、安卓應用生態、IoT等行業，並根據產業實際痛點和研究推出了智慧網聯汽車資訊保安行業解決方案。護航各行業數字化變革，守護全網使用者的資訊保安是騰訊科恩實驗室的使命。