基于多關(guān)系結(jié)構(gòu)圖神經(jīng)網(wǎng)絡(luò)的代碼漏洞檢測*

潘禺涵，舒遠(yuǎn)仲，洪 晟，羅 斌，聶云峰

(1.南昌航空大學(xué) 信息工程學(xué)院，江西 南昌 330000；2.北京航空航天大學(xué) 網(wǎng)絡(luò)安全空間學(xué)院，北京 100191)

0 引言

軟件漏洞是許多系統(tǒng)攻擊[1]和數(shù)據(jù)泄露事件[2]的原因。在軟件產(chǎn)品開發(fā)中，源代碼靜態(tài)分析技術(shù)[3]被廣泛用于檢測漏洞。傳統(tǒng)的檢測方法主要通過一些由人類專家定義的度量標(biāo)準(zhǔn)所實(shí)現(xiàn)。這些方法目前取得的成果非常有限，因?yàn)樗鼈儫o法避免人類專家在特征提取方面的繁重工作[4]，而且用手工制作的特征來覆蓋所有漏洞是不切實(shí)際的。深度學(xué)習(xí)由于具有處理大量軟件代碼和漏洞數(shù)據(jù)的強(qiáng)大能力，被引入到代碼漏洞檢測領(lǐng)域。然而，現(xiàn)有基于深度學(xué)習(xí)的方法中不同形式的代碼表示方式只能保留部分語法或語義信息，這樣就不能覆蓋到每種漏洞并會限制模型檢測的效果。其次這些方法所使用的網(wǎng)絡(luò)模型仍存在一些局限性：例如，卷積神經(jīng)網(wǎng)絡(luò)(Convolutional Neural Network，CNN)和循環(huán)神經(jīng)網(wǎng)絡(luò)(Recurrent Neural Network，RNN)是基于序列的模型，不能處理代碼的非序列特征，只能捕獲源代碼文本的淺表結(jié)構(gòu)，無法利用程序結(jié)構(gòu)豐富且定義良好的語義。

為了應(yīng)對上述的問題，本文提出了一種基于多關(guān)系結(jié)構(gòu)圖神經(jīng)網(wǎng)絡(luò)的漏洞檢測方法。采用多關(guān)系結(jié)構(gòu)圖進(jìn)行代碼的圖形表示，獲取全面的程序結(jié)構(gòu)信息。然后結(jié)合雙向圖神經(jīng)網(wǎng)絡(luò)以及關(guān)系結(jié)構(gòu)圖注意力機(jī)制進(jìn)行表示學(xué)習(xí)得到最終的代碼圖全局特征向量，并使用softmax分類器進(jìn)行分類。

1 相關(guān)工作

深度學(xué)習(xí)中用于漏洞檢測的模型大致分為兩類：基于token的模型和基于圖的模型。在基于token的模型中，代碼被認(rèn)為是token序列。例如，Li等人提出了一種基于雙向長短期記憶網(wǎng)絡(luò)(Bidirectional Long-Short-Term Memory，BiLSTM)的 模 型SySeVR[5]來表示漏洞相關(guān)的語法和語義信息，其中使用程序切片技術(shù)生成更小的代碼段(即一些可能不連續(xù)但語義性和漏洞相關(guān)的語句)，使其適用于深度學(xué)習(xí)網(wǎng)絡(luò)。雖然SySeVR提高了深度學(xué)習(xí)模型在漏洞檢測方面的性能，但它仍然存在問題：SySeVR采用序列神經(jīng)網(wǎng)絡(luò)(如Long Short Term Memory，LSTM)對代碼段進(jìn)行純文本編碼，可能會丟失代碼段中語句之間復(fù)雜的結(jié)構(gòu)信息和依賴關(guān)系，無法有效地學(xué)習(xí)代碼的非序列特征。如圖1所示代碼，第2行代碼與第10行代碼雖然都在token序列中，但是由于相隔較遠(yuǎn)，它們之間的依賴關(guān)系已經(jīng)丟失,這將會導(dǎo)致模型的高漏報(bào)率和高誤報(bào)率。

圖1 漏洞代碼示例

