整理 | 高衛華
頭圖 | CSDN付費下載自視覺中國
近日，在NeurIPS 2020正式發布的論文入選名單中，騰訊安全科恩實驗室聚焦解決二進位安全問題的《CodeCMR: Cross-Modal Retrieval For Function-Level Binary Source Code Matching》，憑藉首次提出基於AI的二進位代碼/原始碼端到端匹配算法的創新研究入選。
該論文首次提出了基於AI的二進位代碼/原始碼端到端匹配算法，與傳統算法相比，準確率大幅提升，並為逆向分析領域提供了新的思路，可提升工業部署方面的效率。
目前，該論文成果已在騰訊安全科恩實驗室研發的代碼檢索工具BinaryAI實現了落地應用。
論文背景
在AAAI 2020中，騰訊安全科恩實驗室提出了一項利用圖神經網絡解決二進位程序函數相似性分析問題的技術，得到廣泛關注。
在此基礎上，科恩實驗室將研究方向擴展到二進位代碼與原始碼的交叉領域，進一步實現其在AI+安全新興方向中的全新探索與突破。
二進位代碼-原始碼匹配是信息安全領域的重點研究方向之一。在給定二進位代碼的情況下，逆向分析的研究人員希望可以找到它對應的原始碼，從而提升逆向分析的效率和準確率。
B2SFinder和BinPro等傳統算法提取原始碼和二進位代碼的字符串、立即數等特徵進行匹配，而函數級別的原始碼與二進位代碼的特徵非常少，匹配準確率不高。另外，要設計合適的特徵，需要大量的專家經驗。
圖1展示了一個函數的原始碼與二進位代碼，從中可看出，除了字符串和立即數特徵，代碼中隱藏的語義特徵也很關鍵。
因此，該論文設計了一種端到端的模型，可以自動提取代碼間的語義特徵，從而提升匹配的準確率。
圖1 - 二進位代碼與對應的原始碼
CodeCMR框架介紹
二進位代碼-原始碼間的檢索任務中，把兩種代碼當作兩個模態輸入，可類比到圖文互搜等跨模態檢索場景。
對此，科恩實驗室基於AI設計了CodeCMR框架，以解決傳統算法匹配準確率不高的問題。
計算最終向量前，兩個模態之間沒有信息傳遞，在實際應用時可預先計算向量，節省線上計算時間以及存儲空間。
圖2 - CodeCMR整體框架
整體結構
模型輸入中有原始碼特徵和二進位代碼特徵兩部分，其中原始碼特徵是字符級別的原始碼、從原始碼中提取的字符串和立即數，二進位代碼特徵是控制流圖、二進位代碼的字符串和立即數。
CodeCMR框架能夠以不同模型對原始碼特徵和二進位代碼特徵兩大模塊的語義、字符串、立即數等三大輸入特徵進行向量計算，並以在基礎上用拼接+BatchNorm方式實現高效匹配轉化，最後使用triplet loss作為損失函數。
語義模型
CodeCMR框架中，字符級原始碼使用DPCNN模型，二進位控制流圖使用端到端的GNN模型。
圖3 - 原始碼與二進位代碼的語義模型
立即數、字符串模型
原始碼和二進位代碼的字符串和立即數不完全相同，需要設計模型進行匹配。
對於立即數，科恩實驗室設計了Integer-LSTM，其輸入有integer token和integer number兩個。integer number作用在LSTM的輸入門和輸出門，以控制信息流動。
對於字符串，其採用分層模型，先通過LSTM模型得到每個字符串的向量，再使用sum pooling方法得到字符串集合的向量。
Norm weighted sampling
在兩大模塊的向量採樣階段，Norm weighted sampling 與隨機採樣和distance weighted採樣方法相比，表現出更好的匹配效果。
distance weighted sampling採樣方法可解決hard樣本在訓練早期收斂到局部極小值的問題，在分布中選擇各個機率的樣本。
在此基礎上，科恩實驗室又提出改進，增加超參數s，幫助調整機率的分布，從而適應不同的任務和數據集。
這種融合訓練的方法可大幅提升二進位代碼/原始碼的匹配準確率，從而提升工業部署效率。
實驗結果
與傳統算法僅提取字符串、立即數特徵進行匹配的做法不同，CodeCMR模型能夠實現對代碼間隱藏語義特徵的自動提取，從而達到提升二進位代碼-原始碼跨模態檢索匹配效率的目的。
從實驗結果中可以看出，在語義模型中，DPCNN+HBMP取得最優效果，表示在二進位側採用端到端訓練優於預訓練的node embedding；此外，與隨機採樣、distance weighted採樣相比，norm weighted的採樣效果更好。
表1 - 實驗結果
train/valid loss曲線也證明了這一點，當s=5時，norm weighted sampling的train loss更高，valid loss更低。
圖4 - 訓練與驗證的損失函數曲線
論文連結：
https://keenlab.tencent.com/zh/whitepapers/neurips-2020-cameraready.pdf