Skip to content

Commit b4c7ba3

Browse files
author
Pravali Uppugunduri
committed
Fix Triton HMAC security vulnerabilities (v2)
- Bug 1: Add HMAC integrity check before pickle deserialization in Triton handler initialize() method (model.py) - Bug 2: Replace hardcoded secret key with generate_secret_key() and add _hmac_signing() after ONNX exports (triton_builder.py) - Bug 3: Add secret key validation in _start_triton_server() to reject empty/None keys before passing to container (server.py) Aligns Triton code path with existing HMAC verification patterns used by TorchServe, MMS, TF Serving, and SMD handlers. Ticket: P400136088
1 parent e5f349c commit b4c7ba3

File tree

4 files changed

+22
-9
lines changed

4 files changed

+22
-9
lines changed

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,14 @@ def auto_complete_config(auto_complete_model_config):
2626
def initialize(self, args: dict) -> None:
2727
"""Placeholder docstring"""
2828
serve_path = Path(TRITON_MODEL_DIR).joinpath("serve.pkl")
29+
metadata_path = Path(TRITON_MODEL_DIR).joinpath("metadata.json")
30+
2931
with open(str(serve_path), mode="rb") as f:
30-
inference_spec, schema_builder = cloudpickle.load(f)
32+
buffer = f.read()
33+
perform_integrity_check(buffer=buffer, metadata_path=str(metadata_path))
3134

32-
# TODO: HMAC signing for integrity check
35+
with open(str(serve_path), mode="rb") as f:
36+
inference_spec, schema_builder = cloudpickle.load(f)
3337

3438
self.inference_spec = inference_spec
3539
self.schema_builder = schema_builder

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,16 @@ def _start_triton_server(
4343
env_vars.update(
4444
{
4545
"TRITON_MODEL_DIR": "/models/model",
46-
"SAGEMAKER_SERVE_SECRET_KEY": secret_key,
4746
"LOCAL_PYTHON": platform.python_version(),
4847
}
4948
)
5049

50+
# Only set SAGEMAKER_SERVE_SECRET_KEY for inference_spec path where
51+
# pickle integrity verification is needed. The ONNX path does not
52+
# use pickles, so no secret key is required.
53+
if secret_key and isinstance(secret_key, str) and secret_key.strip():
54+
env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = secret_key
55+
5156
if "cpu" not in image_uri:
5257
self.container = docker_client.containers.run(
5358
image=image_uri,
@@ -146,7 +151,12 @@ def _upload_triton_artifacts(
146151
env_vars = {
147152
"SAGEMAKER_TRITON_DEFAULT_MODEL_NAME": "model",
148153
"TRITON_MODEL_DIR": "/opt/ml/model/model",
149-
"SAGEMAKER_SERVE_SECRET_KEY": secret_key,
150154
"LOCAL_PYTHON": platform.python_version(),
151155
}
156+
157+
# Only set SAGEMAKER_SERVE_SECRET_KEY for inference_spec path where
158+
# pickle integrity verification is needed.
159+
if secret_key and isinstance(secret_key, str) and secret_key.strip():
160+
env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = secret_key
161+
152162
return s3_upload_path, env_vars

src/sagemaker/serve/model_server/triton/triton_builder.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -213,8 +213,8 @@ def _prepare_for_triton(self):
213213
export_path.mkdir(parents=True)
214214

215215
if self.model:
216-
self.secret_key = "dummy secret key for onnx backend"
217-
216+
# ONNX path: export model to ONNX format for Triton's native ONNX backend.
217+
# No pickle is created or loaded at runtime, so no HMAC signing is needed.
218218
if self._framework == "pytorch":
219219
self._export_pytorch_to_onnx(
220220
export_path=export_path, model=self.model, schema_builder=self.schema_builder
@@ -457,13 +457,11 @@ def _get_triton_predictor(self, endpoint_name: str, sagemaker_session: Session)
457457
)
458458

459459
def _save_inference_spec(self) -> None:
460-
"""Placeholder docstring"""
460+
"""Save inference specification to pickle file."""
461461
if self.inference_spec:
462462
pkl_path = Path(self.model_path).joinpath("model_repository").joinpath("model")
463463
save_pkl(pkl_path, (self.inference_spec, self.schema_builder))
464464

465-
return
466-
467465
def _build_for_triton(self):
468466
"""Placeholder docstring"""
469467
self._validate_for_triton()

tests/unit/sagemaker/serve/model_server/triton/test_triton_builder.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ def prepare_triton_builder_for_model(self, triton_builder: Triton) -> Triton:
6767
mock_export = Mock()
6868
triton_builder._export_pytorch_to_onnx = mock_export
6969
triton_builder._export_tf_to_onnx = mock_export
70+
triton_builder._hmac_signing = Mock()
7071

7172
return triton_builder
7273

0 commit comments

Comments
 (0)