基于代碼序列與圖結構的源代碼漏洞檢測方案

王守梁

(中北大學 計算機科學與技術學院,山西 太原 030051)

0 引 言

軟件源代碼漏洞檢測一直是軟件安全領域的一個重要挑戰[1-3]。隨著大數據時代的來臨,軟件代碼數量以量級攀升,軟件更加復雜與多樣,任何微小的漏洞都可能導致整個系統的崩潰。據統計,2022年,缺陷名錄網站(NVD)全年的漏洞總數為6 515個,預期年增長率提升了170%; 公共漏洞披露平臺(CVE)全年的漏洞總數為10 703個,預期年增長率為70%。政府單位、銀行以及科技企業已成為漏洞攻擊的重點[4],為了維護網絡安全,需對軟件源代碼漏洞檢測相關工作提出更高的要求。

傳統的漏洞檢測主要使用3種技術:靜態分析技術[5]、動態分析技術[6]及機器學習技術[7]。靜態分析指在不執行程序的情況下,通過掃描源代碼,分析代碼詞法、語法、數據流及控制流等信息來篩選漏洞,耗時耗力。動態分析指在執行程序的情況下,通過聚焦執行路徑、程序狀態等信息來篩選漏洞,可擴展性差。機器學習技術借助預設特征的機器學習算法來檢測漏洞,在復雜數據集上效果較差。

近年來,該領域研究方向有所轉變,由被動檢測轉向漏洞代碼的主動預測。同時,深度學習與漏洞檢測的結合成為研究熱點?，F有的基于深度學習的漏洞檢測成果表明,深度學習較傳統技術檢測具有更優越的檢測性能?，F有的檢測方案亦存在一定的局限性。一類研究[8-11]將源代碼看作特殊的自然語言,使用自然語言處理領域的成熟模型解決漏洞檢測任務,未考慮代碼的結構信息; 另一類研究[12-13]抽取代碼對應的圖表示并借助圖神經網絡解決漏洞檢測任務,現有的相關方案一方面未考慮異類邊的處理,另一方面忽略了節點內語句的局部語義信息。此外,現有研究多使用以word2vec[14]、glove[15]、fasttext[16]及doc2vec[17]為代表的傳統詞向量模型為標記生成固定的向量表示,未考慮代碼上下文差異。

為解決以上問題,本文提出了一種函數級源代碼漏洞檢測方案,綜合考慮源代碼基于圖的中間表示與基于序列的中間表示來輔助漏洞檢測。

本文方案的實現步驟過程如下:1) 對源代碼進行預處理,并抽取其代碼屬性圖(Code Property Graph,CPG)[18]。受文獻[19-20]中方法的啟發,通過增加自定義類型的邊對CPG圖進行擴展,得到CPG’,其中,代碼語句被解析為圖節點,語句之間的關系被解析為圖邊。對于節點,分別對其類型與語句進行編碼; 對于邊,分別對其類型與方向進行編碼。同時,采用關系圖卷積神經網絡(RGCN)聚合信息來生成最終圖表示。2) 切分預處理后的源代碼得到對應的標記序列,并使用基于程序語言的預訓練模型CodBert[21]編碼器,生成源代碼對應的詞向量矩陣。3) 集成以上兩種表示并采用由全連接神經網絡(FCNN)與sigmoid層組成的分類器實現漏洞檢測。

1 相關工作

在源代碼漏洞檢測領域,現有的綜述已從不同的角度對漏洞檢測領域的研究進行了全面的總結[22-26]。本節將聚焦本文研究工作的三個關鍵方面:源代碼中間表示、詞向量技術與神經網絡模型,并從以上三個角度對已有的漏洞檢測工作分別進行總結回顧。

1) 源代碼中間表示形式主要分為兩類:一是序列形式,即將源代碼轉換成線性標記序列。Wang[27]、Russell[28]及Li[29]等在源代碼序列基礎上借助不同的神經網絡模型實現了漏洞檢測。二是圖形式,即將源代碼轉換成結構圖。Chakraborty[30]、Li[31]及Dong[32]等將源代碼轉換成CPG圖,Li等[33]將源代碼轉換成PDG圖,Fan等[19]將源代碼轉換成AST、CFG、DFG及NCS四種圖。成果匯總見表1。本文按照圖種類與所抽取的信息類型,對現有研究進行梳理,成果如表2 所示。其中,CPG圖因其含有更多類型的邊關系,能反映更豐富的程序結構信息而被多數研究采用。

