Pickle Load on Pmdarima Model Load Leading to Code Execution

Products Impacted

This vulnerability was introduced in version 1.24.0 of MLflow.

CVSS Score: 8.8

AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H

CWE Categorization

CWE-502: Deserialization of Untrusted Data.

Details

The vulnerability exists within the pmdarima/__init__.py file, within the function _load_model. This is called when the mlflow.pmdarima.load_model function is called.

def _load_model(path):
	with open(path, "rb") as pickled_model:
    		return pickle.load(pickled_model)

An attacker can exploit this by injecting a pickle object that will execute arbitrary code when deserialized into a model. The attacker can then call the pmdarima.log_model() function to serialize this model and log it to the tracking server. In the below example, the malicious pickle object has been injected into the init method within the _PmdarimaModelWrapper class in the file pmdarima/__init__.py, which is called through the auto_arima function. Before logging the model, editing the save_model function to force MLflow to pickle the attacker’s model is also required.

with mlflow.start_run():
	# Create the model
	model = pmdarima.auto_arima(train["sales"], seasonal=True, m=12)
	...
	# Log model
	mlflow.pmdarima.log_model(model, ARTIFACT_PATH, registered_model_name="PmdarimaTestModel")

When the model is loaded by the victim (example code snippet below), the arbitrary code is executed on their machine:

import mlflow
...
logged_model = "models:/PmdarimaTestModel/1"
loaded_model = mlflow.pmdarima.load_model(logged_model, dst_path='/tmp/pmdarima_model')