code embedding研究系列一-基於token的embedding

I still …發表於2020-12-26

NLPembedding無處不在,embedding將大型稀疏向量轉換為保留語義關係的低維空間。在原始碼研究領域(原始碼分類,原始碼克隆檢測,變數,方法等重新命名等)也會需要用到embedding,即code embedding

code embedding可以分成很多種:
1.把原始碼文字當成純文字進行embedding (text embedding)
2.把原始碼經過詞法分析過後生成的token序列進行embedding (token embedding)
3.利用原始碼解析成的抽象語法樹**(AST), 從AST上提取資訊進行embedding
4.利用原始碼的控制流圖
(CFG)**和資料依賴圖(PDG)進行embedding (基於graph的embedding)

這裡Zimin Chen等人做了一個關於code embedding的統計研究

這裡列舉一些應用了token embedding的論文:

1.Automated software vulnerability detection with machine learning

這篇文章結局的問題是c/c++程式碼的bug檢測,任務則是二分類,程式碼分為GoodBuggy,檢測的粒度是function (即分析程式碼中的某個functionGood還是Buggy)

採用的方法結合了build-based(先編譯C程式碼)和source-based(直接分析原始碼文字),build-based分析了編譯後的程式碼的CFG(控制流圖),source-based則是基於token embedding的分析,這裡重點討論source-based (資料集中有一部分不能編譯成功)

資料集來源

資料集來自2部分:
1.Debian包
2.Github

資料集的統計資訊如下:

專案DebianGithubSource
Train Good429807250436754150
Train Buggy14616955631790
Test Good661003850094268
Test Buggy187812203724
Val Good551023210994268
val Buggy185712033724

其中Source為Debian和Github的總和

資料集預處理

1.資料下下來後首先要去除重複的function**(長得一模一樣)**
2.用Clangstatic analyzer (SA) 分析,將分析的輸出去除和漏洞,bug無關的警告。去除過後沒有警告資訊的被標記為Good,有資訊的被標記為Buggy。

分類方法

獲取token

將原始碼文字用自定義文法分析器解析成token序列,比如:

for (int i = 0; i <= 10; ++i){
	arr[i] += i;
}

生成的token序列如下:

for int i = 0 ; i <= 10 ; ++ i { arr [ i ] += i ; }

生成的token分類為:字串值, 單字元值, 字元陣列, 數值, 操作符, names等, names又細分為key(if for等等),系統呼叫(malloc),型別(bool, int, char),變數名。變數名又會被都對映為同一個識別符號(具體參考論文),string值,float值都會被對映。

向量化表示

上述對映過程後會生成一個token序列,之後便是要對每個token進行向量化,用到的方法有:
1.詞袋模型
2.word2vec

這樣便可以生成非監督下的embedding表示,一個word2vec的token表示二維化後如下:
在這裡插入圖片描述

分類模型

論文設計的如下:

word2vec
TextCNN Conv layer
TextCNN fully connected layer
Extra Tree Classifier
Bag of Words

在提取到token序列後,作者採用了3種策略來分類:
1.word2vec向量化 + 卷積提取特徵 + 全連線層分類
2.word2vec向量化 + 卷積提取特徵 + 樹分類器分類
3.Bag of Words向量化 + 樹分類器分類

實驗結果

評測標準採用ROC AUC和PR AUC。結果如下:

ModelROC AUCP-R AUC
BoW + ET0.850.44
word2vec + CNN0.870.45
word2vec + CNN feat. + ET0.870.49

build-based部分

這篇論文還有一部分關於build-based就不提了,大致過程就是根據彙編程式碼生成控制流圖(CFG),一個簡單的程式碼片段CFG如下:
在這裡插入圖片描述

CFG由一系列的基本塊(block)構成,每個基本塊相當於圖的一個node,每個block包括一些彙編程式碼(opcode),這些opcode不包括分支語句。整個CFG的特徵由2部分組成:

1.use-def矩陣
2.op-vec

use-def矩陣表示變數定義和使用,假設一個use-def矩陣matrix。如果一個區域性變數在基本塊i中定義,基本塊j中呼叫,則matrix[j][i]設為i

