融合注意力機制和Child-Sum Tree-LSTM的二進制代碼相似性檢測

李 濤，王金雙

(中國人民解放軍陸軍工程大學 指揮控制工程學院，江蘇 南京 210007)

0 引言

代碼復用在軟件開發過程中非常普遍。大量代碼和庫被復用到多個體系架構的二進制文件中，其中存在的一些脆弱代碼可能會導致安全隱患。二進制代碼相似性檢測技術可以評估兩個或者多個二進制代碼間的相似性和差異性，被廣泛應用于漏洞發現[1-2]、代碼剽竊檢測[3]、惡意軟件檢測[4]等領域。

二進制代碼相似性檢測的任務是尋找與給定查詢二進制函數func最相似的目標函數func′。根據檢測對象的不同，主要分為以下兩種類型：(1)一對一比較：直接計算兩段二進制代碼之間的相似性。這種方法通常需要設置一個合適的閾值，閾值的選擇將直接影響檢測的準確度。(2)一對多比較。該場景將一個源函數與目標函數池中的函數進行相似度評分并排序。目標函數的排序越靠前，說明模型對于函數語義的區分度越好。

編譯選項的差異(如編譯優化O0-O3、-fno-inline選項等)、代碼混淆技術(如Obfuscator-LLVM[5])的使用給研究人員提取二進制代碼語義帶來了諸多挑戰。

基于深度學習的二進制代碼相似度檢測方法可以自動化地提取代碼語義，從而避免了人工分析可能引入的偏差。常用的神經網絡主要有兩類：(1)使用處理結構特征的圖神經網絡，如Structure2Vec[6]、GCN[7]等；(2)使用處理文本序列的神經網絡，如LSTM(Long Short Term Memory)[8]、Transformer[9]等。

Genius[10]結合基本塊語法特征和控制流圖(Control Flow Graph，CFG)生成屬性控制流圖用于相似性檢測，之后的Gemini[11]、VulSeeker[1]等均使用了類似的思想。

在二進制代碼相似度檢測領域，Luo等人[12]利用LSTM和孿生網絡學習匯編代碼的語義表示。Asm2vec[13]使用PV-DM[14]模型學習函數和指令符號的嵌入。jTrans[15]在文獻[16]預訓練任務的基礎上，將控制流信息融入預訓練任務，取得了較好的效果。但是這些方法無法適用于跨指令集架構的檢測場景。

觀察來自庫findutils中的closs_stream函數，如圖1所示，在保持編譯配置相同(編譯器Clang7.0，編譯優化O0，其他編譯器選項相同)的情況下，不同指令集架構(x86和ARM)二進制代碼間的CFG具有一定的差異，而其二進制代碼對應的AST特征則只出現了一個節點的差異。

圖1 二進制代碼的CFG和AST特征對比(bigram工程closs_stream函數O0優化)

Tai[17]等人提出了兩種Tree-LSTM網絡(Child-Sum Tree-LSTM和N-ary Tree-LSTM)用于處理自然語言處理領域的結構化數據，其在情感分類和語義相關任務中的性能優于普通LSTM。其中，N-ary Tree-LSTM適合于二叉樹，而Child-Sum Tree-LSTM適用于多孩子的無序樹，并且具有更高的計算效率。

在此工作的基礎上，Yang等人[18]提出了一種基于AST編碼的跨指令集架構的檢測方法Asteria。通過提取二進制函數的AST作為跨指令集架構的二進制特征來源，并利用能夠處理樹形數據的Tree-LSTM網絡學習二進制代碼的AST語義嵌入。但是原生Tree-LSTM難以捕獲子節點對整棵樹的語義貢獻度，影響訓練精度。

AST的節點代表了代碼中的表達式或語句，不同種類的節點承載著不同的信息量。然而Child-Sum Tree-LSTM對子節點隱藏狀態累加以表示父節點，未能考慮不同類別AST節點對整棵樹特征表示的貢獻程度。例如，AST中包含的num等節點通常和expr、block等節點具有一定的從屬關系，并且通常在樹中處于不同的位置。進行節點向量傳播的過程中應當對這些節點的語義信息對進行重要性區分，使得在模型訓練過程中節點對整棵AST提供適當的語義貢獻度。

