大模型漏洞修復(fù)插件是騰訊朱雀實驗室在安全垂類場景的一個重要實踐。我們希望通過AI大模型，實現(xiàn)研發(fā)安全場景的漏洞自動修復(fù)，給出修復(fù)建議并提供修復(fù)代碼，幫助更多開發(fā)人員提高研發(fā)效率。在騰訊混元大模型的支持下，漏洞修復(fù)插件通過精調(diào)后部署的私有化模型，實現(xiàn)了在帳密硬編碼、SQL注入、命令注入等漏洞類型的修復(fù)建議輸出和修復(fù)代碼生成等功能，實現(xiàn)安全左移，更有效地在編程中使用插件收斂漏洞風(fēng)險。

圖1. IDE插件示意圖

一、研發(fā)安全場景的現(xiàn)狀和挑戰(zhàn)

利用傳統(tǒng)方法做漏洞修復(fù)提效，只適用于比較簡單的場景，比如根據(jù)版本號判斷使用的開源組件是否存在漏洞，更多高危險的如導(dǎo)致數(shù)據(jù)泄露的注入類漏洞/賬密類等，該方案難以通用。主要原因總結(jié)如下：

• 規(guī)則限制：傳統(tǒng)靜態(tài)分析技術(shù)通常基于預(yù)定義的規(guī)則和模式進(jìn)行漏洞檢測和修復(fù)，無法覆蓋所有的漏洞類型和場景;
• 上下文和語義理解限制：傳統(tǒng)靜態(tài)分析技術(shù)通常難以理解代碼的上下文和語義信息，導(dǎo)致無法準(zhǔn)確地理解代碼的含義和邏輯;
• 創(chuàng)造性限制：傳統(tǒng)靜態(tài)分析技術(shù)通常只能分析已有的代碼，無法創(chuàng)造新的代碼片段來修復(fù)漏洞，限制了漏洞修能力。

相比傳統(tǒng)程序分析技術(shù)，大模型具備強(qiáng)大的推理能力，尤其是在代碼生成方面表現(xiàn)突出，可通過訓(xùn)練來學(xué)習(xí)漏洞修復(fù)的模式和規(guī)律。

二、為何在研發(fā)安全場景引入大模型?

1. 大模型可解釋性

大模型擁有豐富的預(yù)訓(xùn)練語料信息，如：書籍文檔、代碼素材、網(wǎng)頁文本等等?；陬A(yù)訓(xùn)練語料信息，并對基座大模型進(jìn)行垂類領(lǐng)域的精調(diào)，大模型能結(jié)合知識信息進(jìn)行整合和解釋，產(chǎn)生有效的回復(fù)。對于研發(fā)安全場景來說，我們不僅需要生成有效的修復(fù)代碼，也希望為更多的開發(fā)同學(xué)提供具體和明確的修復(fù)分析指引，從而更好地提效。

2. 大模型代碼生成能力

各類基座大模型和代碼類Codex大模型，具有強(qiáng)大的代碼生成能力。這些模型可以理解自然語言指令，并根據(jù)這些指令生成相應(yīng)的代碼。這種能力使得它們可以用于各種編程任務(wù)，包括但不限于編寫新的函數(shù)、修復(fù)代碼中的錯誤、優(yōu)化現(xiàn)有代碼等。這些模型的代碼生成能力基于其在大量代碼庫上的訓(xùn)練。在訓(xùn)練過程中，模型學(xué)習(xí)了各種編程語言的語法和語義以及如何將自然語言指令轉(zhuǎn)化為代碼。因此，只要給出清晰的指令，這些模型就能生成相應(yīng)的代碼。但這些生成的代碼中仍存在一些安全問題，即：潛在的不正確、含漏洞的代碼等。因此，我們需要利用大模型的代碼生成能力，同時希望通過精調(diào)等操作，使大模型能有效修復(fù)漏洞代碼，盡量生成安全無漏洞的代碼片段。

