PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Bo Lin, Shangwen Wang, Yihao Qin, Liqian Chen, Xiaoguang Mao

arXiv:2504.16429v1cs.CRcs.SE

TL;DR

Existing RACG systems often prioritize functional correctness over security, leaving a gap in defenses against poisoned knowledge bases. CodeGuarder adds security knowledge to code-generation prompts through a retrieval and filtering framework. It improves security in standard and poisoned RACG settings without compromising functional correctness.

  • Problem

    Existing RACG systems often prioritize functional correctness while overlooking security, and dedicated defenses for risks from poisoned retrieved code are lacking.

    Source: §1 · p. 1, §2.2 · p. 3

  • Method

    CodeGuarder retrieves security knowledge alongside functional examples and injects it into prompts, using query decomposition and vulnerability-susceptibility-based re-ranking and filtering.

    Source: §1 · p. 2, §4.1.1 · p. 4

  • Results

    Security improves by 20.12% in standard RACG and by 31.53% and 21.91% in two poisoning scenarios, without compromising functional correctness.

    Source: §1 · p. 2

  • Takeaways & Limitations

    CodeGuarder is a security-hardening framework for RACG that addresses knowledge-base poisoning while maintaining generated-code functionality.

    Source: §8 · p. 13

    Security measurement relies on a detector configured for 50 CWEs and may not capture the full security posture or all potential flaws.

    Source: §7.5 · p. 12

Abstract

from arXiv · show

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.

1 Introduction

Existing RACG systems prioritize functional correctness while often overlooking security, and poisoned knowledge bases can expose generated code to vulnerabilities. CodeGuarder adds security knowledge to retrieval-augmented prompts; evaluations report improved security across standard and poisoned settings while preserving functionality.

1 Introduction

  • A prior study found that even one poisoned sample could lead to vulnerabilities in 48% of generated code, while current RACG systems lack dedicated defenses against the threat [24].Source: §1 · p. 2
  • CodeGuarder combines functional code examples with security knowledge, including secure code samples and annotated vulnerability root causes, in generation prompts.It decomposes queries into sub-tasks, retrieves relevant guidance, then re-ranks and filters it according to LLM susceptibility to vulnerability types.Source: §1 · p. 2
  • Security improves by 20.12% in standard RACG and by 31.53% and 21.91% in two poisoning scenarios, without compromising functional correctness.The evaluation covers various LLMs and languages.Source: §1 · p. 2
  • The evaluation spans standard RACG, two poisoning setups, and direct code generation, and reports generalization across scenarios.Source: §1 · p. 2

2 Background and Related Works

RACG uses retrieved programming resources to support code generation, while research on code security has identified persistent vulnerabilities and focused largely on non-retrieval settings. The interaction between a model’s internal knowledge and runtime-retrieved knowledge, especially when the latter is poisoned, remains insufficiently addressed.

2.1 Retrieval Augmented Code Generation

  • RACG augments code generation with retrieved code snippets or other programming resources to improve efficiency and quality on complex tasks [9] [49].Source: §2.1 · p. 2

2.2 Security of LLM-Generated Code

  • Studies report persistent security problems in LLM-generated code, with one finding approximately 40% of generated code contained vulnerabilities [34] [44].Traditional post-generation static-analysis scanning adds latency because it requires a separate analysis step.Source: §2.2 · pp. 2–3
  • Prior defenses such as SafeCoder and SVEN use finetuning, while CoSec uses supervised co-decoding; these methods mainly target non-retrieval generation [14] [13].Source: §2.2 · p. 3
  • RACG introduces a distinct security challenge through interactions between internal and runtime-retrieved knowledge, but security hardening for that setting remains insufficiently investigated, particularly under poisoning.Source: §2.2 · p. 3

3 Motivating Examples

The motivating example shows how a poisoned retrieval example can steer generated code toward an insecure function. Adding security guidance can instead direct the model to a safer alternative while preserving the task’s functionality.

3 Motivating Examples

  • In the example, a poisoned semantically similar code sample uses sprintf, and the generated code follows that reference, risking a buffer overflow.Source: §3 · p. 3
  • With security knowledge warning about sprintf, GPT-4o uses snprintf instead, mitigating buffer-overflow risk while preserving functionality.Source: §3 · p. 3
  • The example motivates shifting from functional examples alone to prompts augmented with relevant security knowledge to guide code generation away from vulnerabilities.Source: §3 · p. 3

4 Methodology

CodeGuarder hardens retrieval-augmented code generation through an offline vulnerability-derived security knowledge base and online retrieval and prompt integration of task-specific security guidance.

