fix: autotag/single checks real paths and DB rows against output dirs
This commit is contained in:
Vito Sansevero
+43
-29
@@ -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):
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user