基于BERT模型的源代碼漏洞檢測技術(shù)研究

羅樂琦 張艷碩 王志強 文 津 薛培陽

(北京電子科技學院 北京 100070)

隨著互聯(lián)網(wǎng)技術(shù)的快速發(fā)展,互聯(lián)網(wǎng)已經(jīng)融入現(xiàn)代社會的方方面面,包括醫(yī)療保健、能源、交通、公共安全、教育、娛樂等,互聯(lián)網(wǎng)已經(jīng)成為民生國計中不可或缺的一部分.然而,互聯(lián)網(wǎng)面臨著越來越嚴重的網(wǎng)絡(luò)攻擊威脅.網(wǎng)絡(luò)安全漏洞可以出現(xiàn)在各種不同的系統(tǒng)和應(yīng)用程序中,包括操作系統(tǒng)、網(wǎng)絡(luò)設(shè)備、應(yīng)用程序、數(shù)據(jù)庫等.這些漏洞可能會被黑客用來竊取敏感信息、破壞系統(tǒng)、分發(fā)惡意軟件、進行勒索攻擊等.例如,勒索軟件WannaCry利用Windows操作系統(tǒng)的漏洞進行傳播,導(dǎo)致數(shù)百萬計算機系統(tǒng)和網(wǎng)絡(luò)設(shè)備癱瘓,波及醫(yī)院、銀行、電信運營商、政府機構(gòu)等.往往只是因為幾行代碼不規(guī)范或者沒有考慮到安全問題,就足以導(dǎo)致嚴重的漏洞并使系統(tǒng)成為容易攻擊的目標.著名的“心臟滴血”漏洞是OpenSSL密碼庫中一個嚴重的緩沖區(qū)溢出漏洞,影響了數(shù)十億互聯(lián)網(wǎng)用戶,但只需要增加2行安全代碼就可以防范.截至2022年,通用漏洞披露(common vulnerabilities and exposures, CVE)已記錄了366815個漏洞,其中危險漏洞占比高達41.82%.這些漏洞的成因復(fù)雜,種類繁多,給漏洞的預(yù)防和檢測帶來了很大的困難,如何高效地進行漏洞檢測是一個迫切且值得研究的課題.

檢測漏洞的經(jīng)典方法是靜態(tài)分析、動態(tài)分析和動靜混合分析,運用這些方法進行漏洞檢測需要人類專家進行大量工作.然而,現(xiàn)代軟件和操作系統(tǒng)的規(guī)模更大且復(fù)雜性也更高,這些方法耗費的人力成本越來越大,漏洞檢測的效果也不佳.

機器學習技術(shù)的出現(xiàn)帶來了漏洞檢測的新思路.機器學習算法可以學習易受攻擊代碼的模式并提取漏洞特征.但是基于機器學習的檢測技術(shù)的性能仍然依賴于人類專家定義的模式的質(zhì)量,這意味著專業(yè)知識水平會對性能產(chǎn)生重大影響.

本文提出了一種基于BERT(bidirectional encoder representations from transformers)模型的源代碼漏洞檢測模型,將需要檢測的源代碼分割為一個個適合訓練的小樣本,將源代碼視為自然語言文本,通過BERT模型實現(xiàn)源代碼中漏洞特征的自動提取,并訓練具有良好性能的分類器實現(xiàn)多種漏洞的檢測.本文的創(chuàng)新點和主要貢獻如下:

1) 提出基于BERT模型的源代碼漏洞檢測模型,不需要依托專家知識即可實現(xiàn)對多種不同類型漏洞的檢測;

2) 提出一種代碼切片的方法實現(xiàn)數(shù)據(jù)的自動化預(yù)處理,便于模型將輸入轉(zhuǎn)換為所需的向量形式;

3) 在Python源代碼數(shù)據(jù)集上開展了大量的實驗和評估,實驗結(jié)果表明,本文模型在不同類型的漏洞中實現(xiàn)了平均99.2%的準確率、97.2%的精確率、96.2%的召回率和96.7%的F1分數(shù)的檢測水平.與現(xiàn)有檢測技術(shù)相比,本文模型有更高的準確率和更好的模型魯棒性.

1 相關(guān)工作

源代碼漏洞檢測的研究方法分為3類:基于代碼指標的漏洞檢測、基于機器學習的漏洞檢測和基于深度學習的漏洞檢測.不同的研究方法沒有絕對的區(qū)別,通常根據(jù)不同的應(yīng)用場景、數(shù)據(jù)類型和規(guī)模進行混用.

1.1 基于代碼指標的漏洞檢測研究

Morrison等人[1]提出使用一系列代碼指標構(gòu)建漏洞檢測模型,其中包括流失指標、復(fù)雜性指標、依賴性度量指標、規(guī)模指標等,這些指標已經(jīng)被廣泛運用于漏洞檢測模型中.

