aslein1413's picture
Upload folder using huggingface_hub
71a6ac6 verified
|
Raw
History Blame Contribute Delete
1.35 kB
# PoC — code execution when loading an MLflow model with pyfunc.load_model()
`mlflow.pyfunc.load_model()` on an untrusted model directory runs code straight from the
model, with no pickle involved. The `MLmodel` file names a `loader_module` and a `code`
directory. On load MLflow prepends the model's `code/` dir to `sys.path` and imports the
named `loader_module`, so the module's top-level code runs before you ever call the model.
Here `MLmodel` points `loader_module: evil_loader` at `code/evil_loader.py`, whose import
runs `id` and drops `/tmp/PWNED_mlflow.txt`.
## Files
- `model/MLmodel` — the model manifest (loader_module + code dir).
- `model/code/evil_loader.py` — imported on load; runs `id` at import time.
- `verify.py` — calls `load_model("model")` and prints the marker.
## Reproduce
```
pip install mlflow
python verify.py
# marker after : True
# uid=0(root) gid=0(root) groups=0(root)
```
Confirmed on a clean python:3.12-slim container with stock mlflow 3.14.0.
## Why it matters
There is no pickle here, so a scanner that only looks for pickle opcodes (ModelScan
reports this dir as clean) sees nothing, yet loading the model runs attacker code as the
host process. Any pipeline that loads a user-supplied MLflow model is affected. It fires
on the default `load_model` call with no flags or environment variables.