基于圖神經(jīng)網(wǎng)絡(luò)的代碼漏洞檢測(cè)方法

陳皓，易平

基于圖神經(jīng)網(wǎng)絡(luò)的代碼漏洞檢測(cè)方法

陳皓，易平

（上海交通大學(xué)網(wǎng)絡(luò)空間安全學(xué)院，上海 200240）

使用神經(jīng)網(wǎng)絡(luò)進(jìn)行漏洞檢測(cè)的方案大多基于傳統(tǒng)自然語(yǔ)言處理的思路，將源代碼當(dāng)作序列樣本處理，忽視了代碼中所具有的結(jié)構(gòu)性特征，從而遺漏了可能存在的漏洞。提出了一種基于圖神經(jīng)網(wǎng)絡(luò)的代碼漏洞檢測(cè)方法，通過(guò)中間語(yǔ)言的控制流圖特征，實(shí)現(xiàn)了函數(shù)級(jí)別的智能化代碼漏洞檢測(cè)。首先，將源代碼編譯為中間表示，進(jìn)而提取其包含結(jié)構(gòu)信息的控制流圖，同時(shí)使用詞向量嵌入算法初始化基本塊向量提取代碼語(yǔ)義信息；然后，完成拼接生成圖結(jié)構(gòu)樣本數(shù)據(jù)，使用多層圖神經(jīng)網(wǎng)絡(luò)對(duì)圖結(jié)構(gòu)數(shù)據(jù)特征進(jìn)行模型訓(xùn)練和測(cè)試。采用開(kāi)源漏洞樣本數(shù)據(jù)集生成測(cè)試數(shù)據(jù)對(duì)所提方法進(jìn)行了評(píng)估，結(jié)果顯示該方法有效提高了漏洞檢測(cè)能力。

漏洞檢測(cè)；圖神經(jīng)網(wǎng)絡(luò)；控制流圖；中間表示

1 引言

隨著各種智能設(shè)備的持續(xù)突破和互聯(lián)網(wǎng)的持續(xù)發(fā)展與廣泛應(yīng)用，網(wǎng)絡(luò)空間促進(jìn)了經(jīng)濟(jì)社會(huì)的繁榮進(jìn)步，也正在全面改變?nèi)藗兊纳a(chǎn)生活方式，深刻影響著人類(lèi)社會(huì)歷史的發(fā)展進(jìn)程，但同時(shí)帶來(lái)了新的安全風(fēng)險(xiǎn)和挑戰(zhàn)。2016年，國(guó)家互聯(lián)網(wǎng)信息辦公室發(fā)布的《國(guó)家網(wǎng)絡(luò)空間安全戰(zhàn)略》[1]指出，網(wǎng)絡(luò)空間安全事關(guān)人類(lèi)共同利益，事關(guān)世界和平與發(fā)展，事關(guān)各國(guó)國(guó)家安全，要夯實(shí)網(wǎng)絡(luò)安全基礎(chǔ)?；ヂ?lián)網(wǎng)行業(yè)競(jìng)爭(zhēng)激烈，軟件開(kāi)發(fā)周期越來(lái)越短，安全問(wèn)題日益突出。隨著軟件數(shù)量的大規(guī)模增長(zhǎng)以及復(fù)雜性的增強(qiáng)，傳統(tǒng)方案面臨著嚴(yán)重的技術(shù)瓶頸，這給軟件安全漏洞的研究帶來(lái)嚴(yán)峻的挑戰(zhàn)，已經(jīng)無(wú)法滿(mǎn)足網(wǎng)絡(luò)空間安全的防護(hù)需要。

面對(duì)挑戰(zhàn)，基于機(jī)器學(xué)習(xí)和人工智能的脆弱性分析方法被應(yīng)用于脆弱性分析領(lǐng)域，開(kāi)始成為研究的熱點(diǎn)方向。近年來(lái)，深度學(xué)習(xí)在處理分類(lèi)問(wèn)題中顯示出了巨大的優(yōu)勢(shì)，已經(jīng)在圖像識(shí)別、自然語(yǔ)言處理等領(lǐng)域取得了突破性進(jìn)展。與計(jì)算機(jī)視覺(jué)和自然語(yǔ)言處理等深度學(xué)習(xí)領(lǐng)域相比，源代碼漏洞檢測(cè)首先要解決的關(guān)鍵問(wèn)題是樣本的預(yù)處理與合適的深度神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計(jì)。與圖片樣本和文本樣本不同，源代碼由于具有豐富的代碼結(jié)構(gòu)信息和語(yǔ)義信息，可以看作一種非歐幾里得數(shù)據(jù)?，F(xiàn)有的研究方案大多把源代碼轉(zhuǎn)化為文本序列，使用自然語(yǔ)言處理中常用的循環(huán)神經(jīng)網(wǎng)絡(luò)進(jìn)行自動(dòng)化特征提取，但這忽略了代碼中由于跳轉(zhuǎn)、循環(huán)、判斷形成的控制流特征，往往會(huì)有較大的信息損失。如何提取源代碼樣本的特征信息，把源代碼樣本轉(zhuǎn)換成適合后續(xù)神經(jīng)網(wǎng)絡(luò)訓(xùn)練的數(shù)據(jù)形式并最大化保留樣本信息是關(guān)鍵問(wèn)題之一。