本文提出了一種融合注意力機制和Child-Sum Tree-LSTM的二進制代碼相似性檢測方法。首先使用二進制分析工具IDA Pro對二進制文件進行反編譯提取AST作為特征來源，輸入到融合注意力機制的Child-Sum Tree-LSTM神經網絡中進行訓練，最后通過實驗表現驗證本方法的有效性。所提方法可用于跨指令集架構、跨代碼混淆、跨編譯優化等級等多種二進制代碼相似性檢測場景。

1 相關工作

1.1 二進制代碼相似度檢測

現有二進制代碼特征提取方法主要可分為動態分析特征提取和靜態分析特征提取。動態分析方法能夠準確獲取代碼的功能語義。動態分析環境搭建復雜度較高，難以適用于大規模分析場景。靜態分析方法利用反匯編或者反編譯技術能夠獲得較高的代碼覆蓋率，但是難以準確獲取程序的實際執行語義。

根據特征來源的不同，現有的靜態二進制代碼相似度檢測技術可以分為以下幾類：基于語法特征、基于結構特征和基于語義特征的檢測技術。

基于語法特征的檢測技術提取二進制代碼中的助記符或操作碼、操作數、指令序列等作為相似性檢測的特征來源。此類方法依賴于對檢測目標、檢測場景的精確設計，易受編譯優化、代碼混淆技術干擾。

基于結構特征的檢測技術通常提取二進制代碼的數據流轉信息，如CFG、數據流圖等作為相似性檢測的特征來源。并且通常需要結合如匯編代碼對二進制函數進行綜合表示。然后利用圖匹配算法等度量兩段二進制代碼的相似度。

基于語義特征的檢測技術獲取代碼的語義信息作為相似性檢測的特征來源。隨著深度學習技術的發展，研究人員利用深度神經網絡學習代碼的語義信息，已經取得了較好的效果。

1.2 循環神經網絡

循環神經網絡被廣泛應用于文本表示、神經機器翻譯、時間序列預測與建模等任務，適用于具有時間順序的任務。

LSTM通過引入門機制用于控制信息傳輸，用于解決標準循環神經網絡不能處理長期依賴而出現的梯度消失或梯度爆炸問題。

為了進一步提升神經網絡的學習能力，研究人員在LSTM的基礎上進行創新，如GRU[19]、BiLSTM[20]、Tree-LSTM、GraphLSTM[21]等。

GRU將LSTM中的輸入門和遺忘門合并為更新門，結合重置門用于控制序列信息的記憶或者遺忘。相較于LSTM，GRU具有更少的參數量，具備更快的收斂速度。BiLSTM是一種雙向的LSTM，相較于LSTM能夠更好地捕獲文本的雙向語義。

1.3 注意力機制

注意力機制最早出現在圖像領域，被廣泛應用在基于循環神經網絡/卷積神經網絡的視覺圖像、自然語言處理任務中。

注意力機制能夠增強模型對輸入序列中不同元素的關注程度。它允許模型在預測和生成任務時，動態地聚焦于任務中的輸入部分。

2 模型框架與訓練流程

本節介紹本方法的框架和訓練流程，并給出相應的符號描述。

2.1 模型概述

圖2描述了模型的工作流，共包含特征提取、模型訓練、相似性檢測3個階段：(1)使用IDA Pro對二進制文件進行反編譯并提取AST。由于AST中包含了節點的名稱、類別等特征，無法直接輸入網絡中，因此需要對AST進行節點映射。(2)使用基于注意力機制的Child-Sum Tree-LSTM作為骨干網絡。AST中的節點通過Child-Sum Tree-LSTM后，以從葉子節點到根子節點的順序進行遍歷，采用孿生網絡架構對模型進行訓練，最終生成整個AST的語義向量。(3)利用語法樹相似度衡量二進制函數的相似度。

圖2 模型框架和訓練流程

圖3展示了對二進制函數相似性檢測的流程，從二進制代碼中提取的AST通過先序遍歷編碼生成初始向量輸入模型編碼后，得到AST嵌入，然后利用向量相似度度量語法樹相似度。

圖3 相似性檢測實例(xorriso.elf的Xorriso_option_extract_cut函數)

語法樹相似度定義如下：給定兩個二進制函數f1和f2及其對應的AST表示T1和T2，使用Tree-LSTM對樹型特征進行編碼得到對應的向量N(T1)和N(T1)。最后使用兩個AST的相似度得分衡量兩段二進制代碼的相似性，AST相似度Sim(T1，T2)計算公式如式(1)所示：

Sim(T1，T2)=

softmax(σ(cat(|N(T1)-N(T2)|，N(T1)⊙N(T2))×W))