3. 大模型在安全防護(hù)領(lǐng)域的應(yīng)用已成為趨勢

2023年，大模型成為了各行各業(yè)的熱門關(guān)注點，其中安全垂類領(lǐng)域的安全大模型也呈現(xiàn)了百花齊放的態(tài)勢，主要覆蓋安全咨詢、安全培訓(xùn)、安全監(jiān)控、安全修復(fù)等能力。微軟2023年3月份正式發(fā)布集成GPT-4的Microsoft Security Copilot，旨在更好地提供安全工具和專家知識，輔助企業(yè)識別和檢測安全風(fēng)險。谷歌2023年4月發(fā)布“谷歌云AI安全工作臺”，利用 Sec-PaLM 來幫助用戶查找、總結(jié)和應(yīng)對安全問題。國內(nèi)公司也推出了自己獨有的安全大模型，主要對于安全場景，提供專家級別的咨詢建議，提效運營效率。

三、混元一站式如何快速定制司內(nèi)自研的研發(fā)安全大模型

1. 騰訊混元大模型能力

(1) 訓(xùn)練優(yōu)化

騰訊混元大模型通過大量實驗，對預(yù)訓(xùn)練數(shù)據(jù)進(jìn)行了語料豐富擴(kuò)充，已覆蓋 100 多種自然語言，32 種編程語言;并經(jīng)過了數(shù)據(jù)清洗、過濾、去重等流程，對各類數(shù)據(jù)進(jìn)行了大量的數(shù)據(jù)配比實驗，保證了較為平穩(wěn)的訓(xùn)練過程。同時，項目組對于長文能力、位置編碼等技術(shù)細(xì)節(jié)進(jìn)行了改進(jìn)和優(yōu)化。

(2) 模型效果

騰訊混元項目組參考業(yè)界主流做法，一共選用數(shù)十個數(shù)據(jù)集合，分為中文NLP、英文NLP任務(wù)、代碼、數(shù)學(xué)、AGIEval、CMMLU、CEval等維度，綜合評估模型在各項維度上的能力。其中，騰訊混元十億級別大模型在代碼能力上的提升更加顯著，這點與我們使用中的精調(diào)體驗也非常一致。

2. 漏洞相關(guān)高質(zhì)量數(shù)據(jù)的收集及模型結(jié)果

通用大模型的崛起，依賴于大規(guī)模的數(shù)據(jù)灌輸;大模型能否在垂直領(lǐng)域落地，同樣依賴是否有高質(zhì)量的領(lǐng)域數(shù)據(jù)。對于研發(fā)安全場景，我們也積累了系統(tǒng)且豐富的高質(zhì)量數(shù)據(jù)處理經(jīng)驗。

(1) 漏洞修復(fù)數(shù)據(jù)的格式

漏洞修復(fù)數(shù)據(jù)，最基礎(chǔ)的是漏洞類型、漏洞代碼、修復(fù)后代碼三個屬性。

圖2. JavaScript SQL注入漏洞修復(fù)樣例

但在實踐過程中僅使用這三部分信息來精調(diào)大模型，效果并不好。為此，我們擴(kuò)至漏洞類型、漏洞代碼、修復(fù)后代碼、代碼描述、漏洞信息、修復(fù)建議、修復(fù)過程七個屬性，組成“漏洞代碼 -> 代碼做了什么 -> 存在什么漏洞 -> 應(yīng)該怎么修復(fù) -> 修復(fù)后代碼 -> 具體修復(fù)過程”的整體邏輯，實際訓(xùn)練效果更好。

圖3. Java SQL注入漏洞修復(fù)示例

(2) 漏洞修復(fù)數(shù)據(jù)源及處理方式

我們采用的漏洞修復(fù)數(shù)據(jù)源，包含GitHub數(shù)據(jù)、公開數(shù)據(jù)、業(yè)務(wù)數(shù)據(jù)等，數(shù)據(jù)質(zhì)量依次上升。下面我們大致介紹各數(shù)據(jù)源的潛在問題和處理過程。

