From e2ef31201d117be59cb089563b322453f888aaa6 Mon Sep 17 00:00:00 2001 From: Mudit Aggarwal Date: Tue, 17 Feb 2026 15:49:58 -0800 Subject: [PATCH 1/9] [TRI-665] Block path traversal attacks via the deployment API for MLflow-Triton --- .../mlflow_triton/deployments.py | 13 +++- qa/L0_mlflow/plugin_test.py | 59 ++++++++++++++++++- qa/L0_mlflow/test.sh | 4 +- 3 files changed, 70 insertions(+), 6 deletions(-) diff --git a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py index bebe559b9e..5c176ea7a4 100755 --- a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py +++ b/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:// :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,12 @@ 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): + if not name: + raise Exception("Please provide a model name for the deployment") + if '/' in name or '\\' in name or '..' in name: # Path traversal protection + raise Exception("Path traversal is not allowed in model's name: {}".format(name)) + def _model_exists(self, name): deploys = self.list_deployments() exists = False diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index a5d87a3c19..740868e749 100755 --- a/qa/L0_mlflow/plugin_test.py +++ b/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 @@ -37,7 +37,6 @@ import test_util as tu from mlflow.deployments import get_deploy_client - class PluginTest(tu.TestResultCollector): def setUp(self): self.client_ = get_deploy_client("triton") @@ -116,6 +115,62 @@ def test_onnx_flavor_with_files(self): filecmp.cmp(config_path, "./models/onnx_model_with_files/config.pbtxt") ) + def test_path_traversal_model_name(self): + + model_uri = "models:/onnx_model_with_files/1" + + model_name_normal = "onnx_model_123" + self.client_.create_deployment(model_name_normal, model_uri, flavor="onnx") + self.client_.delete_deployment(model_name_normal) + + model_name_empty = "" + with self.assertRaises(Exception) as e: + self.client_.create_deployment(model_name_empty, model_uri, flavor="onnx") + self.assertIn( + "Please provide a model name for the deployment", + str(e.exception), + ) + + model_name_path_traversal_1 = "/opt/sys/" + with self.assertRaises(Exception) as e: + self.client_.create_deployment( + model_name_path_traversal_1, + model_uri, + flavor="onnx") + self.assertIn( + "Path traversal is not allowed in model's name: {}".format( + model_name_path_traversal_1 + ), + str(e.exception), + ) + + model_name_path_traversal_2 = "../../etc/passwd" + with self.assertRaises(Exception) as e: + self.client_.create_deployment( + model_name_path_traversal_2, + model_uri, + flavor="onnx" + ) + self.assertIn( + "Path traversal is not allowed in model's name: {}".format( + model_name_path_traversal_2 + ), + str(e.exception), + ) + + model_name_path_traversal_3 = "onnx_model\\abc" + with self.assertRaises(Exception) as e: + self.client_.create_deployment( + model_name_path_traversal_3, + model_uri, + flavor="onnx", + ) + self.assertIn( + "Path traversal is not allowed in model's name: {}".format( + model_name_path_traversal_3 + ), + str(e.exception), + ) if __name__ == "__main__": unittest.main() diff --git a/qa/L0_mlflow/test.sh b/qa/L0_mlflow/test.sh index 1bf7bc18f8..fab1671a01 100755 --- a/qa/L0_mlflow/test.sh +++ b/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***" From 3d1b780c3fd62b1b6736ea63ac9454be8ab7205c Mon Sep 17 00:00:00 2001 From: Mudit Aggarwal Date: Thu, 19 Feb 2026 11:07:28 -0800 Subject: [PATCH 2/9] Addressed the PR feedback --- .../mlflow_triton/deployments.py | 2 +- qa/L0_mlflow/plugin_test.py | 20 +------------------ 2 files changed, 2 insertions(+), 20 deletions(-) diff --git a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py index 5c176ea7a4..5943643907 100755 --- a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py +++ b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py @@ -519,7 +519,7 @@ def _validate_flavor(self, flavor): def _validate_model_name(self, name): if not name: raise Exception("Please provide a model name for the deployment") - if '/' in name or '\\' in name or '..' in name: # Path traversal protection + if '/' in name or '..' in name: # Path traversal protection raise Exception("Path traversal is not allowed in model's name: {}".format(name)) def _model_exists(self, name): diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index 740868e749..3b06146c2c 100755 --- a/qa/L0_mlflow/plugin_test.py +++ b/qa/L0_mlflow/plugin_test.py @@ -115,14 +115,10 @@ def test_onnx_flavor_with_files(self): filecmp.cmp(config_path, "./models/onnx_model_with_files/config.pbtxt") ) - def test_path_traversal_model_name(self): + def test_invalid_path_traversal_model_name(self): model_uri = "models:/onnx_model_with_files/1" - model_name_normal = "onnx_model_123" - self.client_.create_deployment(model_name_normal, model_uri, flavor="onnx") - self.client_.delete_deployment(model_name_normal) - model_name_empty = "" with self.assertRaises(Exception) as e: self.client_.create_deployment(model_name_empty, model_uri, flavor="onnx") @@ -158,19 +154,5 @@ def test_path_traversal_model_name(self): str(e.exception), ) - model_name_path_traversal_3 = "onnx_model\\abc" - with self.assertRaises(Exception) as e: - self.client_.create_deployment( - model_name_path_traversal_3, - model_uri, - flavor="onnx", - ) - self.assertIn( - "Path traversal is not allowed in model's name: {}".format( - model_name_path_traversal_3 - ), - str(e.exception), - ) - if __name__ == "__main__": unittest.main() From fb05d7b097b7a5383114463451176be7c80ff8ba Mon Sep 17 00:00:00 2001 From: Mudit Aggarwal Date: Thu, 19 Feb 2026 13:52:03 -0800 Subject: [PATCH 3/9] Making pre-commit happy --- deploy/mlflow-triton-plugin/mlflow_triton/deployments.py | 4 +++- qa/L0_mlflow/plugin_test.py | 9 +++------ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py index 5943643907..2d52185dbc 100755 --- a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py +++ b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py @@ -520,7 +520,9 @@ def _validate_model_name(self, name): if not name: raise Exception("Please provide a model name for the deployment") if '/' in name or '..' in name: # Path traversal protection - raise Exception("Path traversal is not allowed in model's name: {}".format(name)) + raise Exception( + "Path traversal is not allowed in model's name: {}".format(name) + ) def _model_exists(self, name): deploys = self.list_deployments() diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index 3b06146c2c..8ed42ba843 100755 --- a/qa/L0_mlflow/plugin_test.py +++ b/qa/L0_mlflow/plugin_test.py @@ -130,9 +130,8 @@ def test_invalid_path_traversal_model_name(self): model_name_path_traversal_1 = "/opt/sys/" with self.assertRaises(Exception) as e: self.client_.create_deployment( - model_name_path_traversal_1, - model_uri, - flavor="onnx") + model_name_path_traversal_1, model_uri, flavor="onnx" + ) self.assertIn( "Path traversal is not allowed in model's name: {}".format( model_name_path_traversal_1 @@ -143,9 +142,7 @@ def test_invalid_path_traversal_model_name(self): model_name_path_traversal_2 = "../../etc/passwd" with self.assertRaises(Exception) as e: self.client_.create_deployment( - model_name_path_traversal_2, - model_uri, - flavor="onnx" + model_name_path_traversal_2, model_uri, flavor="onnx" ) self.assertIn( "Path traversal is not allowed in model's name: {}".format( From 88c5f478a860e03fe4b20cd95180769449173736 Mon Sep 17 00:00:00 2001 From: Mudit Aggarwal Date: Thu, 19 Feb 2026 13:56:03 -0800 Subject: [PATCH 4/9] Making pre-commit happier --- deploy/mlflow-triton-plugin/mlflow_triton/deployments.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py index 2d52185dbc..923c554226 100755 --- a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py +++ b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py @@ -519,7 +519,7 @@ def _validate_flavor(self, flavor): def _validate_model_name(self, name): if not name: raise Exception("Please provide a model name for the deployment") - if '/' in name or '..' in name: # Path traversal protection + if "/" in name or ".." in name: # Path traversal protection raise Exception( "Path traversal is not allowed in model's name: {}".format(name) ) From 65ca276019d3ffb74f891f33b7a36c930768d5cd Mon Sep 17 00:00:00 2001 From: Mudit Aggarwal Date: Thu, 19 Feb 2026 14:13:44 -0800 Subject: [PATCH 5/9] Making pre-commit super happy --- qa/L0_mlflow/plugin_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index 8ed42ba843..27e3f6b6c5 100755 --- a/qa/L0_mlflow/plugin_test.py +++ b/qa/L0_mlflow/plugin_test.py @@ -37,6 +37,7 @@ import test_util as tu from mlflow.deployments import get_deploy_client + class PluginTest(tu.TestResultCollector): def setUp(self): self.client_ = get_deploy_client("triton") @@ -116,7 +117,6 @@ def test_onnx_flavor_with_files(self): ) def test_invalid_path_traversal_model_name(self): - model_uri = "models:/onnx_model_with_files/1" model_name_empty = "" @@ -151,5 +151,6 @@ def test_invalid_path_traversal_model_name(self): str(e.exception), ) + if __name__ == "__main__": unittest.main() From b84b9a64b7cbd8500c0c19dd571420c3b42f4f72 Mon Sep 17 00:00:00 2001 From: Mudit Aggarwal Date: Mon, 23 Feb 2026 11:23:29 -0800 Subject: [PATCH 6/9] Changed the test name --- qa/L0_mlflow/plugin_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index 27e3f6b6c5..ee47337a1f 100755 --- a/qa/L0_mlflow/plugin_test.py +++ b/qa/L0_mlflow/plugin_test.py @@ -116,7 +116,7 @@ def test_onnx_flavor_with_files(self): filecmp.cmp(config_path, "./models/onnx_model_with_files/config.pbtxt") ) - def test_invalid_path_traversal_model_name(self): + def test_invalid_model_name(self): model_uri = "models:/onnx_model_with_files/1" model_name_empty = "" From d329827576c99cd2083d9af6482b74882f1b409e Mon Sep 17 00:00:00 2001 From: Mudit Aggarwal Date: Thu, 26 Feb 2026 00:21:54 -0800 Subject: [PATCH 7/9] Added more checks for model name --- .../mlflow_triton/deployments.py | 8 +- qa/L0_mlflow/plugin_test.py | 87 ++++++++++++------- 2 files changed, 61 insertions(+), 34 deletions(-) diff --git a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py index 923c554226..3082ae2597 100755 --- a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py +++ b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py @@ -517,9 +517,11 @@ def _validate_flavor(self, flavor): raise Exception("{} model flavor not supported by Triton".format(flavor)) def _validate_model_name(self, name): - if not name: - raise Exception("Please provide a model name for the deployment") - if "/" in name or ".." in name: # Path traversal protection + # Check if the model name is empty or only contains whitespace, tabs, or newlines + if name.strip() == "": + raise Exception("Please provide a valid model name for the deployment") + # Path traversal protection + if "/" in name or name == "..": raise Exception( "Path traversal is not allowed in model's name: {}".format(name) ) diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index ee47337a1f..18e0284b2e 100755 --- a/qa/L0_mlflow/plugin_test.py +++ b/qa/L0_mlflow/plugin_test.py @@ -45,7 +45,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 @@ -116,40 +116,65 @@ def test_onnx_flavor_with_files(self): filecmp.cmp(config_path, "./models/onnx_model_with_files/config.pbtxt") ) - def test_invalid_model_name(self): - model_uri = "models:/onnx_model_with_files/1" - - model_name_empty = "" - with self.assertRaises(Exception) as e: - self.client_.create_deployment(model_name_empty, model_uri, flavor="onnx") - self.assertIn( - "Please provide a model name for the deployment", - str(e.exception), - ) + def test_model_name(self): + EMPTY_MODEL_NAMES = [ + "", + " ", + " ", + "\t\n", + ] + INVALID_PATH_TRAVERSAL_NAMES = [ + "/opt/sys/", + "../../etc/passwd", + "../outside/repo", + "test_models/../identity_py", + "..", + ] + VALID_MODEL_NAMES = [ + "model123", + "model OAI", + "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( + "Please provide a valid model name for the deployment", + str(e.exception), + ) - model_name_path_traversal_1 = "/opt/sys/" - with self.assertRaises(Exception) as e: - self.client_.create_deployment( - model_name_path_traversal_1, model_uri, flavor="onnx" + 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), ) - self.assertIn( - "Path traversal is not allowed in model's name: {}".format( - model_name_path_traversal_1 - ), - str(e.exception), - ) - model_name_path_traversal_2 = "../../etc/passwd" - with self.assertRaises(Exception) as e: - self.client_.create_deployment( - model_name_path_traversal_2, model_uri, flavor="onnx" + for model_name in VALID_MODEL_NAMES: + # _validate_model_name should not raise an exception + model_uri = f"models:/{model_name}/1" + with self.assertRaises(Exception) as e: + self.client_.create_deployment(model_name, model_uri, flavor="onnx") + self.assertNotIn( + "Please provide a valid model name for the deployment", + str(e.exception), + ) + self.assertNotIn( + f"Path traversal is not allowed in model's name: {model_name}", + str(e.exception), + ) + self.assertIn( + f"Registered Model with name={model_name} not found", + str(e.exception), ) - self.assertIn( - "Path traversal is not allowed in model's name: {}".format( - model_name_path_traversal_2 - ), - str(e.exception), - ) if __name__ == "__main__": From df77b0b247af31ef6c944408b39ffc2c3808f3ae Mon Sep 17 00:00:00 2001 From: Mudit Aggarwal Date: Mon, 2 Mar 2026 23:28:44 -0800 Subject: [PATCH 8/9] Incorporated feedback --- .../mlflow_triton/deployments.py | 4 +- qa/L0_mlflow/plugin_test.py | 39 +++++++++---------- 2 files changed, 21 insertions(+), 22 deletions(-) diff --git a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py index 3082ae2597..ee56faf4af 100755 --- a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py +++ b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py @@ -519,7 +519,9 @@ def _validate_flavor(self, 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("Please provide a valid model name for the deployment") + raise Exception( + "Model name cannot be empty. Please enter a valid name to deploy." + ) # Path traversal protection if "/" in name or name == "..": raise Exception( diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index 18e0284b2e..e225ed5881 100755 --- a/qa/L0_mlflow/plugin_test.py +++ b/qa/L0_mlflow/plugin_test.py @@ -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 @@ -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" @@ -121,7 +119,11 @@ def test_model_name(self): "", " ", " ", - "\t\n", + "\n", + "\t", + "\r", + "\v", + "\f", ] INVALID_PATH_TRAVERSAL_NAMES = [ "/opt/sys/", @@ -132,7 +134,7 @@ def test_model_name(self): ] VALID_MODEL_NAMES = [ "model123", - "model OAI", + # "model OAI", TRI-769: Fix this test case "model.version1", "...", "..my_model", @@ -145,7 +147,7 @@ def test_model_name(self): with self.assertRaises(Exception) as e: self.client_.create_deployment(model_name, model_uri, flavor="onnx") self.assertIn( - "Please provide a valid model name for the deployment", + "Model name cannot be empty. Please enter a valid name to deploy.", str(e.exception), ) @@ -159,23 +161,18 @@ def test_model_name(self): ) for model_name in VALID_MODEL_NAMES: - # _validate_model_name should not raise an exception - model_uri = f"models:/{model_name}/1" - with self.assertRaises(Exception) as e: - self.client_.create_deployment(model_name, model_uri, flavor="onnx") - self.assertNotIn( - "Please provide a valid model name for the deployment", - str(e.exception), + model = onnx.load( + "./mlflow-triton-plugin/examples/onnx_float32_int32_int32/1/model.onnx" ) - self.assertNotIn( - f"Path traversal is not allowed in model's name: {model_name}", - str(e.exception), - ) - self.assertIn( - f"Registered Model with name={model_name} not found", - str(e.exception), + + # 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() From 41eb77243602fb2355351746a1f77d4cfc21000a Mon Sep 17 00:00:00 2001 From: Mudit Aggarwal Date: Tue, 3 Mar 2026 10:23:14 -0800 Subject: [PATCH 9/9] Fix CI test --- qa/L0_mlflow/test.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qa/L0_mlflow/test.sh b/qa/L0_mlflow/test.sh index fab1671a01..5e863a8988 100755 --- a/qa/L0_mlflow/test.sh +++ b/qa/L0_mlflow/test.sh @@ -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***"