chore: harden S3 download path handling

Align master-v2 with master by validating that files downloaded from
S3 resolve within the intended destination directory before writing.
Adds a shared validate_path_within_directory() helper in
sagemaker.utils and wires it into the S3 download paths in
_download_files_under_prefix, Session.download_data,
mlflow.utils._download_s3_artifacts, and the local transform job.

Includes unit tests for the new validation behavior.

Author: lucasjia-aws

src/sagemaker/local/entities.py

@@ -28,7 +28,12 @@
 
 from sagemaker.local.image import _SageMakerContainer
 from sagemaker.local.utils import copy_directory_structure, move_to_destination, get_docker_host
-from sagemaker.utils import DeferredError, get_config_value, format_tags
+from sagemaker.utils import (
+    DeferredError,
+    get_config_value,
+    format_tags,
+    validate_path_within_directory,
+)
 from sagemaker.local.exceptions import StepExecutionException
 
 logger = logging.getLogger(__name__)
@@ -511,9 +516,12 @@ def _perform_batch_inference(self, input_data, output_data, **kwargs):
 
             relative_path = os.path.dirname(os.path.relpath(fn, dataset_dir))
             filename = os.path.basename(fn)
-            copy_directory_structure(working_dir, relative_path)
             destination_path = os.path.join(working_dir, relative_path, filename + ".out")
 
+            validate_path_within_directory(destination_path, working_dir, source_description=fn)
+
+            copy_directory_structure(working_dir, relative_path)
+
             with open(destination_path, "wb") as f:
                 for item in batch_provider.pad(fn, max_payload):
                     # call the container and add the result to inference.

src/sagemaker/serve/model_format/mlflow/utils.py

@@ -21,6 +21,7 @@
 import os
 
 from sagemaker import Session, image_uris
+from sagemaker.utils import validate_path_within_directory
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.detector.image_detector import _cast_to_compatible_version
 from sagemaker.serve.model_format.mlflow.constants import (
@@ -260,6 +261,8 @@ def _download_s3_artifacts(s3_path: str, dst_path: str, session: Session) -> Non
             rel_path = os.path.relpath(key, s3_key)
             local_file_path = os.path.join(dst_path, rel_path)
 
+            validate_path_within_directory(local_file_path, dst_path, source_description=key)
+
             if not key.endswith("/"):
                 local_file_dir = os.path.dirname(local_file_path)
                 os.makedirs(local_file_dir, exist_ok=True)

src/sagemaker/session.py

@@ -137,6 +137,7 @@
     format_tags,
     Tags,
     TagsDict,
+    validate_path_within_directory,
 )
 from sagemaker import exceptions
 from sagemaker.session_settings import SessionSettings
@@ -551,12 +552,16 @@ def download_data(self, path, bucket, key_prefix="", extra_args=None):
             download_extra_args["ExpectedBucketOwner"] = expected_owner
         downloaded_paths = []
         for dir_path in directories:
+            validate_path_within_directory(dir_path, path)
             os.makedirs(os.path.dirname(dir_path), exist_ok=True)
         for key in keys:
             tail_s3_uri_path = os.path.basename(key)
             if not os.path.splitext(key_prefix)[1]:
                 tail_s3_uri_path = os.path.relpath(key, key_prefix)
             destination_path = os.path.join(path, tail_s3_uri_path)
