Machine Learning for CybersecuritySource on GitHub

Unit 10 · Lesson 10.3

Securing retrieval, agents and the model supply chain

Lesson 10.2 concluded that a language model cannot be relied on to ignore instructions in its input, and that the guarantees have to come from the system around it. This lesson builds that system. It starts with retrieval-augmented generation, a common way to give a model access to an organization's documents, and its attack surface. It then turns to agents, models that call tools, where an injected instruction can become an action, and to the controls that limit what such an action can do. It ends with the model files themselves, which can carry code, backdoors or a false history.

Retrieval-augmented generation

A model's parameters hold what it learned during training and nothing newer. Retrieval-augmented generation (RAG) gives it access to a document collection at answer time. Lewis et al. (2020) combined a pretrained generator with a dense vector index of documents and a neural retriever. The retriever encodes the query \(x\) and every document \(z\) as vectors \(q(x)\) and \(d(z)\) with two trained encoders and scores each document by

\[ p(z \mid x) \propto \exp\!\left(d(z)^\top q(x)\right), \label{eq:retriever} \]

so retrieving the \(k\) best documents means finding the \(k\) largest inner products, a maximum inner product search that can be approximated in sublinear time. The retrieved passages are pasted into the model's context together with the question. Many deployed systems keep this structure, often with a general-purpose embedding model in place of the trained encoders and with cosine similarity (the inner product divided by the product of the two vectors' lengths) in place of the raw inner product.

Retrieval ranks documents by how close their embeddings are to the query's, and an attacker who can add a document to the collection can write one that is close. Take a query embedding \(q = (0.6, 0.8)\) and three documents with embeddings \(d_1 = (0.8, 0.6)\), a relevant page, \(d_2 = (0, 1)\), a loosely related one, and \(d_3 = (0.62, 0.78)\), a page the attacker wrote by restating the expected question. The inner products are \(0.96\), \(0.80\) and \(0.996\), so the attacker's page ranks first. No access to the model, the retriever's weights or the index is needed, only the ability to publish text that will be indexed. Figure 10.6 marks the points where attacks enter such a pipeline.

Figure 10.6. A retrieval-augmented generation pipeline and the places where attacks enter (red markers). 1, documents written or altered by an attacker in a collection that the system indexes. 2, the vector index, whose embeddings can leak the text they encode and which must enforce the same access

Figure 10.6. A retrieval-augmented generation pipeline and the places where attacks enter (red markers). 1, documents written or altered by an attacker in a collection that the system indexes. 2, the vector index, whose embeddings can leak the text they encode and which must enforce the same access rights as the documents. 3, retrieved passages that carry injected instructions into the model's context. 4, a tampered embedding model or model file in the supply chain. 5, model output passed to users or tools without validation.

Knowledge poisoning. Zou et al. (2025) formalized poisoning of a retrieval collection. The attacker picks a target question and a target answer and crafts malicious texts that satisfy two conditions: a retrieval condition (the text is retrieved for the target question) and a generation condition (once retrieved, it leads the model to the target answer). Injecting five malicious texts per target question into a collection of millions of texts achieved an attack success rate of 90%, measured as the share of target questions for which the model gave the attacker's target answer, and the defenses they evaluated were insufficient. For a security team, the target question might be "Is the domain update-cdn.example malicious?" and the planted answer "No, it belongs to our software vendor".

Embedding weaknesses. An embedding looks like an opaque list of numbers, but it is a function of the text. Morris et al. (2023) trained a model that iteratively corrects and re-embeds a guess, and recovered 92% of 32-token inputs exactly from their embeddings; on embeddings of clinical notes it recovered most of the full names. An index of embeddings must therefore be protected like the documents it was computed from. In a shared index, the permission check has to happen at retrieval time, so that one user's query cannot return passages from documents that user may not read. The OWASP Top 10 for LLM Applications lists these risks as LLM08:2025, Vector and Embedding Weaknesses (OWASP Foundation, 2024).

Injection through retrieval. Every retrieved passage is untrusted text in the sense of Lesson 10.2. The defenses of that lesson apply, and so do the questions of provenance. Three questions matter more than any prompt wording: who can write to the indexed collection, which sources are indexed automatically, and whether answers cite the passages they used so that a reader can check them.

Show the answer

The scraped advisories are attacker-writable, since anyone who can publish a page that the scraper visits, or compromise a vendor's site, controls their text. Ticket history is partly attacker-writable too, because tickets quote emails, alerts and log lines that external parties produced. Only the wiki is written by staff, and even it may quote external material. Reasonable changes are to tag every passage with its source and show the sources with each answer, to spotlight passages from external sources (Lesson 10.2), to restrict what the system may do on the basis of external passages (no tool calls, no changes to tickets), and to monitor for passages that are retrieved unusually often across unrelated questions, a symptom of content written to match many queries.

