fix: autotag endpoints no longer read arbitrary paths or run on GET

py/api/autotag_routes.py

@@ -2,11 +2,23 @@
 
 import asyncio
 import json
+import os
 from pathlib import Path
+from typing import Optional
 
 from aiohttp import web
 
 
+def _public_error(exc: Exception) -> str:
+    """Describe an exception without leaking absolute server paths."""
+    if isinstance(exc, OSError) and exc.strerror:
+        filename = exc.filename
+        if filename:
+            return f"{exc.strerror}: {os.path.basename(str(filename))}"
+        return exc.strerror
+    return str(exc)
+
+
 class AutotagRoutesMixin:
     """Mixin providing auto-tagging API endpoints."""
 
@@ -15,11 +27,11 @@ class AutotagRoutesMixin:
         async def get_autotag_models_route(request):
             return await self.get_autotag_models(request)
 
-        @routes.get("/prompt_manager/autotag/download/{model_type}")
+        @routes.post("/prompt_manager/autotag/download/{model_type}")
         async def download_autotag_model_route(request):
             return await self.download_autotag_model(request)
 
-        @routes.get("/prompt_manager/autotag/start")
+        @routes.post("/prompt_manager/autotag/start")
         async def start_autotag_route(request):
             return await self.start_autotag(request)
 
@@ -116,17 +128,39 @@ class AutotagRoutesMixin:
         await response.write_eof()
         return response
 
+    async def _read_start_params(self, request) -> Optional[dict]:
+        """Merge query-string and optional JSON-body parameters for autotag/start.
+
+        Returns None when a body is present but is not a JSON object.
+        """
+        params = dict(request.query)
+        if request.can_read_body:
+            try:
+                body = await request.json()
+            except (json.JSONDecodeError, UnicodeDecodeError):
+                return None
+            if not isinstance(body, dict):
+                return None
+            params.update(body)
+        return params
+
     async def start_autotag(self, request):
-        """Start batch auto-tagging with streaming progress."""
-        model_type = request.query.get("model_type", "gguf")
-        custom_prompt = request.query.get("prompt", "")
-        skip_tagged = request.query.get("skip_tagged", "true").lower() == "true"
-        keep_in_memory = request.query.get("keep_in_memory", "true").lower() == "true"
+        """Start batch auto-tagging with streaming progress (POST only)."""
+        params = await self._read_start_params(request)
+        if params is None:
+            return web.json_response(
+                {"success": False, "error": "Request body must be a JSON object"},
+                status=400,
+            )
+        model_type = params.get("model_type", "gguf")
+        custom_prompt = params.get("prompt", "")
+        skip_tagged = str(params.get("skip_tagged", "true")).lower() == "true"
+        keep_in_memory = str(params.get("keep_in_memory", "true")).lower() == "true"
         use_gpu = True
 
         # WD14 threshold params
-        general_threshold = request.query.get("general_threshold")
-        character_threshold = request.query.get("character_threshold")
+        general_threshold = params.get("general_threshold")
+        character_threshold = params.get("character_threshold")
         try:
             if general_threshold is not None:
                 general_threshold = float(general_threshold)
@@ -326,11 +360,90 @@ class AutotagRoutesMixin:
         await response.write_eof()
         return response
 
+    def _resolve_allowed_image_path(self, raw_path: str) -> Optional[Path]:
+        """Return the resolved path when it lies inside an output directory."""
+        try:
+            resolved = Path(raw_path).resolve()
+        except (OSError, RuntimeError):
+            return None
+        for output_dir in self._get_all_output_dirs():
+            try:
+                if resolved.is_relative_to(Path(output_dir).resolve()):
+                    return resolved
+            except (OSError, RuntimeError):
+                continue
+        return None
+
+    async def _select_autotag_target(self, data: dict):
+        """Pick the image to tag from an autotag/single body.
+
+        Returns ``(image_path, prompt_id, None)`` on success or
+        ``(None, None, error_response)`` when the request must be rejected.
+        """
+        image_id = data.get("image_id")
+        if image_id is not None:
+            if isinstance(image_id, bool) or not isinstance(image_id, (int, str)):
+                return None, None, self._bad_request("image_id must be an integer")
+            try:
+                image_id = int(image_id)
+            except (TypeError, ValueError):
+                return None, None, self._bad_request("image_id must be an integer")
+            record = await self._run_in_executor(self.db.get_image_by_id, image_id)
+            if not record or not record.get("image_path"):
+                return (
+                    None,
+                    None,
+                    web.json_response(
+                        {"success": False, "error": f"Image {image_id} not found"},
+                        status=404,
+                    ),
+                )
+            return record["image_path"], record.get("prompt_id"), None
+
+        raw_path = data.get("path") or data.get("image_path")
+        if not raw_path or not isinstance(raw_path, str):
+            return None, None, self._bad_request("image_id or path is required")
+
+        resolved = self._resolve_allowed_image_path(raw_path)
+        if resolved is None:
+            return (
+                None,
+                None,
+                web.json_response(
+                    {
+                        "success": False,
+                        "error": (
+                            f"{os.path.basename(raw_path)} is outside the "
+                            "configured output directories"
+                        ),
+                    },
+                    status=403,
+                ),
+            )
+
+        prompt_id = None
+        try:
+            prompt_id = await self._run_in_executor(
+                self.db.get_prompt_id_for_image, raw_path
+            )
+            if prompt_id is None and str(resolved) != raw_path:
+                prompt_id = await self._run_in_executor(
+                    self.db.get_prompt_id_for_image, str(resolved)
+                )
+        except Exception as e:
+            self.logger.warning(f"Could not find linked prompt: {e}")
+        return str(resolved), prompt_id, None
+
+    @staticmethod
+    def _bad_request(message: str):
+        return web.json_response({"success": False, "error": message}, status=400)
+
     async def autotag_single(self, request):
