diff --git a/py/api/images.py b/py/api/images.py index 3335773..f7902fa 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -602,10 +602,19 @@ class ImageRoutesMixin: return [] from ..lora_utils import find_lora_directories, get_lora_image_cache_dir - dirs = [ - Path(d) - for d in find_lora_directories(IntegrationConfig.LORA_MANAGER_PATH) - ] + lora_root = IntegrationConfig.LORA_MANAGER_PATH or "" + if lora_root.strip(): + # A configured path comes from config.json, which can be hand + # edited: only a LoraManager install under custom_nodes may + # widen the serving allow-list. An empty path keeps the + # auto-detection, which itself only scans custom_nodes. + from ..lora_utils import resolve_lora_manager_path + + lora_root = resolve_lora_manager_path(lora_root) + if not lora_root: + return [] + + dirs = [Path(d) for d in find_lora_directories(lora_root)] dirs.append(Path(get_lora_image_cache_dir())) return dirs except Exception: @@ -1139,6 +1148,28 @@ class ImageRoutesMixin: {"success": False, "error": _safe_error(e)}, status=500 ) + @staticmethod + def _resolve_client_media_path(raw_path, allowed_dirs): + """Absolute path for a client-supplied media path, or None when it + does not name a file inside an allowed root. + + Absolute paths are used as given (the caller still runs the access + check). Relative paths are the forms this API hands out: relative to + a gallery root (``sub/a.png``) or to a ComfyUI anchor + (``output/sub/a.png``), and must resolve inside an allowed root. + """ + if os.path.isabs(raw_path): + return raw_path + from ..config import GalleryConfig + + for base in [*allowed_dirs, *GalleryConfig.path_anchors()]: + candidate = Path(base) / raw_path + if candidate.is_file() and any( + _is_within(candidate, root) for root in allowed_dirs + ): + return str(candidate) + return None + async def link_image_to_prompt(self, request): """Link a generated image (inside an allowed directory) to a prompt.""" try: @@ -1151,8 +1182,20 @@ class ImageRoutesMixin: if not prompt_id or not isinstance(image_path, str) or not image_path: return bad_request("prompt_id and image_path are required") + if not os.path.isabs(image_path) and _has_traversal(image_path): + return _forbidden() - denied = _media_access_error(image_path, self._get_all_output_dirs()) + allowed_dirs = list(self._get_all_output_dirs()) + resolved = self._resolve_client_media_path(image_path, allowed_dirs) + if resolved is None: + if Path(image_path).suffix.lower() not in IMAGE_EXTENSIONS: + return _forbidden("Only media files can be served") + return web.json_response( + {"success": False, "error": "Image file not found"}, status=404 + ) + image_path = resolved + + denied = _media_access_error(image_path, allowed_dirs) if denied is not None: return denied diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 7f879ea..0798a63 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -10,6 +10,7 @@ import os import sys import tempfile import unittest +from contextlib import ExitStack from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -1005,6 +1006,95 @@ class TestServeErrorBranches(ImageRouteCoverageCase): self.assertEqual(resp.status, 403) +class TestLoraManagerPathAllowList(ImageRouteCoverageCase): + """A hand-edited LoraManager path must not widen what serve_image reads.""" + + def _enter_lora(self, stack, lora, cache, path, resolved): + from py.config import IntegrationConfig + + stack.enter_context( + patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True) + ) + stack.enter_context(patch.object(IntegrationConfig, "LORA_MANAGER_PATH", path)) + stack.enter_context( + patch("py.lora_utils.find_lora_directories", return_value=[lora]) + ) + stack.enter_context( + patch("py.lora_utils.get_lora_image_cache_dir", return_value=cache) + ) + return stack.enter_context( + patch( + "py.lora_utils.resolve_lora_manager_path", + return_value=resolved, + create=True, + ) + ) + + async def test_configured_path_outside_custom_nodes_is_ignored(self): + with ExitStack() as stack: + lora = stack.enter_context(tempfile.TemporaryDirectory()) + cache = stack.enter_context(tempfile.TemporaryDirectory()) + image_id = self._link_image(make_png(Path(lora) / "preview.png")) + self._enter_lora(stack, lora, cache, lora, None) + + self.assertEqual(self.api._lora_image_dirs(), []) + resp = await self.client.request( + "GET", f"/prompt_manager/images/{image_id}/file" + ) + self.assertEqual(resp.status, 403) + + async def test_configured_custom_nodes_child_is_served(self): + configured = "custom_nodes/ComfyUI-Lora-Manager" + with ExitStack() as stack: + lora = stack.enter_context(tempfile.TemporaryDirectory()) + cache = stack.enter_context(tempfile.TemporaryDirectory()) + image_id = self._link_image(make_png(Path(lora) / "preview.png")) + resolve = self._enter_lora(stack, lora, cache, configured, lora) + + resp = await self.client.request( + "GET", f"/prompt_manager/images/{image_id}/file" + ) + self.assertEqual(resp.status, 200) + resolve.assert_called_once_with(configured) + + +class TestLinkImageRelativePaths(ImageAPITestCase): + """POST /images/link accepts the path forms the API hands out.""" + + URL = "/prompt_manager/images/link" + + async def setUpAsync(self): + await super().setUpAsync() + self.png = make_png(self.output_dir / "sub" / "gen.png") + + async def _link(self, image_path): + prompt_id = self._save_prompt(f"link {image_path}") + resp = await self.client.request( + "POST", self.URL, json={"prompt_id": prompt_id, "image_path": image_path} + ) + return resp.status, await resp.json() + + async def test_root_relative_and_base_relative_paths_link(self): + for form in ("sub/gen.png", "output/sub/gen.png"): + status, data = await self._link(form) + + self.assertEqual(status, 200, form) + image = self.api.db.get_image_by_id(data["image_id"]) + self.assertEqual(Path(image["image_path"]), self.png, form) + + async def test_relative_traversal_is_forbidden(self): + make_png(self.comfy_dir / "secret.png") # removed with the temp tree + + status, _ = await self._link("../secret.png") + + self.assertEqual(status, 403) + + async def test_unknown_relative_path_is_404(self): + status, _ = await self._link("sub/missing.png") + + self.assertEqual(status, 404) + + class TestGalleryListing(ImageRouteCoverageCase): async def test_subfolders_with_and_without_ancestors(self):