4.1.1 Security Knowledge Definition.

  • Each entry links functionality for retrieval with a root-cause explanation and vulnerable-code example, plus a fixing-pattern explanation and corrected-code example.Source: §4.1.1 · p. 4

4.1 Automated Offline Construction of the Security Knowledge Base

  • Offline, CodeGuarder processes historical CVE descriptions, CWE labels, and vulnerable/fixed-code diffs to extract functionality, root causes, and fixing patterns into its security knowledge base.Source: §4.1.2 · p. 4, §4.1.1 · p. 4

4.2 Context-Aware Fine-Grained Knowledge Online Retrieval

  • Online retrieval decomposes requests into fine-grained sub-tasks because a single query can involve multiple vulnerabilities, then seeks relevant security knowledge for each task.Source: §4.2.1 · p. 4, §4.2 · p. 4

4.2.1 Query Decomposition.

  • The LLM breaks a query into semantically finer sub-tasks, which become the units used in the subsequent knowledge-retrieval stage.Source: §4.2.1 · p. 4

4.2.2 Security Knowledge Retrieval.

  • Sub-task and knowledge descriptions are embedded for cosine-similarity retrieval of top-k′ entries, although fixed-size retrieval can cause prompt truncation or attention dilution.Source: §4.2.2 · p. 5, §4.2.3 · p. 5

4.2.3 Security Knowledge Re-ranking and Filtering.

  • Re-ranking estimates each task’s risk by summing retrieved vulnerability-frequency weights, drawing on violation prevalence in LLM-generated code.CWE-476 receives weight 0.40 from a 40.24% NULL-pointer frequency; unlisted CWEs receive 0.01.Source: §4.2.3 · p. 5
  • Tasks are ranked by aggregate risk and the top five retained; their retrieved entries form the query-specific security knowledge supplied for generation.Source: §4.2.3 · p. 5

4.3 Security-Augmented Code Generation

  • The generation prompt pairs the functional query and examples with security guidance for selected sub-tasks, such as instructions for safer string copying.Source: §4.3 · p. 6

5 Study Design

The study evaluates CodeGuarder in standard and poisoned RACG settings, using distinct functional, security, and vulnerable-code bases. It tests several LLMs and retrievers, measuring security and functionality under controlled generation settings.

5.1 Investigated Scenarios

  • The study covers standard RACG, intent-exposed targeted poisoning, and intent-agnostic generalized poisoning, with five poisoned samples or a 10% poisoning proportion in the respective attacks.The generalized scenario poisons representative functional examples with corresponding vulnerable counterparts.Source: §5.1 · p. 6, §5.1.2 · p. 6, §5.1.3 · p. 6

5.2 Benchmark and Knowledge Bases

  • CyberSecEval supplies 1,916 cases across 50 CWE types and an insecure-code detector reported to have 96% precision.The study uses ReposVul for security knowledge and vulnerability-free functional examples, while CyberSecEval vulnerable code forms the attacker’s source.Source: §5.2.1 · p. 6, §5.2.2 · pp. 6–7

5.2.2 Knowledge Bases Construction.

  • The three knowledge bases serve separate roles: S supplies security guidance, K provides functional examples, and V supplies vulnerable code for poisoning.S derives from ReposVul vulnerability-fix pairs; K uses its fixed code, while V uses CyberSecEval vulnerable examples to avoid overlapping fixes.Source: §5.2.2 · pp. 6–7

5.3 Studied LLMs

  • The evaluation includes GPT-4o, DeepSeek-V3, CodeLlama-13B, and DeepSeek-Coder-V2-16B, spanning general-purpose and code-oriented models.The models include both API-accessed and open-source systems.Source: §5.3 · p. 7

5.4 Retriever

  • The code retriever fetches examples from K, the knowledge retriever retrieves security knowledge from S for decomposed query subtasks, and an external retriever simulates poisoning.The code and knowledge retrievers use jina-embeddings-v3; the poisoning retriever uses text-embedding-3-large.Source: §5.4 · p. 7

5.5 Metrics

  • Security Rate measures the percentage of generated code verified secure; functionality is assessed with code similarity and, separately, test-based benchmarks.Because CyberSecEval lacks test cases, CodeBLEU serves as a partial functionality proxy, with MBPP and HumanEval used for test-based assessment.Source: §5.5 · pp. 7–8

5.6 Implementation Details

  • Experiments use an A100 server, temperature 0, a 4,096-token generation limit, and an 8,192-token context window.DeepSeek-V3 supports security-knowledge construction and query decomposition.Source: §5.6 · p. 8