表1 漏洞檢測源代碼中間表示研究匯總

表2 基于圖表示的漏洞檢測研究匯總

2) 詞向量技術用于生成神經網絡可處理的實值向量,其選取決定著漏洞代碼的上下文信息的提取,進而影響模型的檢測性能。Pradel和Sen[37]使用Word2Vec生成代碼向量,用于訓練深度學習模型,以檢測JavaScript代碼中的漏洞。Henkel等[38]應用GloVe從C程序的抽象符號軌跡中學習向量表示。Fang等[39]應用FastText進行基于集成機器學習模型的漏洞預測。Kanade等[40]提出了CuBERT,它通過在軟件源代碼上訓練BERT模型來生成上下文嵌入。Karampatsis和Sutton[41]提出了一個名為SCELMo的模型來生成上下文代碼向量。詳細成果匯總見表3。

表3 詞向量轉化技術研究匯總

3) 神經網絡模型主要分為兩類,分別對應兩種中間表示形式。包括卷積神經網絡(CNN)、多層感知器(MLP)、深度置信網絡(DBN)、長短期記憶網絡(LSTM)及門控單元網絡(GRU)等在內的多種神經網絡模型被應用在基于序列的漏洞檢測中并取得了良好的檢測效果; 包括門控圖神經網絡(GGNN)、圖卷積神經網絡(GCN)等在內的多種圖神經網絡模型被應用在基于圖的漏洞檢測中,用于生成圖向量。成果匯總見表4。

表4 漏洞檢測神經網絡模型匯總

2 模型設計

2.1 問題定義與整體架構

2.1.1 問題定義

2.1.2 整體架構

如圖1 所示,該漏洞檢測方案將函數級源程序樣本作為輸入,包括以下兩大模塊:特征抽取與特征學習。特征抽取階段,使用兩種方式對源代碼進行處理:① 抽取函數源代碼對應的多類型邊結構圖。② 抽取函數源代碼對應標記序列。特征學習階段,分別生成以上兩種形式的中間表示對應的實值向量,并將二者集成拼接。最后,使用全連接網絡輸出漏洞檢測結果。詳細原理將在后續章節中進行闡述。

圖1 本研究漏洞檢測方案的整體架構

2.2 基于圖的特征學習

2.2.1 圖的生成

為學習源代碼中的結構信息并提取漏洞特征,需要將其轉換成合適的結構圖。結構圖的質量在很大程度上決定了從訓練數據中學習到的特征數量,進而影響模型的漏洞檢測性能。代碼屬性圖(簡稱CPG) 提供了控制流圖(簡稱CFG)和數據流圖(簡稱DFG)中的元素以及抽象語法樹(簡稱AST)和程序依賴圖(簡稱PDG)的代碼組合,是現存最綜合全面的程序圖表示,本文選擇其作為源代碼的初始結構圖。

源代碼函數對應的CPG按照以下流程生成:① 解析源代碼函數生成AST。② 應用控制流分析生成CFG。③ 應用依賴性分析生成PDG。④ 在AST中添加CFG與PDG邊生成CPG。以上過程使用工具Joern實現。

CPG在富含多種結構信息的同時,缺失了語句中的順序信息,為解決此問題,本文在初始結構圖CPG基礎上增加自然代碼序列邊(NCS)[50]并對其進行了擴展,通過在AST中的相鄰葉節點之間從左到右添加有向邊,增加了代碼編輯的時序信息,輔助漏洞檢測。本文將最終得到的結構圖記為CPG’=〈V,E〉,其中,V={v|v∈ASTnode}表示圖中的節點集合,ASTnode表示函數的AST中的節點,每個節點包含節點類型與節點代碼兩種元素;E={e|e∈CFGedge或e∈DDGedge或e∈CDGedge或e∈ASTedge或e∈NCSedge}表示圖中的邊集合,這些邊分別反映了節點之間的控制流、數據依賴、控制依賴、語法結構與時序信息。圖2對應一個緩沖區溢出漏洞的簡單樣例,其main函數對應的CPG’如圖3 所示。