本文通過(guò)代碼中間語(yǔ)言的控制流圖提取代碼樣本的結(jié)構(gòu)特征，使用控制流圖基本塊中的代碼序列提取代碼樣本的語(yǔ)義特征，將代碼樣本轉(zhuǎn)化為圖結(jié)構(gòu)數(shù)據(jù)，以解決序列樣本丟失代碼結(jié)構(gòu)特征信息的問(wèn)題。另一個(gè)關(guān)鍵問(wèn)題是針對(duì)具有豐富語(yǔ)法特征和語(yǔ)義特征的源代碼樣本，選用哪種合適的深度神經(jīng)網(wǎng)絡(luò)模型。圖神經(jīng)網(wǎng)絡(luò)由于其設(shè)計(jì)原理更適合處理圖結(jié)構(gòu)數(shù)據(jù)，在圖級(jí)別的分類(lèi)任務(wù)上也取得了非常好的效果[2]。使用基于中間語(yǔ)言的控制流圖信息對(duì)源代碼進(jìn)行漏洞檢測(cè)的過(guò)程可以抽象為在圖結(jié)構(gòu)數(shù)據(jù)樣本上進(jìn)行圖級(jí)別的分類(lèi)任務(wù)，因此，圖神經(jīng)網(wǎng)絡(luò)相較于其他模型更有優(yōu)勢(shì)。

本文的主要貢獻(xiàn)如下。

1) 設(shè)計(jì)了一種基于中間語(yǔ)言控制流圖的代碼特征提取方案。通過(guò)提取源代碼的控制流圖結(jié)構(gòu)信息，同時(shí)使用基本塊的中間語(yǔ)言提取代碼語(yǔ)義信息，嵌入圖節(jié)點(diǎn)中，將源代碼轉(zhuǎn)化為圖結(jié)構(gòu)數(shù)據(jù)樣本。

2) 使用圖神經(jīng)網(wǎng)絡(luò)模型處理代碼圖結(jié)構(gòu)樣本?；趫D結(jié)構(gòu)的數(shù)據(jù)樣本，設(shè)計(jì)了一種基于圖卷積的深度神經(jīng)網(wǎng)絡(luò)模型，用來(lái)對(duì)樣本數(shù)據(jù)進(jìn)行特征提取和樣本分類(lèi)。

3) 基于設(shè)計(jì)方案設(shè)計(jì)了實(shí)驗(yàn)方案并對(duì)結(jié)果進(jìn)行了分析。實(shí)驗(yàn)結(jié)果顯示，本文設(shè)計(jì)的方案在基準(zhǔn)數(shù)據(jù)集上取得了比傳統(tǒng)靜態(tài)掃描更高的準(zhǔn)確率。

2 相關(guān)工作

近年來(lái)，深度學(xué)習(xí)與軟件漏洞檢測(cè)的交叉研究逐漸成為軟件安全和人工智能領(lǐng)域共同的研究熱點(diǎn)。由于傳統(tǒng)軟件漏洞檢測(cè)有了較長(zhǎng)時(shí)間的研究，各種檢測(cè)方法較為成熟，其中的優(yōu)劣勢(shì)也較為明顯。人工智能的加入大多著眼于與現(xiàn)有的軟件檢測(cè)方案結(jié)合，以解決現(xiàn)有檢測(cè)方案的問(wèn)題。傳統(tǒng)的源代碼靜態(tài)檢測(cè)技術(shù)主要分為基于邏輯推理的漏洞檢測(cè)、基于中間表示的漏洞檢測(cè)和基于相似性的漏洞檢測(cè)?；谶壿嬐评淼穆┒礄z測(cè)將源代碼進(jìn)行形式化描述，然后利用數(shù)學(xué)推理證明等方法驗(yàn)證形式化描述的性質(zhì)，從而推斷是否含有某類(lèi)漏洞?；谥虚g表示的漏洞檢測(cè)則將源代碼轉(zhuǎn)換為易于分析檢測(cè)的中間表示，然后檢查其對(duì)應(yīng)的漏洞規(guī)則，具體的檢測(cè)方法有數(shù)據(jù)流分析、控制流分析、污點(diǎn)分析、靜態(tài)符號(hào)執(zhí)行等?；谙嗨菩缘穆┒礄z測(cè)主要檢測(cè)由于代碼復(fù)用等產(chǎn)生的漏洞。2010年，Pham等[3]提出并開(kāi)發(fā)了SecureSync工具，其將源代碼轉(zhuǎn)化為擴(kuò)展的抽象語(yǔ)法樹(shù)的中間表示，用于檢測(cè)代碼相似性進(jìn)而檢測(cè)漏洞。同樣進(jìn)行相似性檢測(cè)，也可以選取不同的中間表示，基于圖的中間表示綜合考慮程序的語(yǔ)法和語(yǔ)義特征，有很好的檢測(cè)能力。2014年，Yamaguchi等[4]提出了一種將經(jīng)典程序分析的概念（抽象語(yǔ)法樹(shù)、控制流圖和程序依賴(lài)圖）合并成為代碼屬性圖的圖結(jié)構(gòu)，進(jìn)而通過(guò)對(duì)常見(jiàn)漏洞的建模識(shí)別源代碼漏洞。

