Skip to content

Commit f832dd2

Browse files
author
Pravali Uppugunduri
committed
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.
1 parent 6a174f4 commit f832dd2

File tree

3 files changed

+14
-3
lines changed

3 files changed

+14
-3
lines changed

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3075,18 +3075,20 @@ def _prepare_for_triton(self):
30753075
export_path.mkdir(parents=True)
30763076

30773077
if self.model:
3078-
self.secret_key = "dummy secret key for onnx backend"
3078+
self.secret_key = generate_secret_key()
30793079

30803080
if self.framework == Framework.PYTORCH:
30813081
self._export_pytorch_to_onnx(
30823082
export_path=export_path, model=self.model, schema_builder=self.schema_builder
30833083
)
3084+
self._hmac_signing()
30843085
return
30853086

30863087
if self.framework == Framework.TENSORFLOW:
30873088
self._export_tf_to_onnx(
30883089
export_path=export_path, model=self.model, schema_builder=self.schema_builder
30893090
)
3091+
self._hmac_signing()
30903092
return
30913093

30923094
raise ValueError("%s is not supported" % self.framework)

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,13 @@ 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")
2930
with open(str(serve_path), mode="rb") as f:
30-
inference_spec, schema_builder = cloudpickle.load(f)
31+
buffer = f.read()
32+
perform_integrity_check(buffer=buffer, metadata_path=str(metadata_path))
3133

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

3437
self.inference_spec = inference_spec
3538
self.schema_builder = schema_builder

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,12 @@ def _start_triton_server(
3636
env_vars: dict,
3737
):
3838
"""Placeholder docstring"""
39+
if not isinstance(secret_key, str) or not secret_key.strip():
40+
raise ValueError(
41+
"A valid secret key is required for Triton deployments. "
42+
"The secret key must be a non-empty string generated by generate_secret_key(). "
43+
f"Received: {type(secret_key).__name__}"
44+
)
3945
self.container_name = "triton" + uuid.uuid1().hex
4046
model_repository = model_path + "/model_repository"
4147
env_vars.update(

0 commit comments

Comments
 (0)