Code Vulnerability Repair with Large Language Model using Context-Aware Prompt Tuning

📄 arXiv: 2409.18395v2 📥 PDF

作者: Arshiya Khan, Guannan Liu, Xing Gao

分类: cs.CR, cs.AI

发布日期: 2024-09-27 (更新: 2025-06-11)

期刊: in 2025 IEEE Security and Privacy Workshops (SPW), San Francisco, CA, USA, 2025, pp. 283-287,

DOI: 10.1109/SPW67851.2025.00040

💡 一句话要点

提出上下文感知Prompt Tuning，提升LLM在代码漏洞修复中的性能

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 代码漏洞修复 大型语言模型 Prompt Tuning 缓冲区溢出 上下文感知

📋 核心要点

  1. 现有LLM在修复复杂代码漏洞（如缓冲区溢出）时，修复成功率较低，存在显著的性能瓶颈。
  2. 通过上下文感知的Prompt Tuning，向LLM注入漏洞相关的领域知识，引导其进行更准确的漏洞修复。
  3. 实验结果表明，该方法能显著提升LLM在缓冲区溢出漏洞修复上的成功率，提升幅度超过四倍。

📝 摘要（中文）

大型语言模型(LLM)在检测和修复代码漏洞方面面临显著挑战，尤其是在处理涉及变量、代码流和代码结构等多个方面的漏洞时。本研究使用GitHub Copilot作为LLM，并专注于缓冲区溢出漏洞。实验表明，Copilot在处理缓冲区溢出漏洞时存在明显不足，漏洞检测率达到76%，但漏洞修复率仅为15%。为了解决这个问题，我们提出了上下文感知prompt tuning技术，旨在提高LLM在修复缓冲区溢出方面的性能。通过注入一系列关于漏洞的领域知识，包括各种安全和代码上下文，我们证明Copilot的成功修复率提高到63%，与没有领域知识的修复相比，提高了四倍以上。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）在代码漏洞修复，特别是缓冲区溢出漏洞修复方面表现不佳的问题。现有的LLM在处理此类漏洞时，往往难以充分理解代码的上下文信息，导致修复成功率较低，无法满足实际应用的需求。

核心思路：论文的核心思路是通过上下文感知的Prompt Tuning，为LLM提供更丰富的领域知识和上下文信息，从而引导LLM更好地理解漏洞的本质，并生成更准确的修复代码。这种方法旨在弥补LLM在处理复杂代码漏洞时对上下文理解不足的缺陷。

技术框架：该方法主要包括以下几个步骤：首先，确定需要修复的缓冲区溢出漏洞代码；然后，构建包含漏洞相关领域知识和上下文信息的Prompt；接着，将Prompt输入到LLM（本研究中使用GitHub Copilot）；最后，评估LLM生成的修复代码的正确性。关键在于Prompt的设计，需要包含足够的上下文信息，例如变量的定义、代码的执行流程、以及缓冲区溢出的常见原因等。

关键创新：该方法最重要的创新点在于上下文感知的Prompt Tuning。与传统的Prompt Tuning方法不同，该方法更加注重Prompt中领域知识和上下文信息的构建，使得LLM能够更好地理解漏洞的本质，从而生成更准确的修复代码。这种方法能够显著提高LLM在复杂代码漏洞修复方面的性能。

关键设计：Prompt的设计是该方法成功的关键。Prompt中需要包含以下信息：漏洞代码片段、漏洞类型的描述（例如，缓冲区溢出）、漏洞可能产生的原因、以及修复漏洞的建议。此外，Prompt还可以包含一些示例代码，展示如何正确地使用缓冲区，避免溢出。论文中没有明确说明损失函数或网络结构，因为该方法主要依赖于LLM自身的预训练能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，通过上下文感知的Prompt Tuning，GitHub Copilot在缓冲区溢出漏洞修复上的成功率从15%提高到63%，提升幅度超过四倍。这表明该方法能够显著提高LLM在代码漏洞修复方面的性能，具有重要的实际应用价值。

🎯 应用场景

该研究成果可应用于自动化代码安全检测与修复，帮助开发者快速定位并修复代码中的漏洞，提高软件的安全性。该方法可以集成到IDE或CI/CD流程中，实现代码的自动安全审查。未来，该方法可以扩展到其他类型的代码漏洞修复，并应用于更复杂的软件系统。

📄 摘要（原文）

Large Language Models (LLMs) have shown significant challenges in detecting and repairing vulnerable code, particularly when dealing with vulnerabilities involving multiple aspects, such as variables, code flows, and code structures. In this study, we utilize GitHub Copilot as the LLM and focus on buffer overflow vulnerabilities. Our experiments reveal a notable gap in Copilot's abilities when dealing with buffer overflow vulnerabilities, with a 76% vulnerability detection rate but only a 15% vulnerability repair rate. To address this issue, we propose context-aware prompt tuning techniques designed to enhance LLM performance in repairing buffer overflow. By injecting a sequence of domain knowledge about the vulnerability, including various security and code contexts, we demonstrate that Copilot's successful repair rate increases to 63%, representing more than four times the improvement compared to repairs without domain knowledge.