RESCUE: Retrieval Augmented Secure Code Generation

TL;DR

RESCUE introduces a retrieval-augmented secure code generation framework that combines a hybrid knowledge base (semantic summaries plus security-focused code slices) with a hierarchical, multi-faceted retrieval strategy. By operating offline to build a CWE-structured knowledge base and online to perform proactive, facet-based retrieval, it guides LLMs to produce secure code with minimal loss in functional accuracy. Comprehensive experiments across four benchmarks and six LLMs demonstrate a substantial improvement in SecurePass@1 (average +4.8 points) and robust ablations validate the effectiveness of both knowledge base construction and hierarchical retrieval, indicating strong potential for practical, scalable secure code generation. The work advances secure software engineering by integrating structured security knowledge with targeted retrieval, enabling more reliable and verifiable code generation in real-world tasks.

Abstract

Despite recent advances, Large Language Models (LLMs) still generate vulnerable code. Retrieval-Augmented Generation (RAG) has the potential to enhance LLMs for secure code generation by incorporating external security knowledge. However, the conventional RAG design struggles with the noise of raw security-related documents, and existing retrieval methods overlook the significant security semantics implicitly embedded in task descriptions. To address these issues, we propose RESCUE, a new RAG framework for secure code generation with two key innovations. First, we propose a hybrid knowledge base construction method that combines LLM-assisted cluster-then-summarize distillation with program slicing, producing both high-level security guidelines and concise, security-focused code examples. Second, we design a hierarchical multi-faceted retrieval to traverse the constructed knowledge base from top to bottom and integrates multiple security-critical facts at each hierarchical level, ensuring comprehensive and accurate retrieval. We evaluated RESCUE on four benchmarks and compared it with five state-of-the-art secure code generation methods on six LLMs. The results demonstrate that RESCUE improves the SecurePass@1 metric by an average of 4.8 points, establishing a new state-of-the-art performance for security. Furthermore, we performed in-depth analysis and ablation studies to rigorously validate the effectiveness of individual components in RESCUE.

Figures (3)

• Figure 1: The overview framework of Rescue.
• Figure 2: A known vulnerability and its fix that demonstrates the replacement of unsafe yaml.load() with yaml.safe_load(), where red-highlighted code shows vulnerable lines removed and green-highlighted code indicates security fixes added (See Appendix \ref{['app:complete-vuln-fix-example']} for the complete example).
• Figure 3: Comparison of SecurePass@1(SP@1) performance across hierarchical and non-hierarchical settings for three LLMs, DeepSeek-Coder-V2-Lite (Deepseek), Llama3.1-8B-Instruct (Llama), Qwen2.5-Coder-7B (Qwen), at varying numbers of selected CWE types (k).