diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index 122bba9460..c3b886145b 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/sagemaker-train/src/sagemaker/train/model_trainer.py @@ -33,7 +33,7 @@ from sagemaker.core.shapes import AlgorithmSpecification, ModelPackageConfig from sagemaker.core.utils.utils import serialize from sagemaker.core.apiutils._boto_functions import to_pascal_case -from sagemaker.core.common_utils import name_from_base +from sagemaker.core.common_utils import name_from_base, _validate_symlinks_within_root from dataclasses import dataclass @@ -1242,6 +1242,11 @@ def create_input_data_channel( input_mode="File", ) else: + # Copying/archiving this directory for upload follows nested + # symlinks, so a symlink pointing outside the directory would pull + # in unrelated content from elsewhere on the file system. Validate + # that nested symlinks resolve inside the directory before staging. + _validate_symlinks_within_root(data_source, "data_source") key_prefix = ( f"{key_prefix}/{channel_name}" if key_prefix diff --git a/sagemaker-train/tests/unit/train/test_model_trainer.py b/sagemaker-train/tests/unit/train/test_model_trainer.py index 2bb0e6246f..254701bc66 100644 --- a/sagemaker-train/tests/unit/train/test_model_trainer.py +++ b/sagemaker-train/tests/unit/train/test_model_trainer.py @@ -2815,3 +2815,37 @@ def test_pipeline_and_interactive_emit_identical_tags(model_trainer): pipeline_tags = _pipeline_trainer([TAG])._create_training_job_args()["Tags"] assert pipeline_tags == _train_wire_tags(model_trainer) + + +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_create_input_data_channel_local_dir_with_in_root_symlink( + mock_default_bucket, mock_upload_data, model_trainer +): + """A local directory whose nested symlinks resolve inside it is staged normally.""" + mock_default_bucket.return_value = DEFAULT_BUCKET + mock_upload_data.return_value = f"s3://{DEFAULT_BUCKET}/code" + with tempfile.TemporaryDirectory() as source_dir: + with open(os.path.join(source_dir, "train.py"), "w") as f: + f.write("print('hi')\n") + os.symlink(os.path.join(source_dir, "train.py"), os.path.join(source_dir, "entry_link.py")) + channel = model_trainer.create_input_data_channel("code", source_dir) + assert channel.channel_name == "code" + mock_upload_data.assert_called_once() + + +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_create_input_data_channel_local_dir_rejects_out_of_root_symlink( + mock_default_bucket, mock_upload_data, model_trainer +): + """A nested symlink resolving outside the local directory is rejected before upload.""" + mock_default_bucket.return_value = DEFAULT_BUCKET + with tempfile.TemporaryDirectory() as outside_dir, tempfile.TemporaryDirectory() as source_dir: + outside_file = os.path.join(outside_dir, "outside.txt") + with open(outside_file, "w") as f: + f.write("data\n") + os.symlink(outside_file, os.path.join(source_dir, "linked.txt")) + with pytest.raises(ValueError, match="resolves outside"): + model_trainer.create_input_data_channel("code", source_dir) + mock_upload_data.assert_not_called()