圖2 緩沖區溢出漏洞代碼樣例

圖3 CPG’

后續任務中基于CPG’生成神經網絡能處理的實值向量至關重要,本文將綜合考慮結構圖中的節點及邊。

2.2.2 節點嵌入

節點嵌入模塊分別對節點類型與節點代碼采用不同嵌入方法來生成初始向量:① 對于節點類型,采用獨熱編碼(one-hot)生成對應向量xi。② 對于節點代碼,首先,完成標準化處理,將源代碼中用戶自定義的變量名稱與函數名稱采用統一標準進行重命名,用VAR與序號的組合替代自定義變量名,用FUN與序號的組合替代自定義函數名,如“VAR1”和“FUN1”,減少語料庫規模與噪聲干擾,提升向量生成的質量; 其次,使用基于編程語言的預訓練模型CodeBert訓練的單詞嵌入權重初始化嵌入層權重,CodeBert模型使用的字節對編碼(BPE)標記器[51]能在很大程度上緩解Out-Of-Vocabulary(OOV)問題; 最后,連接雙向長短期記憶網絡(Bi-LSTM)聚合局部語義信息,獲得最終向量表示yi。最終,連接xi,yi得到節點vi對應的向量表示gi。

上述過程的形式化表述如下:

對于節點類型v[type,i],有

xi=one-hot(v[type,i])。

(1)

對于源代碼S={s1,s2,s3,…,sn},其中,si對應節點vi中的代碼。首先,進行標準化處理

S′=normalization(S)。

(2)

其次,通過CodeBert預訓練的BPE標記器進行標記并使用CodeBert嵌入層權重來初始化自嵌入權重,以獲得每個標記的向量表示eij∈Ei={ei1,ei2,ei3,…,ein}。

(3)

Ei=Codebert(tokensi)。

(4)

最后,連接Bi-LSTM獲取與節點vi對應的語句si的向量表示yi。

(5)

(6)

最終,連接xi,yi得到節點vi的最終向量表示

gi=CONCAT(xi,yi)。

(7)

2.2.3 邊嵌入

邊嵌入模塊分別對邊方向與邊類型采用不同嵌入方法來生成初始向量:① 對于邊方向,采用深度優先遍歷(DFS)將節點排序并從0開始編號,以節點編號數對的形式編碼邊方向; ② 對于邊類型,采用標簽編碼(label encoding)將有限邊類型進行編碼。

上述過程的形式化表述如下:函數源代碼對應的結構圖記為CPG’=〈V,E〉,其中,V代表節點集合,E代表邊集合。

對于邊方向,有

IE={〈vi.order,vj.order〉|e=

〈vi,vj〉,e∈E}。

(8)

對于邊類型,有

TE={te|e∈E},

(9)

(10)

圖4 為一個函數源代碼對應的邊嵌入過程的示例。

圖4 邊嵌入策略實例說明

圖4 中,節點按照深度優先遍歷算法進行編號,5條有向邊分別對應5種不同的邊類型,還給出了邊方向與邊類型對應的向量形式。

2.2.4 圖神經網絡

上述環節孤立考慮每個節點,僅根據單個節點包含的元素進行向量初始化,缺乏相鄰節點的狀態信息,無法較好地反映源代碼的結構信息。本研究采用關系圖卷積網絡(RGCN)完成鄰域聚合,進而更新節點狀態。

圖卷積網絡(GCN)是CNN在圖結構上的擴展,是在給定圖上應用圖卷積運算后,通過聚合其相鄰節點的所有特征來更新節點。與連續計算節點狀態的遞歸圖神經網絡(如GGNN)相比,GCN是每層更新一次節點向量,因此效率更高,可擴展性更強。RGCN是在GCN的基礎上進行了增量改進,其通過在每種關系類型下用相鄰節點更新節點來區別對待不同類型的邊,更適合處理本文所抽取的多類型邊異構圖CPG’。如圖5 所示,在每個RGCN層中,更新單個節點需要按邊類型積累其所有相鄰節點的信息,這些信息隨后與該節點本身的特征相結合,通過激活函數(本文采用ReLU)傳遞,得到最終表示。

