From 781d2a19721f503dc7a4f75d80f496b4478fbf2f Mon Sep 17 00:00:00 2001 From: Raykosan <71210053+Raykosan@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:00:00 +0300 Subject: [PATCH] Harden output-path handling via explicit allowlist (RS_EXTRA_OUTPUT_ROOTS) --- README.md | 146 ++++++++++-- Rayko_VAE_Save_Image.py | 187 ++++++++++++--- tests/test_rayko_vae_save_image_security.py | 251 ++++++++++++++++++++ 3 files changed, 530 insertions(+), 54 deletions(-) create mode 100644 tests/test_rayko_vae_save_image_security.py diff --git a/README.md b/README.md index 3428313..af017e6 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # 🦊 ComfyUI_RaykoStudio Set of custom nodes for ComfyUI providing additional image processing capabilities --- ---- +---
🖥️ Performance has been tested for @@ -26,7 +26,7 @@ Set of custom nodes for ComfyUI providing additional image processing capabiliti
🛠 Installation - + Set of nodes can be installed in several ways: - Clone repository to `ComfyUI/custom_nodes/` folder: ``` @@ -53,7 +53,7 @@ git clone https://github.com/Raykosan/ComfyUI_RaykoStudio.git # NODES
🦊 RS Label - + # 🦊 RS Label **A highly customizable floating label node for ComfyUI.** Perfect for adding annotations, titles, watermarks, or text overlays directly on the canvas with full control over typography and styling. @@ -105,7 +105,7 @@ Fonts are auto-detected on editor open
🦊 RS Label Image - + # 🦊 RS Label Image **A versatile floating image label node for ComfyUI.** Perfect for adding visual markers, logos, icons, or annotations directly onto your canvas with full control over styling and layout. @@ -150,7 +150,7 @@ Perfect for adding visual markers, logos, icons, or annotations directly onto yo
🦊 RS Collage Node - + # 🦊 RS Collage Node **Interactive node for overlaying images with real-time positioning, scaling, rotation, and edge feathering directly on the canvas.** @@ -249,7 +249,7 @@ In Chrome and Edge, the button works immediately (requires a one-time confirmati In Firefox, the Paste button will additionally pop up, when clicked, the path will be inserted. **These are browser security principles and cannot be circumvented.** In any case, you can always paste the path using the keyboard shortcut Ctrl+V. - + ### Inputs | Parameter | Type | Description | | :--- | :--- | :--- | @@ -270,7 +270,7 @@ In any case, you can always paste the path using the keyboard shortcut Ctrl+V.
🦊 RS Image-Prompt - + # 🦊 RS Image-Prompt **A ComfyUI node that extracts the positive prompt from image metadata — instantly, without running the queue.** Load a generated image and `RS Image-Prompt` reads back the prompt that was used to create it. It understands metadata written by ComfyUI, Automatic1111, Forge, and the rest of the RaykoStudio suite. @@ -358,6 +358,7 @@ The prompt is extracted from a wide range of metadata formats, in priority order * 🟠 **Orange** - the group is partially in the bypass. * ⚪ **Gray** — the node/group is active. + **Smart State saving** - The bypass status is saved directly in the JSON workflow. No data is lost when restarting ComfyUI, switching tabs, or sharing PNG/JSON. **Dynamic size** - The node automatically adjusts its height to the number of mounted elements. **Advanced UX**: @@ -372,6 +373,7 @@ If there are a large number of nodes, use the search bar to filter. Click on groups or nodes to switch their state (Bypass/Active). * *Tip: Clicking on the name of the group bypasses it entirely. Clicking on the arrow (▶) will expand the group to select individual nodes.* + When you're done, click anywhere outside the node or press the Escape key to close the menu. To turn off or turn on the bypass of an element inserted into the interface, use the personal toggle. To turn off or enable bypass for all elements, use the TOGGLE ALL switch. @@ -384,7 +386,7 @@ Also, when using bypass using comfi's own methods (the context menu is bypass, b
🦊 RS Decode Save Image - + # 🦊 RS Decode Save Image **A lightweight ComfyUI custom node that combines native VAE Decode and Save Image into a single, streamlined unit.** Decodes latent samples to image tensors and saves them to disk in one step, while preserving the decoded `IMAGE` output for downstream preview or processing. @@ -407,10 +409,110 @@ Decodes latent samples to image tensors and saves them to disk in one step, whil | **PREFIX** | The base name for your saved files. Files will be saved as `prefix_00001.png`, `prefix_00002.png`, etc. | | **FORMAT** | The output file format selector: **PNG**, **JPG**, or **WEBP**. | +## Output directory & write boundary + +`RS_VAE_Decode_Save` supports two kinds of `save_path` values: + +| Input in `save_path` | Behavior | +| --------------------------------------------- | ------------------------------------------------------------ | +| *(empty)* | Saves into ComfyUI's default `output/` directory. | +| `project_v2` (relative) | Saves into `output/project_v2/`. The subfolder is created if missing. | +| `/mnt/ssd/renders` or `I:/Renders` (absolute) | Saves to that exact path — **only if it is inside one of the allowed roots** (see below). | + +### Write boundary + +Every write is resolved through `os.path.realpath()` and verified with `os.path.commonpath()` against an explicit allowlist of roots. If the resolved path falls outside every allowed root, the node raises `PermissionError` and no file is written. + +The default allowlist contains exactly one root: + +- `folder_paths.get_output_directory()` (ComfyUI's `output/` folder). + +This means: **out of the box, the node cannot write anywhere outside ComfyUI's `output/` directory.** Symlinks, `..` traversal, and absolute paths outside the allowlist are all rejected. + +### Enabling additional output roots + +To let the node write to user-chosen locations (e.g. a dedicated renders drive), add those roots to the allowlist at ComfyUI startup via the `RS_EXTRA_OUTPUT_ROOTS` environment variable. The value is a path-separator-delimited list (`:` on Linux/macOS, `;` on Windows), same convention as `PATH`. + +**Linux / macOS:** + +bash + +``` +export RS_EXTRA_OUTPUT_ROOTS="/mnt/ssd/renders:/media/user/backup" +python main.py +``` + + + +**Windows (cmd):** + +bat + +``` +set RS_EXTRA_OUTPUT_ROOTS=I:\Renders;D:\Projects\ComfyOut +python main.py +``` + + + +**Windows (PowerShell):** + +powershell + +``` +$env:RS_EXTRA_OUTPUT_ROOTS = "I:\Renders;D:\Projects\ComfyOut" +python main.py +``` + + + +**comfy-launch / docker-compose:** + +yaml + +``` +environment: + - RS_EXTRA_OUTPUT_ROOTS=/mnt/ssd/renders:/media/user/backup +``` + + + +After setting the variable, absolute paths inside those roots will be accepted by the node. Paths outside every root will still raise `PermissionError`. + +### Programmatic configuration (optional) + +If you prefer configuring the allowlist from Python instead of the environment, you can set it on `comfy.options` before the node is used: + +python + +``` +import comfy.options +comfy.options.rs_extra_output_roots = [ + "/mnt/ssd/renders", + "/media/user/backup", +] +``` + + + +The environment variable and `comfy.options.rs_extra_output_roots` are both additive — the output directory is always included, and any roots from either source are merged into the allowlist. + +### Notes + +- Roots are resolved with `os.path.realpath()` before comparison. A symlinked path is treated as its target, so putting a symlink inside an allowed root that points outside the root will **not** grant access to the target. +- Adding a root grants write access to **everything under it**, including subdirectories that don't exist yet — they will be created on demand. +- Removing a root from `RS_EXTRA_OUTPUT_ROOTS` on the next launch immediately revokes write access to it. Existing files are not touched. +- The `_sanitize_path_component()` helper inside the node is a cosmetic cleaner for relative subfolder names. It is **not** a security boundary and should not be relied on as one; the boundary is enforced by `_resolve_target_dir()`. + +### Security rationale + +The node intentionally supports user-chosen absolute output paths — that's the point of the `save_path` field and the folder picker in the UI. Rather than disabling that feature or relying on blacklist sanitization (which is fragile against encoding tricks, symlinks, and platform-specific quirks), the boundary is expressed as an explicit allowlist checked with `realpath()` + `commonpath()` at a single choke point (`_resolve_target_dir`). The default posture is safe; enabling broader write access is an explicit, auditable opt-in. +
+
🦊 RS rgb2rgba - + # 🦊 RS rgb2rgba **A lightweight ComfyUI custom node that loads images while preserving the alpha channel (RGBA). Ideal for workflows that require transparency handling in PNG, WebP, and TIFF formats.** @@ -436,12 +538,12 @@ Connect the `rgba` output to any node that accepts ComfyUI image tensors (e.g.,
🦊 RS Any Switch - + # 🦊 RS Any Switch **A dynamic switch node for ComfyUI that allows you to switch between multiple inputs of ANY data type with an intuitive toggle interface.** Screenshot_1 - + ### 🔥 Features **Universal Type Support** - Accepts any data type (IMAGE, LATENT, MODEL, AUDIO, VIDEO, TEXT, etc.) **Dynamic Inputs** - Inputs are created automatically as you connect nodes (up to 20) @@ -505,12 +607,12 @@ Upscale Method 3 ──→ input_3 RS Any Switch ↓ Final Output -``` +```
🦊 RS LoRA Tester - + # 🦊 RS LoRA Tester **A ComfyUI custom node for **batch testing LoRAs** — queue one LoRA at a time instead of loading them all at once. Designed for people who test dozens (or hundreds) of LoRAs and want to generate a comparison batch without manually editing the node between each run.** @@ -550,7 +652,7 @@ Use the ❌ to remove a single LoRA
🦊 RS Save Image LoRA - + # 🦊 RS Save Image LoRA **A drop-in replacement for `SaveImage` that automatically puts the names and strengths of all LoRAs used in the generation right into the output file name.** No more guessing which LoRA produced which image. Your files name themselves. @@ -581,7 +683,7 @@ No more guessing which LoRA produced which image. Your files name themselves.
🦊 RS Upscaler - + # 🦊 RS Upscaler **A node for that combines upscale model loading and image upscaling into a single, compact node with a custom UI.** @@ -602,7 +704,7 @@ Tested with ComfyUI core upscaler models including:
🦊 RS Upscale & Resize - + # 🦊 RS Upscale & Resize **Hybrid upscaler and resizer. Uses 'scale_by' by default. When width/height inputs are connected or set manually (>0), it switches to exact external dimensions.** @@ -632,7 +734,7 @@ Set method to keep proportion or pad to prevent distortion when forcing specific
🦊 RS Tile Adjustments - + # 🦊 RS Tile Adjustments **A powerful node that calculates optimal tile dimensions based on image factors and overlap rates, then splits the input image into a batch of tiles with precise coordinate metadata.** @@ -652,7 +754,7 @@ Use the IMAGES output for processing (upscaling/inpainting) and pass the metadat
🦊 RS Tile Assemble - + # 🦊 RS Tile Assemble **The companion node to RS Tile Adjustments. It reconstructs a full-resolution image from a batch of processed tiles using calculated positions and advanced feathered blending to eliminate visible seams.** @@ -676,7 +778,7 @@ The FINAL_IMAGE output is ready for saving or further processing
🦊 RS Ref Encode - + # 🦊 RS Ref Encode **A node that combines VAE Encode and Set Reference Latent into a single efficient operation.** @@ -756,7 +858,7 @@ The node pauses the queue until you confirm the settings, then outputs a ready-t **Example Workflow**: ``` Load Image → RS Outpaint → [Your Model] → Save Image -``` +``` ### Mode 2: Single Image Batch @@ -777,7 +879,7 @@ Load Image → RS Outpaint → [Your Model] → Save Image **Example Workflow**: ``` Load Image → RS Outpaint → [Multiple Model Variations] → Save Image (×N) -``` +``` ### Mode 3: Multi-Image Batch @@ -1746,7 +1848,7 @@ Toggle **ENABLE SHADOW** to activate.
🦊 RS Image Adjustments - + # 🦊 RS Image Adjustments **A powerful interactive image adjustment node with real-time preview and professional-grade controls.** Inspired by industry standards like Adobe Lightroom and Photoshop, it offers real-time previews, advanced color grading, and an efficient batch processing mode. diff --git a/Rayko_VAE_Save_Image.py b/Rayko_VAE_Save_Image.py index 34ec043..6636024 100644 --- a/Rayko_VAE_Save_Image.py +++ b/Rayko_VAE_Save_Image.py @@ -53,15 +53,31 @@ class RS_VAE_Decode_Save: FUNCTION = "decode_and_save" OUTPUT_NODE = True CATEGORY = "🦊 RaykoStudio" - DESCRIPTION = "Combines native VAE Decode and Save Image into a single node." + DESCRIPTION = ( + "Combines native VAE Decode and Save Image into a single node. " + "Write boundary: by default only ComfyUI's output directory is " + "writable; additional roots can be enabled via the " + "RS_EXTRA_OUTPUT_ROOTS environment variable (path-separator-delimited)." + ) PNG_COMPRESSION = 1 JPG_QUALITY = 90 WEBP_QUALITY = 90 EMBED_WORKFLOW = True + # ------------------------------------------------------------------ + # Path handling / output-directory boundary + # ------------------------------------------------------------------ + @staticmethod def _sanitize_path_component(component: str) -> str: + """ + Convenience sanitizer for *relative* subfolder names. + + This is NOT a security boundary. It only cleans up cosmetic issues + (backslashes, redundant "./", filesystem-illegal characters). + The actual write-boundary is enforced by `_resolve_target_dir()`. + """ if not isinstance(component, str): component = str(component) component = component.replace("\\", "/") @@ -70,6 +86,91 @@ class RS_VAE_Decode_Save: component = re.sub(r'[<>:"|?*]', '_', component) return component.strip() + def _get_allowed_roots(self) -> list: + """ + Return the list of directories this node is allowed to write into. + + Always includes ComfyUI's output directory. Additional roots can be + enabled at runtime via the RS_EXTRA_OUTPUT_ROOTS environment variable + (path-separator-delimited, i.e. os.pathsep) or via comfy.options. + """ + roots = {os.path.realpath(self.output_dir)} + + env = os.environ.get("RS_EXTRA_OUTPUT_ROOTS", "") + for part in env.split(os.pathsep): + part = part.strip() + if part: + roots.add(os.path.realpath(part)) + + try: + import comfy.options # type: ignore + extra = getattr(comfy.options, "rs_extra_output_roots", None) or [] + for part in extra: + if isinstance(part, str) and part.strip(): + roots.add(os.path.realpath(part.strip())) + except Exception: + pass + + return sorted(roots) + + @staticmethod + def _is_within(candidate: str, roots) -> bool: + """ + True if realpath(candidate) lies inside one of the given roots. + + Uses os.path.commonpath, which is the correct containment primitive + on both POSIX and Windows (unlike str.startswith, which has + prefix-collision bugs such as /out vs /output-evil). + """ + real = os.path.realpath(candidate) + for root in roots: + try: + if os.path.commonpath([real, root]) == root: + return True + except ValueError: + # Different drives on Windows. + continue + return False + + def _resolve_target_dir(self, requested: str) -> str: + """ + Resolve the user-provided save path into a safe absolute directory. + + Enforces the write boundary: the resolved path must lie inside one of + the roots returned by `_get_allowed_roots()`. Raises PermissionError + otherwise. + """ + allowed = self._get_allowed_roots() + requested = (requested or "").strip() + + is_absolute = os.path.isabs(requested) or ( + len(requested) > 1 and requested[1] == ":" + ) + + if not requested: + candidate = self.output_dir + elif is_absolute: + candidate = os.path.normpath(requested) + else: + safe = self._sanitize_path_component(requested) + candidate = ( + os.path.join(self.output_dir, safe) if safe else self.output_dir + ) + + if not self._is_within(candidate, allowed): + raise PermissionError( + f"[RS] Refusing to write outside allowed roots. " + f"Requested: {requested!r}. " + f"Allowed roots: {allowed}. " + f"Set RS_EXTRA_OUTPUT_ROOTS to allow additional locations." + ) + + return os.path.realpath(candidate) + + # ------------------------------------------------------------------ + # Image normalization + # ------------------------------------------------------------------ + @staticmethod def _normalize_images(images): images = images.squeeze() @@ -91,12 +192,18 @@ class RS_VAE_Decode_Save: elif images.shape[0] == 1: images = images.squeeze(0) else: - images = images.reshape(images.shape[0], -1, images.shape[-2], images.shape[-1]) + images = images.reshape( + images.shape[0], -1, images.shape[-2], images.shape[-1] + ) if images.dim() != 4: - try: - images = images.view(images.shape[0], images.shape[-2], images.shape[-1], -1) - except: - raise ValueError(f"[RS] Cannot normalize tensor with shape: {images.shape}") + try: + images = images.view( + images.shape[0], images.shape[-2], images.shape[-1], -1 + ) + except Exception: + raise ValueError( + f"[RS] Cannot normalize tensor with shape: {images.shape}" + ) if images.shape[-1] == 1: images = images.repeat(1, 1, 1, 3) if images.dtype != torch.float32: @@ -104,9 +211,15 @@ class RS_VAE_Decode_Save: images = torch.clamp(images, 0.0, 1.0) return images + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + def _get_next_counter(self, directory: str, prefix: str, extension: str) -> int: try: - pattern = re.compile(rf'^{re.escape(prefix)}_(\d{{5}})\.{re.escape(extension)}$') + pattern = re.compile( + rf'^{re.escape(prefix)}_(\d{{5}})\.{re.escape(extension)}$' + ) max_num = 0 if os.path.exists(directory): for f in os.listdir(directory): @@ -124,21 +237,25 @@ class RS_VAE_Decode_Save: temp_dir = folder_paths.get_temp_directory() if not os.path.exists(temp_dir): return - + now = time.time() - cutoff = now - 7200 - + cutoff = now - 7200 + for filename in os.listdir(temp_dir): if filename.startswith("rs_prev_"): filepath = os.path.join(temp_dir, filename) try: if os.path.getmtime(filepath) < cutoff: os.remove(filepath) - except: + except Exception: pass except Exception as e: print(f"[RS] Cleanup error: {e}") + # ------------------------------------------------------------------ + # Main entry point + # ------------------------------------------------------------------ + def decode_and_save(self, samples, vae, save_path, file_prefix, format, node_data, prompt=None, extra_pnginfo=None): images = vae.decode(samples["samples"]) @@ -146,19 +263,14 @@ class RS_VAE_Decode_Save: self._cleanup_temp() - clean_path = save_path.strip() - is_absolute = os.path.isabs(clean_path) or (len(clean_path) > 1 and clean_path[1] == ':') + target_dir = self._resolve_target_dir(save_path) - if is_absolute: - target_dir = os.path.normpath(clean_path) - else: - safe_path = self._sanitize_path_component(clean_path) - target_dir = os.path.join(self.output_dir, safe_path) if safe_path else self.output_dir - try: os.makedirs(target_dir, exist_ok=True) except Exception as e: - raise PermissionError(f"[RS] Cannot create directory: {target_dir}. Error: {e}") + raise PermissionError( + f"[RS] Cannot create directory: {target_dir}. Error: {e}" + ) safe_prefix = self._sanitize_path_component(file_prefix) if not safe_prefix: @@ -167,22 +279,28 @@ class RS_VAE_Decode_Save: batch_size = images.shape[0] ext_map = {"png": "png", "jpg": "jpg", "webp": "webp"} extension = ext_map.get(format, "png") - + start_counter = self._get_next_counter(target_dir, safe_prefix, extension) saved_files = [] temp_dir = folder_paths.get_temp_directory() - if self.last_temp_file and os.path.exists(os.path.join(temp_dir, self.last_temp_file)): - try: os.remove(os.path.join(temp_dir, self.last_temp_file)) - except: pass + if self.last_temp_file and os.path.exists( + os.path.join(temp_dir, self.last_temp_file) + ): + try: + os.remove(os.path.join(temp_dir, self.last_temp_file)) + except Exception: + pass for i in range(batch_size): current_counter = start_counter + i filename_main = f"{safe_prefix}_{current_counter:05}.{extension}" filepath_main = os.path.join(target_dir, filename_main) - - img_array = np.clip(255.0 * images[i].cpu().numpy(), 0, 255).astype(np.uint8) + + img_array = np.clip( + 255.0 * images[i].cpu().numpy(), 0, 255 + ).astype(np.uint8) img = Image.fromarray(img_array) save_kwargs = {} @@ -193,8 +311,10 @@ class RS_VAE_Decode_Save: metadata.add_text("prompt", json.dumps(prompt)) if extra_pnginfo: for key, value in extra_pnginfo.items(): - try: metadata.add_text(key, json.dumps(value)) - except: pass + try: + metadata.add_text(key, json.dumps(value)) + except Exception: + pass save_kwargs["pnginfo"] = metadata elif format == "jpg": save_kwargs["quality"] = self.JPG_QUALITY @@ -203,20 +323,20 @@ class RS_VAE_Decode_Save: elif format == "webp": save_kwargs["quality"] = self.WEBP_QUALITY save_kwargs["method"] = 4 - + img.save(filepath_main, **save_kwargs) temp_name = f"rs_prev_{int(time.time() * 1000)}_{i}.png" filepath_temp = os.path.join(temp_dir, temp_name) img.save(filepath_temp, compress_level=1) - + if i == batch_size - 1: self.last_temp_file = temp_name saved_files.append({ "filename": temp_name, "subfolder": "", - "type": "temp" + "type": "temp", }) return {"ui": {"images": saved_files}, "result": (images,)} @@ -233,7 +353,10 @@ async def get_output_folders(request): if os.path.isdir(full_path): real_path = os.path.realpath(full_path) real_output = os.path.realpath(output_dir) - if real_path.startswith(real_output + os.sep) or real_path == real_output: + if ( + real_path.startswith(real_output + os.sep) + or real_path == real_output + ): subfolders.append(item) return web.json_response({"subfolders": sorted(subfolders)}) except Exception as e: diff --git a/tests/test_rayko_vae_save_image_security.py b/tests/test_rayko_vae_save_image_security.py new file mode 100644 index 0000000..2e3fb4e --- /dev/null +++ b/tests/test_rayko_vae_save_image_security.py @@ -0,0 +1,251 @@ +# tests/test_rayko_vae_save_image_security.py +""" +Regression tests for the output-directory boundary in RS_VAE_Decode_Save. + +The node intentionally supports user-configured output directories. To keep +that feature safe, every write is routed through `_resolve_target_dir()`, +which resolves the final path with `os.path.realpath()` and verifies it lies +inside one of the roots returned by `_get_allowed_roots()`. + +These tests exercise the boundary directly, without invoking VAE decode or +image saving. +""" + +import os +import sys +import types +from pathlib import Path + +import pytest + + +# --------------------------------------------------------------------------- +# Import-time stubs +# +# ComfyUI's pytest.ini is picked up from the installation root, which means +# ComfyUI's own conftest may try to load the plugin before our test module is +# imported. We install stubs for the runtime-only ComfyUI modules at *module +# import time* so that any plugin import during collection also sees them. +# +# Only modules that are unavailable outside a live ComfyUI process are stubbed: +# - folder_paths (needs ComfyUI's paths machinery) +# - server (needs a running PromptServer instance) +# - aiohttp.web (its json_response is only used inside the route handler) +# +# torch, numpy and PIL are real, so we do NOT stub them. +# --------------------------------------------------------------------------- + + +def _fresh_module(name: str) -> types.ModuleType: + """Install a blank module under *name*, removing any earlier version.""" + for existing in list(sys.modules): + if existing == name or existing.startswith(name + "."): + sys.modules.pop(existing, None) + mod = types.ModuleType(name) + sys.modules[name] = mod + return mod + + +def _install_import_stubs(output_dir: Path, temp_dir: Path) -> None: + # ---- folder_paths ----------------------------------------------------- + fp = _fresh_module("folder_paths") + fp.get_output_directory = lambda: str(output_dir) + fp.get_temp_directory = lambda: str(temp_dir) + fp.get_save_image_path = lambda *a, **k: (str(output_dir), "img", 1, "", "") + + # ---- server.PromptServer --------------------------------------------- + # `from server import PromptServer` must yield a *class* with an + # `instance` attribute that exposes `.routes.get(...)` as a decorator. + server_mod = _fresh_module("server") + + class _Routes: + def get(self, *args, **kwargs): + def decorator(fn): + return fn + return decorator + + class PromptServer: # noqa: N801 - matches ComfyUI's class name + instance = types.SimpleNamespace(routes=_Routes()) + + server_mod.PromptServer = PromptServer + + # ---- aiohttp.web ------------------------------------------------------ + aiohttp_mod = _fresh_module("aiohttp") + aiohttp_web = types.ModuleType("aiohttp.web") + aiohttp_web.json_response = lambda *a, **k: None + aiohttp_mod.web = aiohttp_web + sys.modules["aiohttp.web"] = aiohttp_web + + +@pytest.fixture +def node_class(tmp_path, monkeypatch): + output_dir = tmp_path / "output" + output_dir.mkdir() + temp_dir = tmp_path / "temp" + temp_dir.mkdir() + + _install_import_stubs(output_dir, temp_dir) + + # Make sure we get a clean import of the node module. + sys.modules.pop("Rayko_VAE_Save_Image", None) + + repo_root = Path(__file__).resolve().parent.parent + monkeypatch.syspath_prepend(str(repo_root)) + + import Rayko_VAE_Save_Image as mod # noqa: E402 + + return mod.RS_VAE_Decode_Save, output_dir + + +# --------------------------------------------------------------------------- +# Sanitizer (convenience, not a security boundary) +# --------------------------------------------------------------------------- + +class TestSanitizePathComponent: + def test_strips_dotdot(self, node_class): + cls, _ = node_class + assert ".." not in cls._sanitize_path_component("../../etc/passwd") + + def test_strips_illegal_chars(self, node_class): + cls, _ = node_class + out = cls._sanitize_path_component('ac:d"e|f?g*h') + assert not any(c in out for c in '<>:"|?*') + + def test_backslashes_become_slashes(self, node_class): + cls, _ = node_class + assert "\\" not in cls._sanitize_path_component(r"a\b\c") + + +# --------------------------------------------------------------------------- +# Containment primitive +# --------------------------------------------------------------------------- + +class TestIsWithin: + def test_path_inside_root(self, tmp_path, node_class): + cls, _ = node_class + root = str(tmp_path / "root") + os.makedirs(root, exist_ok=True) + inside = os.path.join(root, "sub", "file.txt") + assert cls._is_within(inside, [root]) is True + + def test_path_outside_root(self, tmp_path, node_class): + cls, _ = node_class + root = str(tmp_path / "root") + outside = str(tmp_path / "other") + os.makedirs(root, exist_ok=True) + os.makedirs(outside, exist_ok=True) + assert cls._is_within(outside, [root]) is False + + def test_prefix_collision_is_not_containment(self, tmp_path, node_class): + # "/out" must not be treated as a parent of "/out-evil" + cls, _ = node_class + root = str(tmp_path / "out") + sibling = str(tmp_path / "out-evil") + os.makedirs(root, exist_ok=True) + os.makedirs(sibling, exist_ok=True) + assert cls._is_within(sibling, [root]) is False + + +# --------------------------------------------------------------------------- +# Boundary enforcement +# --------------------------------------------------------------------------- + +class TestResolveTargetDir: + def test_empty_path_returns_output_dir(self, node_class): + cls, output_dir = node_class + node = cls() + assert os.path.realpath(node._resolve_target_dir("")) == os.path.realpath( + str(output_dir) + ) + + def test_relative_subfolder_inside_output(self, node_class): + cls, output_dir = node_class + node = cls() + resolved = node._resolve_target_dir("project_v2") + assert resolved == os.path.realpath( + os.path.join(str(output_dir), "project_v2") + ) + + def test_traversal_stays_inside_output(self, node_class): + # After sanitization this must resolve under output_dir, never above it. + cls, output_dir = node_class + node = cls() + resolved = node._resolve_target_dir("../../etc/passwd") + assert resolved.startswith(os.path.realpath(str(output_dir))) + + def test_absolute_path_outside_roots_is_rejected(self, tmp_path, node_class): + cls, _ = node_class + node = cls() + outside = str(tmp_path / "outside") + os.makedirs(outside, exist_ok=True) + with pytest.raises(PermissionError): + node._resolve_target_dir(outside) + + def test_absolute_path_inside_extra_root_is_allowed( + self, tmp_path, monkeypatch, node_class + ): + cls, _ = node_class + extra = tmp_path / "extra" + extra.mkdir() + monkeypatch.setenv("RS_EXTRA_OUTPUT_ROOTS", str(extra)) + node = cls() + resolved = node._resolve_target_dir(str(extra)) + assert resolved == os.path.realpath(str(extra)) + + def test_multiple_extra_roots(self, tmp_path, monkeypatch, node_class): + cls, _ = node_class + a = tmp_path / "a" + b = tmp_path / "b" + a.mkdir() + b.mkdir() + monkeypatch.setenv("RS_EXTRA_OUTPUT_ROOTS", f"{a}{os.pathsep}{b}") + node = cls() + assert node._resolve_target_dir(str(a)) == os.path.realpath(str(a)) + assert node._resolve_target_dir(str(b)) == os.path.realpath(str(b)) + + def test_symlink_escape_is_rejected(self, tmp_path, node_class): + cls, output_dir = node_class + node = cls() + outside = tmp_path / "outside" + outside.mkdir() + link = output_dir / "escape" + try: + os.symlink(str(outside), str(link), target_is_directory=True) + except (OSError, NotImplementedError): + pytest.skip("symlinks not supported on this platform") + with pytest.raises(PermissionError): + node._resolve_target_dir("escape") + + def test_windows_drive_outside_is_rejected(self, node_class): + if os.name != "nt": + pytest.skip("Windows-only case") + cls, _ = node_class + node = cls() + with pytest.raises(PermissionError): + node._resolve_target_dir("Z:/definitely-not-allowed") + + +# --------------------------------------------------------------------------- +# Allowlist composition +# --------------------------------------------------------------------------- + +class TestAllowedRoots: + def test_output_dir_always_allowed(self, node_class): + cls, output_dir = node_class + node = cls() + assert os.path.realpath(str(output_dir)) in node._get_allowed_roots() + + def test_env_var_adds_root(self, tmp_path, monkeypatch, node_class): + cls, _ = node_class + extra = tmp_path / "extra" + extra.mkdir() + monkeypatch.setenv("RS_EXTRA_OUTPUT_ROOTS", str(extra)) + node = cls() + assert os.path.realpath(str(extra)) in node._get_allowed_roots() + + def test_env_var_empty_is_ignored(self, monkeypatch, node_class): + cls, _ = node_class + monkeypatch.setenv("RS_EXTRA_OUTPUT_ROOTS", "") + node = cls() + # Only the output dir should be present. + assert len(node._get_allowed_roots()) == 1 \ No newline at end of file