在基于圖的漏洞檢測模型中，為了保留源代碼復(fù)雜的結(jié)構(gòu)和語義信息，并有效地學(xué)習(xí)這些源代碼的非序列特征，Zhou等人[6]提出了一種基于門控圖神經(jīng)網(wǎng)絡(luò)[7](Gated Graph Sequence Neural Networks，GGNN)的漏洞檢測模型Devign，該模型首先將源代碼的語法和依賴信息集成到一個組合圖表示中，然后借助GGNN從組合圖中提取與漏洞相關(guān)的特征。然而，這會限制模型的表示能力。這種方法是通過在抽象語法樹(Abstract Syntax Tree，AST)上添加語義邊緣來構(gòu)建圖的。然而，程序表示仍然高度依賴語法，語義相對不足。此外，來自語法、數(shù)據(jù)流和控制流等不同方面的信息常?；旌显谝黄?，形成一個單一的視圖，使得很難單獨(dú)提取關(guān)鍵信息。而在本文中，采取了多關(guān)系結(jié)構(gòu)圖的方法進(jìn)行代碼表示，該方法將基于不同的語義關(guān)系把源代碼劃分為5種圖形結(jié)構(gòu)，可以從多個方面和層次來表示源代碼的結(jié)構(gòu)信息,并且更加強(qiáng)調(diào)了語義信息。

2 圖神經(jīng)網(wǎng)絡(luò)模型設(shè)計(jì)

本文的目標(biāo)是設(shè)計(jì)一個用于自動化漏洞檢測的網(wǎng)絡(luò)模型，圖2顯示了模型的框架。它包括四個部分：第一部分，將代碼轉(zhuǎn)換為多個關(guān)系結(jié)構(gòu)圖，以保留程序代碼中的語法和語義關(guān)系(例如，數(shù)據(jù)依賴和控制依賴)；然后第二部分，使用word2vec模型[8]將圖形節(jié)點(diǎn)編碼成向量，作為網(wǎng)絡(luò)模型的輸入；第三部分，對雙向圖神經(jīng)網(wǎng)絡(luò)進(jìn)行訓(xùn)練，得到每個關(guān)系結(jié)構(gòu)圖的向量表示；第四部分，通過關(guān)系結(jié)構(gòu)圖注意力機(jī)制獲得程序代碼的最終全局嵌入向量，接著使用softmax分類器分類獲得最終結(jié)果：代碼是否存在漏洞。

圖2 模型框架圖

2.1 代碼的圖形表示

(1)關(guān)系結(jié)構(gòu)圖劃分

CVE-2019-190753是一種內(nèi)存泄漏漏洞，ca8210_get_platform_data函數(shù)執(zhí)行失敗可能會導(dǎo)致pdata無法釋放內(nèi)存空間，這使得攻擊者通過觸發(fā)ca8210_get_platform_data函數(shù)報(bào)錯來導(dǎo)致拒絕服務(wù)攻擊。如圖3所示，假如priv-?spi-?dev.platform_data=pdata；語句在if語句之前執(zhí)行就不會產(chǎn)生危害性。這種類型的漏洞不能通過單一的組合圖來進(jìn)行檢測，在該漏洞代碼特征中既包含了控制依賴也有數(shù)據(jù)依賴以及操作數(shù)讀寫關(guān)系，需要結(jié)合不同類型的關(guān)系圖對代碼語義進(jìn)行細(xì)致分析?，F(xiàn)實(shí)情況下，這種類似的漏洞代碼結(jié)構(gòu)還有很多。所以為了捕獲更加全面準(zhǔn)確的語義信息，采用多關(guān)系結(jié)構(gòu)圖來表示程序代碼，分別基于以下4種節(jié)點(diǎn)之間的邊的類型來對源代碼進(jìn)行結(jié)構(gòu)圖劃分：

圖3 CVE-2019-19075代碼片段

①Jump連接存在控制依賴關(guān)系變量節(jié)點(diǎn)的邊；

②RwFrom連接含有交互的操作數(shù)的邊；

③Dfuses連接存在數(shù)據(jù)依賴關(guān)系節(jié)點(diǎn)的邊；

④Next按操作順序連接AST上的葉節(jié)點(diǎn)的邊。

具體來說，首先通過開源工具Joern對源代碼進(jìn)行解析生成代碼屬性圖(CPG)，如圖4所示，然后基于語義關(guān)系將代碼屬性圖轉(zhuǎn)化分割為多個分離的有向關(guān)系結(jié)構(gòu)圖，接著將這些關(guān)系結(jié)構(gòu)圖與CPG同時作為最終的程序代碼圖形表示。

圖4 代碼屬性圖

算法1：關(guān)系結(jié)構(gòu)圖的分割輸入：源代碼P輸出：CPG、GJump、GRwFrom、GDfuses、GNext Generate a CPG for P；for edge e in CPG do if e is Jump/RwFrom edge then Add node src(e),end(e)and edge e to GJump/GRwFrom end if else if e is Dfuses/Next edge then Add node src(e),end(e)and edge e to GDfuses/GNext end if end for