圖5 節點更新過程

上述過程的形式化表述如下:

函數源代碼對應的結構圖記為CPG’=〈V,E〉,其中,V代表節點集合,E代表邊集合。單個節點向量的更新公式為

(11)

2.2.5 圖嵌入

圖嵌入模塊為基于圖的特征學習的最后一步,目標是得到最終的圖向量表示。本文通過聚合所有更新后的節點向量生成代表整個CPG’的單個向量,即

(12)

2.3 基于序列的特征學習

基于序列的特征學習是本方案理解源代碼全局語義信息的關鍵過程?，F有研究普遍采用以Word2vec,GloVe及FastText為代表的詞向量模型生成源代碼對應的實值向量,這些模型無法基于不同的代碼上下文為相同標記生成不同向量,即僅為相同標記生成唯一向量,不考慮標記對應變量值的變化,忽略了上下文信息在漏洞檢測工作中的關鍵作用。

本研究采用基于代碼語言的預訓練模型CodeBert來生成源代碼的特征向量,充分考慮了上下文信息。相較傳統詞向量模型,CodeBert的模型結構要更復雜,其結構繼承了BERT的架構,基于一個12層雙向transformer,每層包含12個自注意頭,每個自注意頭的大小設為64,隱藏維度設為768,transformer的多頭注意力機制能夠集中數據流的多個關鍵變量,有助于分析和跟蹤潛在的漏洞代碼數據,幫助學習需長期依賴性分析的漏洞代碼模式。此外,其內置字節對編碼(BPE)標記器在很大程度上緩解了OOV問題,預訓練技術也克服了數據集中漏洞樣本過少或漏洞樣本與非漏洞樣本的不平衡所帶來的過擬合問題。最后,對于每一個函數級源代碼樣本,CodeBert編碼器生成對應的向量矩陣,并封裝其語義含義。

上述過程的形式化表述如下:fi代表一個函數級源代碼樣本輸入,其向量矩陣為

Pi=CodeBert(fi)。

(13)

2.4 特征融合

基于以上流程,分別獲得了所有函數級源代碼樣本fi的基于結構圖生成的向量表示G與基于序列生成的向量表示Pi。本研究采用連接操作組合二者,進而集成兩種特征,生成所有源代碼函數對應的最終向量表示Hi,其公式為

Hi=CONCAT(G,Pi)。

(14)

2.5 分類器

鑒于全連接神經網絡(FCNN)的卓越性能和在頂級分類器中的主導地位,本文使用全連接網絡作為漏洞檢測分類器,其由輸入層、隱藏層與輸出層三部分組成。

上述特征融合模塊的輸出被作為輸入層,最終,輸出層將利用隱藏層特征進行漏洞預測。

(15)

(16)

式中:σ代表隱層激活函數;W1,W2,b1與b2代表可學習的參數。

3 實驗設計

3.1 數據集

本文基于7個不同的c/c++源代碼數據集開展實驗。其中,4個為簡單數據集,它們分別為CWE-362,CWE-476,CWE-754與CWE-758,其從代碼結構和功能相對簡單的軟件中收集,包含可能更容易識別的漏洞模式; 剩余3個為復雜數據集,分別為Big-Vul[52],Reveal[30]與FFMPeg+Qemu[50],而Fan等[52]數據集涵蓋了2002年至2019年的CWE。在函數粒度上,數據集包含超過10 000個漏洞函數。Reveal[30]數據集包含超過18 000個函數,其中9.16%是漏洞函數。FFMPeg+Qemu[50]數據集包含超過22 000個函數,其中45.0%的為漏洞函數。表5 展示了數據集的相關細節。

表5 7個數據集

3.2 研究問題

實驗結果將以回答以下研究問題的形式呈現:

研究問題1:本文提出的漏洞檢測方案與現有軟件漏洞檢測工具相比性能如何?

對照組1:AVDetect[28],VulDeePecker[29],Codebert[21];

對照組2:Devign[50],Reveal[30],IVDetect[33],FUNDED[20];

