fix: Add HMAC integrity verification for Triton inference handler

- Add HMAC integrity check before pickle deserialization in TritonPythonModel.initialize()
- Replace hardcoded secret key with generate_secret_key() in _prepare_for_triton() ONNX path
- Add _hmac_signing() after ONNX export for both PyTorch and TensorFlow frameworks
- Add secret key validation in _start_triton_server() to reject None/empty keys

Fixes RCE vulnerabilities in Triton handler by aligning with HMAC verification
patterns used by TorchServe, MMS, TF Serving, and SMD handlers.

Author: Pravali Uppugunduri

sagemaker-serve/src/sagemaker/serve/model_builder_utils.py

@@ -2879,10 +2879,12 @@ def _is_gpu_instance(self, instance_type: str) -> bool:
         return instance_family in GPU_INSTANCE_FAMILIES
 
     def _save_inference_spec(self) -> None:
-        """Save inference specification to pickle file."""
+        """Save inference specification or model to pickle file."""
+        pkl_path = Path(self.model_path).joinpath("model_repository").joinpath("model")
         if self.inference_spec:
-            pkl_path = Path(self.model_path).joinpath("model_repository").joinpath("model")
             save_pkl(pkl_path, (self.inference_spec, self.schema_builder))
+        elif self.model:
+            save_pkl(pkl_path, (self.model, self.schema_builder))
 
     def _hmac_signing(self):
         """Perform HMAC signing on picke file for integrity check"""
@@ -3075,18 +3077,20 @@ def _prepare_for_triton(self):
             export_path.mkdir(parents=True)
 
         if self.model:
-            self.secret_key = "dummy secret key for onnx backend"
+            self.secret_key = generate_secret_key()
 
             if self.framework == Framework.PYTORCH:
                 self._export_pytorch_to_onnx(
                     export_path=export_path, model=self.model, schema_builder=self.schema_builder
                 )
+                self._hmac_signing()
                 return
 
             if self.framework == Framework.TENSORFLOW:
                 self._export_tf_to_onnx(
                     export_path=export_path, model=self.model, schema_builder=self.schema_builder
                 )
+                self._hmac_signing()
                 return
 
             raise ValueError("%s is not supported" % self.framework)

sagemaker-serve/src/sagemaker/serve/model_server/triton/model.py

@@ -26,10 +26,13 @@ def auto_complete_config(auto_complete_model_config):
     def initialize(self, args: dict) -> None:
         """Placeholder docstring"""
         serve_path = Path(TRITON_MODEL_DIR).joinpath("serve.pkl")
+        metadata_path = Path(TRITON_MODEL_DIR).joinpath("metadata.json")
         with open(str(serve_path), mode="rb") as f:
-            inference_spec, schema_builder = cloudpickle.load(f)
+            buffer = f.read()
+        perform_integrity_check(buffer=buffer, metadata_path=str(metadata_path))
 
-        # TODO: HMAC signing for integrity check
+        with open(str(serve_path), mode="rb") as f:
+            inference_spec, schema_builder = cloudpickle.load(f)
 
         self.inference_spec = inference_spec
         self.schema_builder = schema_builder

sagemaker-serve/src/sagemaker/serve/model_server/triton/server.py

@@ -36,6 +36,12 @@ def _start_triton_server(
         env_vars: dict,
     ):
         """Placeholder docstring"""
+        if not isinstance(secret_key, str) or not secret_key.strip():
+            raise ValueError(
+                "A valid secret key is required for Triton deployments. "
+                "The secret key must be a non-empty string generated by generate_secret_key(). "
+                f"Received: {type(secret_key).__name__}"
+            )
         self.container_name = "triton" + uuid.uuid1().hex
         model_repository = model_path + "/model_repository"
         env_vars.update(

sagemaker-serve/tests/unit/test_model_builder_utils_triton.py

@@ -81,8 +81,9 @@ class TestPrepareForTriton(unittest.TestCase):
     """Test _prepare_for_triton method."""
 
     @patch('shutil.copy2')
+    @patch.object(_ModelBuilderUtils, '_hmac_signing')
     @patch.object(_ModelBuilderUtils, '_export_pytorch_to_onnx')
-    def test_prepare_for_triton_pytorch(self, mock_export, mock_copy):
+    def test_prepare_for_triton_pytorch(self, mock_export, mock_hmac, mock_copy):
         """Test preparing PyTorch model for Triton."""
         utils = _ModelBuilderUtils()
         utils.framework = Framework.PYTORCH
@@ -94,10 +95,12 @@ def test_prepare_for_triton_pytorch(self, mock_export, mock_copy):
             utils._prepare_for_triton()
 
             mock_export.assert_called_once()
+            mock_hmac.assert_called_once()
 
     @patch('shutil.copy2')
+    @patch.object(_ModelBuilderUtils, '_hmac_signing')
     @patch.object(_ModelBuilderUtils, '_export_tf_to_onnx')
-    def test_prepare_for_triton_tensorflow(self, mock_export, mock_copy):
+    def test_prepare_for_triton_tensorflow(self, mock_export, mock_hmac, mock_copy):
         """Test preparing TensorFlow model for Triton."""
         utils = _ModelBuilderUtils()
         utils.framework = Framework.TENSORFLOW
@@ -109,6 +112,7 @@ def test_prepare_for_triton_tensorflow(self, mock_export, mock_copy):
             utils._prepare_for_triton()
 
             mock_export.assert_called_once()
+            mock_hmac.assert_called_once()
 
     @patch('shutil.copy2')
     @patch.object(_ModelBuilderUtils, '_generate_config_pbtxt')
@@ -259,6 +263,22 @@ def test_save_inference_spec(self):
 
             # Check that serve.pkl was created
             self.assertTrue(os.path.exists(os.path.join(pkl_path, "serve.pkl")))
+    def test_save_inference_spec_with_model(self):
+        """Test saving model when inference_spec is None."""
+        utils = _ModelBuilderUtils()
+        utils.inference_spec = None
+        utils.model = Mock()
+        utils.schema_builder = Mock()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            utils.model_path = tmpdir
+            pkl_path = os.path.join(tmpdir, "model_repository", "model")
+            os.makedirs(pkl_path, exist_ok=True)
+
+            utils._save_inference_spec()
+
+            # Check that serve.pkl was created for the model path
+            self.assertTrue(os.path.exists(os.path.join(pkl_path, "serve.pkl")))
 
 
 class TestHMACSignin(unittest.TestCase):