(2)雙向邊連接

由于源代碼比自然語言更具邏輯性和結(jié)構(gòu)性，含有漏洞的代碼片段通常與其上下文相關(guān)。選擇用于學(xué)習(xí)的神經(jīng)網(wǎng)絡(luò)模型應(yīng)該能夠處理向前和向后的序列，以獲得更精細(xì)的表示，并實(shí)現(xiàn)更好的魯棒性。本文通過添加后向邊來構(gòu)造一個雙向圖神經(jīng)網(wǎng)絡(luò)。通過轉(zhuǎn)置鄰接矩陣，對于一對連接邊為ft的節(jié)點(diǎn)[u，v]，可以同時考慮它的前向邊[u，v]以及后向邊[v，u]。因此，節(jié)點(diǎn)信息可以向前和向后傳播，這有助于節(jié)點(diǎn)之間的信息傳遞。

2.2 構(gòu)造圖節(jié)點(diǎn)初始特征向量

為了生成代碼的初始特征向量，使用Word2vec詞向量模型[7]對關(guān)系結(jié)構(gòu)圖中每個節(jié)點(diǎn)的內(nèi)容信息進(jìn)行編碼。首先對每個語句執(zhí)行詞法分析，并通過預(yù)先訓(xùn)練的word2vec模型創(chuàng)建詞法標(biāo)記向量。代碼中的每個符號和word都會被映射到一個100維的詞向量之中。其中經(jīng)常出現(xiàn)在一起的單詞會被映射到向量空間中接近的整數(shù)值，這樣做可以捕獲代碼結(jié)構(gòu)的大部分句法關(guān)系。例如，模型能夠?qū)W習(xí)到if語句必須在else語句之前。為了進(jìn)一步反映語料庫中word之間的抽象語義關(guān)系，還需要考慮每個節(jié)點(diǎn)的類型信息。例如，與API函數(shù)調(diào)用相關(guān)的Call-Expression類型的語句節(jié)點(diǎn)可能包含漏洞相關(guān)的更多信息。對于節(jié)點(diǎn)類型信息使用one-hot編碼，生成60維的類型特征向量。最后將內(nèi)容特征向量與類型特征向量線性聯(lián)結(jié)在一起，以獲得節(jié)點(diǎn)的初始嵌入特征向量。

2.3 特征提取

(1)消息傳遞與鄰域聚合

獲得的初始節(jié)點(diǎn)特征向量只孤立的考慮了每個節(jié)點(diǎn)，缺少了整個圖形結(jié)構(gòu)的信息，因此使用GGNN通過消息傳遞和領(lǐng)域聚合機(jī)制來學(xué)習(xí)聚合信息生成全局圖表示。首先通過消息傳遞機(jī)制來聚合節(jié)點(diǎn)v的鄰域嵌入消息。然后，應(yīng)用GRU[9]單元來聚合和更新關(guān)系結(jié)構(gòu)圖中節(jié)點(diǎn)的狀態(tài)。模型使用前向傳播來更新圖頂點(diǎn)v的狀態(tài)，以獲得新狀態(tài)，傳 播公式如下：

(2)關(guān)系結(jié)構(gòu)圖注意力機(jī)制

經(jīng)過多輪消息傳遞與鄰域聚合后，各個節(jié)點(diǎn)狀態(tài)將包含足夠多的結(jié)構(gòu)圖的信息，在獲得最后一輪迭代生成的節(jié)點(diǎn)嵌入后，使用最大池化層對節(jié)點(diǎn)進(jìn)行集成，以獲得關(guān)系結(jié)構(gòu)圖的全局向量表示：

由于在代碼結(jié)構(gòu)里的不同特征的重要性因漏洞而異，例如，在與數(shù)組的使用相關(guān)漏洞中數(shù)據(jù)依賴可能比控制依賴更重要，因此在此引入注意力機(jī)制來生成最終的代碼全局表示向量。這樣可以關(guān)注關(guān)鍵的漏洞特征提高分類效率，計(jì)算公式如下：

其中t表示與關(guān)系結(jié)構(gòu)圖相關(guān)的類型，at和分別代表注意力得分和對應(yīng)關(guān)系圖的特征向量，SCPG是組合圖CPG的向量表示，Wa是要學(xué)習(xí)的權(quán)重矩陣。然后，通過softmax函數(shù)獲得關(guān)系圖的注意權(quán)重。最后，關(guān)系結(jié)構(gòu)圖特征向量通過加權(quán)求和進(jìn)行線性組合，加權(quán)求和的計(jì)算如下：