對照組3:DeepVulSeeker[34]。

以上實驗方案均采用基于深度學習的漏洞檢測技術,其中對照組1中的模型采用基于序列的方法,對照組2中的模型采用基于圖的方法,而對照組3中的模型采用基于圖與序列相結合的方法。

研究問題2:綜合考慮圖與序列信息是否有助于提升模型的漏洞檢測性能?

本研究的漏洞檢測方案抽取的信息類型包括圖節點(節點類型與節點元素)、圖邊(邊類型與邊方向及序列信息),為驗證綜合考慮多類型信息對漏洞檢測模型的有效性,采取控制變量的方法,減少對照組中抽取的信息類型。

對照組1:僅抽取圖信息; 對照組2:僅抽取序列信息。

研究問題3:不同的節點初始化嵌入方法是否影響方案的漏洞檢測性能?

本文漏洞檢測方案采用預訓練模型CodeBert對節點元素進行初始化,并通過再次訓練微調嵌入權重。通過設計對照實驗與其他圖節點向量化方法進行比較,具體對照組設置如下:

對照組1:Word2vec; 對照組2:Dov2vec。

研究問題4:不同的序列初始化嵌入方法是否影響方案的漏洞檢測性能?

本文漏洞檢測方案采用預訓練模型CodeBert對序列標記進行初始化,并通過再次訓練微調嵌入權重。通過設計對照實驗與傳統的詞向量模型進行比較,具體對照組設置如下:

對照組1:Word2vec; 對照組2:GloVe; 對照組3:Fasttext。

3.3 實驗環境

本文漏洞檢測系統開發環境為Linux操作系統,使用Python語言進行開發,深度學習庫選擇以Tensorflow-gpu為后端的Keras框架開發,其具體內置版本信息如表6 所示。

表6 實驗環境參數

3.4 評估指標

本文采用4個指標對模型進行評估,即漏報率(FNR)、誤報率(FPR)、準確率(Accuracy)及F1指數(F1_score),具體計算公式分別為

其中,

式中:NTP表示將正樣本預測為正樣本的數量;NTN表示將正樣本預測為負樣本的數量;NFP表示將負樣本預測為正樣本的數量,即誤報;NFN表示將正樣本預測為負樣本的數量,即漏報。因此,漏報率(FNR)表示漏報的樣本占正樣本數量的比例,誤報率(FPR)表示誤報的樣本占負樣本數量的比例,準確率(Accuracy)表示預測正確的數據占總樣本的比例,F1指數(F1_score)是根據準確率(Precision)與召回率(Recall)給出的一種綜合評價,是二者的調和平均值。

具有高準確率(Accuracy)、高F1指數(F1_score)和低漏報率(FNR)、低誤報率(FPR)的模型被認為是高效的。

4 實驗結果及分析

4.1 研究問題1的實驗結果

表7 與表8 展示了問題1所涉及的對照組1的實驗結果?？梢钥闯?在簡單數據集上,本研究模型在 4個數據集上的平均準確率與F1值均高于對照組所有模型; 在真實世界數據集上,實驗組除在FFMPeg+Qemu上的F1值略低于VulDeePecker[29]與codebert[21]以外,其他各項結果均高于對照組結果。

表7 對照組1的實驗結果1

表8 對照組1的實驗結果2

表9 與表10 展示了問題1所涉及的對照組2的實驗結果?？梢钥闯?在簡單數據集上,本研究模型在4個數據集上的平均準確率與F1值均高于對照組所有模型; 在真實世界數據集上,實驗組除在FFMPeg+Qemu數據集上的F1值略低于IVDetect[33]與FUNDED[20]以外,其他各項結果均高于對照組結果。

表9 對照組2的實驗結果1

表10 對照組2的實驗結果2

表11 與表12 展示了問題1所涉及的對照組3 的實驗結果?？梢钥闯?在簡單數據集上,本研究模型在CWE-362與CWE-476上的兩項指標均高于對照組模型,在CWE-754與CWE-758上的兩項指標雖然相對略低,但兩項指標均值仍分別達到0.980 8與0.972 4; 在真實世界數據集上,實驗組除在FFMPeg+Qemu上的F1值略低以外,其他各項結果均高于對照組結果。

