Cloudpickle Load on LightGBM SciKit Learn Model Leading to Code Execution

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