DRAINCODE: Stealthy Energy Consumption Attacks on Retrieval-Augmented Code Generation via Context Poisoning

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.

💡 Research Summary

The paper introduces DrainCode, the first adversarial attack that targets the computational efficiency of retrieval‑augmented generation (RAG) based code generation systems. While prior work on code‑generation attacks has focused on functional corruption—e.g., causing the model to emit buggy or malicious code—DrainCode exploits a non‑functional vulnerability: the relationship between output length and inference cost. In a typical RAG pipeline, a retriever fetches relevant code snippets from an external corpus, and a large language model (LLM) generates code conditioned on the retrieved context and the user’s unfinished query. The LLM continues generating tokens until it emits an End‑of‑Sequence (EOS) token, so longer outputs directly translate into higher GPU latency, greater energy consumption, and increased operational cost.

DrainCode’s core idea is to poison the retrieval corpus with short, learnable token sequences—called “triggers”—that are syntactically valid but semantically inert. When a poisoned snippet is retrieved, the trigger subtly manipulates the LLM’s probability distribution, suppressing the likelihood of early EOS emission and encouraging token diversity. The attack is optimized using a dual‑loss objective: (1) an EOS loss that penalizes early termination, and (2) a KL‑divergence constraint that forces the output distribution of the poisoned context to stay close to that of the clean context. This ensures that the generated code remains functionally correct (95‑99 % test‑case pass rate) while inflating the token count, often by a factor of three to ten.

Three technical contributions make DrainCode both effective and stealthy:

1. Query‑agnostic poisoning – The authors introduce a hypothetical query generation step that uses the LLM itself to synthesize plausible unfinished code snippets. This allows the attacker to craft poisoned contexts without knowledge of actual user queries, removing the restrictive assumption that attackers must pre‑define target queries.

2. Efficient mutation via multi‑position mutation and an attack buffer pool – Instead of searching the entire input space, the method restricts mutation to the trigger tokens, identifies the most influential positions via gradient information, and re‑uses promising mutations from a buffer. This reduces the search space dramatically and yields a >3× speed‑up compared with earlier mutation‑based attacks.

3. Stealth through KL‑regularization – By limiting the KL divergence between clean and poisoned outputs, the attack preserves the overall style and semantics of the generated code, making the extra verbosity hard to detect by perplexity‑based or classifier‑based defenses.

The authors evaluate DrainCode on several state‑of‑the‑art code generation models (e.g., CodeGen‑2B, StarCoder‑6B, GPT‑NeoX‑20B) and across four prompting strategies (function implementation, code completion, bug fixing, and test generation). Results show:

• Output length increases by 3‑10× (up to 12× in extreme cases).
• GPU inference latency rises up to 85 % (average ≈ 60 %).
• Measured energy consumption grows by an average of 49 % (peak ≈ 62 %).
• Functional correctness remains high, with 95‑99 % of generated snippets passing unit tests.

DrainCode also outperforms prior energy‑efficiency attacks such as LLMEffChecker and NMTSloth‑style prompt perturbations, delivering 25‑32 % more overhead while poisoning up to 3.5× faster. The attack remains effective in black‑box settings, transfers across models, and evades both classifier‑based and perplexity‑based defenses.

The broader implication is a new class of LLM‑Denial‑of‑Service (LLM‑DoS) attacks that exhaust computational resources rather than corrupt output. In real‑world deployments—IDE plugins, cloud‑based code generation APIs, CI/CD pipelines—continuous invocation of poisoned retrieval contexts could dramatically increase operational costs and carbon emissions. The paper therefore calls for defenses such as integrity verification of retrieval corpora, trigger detection mechanisms, token‑length caps, and real‑time energy‑usage monitoring. Future work is suggested in (i) automated detection of stealthy triggers, (ii) anomaly detection based on energy consumption patterns, and (iii) extending the attack to multi‑modal and multi‑language code generation models.

In summary, DrainCode demonstrates that by subtly poisoning the retrieval component of RAG systems, an adversary can launch a covert, resource‑exhaustion attack that is both highly effective and difficult to detect, opening a new research frontier in the security of AI‑assisted software engineering.