Harden output-path handling via explicit allowlist (RS_EXTRA_OUTPUT_ROOTS)
This commit is contained in:
Raykosan
@@ -1,7 +1,7 @@
|
||||
# 🦊 ComfyUI_RaykoStudio
|
||||
Set of custom nodes for ComfyUI providing additional image processing capabilities
|
||||
---
|
||||
---
|
||||
---
|
||||
|
||||
<details>
|
||||
<summary>🖥️ Performance has been tested for</summary>
|
||||
@@ -26,7 +26,7 @@ Set of custom nodes for ComfyUI providing additional image processing capabiliti
|
||||
</details>
|
||||
<details>
|
||||
<summary>🛠 Installation</summary>
|
||||
|
||||
|
||||
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
|
||||
<details>
|
||||
<summary>🦊 RS Label</summary>
|
||||
|
||||
|
||||
# 🦊 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
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Label Image</summary>
|
||||
|
||||
|
||||
# 🦊 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
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Collage Node</summary>
|
||||
|
||||
|
||||
# 🦊 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.
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Image-Prompt</summary>
|
||||
|
||||
|
||||
# 🦊 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
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Decode Save Image</summary>
|
||||
|
||||
|
||||
# 🦊 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.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>🦊 RS rgb2rgba</summary>
|
||||
|
||||
|
||||
# 🦊 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.,
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Any Switch</summary>
|
||||
|
||||
|
||||
# 🦊 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.**
|
||||
|
||||
<img width="511" height="395" alt="Screenshot_1" src="https://github.com/user-attachments/assets/f5b4ee5c-d342-4346-b101-d364d2f714f4" />
|
||||
|
||||
|
||||
### 🔥 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
|
||||
```
|
||||
```
|
||||
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS LoRA Tester</summary>
|
||||
|
||||
|
||||
# 🦊 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
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Save Image LoRA</summary>
|
||||
|
||||
|
||||
# 🦊 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.
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Upscaler</summary>
|
||||
|
||||
|
||||
# 🦊 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:
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Upscale & Resize</summary>
|
||||
|
||||
|
||||
# 🦊 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
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Tile Adjustments</summary>
|
||||
|
||||
|
||||
# 🦊 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
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Tile Assemble</summary>
|
||||
|
||||
|
||||
# 🦊 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
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Ref Encode</summary>
|
||||
|
||||
|
||||
# 🦊 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.
|
||||
</details>
|
||||
<details>
|
||||
<summary>🦊 RS Image Adjustments</summary>
|
||||
|
||||
|
||||
# 🦊 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.
|
||||
|
||||
+155
-32
@@ -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:
|
||||
|
||||
@@ -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('a<b>c: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
|
||||
Reference in New Issue
Block a user