隨著深度學(xué)習(xí)、人工神經(jīng)網(wǎng)絡(luò)等算法的不斷發(fā)展進(jìn)步，各種新型算法逐步與傳統(tǒng)檢測(cè)算法先后結(jié)合，涌現(xiàn)出了一批具有交叉方案的研究成果。Lin等[5]提出了針對(duì)函數(shù)級(jí)別的跨項(xiàng)目場(chǎng)景下的漏洞檢測(cè)，使用抽象語(yǔ)法樹(shù)AST表示函數(shù)，通過(guò)雙向LSTM神經(jīng)網(wǎng)絡(luò)進(jìn)行特征提取和學(xué)習(xí)。Russell等[6]建立了包含C/C++開(kāi)源代碼的用于深度學(xué)習(xí)訓(xùn)練的函數(shù)級(jí)別漏洞數(shù)據(jù)集，并基于深度表示學(xué)習(xí)開(kāi)發(fā)了一個(gè)大規(guī)模漏洞檢測(cè)系統(tǒng)。在二進(jìn)制代碼檢測(cè)領(lǐng)域，Xu等[7]提出了一種基于神經(jīng)網(wǎng)絡(luò)的跨平臺(tái)的二進(jìn)制代碼相似性檢測(cè)模型，通過(guò)計(jì)算使用神經(jīng)網(wǎng)絡(luò)嵌入得到的函數(shù)向量距離來(lái)進(jìn)行相似性檢測(cè)。李珍等[8]提出了基于圖神經(jīng)網(wǎng)絡(luò)的圖嵌入算法，并在此基礎(chǔ)上提出一種圖匹配網(wǎng)絡(luò)模型來(lái)進(jìn)行代碼控制流圖的功能相似性檢測(cè)。Yu等[9]在文獻(xiàn)[7]的基礎(chǔ)上提出了一種結(jié)合圖神經(jīng)網(wǎng)絡(luò)和NLP算法提取二進(jìn)制程序控制流圖進(jìn)行相似性檢測(cè)的方法，一方面采用MPNN（message passing neural network）提取控制流圖語(yǔ)法結(jié)構(gòu)特征，另一方面采用BERT預(yù)訓(xùn)練框架提取代碼語(yǔ)義特征。Duan等[10]則提出了一種無(wú)監(jiān)督的代碼表示學(xué)習(xí)技術(shù)解決二進(jìn)制相似性分析的問(wèn)題，也是處理代碼語(yǔ)義信息和控制流信息生成基本塊嵌入。李珍等[11]研究了近年來(lái)使用深度學(xué)習(xí)方案進(jìn)行漏洞檢測(cè)的方案，并以基于源代碼相似性的漏洞檢測(cè)系統(tǒng)和面向源代碼的軟件漏洞智能檢測(cè)系統(tǒng)兩個(gè)具體方案為例詳細(xì)介紹了基于深度學(xué)習(xí)的漏洞檢測(cè)過(guò)程，指出了目前存在的數(shù)據(jù)樣本標(biāo)注、漏洞模式和模型選擇以及可解釋性方面存在的問(wèn)題。

圖神經(jīng)網(wǎng)絡(luò)（GNN）是一類(lèi)處理圖數(shù)據(jù)的神經(jīng)網(wǎng)絡(luò)的統(tǒng)稱(chēng)。圖是一種數(shù)據(jù)結(jié)構(gòu)，它對(duì)一組對(duì)象（節(jié)點(diǎn)）及其關(guān)系（邊）進(jìn)行建模。近年來(lái)，由于圖結(jié)構(gòu)的強(qiáng)大表現(xiàn)力，用機(jī)器學(xué)習(xí)方法和深度學(xué)習(xí)方法分析圖的研究越來(lái)越受到重視。圖神經(jīng)網(wǎng)絡(luò)是一類(lèi)基于深度學(xué)習(xí)的處理圖域信息的方法。由于其較好的性能和可解釋性，GNN 最近已成為一種廣泛應(yīng)用的圖分析方法。圖神經(jīng)網(wǎng)絡(luò)的概念最早由 Gori 等[12]提出，并由Scarselli等[13]進(jìn)一步闡明。這些早期的研究以迭代的方式通過(guò)循環(huán)神經(jīng)架構(gòu)傳播鄰近信息來(lái)學(xué)習(xí)目標(biāo)節(jié)點(diǎn)的表示，直到達(dá)到穩(wěn)定的固定點(diǎn)。Gilmer等[14]提出了一種MPNN架構(gòu)用于化學(xué)領(lǐng)域分子特征的預(yù)測(cè)，創(chuàng)造性地將圖信息傳遞的過(guò)程分為傳播、更新、讀出3個(gè)階段，Yu等[9]的研究基于MPNN在代碼領(lǐng)域取得了比較好的效果。

