核心概念
DFEPT 是一種新的資料流嵌入技術,旨在增強預訓練模型在漏洞檢測任務中的效能,通過解析程式碼中的資料流圖並將其嵌入到預訓練模型中,DFEPT 能夠有效地識別程式碼中的漏洞。
研究目標:
本研究旨在解決現有基於深度學習的漏洞檢測方法在實際應用中存在的局限性,提出了一種名為 DFEPT 的資料流嵌入技術,以增強預訓練模型在漏洞檢測任務中的效能。
方法:
DFEPT 的核心思想是將程式碼的資料流圖嵌入到預訓練模型中,具體步驟如下:
資料流圖構建: 利用程式碼解析工具(如 tree-sitter)從函數級別的原始碼中解析出抽象語法樹(AST),並根據 AST 中變數之間的依賴關係構建資料流圖(DFG)。
資料流嵌入: 使用變數的資料類型作為節點特徵,並採用圖神經網路(GNN)對 DFG 進行嵌入,從而提取程式碼的結構和語義資訊。
正弦位置編碼: 為了保留程式碼中節點的相對位置資訊,將正弦位置編碼應用於嵌入後的圖特徵向量。
程式碼序列嵌入: 使用預訓練模型(如 CodeBERT、GraphCodeBERT、CodeT5 或 UniXcoder)將程式碼片段視為自然語言序列,並生成表示整個程式碼塊語義的句子向量。
分類與檢測: 將資料流嵌入和程式碼序列嵌入串聯起來,輸入到一個多層感知器(MLP)中進行分類和檢測。
主要發現:
在 Devign 和 Reveal 兩個公開漏洞資料集上的實驗結果表明,DFEPT 能夠有效地提高多種預訓練模型在漏洞檢測任務中的效能,特別是與 CodeT5 結合使用時,DFEPT 在兩個資料集上均取得了最佳的 F1 分數。
主要結論:
DFEPT 是一種有效的漏洞檢測方法,它能夠克服現有基於深度學習的方法的局限性,並取得更佳的效能。
意義:
本研究提出了一種新穎的資料流嵌入技術,為基於深度學習的漏洞檢測方法提供了一種新的思路,並為開發更準確、更可靠的漏洞檢測工具奠定了基礎。
局限性和未來研究方向:
DFEPT 的效能受限於預訓練模型本身和嵌入層的選擇。
未來可以探索更有效的圖嵌入和池化方法,以進一步提高 DFEPT 的效能。
可以將 DFEPT 應用於其他程式語言和漏洞類型。
統計資料
DFEPT 在 Devign 資料集上達到了 64.53% 的準確率和 61.22% 的 F1 分數。
DFEPT 在 Reveal 資料集上達到了 92.92% 的準確率和 47.9% 的 F1 分數。