①GitHub數(shù)據(jù)

Github是主流開源社區(qū)之一，其歷史數(shù)據(jù)中含有大量的漏洞代碼及修復(fù)的記錄，但不同開發(fā)者的編程水平參差不齊。我們根據(jù)關(guān)鍵詞召回與漏洞相關(guān)的數(shù)據(jù)，但實踐中發(fā)現(xiàn)大部分(>90%)召回數(shù)據(jù)存在與漏洞修復(fù)無關(guān)、錯誤的修復(fù)方式、代碼質(zhì)量低等等問題，因此我們需要一些自動化的方法來過濾這些數(shù)據(jù)。

缺點：

• 不可用數(shù)據(jù)的占比較高，人工審核成本上升;
• 單個樣本可能存在代碼塊分散、缺失上下文信息等問題;

優(yōu)點：

• 數(shù)據(jù)量級較大;
• 代碼多樣性豐富;
• 可按需召回指定類型的漏洞;

②公開數(shù)據(jù)

公開數(shù)據(jù)，更多指的是學(xué)術(shù)界開源的數(shù)據(jù)。這類數(shù)據(jù)，已經(jīng)被大量學(xué)術(shù)工作借鑒使用，本身具備一定的可信度。目前，我們使用的主要是CVE-Fixes數(shù)據(jù)集。

圖4. CVE-Fixes Top-5漏洞與SQL漏洞占比圖

缺點：

• 數(shù)據(jù)量級不大;
• 漏洞類型分布與實際業(yè)務(wù)漏洞類型分布差距大;

優(yōu)點：

• 數(shù)據(jù)質(zhì)量有一定保障;
• 代碼多樣性豐富;

③業(yè)務(wù)數(shù)據(jù)

業(yè)務(wù)數(shù)據(jù)，即收集公司業(yè)務(wù)歷史修復(fù)數(shù)據(jù)，清洗后作為訓(xùn)練數(shù)據(jù)。業(yè)務(wù)數(shù)據(jù)優(yōu)點非常多，最貼近模型使用的場景、含有更符合公司業(yè)務(wù)場景的修復(fù)方式。但業(yè)務(wù)數(shù)據(jù)與GitHub數(shù)據(jù)有相同的缺點，清洗難度較大。

缺點：

• 歷史誤報、錯誤修復(fù)的數(shù)據(jù)占比較高，人工審核壓力大;
• 符合實際訓(xùn)練需求的高質(zhì)量數(shù)據(jù)量級有限;

優(yōu)點：

• 代碼多樣性豐富;
• 與落地場景完全契合;

(3) 訓(xùn)練數(shù)據(jù)構(gòu)造

基于上文獲取的漏洞數(shù)據(jù)，我們進(jìn)行了大量的實驗?？偨Y(jié)得到，在數(shù)據(jù)總量和質(zhì)量不變的情況下，“Prompt工程”和“數(shù)據(jù)配比”對實驗結(jié)果會產(chǎn)生影響。同時，我們指出：不同垂直領(lǐng)域下得到的實驗結(jié)論大概率不同，例如我們在數(shù)據(jù)配比時配比了部分通用數(shù)據(jù)和代碼數(shù)據(jù)，而其他團(tuán)隊可能只是用垂直領(lǐng)域的數(shù)據(jù)。所以，我們建議大家在聚焦的垂直領(lǐng)域多實驗、再總結(jié)。

(4) Prompt 工程

通過調(diào)整對模型提問的方式，使得模型的回答效果更好。具體的，我們通過業(yè)界的論文以及經(jīng)驗分享，總結(jié)了一個相對通用的模板：expert+COT+輸出格式+變量。如圖5所示。其中，expert是一種增強(qiáng)提示策略，用于指示LLM像專家一樣回答問題。COT，即思維鏈，通過增加一系列中間推理步驟，能顯著提高大型語言模型復(fù)雜推理的能力。