Shin等人[2]采用代碼復(fù)雜性指標對Mozilla應(yīng)用程序框架中JavaScript引擎的9個復(fù)雜性指標進行漏洞度量,使用2元邏輯回歸模型實現(xiàn)了80%召回率和25%誤報率的預(yù)測結(jié)果.

Chowdhury等人[3]在軟件早期開發(fā)階段運用了復(fù)雜性、耦合性和內(nèi)聚性相關(guān)的軟件指標自動預(yù)測Mozilla Firefox中容易出現(xiàn)漏洞的文件,利用決策樹、隨機森林、邏輯回歸和樸素貝葉斯等技術(shù),實現(xiàn)了整體預(yù)測準確率約為74%的預(yù)測效果.

Zhou等人[4]設(shè)計了K-fold(K折交叉切分)堆疊分類器,基于提交信息和錯誤報告中對某一問題的代碼修改,結(jié)合自然語言處理和機器學習技術(shù),實現(xiàn)了83%精確率和74%召回率的預(yù)測效果.

1.2 基于機器學習的漏洞檢測研究

機器學習技術(shù)已經(jīng)被用于許多編程語言的漏洞檢測中,通過分析有漏洞的代碼段即可定義漏洞模式用于漏洞檢測.

Pang等人[5]使用特征選擇和分析將整個Java類分類為易受攻擊或不易受攻擊,對4種不同類型的Java Android應(yīng)用程序進行預(yù)測,實現(xiàn)了約92%準確率、96%精確率和87%召回率的預(yù)測效果.

Hovsepyan等人[6]將Java Android項目的源代碼視為一種自然語言.他們使用詞袋技術(shù)將源代碼轉(zhuǎn)換為特征向量.然后使用支持向量機算法對這些特征進行2元分類訓練,達到了87%準確率、85%精確率以及88%召回率的預(yù)測效果.

1.3 基于深度學習的漏洞檢測研究

深度學習技術(shù)相比機器學習技術(shù)不需要手動定義漏洞模式,可以通過深度神經(jīng)網(wǎng)絡(luò)提取易受攻擊的特征.

Nguyen等人[7]將漏洞檢測視為歸納性的文本分類問題,通過將源代碼視為標記序列的方式構(gòu)建圖,使用預(yù)訓練編程語言模型的標記嵌入層進行初始化,然后利用圖神經(jīng)網(wǎng)絡(luò)進行訓練,獲得較好的分類效果.

陳傳濤等人[8]提取程序的抽象語法樹信息,并對抽象語法樹進行壓縮編碼,在保留程序結(jié)構(gòu)和語義特征的同時減少了語法樹的復(fù)雜性,有效提升漏洞檢測的效率.

Wartschinski等人[9]從Github中收集一個大型的Python源代碼提交數(shù)據(jù)集,通過詞嵌入模型將源代碼轉(zhuǎn)換為子向量,通過訓練長短期記憶網(wǎng)絡(luò)模型,實現(xiàn)了96.8%準確率、83%召回率、91%精確率和87%F1分數(shù)的預(yù)測效果.

1.4 存在的問題

基于代碼指標的漏洞檢測無法保留源代碼中的語義信息,限制了其在漏洞檢測方面的表現(xiàn),例如文獻[2-4].基于機器學習的漏洞檢測研究可以較好地提取出漏洞特征,但是需要大量的專家知識對漏洞模式進行定義,例如文獻[5].基于深度學習的漏洞檢測能自動地提取漏洞特征,然而其性能受限于數(shù)據(jù)集能否轉(zhuǎn)化成AST或者圖結(jié)構(gòu)等便于提取的模式,例如文獻[7-8].同時如文獻[6,9]表明,將源代碼視為自然語言似乎能夠提取更深層次的信息.

本文在以上研究的基礎(chǔ)上,提出基于BERT模型的漏洞檢測模型.通過將源代碼視為自然語言進行訓練,降低了數(shù)據(jù)集收集與處理的難度.本文通過BERT模型從代碼中自動提取漏洞語義特征,在保留源代碼語義信息的同時,也解決了提取特征需要大量專家知識的問題.本文還對模型進行優(yōu)化,使得檢測性能優(yōu)于現(xiàn)有研究的水平.

2 模型設(shè)計

2.1 BERT模型

BERT模型是Google團隊于2018年提出的一種預(yù)訓練自然語言處理模型.BERT模型通過在大規(guī)模的文本數(shù)據(jù)上進行無監(jiān)督的預(yù)訓練來學習通用的語言表示,并將這些通用的語言表示用于各種下游自然語言處理任務(wù),如文本分類、命名實體識別、語義相似度計算等.BERT模型的基本結(jié)構(gòu)如圖1所示.

圖1 BERT模型的基本結(jié)構(gòu)