Agents and excessive agency

An agent is a language model run in a loop in which it can call tools: search, read a mailbox, query a database, open a ticket, run code. The model's output is parsed for tool calls, the application executes them, and their results are appended to the context for the next step. Tool results are a retrieval channel like any other, so an injected instruction in a web page or an email can now become an action. Debenedetti et al. (2024) built AgentDojo to measure this, with 97 realistic tasks (managing email, e-banking, travel booking) and 629 security test cases. They found that capable models fail many tasks even without attacks, and that existing injection attacks break some security properties but not all.

The OWASP list names the underlying risk LLM06:2025, Excessive Agency, and traces it to three root causes (OWASP Foundation, 2024). Excessive functionality means the agent has tools or tool features it does not need, such as an extension that can also modify and delete documents when the task only needs to read them. Excessive permissions means the tools act with more rights than the task requires, such as a database identity with write access for a read-only feature. Excessive autonomy means high-impact actions run without independent verification or approval. Each cause suggests a control that lives outside the model and therefore does not depend on the model resisting injection.

Least privilege. Give each agent the smallest set of tools, the narrowest functionality within each tool and the weakest credentials that its task needs, and run tools with the user's own permissions instead of a shared privileged identity. A capability that does not exist cannot be abused.

Allowlists per task. Decide in advance which tools each kind of task may call, and refuse others whatever the model asks for. An agent that summarizes a web page has no reason to send email.

Argument checks. Validate the arguments of every call with ordinary code: recipients must belong to the organization's domain; file paths must stay inside a working folder after normalization; and addresses to block must not be the organization's own servers. These are the same input-validation rules a web application applies to user input, now applied to the model's output.

Confirmation for consequential actions. Ask a person to approve actions with side effects that are hard to undo. The person sees the proposed action, not the model's reasoning, and can judge whether its target matches what they asked for. Confirmation has two limits. A reviewer cannot judge content they do not read, and a reviewer who approves many routine prompts learns to approve without reading.

Isolation and logging. Keep untrusted content away from the component that decides on actions, as in the dual model pattern and CaMeL of Lesson 10.2. Record every tool call with its arguments, the task it served and the context that produced it, so that an incident can be reconstructed, and set budgets on calls and tokens per task, which also addresses LLM10:2025, Unbounded Consumption.

The controls combine in series. If an injected instruction is obeyed with probability \(p_0\) and each of \(J\) layers, numbered \(j = 1, \dots, J\), independently stops a fraction \(b_j\) of the obeyed injected actions that reach it, the probability that an injected document causes a harmful action is

\[ P(\text{harm}) = p_0 \prod_{j=1}^{J} (1 - b_j). \label{eq:layers} \]

With \(p_0 = 0.2\), an allowlist that stops half of the injected actions, argument checks that stop 60% of the rest and a reviewer who rejects 90% of what remains, equation \eqref{eq:layers} gives

\[ P(\text{harm}) = 0.2 \times 0.5 \times 0.4 \times 0.1 = 0.004, \]

or about four harmful actions per thousand injected documents. The independence assumption is the weak point. An injected action that uses the right tool on the right target with malicious content passes the allowlist, the argument checks and a reviewer who checks targets, all for the same reason, so the layers fail together and the product in \eqref{eq:layers} is optimistic for that class of attack.

Lab 10b measures these controls on a deterministic toy agent whose planner obeys every tool request it finds in its context, the worst case for the model. Figure 10.7 summarizes the results.

Figure 10.7. Results of Lab 10b on 24 synthetic security operations tasks, 19 of which contain injected tool requests in an alert, an email, a web page or a ticket. (a) Injected actions executed and requested actions completed under three configurations: no policy, an allowlist of tools per task, an

Figure 10.7. Results of Lab 10b on 24 synthetic security operations tasks, 19 of which contain injected tool requests in an alert, an email, a web page or a ticket. (a) Injected actions executed and requested actions completed under three configurations: no policy, an allowlist of tools per task, and the allowlist with argument checks and confirmation by a simulated reviewer for actions with side effects. (b) Injected actions executed under the full configuration when the reviewer approves a given share of prompts without reading them (mean over 200 random repetitions).

With no policy, all 20 injected actions run. The allowlist stops those that use tools the task does not need and leaves 12, without blocking any of the 54 requested actions. Argument checks and confirmation leave 3, and they require 36 confirmation prompts over 24 tasks. The surviving actions belong to exactly the class that \eqref{eq:layers} handles badly. They are tickets and a message created through the right tool, for the right target, with content the attacker wrote; the message tells the team lead that an alert is a false positive. When the reviewer approves half of the prompts unread, the number of injected actions executed rises to 4.5 on average over the 200 repetitions.

