Skip to content

Commit 3bf5912

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 3bf5912

File tree

3 files changed

+16
-3
lines changed

3 files changed

+16
-3
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: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,13 @@ def _start_triton_server(
3838
env_vars: dict,
3939
):
4040
"""Placeholder docstring"""
41+
if not isinstance(secret_key, str) or not secret_key.strip():
42+
raise ValueError(
43+
"A valid secret key is required for Triton deployments. "
44+
"The secret key must be a non-empty string generated by generate_secret_key(). "
45+
f"Received: {type(secret_key).__name__}"
46+
)
47+
4148
self.container_name = "triton" + uuid.uuid1().hex
4249
model_repository = model_path + "/model_repository"
4350
env_vars.update(

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,18 +213,20 @@ 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"
216+
self.secret_key = generate_secret_key()
217217

218218
if self._framework == "pytorch":
219219
self._export_pytorch_to_onnx(
220220
export_path=export_path, model=self.model, schema_builder=self.schema_builder
221221
)
222+
self._hmac_signing()
222223
return
223224

224225
if self._framework == "tensorflow":
225226
self._export_tf_to_onnx(
226227
export_path=export_path, model=self.model, schema_builder=self.schema_builder
227228
)
229+
self._hmac_signing()
228230
return
229231

230232
raise ValueError("%s is not supported" % self._framework)

0 commit comments

Comments
 (0)