表11 對照組3的實驗結果1

表12 對照組3的實驗結果2

綜上所述,可以得出

結論1:與現有方案相比,在兩類數據集上,本研究的漏洞檢測方案都有更卓越的性能。

4.2 研究問題2的實驗結果

表13 與表14 展示了問題2對應的實驗結果?？梢钥闯?對照組1(即單獨抽取圖表示)的整體性能比對照組2(即單獨抽取序列表示)更好,但實驗組在兩類數據集上的準確率與F1值均有更明顯的提高。以上差距驗證了圖結構信息與序列語義信息對漏洞檢測的重要性。

表13 研究問題2的實驗結果1

表14 研究問題2的實驗結果2

基于以上分析,可以得出

結論2:較單獨考慮圖或單獨考慮序列,綜合考慮二者的漏洞檢測方案有著更好的檢測性能。

4.3 研究問題3的實驗結果

表15 與表16 展示了問題3對應的實驗結果?？梢钥闯?在兩類數據集上,實驗組都具有更高的準確率與F1值,其中,較對照組1,在簡單數據集上,準確率差異的平均值達到0.172 3,F1值差異的平均值達到0.174 4; 在復雜數據集上,準確率差異的平均值達到0.139 4,F1值差異的平均值達到0.116 5。較對照組2,在簡單數據集上,準確率差異的平均值達到0.130 1,F1值差異的平均值達到0.133 7; 在復雜數據集上,準確率差異的平均值達到0.110 1,F1值的差異的平均值達到0.091 5。以上差距驗證了本文所使用的節點嵌入方式能有效替代傳統的詞向量方法,具有一定的合理性。

表15 研究問題3的實驗結果1

表16 研究問題3的實驗結果2

基于以上分析,可以得出

結論3:與傳統的節點源代碼初始化嵌入方式相比,本研究所采用的方法能更好地提取局部語義信息,進而提升方案的漏洞檢測性能。

4.4 研究問題4的實驗結果

表17 與表18 展示了問題4對應的實驗結果?？梢钥闯?所有對照組中,對照組2(即使用GloVe模型)的整體性能最好,然而,實驗組的結果顯示,在簡單數據集上,準確率平均提高了0.061 1,F1值平均提高了0.070 7; 在真實世界數據集上,準確率平均提高了0.065 8,F1值平均提高了0.066 3。以上差距驗證了預訓練模型的引入較傳統詞向量模型對序列的嵌入有更好的效果,具有一定的合理性。

表17 研究問題4的實驗結果1

表18 研究問題4的實驗結果2

基于以上分析,可以得出

結論4:與傳統的源代碼序列初始化嵌入方式相比,本研究所采用的針對代碼的預訓練模型能更好地提取抽取代碼全局語義信息,進而提升模型的整體漏洞檢測性能。

5 結 論

本文主要結論如下:1) 提出了一個函數級源代碼漏洞檢測方案,綜合學習程序語義和結構信息來識別漏洞。本方案以自動化的方式捕獲了比現有方法更綜合的特征。2) 在多環節引入基于代碼的預訓練語言模型來代替傳統的詞向量模型進行信息編碼。使用CodeBert生成源代碼序列的詞向量矩陣,使用CodeBert進行節點元素初始化并用BiLSTM聚合節點內部的局部語義信息。3) 分別在簡單與復雜兩類數據集上與現有檢測方案進行了對照實驗。實驗結果表明,該方案在多個評估指標上顯著優于現有的漏洞檢測方法。此外,對照實驗也驗證了各環節的有效性,并且表明本研究對節點、邊以及序列信息的處理策略優于現有其他方案。

本研究尚存在以下三點局限性:1) 預訓練模型CodeBert含12個編碼器層,參數多達1.1億,訓練與部署成本高昂。2) 數據集均為C/C++語言書寫,語言類型單一。3) 僅關注漏洞識別,為二分類探究,無法準確區分漏洞類型。未來研究工作需要完善以上局限:1) 提出一個輕量級預訓練模型。2) 擴展數據集以適應更多種類的程序設計語言。3) 擴展本方案以適應多類型的漏洞檢測。