Allowing for sym-links, better refactoring

Author: aviruthen

src/sagemaker/iterators.py

@@ -17,6 +17,7 @@
 import io
 
 from sagemaker.exceptions import ModelStreamError, InternalStreamFailure
+from sagemaker.utils import _MAX_BUFFER_SIZE
 
 
 def handle_stream_errors(chunk):
@@ -114,9 +115,6 @@ def __next__(self):
 class LineIterator(BaseIterator):
     """A helper class for parsing the byte Event Stream input to provide Line iteration."""
 
-    # Maximum buffer size to prevent unbounded memory consumption (10 MB)
-    MAX_BUFFER_SIZE = 10 * 1024 * 1024
-
     def __init__(self, event_stream):
         """Initialises a LineIterator Iterator object
 
@@ -189,9 +187,9 @@ def __next__(self):
             # Check buffer size before writing to prevent unbounded memory consumption
             chunk_size = len(chunk["PayloadPart"]["Bytes"])
             current_size = self.buffer.getbuffer().nbytes
-            if current_size + chunk_size > self.MAX_BUFFER_SIZE:
+            if current_size + chunk_size > _MAX_BUFFER_SIZE:
                 raise RuntimeError(
-                    f"Line buffer exceeded maximum size of {self.MAX_BUFFER_SIZE} bytes. "
+                    f"Line buffer exceeded maximum size of {_MAX_BUFFER_SIZE} bytes. "
                     f"No newline found in stream."
                 )
 

src/sagemaker/local/data.py

@@ -26,6 +26,7 @@
 import sagemaker.amazon.common
 import sagemaker.local.utils
 import sagemaker.utils
+from sagemaker.utils import _SENSITIVE_SYSTEM_PATHS
 
 
 def get_data_source_instance(data_source, sagemaker_session):
@@ -118,28 +119,13 @@ def get_root_dir(self):
 class LocalFileDataSource(DataSource):
     """Represents a data source within the local filesystem."""
 
-    # Blocklist of sensitive directories that should not be accessible
-    RESTRICTED_PATHS = [
-        os.path.abspath(os.path.expanduser("~/.aws")),
-        os.path.abspath(os.path.expanduser("~/.ssh")),
-        os.path.abspath(os.path.expanduser("~/.kube")),
-        os.path.abspath(os.path.expanduser("~/.docker")),
-        os.path.abspath(os.path.expanduser("~/.config")),
-        os.path.abspath(os.path.expanduser("~/.credentials")),
-        "/etc",
-        "/root",
-        "/home",
-        "/var/lib",
-        "/opt/ml/metadata",
-    ]
-
     def __init__(self, root_path):
         super(LocalFileDataSource, self).__init__()
 
         self.root_path = os.path.abspath(root_path)
 
         # Validate that the path is not in restricted locations
-        for restricted_path in self.RESTRICTED_PATHS:
+        for restricted_path in _SENSITIVE_SYSTEM_PATHS:
             if self.root_path.startswith(restricted_path):
                 raise ValueError(
                     f"Local Mode does not support mounting from restricted system paths. "

src/sagemaker/utils.py

@@ -76,6 +76,21 @@
 WAITING_DOT_NUMBER = 10
 MAX_ITEMS = 100
 PAGE_SIZE = 10
+_MAX_BUFFER_SIZE = 100 * 1024 * 1024  # 100 MB - Maximum buffer size for streaming iterators
+
+_SENSITIVE_SYSTEM_PATHS = [
+    abspath(os.path.expanduser("~/.aws")),
+    abspath(os.path.expanduser("~/.ssh")),
+    abspath(os.path.expanduser("~/.kube")),
+    abspath(os.path.expanduser("~/.docker")),
+    abspath(os.path.expanduser("~/.config")),
+    abspath(os.path.expanduser("~/.credentials")),
+    "/etc",
+    "/root",
+    "/home",
+    "/var/lib",
+    "/opt/ml/metadata",
+]
 
 logger = logging.getLogger(__name__)
 
@@ -616,35 +631,17 @@ def _validate_source_directory(source_directory):
         # S3 paths and None are safe
         return
 
-    abs_source = abspath(source_directory)
-
-    # Blocklist of sensitive directories that should not be accessible
-    sensitive_paths = [
-        abspath(os.path.expanduser("~/.aws")),
-        abspath(os.path.expanduser("~/.ssh")),
-        abspath(os.path.expanduser("~/.kube")),
-        abspath(os.path.expanduser("~/.docker")),
-        abspath(os.path.expanduser("~/.config")),
-        abspath(os.path.expanduser("~/.credentials")),
-        "/etc",
-        "/root",
-        "/home",
-        "/var/lib",
-        "/opt/ml/metadata",
-    ]
+    # Resolve symlinks to get the actual path
+    abs_source = abspath(realpath(source_directory))
 
     # Check if the source path is under any sensitive directory
-    for sensitive_path in sensitive_paths:
+    for sensitive_path in _SENSITIVE_SYSTEM_PATHS:
         if abs_source.startswith(sensitive_path):
             raise ValueError(
                 f"source_directory cannot access sensitive system paths. "
                 f"Got: {source_directory} (resolved to {abs_source})"
             )
 
-    # Check for symlinks to prevent symlink-based escapes
-    if os.path.islink(abs_source):
-        raise ValueError(f"source_directory cannot be a symlink: {source_directory}")
-
 
 def _validate_dependency_path(dependency):
     """Validate that a dependency path is safe to use.
@@ -660,35 +657,17 @@ def _validate_dependency_path(dependency):
     if not dependency:
         return
 
-    abs_dependency = abspath(dependency)
-
-    # Blocklist of sensitive directories that should not be accessible
-    sensitive_paths = [
-        abspath(os.path.expanduser("~/.aws")),
-        abspath(os.path.expanduser("~/.ssh")),
-        abspath(os.path.expanduser("~/.kube")),
-        abspath(os.path.expanduser("~/.docker")),
-        abspath(os.path.expanduser("~/.config")),
-        abspath(os.path.expanduser("~/.credentials")),
-        "/etc",
-        "/root",
-        "/home",
-        "/var/lib",
-        "/opt/ml/metadata",
-    ]
+    # Resolve symlinks to get the actual path
+    abs_dependency = abspath(realpath(dependency))
 
     # Check if the dependency path is under any sensitive directory
-    for sensitive_path in sensitive_paths:
+    for sensitive_path in _SENSITIVE_SYSTEM_PATHS:
         if abs_dependency.startswith(sensitive_path):
             raise ValueError(
                 f"dependency path cannot access sensitive system paths. "
                 f"Got: {dependency} (resolved to {abs_dependency})"
             )
 
-    # Check for symlinks to prevent symlink-based escapes
-    if os.path.islink(abs_dependency):
-        raise ValueError(f"dependency path cannot be a symlink: {dependency}")
-
 
 def _create_or_update_code_dir(
     model_dir, inference_script, source_directory, dependencies, sagemaker_session, tmp