fix: Block path traversal attacks via the deployment API for MLflow-Triton (#8665)

Author: mudit-eng

deploy/mlflow-triton-plugin/mlflow_triton/deployments.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -90,7 +90,7 @@ def create_deployment(self, name, model_uri, flavor=None, config=None):
         """
         Deploy the model at the model_uri to the Triton model repo. Associated config.pbtxt and *labels* files will be deployed.
 
-        :param name: Name of the of the model
+        :param name: Name of the model
         :param model_uri: Model uri in format model:/<model-name>/<version-or-stage>
         :param flavor: Flavor of the deployed model
         :param config: Configuration parameters
@@ -99,6 +99,9 @@ def create_deployment(self, name, model_uri, flavor=None, config=None):
         """
         self._validate_flavor(flavor)
 
+        # Validate model name
+        self._validate_model_name(name)
+
         # Verify model does not already exist in Triton
         if self._model_exists(name):
             raise Exception(
@@ -513,6 +516,18 @@ def _validate_flavor(self, flavor):
         if flavor not in self.supported_flavors:
             raise Exception("{} model flavor not supported by Triton".format(flavor))
 
+    def _validate_model_name(self, name):
+        # Check if the model name is empty or only contains whitespace, tabs, or newlines
+        if name.strip() == "":
+            raise Exception(
+                "Model name cannot be empty. Please enter a valid name to deploy."
+            )
+        # Path traversal protection
+        if "/" in name or name == "..":
+            raise Exception(
+                "Path traversal is not allowed in model's name: {}".format(name)
+            )
+
     def _model_exists(self, name):
         deploys = self.list_deployments()
         exists = False

qa/L0_mlflow/plugin_test.py

@@ -1,6 +1,6 @@
 #!/usr/bin/python
 
-# Copyright 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -33,7 +33,9 @@
 import json
 import unittest
 
+import mlflow.onnx
 import numpy as np
+import onnx
 import test_util as tu
 from mlflow.deployments import get_deploy_client
 
@@ -45,7 +47,7 @@ def setUp(self):
     def _validate_deployment(self, model_name):
         # create
         self.client_.create_deployment(
-            model_name, "models:/{}/1".format(model_name), flavor="onnx"
+            model_name, f"models:/{model_name}/1", flavor="onnx"
         )
 
         # list
@@ -79,8 +81,6 @@ def _validate_deployment(self, model_name):
 
     def test_onnx_flavor(self):
         # Log the ONNX model to MLFlow
-        import mlflow.onnx
-        import onnx
 
         model = onnx.load(
             "./mlflow-triton-plugin/examples/onnx_float32_int32_int32/1/model.onnx"
@@ -92,8 +92,6 @@ def test_onnx_flavor(self):
 
     def test_onnx_flavor_with_files(self):
         # Log the ONNX model and additional Triton config file to MLFlow
-        import mlflow.onnx
-        import onnx
 
         model = onnx.load(
             "./mlflow-triton-plugin/examples/onnx_float32_int32_int32/1/model.onnx"
@@ -116,6 +114,65 @@ def test_onnx_flavor_with_files(self):
             filecmp.cmp(config_path, "./models/onnx_model_with_files/config.pbtxt")
         )
 
+    def test_model_name(self):
+        EMPTY_MODEL_NAMES = [
+            "",
+            "     ",
+            " ",
+            "\n",
+            "\t",
+            "\r",
+            "\v",
+            "\f",
+        ]
+        INVALID_PATH_TRAVERSAL_NAMES = [
+            "/opt/sys/",
+            "../../etc/passwd",
+            "../outside/repo",
+            "test_models/../identity_py",
+            "..",
+        ]
+        VALID_MODEL_NAMES = [
+            "model123",
+            # "model  OAI",   TRI-769: Fix this test case
+            "model.version1",
+            "...",
+            "..my_model",
+            "model..1",
+            "model....1",
+        ]
+
+        for model_name in EMPTY_MODEL_NAMES:
+            model_uri = f"models:/{model_name}/1"
+            with self.assertRaises(Exception) as e:
+                self.client_.create_deployment(model_name, model_uri, flavor="onnx")
+            self.assertIn(
+                "Model name cannot be empty. Please enter a valid name to deploy.",
+                str(e.exception),
+            )
+
+        for model_name in INVALID_PATH_TRAVERSAL_NAMES:
+            model_uri = f"models:/{model_name}/1"
+            with self.assertRaises(Exception) as e:
+                self.client_.create_deployment(model_name, model_uri, flavor="onnx")
+            self.assertIn(
+                f"Path traversal is not allowed in model's name: {model_name}",
+                str(e.exception),
+            )
+
+        for model_name in VALID_MODEL_NAMES:
+            model = onnx.load(
+                "./mlflow-triton-plugin/examples/onnx_float32_int32_int32/1/model.onnx"
+            )
+
+            # Use a different name to ensure the plugin operates on correct model
+            mlflow.onnx.log_model(
+                model, "triton", registered_model_name=f"{model_name}"
+            )
+
+            # Validate deployment functionalities - create, list, get, predict, delete
+            self._validate_deployment(model_name)
+
 
 if __name__ == "__main__":
     unittest.main()

qa/L0_mlflow/test.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -187,7 +187,7 @@ if [ $? -ne 0 ]; then
     echo -e "\n***\n*** Python Test Failed\n***"
     RET=1
 else
-    check_test_results $TEST_RESULT_FILE 2
+    check_test_results $TEST_RESULT_FILE 3
     if [ $? -ne 0 ]; then
         cat $PY_LOG
         echo -e "\n***\n*** Test Result Verification Failed\n***"
@@ -256,7 +256,7 @@ if [ $? -ne 0 ]; then
     echo -e "\n***\n*** Python Test Failed\n***"
     RET=1
 else
-    check_test_results $TEST_RESULT_FILE 2
+    check_test_results $TEST_RESULT_FILE 3
     if [ $? -ne 0 ]; then
         cat $PY_LOG
         echo -e "\n***\n*** Test Result Verification Failed\n***"