Add input validation and resource management improvements

Author: aviruthen

src/sagemaker/iterators.py

@@ -114,6 +114,9 @@ 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
 
@@ -182,5 +185,15 @@ def __next__(self):
                 # print and move on to next response byte
                 print("Unknown event type:" + chunk)
                 continue
+            
+            # 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:
+                raise RuntimeError(
+                    f"Line buffer exceeded maximum size of {self.MAX_BUFFER_SIZE} bytes. "
+                    f"No newline found in stream."
+                )
+            
             self.buffer.seek(0, io.SEEK_END)
             self.buffer.write(chunk["PayloadPart"]["Bytes"])

src/sagemaker/local/data.py

@@ -118,10 +118,34 @@ 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:
+            if self.root_path.startswith(restricted_path):
+                raise ValueError(
+                    f"Local Mode does not support mounting from restricted system paths. "
+                    f"Got: {root_path}"
+                )
+        
         if not os.path.exists(self.root_path):
             raise RuntimeError("Invalid data source: %s does not exist." % self.root_path)
 

src/sagemaker/local/utils.py

@@ -48,10 +48,7 @@ def copy_directory_structure(destination_directory, relative_path):
             destination_directory
     """
     full_path = os.path.join(destination_directory, relative_path)
-    if os.path.exists(full_path):
-        return
-
-    os.makedirs(destination_directory, relative_path)
+    os.makedirs(full_path, exist_ok=True)
 
 
 def move_to_destination(source, destination, job_name, sagemaker_session, prefix=""):

src/sagemaker/utils.py

@@ -601,6 +601,95 @@ def _save_model(repacked_model_uri, tmp_model_path, sagemaker_session, kms_key):
         shutil.move(tmp_model_path, repacked_model_uri.replace("file://", ""))
 
 
+def _validate_source_directory(source_directory):
+    """Validate that source_directory is safe to use.
+
+    Ensures the source directory path does not access restricted system locations.
+
+    Args:
+        source_directory (str): The source directory path to validate.
+
+    Raises:
+        ValueError: If the path is not allowed.
+    """
+    if not source_directory or source_directory.lower().startswith("s3://"):
+        # 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",
+    ]
+
+    # Check if the source path is under any sensitive directory
+    for sensitive_path in sensitive_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.
+
+    Ensures the dependency path does not access restricted system locations.
+
+    Args:
+        dependency (str): The dependency path to validate.
+
+    Raises:
+        ValueError: If the path is not allowed.
+    """
+    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",
+    ]
+
+    # Check if the dependency path is under any sensitive directory
+    for sensitive_path in sensitive_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
 ):
@@ -614,6 +703,8 @@ def _create_or_update_code_dir(
             custom_extractall_tarfile(t, code_dir)
 
     elif source_directory:
+        # Validate source_directory for security
+        _validate_source_directory(source_directory)
         if os.path.exists(code_dir):
             shutil.rmtree(code_dir)
         shutil.copytree(source_directory, code_dir)
@@ -646,6 +737,8 @@ def _create_or_update_code_dir(
                 )
 
     for dependency in dependencies:
+        # Validate dependency path for security
+        _validate_dependency_path(dependency)
         lib_dir = os.path.join(code_dir, "lib")
         if os.path.isdir(dependency):
             shutil.copytree(dependency, os.path.join(lib_dir, os.path.basename(dependency)))
@@ -1620,6 +1713,38 @@ def _get_safe_members(members):
             yield file_info
 
 
+def _validate_extracted_paths(extract_path):
+    """Validate that extracted paths remain within the expected directory.
+
+    Performs post-extraction validation to ensure all extracted files and directories
+    are within the intended extraction path.
+
+    Args:
+        extract_path (str): The path where files were extracted.
+
+    Raises:
+        ValueError: If any extracted file is outside the expected extraction path.
+    """
+    base = _get_resolved_path(extract_path)
+
+    for root, dirs, files in os.walk(extract_path):
+        # Check directories
+        for dir_name in dirs:
+            dir_path = os.path.join(root, dir_name)
+            resolved = _get_resolved_path(dir_path)
+            if not resolved.startswith(base):
+                logger.error("Extracted directory escaped extraction path: %s", dir_path)
+                raise ValueError(f"Extracted path outside expected directory: {dir_path}")
+
+        # Check files
+        for file_name in files:
+            file_path = os.path.join(root, file_name)
+            resolved = _get_resolved_path(file_path)
+            if not resolved.startswith(base):
+                logger.error("Extracted file escaped extraction path: %s", file_path)
+                raise ValueError(f"Extracted path outside expected directory: {file_path}")
+
+
 def custom_extractall_tarfile(tar, extract_path):
     """Extract a tarfile, optionally using data_filter if available.
 
@@ -1640,6 +1765,8 @@ def custom_extractall_tarfile(tar, extract_path):
         tar.extractall(path=extract_path, filter="data")
     else:
         tar.extractall(path=extract_path, members=_get_safe_members(tar))
+        # Re-validate extracted paths to catch symlink race conditions
+        _validate_extracted_paths(extract_path)
 
 
 def can_model_package_source_uri_autopopulate(source_uri: str):

tests/unit/sagemaker/local/test_local_utils.py

@@ -22,12 +22,10 @@
 from sagemaker.session_settings import SessionSettings
 
 
-@patch("sagemaker.local.utils.os.path")
 @patch("sagemaker.local.utils.os")
-def test_copy_directory_structure(m_os, m_os_path):
-    m_os_path.exists.return_value = False
+def test_copy_directory_structure(m_os):
     sagemaker.local.utils.copy_directory_structure("/tmp/", "code/")
-    m_os.makedirs.assert_called_with("/tmp/", "code/")
+    m_os.makedirs.assert_called_with("/tmp/code/", exist_ok=True)
 
 
 @patch("shutil.rmtree", Mock())