圖5. prompt構(gòu)造

(5) 數(shù)據(jù)配比

除了漏洞數(shù)據(jù)外，我們還收集了開源的通用、代碼、數(shù)學(xué)sft數(shù)據(jù)集，并進(jìn)行了相應(yīng)的配比實驗。我們發(fā)現(xiàn)，數(shù)學(xué)數(shù)據(jù)集對漏洞修復(fù)準(zhǔn)確率的提升基本無影響，而適當(dāng)比例的代碼、通用數(shù)據(jù)，有助于提升模型的漏洞修復(fù)準(zhǔn)確率。

經(jīng)過調(diào)優(yōu)，我們的最佳數(shù)據(jù)配比如圖6所示。當(dāng)我們代碼：通用：漏洞=1:2:4的時候，結(jié)果達(dá)到最優(yōu)。

圖6. 最佳數(shù)據(jù)配比圖

我們將開源的通用(general)、代碼(code)、數(shù)學(xué)(math)、漏洞(vulnerability)數(shù)據(jù)集，輸入到模型中，抽取倒數(shù)第二層對應(yīng)的embedding，使用t-sne算法對其進(jìn)行可視化，得到的結(jié)果如圖7所示?？梢钥吹綌?shù)學(xué)數(shù)據(jù)集距離漏洞數(shù)據(jù)集最遠(yuǎn)，這也解釋了為什么數(shù)學(xué)數(shù)據(jù)集對漏洞修復(fù)準(zhǔn)確率的提升基本無影響。

圖7. t-sne結(jié)果圖

(6) 實驗結(jié)果

我們將實驗結(jié)果與gpt3.5進(jìn)行比較，結(jié)果如表1所示。從表1中可以看出，我們的平均準(zhǔn)確率比gpt3.5高1.67%，具體到每個漏洞類型，均超過或持平gpt3.5。

表1 自建模型與gpt3.5結(jié)果對比表

(7) 上下文長度對結(jié)果的影響

在前文介紹處理業(yè)務(wù)數(shù)據(jù)時，我們提出借助于污點傳播技術(shù)，只摘取漏洞觸發(fā)點所在行的上下文代碼。那么，這個上下文取多少行合適呢?這里我們在推理階段取不同大小的上下文進(jìn)行了測試，如圖8所示。從結(jié)果中可以看出，不同大小的上下文確實對結(jié)果有影響，并且對于不同漏洞，影響不太相同。

對于SQL注入以及命令注入漏洞來說，適中大小的上下文有助于提升測試準(zhǔn)確率。當(dāng)上下文過少時，模型抽取的與漏洞相關(guān)的語義信息不足，通過這些信息無法判斷是否存在漏洞，導(dǎo)致準(zhǔn)確率下降;相反，當(dāng)上下文過多時，模型定位漏洞所在行可能存在困難，導(dǎo)致準(zhǔn)確率下降。

對于帳密硬編碼漏洞來說，測試準(zhǔn)確率隨著上下文行數(shù)的增多而減少，猜測可能是因為對于該漏洞類型，較少的行數(shù)就足以讓模型判斷存在帳密硬編碼，也足以讓模型進(jìn)行修復(fù)。相反，當(dāng)測試上下文過多時，模型定位漏洞所在行可能存在困難，導(dǎo)致準(zhǔn)確率下降。

圖8. 不同漏洞類型修復(fù)準(zhǔn)確率隨上下文長度的變換曲線圖

3. 插件支持漏洞檢測和修復(fù)

