diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 660793e..80e1615 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -19,6 +19,22 @@ def _public_error(exc: Exception) -> str: return str(exc) +def path_is_within(candidate, root) -> bool: + """True when ``candidate`` is strictly inside ``root`` on the real filesystem. + + Both sides go through ``os.path.realpath`` (symlinks, ``..``) and + ``os.path.normcase`` (Windows case/separator folding) before the + ``Path.is_relative_to`` check, so ``/out2`` is never treated as inside + ``/out`` and a symlink inside ``root`` pointing elsewhere is rejected. + """ + try: + child = Path(os.path.normcase(os.path.realpath(str(candidate)))) + parent = Path(os.path.normcase(os.path.realpath(str(root)))) + except (OSError, RuntimeError, ValueError): + return False + return child != parent and child.is_relative_to(parent) + + class AutotagRoutesMixin: """Mixin providing auto-tagging API endpoints.""" @@ -360,20 +376,27 @@ 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 + def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]: + """Return the real path when it is a file inside an output directory.""" 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 + if path_is_within(raw_path, output_dir): + real = os.path.realpath(raw_path) + return real if os.path.isfile(real) else None return None + @staticmethod + def _forbidden(raw_path: str): + return web.json_response( + { + "success": False, + "error": ( + f"{os.path.basename(str(raw_path).rstrip('/' + os.sep))} is not " + "an image inside the configured output directories" + ), + }, + status=403, + ) + async def _select_autotag_target(self, data: dict): """Pick the image to tag from an autotag/single body. @@ -398,7 +421,11 @@ class AutotagRoutesMixin: status=404, ), ) - return record["image_path"], record.get("prompt_id"), None + stored_path = record["image_path"] + resolved = self._resolve_allowed_image_path(stored_path) + if resolved is None: + return None, None, self._forbidden(stored_path) + return resolved, 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): @@ -406,33 +433,20 @@ class AutotagRoutesMixin: 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, - ), - ) + return None, None, self._forbidden(raw_path) 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: + if prompt_id is None and resolved != raw_path: prompt_id = await self._run_in_executor( - self.db.get_prompt_id_for_image, str(resolved) + self.db.get_prompt_id_for_image, resolved ) except Exception as e: self.logger.warning(f"Could not find linked prompt: {e}") - return str(resolved), prompt_id, None + return resolved, prompt_id, None @staticmethod def _bad_request(message: str): diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 3c151fb..dc77133 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -234,6 +234,58 @@ class TestAutotagSingleContainment(AutotagAPITestCase): self.assertEqual(resp.status, 403) self.assertEqual(self.service.generate_calls, []) + async def test_sibling_prefix_dir_is_forbidden(self): + sibling = self.tmp_root / (self.output_dir.name + "2") + sibling.mkdir() + image = self._make_image(sibling, "x.png") + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"path": str(image)} + ) + + self.assertEqual(resp.status, 403) + self.assertEqual(self.service.generate_calls, []) + + async def test_symlink_escaping_output_dir_is_forbidden(self): + secret = self._make_image(self.secret_dir, "password.png") + link = self.output_dir / "link.png" + try: + link.symlink_to(secret) + except (OSError, NotImplementedError): + self.skipTest("symlinks not supported on this platform") + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"path": str(link)} + ) + + self.assertEqual(resp.status, 403) + self.assertEqual(self.service.generate_calls, []) + + async def test_directory_path_is_forbidden(self): + sub = self.output_dir / "subdir" + sub.mkdir() + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"path": str(sub)} + ) + + self.assertEqual(resp.status, 403) + self.assertEqual(self.service.generate_calls, []) + + async def test_image_id_row_pointing_outside_is_forbidden(self): + prompt_id = self._save_prompt() + secret = self._make_image(self.secret_dir, "password.png") + image_id = self._link_image(prompt_id, secret) + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"image_id": image_id} + ) + + self.assertEqual(resp.status, 403) + data = await resp.json() + self.assertNotIn(str(self.secret_dir), data["error"]) + 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) @@ -253,6 +305,42 @@ class TestAutotagSingleContainment(AutotagAPITestCase): self.assertEqual(resp.status, 400) +class TestPathContainmentHelper(unittest.TestCase): + """Pure checks on the containment predicate used by autotag/single.""" + + def setUp(self): + from py.api.autotag_routes import path_is_within + + self.path_is_within = path_is_within + self.tmp = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmp, True) + self.root = self.tmp / "out" + self.root.mkdir() + + def test_child_inside_root(self): + self.assertTrue(self.path_is_within(self.root / "a" / "b.png", self.root)) + + def test_root_itself_is_not_a_child(self): + self.assertFalse(self.path_is_within(self.root, self.root)) + + def test_prefix_sibling_is_outside(self): + self.assertFalse(self.path_is_within(self.tmp / "out2" / "b.png", self.root)) + + def test_dotdot_is_outside(self): + self.assertFalse( + self.path_is_within(self.root / ".." / "secret" / "b.png", self.root) + ) + + def test_case_differences_are_folded_by_normcase(self): + upper = Path(str(self.root).upper()) / "B.PNG" + with patch("os.path.normcase", side_effect=lambda p: p.lower()): + self.assertTrue(self.path_is_within(upper, self.root)) + + def test_separator_differences_do_not_matter(self): + mixed = str(self.root / "a" / "b.png").replace(os.sep, "/") + self.assertTrue(self.path_is_within(mixed, self.root)) + + class TestAutotagSideEffectRoutesArePostOnly(AutotagAPITestCase): async def test_get_download_is_405(self):