From 22f25842b4214819377bef9ff7af2464043905cf Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Mon, 15 Dec 2025 21:48:05 -0800 Subject: [PATCH] [sec]Adding validation for resolved path to deny sensitive directories --- .../src/sagemaker/core/common_utils.py | 4 ++ .../tests/unit/test_common_utils.py | 40 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/sagemaker-core/src/sagemaker/core/common_utils.py b/sagemaker-core/src/sagemaker/core/common_utils.py index b5dd2ecbef..34553c9ec8 100644 --- a/sagemaker-core/src/sagemaker/core/common_utils.py +++ b/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) diff --git a/sagemaker-core/tests/unit/test_common_utils.py b/sagemaker-core/tests/unit/test_common_utils.py index 4052c46f26..2b6f7bbfae 100644 --- a/sagemaker-core/tests/unit/test_common_utils.py +++ b/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()