PromSec: Prompt Optimization for Secure Generation of Functional Source Code with Large Language Models (LLMs)

📄 arXiv: 2409.12699v1 📥 PDF

作者: Mahmoud Nazzal, Issa Khalil, Abdallah Khreishah, NhatHai Phan

分类: cs.SE, cs.LG

发布日期: 2024-09-19

备注: 15 pages, 19 figures, CCS 2024

DOI: 10.1145/3658644.3690298

💡 一句话要点

PromSec：通过提示优化，安全生成大语言模型的功能性源代码

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 代码生成 大型语言模型 安全漏洞 提示优化 图神经网络 生成对抗网络 对比学习

📋 核心要点

  1. 现有方法生成的代码存在安全漏洞，源于LLM训练数据中的不安全代码，缺乏有效保障代码安全性的机制。
  2. PromSec通过结合生成对抗图神经网络（gGAN）的代码漏洞清除和LLM的代码生成，形成交互循环，优化提示以生成安全代码。
  3. 实验表明，PromSec能有效解决代码漏洞，显著降低操作时间、LLM查询次数和安全分析成本，并具备跨语言和漏洞的泛化能力。

📝 摘要（中文）

利用大型语言模型（LLM）生成高质量源代码能够缩短软件开发时间和成本。然而，由于LLM使用不安全的开源数据进行训练，生成的代码通常包含安全漏洞。因此，确保代码生成过程的安全性和功能性至关重要。本文提出了一种名为PromSec的算法，该算法通过提示优化，利用LLM生成安全且功能正常的代码。PromSec结合了1) 使用生成对抗图神经网络（gGAN）进行代码漏洞清除，以修复和减少生成代码中的安全漏洞；2) 使用LLM进行代码生成，并将二者置于一个交互循环中，gGAN的输出结果通过增强的提示驱动LLM生成更安全的代码，同时保持其功能性。通过在gGAN中引入一种新的对比学习方法，我们将代码清除和生成过程构建为一个双目标优化问题，从而显著减少LLM的推理次数。PromSec为生成安全、功能正常的代码提供了一种经济高效且实用的解决方案。在Python和Java代码数据集上进行的大量实验证实，PromSec能够有效提高代码安全性，同时保持其预期功能。实验表明，虽然现有方法未能解决所有代码漏洞，但PromSec能够有效解决这些问题。此外，PromSec在操作时间、LLM查询次数和安全分析成本方面实现了超过一个数量级的降低。更重要的是，使用PromSec优化的提示具有跨编程语言的LLM迁移性和对训练中未见漏洞的泛化性。本研究是提高LLM在安全和功能性代码生成方面可信度的一步，支持将其集成到实际软件开发中。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）生成代码时存在的安全漏洞问题。现有方法生成的代码由于训练数据的缺陷，往往包含安全隐患，并且缺乏有效的漏洞修复和安全保障机制，导致软件开发面临潜在的安全风险。

核心思路：论文的核心思路是将代码漏洞清除和代码生成过程结合起来，形成一个交互式的优化循环。通过生成对抗图神经网络（gGAN）识别并修复代码中的安全漏洞，并将修复结果反馈给LLM，从而优化LLM的提示，引导其生成更安全的代码。这种方法旨在从源头上减少漏洞的产生，并提高代码的整体安全性。

技术框架：PromSec的技术框架包含两个主要模块：1) 基于生成对抗图神经网络（gGAN）的代码漏洞清除模块；2) 基于LLM的代码生成模块。这两个模块通过一个交互循环连接在一起。首先，LLM根据初始提示生成代码；然后，gGAN对生成的代码进行安全分析和漏洞修复；接着，gGAN的修复结果被用于优化LLM的提示；最后，LLM根据优化后的提示重新生成代码。这个过程迭代进行，直到生成的代码满足安全性和功能性要求。

关键创新：PromSec的关键创新在于将代码漏洞清除和代码生成过程整合为一个双目标优化问题，并引入了一种新的对比学习方法来训练gGAN。这种方法能够更有效地识别和修复代码中的安全漏洞，并减少LLM的推理次数，从而提高代码生成的效率和安全性。此外，PromSec还具有跨编程语言的LLM迁移性和对训练中未见漏洞的泛化性。

关键设计：gGAN采用图神经网络结构，将代码表示为图的形式，节点表示代码元素，边表示代码元素之间的关系。gGAN包含一个生成器和一个判别器，生成器负责生成修复后的代码，判别器负责判断修复后的代码是否安全。对比学习方法用于训练gGAN，通过对比安全代码和不安全代码的特征，提高gGAN的漏洞识别和修复能力。此外，PromSec还设计了一种提示优化策略，根据gGAN的修复结果调整LLM的提示，引导LLM生成更安全的代码。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，PromSec能够有效解决代码漏洞，优于现有方法。PromSec在操作时间、LLM查询次数和安全分析成本方面实现了超过一个数量级的降低。更重要的是，使用PromSec优化的提示具有跨编程语言的LLM迁移性和对训练中未见漏洞的泛化性。例如，PromSec能够将LLM查询次数减少超过10倍，同时保证代码的安全性。

🎯 应用场景

PromSec可应用于软件开发的各个阶段，例如代码自动生成、代码审查和漏洞修复等。它可以帮助开发人员快速生成安全可靠的代码，减少人工审查和修复漏洞的工作量，从而提高软件开发的效率和质量。此外，PromSec还可以用于安全教育和培训，帮助开发人员了解常见的代码安全漏洞和防御方法，提高其安全意识。

📄 摘要（原文）

The capability of generating high-quality source code using large language models (LLMs) reduces software development time and costs. However, they often introduce security vulnerabilities due to training on insecure open-source data. This highlights the need for ensuring secure and functional code generation. This paper introduces PromSec, an algorithm for prom optimization for secure and functioning code generation using LLMs. In PromSec, we combine 1) code vulnerability clearing using a generative adversarial graph neural network, dubbed as gGAN, to fix and reduce security vulnerabilities in generated codes and 2) code generation using an LLM into an interactive loop, such that the outcome of the gGAN drives the LLM with enhanced prompts to generate secure codes while preserving their functionality. Introducing a new contrastive learning approach in gGAN, we formulate code-clearing and generation as a dual-objective optimization problem, enabling PromSec to notably reduce the number of LLM inferences. PromSec offers a cost-effective and practical solution for generating secure, functional code. Extensive experiments conducted on Python and Java code datasets confirm that PromSec effectively enhances code security while upholding its intended functionality. Our experiments show that while a state-of-the-art approach fails to address all code vulnerabilities, PromSec effectively resolves them. Moreover, PromSec achieves more than an order-of-magnitude reduction in operation time, number of LLM queries, and security analysis costs. Furthermore, prompts optimized with PromSec for a certain LLM are transferable to other LLMs across programming languages and generalizable to unseen vulnerabilities in training. This study is a step in enhancing the trustworthiness of LLMs for secure and functional code generation, supporting their integration into real-world software development.