Skip to content

fix: Replace HMAC with SHA based integrity checks in model builder. - #5656

Merged
nargokul merged 1 commit into
aws:master-v2from
pravali96:fix/triton-hmac-security-v2-all-bugs
Apr 1, 2026
Merged

fix: Replace HMAC with SHA based integrity checks in model builder.#5656
nargokul merged 1 commit into
aws:master-v2from
pravali96:fix/triton-hmac-security-v2-all-bugs

Conversation

@pravali96

@pravali96 pravali96 commented Mar 19, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Issue #, if available:

Description of changes:
Replace HMAC with SHA based integrity checks in model builder.

Issue

Fix

Switch from HMAC-SHA256 (requires a secret key) to plain SHA-256 (no key needed)

Changes

check_integrity.py

  • Removed 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

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

@pravali96
pravali96 requested a review from a team as a code owner March 19, 2026 22:58
@pravali96
pravali96 requested a review from mollyheamazon March 19, 2026 22:58
@pravali96
pravali96 force-pushed the fix/triton-hmac-security-v2-all-bugs branch from 6e9283e to 8fe96c8 Compare March 20, 2026 17:48
@pravali96
pravali96 force-pushed the fix/triton-hmac-security-v2-all-bugs branch from c8d05ec to 079e1f2 Compare March 20, 2026 19:08
@pravali96 pravali96 changed the title fix: Security fixes for Triton HMAC key exposure and missing integrit… fix: fixes for Triton HMAC key exposure and missing integrit… Mar 31, 2026
@pravali96 pravali96 changed the title fix: fixes for Triton HMAC key exposure and missing integrit… fix: Replace HMAC with SHA based integrity checks in model builder. Mar 31, 2026
nargokul
nargokul requested changes Mar 31, 2026
View reviewed changes

@nargokul nargokul left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please fix integ test failures.

Comment thread src/sagemaker/serve/model_server/multi_model_server/prepare.py Outdated
metadata.write(_MetaData(hash_value).to_json())

return secret_key
return ""

@nargokul nargokul Mar 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: We need not return anything here

@pravali96 pravali96 Mar 31, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The 11 failing tests are in JumpStart, not model builder. They are unrelated to these changes. is that a blocker?

@pravali96 pravali96 Mar 31, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's not needed for long-term but its better to keep return "" for backward compatibility in this security fix. Callers (self.secret_key = prepare_for_mms(...)) still assign this return value, so it cannot be none.

@pravali96 pravali96 Apr 1, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the return statements.

@pravali96
pravali96 force-pushed the fix/triton-hmac-security-v2-all-bugs branch 2 times, most recently from 079e1f2 to 334346a Compare March 31, 2026 22:15
@pravali96
pravali96 force-pushed the fix/triton-hmac-security-v2-all-bugs branch from 95dc333 to a809727 Compare March 31, 2026 22:47
Pravali Uppugunduri
fix: Security fixes for Triton HMAC key exposure and missing integrit…
d10bf43 
…y check (v2)

Backport of v3 security fixes for P400136088 and V2146375387.

1. check_integrity.py: Switch from HMAC-SHA256 to plain SHA-256.
   Remove generate_secret_key, remove env var dependency.

2. triton/model.py: Add integrity check in initialize() BEFORE
   cloudpickle deserialization.

3. triton/server.py: Remove SAGEMAKER_SERVE_SECRET_KEY from
   container environment variables.

4. triton/triton_builder.py: Remove hardcoded dummy secret key
   for ONNX path. Rename _hmac_signing to _compute_integrity_hash.
   Use plain SHA-256.

5. All prepare.py files (torchserve, mms, tf_serving, smd):
   Remove generate_secret_key usage, switch to plain SHA-256.