其中Attentiont是對應(yīng)關(guān)系結(jié)構(gòu)圖的注意力權(quán)重,SSG是關(guān)系結(jié)構(gòu)圖加權(quán)求和后的特征向量。

其中SEG是最終的圖全局特征向量，⊕表示連接操作。

最后將最終的圖全局特征向量SEG輸入全連接神經(jīng)網(wǎng)絡(luò)通過softmax函數(shù)分類輸出最終的預(yù)測結(jié)果，1為有漏洞，0為無漏洞。

綜上所述，本文提出的方法如算法2所示。

算法2：本文方法輸入：源代碼P輸出：代碼的圖嵌入表示Generate a CPG for P Construct GJump，GRwFrom，GDfuses，GNext based on CPG for Gi∈{GJump，GNext，GRwFrom，GDfuses，CPG}do for token in tokenized(v.code)do embeddings.append(word2vec(token))end for Dv=average(embeddings)for each node v in Vi do Tv=onehot(v.type)end for hv=concat(Tv，Dv)xu 0.append(hv)Iteratively update node state with Bi-GGNN for k steps(Eq.2)；Compute the graph representation St as Eq.3；end for Obtain the program representation SEG as Eq.4-6；Feed SEG to downstream tasks

3 實(shí)驗(yàn)設(shè)計(jì)與結(jié)果分析

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

本文選用Lu等人提供的CodeXGLUE[10]的標(biāo)準(zhǔn)真實(shí)數(shù)據(jù)集進(jìn)行漏洞檢測實(shí)驗(yàn)。其中包括27 318個手動標(biāo)記的漏洞樣本集以及非漏洞樣本集，這些樣本是從兩個大型且功能多樣化的C、C++編程語言開源項(xiàng)目QEMU和FFmpeG中提取的。在非漏洞修復(fù)提交中提取修改前的源碼函數(shù)作為非漏洞樣本集，從漏洞修復(fù)提交中提取修改前的源碼函數(shù)作為漏洞樣本集，如表1所示。

表1 數(shù)據(jù)集信息

3.2 實(shí)驗(yàn)設(shè)置

現(xiàn)實(shí)情況下含有漏洞的程序代碼往往在整個項(xiàng)目中占據(jù)非常小的比例，因此在真實(shí)代碼數(shù)據(jù)集中存在著嚴(yán)重的類別不平衡現(xiàn)象。為了處理漏洞樣本和非漏洞樣本數(shù)量的不平衡，采用SMOTE[11](Synthetic Minority Oversampling Technique)算法對多數(shù)類進(jìn)行子采樣，同時對少數(shù)類進(jìn)行超級采樣(通過創(chuàng)建合成樣本)，直到所有類具有相同的比例。SMOTE已被證實(shí)在許多數(shù)據(jù)集不平衡的領(lǐng)域有效[12]。

采用分批次訓(xùn)練的方式，使用Adam優(yōu)化器和SGD以及交叉熵?fù)p失函數(shù)訓(xùn)練模型，學(xué)習(xí)率為0.001，batch size設(shè)置為128。當(dāng)損失小于0.005或達(dá)到最大100個epoch時，訓(xùn)練 終止。dropout設(shè)置為0.2以防止過擬合。

3.3 實(shí)驗(yàn)與結(jié)果分析

采用準(zhǔn)確率(Accuracy，Acc)、精度(Precision，P)、召回率(Recall，R)和F1值(F1-measure)作為最終模型在測試集上效果的評估指標(biāo)評價(jià)。其中含有漏洞的樣本表示為正類，不含漏洞的樣本表示為負(fù)類。TP(True Positive)是被模型預(yù)測為正類的正樣本數(shù)量，F(xiàn)P(False Positive)是被模型預(yù)測為負(fù)類的正樣本數(shù)量，TN(True Negative)是被模型預(yù)測為負(fù)類的負(fù)樣本數(shù)量，F(xiàn)N(False Negative)是被模型預(yù)測為正類的負(fù)樣本數(shù)量。

(1)Acc表示正確分類的概率，計(jì)算公式如下：

(2)P表示被判定為漏洞的函數(shù)中真正的漏洞的比例，計(jì)算式如下：

(3)R則表示樣本中有多少漏洞被正確的預(yù)測了，計(jì)算公式如下：

(4)F1值是精度和召回率的加權(quán)平均值，計(jì)算公式如下：

上述四個指標(biāo)的度量值越高，表示模型越理想。