6 Evaluation

In standard RACG, CodeGuarder improves security across the evaluated models and languages while maintaining or slightly improving functionality. The evaluation frames this test alongside questions about poisoning resilience and generalization.

6.1 Research Questions

  • The evaluation poses questions about performance in standard RACG, resilience to poisoning, and generalization across conditions including missing target-language security knowledge.The standard-setting question concerns a knowledge base containing functional examples without poisoning.Source: §6.1 · p. 8, §6.2 · p. 8

6.2 RQ1: Performance in the Standard RACG

  • CodeGuarder maintains or slightly improves similarity, with average Sim values increasing from 22.67, 24.23, 28.13, and 21.29 to 23.64, 25.35, 28.46, and 23.03 for C, C++, Java, and Python.The authors also report test-based MBPP and HumanEval results as confirming preserved functionality.Source: §6.2 · p. 8
  • 27.20% for C, 9.37% for C++, 47.72% for Java, and 12.35% for Python are the reported average Security Rate improvements in standard RACG.Across four LLMs, the reported range is 9.37% to 47.72%; the table presents standard-scenario performance.Source: §6.2 · p. 8, Table 4 · p. 8

6.3 RQ2: Resilience to Poisoning Attacks

CodeGuarder improves security under both targeted poisoning with exposed programming intents and generalized poisoning without direct access to intents. The reported gains accompany maintained or slightly improved code similarity.

6.3.1 Poisoning Scenario I.

  • Under targeted poisoning, CodeGuarder improves Security Rate by an average of 31.53% despite a reported baseline drop from 60.84 to 50.57.The attacker injects five semantically relevant vulnerable examples while having access to programming intents; average Sim rises from 28.70 to 29.06.Source: §6.3.1 · p. 9

6.3.2 Poisoning Scenario II.

  • Under generalized poisoning, CodeGuarder improves Security Rate by an average of 21.91% when the attacker lacks direct query access and poisons representative examples.The attack uses the top 10% most representative examples; average Sim rises from 21.10 to 22.54 with CodeGuarder.Source: §6.3.2 · pp. 9–10

6.3 RQ2: Resilience to Poisoning Attacks

  • 31.53% and 21.91% are the average Security Rate improvements under targeted poisoning and generalized poisoning, respectively.The authors report that functionality is not compromised across the poisoning scenarios.Source: Table 5 · p. 9, Table 6 · p. 9, §6.3.2 · p. 10

6.4 RQ3: Generalization of CodeGuarder

CodeGuarder generalizes beyond retrieval-augmented settings: it improves security in non-func-retrieval and when target-language security knowledge is absent. These gains are evaluated alongside functionality preservation and comparisons with existing methods.

6.4.1 Performance in the Non-Func-Retrieval Scenario.

  • CodeGuarder is evaluated on C, C++, and Python with Mistral-7B and CodeLlama-7B in the non-func-retrieval setting.The authors select these models because baseline reproduction requires additional training, and avoid other languages to reduce potential evaluation bias.Source: §6.4.1 · p. 10
  • 9.80% more secure code on average than SafeCoder (72.31 →79.40) with Mistral-7B; CodeGuarder also outperforms investigated baselines in security and functionality in non-func-retrieval.The relative SR improvements over SafeCoder on Mistral-7B are 11.87% for C, 10.77% for C++, and 7.25% for Python.Source: §6.4.1 · p. 10, Table 7 · p. 10

6.4.2 Performance with Language-Specific Knowledge Absence.

  • 15.69%, 21.26%, and 13.50% average SR improvements in the standard and two poisoning scenarios, respectively, without target-language security knowledge.For poisoning scenario I, the average improvement is 21.26% without language-specific knowledge versus 31.53% for C, C++, Java, and Python with corresponding knowledge.Source: §6.4.2 · p. 11, Table 8 · p. 11
  • Without language-specific security knowledge, CodeGuarder still improves security by relying on general security patterns, but its poisoning mitigation is weaker than with tailored knowledge.The authors note that mitigation becomes more challenging without language-specific information.Source: §6.4.2 · p. 11

7 Discussion

The discussion examines functional correctness, vulnerability-specific prevention, component ablations, knowledge-volume tuning, and validity threats. Results indicate benefits alongside limits in vulnerability coverage and measurement.

7.1 Functional Correctness Maintenance

  • CodeGuarder slightly enhances functionality across four LLMs and both benchmarks; GPT-4o’s MBPP pass@1 rises from 72.8 to 73.2 and pass@5 from 77.4 to 78.8.The authors report a similar trend on HumanEval and across the other evaluated LLMs.Source: §7.1 · p. 11

