From 824157d4cf9aa21046d2fa0645a7f67d1842a5d9 Mon Sep 17 00:00:00 2001 From: JasonHoku Date: Mon, 16 Feb 2026 17:36:15 -0800 Subject: [PATCH] use comfy built-in endpoitns to comply with comfy best practices --- ProjectStructure.md | 10 +-- __init__.py | 188 +++++++++++++-------------------------- directory_scanner.py | 34 ++++--- resources/logic_utils.js | 10 --- 4 files changed, 83 insertions(+), 159 deletions(-) diff --git a/ProjectStructure.md b/ProjectStructure.md index 8aa10f6..f8d2669 100644 --- a/ProjectStructure.md +++ b/ProjectStructure.md @@ -12,8 +12,10 @@ are hidden from ComfyUI Manager. **All contributors must follow these rules:** only, never called at runtime. 4. **No code obfuscation** — No encoded strings, multi-statement tricks, or undefined variable patterns. -5. **No custom file-serving endpoints** — Use ComfyUI's built-in `WEB_DIRECTORY` mechanism - for serving JS. Never create custom routes that serve arbitrary files from disk. +5. **No custom file-serving endpoints** — Use ComfyUI's built-in `WEB_DIRECTORY` for JS and + the built-in `/view` endpoint for images. Never use `web.FileResponse` or create custom + routes that serve files from disk. For external images, create symlinks into the output + directory and use `/view?filename=X&type=output&subfolder=Y`. 6. **Path containment on all filesystem operations** — Every endpoint that touches the filesystem must validate that resolved paths stay within the expected base directory (e.g., `output/benchmarks/`). Use `os.path.realpath()` + `startswith()` checks. @@ -139,9 +141,7 @@ All routes registered in `__init__.py` via `PromptServer.instance.routes`. | POST | `/config_tester/save_manifest` | Full merge save `{session_name, manifest}` | | POST | `/config_tester/get_session_html` | Generate dashboard HTML on-demand `{session_name, node_id}` | | POST | `/config_tester/export_favorites` | Export favorited images `{session_name, pack_metadata, organize_by_prompt, organize_by_lora}` | -| POST | `/config_tester/scan_directory` | Scan external dir for images `{directory_path, session_name}` | -| GET | `/config_tester/view_external` | Serve whitelisted external images `?path=...` | -| POST | `/config_tester/whitelist_directory` | Re-whitelist dir after restart `{directory_path}` | +| POST | `/config_tester/scan_directory` | Scan external dir, create symlinks, use built-in `/view` `{directory_path, session_name}` | --- diff --git a/__init__.py b/__init__.py index a185ba5..7ecc878 100644 --- a/__init__.py +++ b/__init__.py @@ -15,9 +15,6 @@ from .json_text_node import SmartJSONTextNode from .metadata_packer import pack_metadata_into_image from .directory_scanner import scan_directory_for_images -# Security whitelist for external image serving (populated by scan_directory route) -_whitelisted_directories = set() - # --- PATH SECURITY HELPERS --- def _get_benchmarks_base(): @@ -427,18 +424,16 @@ async def export_favorites(request): file_path = item.get("file", "") # Parse filename from URL format: /view?filename=img_123.webp&type=output&subfolder=benchmarks/Session/images - if file_path.startswith("/config_tester/view_external?"): - # External directory image - parse absolute path from query param - parsed_url = urllib.parse.urlparse(file_path) - url_params = urllib.parse.parse_qs(parsed_url.query) - source_path = url_params.get("path", [""])[0] - filename = os.path.basename(source_path) - elif file_path.startswith("/view?"): - # Extract filename from URL + if file_path.startswith("/view?"): + # Extract filename and subfolder from URL parsed_url = urllib.parse.urlparse(file_path) url_params = urllib.parse.parse_qs(parsed_url.query) filename = url_params.get("filename", [""])[0] - source_path = os.path.join(images_dir, filename) + subfolder = url_params.get("subfolder", [""])[0] + if subfolder: + source_path = os.path.join(folder_paths.get_output_directory(), subfolder, filename) + else: + source_path = os.path.join(images_dir, filename) elif file_path.startswith("./images/"): # Relative path format filename = file_path[9:] # Remove ./images/ @@ -528,6 +523,8 @@ async def scan_directory_route(request): """ Scan an external directory for images, extract PNG metadata, and generate a manifest compatible with the dashboard. + Creates symlinks in the output directory so ComfyUI's built-in + /view endpoint can serve the images (no custom file serving needed). """ try: data = await request.json() @@ -559,16 +556,33 @@ async def scan_directory_route(request): print(f"[DirScanner] Scanning directory: {directory_path}") print(f"[DirScanner] Session name: {session_name}") + # Prepare the symlink directory within the output folder + # so ComfyUI's built-in /view endpoint can serve the images + base_dir = os.path.join(folder_paths.get_output_directory(), "benchmarks", session_name) + + if not _is_path_within(base_dir, _get_benchmarks_base()): + return web.json_response({"error": "Forbidden: path outside benchmarks directory"}, status=403) + + link_dir = os.path.join(base_dir, "external_images") + os.makedirs(link_dir, exist_ok=True) + + # The subfolder path for /view URLs (relative to output dir) + view_subfolder = f"benchmarks/{session_name}/external_images" + # Run the scan in a thread pool to avoid blocking the event loop loop = asyncio.get_event_loop() items, stats = await loop.run_in_executor( - None, scan_directory_for_images, directory_path + None, scan_directory_for_images, directory_path, 5000, view_subfolder ) from_manifest = stats.get("from_manifest", 0) manifest_info = f", {from_manifest} from existing manifest" if from_manifest > 0 else "" print(f"[DirScanner] Found {stats['total']} images ({stats['with_metadata']} with metadata, {stats['skipped']} skipped{manifest_info})") + # Create symlinks (or copy on Windows if symlinks not available) + # for each image so /view can serve them + _create_image_links(directory_path, link_dir, items) + # Build manifest manifest = { "session_name": session_name, @@ -582,12 +596,6 @@ async def scan_directory_route(request): } # Save manifest to benchmarks session directory - base_dir = os.path.join(folder_paths.get_output_directory(), "benchmarks", session_name) - - if not _is_path_within(base_dir, _get_benchmarks_base()): - return web.json_response({"error": "Forbidden: path outside benchmarks directory"}, status=403) - - os.makedirs(base_dir, exist_ok=True) manifest_path = os.path.join(base_dir, "manifest.json") # If manifest already exists, preserve user tags (favorited/rejected/notes) @@ -616,16 +624,7 @@ async def scan_directory_route(request): with open(manifest_path, "w") as f: json.dump(manifest, f, indent=2) - # Whitelist the directory for the view_external route - real_dir = os.path.realpath(directory_path) - _whitelisted_directories.add(real_dir) - # Also whitelist images/ subdirectory if it exists (generated sessions store images there) - images_subdir = os.path.realpath(os.path.join(directory_path, "images")) - if os.path.isdir(images_subdir): - _whitelisted_directories.add(images_subdir) - print(f"[DirScanner] Whitelisted directory: {real_dir} (+ images/)") - else: - print(f"[DirScanner] Whitelisted directory: {real_dir}") + print(f"[DirScanner] Created image links in: {link_dir}") from_manifest = stats.get("from_manifest", 0) return web.json_response({ @@ -646,109 +645,48 @@ async def scan_directory_route(request): return web.json_response({"error": str(e)}, status=500) -# ============================================================================= -# API: SERVE EXTERNAL IMAGES (Security-Whitelisted) -# ============================================================================= - -_IMAGE_CONTENT_TYPES = { - ".png": "image/png", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".webp": "image/webp", - ".bmp": "image/bmp", - ".gif": "image/gif", -} - -@server.PromptServer.instance.routes.get("/config_tester/view_external") -async def view_external_image(request): +def _create_image_links(source_dir, link_dir, items): """ - Serve image files from whitelisted directories. - Only directories registered via scan_directory are allowed. + Create symlinks (or copies on Windows) from source images into + the output directory so ComfyUI's built-in /view endpoint can serve them. """ - try: - encoded_path = request.query.get("path", "") - if not encoded_path: - return web.Response(status=400, text="Missing path parameter") + IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} - file_path = urllib.parse.unquote(encoded_path) + # Collect all source files referenced by items + filenames = set() + for item in items: + src = item.get("source_file") + if src: + filenames.add(src) - # Security: reject directory traversal attempts - if ".." in file_path: - return web.Response(status=403, text="Forbidden: directory traversal") + # Scan source directory and subdirectories for matching files + scan_dirs = [source_dir] + images_subdir = os.path.join(source_dir, "images") + if os.path.isdir(images_subdir): + scan_dirs.append(images_subdir) - # Resolve to real path - real_path = os.path.realpath(file_path) + for scan_dir in scan_dirs: + try: + for entry in os.scandir(scan_dir): + if not entry.is_file(): + continue + ext = os.path.splitext(entry.name)[1].lower() + if ext not in IMAGE_EXTENSIONS: + continue + if entry.name not in filenames: + continue - # Security: check file extension is an image - ext = os.path.splitext(real_path)[1].lower() - if ext not in _IMAGE_CONTENT_TYPES: - return web.Response(status=403, text="Forbidden: not an image file") + link_path = os.path.join(link_dir, entry.name) + if os.path.exists(link_path): + continue - # Security: check the file's directory is whitelisted - real_dir = os.path.dirname(real_path) - is_allowed = any( - real_dir == wl_dir or real_dir.startswith(wl_dir + os.sep) - for wl_dir in _whitelisted_directories - ) - - if not is_allowed: - return web.Response(status=403, text="Forbidden: directory not whitelisted") - - if not os.path.isfile(real_path): - return web.Response(status=404, text="File not found") - - content_type = _IMAGE_CONTENT_TYPES.get(ext, "application/octet-stream") - - return web.FileResponse( - real_path, - headers={ - "Content-Type": content_type, - "Cache-Control": "public, max-age=3600", - } - ) - - except Exception as e: - return web.Response(status=500, text=str(e)) - - -# ============================================================================= -# API: RE-WHITELIST DIRECTORY (for loading previously scanned sessions) -# ============================================================================= - -@server.PromptServer.instance.routes.post("/config_tester/whitelist_directory") -async def whitelist_directory_route(request): - """ - Re-whitelist a directory for serving images after server restart. - Called automatically when loading a scan-type session. - Only image files (checked by extension) can be served from whitelisted dirs. - """ - try: - data = await request.json() - directory_path = data.get("directory_path", "").strip() - - if not directory_path: - return web.Response(status=400, text="Missing directory_path") - - directory_path = os.path.normpath(directory_path) - - if not os.path.isabs(directory_path): - return web.Response(status=400, text="Path must be absolute") - - if not os.path.isdir(directory_path): - return web.Response(status=404, text="Directory not found") - - # Cap whitelist size to prevent unbounded growth - if len(_whitelisted_directories) >= 50: - return web.Response(status=429, text="Too many whitelisted directories (max 50)") - - real_dir = os.path.realpath(directory_path) - _whitelisted_directories.add(real_dir) - print(f"[DirScanner] Re-whitelisted directory: {real_dir}") - - return web.Response(status=200, text="Whitelisted") - - except Exception as e: - return web.Response(status=500, text=str(e)) + try: + os.symlink(entry.path, link_path) + except (OSError, NotImplementedError): + # Windows may need admin privileges for symlinks; fall back to copy + shutil.copy2(entry.path, link_path) + except PermissionError: + print(f"[DirScanner] Warning: Permission denied scanning {scan_dir}") # ============================================================================= diff --git a/directory_scanner.py b/directory_scanner.py index 6caf126..3b47a21 100644 --- a/directory_scanner.py +++ b/directory_scanner.py @@ -8,7 +8,6 @@ import re import json import time import random -import urllib.parse from PIL import Image from .metadata_packer import extract_metadata_from_image @@ -248,21 +247,16 @@ def parse_a1111_parameters(parameters_text): def _extract_filename_from_url(file_url): """ Extract the image filename from a manifest item's file URL. - Handles both generated URLs (/view?filename=xxx&...) and external URLs (/config_tester/view_external?path=...). + Handles /view?filename=xxx&... URLs (both generated and external scan). """ if not file_url: return None - # Generated session format: /view?filename=img_xxx.webp&type=output&subfolder=... + # Standard ComfyUI format: /view?filename=img_xxx.webp&type=output&subfolder=... match = re.search(r'filename=([^&]+)', file_url) if match: - return urllib.parse.unquote(match.group(1)) - - # External scan format: /config_tester/view_external?path=... - match = re.search(r'path=([^&]+)', file_url) - if match: - decoded = urllib.parse.unquote(match.group(1)) - return os.path.basename(decoded) + from urllib.parse import unquote + return unquote(match.group(1)) return None @@ -325,13 +319,15 @@ def _load_existing_manifest_items(directory_path): IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} -def scan_directory_for_images(directory_path, max_items=5000): +def scan_directory_for_images(directory_path, max_items=5000, view_subfolder=""): """ Scan a directory for image files and build manifest-compatible items. Args: directory_path: Absolute path to the directory to scan max_items: Maximum number of images to process + view_subfolder: Subfolder path for ComfyUI's /view endpoint + (e.g. "benchmarks/session/external_images") Returns: tuple: (items_list, stats_dict) @@ -377,9 +373,9 @@ def scan_directory_for_images(directory_path, max_items=5000): # If we have manifest data for this image, use it (much richer metadata) if filename in manifest_items: item = manifest_items[filename].copy() - # Rewrite file URL to use the external view endpoint - encoded_path = urllib.parse.quote(file_path, safe="") - item["file"] = f"/config_tester/view_external?path={encoded_path}" + # Use ComfyUI's built-in /view endpoint with subfolder + from urllib.parse import quote + item["file"] = f"/view?filename={quote(filename, safe='')}&type=output&subfolder={quote(view_subfolder, safe='/')}" item["source_file"] = filename # Ensure it has an ID if "id" not in item: @@ -390,7 +386,7 @@ def scan_directory_for_images(directory_path, max_items=5000): # No manifest entry — fall through to normal image processing try: - item = _process_single_image(file_path, directory_path) + item = _process_single_image(file_path, directory_path, view_subfolder) if item: items.append(item) except Exception as e: @@ -409,7 +405,7 @@ def scan_directory_for_images(directory_path, max_items=5000): return items, stats -def _process_single_image(file_path, directory_path): +def _process_single_image(file_path, directory_path, view_subfolder=""): """ Process a single image file into a manifest-compatible item dict. """ @@ -418,9 +414,9 @@ def _process_single_image(file_path, directory_path): # Generate unique ID (matching existing manifest pattern) ts = int(time.time() * 100000) + random.randint(0, 1000) - # Build the external view URL - encoded_path = urllib.parse.quote(file_path, safe="") - file_url = f"/config_tester/view_external?path={encoded_path}" + # Build URL using ComfyUI's built-in /view endpoint + from urllib.parse import quote + file_url = f"/view?filename={quote(filename, safe='')}&type=output&subfolder={quote(view_subfolder, safe='/')}" # Get image dimensions from PIL try: diff --git a/resources/logic_utils.js b/resources/logic_utils.js index 5abb6a1..0abef20 100644 --- a/resources/logic_utils.js +++ b/resources/logic_utils.js @@ -91,16 +91,6 @@ async function loadSession() { activeData = fullManifest.items || []; meta = fullManifest.meta || {}; - // Auto-whitelist source directory for scan-type sessions - // (images won't load after server restart without this) - if (meta.scan_type === 'directory' && meta.source_directory) { - fetch('/config_tester/whitelist_directory', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ directory_path: meta.source_directory }) - }).catch(e => console.warn('[Load] Could not whitelist directory:', e)); - } - console.log(`[Load] 📦 activeData now has ${activeData.length} items`); // 5. Reset indices