-        """Generate tags for a single image."""
+        """Generate tags for a single image selected by image_id or gallery path."""
         try:
             data = await request.json()
-            image_path = data.get("image_path")
+            if not isinstance(data, dict):
+                return self._bad_request("Request body must be a JSON object")
             model_type = data.get("model_type", "gguf")
             custom_prompt = data.get("prompt")
             use_gpu = data.get("use_gpu", True)
@@ -342,18 +455,13 @@ class AutotagRoutesMixin:
                 if character_threshold is not None:
                     character_threshold = float(character_threshold)
             except (ValueError, TypeError):
-                return web.json_response(
-                    {
-                        "success": False,
-                        "error": "general_threshold and character_threshold must be numeric",
-                    },
-                    status=400,
+                return self._bad_request(
+                    "general_threshold and character_threshold must be numeric"
                 )
 
-            if not image_path:
-                return web.json_response(
-                    {"success": False, "error": "image_path is required"}, status=400
-                )
+            image_path, prompt_id, error = await self._select_autotag_target(data)
+            if error is not None:
+                return error
 
             from ..autotag import get_autotag_service
 
@@ -363,44 +471,34 @@ class AutotagRoutesMixin:
                 not service.is_model_loaded()
                 or service.get_loaded_model_type() != model_type
             ):
-                loop = asyncio.get_event_loop()
-                await loop.run_in_executor(
-                    None, lambda: service.load_model(model_type, use_gpu)
-                )
+                await self._run_in_executor(service.load_model, model_type, use_gpu)
 
             if custom_prompt:
                 service.custom_prompt = custom_prompt
 