(1)

其中T=，V和E為AST的點與邊的合集。N(T)為使用模型對T編碼后的向量。

2.2 特征提取階段

特征提取階段使用二進制分析工具IDA Pro7.0的Hex-rays插件對二進制文件進行反編譯提取AST。由于7.0版本的Hex-rays插件僅支持ARM和x86指令集的反編譯，在后續的實驗中本文使用ARM、x86作為模型訓練和測試的指令集架構來源。

從二進制代碼中提取的AST節點是字符串形式的數據，無法直接輸入到神經網絡進行訓練，通過將節點映射到固定的整型數值實現對節點名稱的編碼[18]。AST的節點被分為語句節點和表達式節點2類。語句節點包括控制函數執行流的節點，而表達式節點則包含各種用于計算的節點，如表1所示。

表1 AST中的標簽映射關系

2.3 模型訓練階段

模型訓練階段用于實現對二進制AST語義嵌入。LSTM能夠捕獲句子中的詞序，并且有效地解決了RNN梯度消失的問題。但是原生LSTM只能處理線性數據，通過反編譯提取到的AST無法直接用于LSTM模型中。

Child-Sum Tree-LSTM能夠利用AST所有子節點的語義向量進行聚合和前向傳播，以生成父節點和整棵樹的語義向量。由于其適用于多孩子節點的情況，并且具有更高的計算效率，因此本文采用此網絡作為對比網絡。

每個Child-Sum Tree-LSTM單元包含輸入輸出門ij和oj以及用于每個孩子k的遺忘門fjk。其內部門狀態由所有子節點的隱藏狀態和更新，如式(2)所示：

(2)

其中，hk表示子節點k的隱藏狀態，C(j)表示節點j的孩子節點。

(3)

(4)

(5)

其中xj表示節點嵌入，由節點Vj嵌入到高維向量生成，W(o)、W(c)、U(i)、U(o)、U(c)、b(i)、b(o)和b(c)為需要學習的參數矩陣。其中σ表示sigmoid函數，能夠將輸入的向量轉化為[0，1]區間內的向量。

Child-Sum Tree-LSTM含了k個遺忘門，用于合并來自每個子節點的單個信息。

fjk=σ(W(f)xj+U(f)hk+b(f))

(6)

然后，多個遺忘門向量輸出與相應的單元狀態相乘后組合得到單個遺忘門向量，模型的最終單元狀態cj如式(7)所示：

(7)

最后，Child-Sum Tree-LSTM的隱藏狀態hj更新方程如式(8)所示：

hj=oj⊙tanh(cj)

(8)

每一個參數矩陣代表輸入xi和組件單元的第k個子節點的隱藏狀態hk之間的相關性。

Child-Sum Tree-LSTM對AST進行從葉子節點到根節點進行遍歷時，由不同組成部分或從屬部分的信息首先被組合起來表示上層的根節點，然后這個根節點再作為子節點繼續遍歷的過程。在Child-Sum Tree-LSTM中，每個節點都會接收它的子節點的隱藏狀態和細胞狀態，并將它們相加得到總的表示。

注意力機制因為在圖像領域取得了很好的效果而備受關注，它允許模型在做下游任務時，能夠對句子中的數據進行重要性區分。受相關研究的啟發，本文將注意力機制和Child-Sum Tree-LSTM進行融合，用于學習AST的語義向量。融合注意力機制的Child-Sum Tree-LSTM的隱藏狀態更新方程如式(11)所示：

Watt=softmax(attention(hj))

(9)

hatt=Watt⊙hj

(10)

hj=oj⊙(tanh(cj)+hatt)

(11)

其中，attention定義為一個線性層，維度為記憶單元向量的維度。注意力層引入的過程使得模型能夠根據注意力權重動態地關注不同子節點，并在計算當前節點的隱藏狀態時加入子節點的注意力加權和。

綜上，使用神經網絡建模二進制代碼AST對的相似性包括兩個步驟：通過一個語義編碼器將兩段二進制代碼表示為語義向量，然后使用這兩個向量進行分類過程。訓練后的模型可以用于二進制代碼相似度檢測。

3 實驗配置

3.1 實驗環境及超參數介紹

