fix: image routes serve only media files inside allowed directories
serve_image and serve_output_image now reject anything whose extension is not in IMAGE_EXTENSIONS, fail closed when no output directory is configured, and answer 403 (not 404) for dot-dot, absolute and symlink-escaping paths. Containment is checked after realpath+normcase so it behaves the same on Windows, macOS and Linux. 500 responses use safe_error_message so OSError text no longer leaks absolute paths.
This commit is contained in:
Vito Sansevero
+125
-65
@@ -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."""
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user