From bbebdd416c41a37232e35932616a8b259174c004 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Thu, 30 Jul 2026 13:25:29 -0700 Subject: [PATCH 1/2] fix(serve): Validate ECR registry host before docker login in LocalContainerMode _is_ecr_image() classified an image as ECR by substring-matching ".dkr.ecr." and ".amazonaws.com" anywhere in the URI, while _pull_image() took the docker login target from image.split("/")[0]. The two used inconsistent values, so a crafted URI such as attacker.com/x.dkr.ecr..amazonaws.com/repo passed the classifier yet caused "docker login -u AWS -p attacker.com", leaking a valid, replayable ECR authorization token to an attacker-controlled host. Parse the registry host first and validate it against a strict ECR endpoint pattern, and reuse that same validated host for docker login, so the classifier and the login target can never disagree. Add unit tests covering valid ECR (incl. China partition), attacker-crafted, and public image URIs. --- .../serve/mode/local_container_mode.py | 37 +++++- .../mode/test_local_container_mode_ecr.py | 125 ++++++++++++++++++ 2 files changed, 158 insertions(+), 4 deletions(-) create mode 100644 sagemaker-serve/tests/unit/mode/test_local_container_mode_ecr.py diff --git a/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py b/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py index 3633b850e6..acbbb3e85d 100644 --- a/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py +++ b/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py @@ -4,8 +4,9 @@ from pathlib import Path import logging import os +import re from datetime import datetime, timedelta -from typing import Dict, Type +from typing import Dict, Optional, Type import base64 import time import subprocess @@ -29,6 +30,15 @@ logger = logging.getLogger(__name__) +# Strict allowlist for a real ECR registry host: <12-digit-account>.dkr.ecr..amazonaws.com +# Also matches AWS China (.amazonaws.com.cn). This is intentionally an exact host match so that the +# "is this an ECR image?" classifier and the "which host do I docker login to?" extractor operate on +# the SAME value, preventing a crafted URI (e.g. attacker.com/x.dkr.ecr..amazonaws.com/repo) +# from being classified as ECR while login is pointed at an attacker-controlled host. +_ECR_HOST_RE = re.compile( + r"^[0-9]{12}\.dkr\.ecr\.[a-z0-9-]+\.amazonaws\.com(\.cn)?$" +) + _PING_HEALTH_CHECK_INTERVAL_SEC = 5 _PING_HEALTH_CHECK_FAIL_MSG = ( @@ -252,7 +262,10 @@ def _pull_image(self, image: str): ) decoded_token = base64.b64decode(encoded_token).decode("utf-8") username, password = decoded_token.split(":") - ecr_uri = image.split("/")[0] + # Reuse the same validated host the classifier accepted, so the ECR credential is + # only ever sent to a verified ECR endpoint (never to an attacker-controlled host + # embedded elsewhere in the image URI). + ecr_uri = self._ecr_registry_host(image) login_command = ["docker", "login", "-u", username, "-p", password, ecr_uri] result = subprocess.run(login_command, check=True, capture_output=True, text=True) @@ -277,7 +290,23 @@ def _pull_image(self, image: str): except docker.errors.APIError as e: raise RuntimeError(f"Failed to pull image '{image}': {e}") from e + def _ecr_registry_host(self, image: str) -> Optional[str]: + """Return the ECR registry host if ``image``'s registry is a valid ECR endpoint, else None. + + The registry host is the first "/"-delimited segment of the image URI -- i.e. the exact + value that would be passed to ``docker login``. It is validated against a strict ECR + endpoint pattern so that classification and the login target are derived from the same + value. + """ + host = image.split("/")[0] + return host if _ECR_HOST_RE.match(host) else None + def _is_ecr_image(self, image: str) -> bool: - """Check if image is from ECR.""" - return ".dkr.ecr." in image and ".amazonaws.com" in image + """Check if image is from ECR. + + Uses the registry host that would actually be used for ``docker login`` and validates it + against a strict ECR endpoint pattern, so the classifier and the login-host extractor can + never disagree. + """ + return self._ecr_registry_host(image) is not None diff --git a/sagemaker-serve/tests/unit/mode/test_local_container_mode_ecr.py b/sagemaker-serve/tests/unit/mode/test_local_container_mode_ecr.py new file mode 100644 index 0000000000..82dac80748 --- /dev/null +++ b/sagemaker-serve/tests/unit/mode/test_local_container_mode_ecr.py @@ -0,0 +1,125 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Tests for ECR host classification / login-target extraction in LocalContainerMode. + +These guard against the parser-confusion vulnerability where the "is this an ECR image?" +classifier and the "which host do I docker login to?" extractor disagreed, allowing a crafted +image URI to leak the ECR authorization token to an attacker-controlled host. +""" +from __future__ import absolute_import + +import unittest +from unittest.mock import MagicMock, patch + +from sagemaker.serve.mode.local_container_mode import LocalContainerMode + + +def _bare_mode(): + """Build a LocalContainerMode without running the heavy constructor. + + ``_is_ecr_image`` / ``_ecr_registry_host`` are pure w.r.t. their ``image`` argument, so we + only need a bound method, not a fully initialized instance. + """ + return LocalContainerMode.__new__(LocalContainerMode) + + +class TestEcrHostClassification(unittest.TestCase): + def setUp(self): + self.mode = _bare_mode() + + def test_valid_ecr_uri_is_recognized(self): + image = "123456789012.dkr.ecr.us-east-1.amazonaws.com/my-repo:latest" + self.assertTrue(self.mode._is_ecr_image(image)) + self.assertEqual( + self.mode._ecr_registry_host(image), + "123456789012.dkr.ecr.us-east-1.amazonaws.com", + ) + + def test_valid_ecr_uri_china_partition(self): + image = "123456789012.dkr.ecr.cn-north-1.amazonaws.com.cn/my-repo:latest" + self.assertTrue(self.mode._is_ecr_image(image)) + self.assertEqual( + self.mode._ecr_registry_host(image), + "123456789012.dkr.ecr.cn-north-1.amazonaws.com.cn", + ) + + def test_attacker_crafted_uri_is_not_ecr(self): + # The malicious host is first; the ECR-looking substrings appear only in a later segment. + image = "attacker.com/x.dkr.ecr.us-east-1.amazonaws.com/repo:tag" + self.assertFalse(self.mode._is_ecr_image(image)) + self.assertIsNone(self.mode._ecr_registry_host(image)) + + def test_ecr_substrings_in_repo_path_do_not_classify_as_ecr(self): + image = "evil.example.com/a.dkr.ecr.b.amazonaws.com:latest" + self.assertFalse(self.mode._is_ecr_image(image)) + self.assertIsNone(self.mode._ecr_registry_host(image)) + + def test_public_image_is_not_ecr(self): + for image in ( + "docker.io/library/nginx:latest", + "nginx:latest", + "ubuntu", + ): + self.assertFalse(self.mode._is_ecr_image(image)) + self.assertIsNone(self.mode._ecr_registry_host(image)) + + def test_bad_account_id_length_is_not_ecr(self): + # 11 digits instead of 12 must not match. + image = "12345678901.dkr.ecr.us-east-1.amazonaws.com/repo:tag" + self.assertFalse(self.mode._is_ecr_image(image)) + + +class TestPullImageLoginTarget(unittest.TestCase): + """Ensure docker login is only ever pointed at the validated ECR host.""" + + @patch("sagemaker.serve.mode.local_container_mode.subprocess.run") + @patch("sagemaker.serve.mode.local_container_mode._get_docker_client") + def test_login_uses_validated_host_for_valid_ecr(self, mock_client, mock_run): + import base64 + + mode = _bare_mode() + mode.client = MagicMock() + mock_client.return_value = mode.client + mode.ecr = MagicMock() + token = base64.b64encode(b"AWS:secret-token").decode("utf-8") + mode.ecr.get_authorization_token.return_value = { + "authorizationData": [{"authorizationToken": token}] + } + + image = "123456789012.dkr.ecr.us-east-1.amazonaws.com/my-repo:latest" + mode._pull_image(image) + + # docker login must target the validated ECR host, never any other segment. + self.assertTrue(mock_run.called) + login_args = mock_run.call_args[0][0] + self.assertEqual(login_args[0:2], ["docker", "login"]) + self.assertEqual(login_args[-1], "123456789012.dkr.ecr.us-east-1.amazonaws.com") + + @patch("sagemaker.serve.mode.local_container_mode.subprocess.run") + @patch("sagemaker.serve.mode.local_container_mode._get_docker_client") + def test_no_login_for_attacker_crafted_uri(self, mock_client, mock_run): + mode = _bare_mode() + mode.client = MagicMock() + mock_client.return_value = mode.client + mode.ecr = MagicMock() + + image = "attacker.com/x.dkr.ecr.us-east-1.amazonaws.com/repo:tag" + mode._pull_image(image) + + # The crafted URI is not ECR: no token fetched, no docker login, no credential leak. + mode.ecr.get_authorization_token.assert_not_called() + mock_run.assert_not_called() + + +if __name__ == "__main__": + unittest.main() From 50cce3f4f5c4e07fb4ce71066ba1fa0d12a1fa5d Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Wed, 5 Aug 2026 15:43:30 -0700 Subject: [PATCH 2/2] fix(serve): validate hub-sourced ECR address as defense-in-depth Add a consumption-time guard for image URIs sourced from a (potentially untrusted) hub document, complementing the sink-side validation in LocalContainerMode._pull_image. An attacker-writable hub could publish an EcrAddress such as attacker.com/x.dkr.ecr..amazonaws.com/repo that looks like ECR under the historical substring classifier while its real registry host is attacker-controlled. validate_hub_ecr_address() rejects exactly that parser-confusion signature -- a URI carrying the ECR-like substrings whose registry host is not a valid ECR endpoint -- at the point EcrAddress / init_kwargs.image_uri is read from the hub document, before it can propagate to docker login. It fails closed ONLY on the attack signature: legitimate non-ECR images (public/DockerHub) and ECR hosts in partitions the strict pattern does not enumerate (ISO c2s.ic.gov, FIPS ecr-fips) pass through untouched rather than raising, so a narrow regex cannot turn into an outage. The strict ECR host pattern (ECR_HOST_RE) is promoted to check_image_uri as the single source of truth and reused by LocalContainerMode, so the classifier, the docker-login target extractor, and the hub-consumption guard all operate on the same value. Add unit tests covering valid ECR (standard + China), attacker-spoofed host, ECR substrings in the repo path, public images, ISO/FIPS pass-through, and the empty-URI no-op. --- .../serve/mode/local_container_mode.py | 17 +++--- .../src/sagemaker/serve/model_builder.py | 8 ++- .../sagemaker/serve/model_builder_servers.py | 6 ++ .../serve/validations/check_image_uri.py | 60 +++++++++++++++++++ .../unit/validations/test_check_image_uri.py | 48 ++++++++++++++- 5 files changed, 128 insertions(+), 11 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py b/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py index acbbb3e85d..f61d80aad4 100644 --- a/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py +++ b/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py @@ -4,7 +4,6 @@ from pathlib import Path import logging import os -import re from datetime import datetime, timedelta from typing import Dict, Optional, Type import base64 @@ -26,18 +25,18 @@ from sagemaker.serve.model_server.tgi.server import LocalTgiServing from sagemaker.serve.model_server.tei.server import LocalTeiServing from sagemaker.serve.model_server.multi_model_server.server import LocalMultiModelServer +from sagemaker.serve.validations.check_image_uri import ECR_HOST_RE from sagemaker.core.helper.session_helper import Session logger = logging.getLogger(__name__) -# Strict allowlist for a real ECR registry host: <12-digit-account>.dkr.ecr..amazonaws.com -# Also matches AWS China (.amazonaws.com.cn). This is intentionally an exact host match so that the -# "is this an ECR image?" classifier and the "which host do I docker login to?" extractor operate on -# the SAME value, preventing a crafted URI (e.g. attacker.com/x.dkr.ecr..amazonaws.com/repo) -# from being classified as ECR while login is pointed at an attacker-controlled host. -_ECR_HOST_RE = re.compile( - r"^[0-9]{12}\.dkr\.ecr\.[a-z0-9-]+\.amazonaws\.com(\.cn)?$" -) +# Strict ECR registry host pattern, shared with the hub-consumption validation in +# check_image_uri so the "is this an ECR image?" classifier, the "which host do I docker login +# to?" extractor, and the hub EcrAddress validation all operate on the SAME value. This is +# intentionally an exact host match, preventing a crafted URI (e.g. +# attacker.com/x.dkr.ecr..amazonaws.com/repo) from being classified as ECR while login is +# pointed at an attacker-controlled host. +_ECR_HOST_RE = ECR_HOST_RE _PING_HEALTH_CHECK_INTERVAL_SEC = 5 diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 11853c8a11..27c75c1f2c 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -79,7 +79,7 @@ from sagemaker.serve.utils.types import ModelServer, ModelHub from sagemaker.serve.detector.image_detector import _get_model_base, _detect_framework_and_version from sagemaker.serve.detector.pickler import save_pkl, save_xgboost -from sagemaker.serve.validations.check_image_uri import is_1p_image_uri +from sagemaker.serve.validations.check_image_uri import is_1p_image_uri, validate_hub_ecr_address from sagemaker.core.inference_config import ResourceRequirements from sagemaker.serve.inference_recommendation_mixin import _InferenceRecommenderMixin from sagemaker.serve.model_builder_utils import _ModelBuilderUtils, SPECULATIVE_DRAFT_MODEL @@ -1253,6 +1253,9 @@ def _fetch_and_cache_recipe_config(self): if self._is_nova_model(): nova_config = self._get_nova_hosting_config(instance_type=self.instance_type) if not self.image_uri: + # Defense-in-depth: reject a hub-sourced image URI that spoofs an ECR host before + # it can propagate to LocalContainerMode's docker login (see check_image_uri). + validate_hub_ecr_address(nova_config["image_uri"]) self.image_uri = nova_config["image_uri"] if self.env_vars: user_overrides = dict(self.env_vars) @@ -1273,6 +1276,9 @@ def _fetch_and_cache_recipe_config(self): if hosting_configs: config = self._select_recipe_hosting_config(hosting_configs) if not self.image_uri: + # Defense-in-depth: reject a hub-sourced image URI that spoofs an ECR host before + # it can propagate to LocalContainerMode's docker login (see check_image_uri). + validate_hub_ecr_address(config.get("EcrAddress")) self.image_uri = config.get("EcrAddress") # Cache environment variables from recipe config. Use `or {}` (not a `{}` default) so a diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py index 08a77e0d1e..71a668bb75 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py @@ -996,6 +996,12 @@ def _build_for_jumpstart(self) -> Model: init_kwargs = get_init_kwargs(**init_kwargs_params) # Configure image URI and environment variables + if not self.image_uri: + # Defense-in-depth: reject a hub-sourced image URI that spoofs an ECR host before it can + # propagate to LocalContainerMode's docker login (see check_image_uri). + from sagemaker.serve.validations.check_image_uri import validate_hub_ecr_address + + validate_hub_ecr_address(init_kwargs.image_uri) self.image_uri = self.image_uri or init_kwargs.image_uri if hasattr(init_kwargs, "env") and init_kwargs.env: diff --git a/sagemaker-serve/src/sagemaker/serve/validations/check_image_uri.py b/sagemaker-serve/src/sagemaker/serve/validations/check_image_uri.py index c11eef78c1..a71e0b2db4 100644 --- a/sagemaker-serve/src/sagemaker/serve/validations/check_image_uri.py +++ b/sagemaker-serve/src/sagemaker/serve/validations/check_image_uri.py @@ -2,6 +2,23 @@ from __future__ import absolute_import +import logging +import re + +logger = logging.getLogger(__name__) + +# Strict allowlist for a real ECR registry host: <12-digit-account>.dkr.ecr..amazonaws.com +# Also matches AWS China (.amazonaws.com.cn). Anchored (^...$) so it matches the WHOLE host, never a +# substring buried elsewhere in the URI. This is the single source of truth shared with +# LocalContainerMode so the "is this an ECR image?" classifier, the "which host do I docker login +# to?" extractor, and the hub-consumption validation below all operate on the SAME value. +ECR_HOST_RE = re.compile(r"^[0-9]{12}\.dkr\.ecr\.[a-z0-9-]+\.amazonaws\.com(\.cn)?$") + +# Loose "looks like ECR" signature matching the historical _is_ecr_image() substring check. Used +# ONLY to detect the parser-confusion attack: a URI whose registry host is NOT a valid ECR host yet +# still carries the ECR-looking substrings elsewhere in the string. +_ECR_LIKE_SUBSTRINGS = (".dkr.ecr.", ".amazonaws.com") + # Generated by running the parse_registry_accounts.py script all_accounts = { # Nova escrow ECR accounts (training, evaluation, inference images per region) @@ -305,3 +322,46 @@ def is_1p_image_uri(image_uri: str) -> bool: """Shows if the given image_uri is owned by a 1st party account""" image_uri_account = image_uri[0:12] return image_uri_account in all_accounts + + +def _ecr_registry_host(image_uri: str): + """Return the registry host if ``image_uri``'s registry is a valid ECR endpoint, else None. + + The registry host is the first "/"-delimited segment -- i.e. the exact value that would be + passed to ``docker login`` -- validated against the strict ECR endpoint pattern. + """ + host = image_uri.split("/")[0] + return host if ECR_HOST_RE.match(host) else None + + +def validate_hub_ecr_address(image_uri: str) -> None: + """Defense-in-depth check for image URIs sourced from a (potentially untrusted) hub document. + + This guards the parser-confusion vulnerability where an attacker-writable hub could publish an + ``EcrAddress`` such as ``attacker.com/x.dkr.ecr..amazonaws.com/repo`` that *looks* like + ECR (the historical substring classifier accepted it) while the real registry host is + attacker-controlled. Left unvalidated, LocalContainerMode would ``docker login`` to that host + and leak the caller's ECR authorization token. + + It intentionally only rejects the attack signature -- a URI that carries the ECR-looking + substrings but whose registry host is NOT a valid ECR endpoint -- so legitimate non-ECR images + (public images) and ECR hosts in partitions the strict pattern does not enumerate (e.g. ISO + ``c2s.ic.gov`` / FIPS ``ecr-fips`` endpoints) pass through untouched rather than failing closed. + + Args: + image_uri: The image URI extracted from a hub document (``EcrAddress`` / nova image_uri). + + Raises: + ValueError: If the URI matches the ECR-substring signature but its registry host is not a + valid ECR endpoint (the parser-confusion attack). + """ + if not image_uri: + return + looks_like_ecr = all(sub in image_uri for sub in _ECR_LIKE_SUBSTRINGS) + if looks_like_ecr and _ecr_registry_host(image_uri) is None: + host = image_uri.split("/")[0] + raise ValueError( + f"Hub content contains an invalid ECR address: registry host '{host}' carries " + "ECR-like substrings but is not a valid AWS ECR endpoint. Refusing to use this image " + "to avoid forwarding ECR credentials to a non-AWS host." + ) diff --git a/sagemaker-serve/tests/unit/validations/test_check_image_uri.py b/sagemaker-serve/tests/unit/validations/test_check_image_uri.py index c978d74a5f..e7871f5573 100644 --- a/sagemaker-serve/tests/unit/validations/test_check_image_uri.py +++ b/sagemaker-serve/tests/unit/validations/test_check_image_uri.py @@ -1,5 +1,51 @@ import unittest -from sagemaker.serve.validations.check_image_uri import is_1p_image_uri, all_accounts +from sagemaker.serve.validations.check_image_uri import ( + is_1p_image_uri, + all_accounts, + validate_hub_ecr_address, +) + + +class TestValidateHubEcrAddress(unittest.TestCase): + """Defense-in-depth: reject hub-sourced image URIs that spoof an ECR host. + + Guards the parser-confusion vulnerability where an untrusted hub could publish an EcrAddress + that looks like ECR but whose real registry host is attacker-controlled, causing + LocalContainerMode to forward the ECR authorization token to that host. + """ + + def test_valid_ecr_uri_passes(self): + validate_hub_ecr_address("123456789012.dkr.ecr.us-east-1.amazonaws.com/my-repo:latest") + + def test_valid_ecr_uri_china_partition_passes(self): + validate_hub_ecr_address("123456789012.dkr.ecr.cn-north-1.amazonaws.com.cn/my-repo:latest") + + def test_attacker_spoofed_host_raises(self): + with self.assertRaises(ValueError): + validate_hub_ecr_address("attacker.com/x.dkr.ecr.us-east-1.amazonaws.com/repo:tag") + + def test_ecr_substrings_in_repo_path_raises(self): + with self.assertRaises(ValueError): + validate_hub_ecr_address("evil.example.com/a.dkr.ecr.b.amazonaws.com:latest") + + def test_public_image_passes(self): + # Not ECR-like at all -> not our concern, must pass through untouched. + for image in ("nginx:latest", "docker.io/library/nginx:latest", "ubuntu"): + validate_hub_ecr_address(image) + + def test_iso_partition_ecr_host_passes(self): + # ISO ECR hosts use a non-amazonaws.com TLD; they are legitimate and must NOT fail closed + # just because the strict pattern does not enumerate them. + validate_hub_ecr_address("123456789012.dkr.ecr.us-iso-east-1.c2s.ic.gov/my-repo:latest") + + def test_fips_endpoint_passes(self): + # FIPS endpoints (ecr-fips) lack the ".dkr.ecr." substring, so they are not ECR-like and + # must pass through rather than fail closed. + validate_hub_ecr_address("123456789012.dkr.ecr-fips.us-east-1.amazonaws.com/my-repo:latest") + + def test_empty_uri_is_noop(self): + validate_hub_ecr_address(None) + validate_hub_ecr_address("") class TestCheckImageUri(unittest.TestCase):