基于殘差門控圖卷積網(wǎng)絡(luò)的源代碼漏洞檢測

張 俊，李山山，李 磊，王浩宇

武漢工程大學(xué) 智能機器人湖北省重點實驗室，武漢430205

21世紀公共漏洞披露CVE（common vulnerabilities and exposures）公開發(fā)表的報告表明，軟件漏洞數(shù)量正在迅速增加，這些軟件漏洞是網(wǎng)絡(luò)空間中遭到安全攻擊的主要原因，對社會造成嚴重的經(jīng)濟損失。如2017年，WanaCry病毒勒索軟件利用操作系統(tǒng)的一個漏洞，對全球造成了約80億美元的損失。并且傳統(tǒng)的漏洞檢測方法面臨著技術(shù)瓶頸，無法保障網(wǎng)絡(luò)空間安全。因此，如何自動化檢測漏洞是一個嚴峻的挑戰(zhàn)。

面對挑戰(zhàn)，目前漏洞檢測方法除了包含模糊測試[1]、靜態(tài)分析[2-3]、動態(tài)分析[4-6]這些方法之外，學(xué)術(shù)界也引入了一些更加新穎的方法，如引入了機器學(xué)習(xí)方法來輔助或者代替人工進行漏洞分析。早期的機器學(xué)習(xí)方法[7-9]主要是人工設(shè)計漏洞特征來檢測漏洞，然而這些方法也存在部分缺點，即設(shè)計漏洞特征往往需要專家花費大量時間，同時也存在主觀性因素，而且設(shè)計的漏洞特征無法包含新產(chǎn)生的漏洞[10]。采用這些機器學(xué)習(xí)方法的靜態(tài)分析工具檢測的結(jié)果會產(chǎn)生較高誤報率和漏報率問題。如Coverity 等商用靜態(tài)分析工具可被接受的最大誤報率為20%[11]，目前已知的應(yīng)用挖掘所得的構(gòu)件行為模型進行漏洞檢測最低誤報率為25%[12]，文獻[13-14]實驗結(jié)果中的誤報率更是高達50%和62.5%。

因此，為了降低現(xiàn)有方法的誤報率和漏報率，同時也避免人類專家在特征提取方面的大量勞動，學(xué)術(shù)界研究如何將深度學(xué)習(xí)應(yīng)用到漏洞檢測中，從而使得檢測方法更加自動化[15-17]。目前，一些做法是將每個源代碼視為一種自然語言序列，并應(yīng)用到自然語言處理的深度學(xué)習(xí)架構(gòu)中，如LSTM（long short-term memory）[18]、BERT（bidirectional encoder representations form transformers）[19]和CNN（convolutional neural network）[20]。基于深度學(xué)習(xí)的漏洞檢測可以自動生成漏洞特征，從而緩解人類專家手工定義漏洞特征的需求。李珍等人[21]提出了第一個使用深度學(xué)習(xí)的系統(tǒng)框架來檢測C/C++源代碼程序中的漏洞。該方法將代碼片段視為一段序列，通過RNN（recurrent neural network）對序列進行分類，但由于源代碼具有豐富的語義結(jié)構(gòu)信息，這種方法忽略這些語義結(jié)構(gòu)特征，存在較大的信息損失。因此如何提取源代碼的語義結(jié)構(gòu)特征，把源代碼轉(zhuǎn)換成具有綜合語義、適合神經(jīng)網(wǎng)絡(luò)訓(xùn)練的數(shù)據(jù)形式是關(guān)鍵問題之一。

最近，學(xué)術(shù)界又將圖神經(jīng)網(wǎng)絡(luò)應(yīng)用到代碼漏洞檢測中。圖神經(jīng)網(wǎng)絡(luò)（graph neural networks，GNN）是處理圖數(shù)據(jù)神經(jīng)網(wǎng)絡(luò)的統(tǒng)稱，是將節(jié)點和圖嵌入到低維連續(xù)向量空間的主要表示方法[22-24]。由于圖結(jié)構(gòu)強大的表示能力，GNN在處理文本分類等任務(wù)時，效果可以達到更佳[25-28]。基于圖的代碼表示方法包含了源代碼豐富的語義結(jié)構(gòu)特征，可以將源代碼編碼成具有綜合語義信息的向量。

