基于多關系結構圖神經網絡的代碼漏洞檢測*

潘禺涵，舒遠仲，洪 晟，羅 斌，聶云峰

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

0 引言

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

為了應對上述的問題，本文提出了一種基于多關系結構圖神經網絡的漏洞檢測方法。采用多關系結構圖進行代碼的圖形表示，獲取全面的程序結構信息。然后結合雙向圖神經網絡以及關系結構圖注意力機制進行表示學習得到最終的代碼圖全局特征向量，并使用softmax分類器進行分類。

1 相關工作

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

圖1 漏洞代碼示例

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

2 圖神經網絡模型設計

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

圖2 模型框架圖

2.1 代碼的圖形表示

(1)關系結構圖劃分

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

圖3 CVE-2019-19075代碼片段

①Jump連接存在控制依賴關系變量節點的邊；

②RwFrom連接含有交互的操作數的邊；

③Dfuses連接存在數據依賴關系節點的邊；

④Next按操作順序連接AST上的葉節點的邊。

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

圖4 代碼屬性圖

算法1：關系結構圖的分割輸入：源代碼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)雙向邊連接

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

2.2 構造圖節點初始特征向量

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

2.3 特征提取

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

獲得的初始節點特征向量只孤立的考慮了每個節點，缺少了整個圖形結構的信息，因此使用GGNN通過消息傳遞和領域聚合機制來學習聚合信息生成全局圖表示。首先通過消息傳遞機制來聚合節點v的鄰域嵌入消息。然后，應用GRU[9]單元來聚合和更新關系結構圖中節點的狀態。模型使用前向傳播來更新圖頂點v的狀態，以獲得新狀態，傳 播公式如下：

(2)關系結構圖注意力機制

經過多輪消息傳遞與鄰域聚合后，各個節點狀態將包含足夠多的結構圖的信息，在獲得最后一輪迭代生成的節點嵌入后，使用最大池化層對節點進行集成，以獲得關系結構圖的全局向量表示：

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

其中t表示與關系結構圖相關的類型，at和分別代表注意力得分和對應關系圖的特征向量，SCPG是組合圖CPG的向量表示，Wa是要學習的權重矩陣。然后，通過softmax函數獲得關系圖的注意權重。最后，關系結構圖特征向量通過加權求和進行線性組合，加權求和的計算如下：

其中Attentiont是對應關系結構圖的注意力權重,SSG是關系結構圖加權求和后的特征向量。

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

最后將最終的圖全局特征向量SEG輸入全連接神經網絡通過softmax函數分類輸出最終的預測結果，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 實驗設計與結果分析

3.1 數據集

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

表1 數據集信息

3.2 實驗設置

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

采用分批次訓練的方式，使用Adam優化器和SGD以及交叉熵損失函數訓練模型，學習率為0.001，batch size設置為128。當損失小于0.005或達到最大100個epoch時，訓練 終止。dropout設置為0.2以防止過擬合。

3.3 實驗與結果分析

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

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

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

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

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

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

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

表2 漏洞檢測方法

表3 實驗結果

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

此外，Flawfinder這類靜態分析方法的檢測效果明顯低于使用深度學習網絡模型的方法。這是由于靜態分析工具的局限性，只能依靠專家定義的規則檢測特定類型的漏洞，從而導致準確率和召回率很低。

為了證實關系結構圖注意力機制的有效性，使用不同的readout方法來進行對比實驗。以常用的CONCAT方式作為基準進行比較。如表4所示，使用關系結構圖注意力機制后，模型可以更加關注代碼的結構和語義特征，比平均地集成所有節點信息的方式更具有優勢，實驗結果表明關系結構圖注意力機制提高了模型的檢測性能。

表4 不同readout方法的對比實驗結果

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

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

圖5 特征向量的可視化

4 結論

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