3 方案設(shè)計(jì)

本文的目標(biāo)是建立一種基于圖神經(jīng)網(wǎng)絡(luò)的源代碼漏洞檢測(cè)方案。方案的整體設(shè)計(jì)包括代碼預(yù)處理與中間語(yǔ)言編譯、控制流圖提取、代碼語(yǔ)義信息提取、圖神經(jīng)網(wǎng)絡(luò)模型設(shè)計(jì)等。方案的整體設(shè)計(jì)如圖1所示。

3.1 代碼預(yù)處理

在傳統(tǒng)的源代碼漏洞掃描方案中，由于不同類(lèi)型編程語(yǔ)言在語(yǔ)法上存在較大差異，往往會(huì)針對(duì)不同類(lèi)型語(yǔ)言開(kāi)發(fā)不同的規(guī)則或者工具。本文采用LLVM Intermediate Representation（IR）的中間語(yǔ)言方案將源代碼預(yù)編譯為中間表示，使檢測(cè)方法能夠兼容不同類(lèi)型的編程語(yǔ)言。

LLVM編譯架構(gòu)如圖2所示。LLVM編譯架構(gòu)遵循前后端分離的設(shè)計(jì)理念，編譯器前端可以把不同類(lèi)型的高級(jí)語(yǔ)言編譯成為L(zhǎng)LVM IR，經(jīng)過(guò)優(yōu)化之后再由不同系統(tǒng)架構(gòu)的編譯器后端將LLVM IR編譯成不同架構(gòu)的二進(jìn)制可執(zhí)行文件。

LLVM IR有兩種表示形式：一種是可讀的匯編語(yǔ)言形式，以.ll文件存儲(chǔ)；另一種是序列化之后的bitcode形式，以.bc文件存儲(chǔ)。

圖1 方案的整體設(shè)計(jì)

Figure 1 Overall design drawing

圖2 LLVM編譯架構(gòu)

Figure 2 LLVM compilation architecture

本文通過(guò)Clang編譯器將數(shù)據(jù)集代碼樣本編譯成為可讀形式的中間語(yǔ)言代碼樣本，作為下一步特征提取的基礎(chǔ)。圖3所示為數(shù)據(jù)集中的一個(gè)樣本示例，可以看到在編譯到LLVM IR之后，源代碼中的變量名標(biāo)識(shí)符（如data、result）均被替換為%{number}形式的變量名，同時(shí)編譯后的LLVM IR代碼具有靜態(tài)單賦值形式（SSA，static single assignment form），即每一個(gè)變量?jī)H能被賦值一次。另外，編譯到IR之后需要對(duì)代碼文本進(jìn)行預(yù)處理，主要包括函數(shù)名替換、特殊符號(hào)替換等。相較于使用源代碼，一方面，中間語(yǔ)言代碼的SSA形式和變量標(biāo)識(shí)符替換有利于后續(xù)代碼語(yǔ)義特征提取和預(yù)訓(xùn)練；另一方面，使用中間語(yǔ)言可以屏蔽不同類(lèi)型高級(jí)語(yǔ)言的語(yǔ)法差異，能夠兼容不同類(lèi)型的高級(jí)語(yǔ)言，起到擴(kuò)展訓(xùn)練樣本的作用。

圖3 源代碼和IR樣本示例

Figure 3 An example of source code and its LLVM IR

3.2 圖結(jié)構(gòu)數(shù)據(jù)的生成

本節(jié)需要實(shí)現(xiàn)代碼樣本控制流圖的抽取、代碼基本塊的預(yù)訓(xùn)練和圖結(jié)構(gòu)數(shù)據(jù)拼接。為了能夠充分提取代碼樣本的結(jié)構(gòu)信息，本文采用代碼的控制流圖作為圖神經(jīng)網(wǎng)絡(luò)輸入樣本的框架，同時(shí)對(duì)控制流圖中的基本塊代碼進(jìn)行預(yù)訓(xùn)練初始化，最后將訓(xùn)練好的基本塊初始化向量填回到控制流圖結(jié)構(gòu)中，完成圖結(jié)構(gòu)數(shù)據(jù)的生成。

LLVM IR代碼的控制流圖抽取采用LLVM opt工具實(shí)現(xiàn)。opt命令是模塊化的LLVM優(yōu)化器和分析器，它將LLVM IR作為輸入，在其上運(yùn)行用于實(shí)現(xiàn)指定優(yōu)化或分析的PASS模塊，然后輸出優(yōu)化的文件或分析結(jié)果。通過(guò)opt調(diào)用生成CFG的PASS模塊可以生成dot文件格式的控制流圖。