鑒于已有的靜態(tài)分析漏洞檢測方法存在準確率、召回率等不高的問題，本文提出了一種基于殘差門控圖卷積網(wǎng)絡(luò)[29]的源代碼漏洞檢測方法。使用殘差門控圖卷積網(wǎng)絡(luò)可以緩解隨著網(wǎng)絡(luò)的深度增加造成的梯度消失問題，在一定程度上可以提高模型性能。方法首先通過提取源代碼的抽象語法樹結(jié)構(gòu)信息，以邊的形式保存在圖數(shù)據(jù)中，同時把抽象語法樹節(jié)點的出度作為圖節(jié)點的初始實值向量，結(jié)合兩者將源代碼轉(zhuǎn)換成圖結(jié)構(gòu)數(shù)據(jù)樣本，然后使用殘差門控圖卷積神經(jīng)網(wǎng)絡(luò)進行表示學(xué)習(xí)后，再使用神經(jīng)網(wǎng)絡(luò)來預(yù)測代碼漏洞。實驗結(jié)果顯示，本文設(shè)計的方案在VDISC數(shù)據(jù)集上取得了比基線方法更高的準確率、F1得分，降低了誤報率和漏報率。

1 方法設(shè)計

本文的方法架構(gòu)如圖1所示，它包含三部分：（1）數(shù)據(jù)處理部分，由于源代碼數(shù)據(jù)集中數(shù)據(jù)存在分布不均的問題，對數(shù)據(jù)進行下采樣操作，使得數(shù)據(jù)分布均勻。（2）提取抽象語法樹部分，通過提取代碼的抽象語法樹，并在抽象語法樹的基礎(chǔ)上增加語義、結(jié)構(gòu)信息，來構(gòu)建代碼圖。（3）殘差門控圖卷積網(wǎng)絡(luò)部分，將代碼圖作為模型的輸入，提取漏洞特征之后，利用神經(jīng)網(wǎng)絡(luò)來預(yù)測源代碼中是否含有漏洞。

圖1 方法架構(gòu)Fig.1 Method structure

本文的模型架構(gòu)如圖2所示，它包含四部分：（1）復(fù)合源代碼語義的代碼圖表示層，將一個函數(shù)的初始源代碼編碼成一個具有綜合程序語義的代碼圖。（2）殘差門控圖卷積網(wǎng)絡(luò)層，該層的作用是聚集和傳遞圖中相鄰節(jié)點的信息來學(xué)習(xí)節(jié)點的特征。（3）圖池化層，作用是聚合圖的全局信息，得到一個可以表示全圖特征的向量。（4）全連接層，作用是對提取的特征向量進行特征整合，最后得到預(yù)測值0 或1，0 代表沒有漏洞，反之1 則代碼含有漏洞。

圖2 模型架構(gòu)Fig.2 Model architecture

1.1 問題定義

漏洞檢測是判別源代碼中的某個函數(shù)是否含有漏洞，因此可以把漏洞檢測視為一個二分類問題。用數(shù)學(xué)語言來描述，把數(shù)據(jù)樣本定義為{(ci,yi)|ci∈C,yi∈Y},i∈{1,2,…,n}，其中C代表源代碼數(shù)據(jù)集，Y={0,1}n代表標簽集，1意味當前樣本含有漏洞，反之0意味當前樣本不具有漏洞，n為樣本的數(shù)目。ci被定義為一個函數(shù)，并為每一個ci構(gòu)建一個圖gi(V,X,A) ∈G。其中V是圖中的m個節(jié)點，X∈Rm×d是圖中節(jié)點的特征矩陣，每個節(jié)點vj∈V由一個d維實值向量xj∈Rd表示。A∈{0,1}m×m是一個鄰接矩陣，Ai,j值為1代表節(jié)點i與節(jié)點j之間有一條邊連接，反之Ai,j值為0 則代表節(jié)點i與節(jié)點j之間沒有邊連接。結(jié)合本文任務(wù)，可以歸納出是學(xué)習(xí)一個從G 到Y(jié)的映射關(guān)系，用函數(shù)表達f:G →Y,由此來預(yù)測一個數(shù)據(jù)樣本是否含有漏洞。映射函數(shù)f可以通過最小化損失函數(shù)的方式來學(xué)習(xí)，如式（1）所示。

其中，L(·) 是交叉熵損失函數(shù)，ω(·) 是正則化，λ是可調(diào)節(jié)權(quán)重。

1.2 代碼圖表示

本文以抽象語法樹為基礎(chǔ)來構(gòu)建代碼圖，使用提取AST的工具為Clang編譯器。圖3所示為提取AST的一個示例，從根節(jié)點開始，代碼被解析成變量聲明、函數(shù)名稱、復(fù)合語句、代碼塊、運算符操作等。如圖3中的節(jié)點1為函數(shù)聲明，節(jié)點2 為類型，節(jié)點5 為復(fù)合語句，節(jié)點7為變量，節(jié)點8為判斷語句，節(jié)點9為返回語句，節(jié)點10為運算操作等。父節(jié)點與子節(jié)點之間有一條邊連接，同級節(jié)點之間有一條虛邊連接，實際轉(zhuǎn)換成代碼圖表示中同級節(jié)點沒有邊相連。