+
+            validate_path_within_directory(destination_path, path, source_description=key)
+
             if not os.path.exists(os.path.dirname(destination_path)):
                 os.makedirs(os.path.dirname(destination_path), exist_ok=True)
             s3.download_file(

src/sagemaker/utils.py

@@ -449,6 +449,31 @@ def download_folder(bucket_name, prefix, target, sagemaker_session):
     _download_files_under_prefix(bucket_name, prefix, target, s3, extra_args=extra_args)
 
 
+def validate_path_within_directory(file_path, target_directory, source_description=""):
+    """Validate that file_path resolves to a location within target_directory.
+
+    Prevents path traversal attacks (CWE-22) by resolving both paths to their
+    canonical forms and checking containment.
+
+    Args:
+        file_path (str): The file path to validate.
+        target_directory (str): The directory that file_path must stay within.
+        source_description (str): Optional description of the source (e.g. S3 key)
+            included in the error message for debugging.
+
+    Raises:
+        ValueError: If file_path resolves to a location outside target_directory.
+    """
+    target_real = os.path.realpath(target_directory)
+    file_real = os.path.realpath(file_path)
+    if not file_real.startswith(target_real + os.sep) and file_real != target_real:
+        source_info = f"'{source_description}' resolves to " if source_description else ""
+        raise ValueError(
+            f"Path traversal detected: {source_info}"
+            f"'{file_real}' which is outside the target directory '{target_real}'"
+        )
+
+
 def _download_files_under_prefix(bucket_name, prefix, target, s3, extra_args=None):
     """Download all S3 files which match the given prefix
 
@@ -469,6 +494,8 @@ def _download_files_under_prefix(bucket_name, prefix, target, s3, extra_args=Non
         s3_relative_path = obj_sum.key[len(prefix) :].lstrip("/")
         file_path = os.path.join(target, s3_relative_path)
 
+        validate_path_within_directory(file_path, target, source_description=obj_sum.key)
+
         try:
             os.makedirs(os.path.dirname(file_path))
         except OSError as exc:

tests/unit/sagemaker/local/test_local_entities.py

@@ -173,6 +173,44 @@ def test_local_transform_job_perform_batch_inference(
     assert "file2.out" in output_files
 
 
+def test_local_transform_job_perform_batch_inference_path_traversal(
+    sagemaker_local_session, tmpdir
+):
+    """Test that file paths resolving outside working_dir are blocked (CWE-22)."""
+    # Build a dedicated job instance to avoid mutating the shared session-scoped fixture.
+    with patch(
+        "sagemaker.local.local_session.LocalSagemakerClient.describe_model"
+    ) as describe_model:
+        describe_model.return_value = {
+            "PrimaryContainer": {"Environment": {}, "Image": "some-image:1.0"}
+        }
+        job = sagemaker.local.entities._LocalTransformJob(
+            "traversal-transform-job", "some-model", sagemaker_local_session
+        )
+
+    working_dir = str(tmpdir.mkdir("working"))
+    dataset_dir = str(tmpdir.mkdir("dataset"))
+
+    job._get_working_directory = Mock(return_value=working_dir)
+
+    data_source = Mock()
+    data_source.get_root_dir.return_value = dataset_dir
+    # A file that resolves outside working_dir via '..' relative path
+    traversal_file = os.path.join(dataset_dir, "..", "..", "..", "etc", "passwd")
+    data_source.get_file_list.return_value = [traversal_file]
+
+    batch_provider = Mock()
+    job._prepare_data_transformation = Mock(return_value=(data_source, batch_provider))
+
+    input_data = {"ContentType": "text/csv"}
+    output_data = {"S3OutputPath": "s3://bucket/output", "Accept": "text/csv"}
+
+    with pytest.raises(ValueError, match="Path traversal detected"):
+        job._perform_batch_inference(
+            input_data, output_data, BatchStrategy="MultiRecord", MaxPayloadInMB="6"
+        )
+
+
 @patch("sagemaker.local.entities._SageMakerContainer", Mock())
 @patch("sagemaker.local.entities.get_docker_host")
 @patch("sagemaker.local.entities._perform_request")

tests/unit/sagemaker/serve/model_format/mlflow/test_mlflow_utils.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import os
+import tempfile
 from pathlib import Path
 from unittest.mock import patch, MagicMock, mock_open
 
@@ -230,6 +231,60 @@ def test_download_s3_artifacts_valid_s3_path(mock_os_makedirs, mock_session):
     )
 
 
+@patch("sagemaker.Session")
+@patch("os.makedirs")
+def test_download_s3_artifacts_path_traversal_via_dotdot_in_key(mock_os_makedirs, mock_session):
+    """Test that S3 keys with '..' traversal sequences are blocked."""
+    mock_s3_client = MagicMock()
+    mock_s3_client.get_paginator.return_value.paginate.return_value = [
+        {"Contents": [{"Key": "model/../../../../etc/passwd"}]}
+    ]
+    mock_session.boto_session.client.return_value = mock_s3_client
+    mock_session._get_account_id_if_default_bucket = lambda bucket: None
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with pytest.raises(ValueError, match="Path traversal detected"):
+            _download_s3_artifacts("s3://my-bucket/model", tmpdir, mock_session)
+
+    mock_s3_client.download_file.assert_not_called()
+
+
+@patch("sagemaker.Session")
+@patch("os.makedirs")
+def test_download_s3_artifacts_path_traversal_overwrite_ssh_keys(mock_os_makedirs, mock_session):
+    """Test the attack scenario from the vulnerability report targeting SSH keys."""
+    mock_s3_client = MagicMock()
+    mock_s3_client.get_paginator.return_value.paginate.return_value = [
+        {"Contents": [{"Key": "mlruns/exp1/model/../../../../Users/alice/.ssh/authorized_keys"}]}
+    ]
+    mock_session.boto_session.client.return_value = mock_s3_client
+    mock_session._get_account_id_if_default_bucket = lambda bucket: None
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with pytest.raises(ValueError, match="Path traversal detected"):
+            _download_s3_artifacts(
+                "s3://shared-bucket/mlruns/exp1/model", tmpdir, mock_session
+            )
+
+    mock_s3_client.download_file.assert_not_called()
+
+
+@patch("sagemaker.Session")
+def test_download_s3_artifacts_safe_keys_are_allowed(mock_session):
+    """Test that normal S3 keys within the target directory are allowed."""
+    mock_s3_client = MagicMock()
+    mock_s3_client.get_paginator.return_value.paginate.return_value = [
+        {"Contents": [{"Key": "model/MLmodel"}, {"Key": "model/subdir/model.pkl"}]}
+    ]
+    mock_session.boto_session.client.return_value = mock_s3_client
+    mock_session._get_account_id_if_default_bucket = lambda bucket: None
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        _download_s3_artifacts("s3://my-bucket/model", tmpdir, mock_session)
+
+    assert mock_s3_client.download_file.call_count == 2
+
+
 @patch("sagemaker.image_uris.retrieve")
 @patch("sagemaker.serve.model_format.mlflow.utils._cast_to_compatible_version")
 @patch("sagemaker.serve.model_format.mlflow.utils._get_framework_version_from_requirements")

tests/unit/test_session.py

@@ -18,6 +18,7 @@
 import json
 import logging
 import os
+import tempfile
 
 import pytest
 import six
@@ -7052,6 +7053,68 @@ def test_download_data_with_file_and_directory(makedirs, sagemaker_session):
     )
 
 
+def test_download_data_path_traversal_in_file_key(sagemaker_session):
+    """Test that S3 keys with '..' traversal sequences are blocked."""
+    sagemaker_session.s3_client = Mock()
+    sagemaker_session.s3_client.list_objects_v2 = Mock(
+        return_value={
+            "Contents": [
+                {"Key": "data/../../../../etc/passwd", "Size": 100},
+            ]
+        }
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with pytest.raises(ValueError, match="Path traversal detected"):
+            sagemaker_session.download_data(
+                path=tmpdir, bucket="foo-bucket", key_prefix="data/"
+            )
+
+    sagemaker_session.s3_client.download_file.assert_not_called()
+
+
+def test_download_data_path_traversal_overwrite_credentials(sagemaker_session):
+    """Test the exact attack scenario from the vulnerability report."""
+    sagemaker_session.s3_client = Mock()
+    sagemaker_session.s3_client.list_objects_v2 = Mock(
+        return_value={
+            "Contents": [
+                {"Key": "data/../../../../Users/alice/.aws/credentials", "Size": 200},
+            ]
+        }
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with pytest.raises(ValueError, match="Path traversal detected"):
+            sagemaker_session.download_data(
+                path=tmpdir, bucket="shared-bucket", key_prefix="data/"
+            )
+
+    sagemaker_session.s3_client.download_file.assert_not_called()
+
+
+@patch("os.makedirs")
+def test_download_data_safe_keys_are_allowed(makedirs, sagemaker_session):
+    """Test that normal S3 keys within the target directory are allowed."""
+    sagemaker_session.s3_client = Mock()
+    sagemaker_session.s3_client.list_objects_v2 = Mock(
+        return_value={
+            "Contents": [
+                {"Key": "data/train.csv", "Size": 100},
+                {"Key": "data/models/model.pkl", "Size": 500},
+            ]
+        }
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        result = sagemaker_session.download_data(
+            path=tmpdir, bucket="foo-bucket", key_prefix="data/"
+        )
+
+    assert len(result) == 2
+    assert sagemaker_session.s3_client.download_file.call_count == 2
+
+
 def test_create_hub(sagemaker_session):
     sagemaker_session.create_hub(
         hub_name="mock-hub-name",