fix: gate the LoraManager serving roots and accept relative link paths
A LoraManager path hand-edited into config.json went straight into find_lora_directories and so into the image-serving allow-list. When a path is configured it must now pass lora_utils.resolve_lora_manager_path (a custom_nodes child whose name contains 'lora'); otherwise, or when the resolver is unavailable, no LoRA directories are served. An empty path keeps the custom_nodes-only auto-detection. POST /images/link now also accepts the relative forms the API hands out (root-relative 'sub/a.png' and ComfyUI-relative 'output/sub/a.png'), resolving them against the allowed roots and anchors; '..' segments in a relative path are refused and unknown files are 404.
This commit is contained in:
Vito Sansevero
+48
-5
@@ -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
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user