本方法部署在一臺搭載Ubuntu 20.04 LTS操作系統的工作站中。處理器為一顆Inter Xeon(R)Gold 5218R CPU@2.10 GHz，內存大小為128 GB，顯卡為一張NVIDIA Corporation 3090，顯存大小為24 GB。神經網絡向量嵌入維度為32，記憶狀態維度為200，訓練epoch為30，學習率為0.01。訓練結果為3次實驗平均值。

3.2 實驗數據集介紹

為了評估本方法的有效性，使用開源數據集BINKIT[22]中的NOLINE和OBFUSCATION兩種類別的項目作為訓練集和測試集。為了實現對跨指令集架構以及跨混淆代碼的檢測，選擇的所有項目都以Clang 7.0作為編譯器，編譯優化為O0-O3，指令集架構為ARM和x86，字長為32，包含bcf、sub和fla三種代碼混淆方式，具體數據集介紹如表2所示。

表2 訓練和測試數據集介紹

函數對池的構造方式如下：以函數名作為函數語義標簽，對無混淆數據和三種混淆數據使用隨機采樣生成訓練數據對。其中，共生成7組函數池，每種函數池中包含4種優化的數據。具體如表3所示。

表3 訓練和測試樣本對構造

訓練集和測試集的構造方式為：將函數對池打亂后隨機采樣函數對，共生成54 000個函數對用于模型訓練過程，6 000個函數對用于模型測試過程。其中正樣本對和負樣本對的比例為1∶1，無混淆數據對和混淆數據對的比例為2∶1，使用1和-1對相似對和不相似對進行標注。

3.3 評價指標介紹

本文選用的評價指標為Accuracy、AUC和MSE。

AUC(Area Under the Curve)：AUC是ROC曲線下面積，AUC越大，模型的辨別能力越好。ROC曲線通過繪制真陽性率與假陽性率之間的關系來表示模型在不同閾值下的表現，能夠說明模型的診斷能力。AUC反映了模型能否正確判斷AST是否為相似對的概率。函數對的相似性被計算為r的分數，并且設置閾值為β，如果得分r大于或者等于β，則該函數對被認為是正結果，否則判斷為負結果。

(12)

(13)

其中，TP、FP、TN、FN分別表示：二進制函數對為正，預測為正；二進制函數對為負，預測為正；二進制函數對標簽為正，預測為負；二進制函數對標簽為負，預測為負。

Accuracy：準確率，用于衡量模型分類的性能。計算模型在所有樣本中正確分類的比例，準確率的值越高表示模型的分類準確性越好。計算公式如式(14)所示：

(14)

MSE(Mean Squared Error)：均方誤差，用于衡量模型的性能。MSE用于計算預測值與真實值之間的平均平方差，其對預測值之間的差異性較為敏感。MSE的值越小表示模型對于觀測值的擬合效果越好。

3.4 實驗結果

實驗結果如表4、圖4至圖6所示，訓練過程中本文方法在測試集上的MSE曲線較低，證明本方法在相似對判別中具有更好的性能。同時本文方法在測試集中的AUC指標為94.1%，比Child-Sum Tree-LSTM網絡93.6%的結果提高了0.5%，Accuracy指標為66.2%，比Child-Sum Tree-LSTM網絡65.8%的結果提高了0.4%。

表4 測試集結果比較

圖4 訓練過程測試集MSE對比

圖5 訓練過程測試集AUC對比

圖6 訓練過程測試集Accuracy對比

3.5 超參數實驗

本節主要用于測試向量維度對于模型性能的影響，包括詞向量的嵌入維度和隱藏狀態維度。消融實驗epoch為5。

從表5可以看出，模型的效果受到輸入向量維度和隱藏狀態維度的影響。實驗結果表明，當隱藏狀態為200時，模型的效果優于100。持續增大隱藏狀態維度會較大地增加算力消耗，所以本文采用的隱藏狀態為200。當輸入向量維度為16和32時，模型的測試結果接近。綜上，本文選擇的輸入維度為32，隱藏狀態維度為200。

表5 超參數驗證

4 結論

本文提出了一種融合注意力機制和Child-Sum Tree-LSTM的二進制代碼相似度檢測技術，用于跨指令集架構、跨代碼混淆技術的檢測場景。首先利用二進制分析工具提取跨指令集架構、跨代碼混淆的二進制代碼AST作為特征；然后構建訓練AST數據對用于模型訓練；最后利用融合注意力機制和Child-Sum Tree-LSTM的神經網絡學習AST的語義信息。實驗表明，注意力機制增強了Child-Sum Tree-LSTM對二進制AST的語義學習能力。