Show the answer

Each task type in the lab needs a small, known set of tools, and many injections ask for tools outside that set (a shell command during page summarization, an email during alert triage). Those calls are refused, and every legitimate call is to an allowed tool. A general-purpose agent has no task type to restrict by; every tool is allowed for every task, so an allowlist removes nothing. This is the argument of Beurer-Kellner et al. (2025) for application-specific agents with clearly defined trust boundaries. The narrower the task, the more a control outside the model can enforce.

The model supply chain

Applications download models, tokenizers and embedding models from public repositories and load them with a single function call. Each of those files is a supply-chain dependency, and OWASP lists the risk as LLM03:2025, Supply Chain (OWASP Foundation, 2024). Three problems deserve separate treatment, because each has a different control.

Code in model files. Python's pickle format, which PyTorch uses to save models, does not store data alone. A pickle is a short program for a stack machine, and it can name any importable function and call it with arguments while the file is loaded. The Python documentation warns that it is possible to construct malicious pickle data that will execute arbitrary code during unpickling, and that data from an untrusted source should never be unpickled (Python Software Foundation, 2026). Since version 2.6, torch.load uses weights_only=True by default, which restricts the unpickler to tensors, primitive types, dictionaries and any types the user explicitly allowlists, and does not let it import anything else, so files that name other functions are refused. The PyTorch documentation is careful about what this achieves. It narrows the surface for remote code execution, but does not guard against denial of service, and memory corruption during loading may still be possible (PyTorch Contributors, 2026).

Weights-only formats. The safetensors format avoids executable content by construction. A file holds an 8-byte length, a JSON header that gives each tensor's name, data type, shape and byte offsets, and then the raw bytes (Hugging Face, 2026). Loading it means parsing JSON and copying memory, so there is no step at which the file can name a function to call. Figure 10.8 contrasts the two layouts. The model in Lab 10a is distributed in this format.

Figure 10.8. What a loader does with two model file formats. (a) A pickle is a sequence of opcodes for a small stack machine; opcodes such as STACK_GLOBAL, which imports a named function (here the built-in print), and REDUCE, which calls it, run during loading. Red boxes mark the opcodes that a mali

Figure 10.8. What a loader does with two model file formats. (a) A pickle is a sequence of opcodes for a small stack machine; opcodes such as STACK_GLOBAL, which imports a named function (here the built-in print), and REDUCE, which calls it, run during loading. Red boxes mark the opcodes that a malicious file uses. (b) A safetensors file is a length, a JSON header and raw tensor bytes; loading parses the header and copies the bytes.

Integrity and provenance. A safe format says that loading a file will not run code. It does not say that the file is the one its publisher released. A cryptographic digest (SHA-256) or, better, a signature, published through a channel the attacker does not control and checked before loading, detects tampering in transit or in a mirror. OWASP observes that published models have no strong provenance assurances, since model cards and documentation give information about a model but no guarantee of its origin, and recommends using models from verifiable sources together with integrity checks based on signatures and file hashes (OWASP Foundation, 2024).

Backdoors in the weights. A signed file in a safe format can still contain a model that misbehaves on purpose. Kurita et al. (2020) showed that pretrained weights can be poisoned so that a backdoor survives the victim's own fine-tuning, letting the attacker change the fine-tuned model's predictions by inserting a chosen keyword into the input. Nothing in the file format or the signature reveals this, because the publisher's weights are exactly the poisoned ones. The defenses are those discussed for backdoors in Lesson 8.1, such as behavioral testing for triggers, together with a preference for publishers whose training process is documented and accountable. A backdoored model also remains bounded by the controls on what its output can do.

These risks appear in NIST's profile for generative AI, NIST AI 600-1 (July 2024), which lists twelve risks that generative AI creates or worsens: chemical, biological, radiological or nuclear (CBRN) information or capabilities; confabulation; dangerous, violent or hateful content; data privacy; environmental impacts; harmful bias or homogenization; human-AI configuration; information integrity; information security; intellectual property; obscene, degrading or abusive content; and value chain and component integration (National Institute of Standards and Technology, 2024a). The last of these, the non-transparent integration of third-party components and data, is the supply-chain risk of this section, and information security covers both attacks on generative AI systems and their misuse for attacks. Unit 11 places the profile within NIST's broader risk management framework.

Show the answer

The safe format addresses code execution during loading, and the matching digest shows that the file was not altered between the publisher's page and the download. Several risks remain. The digest on the page proves nothing if the attacker controls the page or the publisher's account, so the publisher's identity needs a stronger anchor, such as a signature tied to a known identity. The weights may contain a backdoor planted by the publisher or by whoever trained the base model, which needs behavioral testing. And the classifier still needs the evaluation of Units 1 and 2 on the organization's own data before anyone relies on it.

