Module 15 — ML/AI system security¶
Why it matters to a software engineer¶
Module 12 taught you to secure an application that calls an LLM. This module flips the lens: the model, its training data, and its serving pipeline are now the asset you are defending, using the exact same trust-boundary and asset/threat/risk/control vocabulary from Module 1. If you build or ship models — recommendation, classification, embeddings, fraud scoring, or an LLM — this is the module that treats them as production systems with their own attack surface, not as a black box someone else secures.
Visual overview¶
flowchart LR
DATA[Training data] --> TRAIN[Training job]
TRAIN --> REG[Model registry]
REG --> SERVE[Serving API]
CALLER[Caller / app] -->|query| SERVE
SERVE -->|prediction| CALLER
DATA -. poisoned examples .-> TRAIN
CALLER -. heavy querying .-> STEAL[Reconstructed copy]
SERVE -. crafted input .-> MISCLASS[Wrong decision, normal-looking input]
Intuition
Redraw this as Module 5's supply-chain diagram with different labels —
developer→source→CI→artifact→registry→runtime becomes
data→training→registry→serving. The shapes of the attacks are
familiar too: poisoning is tampering with the "source," extraction is
theft via the public "runtime" API, and a crafted input is Module 4's
"data becomes code" pattern with the model as the unsafe interpreter.
| Attack | Targets | Looks like | Prevent / raise cost | Detect / attribute |
|---|---|---|---|---|
| Data poisoning | Training pipeline integrity | A backdoored or biased model, discovered late | Provenance + eval on trusted held-out data | Eval drift, canary triggers |
| Model extraction | Confidentiality/availability of the model | A very active, very ordinary-looking API client | Authn, quotas, rate limits, restrict high-information outputs | Query auditing, behavioral detection, fingerprinting, watermarking where applicable |
| Adversarial input | Serving-time integrity | A normal-looking input, wrong output | Robustness testing, not encryption | Anomalous-input / decision monitoring |
| Excessive agency | Blast radius of a wrong output | A correct-sounding action with real consequences | Tool-scoped policy + human approval (Module 12) | Tool-call audit |
Hint
"Encrypt the model file" answers a question nobody asked. Model extraction only requires query access to a public API — the file never has to leave the server. Do not treat rate limits, auditing, and watermarking as one "query path defense." Raise the cost of copying (authentication, quotas, rate limits, restrict high-information outputs) separately from detecting or attributing it (query audit, fingerprinting, watermarking). Neither list is the other.
Learning objectives¶
- Apply the Module 1 asset/boundary/threat/control lens to a model instead of an API.
- Distinguish attacks on training data, the model artifact, and the serving API, and for extraction name both a prevent/raise-cost control and a detect/attribute control (they are not the same).
- Explain why a model's outputs are untrusted input to whatever reads them next — the same "data becomes code" pattern from Module 4.
- Review a tool-using agent's configuration for supply-chain and excessive-agency risk without adding new lab infrastructure.
Key concepts¶
The ML pipeline has the same shape as any other supply chain. Data source → collection → labeling → training → evaluation → registry → serving → monitoring. Module 5's supply-chain diagram (developer → source → CI identity → artifact → signature → registry → runtime) maps directly: swap "source code" for "training data" and "build" for "training run." A compromise at any stage can arrive as an apparently normal deployment.
Training-data poisoning. An attacker who can influence training or fine-tuning data biases the model's future behavior — a backdoor trigger phrase, a systematically mislabeled class, a skewed recommendation. The defense is provenance (know where every training example came from) and evaluation on held-out, trusted data before promotion, the same "verify before trust" instinct as signature checking in Module 5.
Model theft / extraction. An adversary with only query access can reconstruct a close approximation of a model by querying it heavily and training a copy on the input/output pairs. That is primarily a confidentiality problem (the model is the secret). Heavy querying can also become unbounded consumption (availability). Encrypting the model file still matters against registry/backup theft — a different path. Query-access extraction is not solved by at-rest encryption.
Controls for this split into two different jobs, and confusing them is how "add rate limiting" ends up as the entire security review:
- Prevent / raise extraction cost (make copying expensive, not impossible): authentication, per-caller quotas, rate limits, and restricting high-information outputs (e.g. truncating raw logits/probabilities).
- Detect / attribute (assume some extraction succeeds, and catch it): query auditing for systematic probing, behavioral detection on query volume/diversity, response fingerprinting, and watermarking where applicable.
Neither list substitutes for the other. Rate limits are not watermarking: raising cost slows a patient attacker but does not tell you copying happened; detection/attribution tells you it happened but does not stop the first successful run.
Adversarial examples. Inputs crafted to be misclassified while looking normal to a human (or normal-looking log lines crafted to look like instructions to an LLM — this is Module 12's prompt injection, restated: the same "interpretation crosses a boundary" pattern from Module 4, with the model as the unsafe interpreter).
Excessive agency and tool misuse. Covered operationally in Module 12; here, review it as a design-time control. A model that can only read is a different risk than one that can write, delete, or call other services. The blast radius of a wrong model output is bounded by what its tools are allowed to do, not by how accurate the model usually is.
Model/data confidentiality vs business value. A model trained on sensitive data can leak fragments of that data through its outputs (membership inference, verbatim regurgitation). Treat "the model has seen this data" as an additional probabilistic exposure surface, not as a second conventional access path. That changes classification and retention decisions from Module 7.
A trained model may memorize and expose training information, creating a
probabilistic read path to sensitive data — not a deterministic one.
For threat modelling, treat model memorization as a potential data-exposure
surface rather than as a conventional database. You cannot SELECT a
specific record out of it, query it with guaranteed recall, delete one row
from it on request, or reason about its access control the way you would a
database table. That difference matters operationally: "the model saw this
data" does not tell you which queries will surface it, and "we deleted
the row" does not mean the model has forgotten it. Plan retention and
deletion requests around the model's training/retraining cycle, not around
a single row's lifecycle.
Architecture connection¶
Draw the same trust-boundary diagram as Module 1, with four new boxes:
Training data → Training job → Model registry → Serving API. Ask the
Module 1 questions at each boundary: what would hurt if this were disclosed,
changed, or unavailable; who can write to it; what does a caller's token
actually authorize once it reaches the model.
Hands-on lab — threat-model an added model, audit an existing agent¶
No new containers. This lab reuses Module 1's method and Module 12's real lab files — the point is to prove the framework transfers, not to stand up new ML infrastructure.
Prerequisites¶
Completed Module 1 (trust-boundary diagram) and Module 12 (agentic SOC lab).
Before you write this¶
Predict: (1) which trust boundaries the smart-search feature adds (2)
which tool in policy.yaml has the worst blast radius (3) why.
Then do the steps. Compare with your notes. If you missed a boundary or a tool, which assumption was wrong?
Steps¶
- Design exercise. Acme Notes is adding a "smart search" feature: notes
are embedded and a similarity model suggests related notes. Draw the
trust-boundary diagram for this addition: where do embeddings get
written, who can query the similarity index, does it cross the
IMDS/egress boundary from Module 1's reference diagram if the embedding model is a hosted API. - Using the STRIDE categories from Module 1, list one concrete threat per category for the smart-search feature (e.g. Tampering: a note author poisons their own note text to manipulate what gets suggested to other users).
- Audit exercise. Open
labs/agentic-soc/policy.yamlandlabs/agentic-soc/agent.py. For each tool the policy allows, write one sentence: what is the worst thing this tool could do if the planner chose it based on a manipulated summary. Compare against Module 12's "excessive agency" note. - Propose one control for the highest-severity item you found in step 3 that does not involve changing the model or prompt (e.g. narrowing the policy allowlist, adding an approval gate, capping call frequency).
Expected observations¶
A written trust-boundary diagram and STRIDE table for smart search,
matching Module 1's format. A short audit note per tool in
policy.yaml naming a concrete worst case, not a generic "could be
misused."
Security lessons¶
The lab is the smart-search diagram plus the tool audit: a model can
create a probabilistic read path into memorized training information,
while a tool-using agent can additionally become an actor. Bound the
search feature by who can write embeddings and who can query them; bound
the agent by what policy.yaml allows, not by model accuracy. Extraction
prevent vs detect (quotas vs watermarking) belongs in the concepts
section above — this lab never runs an extraction client. The Module 4
pattern (untrusted input crossing an interpreter) did not change; the
interpreter did.
Common mistakes¶
- Treating "the model is a black box" as a reason to skip the trust- boundary exercise instead of a reason to do it more carefully.
- Proposing "make the model more accurate" as a security control. Accuracy and safety are different properties; a highly accurate model with unbounded tool access is still a high-severity design.
- Confusing model theft (confidentiality/availability of the model) with data poisoning (integrity of training data) — they need different controls.
Cleanup¶
None.
Knowledge check¶
- Why doesn't encrypting a model file stop model extraction?
- Name the ML-pipeline equivalent of Module 5's "signed artifact."
- How is a poisoned training example different from a prompt-injected log line, and how are they the same?
- What bounds the blast radius of a wrong model output?
- Why does "the model has seen this data" change a data-classification decision?
- Why is watermarking not the same kind of control as a rate limit?
Answers: (1) The model is served through a public API; extraction only needs query access, not the file. (2) A trained model with recorded provenance and an evaluation signoff before it enters the registry. (3) Different injection point (training time vs. inference time), same pattern: untrusted content shapes future behavior. (4) The tool permissions granted to whatever acts on the model's output, not the model's accuracy. (5) The data now has an additional probabilistic read path (memorized training information in the model's outputs) — not a conventional database you can query or delete a row from. (6) Rate limits (with authn, quotas, and restricted high-information outputs) raise extraction cost. Watermarking, query auditing, and fingerprinting detect or attribute copying after queries already succeeded.
Engineering assignment¶
Extend your Module 1 trust-boundary diagram to include the smart-search addition from step 1. Submit it alongside one paragraph naming the two highest-severity threats and their controls.
Self-check¶
Answer before expanding. These are the assessment moves on this module's Acme Notes lab, not trivia.
Explain: Why doesn't encrypting the model file stop query-access extraction?
Extraction needs the serving API, not the registry file. At-rest encryption still matters against backup theft — a different path.
Predict: How is a poisoned training example different from a prompt-injected log line, and how is it the same?
Different injection point (training vs inference). Same Module 4 pattern: untrusted content shapes future behavior. Accuracy is not the control.
Diagnose: simulate_action is the only respond tool in policy.yaml. What is the worst case if the planner trusts a manipulated summary?
It can still request an allowlisted action (disable_lab_mode, revoke_token_notice, block_actor, snapshot_logs). Approval must fail closed. The blast radius is the tool list, not model quality.
Design: Name one prevent/raise-cost control and one detect/attribute control for extraction — and why they are not interchangeable.
Prevent: authn, quotas, rate limits, truncated logits. Detect: query audit, watermarking, fingerprinting. Rate limits do not tell you copying happened; watermarks do not stop the first successful run.
Defend: Residual risk of 'the model has seen this notes corpus.'
A probabilistic read path remains after you delete a sqlite row.
Containment is restrict who can query, log queries, retrain on a
cycle — not DELETE FROM notes. "Make the model more accurate" does
not bound that surface.
Before you leave¶
- Predict — write expected findings (what appears, what does not, and why) before the threat-model and tool audit.
- Diagnose — name the interpreter or tool-permission miss, not "the model is a black box."
- Build — smart-search trust-boundary diagram plus one worst-case sentence per allowlisted tool.
- Defend — state containment and residual risk in one sentence each.
- Exit criteria — the course pass bar: Explain → Predict → Diagnose → Design → Defend.