Security Hardening: Path Traversal Prevention in RS_VAE_Decode_Save

Rayko_VAE_Save_Image.py

@@ -56,6 +56,21 @@ class RS_VAE_Decode_Save:
     EMBED_WORKFLOW = True
     OVERWRITE_EXISTING = False
 
+    @staticmethod
+    def _sanitize_path_component(component: str) -> str:
+        """Sanitizes a string to be safe for use in file paths."""
+        if not isinstance(component, str):
+            component = str(component)
+        component = component.replace("\\", "/")
+        while "../" in component or "./" in component:
+            component = component.replace("../", "").replace("./", "")
+        if os.path.isabs(component):
+            component = os.path.basename(component)
+        component = re.sub(r'[<>:"|?*]', '_', component)
+        if not component.strip():
+            component = "sanitized_empty"
+        return component.strip()
+
     @staticmethod
     def _normalize_images(images):
         while images.dim() > 4 and images.shape[0] == 1:
@@ -82,15 +97,28 @@ class RS_VAE_Decode_Save:
 
     def _get_next_filename(self, directory, base_name, extension, is_batch, batch_index=None):
         try:
+            safe_base_name = self._sanitize_path_component(base_name)
+            safe_extension = self._sanitize_path_component(extension).lower()
+
+            test_path = os.path.join(directory, f"{safe_base_name}.{safe_extension}")
+            real_dir = os.path.realpath(directory)
+            real_path = os.path.realpath(test_path)
+
+            if not real_path.startswith(real_dir + os.sep) and real_path != real_dir:
+                print(f"[RS] Security Warning: Path traversal detected in '{base_name}'. Falling back to safe name.")
+                safe_base_name = "unsafe_input_sanitized"
+
             if self.OVERWRITE_EXISTING:
                 if is_batch and batch_index is not None:
-                    return os.path.join(directory, f"{base_name}_{batch_index:03d}.{extension}")
-                return os.path.join(directory, f"{base_name}_001.{extension}")
+                    return os.path.join(directory, f"{safe_base_name}_{batch_index:03d}.{safe_extension}")
+                return os.path.join(directory, f"{safe_base_name}_001.{safe_extension}")
+
             if is_batch and batch_index is not None:
-                base_with_index = f"{base_name}_{batch_index:03d}"
-                pattern = re.compile(rf'^{re.escape(base_with_index)}_(\d+)\.{re.escape(extension)}$')
+                base_with_index = f"{safe_base_name}_{batch_index:03d}"
+                pattern = re.compile(rf'^{re.escape(base_with_index)}_(\d+)\.{re.escape(safe_extension)}$')
             else:
-                pattern = re.compile(rf'^{re.escape(base_name)}_(\d+)\.{re.escape(extension)}$')
+                pattern = re.compile(rf'^{re.escape(safe_base_name)}_(\d+)\.{re.escape(safe_extension)}$')
+
             max_num = 0
             if os.path.exists(directory):
                 for f in os.listdir(directory):
@@ -99,16 +127,18 @@ class RS_VAE_Decode_Save:
                         num = int(match.group(1))
                         if num > max_num:
                             max_num = num
+
             next_num = max_num + 1
             if is_batch and batch_index is not None:
-                filename = f"{base_name}_{batch_index:03d}_{next_num:03d}.{extension}"
+                filename = f"{safe_base_name}_{batch_index:03d}_{next_num:03d}.{safe_extension}"
             else:
-                filename = f"{base_name}_{next_num:03d}.{extension}"
+                filename = f"{safe_base_name}_{next_num:03d}.{safe_extension}"
+
             return os.path.join(directory, filename)
         except Exception as e:
             print(f"[RS] Filename error: {e}")
             ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-            return os.path.join(directory, f"{base_name}_{ts}.{extension}")
+            return os.path.join(directory, f"error_fallback_{ts}.png")
 
     def _build_workflow_json(self, prompt, extra_pnginfo):
         try:
@@ -173,8 +203,10 @@ class RS_VAE_Decode_Save:
         images = vae.decode(samples["samples"])
         images = self._normalize_images(images)
 
+        safe_prefix = self._sanitize_path_component(filename_prefix)
+
         full_output_folder, fn_prefix, counter, subfolder, _ = \
-            folder_paths.get_save_image_path(filename_prefix, self.output_dir, 
+            folder_paths.get_save_image_path(safe_prefix, self.output_dir,
                                              images[0].shape[1], images[0].shape[0])
 
         batch_size = images.shape[0]
@@ -185,7 +217,7 @@ class RS_VAE_Decode_Save:
         saved_files = []
         for i in range(batch_size):
             batch_index = i + 1 if is_batch else None
-            
+
             file_counter = counter + i
             filename = f"{fn_prefix}_{file_counter:05}.{extension}"
             filepath = os.path.join(full_output_folder, filename)
@@ -228,6 +260,7 @@ class RS_VAE_Decode_Save:
 
         return {"ui": {"images": saved_files}, "result": (images,)}
 
+
 NODE_CLASS_MAPPINGS = {
     "RS_VAE_Decode_Save": RS_VAE_Decode_Save,
 }

__init__.py

@@ -29,7 +29,7 @@ IGNORED_MODULES: Set[str] = {'__init__.py'}
 REQUIRED_MODULES: Set[str] = set()
 
 PACKAGE_DISPLAY_NAME = "ComfyUI_RaykoStudio"
-__version__ = "0.36.1"
+__version__ = "0.36.2"
 
 def load_modules():
     print(f"🦊 {PACKAGE_DISPLAY_NAME}: Version: {__version__}")

package.json

@@ -1,6 +1,6 @@
 {
   "name": "ComfyUI_RaykoStudio",
-  "version": "0.36.1",
+  "version": "0.36.2",
   "description": "Set of custom nodes for ComfyUI providing additional image processing capabilities",
   "author": "Raykosan",
   "license": "Apache License 2.0",

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI_RaykoStudio"
 description = "Set of custom nodes for ComfyUI providing additional image processing capabilities"
-version = "0.36.1"
+version = "0.36.2"
 license = { file = "LICENSE" }
 dependencies = ["torch", "numpy", "Pillow>=10.0.0", "freetype-py>=2.5.1", "pycairo>=1.29.0", "scipy>=1.10.0"]