表2總結(jié)了目前幾種具有代表性的漏洞檢測方法。其中包含經(jīng)典的靜態(tài)漏洞檢測工具Flawfinder[13]，兩種基于切片的方法VulDeePecker[14]和SySeVR，以及兩種基于圖形的方法Devign和VdoRGCN[15]。為了證實(shí)本文方法的有效性，使用上述的方法作為基準(zhǔn)進(jìn)行實(shí)驗(yàn)對比，結(jié)果如表3所示。

表2 漏洞檢測方法

表3 實(shí)驗(yàn)結(jié)果

如表3中結(jié)果所示，與VulDeePecker相比，本文方法在準(zhǔn)確率和F1分?jǐn)?shù)上分別提高了5.5%和14.9%，與SySeVR相比在精度，召回率以及F1分?jǐn)?shù)方面都有較大優(yōu)勢。SySeVR和VulDeePecker這類基于token的方法中使用了多個RNN(如LSTM、GRU)來訓(xùn)練檢測模型。由于RNN處理數(shù)據(jù)的不連續(xù)特征的性能較差，一些重要的特征可能會被忽略或掩蓋。相比之下，雙向圖神經(jīng)網(wǎng)絡(luò)在學(xué)習(xí)基于圖的數(shù)據(jù)的特征和通過其雙向邊緣容納上下文信息方面更加高效。與基于圖形的方法相比，本文方法也是更有優(yōu)勢。Devign和VdoRGCN都只使用一個組合圖來包含代碼不同的信息，這會更加側(cè)重語法而忽略程序語義的重要性，限制模型的表示能力。而多關(guān)系結(jié)構(gòu)圖的使用涵蓋了多方面的程序表示提高了分類性能。

此外，F(xiàn)lawfinder這類靜態(tài)分析方法的檢測效果明顯低于使用深度學(xué)習(xí)網(wǎng)絡(luò)模型的方法。這是由于靜態(tài)分析工具的局限性，只能依靠專家定義的規(guī)則檢測特定類型的漏洞，從而導(dǎo)致準(zhǔn)確率和召回率很低。

為了證實(shí)關(guān)系結(jié)構(gòu)圖注意力機(jī)制的有效性，使用不同的readout方法來進(jìn)行對比實(shí)驗(yàn)。以常用的CONCAT方式作為基準(zhǔn)進(jìn)行比較。如表4所示，使用關(guān)系結(jié)構(gòu)圖注意力機(jī)制后，模型可以更加關(guān)注代碼的結(jié)構(gòu)和語義特征，比平均地集成所有節(jié)點(diǎn)信息的方式更具有優(yōu)勢，實(shí)驗(yàn)結(jié)果表明關(guān)系結(jié)構(gòu)圖注意力機(jī)制提高了模型的檢測性能。

表4 不同readout方法的對比實(shí)驗(yàn)結(jié)果

現(xiàn)有的漏洞檢測方法都是通過網(wǎng)絡(luò)模型將源代碼轉(zhuǎn)換為數(shù)值特征向量。所以漏洞檢測的效率取決于正類樣本與負(fù)類樣本特征向量的可分離性，類別可分離性越大，模型就越容易區(qū)分代碼中是否含有漏洞。為了直觀地顯示本文方法可以獲取更好的程序代碼嵌入表示，使用UMAP算法[16]來可視化實(shí)驗(yàn)中模型生成的表示向量，如圖5所示。

從圖5可以看出，所有的方法在兩個類別之間的特征向量空間上都有很大程度的重疊，但是本文的方法明顯可以給出更清晰的兩種類別的決策邊界，這意味著本文方法比其他方法具有更高的表示能力，在漏洞檢測方面更加有效。

圖5 特征向量的可視化

4 結(jié)論

本文提出了一種基于多關(guān)系結(jié)構(gòu)圖和雙向圖神經(jīng)網(wǎng)絡(luò)的漏洞檢測方法。通過考慮多個代碼圖形表示來更加全面地保留語法語義信息，接著采用雙向圖神經(jīng)網(wǎng)絡(luò)對編碼后的圖進(jìn)行表示學(xué)習(xí)，然后利用關(guān)系結(jié)構(gòu)圖注意力機(jī)制關(guān)注關(guān)鍵的漏洞特征，最后通過全連接神經(jīng)網(wǎng)絡(luò)進(jìn)行分類預(yù)測。在真實(shí)數(shù)據(jù)集上進(jìn)行的對比實(shí)驗(yàn)結(jié)果表明，本文方法具有更高準(zhǔn)確度和召回率。未來的工作將研究不同編程語言的通用漏洞檢測模型。