Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions sagemaker-core/src/sagemaker/core/common_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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("/"):
Expand Down
15 changes: 15 additions & 0 deletions sagemaker-core/tests/unit/test_common_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Loading