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..f61d80aad4 100644 --- a/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py +++ b/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py @@ -5,7 +5,7 @@ import logging import os from datetime import datetime, timedelta -from typing import Dict, Type +from typing import Dict, Optional, Type import base64 import time import subprocess @@ -25,10 +25,19 @@ 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 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 _PING_HEALTH_CHECK_FAIL_MSG = ( @@ -252,7 +261,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 +289,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/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/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() 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):