What the lab shows

Lab 10b

Least privilege for a tool-using agent

Lab 10b implements the three agent configurations and the reviewer-fatigue experiment of Figure 10.7, then demonstrates the supply-chain points with harmless code. It builds a pickle whose loading prints a message, shows the opcodes that cause it, loads an object with the same payload through torch.load with weights_only=True (refused) and with weights_only=False (the message prints), writes and reads a safetensors file with a few lines of Python and NumPy, without the safetensors library, and detects a one-byte change with a SHA-256 digest.

Summary

Retrieval gives a model access to documents, and anyone who can write to the indexed collection can place text that ranks high for chosen queries, plants false answers or carries instructions into the context; embeddings of sensitive text are themselves sensitive. Agents turn obeyed instructions into actions, so their security rests on controls outside the model: least privilege, allowlists of tools per task, argument checks written in ordinary code, confirmation for consequential actions with attention to reviewer fatigue, isolation of untrusted content, and logging. These layers multiply the protection only for attacks that some layer can recognize; injected content delivered through a legitimate action passes them together. Model files are dependencies that can run code when loaded, can be altered in transit and can carry backdoors in their weights, which call for weights-only formats, verified signatures and behavioral testing respectively.

References

  1. Beurer-Kellner, L., Buesser, B., Creţu, A.-M., Debenedetti, E., Dobos, D., Fabian, D., Fischer, M., Froelicher, D., Grosse, K., Naeff, D., Ozoani, E., Paverd, A., Tramèr, F., and Volhejn, V. (2025). Design patterns for securing LLM agents against prompt injections. arXiv preprint arXiv:2506.08837. https://arxiv.org/abs/2506.08837
  2. Debenedetti, E., Zhang, J., Balunović, M., Beurer-Kellner, L., Fischer, M., and Tramèr, F. (2024). AgentDojo: A dynamic environment to evaluate prompt injection attacks and defenses for LLM agents. Advances in Neural Information Processing Systems (Datasets and Benchmarks Track), 37, 82895-82920. https://proceedings.neurips.cc/paper_files/paper/2024/hash/97091a5177d8dc64b1da8bf3e1f6fb54-Abstract-Datasets_and_Benchmarks_Track.html
  3. Hugging Face (2026). safetensors: Simple, safe way to store and distribute tensors. Repository README and documentation; accessed 15 October 2026. https://github.com/huggingface/safetensors
  4. Kurita, K., Michel, P., and Neubig, G. (2020). Weight poisoning attacks on pretrained models. Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, 2793-2806. https://doi.org/10.18653/v1/2020.acl-main.249
  5. Lewis, P., Perez, E., Piktus, A., Petroni, F., Karpukhin, V., Goyal, N., Küttler, H., Lewis, M., Yih, W.-t., Rocktäschel, T., Riedel, S., and Kiela, D. (2020). Retrieval-augmented generation for knowledge-intensive NLP tasks. Advances in Neural Information Processing Systems, 33, 9459-9474. https://proceedings.neurips.cc/paper_files/paper/2020/hash/6b493230205f780e1bc26945df7481e5-Abstract.html
  6. Morris, J., Kuleshov, V., Shmatikov, V., and Rush, A. (2023). Text embeddings reveal (almost) as much as text. Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, 12448-12460. https://doi.org/10.18653/v1/2023.emnlp-main.765
  7. National Institute of Standards and Technology (2024a). Artificial intelligence risk management framework: Generative artificial intelligence profile. National Institute of Standards and Technology. NIST AI 600-1, July 2024. https://doi.org/10.6028/NIST.AI.600-1
  8. OWASP Foundation (2024). OWASP Top 10 for LLM Applications 2025. OWASP Gen AI Security Project. Version 2025, released 18 November 2024. https://genai.owasp.org/resource/owasp-top-10-for-llm-applications-2025/
  9. Python Software Foundation (2026). pickle: Python object serialization. Python 3 standard library documentation; accessed 15 October 2026. https://docs.python.org/3/library/pickle.html
  10. PyTorch Contributors (2026). Serialization semantics, and torch.load. PyTorch 2.14 documentation; accessed 15 October 2026. https://docs.pytorch.org/docs/2.14/notes/serialization.html
  11. Zou, W., Geng, R., Wang, B., and Jia, J. (2025). PoisonedRAG: Knowledge corruption attacks to retrieval-augmented generation of large language models. Proceedings of the 34th USENIX Security Symposium, 3827-3844. https://www.usenix.org/conference/usenixsecurity25/presentation/zou-poisonedrag