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

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

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

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

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

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

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

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

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

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

1 相關工作

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

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

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

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

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

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

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

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

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

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

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

深度學習技術相比機器學習技術不需要手動定義漏洞模式,可以通過深度神經網絡提取易受攻擊的特征.

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

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

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

1.4 存在的問題

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

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

2 模型設計

2.1 BERT模型

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

圖1 BERT模型的基本結構

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

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

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

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

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

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

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

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

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

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

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

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

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

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

3 實驗與結果分析

3.1 數據預處理

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

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

圖3 SQL注入漏洞代碼片段

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

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

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

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

以上述代碼為例,可以生成如圖4的數據集:

圖4 生成數據集示例

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

表1 7種漏洞的數據集和驗證集

3.2 評判指標

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

表2 符號定義

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

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

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

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

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

3.3 實驗與結果分析

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

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

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

圖5 本文模型采用不同優化器的準確率

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

圖6 本文模型采用不同優化器的F1分數

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

表3 對7種漏洞檢測結果 %

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

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

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

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

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

4 結 語

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

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