BERT模型將輸入的文本轉(zhuǎn)化為3種向量:1)字向量.每個字在字向量表中對應(yīng)的1維向量.2)文本向量.該向量的取值在模型訓練過程中自動學習,用于刻畫文本的全局語義信息,并與單字/詞的語義信息相融合.3)位置向量.用于區(qū)分由于出現(xiàn)在文本不同位置的字/詞所攜帶的語義信息的差異.

BERT模型將字向量、文本向量和位置向量的加和作為輸入,通過多層Transformer編碼器[10]架構(gòu)進行漏洞語義特征的提取.

Transformer編碼器結(jié)合多頭注意力機制、前饋神經(jīng)網(wǎng)絡(luò)與殘差連接等技術(shù),通過雙向訓練及多層Transformer編碼器堆疊的方式,可以準確地捕捉文本中長距離的依賴關(guān)系和上下文信息,進而生成高質(zhì)量的表示向量.

2.2 基于BERT模型的源代碼漏洞檢測模型

本文在BERT模型的基礎(chǔ)上,設(shè)計了用于源代碼漏洞檢測的總體模型,模型主要分為3個部分:數(shù)據(jù)預(yù)處理、訓練模型及分類器.如圖2所示.

圖2 基于BERT的源代碼漏洞檢測總體模型

訓練一個高效的漏洞檢測模型需要大量的源代碼作為數(shù)據(jù)集.數(shù)據(jù)預(yù)處理部分對從網(wǎng)上收集的大量源代碼進行切片操作,切割后的源代碼成為一個個小樣本,每個小樣本都會打上有無漏洞的標簽.將一個小樣本與它對應(yīng)標簽作為模型訓練的輸入,從而形成訓練數(shù)據(jù)集.

訓練模型主要由BERT模型、隨機失活層、線性層與Sigmoid激活函數(shù)組成.

1) BERT模型.基本架構(gòu)是12層Transformer編碼器結(jié)構(gòu),采用掩碼語言模型[11]的方法進行預(yù)訓練,可用于提取代碼中的語義信息.

2) 隨機失活層.隨機地將一些神經(jīng)元的輸出設(shè)置為0,防止模型在訓練時出現(xiàn)過擬合的現(xiàn)象,隨機失活層參數(shù)設(shè)置為0.3,也就是有30%左右的神經(jīng)元會失活.

3) 線性層.用于將所有提取出的特征映射到一個維度為1的張量上,這個張量表示提取出來的整體特征.

4) Sigmoid激活函數(shù).將提取出的輸出映射到介于0和1之間的值,也就是輸入樣本是否存在漏洞的概率.

分類器根據(jù)輸入的概率判斷樣本是否存在漏洞.本文設(shè)置分類器閾值為0.5,即認為輸出概率大于等于50%的樣本是存在漏洞的,而輸出概率小于50%的樣本不存在漏洞.

本文通過數(shù)據(jù)預(yù)處理將源代碼處理為近似自然語言的形式,使用本文提出的模型對源代碼進行學習和特征提取,訓練出高效的源代碼漏洞檢測分類器.

3 實驗與結(jié)果分析

3.1 數(shù)據(jù)預(yù)處理

本文采用的代碼集來自于Wartschinski等人[9]的研究,他們從Github上收集了大量高分Python項目修復(fù)漏洞的補丁信息,補丁中記錄了修補漏洞的類型以及新舊版本的源代碼差異.

代碼集含有跨站腳本攻擊、SQL注入、跨站請求偽造、遠程代碼執(zhí)行、路徑泄露、開放重定向和命令注入7種漏洞修復(fù)的補丁信息集.本文以SQL注入漏洞為例展示本文所用數(shù)據(jù)集的生成過程,如圖3所示.

圖3 SQL注入漏洞代碼片段

代碼片段來自于SQL漏洞的修復(fù)補丁信息,紅色代碼片段是舊代碼相比于新代碼被修改的片段.由于這次補丁提交修改的目的是修復(fù)SQL漏洞,所以可以認為被修改的舊代碼片段是存在SQL注入漏洞的.據(jù)此,本文進行以下操作生成訓練與驗證的數(shù)據(jù)集:

1) 將代碼集按空格、“ ”“ ”等間隔符號分割成一定長度可訓練的小樣本;

2) 將代碼中多余的空格以及“ ”“ ”等間隔符號刪除,并將所有字母轉(zhuǎn)化為小寫,讓代碼更接近自然語言;

3) 根據(jù)修復(fù)補丁的信息,給每個小樣本打上標簽,“1”代表存在漏洞,“0”代表不存在漏洞.

以上述代碼為例,可以生成如圖4的數(shù)據(jù)集:

圖4 生成數(shù)據(jù)集示例

按上述方法,本文以8∶2的比例生成所用的7種漏洞的訓練集和驗證集,如表1所示:

