fix: address review comments (iteration #1)

Author: aviruthen

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

@@ -647,7 +647,10 @@ def _validate_source_directory(source_directory):
 
     # Check if the source path is under any sensitive directory
     for sensitive_path in _SENSITIVE_SYSTEM_PATHS:
-        if abs_source != "/" and os.path.commonpath([abs_source, sensitive_path]) == sensitive_path:
+        if abs_source != "/" and (
+            os.path.commonpath([abs_source, sensitive_path])
+            == sensitive_path
+        ):
             raise ValueError(
                 f"source_directory cannot access sensitive system paths. "
                 f"Got: {source_directory} (resolved to {abs_source})"
@@ -673,7 +676,10 @@ def _validate_dependency_path(dependency):
 
     # Check if the dependency path is under any sensitive directory
     for sensitive_path in _SENSITIVE_SYSTEM_PATHS:
-        if abs_dependency != "/" and os.path.commonpath([abs_dependency, sensitive_path]) == sensitive_path:
+        if abs_dependency != "/" and (
+            os.path.commonpath([abs_dependency, sensitive_path])
+            == sensitive_path
+        ):
             raise ValueError(
                 f"dependency path cannot access sensitive system paths. "
                 f"Got: {dependency} (resolved to {abs_dependency})"
@@ -686,10 +692,13 @@ def _create_or_update_code_dir(
     """Placeholder docstring"""
     code_dir = os.path.join(model_dir, "code")
     resolved_code_dir = _get_resolved_path(code_dir)
-    
+
     # Validate that code_dir does not resolve to a sensitive system path
     for sensitive_path in _SENSITIVE_SYSTEM_PATHS:
-        if resolved_code_dir != "/" and os.path.commonpath([resolved_code_dir, sensitive_path]) == sensitive_path:
+        if resolved_code_dir != "/" and (
+            os.path.commonpath([resolved_code_dir, sensitive_path])
+            == sensitive_path
+        ):
             raise ValueError(
                 f"Invalid code_dir path: {code_dir} resolves to sensitive system path {resolved_code_dir}"
             )
@@ -1688,7 +1697,8 @@ def _is_bad_path(path, base):
         bool: True if the path is not rooted under the base directory, False otherwise.
     """
     # joinpath will ignore base if path is absolute
-    return not _get_resolved_path(joinpath(base, path)).startswith(base)
+    resolved = _get_resolved_path(joinpath(base, path))
+    return os.path.commonpath([resolved, base]) != base
 
 
 def _is_bad_link(info, base):
@@ -1708,19 +1718,18 @@ def _is_bad_link(info, base):
     return _is_bad_path(info.linkname, base=tip)
 
 
-def _get_safe_members(members):
+def _get_safe_members(members, base):
     """A generator that yields members that are safe to extract.
 
     It filters out bad paths and bad links.
 
     Args:
-        members (list): A list of members to check.
+        members (list): A list of TarInfo members to check.
+        base (str): The resolved base directory for extraction.
 
     Yields:
         tarfile.TarInfo: The tar file info.
     """
-    base = _get_resolved_path("")
-
     for file_info in members:
         if _is_bad_path(file_info.name, base):
             logger.error("%s is blocked (illegal path)", file_info.name)
@@ -1783,7 +1792,11 @@ def custom_extractall_tarfile(tar, extract_path):
     if hasattr(tarfile, "data_filter"):
         tar.extractall(path=extract_path, filter="data")
     else:
-        tar.extractall(path=extract_path, members=_get_safe_members(tar))
+        base = _get_resolved_path(extract_path)
+        tar.extractall(
+            path=extract_path,
+            members=_get_safe_members(tar.getmembers(), base),
+        )
         # Re-validate extracted paths to catch symlink race conditions
         _validate_extracted_paths(extract_path)