fix(serve): validate ECR registry host before docker login

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.<region>.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
 

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

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:

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.<region>.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.<region>.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."
+        )

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()

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):