Give LLMs a Security Course: Securing Retrieval-Augmented Code Generation via Knowledge Injection

TL;DR

This work tackles security gaps in Retrieval-Augmented Code Generation by introducing CodeGuarder, a security-hardening framework that injects CVE-derived security knowledge into RACG prompts. It offline-constructs a security knowledge base capturing functionality, root cause, and fixing patterns, then online-decomposes queries into sub-tasks, retrieves and re-ranks security knowledge, and injects it into generation to guide secure code synthesis. Across standard and poisoned scenarios, CodeGuarder consistently improves code security (average SR gains up to 47.72%) without sacrificing functional correctness, and generalizes well across languages and model types, even when language-specific security knowledge is absent. The approach demonstrates strong resilience to both targeted and broad poisoning attacks and offers a practical path toward secure, trustworthy LLM-driven software development.

Abstract

Retrieval-Augmented Code Generation (RACG) leverages external knowledge to enhance Large Language Models (LLMs) in code synthesis, improving the functional correctness of the generated code. However, existing RACG systems largely overlook security, leading to substantial risks. Especially, the poisoning of malicious code into knowledge bases can mislead LLMs, resulting in the generation of insecure outputs, which poses a critical threat in modern software development. To address this, we propose a security-hardening framework for RACG systems, CodeGuarder, that shifts the paradigm from retrieving only functional code examples to incorporating both functional code and security knowledge. Our framework constructs a security knowledge base from real-world vulnerability databases, including secure code samples and root cause annotations. For each code generation query, a retriever decomposes the query into fine-grained sub-tasks and fetches relevant security knowledge. To prioritize critical security guidance, we introduce a re-ranking and filtering mechanism by leveraging the LLMs' susceptibility to different vulnerability types. This filtered security knowledge is seamlessly integrated into the generation prompt. Our evaluation shows CodeGuarder significantly improves code security rates across various LLMs, achieving average improvements of 20.12\% in standard RACG, and 31.53\% and 21.91\% under two distinct poisoning scenarios without compromising functional correctness. Furthermore, CodeGuarder demonstrates strong generalization, enhancing security even when the targeted language's security knowledge is lacking. This work presents CodeGuarder as a pivotal advancement towards building secure and trustworthy RACG systems.

Figures (4)

• Figure 1: An example of code generated by GPT-4o with and without additional security knowledge.
• Figure 2: The overall workflow of Code-Guard-er
• Figure 3: Vulnerable code generated by GPT-4o along with its corresponding decomposed queries and retrieved knowledge.
• Figure 4: An example of security knowledge extracted from vulnerability