diff --git a/py/api/images.py b/py/api/images.py index 320a583..cdb56c3 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -11,6 +11,63 @@ from pathlib import Path from aiohttp import web from PIL import Image +from .prompts import safe_error_message as _safe_error + +# Only these file types are ever served by the image routes, regardless of +# what sits inside an allowed directory (defence in depth for the gallery). +IMAGE_EXTENSIONS = frozenset( + { + ".png", + ".jpg", + ".jpeg", + ".webp", + ".gif", + ".bmp", + ".tiff", + ".tif", + ".mp4", + ".webm", + ".mov", + } +) + + +def _canonical(path): + """Resolve symlinks and normalise case so paths compare on every OS.""" + return Path(os.path.normcase(os.path.realpath(str(path)))) + + +def _is_within(path, root): + """True when *path* (after realpath) lives under *root* (after realpath).""" + return _canonical(path).is_relative_to(_canonical(root)) + + +def _has_traversal(relative): + """True for absolute paths, drive-qualified paths or any ``..`` segment.""" + if not relative or os.path.isabs(relative) or os.path.splitdrive(relative)[0]: + return True + if relative[0] in ("/", "\\"): + return True + return ".." in re.split(r"[\\/]+", relative) + + +def _forbidden(message="Access denied"): + return web.json_response({"success": False, "error": message}, status=403) + + +def _media_access_error(path, allowed_dirs): + """Return a 403 response when *path* must not be served, else None. + + Fails closed: an empty *allowed_dirs* denies everything. + """ + if Path(path).suffix.lower() not in IMAGE_EXTENSIONS: + return _forbidden("Only media files can be served") + if not allowed_dirs: + return _forbidden("No allowed image directories configured") + if not any(_is_within(path, root) for root in allowed_dirs): + return _forbidden() + return None + class ImageRoutesMixin: """Mixin providing image and gallery-related API endpoints.""" @@ -302,47 +359,19 @@ class ImageRoutesMixin: {"success": False, "error": "Image not found"}, status=404 ) - image_path = Path(image["image_path"]).resolve() + image_path = Path(image["image_path"]) + allowed_dirs = list(self._get_all_output_dirs()) + self._lora_image_dirs() - # Validate path is within any allowed directory - allowed_dirs = list(self._get_all_output_dirs()) + denied = _media_access_error(image_path, allowed_dirs) + if denied is not None: + return denied - # Also allow LoRA directories when integration is enabled - try: - from ..config import IntegrationConfig - - if IntegrationConfig.LORA_MANAGER_ENABLED: - from ..lora_utils import ( - find_lora_directories, - get_lora_image_cache_dir, - ) - - lora_dirs = find_lora_directories( - IntegrationConfig.LORA_MANAGER_PATH - ) - allowed_dirs.extend(Path(d) for d in lora_dirs) - allowed_dirs.append(get_lora_image_cache_dir()) - except Exception: - # LoRA integration is optional — skip if unavailable - pass - - if allowed_dirs: - allowed = any( - image_path.is_relative_to(d.resolve()) for d in allowed_dirs - ) - if not allowed: - return web.json_response( - {"success": False, "error": "Access denied"}, status=403 - ) - - if not image_path.exists(): + if not image_path.is_file(): return web.json_response( {"success": False, "error": "Image file not found"}, status=404 ) - response = web.FileResponse(image_path) - response.headers["Cache-Control"] = "public, max-age=3600" - return response + return self._file_response(image_path) except ValueError: return web.json_response( @@ -350,44 +379,73 @@ class ImageRoutesMixin: ) except Exception as e: self.logger.error(f"Serve image error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) + + def _lora_image_dirs(self): + """Extra directories served when the LoRA integration is enabled.""" + try: + from ..config import IntegrationConfig + + if not IntegrationConfig.LORA_MANAGER_ENABLED: + 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) + ] + dirs.append(Path(get_lora_image_cache_dir())) + return dirs + except Exception: + # LoRA integration is optional; never widen access on failure. + return [] + + @staticmethod + def _file_response(path): + response = web.FileResponse(path) + response.headers["Cache-Control"] = "public, max-age=3600" + return response + + @staticmethod + def _select_root(output_dirs, root_param): + """Narrow *output_dirs* to the ``?root=N`` entry when it is valid.""" + if root_param is None: + return output_dirs + try: + idx = int(root_param) + except ValueError: + return output_dirs + if 0 <= idx < len(output_dirs): + return [output_dirs[idx]] + return output_dirs async def serve_output_image(self, request): - """Serve image file directly from ComfyUI output folder using streamed FileResponse.""" + """Serve a media file from the ComfyUI output folder(s).""" try: filepath = request.match_info["filepath"] + if _has_traversal(filepath): + return _forbidden() + if Path(filepath).suffix.lower() not in IMAGE_EXTENSIONS: + return _forbidden("Only media files can be served") - # Search all configured output directories - output_dirs = self._get_all_output_dirs() + output_dirs = self._select_root( + list(self._get_all_output_dirs()), request.query.get("root") + ) if not output_dirs: - return web.json_response( - {"success": False, "error": "ComfyUI output directory not found"}, - status=404, - ) + return _forbidden("No allowed image directories configured") - # If a root index is provided, use only that root to avoid - # path collisions when multiple roots share the same relative path. - root_idx = request.query.get("root") - if root_idx is not None: - try: - idx = int(root_idx) - if 0 <= idx < len(output_dirs): - output_dirs = [output_dirs[idx]] - except (ValueError, IndexError): - pass - - # Try each root directory until the file is found for output_path in output_dirs: - image_path = (output_path / filepath).resolve() - - # Security check: must be within this output directory - if not image_path.is_relative_to(output_path.resolve()): + image_path = output_path / filepath + if not image_path.exists(): continue - - if image_path.exists(): - response = web.FileResponse(image_path) - response.headers["Cache-Control"] = "public, max-age=3600" - return response + # A symlink that points outside the root is a hard deny, not + # a fall-through to the next root. + if not _is_within(image_path, output_path): + return _forbidden() + if image_path.is_file(): + return self._file_response(image_path) return web.json_response( {"success": False, "error": "Image file not found"}, status=404 @@ -395,7 +453,9 @@ class ImageRoutesMixin: except Exception as e: self.logger.error(f"Serve output image error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) async def get_gallery_subfolders(self, request): """Get distinct subfolders from gallery output directories.""" diff --git a/py/api/prompts.py b/py/api/prompts.py index 5d0efdf..0c0a57d 100644 --- a/py/api/prompts.py +++ b/py/api/prompts.py @@ -2,9 +2,25 @@ import datetime import json +import os from aiohttp import web + +def safe_error_message(exc): + """Describe *exc* for a client without leaking absolute server paths. + + OSError carries the offending path in ``filename``; only its basename + is echoed back together with ``strerror``. + """ + if isinstance(exc, OSError): + reason = exc.strerror or exc.__class__.__name__ + if exc.filename: + return f"{reason}: {os.path.basename(str(exc.filename))}" + return reason + return str(exc) or exc.__class__.__name__ + + try: from ...utils.validators import ( validate_prompt_text, diff --git a/tests/test_api_images.py b/tests/test_api_images.py new file mode 100644 index 0000000..0bd34e8 --- /dev/null +++ b/tests/test_api_images.py @@ -0,0 +1,280 @@ +""" +Route tests for the image API (py/api/images.py). + +Uses aiohttp's test client against a real PromptManagerAPI wired to a +temporary SQLite database and a temporary output directory. +""" + +import os +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +# ComfyUI-only modules are stubbed so the API can be imported standalone. +sys.modules.setdefault("folder_paths", MagicMock()) +sys.modules.setdefault("server", MagicMock()) + +from aiohttp import web # noqa: E402 +from aiohttp.test_utils import AioHTTPTestCase # noqa: E402 +from PIL import Image # noqa: E402 + +from database.operations import PromptDatabase # noqa: E402 +from py.api import PromptManagerAPI # noqa: E402 +from utils.hashing import generate_prompt_hash # noqa: E402 + + +def make_png(path, size=(4, 4), color=(255, 0, 0)): + """Write a tiny but valid PNG to *path* and return it as a Path.""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with Image.new("RGB", size, color) as img: + img.save(path, "PNG") + return path + + +class ImageAPITestCase(AioHTTPTestCase): + """App with PromptManager routes, a temp DB and a temp output directory.""" + + async def get_application(self): + self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") + self._temp_db.close() + + self._output_tmp = tempfile.TemporaryDirectory() + self.output_dir = Path(os.path.realpath(self._output_tmp.name)) + + app = web.Application() + routes = web.RouteTableDef() + + self.api = PromptManagerAPI() + self.api.db = PromptDatabase(self._temp_db.name) + self.api._get_all_output_dirs = lambda: [self.output_dir] + self.api._find_comfyui_output_dir = lambda: str(self.output_dir) + self.api.add_routes(routes) + app.router.add_routes(routes) + return app + + async def tearDownAsync(self): + for path in ( + self._temp_db.name, + self._temp_db.name + "-wal", + self._temp_db.name + "-shm", + ): + if os.path.exists(path): + os.unlink(path) + self._output_tmp.cleanup() + + # -- helpers --------------------------------------------------------- + + def _save_prompt(self, text="Test prompt", **kwargs): + return self.api.db.save_prompt( + text=text, prompt_hash=generate_prompt_hash(text), **kwargs + ) + + def _link_image(self, image_path, text="Linked prompt"): + """Create a prompt and link *image_path* to it; return the image id.""" + prompt_id = self._save_prompt(text) + image_id = self.api.db.link_image_to_prompt(prompt_id, str(image_path)) + self.assertGreater(image_id, 0, "fixture: image link failed") + return image_id + + def _make_symlink(self, target, link): + """Create a symlink or skip the test when the platform refuses.""" + try: + os.symlink(str(target), str(link)) + except (OSError, NotImplementedError, AttributeError) as exc: + self.skipTest(f"symlinks unavailable: {exc}") + + +class TestServeOutputImage(ImageAPITestCase): + """GET /prompt_manager/images/serve/{filepath}""" + + async def test_serves_png_inside_output_dir(self): + make_png(self.output_dir / "sub" / "pic.png") + + resp = await self.client.request( + "GET", "/prompt_manager/images/serve/sub/pic.png" + ) + + self.assertEqual(resp.status, 200) + self.assertEqual(resp.headers["Content-Type"], "image/png") + body = await resp.read() + self.assertTrue(body.startswith(b"\x89PNG")) + + async def test_non_media_extensions_are_forbidden(self): + for name in ("secret.py", "prompts.db", "notes.txt"): + (self.output_dir / name).write_text("not an image") + resp = await self.client.request( + "GET", f"/prompt_manager/images/serve/{name}" + ) + self.assertEqual(resp.status, 403, name) + data = await resp.json() + self.assertFalse(data["success"]) + + async def test_extension_check_is_case_insensitive(self): + make_png(self.output_dir / "UPPER.PNG") + + resp = await self.client.request( + "GET", "/prompt_manager/images/serve/UPPER.PNG" + ) + + self.assertEqual(resp.status, 200) + + async def test_empty_allowed_dirs_fails_closed(self): + make_png(self.output_dir / "pic.png") + self.api._get_all_output_dirs = lambda: [] + + resp = await self.client.request("GET", "/prompt_manager/images/serve/pic.png") + + self.assertEqual(resp.status, 403) + data = await resp.json() + self.assertFalse(data["success"]) + + async def test_dot_dot_traversal_is_forbidden(self): + outside = Path(self._output_tmp.name).parent / "pm_outside_probe.png" + make_png(outside) + self.addCleanup(lambda: outside.exists() and outside.unlink()) + + # "..%2f" survives URL normalisation and reaches the handler as "../". + resp = await self.client.request( + "GET", "/prompt_manager/images/serve/..%2fpm_outside_probe.png" + ) + + self.assertEqual(resp.status, 403) + + async def test_absolute_path_is_forbidden(self): + outside = Path(self._output_tmp.name).parent / "pm_abs_probe.png" + make_png(outside) + self.addCleanup(lambda: outside.exists() and outside.unlink()) + + encoded = str(outside).replace(os.sep, "%2f").replace("/", "%2f") + resp = await self.client.request( + f"GET", f"/prompt_manager/images/serve/{encoded}" + ) + + self.assertEqual(resp.status, 403) + + async def test_symlink_escaping_output_dir_is_forbidden(self): + with tempfile.TemporaryDirectory() as other: + target = make_png(Path(other) / "elsewhere.png") + self._make_symlink(target, self.output_dir / "escape.png") + + resp = await self.client.request( + "GET", "/prompt_manager/images/serve/escape.png" + ) + + self.assertEqual(resp.status, 403) + + async def test_missing_file_is_404(self): + resp = await self.client.request("GET", "/prompt_manager/images/serve/nope.png") + + self.assertEqual(resp.status, 404) + data = await resp.json() + self.assertFalse(data["success"]) + + async def test_root_index_selects_single_root(self): + with tempfile.TemporaryDirectory() as second: + second_dir = Path(os.path.realpath(second)) + make_png(second_dir / "only_here.png") + self.api._get_all_output_dirs = lambda: [self.output_dir, second_dir] + + hit = await self.client.request( + "GET", "/prompt_manager/images/serve/only_here.png?root=1" + ) + miss = await self.client.request( + "GET", "/prompt_manager/images/serve/only_here.png?root=0" + ) + bad_index = await self.client.request( + "GET", "/prompt_manager/images/serve/only_here.png?root=abc" + ) + + self.assertEqual(hit.status, 200) + self.assertEqual(miss.status, 404) + self.assertEqual(bad_index.status, 200) + + +class TestServeImageById(ImageAPITestCase): + """GET /prompt_manager/images/{image_id}/file""" + + async def test_serves_linked_png(self): + png = make_png(self.output_dir / "linked.png") + image_id = self._link_image(png) + + resp = await self.client.request( + "GET", f"/prompt_manager/images/{image_id}/file" + ) + + self.assertEqual(resp.status, 200) + self.assertEqual(resp.headers["Content-Type"], "image/png") + + async def test_non_media_record_is_forbidden(self): + txt = self.output_dir / "linked.txt" + txt.write_text("plain") + image_id = self._link_image(txt) + + resp = await self.client.request( + "GET", f"/prompt_manager/images/{image_id}/file" + ) + + self.assertEqual(resp.status, 403) + + async def test_empty_allowed_dirs_fails_closed(self): + png = make_png(self.output_dir / "linked.png") + image_id = self._link_image(png) + self.api._get_all_output_dirs = lambda: [] + + resp = await self.client.request( + "GET", f"/prompt_manager/images/{image_id}/file" + ) + + self.assertEqual(resp.status, 403) + + async def test_record_outside_allowed_dirs_is_forbidden(self): + with tempfile.TemporaryDirectory() as other: + png = make_png(Path(other) / "outside.png") + image_id = self._link_image(png) + + resp = await self.client.request( + "GET", f"/prompt_manager/images/{image_id}/file" + ) + + self.assertEqual(resp.status, 403) + + async def test_symlink_record_escaping_is_forbidden(self): + with tempfile.TemporaryDirectory() as other: + target = make_png(Path(other) / "elsewhere.png") + link = self.output_dir / "escape.png" + self._make_symlink(target, link) + image_id = self._link_image(link) + + resp = await self.client.request( + "GET", f"/prompt_manager/images/{image_id}/file" + ) + + self.assertEqual(resp.status, 403) + + async def test_missing_file_is_404(self): + image_id = self._link_image(self.output_dir / "gone.png") + + resp = await self.client.request( + "GET", f"/prompt_manager/images/{image_id}/file" + ) + + self.assertEqual(resp.status, 404) + + async def test_unknown_id_is_404(self): + resp = await self.client.request("GET", "/prompt_manager/images/424242/file") + + self.assertEqual(resp.status, 404) + + async def test_non_numeric_id_is_400(self): + resp = await self.client.request("GET", "/prompt_manager/images/abc/file") + + self.assertEqual(resp.status, 400) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_api_prompts.py b/tests/test_api_prompts.py new file mode 100644 index 0000000..83edbad --- /dev/null +++ b/tests/test_api_prompts.py @@ -0,0 +1,82 @@ +""" +Route tests for the prompt API (py/api/prompts.py). +""" + +import os +import sys +import tempfile +import unittest +from unittest.mock import MagicMock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +sys.modules.setdefault("folder_paths", MagicMock()) +sys.modules.setdefault("server", MagicMock()) + +from aiohttp import web # noqa: E402 +from aiohttp.test_utils import AioHTTPTestCase # noqa: E402 + +from database.operations import PromptDatabase # noqa: E402 +from py.api import PromptManagerAPI # noqa: E402 +from py.api.prompts import safe_error_message # noqa: E402 +from utils.hashing import generate_prompt_hash # noqa: E402 + + +class PromptAPITestCase(AioHTTPTestCase): + """App with PromptManager routes and a temporary SQLite database.""" + + async def get_application(self): + self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") + self._temp_db.close() + + app = web.Application() + routes = web.RouteTableDef() + + self.api = PromptManagerAPI() + self.api.db = PromptDatabase(self._temp_db.name) + self.api.add_routes(routes) + app.router.add_routes(routes) + return app + + async def tearDownAsync(self): + for path in ( + self._temp_db.name, + self._temp_db.name + "-wal", + self._temp_db.name + "-shm", + ): + if os.path.exists(path): + os.unlink(path) + + def _save_prompt(self, text="Test prompt", **kwargs): + return self.api.db.save_prompt( + text=text, prompt_hash=generate_prompt_hash(text), **kwargs + ) + + +class TestSafeErrorMessage(unittest.TestCase): + """Error strings returned to clients never contain absolute server paths.""" + + def test_oserror_uses_strerror_and_basename(self): + secret = os.path.join(os.sep, "srv", "comfy", "output", "hidden.png") + exc = FileNotFoundError(2, "No such file or directory", secret) + + message = safe_error_message(exc) + + self.assertNotIn(os.path.dirname(secret), message) + self.assertIn("No such file or directory", message) + self.assertIn("hidden.png", message) + + def test_oserror_without_filename(self): + exc = PermissionError(13, "Permission denied") + + self.assertEqual(safe_error_message(exc), "Permission denied") + + def test_plain_exception_uses_str(self): + self.assertEqual(safe_error_message(ValueError("bad value")), "bad value") + + def test_empty_message_falls_back_to_type_name(self): + self.assertEqual(safe_error_message(RuntimeError()), "RuntimeError") + + +if __name__ == "__main__": + unittest.main()