From fc9bd30cfad0e03e2ee8508414b1464c0dfd9540 Mon Sep 17 00:00:00 2001 From: Raykosan <71210053+Raykosan@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:15:30 +0300 Subject: [PATCH] Security Hardening: Path Traversal Prevention in RS_VAE_Decode_Save --- Rayko_VAE_Save_Image.py | 53 +++++++++++++++++++++++++++++++++-------- __init__.py | 2 +- package.json | 2 +- pyproject.toml | 2 +- 4 files changed, 46 insertions(+), 13 deletions(-) diff --git a/Rayko_VAE_Save_Image.py b/Rayko_VAE_Save_Image.py index 42caf95..0beff54 100644 --- a/Rayko_VAE_Save_Image.py +++ b/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, } diff --git a/__init__.py b/__init__.py index 6d97dd0..3c4f50d 100644 --- a/__init__.py +++ b/__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__}") diff --git a/package.json b/package.json index d496449..f5e1de1 100644 --- a/package.json +++ b/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", diff --git a/pyproject.toml b/pyproject.toml index d56a4f9..8d7368c 100644 --- a/pyproject.toml +++ b/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"]