YAML Metadata Warning:empty or missing yaml metadata in repo card
Check out the documentation for more information.
MLflow PyFunc Remote Model Loading Code Execution Demonstration
This repository demonstrates arbitrary Python code execution during mlflow.pyfunc.load_model() when a PyFunc model stored in a shared remote artifact store contains an attacker-controlled loader module.
The proof of concept uses a mocked S3 environment (moto) to simulate a shared MLflow artifact bucket.
Execution occurs as part of the model loading workflow after MLflow processes the artifact.
Attack Overview
Modify an existing PyFunc model artifact stored in a shared remote bucket.
Replace the loader module referenced in
MLmodel.Load the model using
mlflow.pyfunc.load_model("s3://...").MLflow imports the configured loader module during model initialization.
Attacker-controlled Python code executes before the model is returned.
Root Cause
PyFunc models support custom Python loaders through the loader_module field in the MLmodel configuration. During loading, MLflow retrieves this configuration from the model artifact and imports the specified module.
The model artifact therefore controls which Python module is imported during initialization. If the artifact contents are modified, the imported module can be replaced with attacker-controlled code.
Prerequisites
- Attacker can modify the contents of a shared model artifact
- Victim loads the model through
mlflow.pyfunc.load_model() - Victim has permission to read the model artifact
Reproduce
pip install -r requirements.txt
python repro.py
Expected Output
=== MLflow S3 Model Store Demo ===
Simulating shared team S3 bucket attack
[+] Created shared S3 bucket: s3://shared-mlflow-models/
[+] Modified model artifact uploaded to S3
[+] s3://shared-mlflow-models/fraud-detector/MLmodel references loader_module: loader
[*] Victim calls: mlflow.pyfunc.load_model('s3://shared-mlflow-models/fraud-detector')
[*] Marker before load: False
[loader] Loader executed — marker created at /tmp/MLFLOW_S3_RCE
[+] Model loaded successfully (looks legitimate)
[+] predict() returned: [1.0, 2.0, 3.0]
[*] Marker after load: True
============================================================
Code execution confirmed during model loading
============================================================
Marker created: /tmp/MLFLOW_S3_RCE
Code executed on victim machine during model load
============================================================