[sec]Adding validation for resolved path to deny sensitive directories

Author: jam-jee

sagemaker-core/src/sagemaker/core/common_utils.py

@@ -74,6 +74,7 @@
 WAITING_DOT_NUMBER = 10
 MAX_ITEMS = 100
 PAGE_SIZE = 10
+SENSITIVE_SYSTEM_ROOTS = ["/", "/etc", "/var"]
 
 logger = logging.getLogger(__name__)
 
@@ -612,6 +613,9 @@ def _create_or_update_code_dir(
 ):
     """Placeholder docstring"""
     code_dir = os.path.join(model_dir, "code")
+    resolved_code_dir = _get_resolved_path(code_dir)
+    if resolved_code_dir in SENSITIVE_SYSTEM_ROOTS:
+        raise ValueError(f"Invalid code_dir path: {code_dir} resolves to sensitive system root {resolved_code_dir}")
     if source_directory and source_directory.lower().startswith("s3://"):
         local_code_path = os.path.join(tmp, "local_code.tar.gz")
         download_file_from_url(source_directory, local_code_path, sagemaker_session)

sagemaker-core/tests/unit/test_common_utils.py

@@ -2209,3 +2209,43 @@ def test_nested_set_dict_multiple_keys(self):
         d = {}
         nested_set_dict(d, ["a", "b", "c"], "value")
         assert d["a"]["b"]["c"] == "value"
+
+
+class TestCreateOrUpdateCodeDir:
+    """Test _create_or_update_code_dir function."""
+
+    def test_create_or_update_code_dir_rejects_root(self, tmp_path):
+        """Test that root path is rejected."""
+        from sagemaker.core.common_utils import _create_or_update_code_dir
+
+        with patch("sagemaker.core.common_utils._get_resolved_path", return_value="/"):
+            with pytest.raises(ValueError, match="sensitive system root"):
+                _create_or_update_code_dir("/", "script.py", None, [], Mock(), str(tmp_path))
+
+    def test_create_or_update_code_dir_rejects_etc(self, tmp_path):
+        """Test that /etc path is rejected."""
+        from sagemaker.core.common_utils import _create_or_update_code_dir
+
+        with patch("sagemaker.core.common_utils._get_resolved_path", return_value="/etc"):
+            with pytest.raises(ValueError, match="sensitive system root"):
+                _create_or_update_code_dir("/etc", "script.py", None, [], Mock(), str(tmp_path))
+
+    def test_create_or_update_code_dir_rejects_var(self, tmp_path):
+        """Test that /var path is rejected."""
+        from sagemaker.core.common_utils import _create_or_update_code_dir
+
+        with patch("sagemaker.core.common_utils._get_resolved_path", return_value="/var"):
+            with pytest.raises(ValueError, match="sensitive system root"):
+                _create_or_update_code_dir("/var", "script.py", None, [], Mock(), str(tmp_path))
+
+    def test_create_or_update_code_dir_allows_safe_path(self, tmp_path):
+        """Test that safe paths are allowed."""
+        from sagemaker.core.common_utils import _create_or_update_code_dir
+
+        model_dir = tmp_path / "model"
+        model_dir.mkdir()
+        script = tmp_path / "script.py"
+        script.write_text("# script")
+
+        _create_or_update_code_dir(str(model_dir), str(script), None, [], Mock(), str(tmp_path))
+        assert (model_dir / "code").exists()