fix: normalize tar member paths when extracting Bedrock model artifacts (#6543)

* fix: normalize tar member paths when extracting Bedrock model artifacts

* fix: report skipped archive members when extracting Bedrock model artifacts

Author: aviruthen

sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py

@@ -16,6 +16,7 @@
 
 import json
 import os
+import re
 import time
 import logging
 from datetime import datetime, timezone
@@ -95,6 +96,35 @@ def _log_bedrock_api_call(api_name: str, params: Dict[str, Any], response: Dict[
     print(f"[BedrockModelBuilder] API call logged to: {filepath}")
 
 
+def _tar_member_relative_key(name: str) -> Optional[str]:
+    """Normalize a tar member name into a key relative to the extraction prefix.
+
+    Leading ``./`` segments and empty segments are removed so that archives
+    created with ``tar -C <dir> .`` map to the same keys as archives created
+    from the directory contents. Names that normalize to the same key (for
+    example ``./a`` and ``a``) map to one key, and the later member wins, as
+    it would when extracting the archive to a directory.
+
+    Only ``/`` is treated as a separator when building the key, so a backslash
+    in a name is kept as a literal character. Names that are rooted or contain
+    a ``..`` segment under either separator are not supported, so that keys
+    stay relative under both POSIX and Windows path conventions.
+
+    Args:
+        name: The tar member name.
+
+    Returns:
+        The normalized relative key, or None if the name is absolute, contains
+        a ``..`` segment, or is empty after normalization.
+    """
+    if not name or name.startswith(("/", "\\")) or re.match(r"^[A-Za-z]:", name):
+        return None
+    if ".." in re.split(r"[\\/]", name):
+        return None
+    relative = "/".join(part for part in name.split("/") if part not in ("", "."))
+    return relative or None
+
+
 class BedrockModelBuilder:
     """Builder class for deploying models to Amazon Bedrock.
 
@@ -1076,21 +1106,38 @@ def _extract_tar_gz_to_s3(self, tar_gz_uri: str) -> str:
         stream.seekable = lambda: False
 
         extracted_count = 0
+        skipped_count = 0
         with tarfile.open(fileobj=stream, mode="r|gz") as tar:
             for member in tar:
                 if not member.isfile():
                     continue
+                relative_key = _tar_member_relative_key(member.name)
+                if relative_key is None:
+                    skipped_count += 1
+                    logger.warning("Skipping archive member with unsupported path: %s", member.name)
+                    continue
                 f = tar.extractfile(member)
                 if f is None:
                     continue
-                dest_key = extract_prefix + member.name
+                dest_key = extract_prefix + relative_key
                 size_mb = member.size / (1024 * 1024)
                 extracted_count += 1
                 logger.info("Extracting [%d]: %s (%.1f MB)", extracted_count, member.name, size_mb)
                 s3_client.put_object(Bucket=bucket, Key=dest_key, Body=f.read())
 
         if extracted_count == 0:
+            if skipped_count:
+                raise RuntimeError(
+                    f"No files extracted from {tar_gz_uri}: {skipped_count} archive member(s) "
+                    "were skipped because their paths are absolute or contain '..' segments."
+                )
             raise RuntimeError(f"No files found in {tar_gz_uri}.")
+        if skipped_count:
+            logger.warning(
+                "Skipped %d archive member(s) with unsupported paths in %s.",
+                skipped_count,
+                tar_gz_uri,
+            )
 
         logger.info("Extracted %d files to %s", extracted_count, extract_uri)
         return extract_uri

sagemaker-serve/tests/unit/test_bedrock_model_builder.py

@@ -12,15 +12,22 @@
 # language governing permissions and limitations under the License.
 """Unit tests for BedrockModelBuilder."""
 
+import io
 import json
+import tarfile
+
 import pytest
 from unittest.mock import Mock, patch
 from botocore.exceptions import ClientError
 
 from sagemaker.core.resources import TrainingJob
 from sagemaker.core.shapes import ModelArtifacts, OutputDataConfig
 from sagemaker.core.utils.utils import Unassigned
-from sagemaker.serve.bedrock_model_builder import BedrockModelBuilder, _is_nova_model
+from sagemaker.serve.bedrock_model_builder import (
+    BedrockModelBuilder,
+    _is_nova_model,
+    _tar_member_relative_key,
+)
 
 MODULE = "sagemaker.serve.bedrock_model_builder"
 
@@ -1625,3 +1632,126 @@ def test_oss_reuse_preserves_user_tags(self):
         assert user_tag in tags
         source_tag = {"key": "sagemaker.amazonaws.com/model-source", "value": "s3://bucket/path/"}
         assert source_tag in tags
+
+
+# ── _extract_tar_gz_to_s3 ───────────────────────────────────────────────────
+
+
+def _tar_gz_bytes(members):
+    """Build an in-memory tar.gz from (name, content) pairs.
+
+    A content of None adds a directory entry instead of a file.
+    """
+    buf = io.BytesIO()
+    with tarfile.open(fileobj=buf, mode="w:gz") as tar:
+        for name, content in members:
+            info = tarfile.TarInfo(name=name)
+            if content is None:
+                info.type = tarfile.DIRTYPE
+                tar.addfile(info)
+            else:
+                info.size = len(content)
+                tar.addfile(info, io.BytesIO(content))
+    return buf.getvalue()
+
+
+def _extract_with_members(members):
+    """Run _extract_tar_gz_to_s3 against a fake S3 client and return put_object keys."""
+    s3_client = Mock()
+    s3_client.list_objects_v2.return_value = {"KeyCount": 0}
+    s3_client.get_object.return_value = {"Body": io.BytesIO(_tar_gz_bytes(members))}
+    b = object.__new__(BedrockModelBuilder)
+    b.boto_session = Mock()
+    b.boto_session.client.return_value = s3_client
+
+    uri = b._extract_tar_gz_to_s3("s3://bucket/jobs/run/output/model.tar.gz")
+
+    assert uri == "s3://bucket/jobs/run/output/extracted/"
+    return {c.kwargs["Key"]: c.kwargs["Body"] for c in s3_client.put_object.call_args_list}
+
+
+class TestTarMemberRelativeKey:
+    @pytest.mark.parametrize(
+        "name,expected",
+        [
+            ("config.json", "config.json"),
+            ("checkpoints/hf/config.json", "checkpoints/hf/config.json"),
+            ("./config.json", "config.json"),
+            ("./checkpoints//hf/./config.json", "checkpoints/hf/config.json"),
+            ("model..safetensors", "model..safetensors"),
+            ("a\\b.json", "a\\b.json"),
+        ],
+    )
+    def test_normalizes_relative_names(self, name, expected):
+        assert _tar_member_relative_key(name) == expected
+
+    @pytest.mark.parametrize(
+        "name",
+        [
+            "",
+            ".",
+            "./",
+            "/config.json",
+            "\\config.json",
+            "C:/config.json",
+            "../config.json",
+            "a/../../config.json",
+            "a\\..\\config.json",
+        ],
+    )
+    def test_rejects_unsupported_names(self, name):
+        assert _tar_member_relative_key(name) is None
+
+
+class TestExtractTarGzToS3:
+    def test_uploads_members_under_extract_prefix(self):
+        keys = _extract_with_members(
+            [("config.json", b"{}"), ("./checkpoints/hf/model.safetensors", b"weights")]
+        )
+
+        assert keys == {
+            "jobs/run/output/extracted/config.json": b"{}",
+            "jobs/run/output/extracted/checkpoints/hf/model.safetensors": b"weights",
+        }
+
+    def test_skips_directory_members(self):
+        keys = _extract_with_members(
+            [("./", None), ("./checkpoints", None), ("./checkpoints/config.json", b"{}")]
+        )
+
+        assert keys == {"jobs/run/output/extracted/checkpoints/config.json": b"{}"}
+
+    def test_names_normalizing_to_same_key_keep_last_member(self):
+        keys = _extract_with_members([("./config.json", b"first"), ("config.json", b"second")])
+
+        assert keys == {"jobs/run/output/extracted/config.json": b"second"}
+
+    def test_skips_members_with_unsupported_paths(self, caplog):
+        with caplog.at_level("WARNING"):
+            keys = _extract_with_members(
+                [("config.json", b"{}"), ("../other/file.txt", b"x"), ("/abs/file.txt", b"y")]
+            )
+
+        assert keys == {"jobs/run/output/extracted/config.json": b"{}"}
+        assert "Skipped 2 archive member(s)" in caplog.text
+
+    def test_raises_with_skipped_count_when_all_members_skipped(self):
+        with pytest.raises(RuntimeError, match="2 archive member\\(s\\) were skipped"):
+            _extract_with_members([("../file.txt", b"x"), ("/abs/file.txt", b"y")])
+
+    def test_raises_when_archive_has_no_files(self):
+        with pytest.raises(RuntimeError, match="No files found"):
+            _extract_with_members([("./", None)])
+
+    def test_skips_extraction_when_prefix_exists(self):
+        s3_client = Mock()
+        s3_client.list_objects_v2.return_value = {"KeyCount": 1}
+        b = object.__new__(BedrockModelBuilder)
+        b.boto_session = Mock()
+        b.boto_session.client.return_value = s3_client
+
+        uri = b._extract_tar_gz_to_s3("s3://bucket/jobs/model.tar.gz")
+
+        assert uri == "s3://bucket/jobs/extracted/"
+        s3_client.get_object.assert_not_called()
+        s3_client.put_object.assert_not_called()