AAAI-20論文解讀:基于圖神經網絡的二進制代碼分析

首頁 > 國際新聞 > 工業互聯網安全應急響應中心
來源:工業互聯網安全應急響應中心 發布日期:2019-12-10 14:31 瀏覽:23次

引言&背景


二進制安全分析是信息安全中重要研究領域之一,其中一類場景是在不訪問源代碼的情況下檢測相似的二進制函數,也就是找到由同一份源代碼所編譯出的不同的二進制代碼。但是同一份源代碼在不同編譯器,不同平臺,不同優化選項的條件下所得到的二進制代碼截然不同。傳統算法使用圖匹配算法解決此問題,但圖匹配算法的速度較慢,且準確率較低。隨著近年來深度學習算法的發展,學者們嘗試在控制流圖(CFG)上使用圖神經網絡算法,取得了不錯的效果。

圖1. 控制流圖以及其低維向量的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),這個特征體現在它們的鄰接矩陣中,即圖6中被虛線方框圈出部分。首先對比圖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.




原文來源:騰訊科恩實驗室

論壇
  閱讀原文
支持0次 | 反對0次  
  用戶評論區,文明評論,做文明人!

通行證: *郵箱+888(如:[email protected]

北京足球单场比分直播 亿客隆时时彩票 下载广西麻将 捕鱼来了官网停运了 广西快乐10分开奖时间 股票质押 比分直播188比 哪里买手机便宜 波音足球指数皇冠正网 内蒙快三专家丨最新预测 ag金拉霸800倍游戏论坛 夏天买西瓜赚钱吗 棒球比分直播7m 3d怎么玩稳赚不赔 上海快3 写字能赚钱吗 东方6+1