目前我們基于騰訊混元大模型，推出了漏洞檢出和修復(fù)功能的插件。插件目前重點發(fā)現(xiàn)三大類研發(fā)過程中對現(xiàn)網(wǎng)產(chǎn)生嚴(yán)重影響的漏洞：SQL注入、命令注入和賬密泄漏，并較業(yè)內(nèi)大幅提升發(fā)現(xiàn)風(fēng)險的準(zhǔn)確性。目前插件已支持在vs code內(nèi)使用，僅需在vs code內(nèi)安裝插件，即可實現(xiàn)在提交代碼前檢測安全風(fēng)險并及時修復(fù)。

(1) 插件的漏洞檢測+修復(fù)示例

圖9. 插件的漏洞檢測+修復(fù)示例圖

(2) 與競品模型回復(fù)的對比

我們將自建模型與gpt3.5模型的回復(fù)進(jìn)行對比，其中一個典型案例如下。gpt3.5與自建模型均能判斷出代碼中存在SQL注入漏洞，但是gpt3.5無法對其進(jìn)行正確修復(fù)，而自建模型可以。

圖10. 自建模型與gpt3.5回復(fù)對比

4. 插件模型版本的自動化測評和迭代

模型評測迭代的過程中，針對漏洞修復(fù)，我們完全自建了benchmark以及對應(yīng)的評判規(guī)則：

• 待評測數(shù)據(jù)通過自研模型，拿到修復(fù)結(jié)果。
• 將修復(fù)前后的代碼，輸入到代碼相似度模型中，防止模型吐出與原始代碼不相關(guān)的代碼。
• 使用算法/工具判斷脫敏后的漏洞代碼是否修復(fù)。
• 根據(jù)實際測試，我們發(fā)現(xiàn)以上流程無法對修復(fù)前后代碼語義是否更改進(jìn)行判斷。如圖11所示，修復(fù)后代碼漏掉了=號，導(dǎo)致語義的改變。因此我們補(bǔ)充設(shè)計了一套規(guī)則，來判斷修復(fù)前后代碼語義是否有更改。

圖11. 修復(fù)前后代碼語義改變示意圖

四、總結(jié)與展望

以研發(fā)安全場景為切入點，擁有均衡能力的騰訊混元大模型未來在安全領(lǐng)域會有更廣泛的應(yīng)用場景和前景，幫助安全業(yè)務(wù)提高安全防護(hù)的效率，可以降低業(yè)務(wù)的投入成本，提高經(jīng)濟(jì)效益。

• 在數(shù)據(jù)安全領(lǐng)域，大模型可以通過對大量數(shù)據(jù)的分析，發(fā)現(xiàn)數(shù)據(jù)中的異常行為，及時預(yù)警并采取措施，防止數(shù)據(jù)泄露、數(shù)據(jù)篡改等安全問題的發(fā)生。
• 在網(wǎng)絡(luò)安全領(lǐng)域，大模型可以通過對網(wǎng)絡(luò)流量的實時監(jiān)控，發(fā)現(xiàn)并阻止網(wǎng)絡(luò)攻擊，保障網(wǎng)絡(luò)的正常運行。
• 在安全運營領(lǐng)域，大模型可以通過對用戶數(shù)據(jù)、設(shè)備狀態(tài)數(shù)據(jù)的持續(xù)監(jiān)控，發(fā)現(xiàn)并預(yù)防各種安全風(fēng)險，提高安全運營的效率。
• 同時，大模型還可以通過對安全審計的支持，幫助企業(yè)更好地遵守各種安全法規(guī)，避免因違規(guī)操作而帶來的法律風(fēng)險。

總的來說，騰訊混元大模型的出現(xiàn)，將會在安全領(lǐng)域開辟出更廣闊的應(yīng)用場景，充分利用好公司內(nèi)已有的安全數(shù)據(jù)，做好業(yè)務(wù)數(shù)據(jù)處理和沉淀，將大大提高安全水位和業(yè)務(wù)效率。

本文轉(zhuǎn)載自騰訊技術(shù)工程，作者：騰訊程序員

原文李鏈接：???https://mp.weixin.qq.com/s/KwyuQPmInzXwqWjV46OmhQ????