DrainCode: Efficiency Attack on RAG Systems

• DrainCode is an adversarial attack framework that poisons the retrieval corpus with learnable triggers to induce verbose outputs in code-generation RAG systems.
• It strategically increases GPU latency and energy consumption by forcing the model to generate extended outputs while maintaining functional accuracy.
• The framework employs gradient-based trigger optimization and dual loss functions to maximize token diversity and delay termination signals.

DrainCode denotes an adversarial attack framework against retrieval-augmented generation (RAG) systems for code completion, targeting the computational efficiency—specifically latency and energy consumption—of LLMs. By poisoning the retrieval corpus with learnable triggers, DrainCode exploits the dependency of code LLMs on retrieved context, coercing the generation of anomalously verbose outputs while preserving functional correctness. This attack paradigm exposes RAG systems to a new vector: the stealthy inflation of non-functional resource use, distinct from prior attacks focused on output semantics or security.

1. RAG-Based Code Generation and the Efficiency Attack Surface

Retrieval-augmented generation (RAG) combines a retriever (e.g., BM25, dense encoder) and a generator (LLM) for scalable code completion tasks. In repository-level workflows, frameworks such as RepoCoder, RLCoder, and CoCoMIC rely on retrieving cross-file code fragments to resolve dependencies or complete missing functions, with the LLM generating completions conditioned on concatenated retrieved contexts with an unfinished snippet. The cost profile of such systems is dominated by GPU latency (inference time) and energy draw (joule consumption), due to both linear token expansion and nonlinear computation overhead associated with large contexts and model dimensions.

DrainCode leverages the fact that extending output length directly scales latency and energy. Beyond functional cost, this also constitutes an attack vector analogous to LLM-Denial-of-Service (LLM-DoS) via “token exhaustion” in IDE assistants or cloud-hosted code agents. Previous adversarial efforts (such as NMTSloth, LLMEffiChecker) address natural language or prompt-level perturbations, often breaking code semantics or assuming white-box access, without consideration for stealthy context poisoning in code-specific RAG systems (Wang et al., 28 Jan 2026).

2. Adversarial Threat Model and Attack Workflow

The adversarial model underpinning DrainCode assumes the capability to inject up to K poisoned snippets into a public corpus (with K ∈ {1, 2, 3}), without influencing the retriever or the user query distribution. The adversary can access open-source LLM weights to perform gradient-based optimization for trigger engineering. The attack objective is to ensure that whenever a poisoned context is retrieved, the LLM emits many additional, benign-looking tokens, thereby surging throughput demands and energy costs while the completion maintains a high pass rate in functional tests.

A query-agnostic workflow is employed: first, a clean snippet is sampled from the corpus, and the LLM generates a “hypothetical query” mimicking real user tasks. A short, learnable subsequence (“trigger”) is injected at salient positions in the code snippet. An iterative mutation-based loop uses gradient-guided optimization to maximize output length and token diversity, constrained by KL-divergence to preserve non-trigger token distributions analogous to the clean case. On deployment, the retriever’s natural recall probability for poison ensures operational impact at inference, without altering overt completion correctness.

3. Attack Objective, Optimization Strategies, and Algorithmic Design

The technical core of DrainCode is an optimization procedure over the concatenated input $x' = [c ~\|~ s ~\|~ u']$, targeting two losses:

• EOS-loss: $$\mathcal{L}_1(x') = \frac{1}{N} \sum_{i=1}^{N} P(y_i = \mathrm{\langle EOS \rangle}~|~y_{<i}, x')$$—penalizes early output termination, shifting the output distribution to delay EOS.
• Token-Diversity-loss: $$\mathcal{L}_2(x') = -\lVert [g_1(x'), g_2(x'), \dots, g_N(x')] \rVert_*$$—maximizes nuclear norm of hidden states across tokens for greater output diversity.

The overall objective is:

$$\mathcal{L} = \mathcal{L}_1 + \lambda\,\mathcal{L}_2$$

subject to a KL-divergence constraint:

$$\forall t \notin \mathcal{T}: D_\mathrm{KL}(P_\mathrm{clean}(y_t) \|\ P_\mathrm{attack}(y_t)) \leq \epsilon$$

The mutation-based trigger search adopts a buffer pool paradigm, initializing multiple random triggers and iteratively refining via backpropagated gradients at top positions, substituting tokens to minimize the attack loss, as detailed in the pseudocode provided in (Wang et al., 28 Jan 2026).

GPU latency L and energy E models are empirically derived:

• $L \approx f(T, D)$, where T is output tokens and D is model dimension.
• $E = P_{avg} \times L$, with GPU power sampled via NVIDIA Management Library (NVML).

4. Experimental Framework and Evaluation Protocol

Evaluations are performed on two principal datasets:

• RepoEval: 373 repository-level completion tasks, 1.7 million GitHub files.
• Odex: 945 open-domain Python problems, with a 34,000-library-document corpus.

Code-centric LLMs (DeepSeek-Coder-7B, CodeQwen1.5-7B) and general LLMs (InternLM2-7B, Llama3-8B) are paired with BM25 retrieval (context window 2048 tokens, max new tokens 1024), with inference on NVIDIA A100 GPUs. Metrics include latency (ms), energy (J), functional accuracy (Pass@1 via test-case execution), and poisoning overhead (seconds per snippet).

5. Empirical Results and Comparative Analysis

DrainCode attains a pronounced increase in computational overhead:

• Output length: 3–10× (e.g., DeepSeek-7B, RepoEval: 313.8 → 1023.9 tokens, +226%).
• Latency: up to +85% (e.g., CodeQwen-7B: 22,221 ms → 41,112 ms).
• Energy: up to +49% (e.g., CodeQwen-7B: 5510 J → 8227 J).
• Pass@1: 95–99% of baseline (≤5% drop).

Against baselines:

• Prompt-injection breaks semantics and produces smaller overhead.
• LLMEffiChecker (white-box full-prompt mutation): DrainCode adds +25–32% extra latency/energy.

Ablation demonstrates the dual loss necessity: removing diversity or EOS loss halves output length and reduces resource impact by up to 40%. Hypothetical query generation as context boosts attack effectiveness compared to random or empty queries. Length caps scale energy linearly, with Pass@1 remaining stable.

Poisoning is computationally efficient: while LLMEffiChecker incurs ~760 s per snippet, DrainCode achieves ~216 s (3.5× speedup).

6. Generalizability, Defenses, and Detection Limitations

DrainCode triggers are broadly effective across prompting modalities (Zero-Shot, Few-Shot, Chain-of-Thought), with output generation saturating at context window maxima and minimal functional accuracy loss (<7%). Triggers generalize in “black-box” transfer: optimization against one model increases length 2–3× and latency/energy 70–300% on others, with ≤9% Pass@1 drop.

Detection approaches (SVM over hidden-state embeddings, perplexity-based ROC-AUC) are largely ineffective (Acc ∼ 0.30–0.37; AUC ~0.41). CodeBERT fine-tuned for detection achieves AUC ~0.82 but is limited to 512-token inputs, indicating substantial blind spots and the ongoing stealthiness of poisoned contexts.

7. Security Implications and Future Research Trajectories

DrainCode reveals a vulnerability in RAG-based code agents: stealthy “computational-efficiency” attacks that inflate non-functional costs without degradating correctness. This exposes energy-consumption as a new operational surface, particularly salient in resource-constrained or multi-tenant deployments.

Directions for mitigation include retrieval corpus vetting, dynamic EOS injection or hard generation caps, energy-aware anomaly monitoring, and adversarial training penalizing verbosity triggers. Prospective research encompasses robust retriever scoring, multi-modal RAG systems incorporating code, documentation, and API queries, as well as the formal integration of energy consumption into threat models and safety benchmarks.

In summary, DrainCode is the inaugural paradigm for context poisoning in code-generation RAG corpora, producing stealthy, functionally valid but computationally expensive outputs while eluding extant defenses. It establishes “efficiency security” as a complementary criterion to traditional correctness and safety evaluations in LLM deployment (Wang et al., 28 Jan 2026).