圖3 抽象語法樹Fig.3 Abstract syntax tree

在代碼圖中，用一個鄰接矩陣來存儲圖數(shù)據(jù)，圖中的節(jié)點為AST中節(jié)點，AST節(jié)點中的父子關(guān)系則為圖的邊，例如鄰接矩陣A,其中Ai,j=1，意味著節(jié)點i與節(jié)點j之間有一條無向邊連接，反之則無邊相連。

在實值向量初始化上，采取的方法是通過節(jié)點的出度作為初始值，圖3中，節(jié)點1的出度為4，節(jié)點2的出度為0，節(jié)點4的出度為2等。

1.3 殘差門控圖卷積神經(jīng)網(wǎng)絡(luò)

自從圖神經(jīng)網(wǎng)絡(luò)（GNN）首次被提出以來，在處理圖結(jié)構(gòu)數(shù)據(jù)上得到了廣泛的應(yīng)用，如提取和發(fā)掘圖結(jié)構(gòu)數(shù)據(jù)的特征和模式，從而滿足聚類、分類、預(yù)測、分割、生成等圖學(xué)習(xí)任務(wù)需求。GNN核心思想是通過聚合圖中節(jié)點vi的特征xi與它的鄰居節(jié)點特征xj來生成節(jié)點vi新表示。基于不同聚合節(jié)點的信息的方法，誕生了圖卷積網(wǎng)絡(luò)（graph convolution networks，GCN）[30]、圖注意力網(wǎng)絡(luò)（graph attention networks，GAN）[31]、門控圖神經(jīng)網(wǎng)絡(luò)（gated graph neural networks，GGNN）[32]，以及殘差門控圖卷積神經(jīng)網(wǎng)絡(luò)（residual gated graph ConvNets，RGGCN）[29]等。本文選擇殘差門控圖卷積神經(jīng)網(wǎng)絡(luò)來學(xué)習(xí)圖節(jié)點特征。用數(shù)學(xué)語言來描述，給定一個圖gk(V,X,A)，對于每個圖中節(jié)點vi∈V都有一個初始D維向量，表示為當前節(jié)點的特征向量，如果實值向量xi的維度小于D時，則用0 填充，即[,0]T,則用式（2）描述為：

其中{:j→i} 表示節(jié)點i的所有鄰接節(jié)點的特征向量f為映射函數(shù)，輸入為當前節(jié)點的特征矩陣和當前節(jié)點所有鄰接節(jié)點的特征矩陣{→i},?為層數(shù)，W為權(quán)重矩陣。可以用式（3）描述為：，

其中⊙為哈達瑪積，ηi,j表示邊的門控機制，ηi,j的計算公式如式（4）所示：

其中σ表示為Sigmoid函數(shù)。

1.4 圖池化

對于圖神經(jīng)網(wǎng)絡(luò)而言，在傳播階段圖的結(jié)構(gòu)不會改變，改變的只是節(jié)點的特征向量，因此需要聚合圖的全局信息，得到一個可以表示全圖特征的向量，后續(xù)對該向量進行分類回歸操作，因此需要對圖進行池化操作。本文選取的是全局池化，其操作公式如式（5）所示：

其中R為池化函數(shù)，本文選取的是max、mean 函數(shù)，如式（6）所示：

其中cat(·) 為拼接函數(shù)，按列拼接，最后采用MLP（multilayer perceptron）進行特征整合，用式（7）描述為：

yi為最后的預(yù)測向量。

2 實驗設(shè)計與結(jié)果分析

2.1 數(shù)據(jù)集

數(shù)據(jù)集部分采用Draper VDISC Dataset-Vulnerability Detection in Source Code[33]數(shù)據(jù)集，該數(shù)據(jù)集包括從開源軟件中挖掘出的127萬個函數(shù)的源代碼，通過靜態(tài)分析對潛在漏洞進行標記。使用標記的工具包括Clang、Cppcheck 和Flawfinder。在標記之后，數(shù)據(jù)集的作者做了大量的工作來清理數(shù)據(jù)集重復(fù)的數(shù)據(jù)并去除錯誤的標簽。數(shù)據(jù)集中存在的漏洞類型為CWE-119（內(nèi)存緩沖區(qū)操作越界）、CWE-120（緩沖區(qū)溢出）、CWE-469（非法指針相減）、CWE-476（空指針引用）以及CWE-OTHERS（變量未初始化，以不正確的長度值訪問緩沖區(qū)等）。由于樣本數(shù)量龐大，有漏洞的函數(shù)占比較少，數(shù)據(jù)分布不均，因此對數(shù)據(jù)采用下采樣的方法，使得數(shù)據(jù)均勻分布。數(shù)據(jù)集樣本統(tǒng)計如表1所示，下采樣后數(shù)據(jù)樣本統(tǒng)計如表2所示。