-            loop = asyncio.get_event_loop()
-            tags = await loop.run_in_executor(
-                None,
-                lambda: service.generate_tags(
-                    image_path,
-                    general_threshold=general_threshold,
-                    character_threshold=character_threshold,
-                ),
+            tags = await self._run_in_executor(
+                service.generate_tags,
+                image_path,
+                general_threshold=general_threshold,
+                character_threshold=character_threshold,
             )
 
-            prompt_id = None
-            try:
-                prompt_id = await self._run_in_executor(
-                    self.db.get_prompt_id_for_image, image_path
-                )
-            except Exception as e:
-                self.logger.warning(f"Could not find linked prompt: {e}")
-
             return web.json_response(
                 {
                     "success": True,
                     "tags": tags,
                     "prompt_id": prompt_id,
-                    "image_path": image_path,
+                    "filename": os.path.basename(image_path),
                 }
             )
 
+        except json.JSONDecodeError:
+            return self._bad_request("Request body must be valid JSON")
         except Exception as e:
             self.logger.error(f"AutoTag single error: {e}")
-            return web.json_response({"success": False, "error": str(e)}, status=500)
+            return web.json_response(
+                {"success": False, "error": _public_error(e)}, status=500
+            )
 
     async def apply_autotag(self, request):
         """Apply selected tags to a prompt."""

tests/test_api_autotag.py

@@ -0,0 +1,296 @@
+"""Route tests for the autotag endpoints.
+
+The ML engine is replaced by a fake service so no model is ever loaded and
+no network request is made.
+"""
+
+import json
+import os
+import shutil
+import sys
+import tempfile
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+if "folder_paths" not in sys.modules:
+    _stub = types.ModuleType("folder_paths")
+    _stub.models_dir = tempfile.gettempdir()
+    _stub.base_path = tempfile.gettempdir()
+    _stub.get_output_directory = lambda: tempfile.gettempdir()
+    sys.modules["folder_paths"] = _stub
+
+from aiohttp import web
+from aiohttp.test_utils import AioHTTPTestCase
+
+from database.operations import PromptDatabase
+from py.api import PromptManagerAPI
+from utils.hashing import generate_prompt_hash
+
+
+class FakeAutotagService:
+    """Stand-in for AutoTagService that records calls and returns fixed tags."""
+
+    def __init__(self, tags=None, downloaded=True):
+        self.tags = tags if tags is not None else ["1girl", "smile"]
+        self.downloaded = downloaded
+        self.generate_calls = []
+        self.load_calls = []
+        self.download_calls = []
+        self.unload_calls = 0
+        self.custom_prompt = ""
+        self.default_prompt = "default prompt"
+        self.wd14_general_threshold = 0.35
+        self.wd14_character_threshold = 0.85
+        self.models_config = {"gguf": {}, "wd14-vit": {}}
+        self._loaded_type = None
+        self.generate_error = None
+        self.download_result = True
+        self.load_error = None
+
+    def get_models_status(self):
+        return {
+            name: {"downloaded": self.downloaded, "model_path": None}
+            for name in self.models_config
+        }
+
+    def is_model_loaded(self):
+        return self._loaded_type is not None
+
+    def get_loaded_model_type(self):
+        return self._loaded_type
+
+    def load_model(self, model_type, use_gpu=True):
+        if self.load_error is not None:
+            raise self.load_error
+        self.load_calls.append((model_type, use_gpu))
+        self._loaded_type = model_type
+        return True
+
+    def unload_model(self):
+        self.unload_calls += 1
+        self._loaded_type = None
+
+    def download_model(self, model_type, progress_callback=None):
+        self.download_calls.append(model_type)
+        if progress_callback:
+            progress_callback("Downloading...", 50)
+        return self.download_result
+
+    def generate_tags(
+        self, image_path, prompt=None, general_threshold=None, character_threshold=None
+    ):
+        self.generate_calls.append(
+            (image_path, prompt, general_threshold, character_threshold)
+        )
+        if self.generate_error is not None:
+            raise self.generate_error
+        return list(self.tags)
+
+
+class AutotagAPITestCase(AioHTTPTestCase):
+    """Stands up the PromptManager routes with a temp DB, temp output dir and a fake engine."""
+
+    async def get_application(self):
+        self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
+        self._temp_db.close()
+        self.tmp_root = Path(tempfile.mkdtemp())
+        self.output_dir = self.tmp_root / "output"
+        self.output_dir.mkdir()
+        self.secret_dir = self.tmp_root / "secret"
+        self.secret_dir.mkdir()
+
+        self.service = FakeAutotagService()
+        self._service_patch = patch(
+            "py.autotag.get_autotag_service", return_value=self.service
+        )
+        self._service_patch.start()
+
+        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.add_routes(routes)
+        app.router.add_routes(routes)
+        return app
+
+    async def tearDownAsync(self):
+        self._service_patch.stop()
+        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)
+        shutil.rmtree(self.tmp_root, ignore_errors=True)
+
+    # ── helpers ────────────────────────────────────────────────────────
+
+    def _save_prompt(self, text="Test prompt", tags=None):
+        return self.api.db.save_prompt(
+            text=text,
+            category=None,
+            tags=tags or [],
+            rating=None,
+            prompt_hash=generate_prompt_hash(text),
+        )
+
+    def _make_image(self, directory, name="img.png"):
+        path = directory / name
+        path.write_bytes(b"\x89PNG\r\n\x1a\n")
+        return path
+
+    def _link_image(self, prompt_id, path):
+        return self.api.db.link_image_to_prompt(prompt_id, str(path))
+
+    async def _post_json(self, url, body):
+        return await self.client.request(
+            "POST",
+            url,
+            data=json.dumps(body),
+            headers={"Content-Type": "application/json"},
+        )
+
+    @staticmethod
+    def _sse_events(text):
+        events = []
+        for line in text.splitlines():
+            if line.startswith("data: "):
+                events.append(json.loads(line[len("data: ") :]))
+        return events
+
+
+class TestAutotagSingleContainment(AutotagAPITestCase):
+
+    async def test_image_id_tags_the_linked_image(self):
+        prompt_id = self._save_prompt()
+        image = self._make_image(self.output_dir)
+        image_id = self._link_image(prompt_id, image)
+
+        resp = await self._post_json(
+            "/prompt_manager/autotag/single", {"image_id": image_id}
+        )
+
+        self.assertEqual(resp.status, 200)
+        data = await resp.json()
+        self.assertTrue(data["success"])
+        self.assertEqual(data["tags"], ["1girl", "smile"])
+        self.assertEqual(data["prompt_id"], prompt_id)
+        self.assertEqual(len(self.service.generate_calls), 1)
+        self.assertEqual(
+            os.path.realpath(self.service.generate_calls[0][0]),
+            os.path.realpath(str(image)),
+        )
+
+    async def test_path_inside_output_dir_is_tagged(self):
+        image = self._make_image(self.output_dir)
+
+        resp = await self._post_json(
+            "/prompt_manager/autotag/single", {"path": str(image)}
+        )
+
+        self.assertEqual(resp.status, 200)
+        data = await resp.json()
+        self.assertTrue(data["success"])
+        self.assertEqual(data["tags"], ["1girl", "smile"])
+        self.assertEqual(len(self.service.generate_calls), 1)
+
+    async def test_path_outside_output_dir_is_forbidden(self):
+        secret = self._make_image(self.secret_dir, "password.png")
+
+        resp = await self._post_json(
+            "/prompt_manager/autotag/single", {"path": str(secret)}
+        )
+
+        self.assertEqual(resp.status, 403)
+        data = await resp.json()
+        self.assertFalse(data["success"])
+        self.assertNotIn(str(self.secret_dir), data["error"])
+        self.assertEqual(self.service.generate_calls, [])
+
+    async def test_traversal_path_is_forbidden(self):
+        self._make_image(self.secret_dir, "password.png")
+        traversal = os.path.join(str(self.output_dir), "..", "secret", "password.png")
+
+        resp = await self._post_json(
+            "/prompt_manager/autotag/single", {"path": traversal}
+        )
+
+        self.assertEqual(resp.status, 403)
+        self.assertEqual(self.service.generate_calls, [])
+
+    async def test_legacy_image_path_field_is_still_contained(self):
+        secret = self._make_image(self.secret_dir, "password.png")
+
+        resp = await self._post_json(
+            "/prompt_manager/autotag/single", {"image_path": str(secret)}
+        )
+
+        self.assertEqual(resp.status, 403)
+        self.assertEqual(self.service.generate_calls, [])
+
+    async def test_missing_selector_is_400(self):
+        resp = await self._post_json("/prompt_manager/autotag/single", {})
+        self.assertEqual(resp.status, 400)
+        self.assertFalse((await resp.json())["success"])
+
+    async def test_unknown_image_id_is_404(self):
+        resp = await self._post_json(
+            "/prompt_manager/autotag/single", {"image_id": 9999}
+        )
+        self.assertEqual(resp.status, 404)
+        self.assertFalse((await resp.json())["success"])
+
+    async def test_non_integer_image_id_is_400(self):
+        resp = await self._post_json(
+            "/prompt_manager/autotag/single", {"image_id": "abc"}
+        )
+        self.assertEqual(resp.status, 400)
+
+
+class TestAutotagSideEffectRoutesArePostOnly(AutotagAPITestCase):
+
+    async def test_get_download_is_405(self):
+        resp = await self.client.request("GET", "/prompt_manager/autotag/download/gguf")
+        self.assertEqual(resp.status, 405)
+        self.assertEqual(self.service.download_calls, [])
+
+    async def test_get_start_is_405(self):
+        resp = await self.client.request("GET", "/prompt_manager/autotag/start")
+        self.assertEqual(resp.status, 405)
+        self.assertEqual(self.service.load_calls, [])
+
+    async def test_post_download_streams_completion(self):
+        resp = await self.client.request(
+            "POST", "/prompt_manager/autotag/download/gguf"
+        )
+
+        self.assertEqual(resp.status, 200)
+        events = self._sse_events(await resp.text())
+        self.assertEqual(events[-1]["type"], "complete")
+        self.assertEqual(self.service.download_calls, ["gguf"])
+
+    async def test_post_start_with_no_images_completes(self):
+        resp = await self.client.request(
+            "POST", "/prompt_manager/autotag/start?model_type=gguf"
+        )
+
+        self.assertEqual(resp.status, 200)
+        events = self._sse_events(await resp.text())
+        self.assertEqual(events[-1]["type"], "complete")
+        self.assertEqual(events[-1]["processed"], 0)
+        self.assertEqual(self.service.load_calls, [("gguf", True)])
+
+    async def test_models_status_stays_get(self):
+        resp = await self.client.request("GET", "/prompt_manager/autotag/models")
+        self.assertEqual(resp.status, 200)
+        self.assertTrue((await resp.json())["success"])
+
+
+if __name__ == "__main__":
+    unittest.main()