表1 7種漏洞的數(shù)據(jù)集和驗證集

3.2 評判指標

首先定義公式中所用到的符號:1)真陽性TP(true positive).被模型預(yù)測為正類的正樣本.2)真陰性TN(true negative).被模型預(yù)測為負類的負樣本.3)假陽性FP(false positive).被模型預(yù)測為正類的負樣本.4)假陰性FN(false negative).被模型預(yù)測為負類的正樣本,如表2所示:

表2 符號定義

為了衡量模型的性能,了解模型在檢測任務(wù)中的表現(xiàn),本文引入準確率、精確率、召回率以及F1分數(shù)4項衡量指標:

1) 準確率指模型預(yù)測正確的樣本數(shù)量占總樣本數(shù)量的比例,公式為

2) 精確率指所有被模型預(yù)測為正例的樣本中真正為正例的樣本占比,公式為

3) 召回率指所有真正為正例的樣本中,被模型正確預(yù)測為正例的樣本占比,公式為

4)F1分數(shù)是精確率和召回率的加權(quán)調(diào)和平均數(shù).它的值越接近1,表示模型性能越好,公式為

3.3 實驗與結(jié)果分析

本文實驗所使用的環(huán)境配置如下:CPU為AMD Ryzen9 3950X 3.5GHz,GPU為NVIDIA GeForce GTX 3090,內(nèi)存為32GB,操作系統(tǒng)為Ubuntu20.04.模型框架搭建是基于開源框架Pytorch.

模型的輸入長度設(shè)置為128B,學習率設(shè)置為1E-05,訓練輪數(shù)設(shè)置為10輪.通過小批量的方式進行訓練與驗證,訓練時批次大小設(shè)置為8,驗證時批次大小設(shè)置為4.

由于目前基于BERT模型進行漏洞檢測任務(wù)的相關(guān)研究較少,為了探究哪種優(yōu)化算法更適合漏洞檢測任務(wù),本文對比了Adadelta(自適應(yīng)增量學習率算法)[12]、Adagrad(自適應(yīng)梯度算法)[13]、Adam(自適應(yīng)矩估計算法)[14]、AdamW(帶權(quán)重衰減的自適應(yīng)矩估計算法)[15]、SGD(隨機梯度下降算法)[16]和RMSPROP(均方根傳遞算法)的優(yōu)化效率,結(jié)果如圖5所示:

圖5 本文模型采用不同優(yōu)化器的準確率

從結(jié)果可以看出,Adadelta算法對本文模型幾乎不起作用,SGD和Adagrad算法優(yōu)化效率較慢,而Adam,RMSPROP,AdamW算法模型收斂較快,在準確率性能方面也很接近.本文接著從F1分數(shù)指標進一步對比Adam,RMSPROP,AdamW算法,如圖6所示:

圖6 本文模型采用不同優(yōu)化器的F1分數(shù)

在F1指標中3種算法的表現(xiàn)也很接近,其中AdamW算法相比其他2個算法有微小的領(lǐng)先,于是本文采用AdamW算法作為模型訓練時的優(yōu)化算法.本文針對7種漏洞分別訓練了分類模型,每個模型取10輪中驗證效果最好的1輪數(shù)據(jù),結(jié)果如表3所示:

表3 對7種漏洞檢測結(jié)果 %

本文模型對7種漏洞進行了檢測,平均準確率指標為99.2%,平均精確率指標為97.2%,平均召回率指標為96.2%,平均F1分數(shù)指標為96.7%.本文模型在漏洞檢測任務(wù)上取得了非常好的效果.

對比文獻[9]中使用長短期記憶網(wǎng)絡(luò)和詞嵌入進行漏洞檢測模型訓練方法,在準確率、精確率、召回率和F1分數(shù)指標上分別有2.4%,5.8%,14%,9.6%的提升,如圖7～10所示:

圖7 7種類型漏洞檢測準確率指標對比

圖8 7種類型漏洞檢測精確率指標對比

圖9 7種類型漏洞檢測召回率指標對比

4 結(jié) 語

本文提出了一種基于BERT模型的源代碼漏洞檢測模型.運用自然語言處理技術(shù),將源代碼視為自然語言文本的方式進行訓練,在跨站腳本攻擊、SQL注入、跨站請求偽造、遠程代碼執(zhí)行、路徑泄露、開放重定向和命令注入7種漏洞檢測中,本文模型較文獻[9]實現(xiàn)了全面提升.

本文提出的檢測模型適用于多種語言,未來的工作可以將模型推廣到Java,C,C++等源代碼的漏洞檢測.也可以考慮將源代碼轉(zhuǎn)換為中間表示的形式,例如抽象語法樹或數(shù)據(jù)流圖等,突出主要語義信息以提高整體性能.