[论文解读] DRAINCODE: Stealthy Energy Consumption Attacks on Retrieval-Augmented Code Generation via Context Poisoning

Yanlin Wang, Jiadong Wu+8 more|arXiv (Cornell University)|Jan 28, 2026

Adversarial Robustness in Machine Learning被引用 0

一句话总结

DrainCode 是基于 RAG 的代码生成攻击，通过污染检索上下文来诱导显著更长的输出，从而增加延迟和能耗，同时保持功能准确性。

ABSTRACT

Large language models (LLMs) have demonstrated impressive capabilities in code generation by leveraging retrieval-augmented generation (RAG) methods. However, the computational costs associated with LLM inference, particularly in terms of latency and energy consumption, have received limited attention in the security context. This paper introduces DrainCode, the first adversarial attack targeting the computational efficiency of RAG-based code generation systems. By strategically poisoning retrieval contexts through a mutation-based approach, DrainCode forces LLMs to produce significantly longer outputs, thereby increasing GPU latency and energy consumption. We evaluate the effectiveness of DrainCode across multiple models. Our experiments show that DrainCode achieves up to an 85% increase in latency, a 49% increase in energy consumption, and more than a 3x increase in output length compared to the baseline. Furthermore, we demonstrate the generalizability of the attack across different prompting strategies and its effectiveness compared to different defenses. The results highlight DrainCode as a potential method for increasing the computational overhead of LLMs, making it useful for evaluating LLM security in resource-constrained environments. We provide code and data at https://github.com/DeepSoftwareAnalytics/DrainCode.

研究动机与目标

• 突出对检索增强代码生成（RAG）的能量和延迟导向攻击的安全风险。
• 将 DrainCode 作为基于变异的中毒框架，定位检索上下文以提高输出长度而不破坏代码功能。
• 开发带有隐蔽约束的梯度引导触发器优化方法，以最大化标记长度和能耗。
• 展示跨模型与跨提示的泛化能力，并与基线和防御进行对比。

提出的方法

• 将被污染的代码上下文注入检索语料库，诱导大型语言模型产生冗长输出。
• 生成一个假设的未完成代码查询，以实现检索片段的提示无关中毒。
• 对对抗性触发器进行梯度引导的变异，结合 EOS 损失和 KL 散度约束以保留非触发行为。
• 通过多位置变异和攻击缓冲池提升中毒效率，加速收敛。
• 在多种代码模型和数据集上评估 DrainCode，并与 RawRAG、Prompt Injection 及 LLMEffiChecker 进行对比。

实验结果

研究问题

• RQ1 DrainCode 在基于 RAG 的代码生成中提高延迟、能耗和输出长度的效果有多大？
• RQ2 DrainCode 的哪些组件对攻击效果贡献最大？
• RQ3 DrainCode 在提示策略和模型间的泛化能力如何？
• RQ4 常见检测方法能否识别 DrainCode 中毒上下文？
• RQ5 在攻击下 DrainCode 对代码的可读性和功能准确性有何影响？

主要发现

• DrainCode 将输出长度提升约 3×–10×，延迟最高提升至 85%，能耗增加约 49%。
• 尽管开销显著，功能准确性仍保持在 95–99% 之间。
• DrainCode 的中毒速度比以往攻击快 3.5×，在开销方面优于基线。
• 攻击效果对提示策略普遍鲁棒，并能转移到黑箱设置。
• 在对分类器和困惑度防御下仍保持隐蔽性，同时在资源受限环境中吞吐量下降。