op-vec,每個基本塊用1個op-vec表示。這個op-vec是一個9維二進位制向量。分別對應 conditional, arrogate,binary, bit binary, conversion, memory address, termination, vector operation, and other 9個特徵,如果基本塊中存在對應該特徵的指令,則該特徵對應的值為1,否則為0。

之後結合CFG的鄰接矩陣每個程式碼塊的op-vec以及CFG的use-def矩陣構建1個116維向量(怎麼構建的就看不太懂了)來作為一個function的向量表示。通過隨機森林對116維向量進行分類。

2.CodeCMR: Cross-Modal Retrieval For Function-Level Binary Source Code Matching

騰訊安全科恩實驗室首次提出了基於AI的二進位制程式碼/原始碼端到端匹配演算法。大概就是給你一個二進位制程式碼,找出儘可能與之匹配的原始碼,對搞逆向的人來說是個好東西,最新論文研究成果也將應用於騰訊安全科恩實驗室研發的程式碼檢索工具BinaryAI

研究背景如下:在給定二進位制程式碼的情況下,逆向分析研究人員希望找到它對應的原始碼,從而提升逆向分析的效率和準確率(以後IDA逆向出來的虛擬碼可讀性會更強)。

當然,這裡我們更加關注token embedding部分。

總體架構

論文的總體架構如下:
在這裡插入圖片描述大致流程就是對一個c語言函式,把它編譯形成的二進位制程式碼與c語言程式碼本身形成一個pair,對這樣一個(c語言,二進位制)pair,分別把原始碼和二進位制程式碼向量化後,計算Triplet Loss,具體可以參考論文。

這裡source code(原始碼)被解析成一個很長的token序列。token序列被embedding後通過DPCNN後形成一個source語義向量,如下圖:
在這裡插入圖片描述
而原始碼的處理不止於此,還需要提取String特徵和Integer特徵,即原始碼文字中所有的字串的值和數值,對於下面的程式碼。
在這裡插入圖片描述if條件中的(void *)0中的0和exit(-1)中的-1會組成0 -1這樣的數值序列輸入Integer LSTM中生成一個Integer向量。

而原始碼中的字串有fatal error in SolveMap_nproc(%p)以及bad input,對於string特徵採用一個hierarchical-LSTM 模型處理string變數。來得到一個string向量。

把上述的token序列向量,integer向量以及string向量拼接在一起就形成了原始碼向量,對於二進位制程式碼向量的處理也類似,不同的是採用HBMP + GGNN + Set2Set 處理二進位制程式碼本身(關於二進位制程式碼就不展開了)。

至於為什麼不採用AST處理原始碼,論文中提到:
1.需要耗費更多時間
2.解析採用的演算法不一定能成功處理整段程式碼,比如用gcc來解析論文中使用的資料集,只有10%的程式碼被成功解析。
3.需要更多的專家知識。

資料集

資料集地址
資料集分成了2部分,一部分由gcc-x64-O0編譯後的二進位制程式碼和原始碼組成pairs,另一部分由clang-arm-O3編譯後的二進位制程式碼和原始碼組成的pairs組成。每部分包含50000個source-binary pairs, 30000用來training, 10000用來validation,10000用來testing。

實驗結果這裡就不提了(都是大佬寫的)。

總結

上述引用2篇論文,講述了以下code embedding之中的token embedding技術。大致流程就是將原始碼文法處理後的token序列作為原始碼的表示,應用到SE(軟體工程)任務。接下來我會分析下基於AST,CFG等的code embedding技術。

參考文獻

Harer J A , Kim L Y , Russell R L , et al. Automated software vulnerability detection with machine learning[J]. 2018.

Zeping Yu, Wenxin Zheng, Jiaqi Wang, Qiyi Tang, Sen Nie, Shi Wu. CodeCMR: Cross-Modal Retrieval For Function-Level Binary Source Code Matching

Chen Z , Monperrus M . A Literature Study of Embeddings on Source Code[J]. 2019.

相關文章