針對(duì)需要根據(jù)控制流圖基本塊內(nèi)的代碼信息，對(duì)基本塊進(jìn)行初始化向量訓(xùn)練，本文方案采用word2vec算法[15]實(shí)現(xiàn)基本塊初始化向量的預(yù)訓(xùn)練?；緣K中的代碼是完全順序結(jié)構(gòu)，不再具有跳轉(zhuǎn)指令，本文把基本塊中的代碼語(yǔ)句經(jīng)過(guò)預(yù)處理之后當(dāng)作word2vec訓(xùn)練語(yǔ)料中的一句sentence，遍歷所有的訓(xùn)練集基本塊代碼生成語(yǔ)料庫(kù)。在該語(yǔ)料庫(kù)上使用word2vec算法訓(xùn)練語(yǔ)料庫(kù)詞典，也就獲得了所有單詞的詞嵌入向量，進(jìn)一步通過(guò)取基本塊中所有詞向量的平均值即可獲得基本塊的初始化特征向量。

本文采用pytorch深度學(xué)習(xí)框架進(jìn)行神經(jīng)網(wǎng)絡(luò)模型開(kāi)發(fā)和訓(xùn)練。PyG（pytorch_geometric）是基于pytorch的一個(gè)圖神經(jīng)網(wǎng)絡(luò)框架，它定義了一種表示圖結(jié)構(gòu)數(shù)據(jù)的數(shù)據(jù)類(lèi)型。本文實(shí)現(xiàn)了從dot形式的文件到PyG圖數(shù)據(jù)類(lèi)型之間的轉(zhuǎn)化，轉(zhuǎn)化后的圖數(shù)據(jù)包括2×num_edges維度的邊特征向量、num_nodes×num_node_features維度的節(jié)點(diǎn)特征向量和一維的全局標(biāo)簽向量。其中，num_edges為控制流圖中邊的數(shù)量，num_nodes為控制流圖中節(jié)點(diǎn)的數(shù)量，也就是基本塊的數(shù)量，num_node_features表示基本塊節(jié)點(diǎn)的特征維度，與前述word2vec算法的輸出有關(guān)。從dot文件到PyG庫(kù)data類(lèi)型對(duì)象的轉(zhuǎn)化過(guò)程如算法1所示。

算法1 圖數(shù)據(jù)類(lèi)型轉(zhuǎn)化算法

輸入 dot Object

輸出 PyG.data Object

算法

for edge in dot.edges do

src, dst = edge.src, edge.dst

src_index = dot.nodes.get_index(src)

dst_index = dot.nodes.get_index(dst)

PyG.data.edge_index.add ([src_index, dst_ index])

for node in dot.nodes do

node_value = preprocess(node.value)

node_vector = Word2Vec.infer (node_ value)

PyG.data.x.add(node_vector)

PyG.data.y = dot.label #1/0

return PyG.data

圖3所示的代碼樣本的dot類(lèi)型圖文件如圖4所示，經(jīng)過(guò)算法1轉(zhuǎn)化之后，生成的PyG.data類(lèi)型的數(shù)據(jù)示例如圖5所示。該樣本有3個(gè)節(jié)點(diǎn)、3條邊，標(biāo)簽為bad，采用256維的word2vec算法進(jìn)行向量初始化，生成的PyG.data數(shù)據(jù)中邊特征向量維度為2×3，節(jié)點(diǎn)特征向量維度為3×256，標(biāo)簽為1，代表該樣本為負(fù)樣本。

圖4 dot類(lèi)型樣本示例

Figure 4 An example of dot sample

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

Gilmer等[14]提出的MPNN架構(gòu)是目前大多數(shù)圖神經(jīng)網(wǎng)絡(luò)的頂層抽象，這種信息傳播模型的公式如下所示。

圖5 PyG.data樣本數(shù)據(jù)示例

Figure 5 An example of PyG.data sample

本文提出的圖神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)由多組圖卷積和池化層組成的基本單元來(lái)完成在圖級(jí)別上的端到端分類(lèi)學(xué)習(xí)問(wèn)題，其網(wǎng)絡(luò)結(jié)構(gòu)如圖6所示。具體來(lái)說(shuō)，在每一組圖卷積與池化層中，通過(guò)應(yīng)用圖卷積層，本文為每個(gè)圖樣本中的每個(gè)節(jié)點(diǎn)獲得了特征向量轉(zhuǎn)化為具有固定數(shù)量維數(shù)的表示，并通過(guò)池化圖中所有節(jié)點(diǎn)的表示向量來(lái)獲得整個(gè)圖的表示。使用多組基本單元來(lái)抽取圖結(jié)構(gòu)中每個(gè)節(jié)點(diǎn)在不同層次上的特征。多組基本單元之間可以直接串聯(lián)連接，也可以通過(guò)殘差方式連接。文獻(xiàn)[16]指出殘差連接方式在多層網(wǎng)絡(luò)上有更好的效果。殘差連接方式如圖7所示，其表達(dá)如式(4)所示。最后，通過(guò)應(yīng)用多層感知器（MLP）和Softmax層，獲得最終的圖標(biāo)簽預(yù)測(cè)。

圖6 神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)

Figure 6 The architecture of our neural network

圖7 多層網(wǎng)絡(luò)殘差連接

Figure 7 Multi-layer network residual connection

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

4.1 建立訓(xùn)練與測(cè)試數(shù)據(jù)集