表1 漏洞數(shù)據(jù)集Table 1 Vulnerability dataset

表2 下采樣后漏洞數(shù)據(jù)集Table 2 Under-sampling vulnerability dataset

2.2 基線方法

為了驗證本文方法的優(yōu)越性，將本文方法與VulDeePecker[15]、TextCNN[20]、GGNN[32]、文獻[34]進行對比實驗，綜合評估本文方法的性能。具體的數(shù)據(jù)處理過程如下所述。

VulDeePecker首先提取庫/API函數(shù)調(diào)用，并進行前向切片，獲取了相關(guān)代碼片段之后，組裝成codegadget并標注該codegadget是否具有漏洞，之后利用word2vec對其向量化，最終將向量輸入到BiLSTM 神經(jīng)網(wǎng)絡(luò)，進行模型訓(xùn)練。

TextCNN 將整個函數(shù)視為單位，利用word2vec 對函數(shù)進行向量化，之后輸入到網(wǎng)絡(luò)結(jié)構(gòu)為一層卷積和一層max-pooling以及一個softmax層的TextCNN網(wǎng)絡(luò)，其中設(shè)置詞向量維度為1 000，序列長度為500，類別為2。

GGNN方法使用本文方法的代碼圖表示方法，將函數(shù)轉(zhuǎn)換成向量之后，輸入到out_channels 為128 的6 層GGNN網(wǎng)絡(luò)中，聚合函數(shù)為add，并經(jīng)過一個Dense層匯聚網(wǎng)絡(luò)信息。

文獻[34]方法也就是本文使用數(shù)據(jù)集中論文的方法，該方法使用分詞器將函數(shù)向量化，就是基于one-hot編碼將函數(shù)向量化，之后將其輸入到帶有高斯噪聲的CNN中進行訓(xùn)練。

2.3 實驗設(shè)置

對于數(shù)據(jù)集而言，采用8∶1∶1 的比例來劃分訓(xùn)練集、驗證集以及測試集，網(wǎng)絡(luò)結(jié)構(gòu)采取3 層的殘差門控圖卷積，BatchSize 設(shè)置為256，學(xué)習(xí)率設(shè)置為0.001，損失函數(shù)為二分類交叉熵，同時選擇Adam優(yōu)化器訓(xùn)練模型，迭代次數(shù)為100。

2.4 實驗結(jié)果

2.4.1 采樣實驗結(jié)果與分析

從表1中可以看到，無漏洞的函數(shù)占比遠遠多于有漏洞的函數(shù)，而本文更關(guān)心的是有漏洞的函數(shù)，這時數(shù)據(jù)分布不均會更加突出。在這種數(shù)據(jù)分布下，訓(xùn)練出的模型顯然更加趨向預(yù)測無漏洞的函數(shù)，在數(shù)據(jù)集中，有漏洞的函數(shù)會被當成噪點或者被忽略，相比無漏洞的函數(shù)，有漏洞的函數(shù)被預(yù)測為無漏洞的函數(shù)可能性非常大。為了驗證這一結(jié)果，使用本文方法對下采樣后數(shù)據(jù)集和未采樣的數(shù)據(jù)集分別進行實驗，實驗結(jié)果如表3所示。

表3 采樣對比實驗結(jié)果Table 3 Results of sampling comparison experiments單位：%

表3中，未采樣的數(shù)據(jù)集雖然在準確率方面均高于下采樣后的數(shù)據(jù)集，但是F1 得分均低于下采樣后的數(shù)據(jù)集。這是由于未采樣的數(shù)據(jù)集無漏洞的函數(shù)占比太大，模型訓(xùn)練之后更傾向預(yù)測無漏洞的函數(shù)，而有漏洞的函數(shù)往往無法被正確地預(yù)測出來，導(dǎo)致了較高的誤報率和漏報率。下采樣后的數(shù)據(jù)集可以使得數(shù)據(jù)均勻分布，在一定程度上，可以避免這種問題，與未采樣的數(shù)據(jù)集相比，下采樣后的數(shù)據(jù)集在使用本文方法進行實驗后，得到了較高的F1得分，有效地降低了誤報率和漏報率。