7.2 Effectiveness Across Vulnerability Types

  • Prevention rates range from 4.76% for CWE-79 to 89.15% for CWE-119; absent JavaScript-specific knowledge limits effectiveness on CWE-79.CWE-200 prevention is 66.67%; CWE-22 also performs less well, at 8.26%.Source: §7.2 · pp. 11–12, Table 10 · p. 12

7.3 Ablation Study

  • Removing query decomposition decreases SR by approximately 10.1% (76.36% to 68.61%) and Sim from 25.56 to 23.04 in the standard scenario.The ablation uses DS-V3 and holds the number of injected knowledge entries consistent across variants.Source: §7.3 · p. 12, Table 11 · p. 12

7.4 Impact of the Number of Injected Knowledge Entries

  • Performance generally improves with more injected knowledge up to a point, then degrades; optimal (k′, k) settings are (2, 7) for DS-V3 and (2, 5) for CodeLlama-13B.The authors suggest smaller models may be overwhelmed by larger knowledge contexts; the main evaluations use (2, 5).Source: §7.4 · p. 12, Table 12 · p. 12

7.5 Threats to Validity

  • Security measurement uses a static-analysis detector targeting 50 CWEs, which may not capture the full security posture or all potential flaws.The authors report 96% detection precision for vulnerabilities within that set and note that identical tools and criteria support fair comparisons.Source: §7.5 · pp. 12–13
  • Across five DS-V3 runs, the maximum SR deviation was 0.38%, despite residual non-determinism not being fully assured at zero temperature.The authors characterize the observed variability as negligible.Source: §7.5 · p. 13

8 Conclusion

CodeGuarder injects security knowledge into RACG prompts to address knowledge-base poisoning. Evaluations show improved code security while maintaining functionality, with generalization across languages and models.

8 Conclusion

  • CodeGuarder is a security-hardening framework for RACG that explicitly incorporates security knowledge into code-generation prompts.The conclusion identifies knowledge-base poisoning as an overlooked threat the framework addresses.Source: §8 · p. 13
  • Across evaluated scenarios, CodeGuarder improves generated-code security while maintaining functionality and demonstrates generalization across languages and models.Source: §8 · p. 13
  • The authors characterize the work as a step toward securing LLM-based software development.Source: §8 · p. 13

A Security Knowledge Example

The security-knowledge example traces a buffer-overflow risk to an insecure string-copy function and illustrates how the knowledge is structured to guide safer generation.

A Security Knowledge Example

  • CVE-2012-6538 involves an insecure function whose use may lead to a buffer overflow.The example identifies strcpy as the unsafe function.Source: A Security Knowledge Example · p. 15, Figure 4 · p. 15
  • The patched example replaces strcpy with strncpy and imposes a specific destination-buffer size constraint.Source: A Security Knowledge Example · p. 15
  • The extracted knowledge represents functionality, root cause, and fixing pattern, combining explanations with code examples.The functionality dimension describes the vulnerable snippet’s intended string-copy operation; the root-cause and fixing-pattern dimensions explain the risk and safer practice.Source: A Security Knowledge Example · p. 15

B Prompt Templates

Appendix B provides prompt templates for extracting security knowledge from fixing commits, decomposing generation queries, and pairing sub-tasks with retrieved security knowledge during generation.

B.1 Security Knowledge Extraction

  • The extraction prompt takes a vulnerability description, CWE type, and function-level diff containing vulnerable and fixed code as input.The diff includes full function context.Source: §B.1 · p. 15
  • It asks for code functionality, vulnerability root cause, and fixing pattern, with the output structured as JSON descriptions and code examples.Source: §B.1 · p. 15

B.2 Query Decomposition

  • The query-decomposition prompt breaks the user's original code-generation query into smaller, detailed sub-tasks described by their functionality.Source: §B.2 · p. 15
  • Its output is a JSON array of sub-task descriptions.Source: §B.2 · p. 15

B.3 Security-Augmented Code Generation

  • The security-augmented generation prompt pairs each sub-task with its corresponding retrieved security knowledge.The template labels these as the N-th sub-task and its corresponding retrieved knowledge.Source: §B.3 · p. 16
  • The example preserves the original code-generation prompt and code examples, then appends injected security knowledge organized by sub-task.Source: §B.3 · p. 16
  • The example knowledge describes an unsafe strcpy buffer-size assumption and includes descriptions and code examples of root causes and fixing patterns.A separate excerpt gives strncpy as safe-copy guidance.Source: §B.3 · p. 16

Related papers

Loading 2504.16429v1…