本文采用的數(shù)據(jù)集是來(lái)源于美國(guó)國(guó)家標(biāo)準(zhǔn)技術(shù)研究所（NIST）的軟件保障參考數(shù)據(jù)集（SARD）[17]。該數(shù)據(jù)集按照CWE編號(hào)收錄了各種源代碼和二進(jìn)制程序漏洞示例，同時(shí)部分具有修復(fù)后的代碼樣本。本文主要采用其C/C++語(yǔ)言的源代碼數(shù)據(jù)集，使用Clang編譯器和相關(guān)模塊獲得其代碼的LLVM IR中間語(yǔ)言表示和控制流圖，進(jìn)而根據(jù)方案設(shè)計(jì)生成圖結(jié)構(gòu)數(shù)據(jù)。本文實(shí)驗(yàn)選用了數(shù)據(jù)集中的整數(shù)溢出（CWE-190，integer overflow or wraparound）、內(nèi)存泄露（CWE-401，memory leak）和釋放非堆內(nèi)存（CWE-590，free of memory not on the heap）三大類(lèi)漏洞代碼數(shù)據(jù)，生成的函數(shù)級(jí)數(shù)據(jù)樣本數(shù)量如表1所示。模型訓(xùn)練與測(cè)試數(shù)據(jù)集按照7:3比例劃分。

表1 數(shù)據(jù)集漏洞樣本數(shù)量

4.2 模型評(píng)估指標(biāo)

本文根據(jù)不同類(lèi)型的數(shù)據(jù)集訓(xùn)練不同類(lèi)型的神經(jīng)網(wǎng)絡(luò)模型，再根據(jù)測(cè)試結(jié)果評(píng)估模型效果。在對(duì)模型進(jìn)行評(píng)估時(shí)，測(cè)試結(jié)果的混淆矩陣定義如表2所示。

表2 模型測(cè)試結(jié)果混淆矩陣定義

根據(jù)混淆矩陣的定義，使用以下指標(biāo)對(duì)模型測(cè)試結(jié)果進(jìn)行評(píng)估。

（1）Acc = (TP+TN)/(TP+TN+FP+FN)

Acc被稱(chēng)為分類(lèi)的準(zhǔn)確率，表示模型檢測(cè)正確的樣本占所有樣本數(shù)的比例，Acc越高越好。

（2）TPR = TP / (TP+FN)

TPR被稱(chēng)為Recall（召回率），表示分類(lèi)正確的漏洞樣本占實(shí)際所有漏洞樣本的比例，TPR越高越好。

（3）FPR = FP / (FP + TN)

FPR被表示錯(cuò)誤檢出的漏洞樣本數(shù)量占所有實(shí)際不含漏洞樣本的比例，F(xiàn)PR越小越好。

（4）Precision = TP / (TP + FP)

Precision被稱(chēng)為精確率，表示在模型檢出的所有樣本中含有漏洞的樣本所占的比例，Precision越高越好。

（5）-1 = 2.Precision.Recall/ (Precision + Recall)

-1是Precision和Recall的調(diào)和平均，綜合了兩者的結(jié)果，-1值越高，代表模型性能越好。

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

本文的模型測(cè)試環(huán)境為CentOS 7.6操作系統(tǒng)，Intel Xeon CPU E5-2630處理器，180 GB內(nèi)存，GeForce GTX 1080顯卡，8 GB顯存。模型使用GPU顯卡加速模型訓(xùn)練。

在圖結(jié)構(gòu)數(shù)據(jù)生成部分，word2vec算法生成的詞嵌入?yún)?shù)大小設(shè)定為256，基本塊初始化向量的大小即num_node_features為256。

在圖神經(jīng)網(wǎng)絡(luò)部分，本文在CWE-190數(shù)據(jù)集上測(cè)試了具有不同層數(shù)網(wǎng)絡(luò)基本單位和不同層連接方式的圖神經(jīng)網(wǎng)絡(luò)的效果。本文選擇了從1到6層圖卷積層加池化層網(wǎng)絡(luò)基本單元，分別使用殘差鏈接和直接連接兩種層間連接方式，其訓(xùn)練結(jié)果如圖8所示。

圖8 不同網(wǎng)絡(luò)結(jié)構(gòu)在測(cè)試集上的準(zhǔn)確率

Figure 8 Accuracy of different network structures on the test dataset

由圖8可知，直接連接方案在大于2層網(wǎng)絡(luò)層之后出現(xiàn)了準(zhǔn)確率的大幅下降，而殘差網(wǎng)絡(luò)連接方案在多層網(wǎng)絡(luò)環(huán)境下具有更好的效果，因此本文選擇了效果最好的3層殘差連接網(wǎng)絡(luò)。

其他網(wǎng)絡(luò)參數(shù)包括中間層圖卷積隱藏層的輸入輸出通道數(shù)，其值均為128，優(yōu)化器選擇Adam，學(xué)習(xí)率為0.001，損失函數(shù)為交叉熵。在完成50個(gè)epoch訓(xùn)練后，測(cè)試集上的損失基本穩(wěn)定，模型能夠成功收斂。模型訓(xùn)練準(zhǔn)確率和損失曲線如圖9所示。

