fix: Add HMAC integrity verification for Triton inference handler

Addresses P400136088 Bug 1 and V2146375387 (Triton path).

Three changes:

1. check_integrity.py: Switch from HMAC-SHA256 to plain SHA-256.
   - Remove generate_secret_key() — no longer needed
   - compute_hash() now uses hashlib.sha256() instead of hmac.new()
   - perform_integrity_check() no longer reads SAGEMAKER_SERVE_SECRET_KEY
     from environment

2. triton/model.py: Add integrity check in initialize() BEFORE
   cloudpickle deserialization. Previously the handler called
   cloudpickle.load() with no verification (acknowledged by a TODO
   comment). Now reads the file into a buffer, runs
   perform_integrity_check(), then deserializes with cloudpickle.loads().

3. triton/server.py: Remove SAGEMAKER_SERVE_SECRET_KEY from container
   environment variables in both local and SageMaker deployment modes.
   The key is no longer needed since integrity checking uses plain
   SHA-256.

4. model_builder_utils.py: Update _hmac_signing() to use plain SHA-256
   and stop generating/storing a secret key. Remove generate_secret_key
   import.

The integrity check still detects accidental corruption of model
artifacts in S3. The HMAC was providing a false sense of security since
the key was exposed via DescribeModel/DescribeEndpointConfig APIs.

Author: Pravali Uppugunduri

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

@@ -131,7 +131,6 @@ def build(self):
 from sagemaker.serve.detector.pickler import save_pkl
 from sagemaker.serve.builder.requirements_manager import RequirementsManager
 from sagemaker.serve.validations.check_integrity import (
-    generate_secret_key,
     compute_hash,
 )
 from sagemaker.core.remote_function.core.serialization import _MetaData
@@ -2884,20 +2883,17 @@ def _save_inference_spec(self) -> None:
             pkl_path = Path(self.model_path).joinpath("model_repository").joinpath("model")
             save_pkl(pkl_path, (self.inference_spec, self.schema_builder))
 
-    def _hmac_signing(self):
-        """Perform HMAC signing on picke file for integrity check"""
-        secret_key = generate_secret_key()
+    def _compute_integrity_hash(self):
+        """Compute SHA-256 hash of serve.pkl and store in metadata.json for integrity check."""
         pkl_path = Path(self.model_path).joinpath("model_repository").joinpath("model")
 
         with open(str(pkl_path.joinpath("serve.pkl")), "rb") as f:
             buffer = f.read()
-        hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+        hash_value = compute_hash(buffer=buffer)
 
         with open(str(pkl_path.joinpath("metadata.json")), "wb") as metadata:
             metadata.write(_MetaData(hash_value).to_json())
 
-        self.secret_key = secret_key
-
     def _generate_config_pbtxt(self, pkl_path: Path):
         """Generate Triton config.pbtxt file."""
         config_path = pkl_path.joinpath("config.pbtxt")
@@ -3100,7 +3096,7 @@ def _prepare_for_triton(self):
 
             self._pack_conda_env(pkl_path=pkl_path)
 
-            self._hmac_signing()
+            self._compute_integrity_hash()
 
             return
 

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

@@ -26,10 +26,14 @@ 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")
-        with open(str(serve_path), mode="rb") as f:
-            inference_spec, schema_builder = cloudpickle.load(f)
+        metadata_path = Path(TRITON_MODEL_DIR).joinpath("metadata.json")
 
-        # TODO: HMAC signing for integrity check
+        # Integrity check BEFORE deserialization to prevent RCE via malicious pickle
+        with open(str(serve_path), "rb") as f:
+            buffer = f.read()
+        perform_integrity_check(buffer=buffer, metadata_path=metadata_path)
+
+        inference_spec, schema_builder = cloudpickle.loads(buffer)
 
         self.inference_spec = inference_spec
         self.schema_builder = schema_builder

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

@@ -41,7 +41,6 @@ def _start_triton_server(
         env_vars.update(
             {
                 "TRITON_MODEL_DIR": "/models/model",
-                "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
                 "LOCAL_PYTHON": platform.python_version(),
             }
         )
@@ -133,7 +132,6 @@ def _upload_triton_artifacts(
         env_vars = {
             "SAGEMAKER_TRITON_DEFAULT_MODEL_NAME": "model",
             "TRITON_MODEL_DIR": "/opt/ml/model/model",
-            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
             "LOCAL_PYTHON": platform.python_version(),
         }
         return s3_upload_path, env_vars

sagemaker-serve/src/sagemaker/serve/validations/check_integrity.py

@@ -1,29 +1,21 @@
-"""Validates the integrity of pickled file with HMAC signing."""
+"""Validates the integrity of pickled file with SHA-256 hash."""
 
 from __future__ import absolute_import
-import secrets
 import hmac
 import hashlib
-import os
 from pathlib import Path
 
 from sagemaker.core.remote_function.core.serialization import _MetaData
 
 
-def generate_secret_key(nbytes: int = 32) -> str:
-    """Generates secret key"""
-    return secrets.token_hex(nbytes)
-
-
-def compute_hash(buffer: bytes, secret_key: str) -> str:
-    """Compute hash value using HMAC"""
-    return hmac.new(secret_key.encode(), msg=buffer, digestmod=hashlib.sha256).hexdigest()
+def compute_hash(buffer: bytes) -> str:
+    """Compute SHA-256 hash of the given buffer."""
+    return hashlib.sha256(buffer).hexdigest()
 
 
 def perform_integrity_check(buffer: bytes, metadata_path: Path):
-    """Validates the integrity of bytes by comparing the hash value"""
-    secret_key = os.environ.get("SAGEMAKER_SERVE_SECRET_KEY")
-    actual_hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+    """Validates the integrity of bytes by comparing the hash value."""
+    actual_hash_value = compute_hash(buffer=buffer)
 
     if not Path.exists(metadata_path):
         raise ValueError("Path to metadata.json does not exist")

sagemaker-serve/tests/unit/test_model_builder_utils_triton.py

@@ -113,7 +113,7 @@ def test_prepare_for_triton_tensorflow(self, mock_export, mock_copy):
     @patch('shutil.copy2')
     @patch.object(_ModelBuilderUtils, '_generate_config_pbtxt')
     @patch.object(_ModelBuilderUtils, '_pack_conda_env')
-    @patch.object(_ModelBuilderUtils, '_hmac_signing')
+    @patch.object(_ModelBuilderUtils, '_compute_integrity_hash')
     def test_prepare_for_triton_inference_spec(self, mock_hmac, mock_pack, mock_config, mock_copy):
         """Test preparing inference spec for Triton."""
         utils = _ModelBuilderUtils()
@@ -262,9 +262,9 @@ def test_save_inference_spec(self):
 
 
 class TestHMACSignin(unittest.TestCase):
-    """Test _hmac_signing method."""
+    """Test _compute_integrity_hash method."""
 
-    def test_hmac_signing(self):
+    def test_compute_integrity_hash(self):
         """Test HMAC signing."""
         utils = _ModelBuilderUtils()
 
@@ -276,7 +276,7 @@ def test_hmac_signing(self):
             # Create dummy serve.pkl
             (pkl_path / "serve.pkl").write_bytes(b"dummy content")
 
-            utils._hmac_signing()
+            utils._compute_integrity_hash()
 
             # Secret key is generated, not mocked
             self.assertIsNotNone(utils.secret_key)