摘要:第三個分塊是對order-aware模塊的對比實驗,基礎CNN模型使用3層CNN以及7、11層的Resnet,CNN_random是對訓練集中控制流圖的節點順序隨機打亂再進行訓練,MPNN_ws是去除控制流圖節點中的語義信息(所有block向量設爲相同的值)再用MPNN訓練。在semantic-aware模塊,模型將控制流圖作爲輸入,使用BERT[2]對token embedding作預訓練,得到block embedding。

騰訊安全科恩實驗室《Order Matters: Semantic-Aware Neural Networks for Binary Code Similarity Detection》論文入選人工智能領域頂級學術會議AAAI-20。研究核心是利用AI算法解決大規模二進制程序函數相似性分析的問題,本文將深入對該論文進行解讀,點擊鏈接獲取完整論文。https://keenlab.tencent.com/en/whitepapers/Ordermatters.pdf

二進制函數相似性比對演示效果:

gif

論文: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等行業,並根據產業實際痛點和研究推出了智能網聯汽車信息安全行業解決方案。護航各行業數字化變革,守護全網用戶的信息安全是騰訊科恩實驗室的使命。

相關文章