圖9 CWE-190模型訓(xùn)練中測(cè)試集準(zhǔn)確率與損失曲線

Figure 9 Accuracy and loss curve of test dataset in CWE-190 model training

模型訓(xùn)練時(shí)間、測(cè)試時(shí)間和靜態(tài)掃描工具掃描時(shí)間如表3所示?？梢钥吹?，模型訓(xùn)練所需要時(shí)間較多，但當(dāng)模型訓(xùn)練完成后，再使用訓(xùn)練好的模型進(jìn)行漏洞測(cè)試時(shí)，運(yùn)行時(shí)間較短，這也是使用學(xué)習(xí)模型方法解決漏洞檢測(cè)問(wèn)題的優(yōu)勢(shì)之一，根據(jù)實(shí)驗(yàn)結(jié)果可以看到模型測(cè)試時(shí)間與傳統(tǒng)靜態(tài)掃描工具時(shí)間基本持平。

將本文提出的方法與現(xiàn)有的開(kāi)源靜態(tài)掃描工具和LSTM網(wǎng)絡(luò)進(jìn)行檢測(cè)效果對(duì)比。目前沒(méi)有廣泛應(yīng)用在中間表示的靜態(tài)掃描工具，本文選擇了Flawfinder[18]和RATS[19]兩種支持C/C++代碼靜態(tài)掃描工具進(jìn)行對(duì)比，兩者在實(shí)驗(yàn)時(shí)均是對(duì)C/C++源代碼樣本進(jìn)行掃描和檢測(cè)，同時(shí)使用LSTM網(wǎng)絡(luò)對(duì)中間語(yǔ)言樣本進(jìn)行了訓(xùn)練和檢測(cè)。測(cè)試集上的評(píng)估指標(biāo)對(duì)比如表4所示。

實(shí)驗(yàn)表明，本文提出的基于控制流圖的特征提取與圖神經(jīng)網(wǎng)絡(luò)模型在開(kāi)源數(shù)據(jù)集上的準(zhǔn)確率均大幅高于靜態(tài)掃描工具，在開(kāi)源工具普遍召回率為10%左右的情況下，本文的方法召回率可以達(dá)到60%至70%，RATS工具下CWE-401上取得較高召回率TPR是由于其FPR較高，即存在較高的誤報(bào)率。與LSTM網(wǎng)絡(luò)模型相比，本文提出的圖網(wǎng)絡(luò)模型在各項(xiàng)指標(biāo)上均更有優(yōu)勢(shì)，同時(shí)本文提出的模型方法在測(cè)試時(shí)間基本持平的情況下取得了更高的準(zhǔn)確率指標(biāo)，因此，本文提出的方法相較于傳統(tǒng)方法效果更優(yōu)。

5 結(jié)束語(yǔ)

本文提出了一種基于圖神經(jīng)網(wǎng)絡(luò)的代碼漏洞檢測(cè)方法，一方面通過(guò)詞向量初始化基本塊節(jié)點(diǎn)，另一方面通過(guò)多層圖卷積神經(jīng)網(wǎng)絡(luò)在控制流圖上進(jìn)行信息傳播和聚合，通過(guò)MLP層對(duì)函數(shù)級(jí)別的中間語(yǔ)言代碼樣本進(jìn)行漏洞檢測(cè)。本文在開(kāi)源數(shù)據(jù)集3類(lèi)漏洞樣本上進(jìn)行分類(lèi)實(shí)驗(yàn)，取得了良好的效果，證明該方法相較于傳統(tǒng)靜態(tài)掃描工具更加有效，通過(guò)引入圖神經(jīng)網(wǎng)絡(luò)進(jìn)行漏洞特征提取和漏洞檢測(cè)的可行性得到了證明。進(jìn)一步的研究包括實(shí)驗(yàn)跨語(yǔ)言模型的檢測(cè)，驗(yàn)證基于中間語(yǔ)言的模型訓(xùn)練是否可遷移到其他高級(jí)語(yǔ)言；改進(jìn)基本塊向量初始化辦法以及實(shí)驗(yàn)其他類(lèi)型的深度圖神經(jīng)網(wǎng)絡(luò)在該方法中的效果，使用更細(xì)粒度的圖結(jié)構(gòu)使圖神經(jīng)網(wǎng)絡(luò)能夠提取更細(xì)粒度的特征等。

表4 不同模型在不同類(lèi)型測(cè)試集上的評(píng)估指標(biāo)對(duì)比

[1]《國(guó)家網(wǎng)絡(luò)空間安全戰(zhàn)略》(全文)[J]. 中國(guó)信息安全, 2017(1): 26-31.

“National Cyberspace Security Strategy” (full-text) [J]. China Information Security, 2017(1):26-31.

[2]WU Z, PAN S, CHEN F, et al. A comprehensive survey on graph neural networks[J]. arXiv:1901.00596, 2019.

[3]PHAM N H, NGUYEN T T, NGUYEN H A, et al. Detection of recurring software vulnerabilities[C]//25th IEEE/ACM International Conference on Automated Software Engineering. 2010.

