fix(train): validate symlinks stay within source root when staging to S3

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

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()