2.4.2 對比實驗結(jié)果與分析

實驗表明，本文提出的基于抽象語法樹的代碼圖提取與殘差門控圖卷積網(wǎng)絡(luò)模型結(jié)合，在VDISC 數(shù)據(jù)集中漏洞檢測的綜合效能高于對比模型。

本文使用準確率和F1 得分來評估模型性能，并用其他模型最高評估指標與本文方法進行比較。表4 展示了所有實驗結(jié)果。從表中可以看到，檢測漏洞類型CWE-119 時，本文方法準確率為76.00%，F(xiàn)1 得分為76.60%，準確率比文獻[34]方法高了3.42 個百分點，F(xiàn)1得分比GGNN 方法高了5.67 個百分點。檢測漏洞類型CWE-120 時，本文方法準確率為69.91%，F(xiàn)1 得分為70.94%，準確率比GGNN方法高了4.42個百分點，F(xiàn)1得分比文獻[34]方法高了3.71 個百分點。檢測漏洞類型CWE-469 時，本文方法準確率為75.42%，F(xiàn)1 得分為74.95%，準確率比GGNN方法高了2.27個百分點，F(xiàn)1得分比文獻[34]方法高了4.03 個百分點。檢測漏洞類型CWE-476 時，本文方法準確率為60.38%，F(xiàn)1 得分為64.86%，準確率比文獻[34]方法高了3.06 個百分點，F(xiàn)1得分比GGNN 方法高了2.03 個百分點。檢測漏洞類型CWE-OTHERS 時，本文方法準確率為67.01%，F(xiàn)1 得分為64.82%，準確率比文獻[34]方法高了1.26 個百分點，F(xiàn)1 得分比文獻[34]方法高了2.43 個百分點。綜合評估指標來看本文方法優(yōu)于基線方法。

表4 基線實驗結(jié)果Table 4 Baseline experimental results單位：%

ROC（receiver operating characteristic）曲線是根據(jù)一系列不同的二分類方式，以真陽性率為縱坐標，假陽性率為橫坐標繪制的曲線。ROC曲線圖是衡量分類結(jié)果好壞的一種評判方法。ROC曲線與坐標軸圍成的面積為AUC值，取值范圍為0.5～1.0之間，對應(yīng)AUC越大，則說明分類器效果越好。圖4～圖8 是本文方法與基線方法的ROC曲線圖。

圖4 CWE-119 ROC曲線Fig.4 CWE-119 ROC curves

圖5 CWE-120 ROC曲線Fig.5 CWE-120 ROC curves

圖6 CWE-469 ROC曲線Fig.6 CWE-469 ROC curves

圖7 CWE-476 ROC曲線Fig.7 CWE-476 ROC curves

圖8 CWE-OTHERS ROC曲線Fig.8 CWE-OTHERS ROC curves

如圖4～圖8所示，在5種類型的漏洞檢測中本文方法的AUC值均高于基線方法，由此可見，本文的代碼圖表示方法和圖學(xué)習(xí)模型在漏洞檢測上更加有效。

本文通過抽象語法樹對代碼進行編碼，使用殘差門控圖卷積網(wǎng)絡(luò)學(xué)習(xí)代碼的漏洞特征，結(jié)合多層感知機來檢測漏洞，提高了自動檢測代碼漏洞的能力。

3 結(jié)束語

本文提出了一種基于殘差門控圖卷積神經(jīng)網(wǎng)絡(luò)的源代碼漏洞檢測方法。首先，構(gòu)建源代碼的抽象語法樹，利用抽象語法樹的層次關(guān)系來構(gòu)建代碼圖，并且結(jié)合抽象語法樹節(jié)點的出度作為圖節(jié)點實值向量的初始值。其次，把第一步得到的源代碼圖表示，作為殘差門控圖卷積神經(jīng)網(wǎng)絡(luò)的輸入，利用該模型來進行信息傳播和聚合，提取源代碼的特征，并通過池化來聚合圖的全局信息。最后，利用MLP 來對得到的特征向量進行特征整合，從而實現(xiàn)源代碼漏洞檢測。實驗表明，在VDISC數(shù)據(jù)集中，本文方法在5 種漏洞類型檢測中，F(xiàn)1 值均高于基線方法，有效提高了準確率以及自動化漏洞檢測能力，進一步降低了誤報率與漏報率，表明了本文方法是有效的。

本文方法只能檢測出函數(shù)中是否含有漏洞，無法檢測出漏洞所在的行，檢測的粒度不夠細，因此未來的工作將在更細粒度的漏洞檢測上展開。