[4]YAMAGUCHI F, GOLDE N, ARP D, et al. Modeling and discovering vulnerabilities with code property graphs[C]//IEEE Symposium on Security and Privacy. 2014.

[5]LIN G, ZHANG J, LUO W, et al. POSTER: vulnerability discovery with function representation learning from unlabeled projects[C]// ACM Sigsac Conference. 2017.

[6]RUSSELL R, KIM L, HAMILTON L, et al. Automated vulnerability detection in source code using deep representation learning[C]// 2018 17th IEEE International Conference on Machine Learning and Applications (ICMLA). 2018: 757-762.

[7]XU X, LIU C, FENG Q, et al. Neural network-based graph embedding for cross-platform binary code similarity detection[C]//Proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security-(CCS′17). 2017: 363-376.

[8]LI Y, GU C, DULLIEN T, et al. Graph matching networks for learning the similarity of graph structured objects[C]//Thirty-sixth International Conference on Machine Learning(ICML 2019). 2019.

[9]YU Z, CAO R, TANG Q, et al. Order matters: semantic-aware neural networks for binary code similarity detection[C]//The Thirty-Fourth AAAI Conference on Artificial Intelligence (AAAI-20). 2020.

[10]DUAN Y, LI X, WANG J, et al. DeepBinDiff: learning program-wide code representations for binary diffing[C]//Proceedings 2020 Network and Distributed System Security Symposium. 2020.

[11]李珍, 鄒德清, 王澤麗, 等. 面向源代碼的軟件漏洞靜態(tài)檢測(cè)綜述[J]. 網(wǎng)絡(luò)與信息安全學(xué)報(bào), 2019, 5(1):1-14.

LI Z, ZOU D Q, WANG Z L, et al. Survey on static software vulnerability detection for source code[J]. Chinese Journal of Network and Information Security, 2019, 5(1): 1-14.

[12]GORI M, GABRIELE M, FRANCO S. A new model for learning in graph domains[C]//IEEE International Joint Conference on Neural Networks. 2005.

[13]SCARSELLI F, GORI M, TSOI A C, et al. The graph neural network model[J]. IEEE Transactions on Neural Networks, 2009, 20(1): 61-80.

[14]GILMER J, SCHOENHOLZ S S, RILEY P F. et al. Neural message passing for Quantum chemistry[C]//Proceedings of the 34th International Conference on Machine Learning-Volume 70 (ICML’17). 2011: 1263–1272.

[15]MIKOLOV T, CHEN K, CORRADO G, et al. Efficient estimation of word representations in vector space[J]. arXiv preprint arXiv:1301.3781, 2013.

[16]KIPF T N, WELLING M. Semi-supervised classification with graph convolutional networks[J]. arXiv preprint arXiv:1609.02907, 2016.

[17]NIST software assurance reference dataset project[EB].

[18]Flawfinder[EB].

[19]Rough-auditing-tool-for-security [EB].

Code vulnerability detection method based on graph neural network

CHENHao, YIPing

School of Cyber Science and Engineering, Shanghai Jiao Tong University, Shanghai 200240, China

The schemes of using neural networks for vulnerability detection are mostly based on traditional natural language processing ideas, processing the code as array samples and ignoring the structural features in the code, which may omit possible vulnerabilities. A code vulnerability detection method based on graph neural network was proposed, which realized function-level code vulnerability detection through the control flow graph feature of the intermediate language. Firstly, the source code was compiled into an intermediate representation, and then the control flow graph containing structural information was extracted. At the same time, the word vector embedding algorithm was used to initialize the vector of basic block to extract the code semantic information. Then both of above were spliced to generate the graph structure sample data. The multilayer graph neural network model was trained and tested on graph structure data features. The open source vulnerability sample data set was used to generate test data to evaluate the method proposed. The results show that the method effectively improves the vulnerability detection ability.

vulnerability detection, graph neural network, control flow graph, intermediate representation

TP309

A

10.11959/j.issn.2096?109x.2021039

2020?10?19；

2020?12?15

易平，yiping@sjtu.edu.cn

國(guó)家重點(diǎn)研發(fā)計(jì)劃（2019YFB1405000, 2017YFB0802900）

The National Key R&D Program of China (2019YFB1405000, 2017YFB0802900)

陳皓, 易平. 基于圖神經(jīng)網(wǎng)絡(luò)的代碼漏洞檢測(cè)方法[J]. 網(wǎng)絡(luò)與信息安全學(xué)報(bào), 2021, 7(3): 37-45.

CHEN H, YI P. Code vulnerability detection method based on graph neural network[J]. Chinese Journal of Network and Information Security, 2021, 7(3): 37-45.

陳皓（1995? ），男，山東濰坊人，上海交通大學(xué)碩士生，主要研究方向?yàn)樯疃葘W(xué)習(xí)與漏洞檢測(cè)。

易平（1969? ），男，河南洛陽(yáng)人，博士，上海交通大學(xué)副教授，主要研究方向?yàn)槿斯ぶ悄馨踩?/p>