diff --git a/sagemaker-core/src/sagemaker/core/common_utils.py b/sagemaker-core/src/sagemaker/core/common_utils.py index 4f865b439f..c6ec385dbb 100644 --- a/sagemaker-core/src/sagemaker/core/common_utils.py +++ b/sagemaker-core/src/sagemaker/core/common_utils.py @@ -424,6 +424,9 @@ def download_folder(bucket_name, prefix, target, sagemaker_session): prefix = prefix.lstrip("/") + if ".." in prefix: + raise ValueError("Traversal components are not allowed in S3 path!") + # Try to download the prefix as an object first, in case it is a file and not a 'directory'. # Do this first, in case the object has broader permissions than the bucket. if not prefix.endswith("/"): diff --git a/sagemaker-core/tests/unit/test_common_utils.py b/sagemaker-core/tests/unit/test_common_utils.py index 291b667191..8aeb496922 100644 --- a/sagemaker-core/tests/unit/test_common_utils.py +++ b/sagemaker-core/tests/unit/test_common_utils.py @@ -945,6 +945,21 @@ def test_download_folder_with_prefix(self): with tempfile.TemporaryDirectory() as tmpdir: download_folder("bucket", "prefix/", tmpdir, mock_session) + def test_download_folder_with_traversal_error(self): + """Test downloading folder with prefix.""" + from sagemaker.core.common_utils import download_folder + + mock_session = Mock() + mock_s3 = Mock() + mock_bucket = Mock() + mock_session.s3_resource = mock_s3 + mock_s3.Bucket.return_value = mock_bucket + mock_bucket.objects.filter.return_value = [] + + with tempfile.TemporaryDirectory() as tmpdir: + with pytest.raises(ValueError): + download_folder("bucket", "/../prefix/", tmpdir, mock_session) + class TestRepackModel: """Test repack_model function."""