Exploring the Security Threats of Knowledge Base Poisoning in Retrieval-Augmented Code Generation. Bo Lin, Ian Chen, Shangwen Wang, Xiaoguang Mao. arXiv:2502.03233v1 [cs.CR], 5 Feb 2025.
研究背景与动机
研究背景
大型语言模型（LLMs）与软件开发的深度融合，特别是通过检索增强代码生成（RACG）系统，极大地改变了该领域。RACG系统通过利用外部知识库中的信息来增强代码生成的质量。然而，RACG系统的安全性影响，特别是知识库中包含易受攻击的代码示例所带来的风险，在很大程度上仍未得到探索。这一风险尤其令人担忧，因为公共代码仓库（通常是RACG系统中知识库收集的来源）通常对社区中的任何人都是可访问的。恶意攻击者可以利用这种可访问性将易受攻击的代码注入知识库，使其变得有害。一旦这些被污染（poisoned）的样本被检索并整合到生成的代码中，它们就会将安全漏洞传播到最终产品中。
研究动机与问题
先前的研究主要关注LLM直接生成的代码的安全性，而对RACG系统生成的代码安全性，尤其是在知识库被攻击者污染的情况下，存在关键的理解空白。随着RACG迅速成为现代基于LLM的系统中的主流范式，这一知识空白变得更加紧迫。因此，本文旨在首次对RACG系统相关的安全风险进行全面研究，重点关注知识库中的易受攻击代码如何危害生成代码的安全性。研究具体探讨了以下两个研究问题（RQs）：
**RQ1：** 易受攻击的代码示例是否会损害LLM生成的代码的安全性？
**RQ2：** RACG中的不同因素在多大程度上影响生成代码的安全性？
论文核心方法和步骤
威胁模型与攻击场景
论文定义了两种现实的攻击场景：
**场景I（暴露的编程意图）：** 攻击者可以在开发者与RACG系统交互之前观察或预测其编程意图（即查询）。攻击者基于查询的语义，有策略地向知识库中注入相关的易受攻击代码示例。
**场景II（隐藏的编程意图）：** 攻击者无法预先知道具体的开发者查询。攻击者转而专注于向常见的编程模式和功能中注入易受攻击的代码示例。
知识库污染攻击的形式化
从攻击者的视角，知识库污染攻击被定义为一个优化问题，目标是构建一组易受攻击的示例 V，使得RACG系统中的LLM在利用从被污染知识库 K∪V中检索到的 r个示例作为上下文时，更可能生成易受攻击的代码。目标函数为：
Vmax∣Q∣1q∈Q∑I(LLM(q;RETRIEVE(q,K∪V)))
其中，I(⋅)是评估LLM生成代码安全性的函数（安全为0，易受攻击为1），∣Q∣是查询集合的大小，RETRIEVE(⋅)函数从被污染知识库中为查询 q检索一组 r个文本。
实验设置与方法
**数据集构建：** 基于ReposVul数据集，该数据集包含C、C++、Java和Python四种语言的安全代码和易受攻击代码对，共计12,053个实例，涵盖236个不同的CWE（通用弱点枚举）类型。对于缺乏明确功能描述（作为查询）的函数，使用LLM（DeepSeek-V2.5）根据安全版本的代码生成查询。
知识库构建与污染：
**场景I：** 知识库 K包含所有安全代码。对于每个查询 q，攻击者使用一个外部的“污染检索器”（基于TE3嵌入模型）从漏洞知识库（所有易受攻击代码）中检索与 q最相似的 m个易受攻击示例，注入 K。
**场景II：** 使用基于K-means的聚类方法对知识库中的代码元素进行聚类。然后根据污染比例 p从每个簇中选择代表性的代码元素，并为每个选中的元素从漏洞知识库中检索最相似的易受攻击代码进行注入。
**实验对象：** 研究了四个代表性的LLM（GPT-4o, Llama-3-8B, CodeLlama-13B, DeepSeek-Coder-V2-16B）和两种检索器（稀疏检索器BM25，稠密检索器JINA）。实验覆盖了两种攻击场景，共16个子场景。
**结果验证（LLM作为评判官）：** 由于生成的代码在语义和格式上差异很大，手动分析耗时且容易出错，而静态分析工具也难以应用。因此，论文采用“LLM-as-a-Judge”的方法进行安全评估。该过程分为两步：
**漏洞知识提取：** 分析数据集中安全代码和易受攻击代码之间的补丁差异，提取漏洞成因模式（Vulnerability Cause Pattern）和修复模式（Fixing Pattern）。
**漏洞检测：** 使用提取的模式来判断生成的代码是否包含这些漏洞模式。评估表明该LLM评判官在不同编程语言上的准确率在77%到81%之间。
评估指标：
**漏洞率（VR, Vulnerability Rate）：** 生成的代码片段中存在安全漏洞的百分比，VR=Nv/Nt。
**相似度（Similarity）：** 使用CrystalBLEU度量生成代码与地面真值（secure version）之间的功能相似性。
**检索代码中的漏洞率（VRRC, Vulnerability Rate in Retrieved Code）：** 在所有检索到的代码中，易受攻击代码的平均比例，VRRC=∣Q∣1∑q∈Qr∣Vq∣。
实验结果与结论
主要发现（对应RQ1）
**知识库污染构成重大威胁：** 知识库污染对RACG系统中LLM生成的代码安全构成显著威胁，特别是对代码LLM（如CodeLlama）。例如，在场景I中，即使只注入一个被污染的样本，使用JINA检索器和CodeLlama时，约48%的生成代码会变得易受攻击。即使在场景II（攻击者不知查询）下，当污染比例达到20%时，CodeLlama生成的代码仍有约36%是易受攻击的。
**代码LLM更易受影响：** 代码LLM（CodeLlama, DS-Coder）比通用LLM（GPT-4o, Llama-3）更容易生成易受攻击的代码。这是因为代码LLM在更大的代码数据集上训练，学习了更多代码模式（包括易受攻击的模式），使其更容易受到输入中易受攻击代码的影响。
**少样本学习增加风险：** 提供更多的演示示例（从1-shot增加到3-shot）通常会提高VR，尽管模型性能保持稳定或有轻微改善。例如，在场景I中使用JINA检索器时，所有模型的平均VR从1-shot的0.46增加到3-shot的0.49（增长6.5%）。这是因为更多的示例增加了检索到易受攻击代码的可能性及其在检索结果中的比例。
影响因素分析（对应RQ2）
**编程语言的影响：** LLM在C++语言中生成的易受攻击代码最多（VR最高），而在Java中表现出更强的抵抗力。C语言的VR也较高，尤其是在JINA检索器下，表明其缺乏安全抽象使其更易受影响。Python和Java由于内置的安全机制，VR相对较低且稳定。
**示例-查询相似度的影响：** 生成易受攻击代码的可能性与检索到的代码和查询之间的语义相似度呈正相关。当相似度高于60%时，漏洞风险显著增加；而较低相似度（0-60%）的影响相对较小。
**漏洞类型的影响：** 漏洞率（VR）因CWE类型而异。在MITRE Top-10弱点中，CWE-352（跨站请求伪造） consistently表现出最高的VR（约0.8），而CWE-434（不受限制上传危险类型文件）的VR最低。
结论与意义
本研究首次全面探讨了RACG系统因知识库污染带来的安全风险。实验结果表明，知识库污染会显著危害生成代码的安全性，即使是单一污染样本也可能导致近半数的生成代码存在漏洞。研究还识别了少样本学习、编程语言、示例-查询相似度等因素对漏洞引入的影响。这些发现强调了提高知识库安全性、设计更健壮的检索策略（如不从最相似的示例中检索）以及隐藏编程意图以增加攻击难度的重要性。本研究为未来旨在保护RACG系统安全、增强其在软件开发中可靠性的工作奠定了基础。