@pravali96
pravali96 force-pushed the fix/triton-hmac-security-v2-all-bugs branch from 032f58f to d10bf43 Compare April 1, 2026 00:01
@nargokul
nargokul merged commit 48f45da into aws:master-v2 Apr 1, 2026
9 of 11 checks passed
@pravali96
pravali96 deleted the fix/triton-hmac-security-v2-all-bugs branch April 1, 2026 17:00
@mohamedzeidan2021 mohamedzeidan2021 mentioned this pull request Oct 9, 2026
Closed
@mohamedzeidan2021 mohamedzeidan2021 mentioned this pull request Oct 23, 2026
Merged
3 of 5 tasks
mohamedzeidan2021 added a commit that referenced this pull request Oct 23, 2026
@mohamedzeidan2021 Mohamed Zeidan
fix(serve): decouple the SMD orchestrator entrypoint from the contain…
4c07a98 
…er's sagemaker (#6200) (#6534)

* fix(serve): decouple the SMD orchestrator entrypoint from the container's sagemaker

Deploying a CustomOrchestrator with ModelBuilder fails its container health
check at startup with `AttributeError: 'NoneType' object has no attribute
'encode'`, so the Inference Component never reaches InService.

The generated `code/inference.py` did `from sagemaker.serve.validations.
check_integrity import perform_integrity_check`, which resolves against the
*container's* installed sagemaker rather than the SDK that produced the
artifacts. `sagemaker-distribution.json` maps SMD inference to 3.2.0 only, and
that image bundles sagemaker 2.245.0, whose `perform_integrity_check` still
reads `SAGEMAKER_SERVE_SECRET_KEY` and calls `.encode()` on it. The host no
longer sets that variable, and must not: the secret key was removed on purpose
in #5654 / #5656 as the fix for CVE-2026-8596, where it was exposed in
plaintext through the Describe APIs. So the missing `return` in
`prepare_for_smd()` is not the defect, and restoring it would reintroduce the
vulnerability -- the version coupling is the defect.

Inline the digest comparison in the entrypoint instead. It is the same plain
SHA-256 check against `metadata.json` that the host writes, so the integrity
posture is unchanged, but the script no longer imports anything from sagemaker
and therefore cannot drift against whatever version the image carries.

A second, independent mismatch sat behind the first: `_save_model_inference_
spec` pickled `(inference_spec, schema_builder)`, and a v3 SchemaBuilder
pickles with references to `sagemaker.core.*` serializer modules that a 2.x
install cannot import, so `cloudpickle.load` failed even once the integrity
check passed. The orchestrator entrypoint discards the second element of that
tuple, so pickle `(inference_spec, None)` for custom orchestrators.

Also drop the six `self.secret_key = prepare_for_*(...)` assignments and their
`# fix out of scope` / pylint-disable comments, and correct the `-> str`
annotations to `-> None` on all four `prepare_for_*` functions. Those functions
record a digest and have had nothing to return since the CVE fix; the lying
annotation is what led two earlier PRs (#6221, #6222) to try to resurrect the
secret key. Every builder already initializes `self.secret_key = ""` before
calling them, so this is not a behavior change.

Verified by generating the artifacts with this branch and executing the
resulting `inference.py` under a sagemaker 2.245.0 interpreter -- the version
bundled in the image the SDK selects. Before: the reported traceback. After:
the orchestrator loads and serves. Tampered `serve.pkl` and missing
`metadata.json` are still rejected.

Closes #6200

* test(serve): import the SMD entrypoint instead of exec-ing its source

Coverage could not attribute the entrypoint's lines when the test exec'd its
source text, so the new code read as unhit (codecov/patch 64%) despite every
line being exercised. The module reads its code directory from
SAGEMAKER_INFERENCE_CODE_DIRECTORY, so it can be imported for real against a
prepared temp dir and reloaded per case, which is both measurable and closer to
how the model server loads it.

Also cover the awaited branch of handler() with an AsyncCustomOrchestrator;
custom_execution_inference.py is now at 100%.

---------

Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com>

This branch was previously deployed

1 inactive deployment
auto-approve — d10bf43c Deployed Apr 1, 2026 by pravali96 via wait-for-approval #119
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@gwang111 gwang111 gwang111 approved these changes
@nargokul nargokul nargokul approved these changes
@mollyheamazon mollyheamazon Awaiting requested review from mollyheamazon mollyheamazon is a code owner automatically assigned from aws/sagemaker-ml-frameworks

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@pravali96 @gwang111 @nargokul