From b6f1e0a2c02343d4da8ae0ed31358ea0f7e874fa Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Wed, 7 Oct 2026 19:50:25 -0700 Subject: [PATCH] feat: capped JSON body reader and single-flight lock for long admin jobs Add _read_json_body(request, max_bytes=1_000_000) to py/api: reads the body in chunks, answers 413 past the cap (declared or streamed) and 400 for invalid or non-object JSON. Every JSON POST in the admin and logging handlers goes through it; it is exposed as PromptManagerAPI._read_json_body for the other route mixins. Add _acquire_long_job(name) / _long_job_busy_response(label) on the API object: one asyncio.Lock per job family, handed out as an async context manager or None when busy, so a handler answers 409 instead of starting a second scan. The duplicate scan and delete share the 'duplicates' family. --- py/api/__init__.py | 88 ++++++++++++++++++++++++ py/api/admin.py | 40 ++++++----- py/api/logging_routes.py | 10 +-- tests/test_api_admin.py | 88 ++++++++++++++++++++++++ tests/test_api_core.py | 141 ++++++++++++++++++++++++++++++++++++++ tests/test_api_logging.py | 12 ++++ 6 files changed, 355 insertions(+), 24 deletions(-) diff --git a/py/api/__init__.py b/py/api/__init__.py index 0afc8a4..6d831a3 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -129,6 +129,64 @@ def _resolve_static_file(root, filepath): return str(candidate) +# ── JSON request bodies ──────────────────────────────────────────────── +JSON_BODY_MAX_BYTES = 1_000_000 +_JSON_READ_CHUNK = 64 * 1024 + + +def _json_error(message, status): + return web.json_response({"success": False, "error": message}, status=status) + + +async def _read_json_body(request, max_bytes=JSON_BODY_MAX_BYTES): + """Read and parse a JSON object body of at most ``max_bytes``. + + The body is read in chunks so an oversized or chunked upload is refused + as soon as it crosses the cap rather than buffered whole. + + Returns: + (data, None) on success, or (None, response) where ``response`` is a + 413 (body too large) or 400 (invalid JSON, or not a JSON object). + """ + too_large = _json_error(f"Request body too large (max {max_bytes} bytes)", 413) + declared = request.content_length + if declared is not None and declared > max_bytes: + return None, too_large + + chunks = [] + size = 0 + async for chunk in request.content.iter_chunked(_JSON_READ_CHUNK): + size += len(chunk) + if size > max_bytes: + return None, too_large + chunks.append(chunk) + + try: + data = json.loads(b"".join(chunks).decode("utf-8")) + except (UnicodeDecodeError, ValueError): + return None, _json_error("Request body must be valid JSON", 400) + if not isinstance(data, dict): + return None, _json_error("Request body must be a JSON object", 400) + return data, None + + +class _LongJobGuard: + """Async context manager handed out by ``_acquire_long_job``.""" + + def __init__(self, lock): + self._lock = lock + + async def __aenter__(self): + # The caller checked ``locked()`` on this same loop iteration, so this + # acquire never waits. + await self._lock.acquire() + return self + + async def __aexit__(self, exc_type, exc, tb): + self._lock.release() + return False + + # ── Gzip compression middleware ──────────────────────────────────────── _GZIP_MIN_SIZE = 1024 # Only compress bodies larger than 1 KB _GZIP_TYPES = frozenset( @@ -212,6 +270,10 @@ class PromptManagerAPI( self._html_cache = {} # Cached HTML file contents keyed by path self._gallery_cache = {} # dict: path_str -> (files, timestamp) self._gallery_cache_ttl = 30 # Cache TTL in seconds + # Single-flight locks for long jobs, keyed by job family; see + # _acquire_long_job(). Created lazily so no lock binds to an event + # loop before the server's loop exists. + self._long_jobs: dict = {} # Run cleanup on initialization to remove any existing duplicates try: @@ -227,6 +289,32 @@ class PromptManagerAPI( _public_path = staticmethod(_public_path) _public_error = staticmethod(_public_error) + _read_json_body = staticmethod(_read_json_body) + + def _acquire_long_job(self, name): + """Claim the single-flight slot for a long job family, or None when busy. + + Usage in a handler:: + + job = self._acquire_long_job("duplicates") + if job is None: + return self._long_job_busy_response("Duplicate scan") + async with job: + ... + + Families are free-form strings ("duplicates", "thumbnails", + "autotag", ...); the lock is released when the ``async with`` block + exits, including on exceptions. Must be called from the event loop. + """ + lock = self._long_jobs.setdefault(name, asyncio.Lock()) + if lock.locked(): + return None + return _LongJobGuard(lock) + + @staticmethod + def _long_job_busy_response(label): + """409 envelope for a long job that is already running.""" + return _json_error(f"{label} already running", 409) async def _run_in_executor(self, func, *args, **kwargs): """Run a blocking function in the default thread pool executor. diff --git a/py/api/admin.py b/py/api/admin.py index 60014d2..64d829a 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -59,24 +59,14 @@ def _thumbnail_rel_path(rel_path, thumbnail_ext): async def _read_json(request): - """Parse a JSON object body. + """Parse a JSON object body through the package's capped reader. Returns: - (data, None) on success, or (None, response) carrying a 400 reply. + (data, None) on success, or (None, response) carrying a 400/413 reply. """ - try: - data = await request.json() - except ValueError: - return None, web.json_response( - {"success": False, "error": "Request body must be valid JSON"}, - status=400, - ) - if not isinstance(data, dict): - return None, web.json_response( - {"success": False, "error": "Request body must be a JSON object"}, - status=400, - ) - return data, None + from . import _read_json_body + + return await _read_json_body(request) def _sse(payload): @@ -253,8 +243,17 @@ class AdminRoutesMixin: async def scan_images_route(request): return await self.scan_images(request) + LONG_JOB_DUPLICATES = "duplicates" + async def scan_duplicates_endpoint(self, request): - """Scan for duplicate images without removing them.""" + """Scan for duplicate images without removing them (single flight).""" + job = self._acquire_long_job(self.LONG_JOB_DUPLICATES) + if job is None: + return self._long_job_busy_response("Duplicate scan") + async with job: + return await self._scan_duplicates(request) + + async def _scan_duplicates(self, request): try: duplicates = await self.find_duplicate_images() @@ -372,7 +371,14 @@ class AdminRoutesMixin: return hash_sha256.hexdigest() async def delete_duplicate_images_endpoint(self, request): - """Delete duplicate image files from disk.""" + """Delete duplicate image files from disk (single flight with the scan).""" + job = self._acquire_long_job(self.LONG_JOB_DUPLICATES) + if job is None: + return self._long_job_busy_response("Duplicate scan or delete") + async with job: + return await self._delete_duplicate_images(request) + + async def _delete_duplicate_images(self, request): try: data, error_response = await _read_json(request) if error_response is not None: diff --git a/py/api/logging_routes.py b/py/api/logging_routes.py index c256b81..77dc261 100644 --- a/py/api/logging_routes.py +++ b/py/api/logging_routes.py @@ -244,13 +244,9 @@ class LoggingRoutesMixin: async def update_log_config(self, request): """Update logging configuration (whitelisted keys only).""" try: - try: - data = await request.json() - except ValueError: - return web.json_response( - {"success": False, "error": "Request body must be valid JSON"}, - status=400, - ) + data, error_response = await self._read_json_body(request) + if error_response is not None: + return error_response clean, error = _validate_log_config(data) if error: diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index beb1a05..4eb7c47 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -711,6 +711,94 @@ class TestDeleteDuplicatesEndpoint(AdminAPITestCase): self.assertNotIn(self.tmpdir, result["failed_files"][0]) +class TestJsonBodyCap(AdminAPITestCase): + """Every JSON POST in the admin API rejects bodies over the shared cap.""" + + JSON_POSTS = ( + "/prompt_manager/settings", + "/prompt_manager/delete_duplicate_images", + "/prompt_manager/diagnostics/test-link", + "/prompt_manager/maintenance", + ) + + async def test_oversized_bodies_are_413(self): + body = json.dumps({"pad": "x" * 1_000_100}).encode() + for path in self.JSON_POSTS: + resp = await self.client.request( + "POST", path, data=body, headers={"Content-Type": "application/json"} + ) + self.assertEqual(resp.status, 413, path) + self.assertFalse((await resp.json())["success"], path) + + +class TestDuplicateJobsAreSingleFlight(AdminAPITestCase): + """A duplicate scan or delete refuses to overlap with a running one.""" + + def _block_scan(self): + import asyncio + + gate = asyncio.Event() + + async def slow_scan(): + await gate.wait() + return [] + + self.api.find_duplicate_images = slow_scan + return gate + + async def test_second_scan_while_one_runs_is_409(self): + import asyncio + + gate = self._block_scan() + first = asyncio.ensure_future( + self.client.request("GET", "/prompt_manager/scan_duplicates") + ) + await asyncio.sleep(0.05) + + second = await asyncio.wait_for( + self.client.request("GET", "/prompt_manager/scan_duplicates"), 5 + ) + self.assertEqual(second.status, 409) + data = await second.json() + self.assertFalse(data["success"]) + self.assertIn("already running", data["error"]) + + gate.set() + self.assertEqual((await first).status, 200) + + async def test_delete_while_scan_runs_is_409_and_scan_after_is_fine(self): + import asyncio + + gate = self._block_scan() + first = asyncio.ensure_future( + self.client.request("GET", "/prompt_manager/scan_duplicates") + ) + await asyncio.sleep(0.05) + + resp = await asyncio.wait_for( + self.client.request( + "POST", + "/prompt_manager/delete_duplicate_images", + json={"image_paths": ["a.png"]}, + ), + 5, + ) + self.assertEqual(resp.status, 409) + + gate.set() + await first + resp = await self.client.request("GET", "/prompt_manager/scan_duplicates") + self.assertEqual(resp.status, 200) + + async def test_lock_is_released_after_a_failed_scan(self): + self.api._find_comfyui_output_dir = _raise + resp = await self.client.request("GET", "/prompt_manager/scan_duplicates") + self.assertEqual(resp.status, 500) + self.api._find_comfyui_output_dir = lambda: None + resp = await self.client.request("GET", "/prompt_manager/scan_duplicates") + self.assertEqual(resp.status, 200) + + class TestSettingsMisc(AdminAPITestCase): async def test_result_timeout_and_display_mode_saved(self): diff --git a/tests/test_api_core.py b/tests/test_api_core.py index d0b8cb6..31ceb33 100644 --- a/tests/test_api_core.py +++ b/tests/test_api_core.py @@ -215,6 +215,147 @@ class TestGzipMiddleware(AioHTTPTestCase): self.assertEqual(len(await resp.read()), 4000) +# ── JSON body reader ────────────────────────────────────────────────── + + +class TestReadJsonBody(AioHTTPTestCase): + """_read_json_body: chunked read with a byte cap, object-only payloads.""" + + async def get_application(self): + async def echo(request): + data, error = await api_module._read_json_body(request) + if error is not None: + return error + return web.json_response({"success": True, "keys": sorted(data)}) + + async def small(request): + data, error = await api_module._read_json_body(request, max_bytes=64) + if error is not None: + return error + return web.json_response({"success": True, "data": data}) + + app = web.Application(client_max_size=4 * 1024 * 1024) + app.router.add_post("/echo", echo) + app.router.add_post("/small", small) + return app + + async def test_valid_object_is_parsed(self): + resp = await self.client.request("POST", "/echo", json={"b": 1, "a": 2}) + self.assertEqual(resp.status, 200) + self.assertEqual((await resp.json())["keys"], ["a", "b"]) + + async def test_default_cap_is_one_megabyte(self): + self.assertEqual(api_module.JSON_BODY_MAX_BYTES, 1_000_000) + body = json.dumps({"pad": "x" * 1_000_100}).encode() + resp = await self.client.request( + "POST", "/echo", data=body, headers={"Content-Type": "application/json"} + ) + self.assertEqual(resp.status, 413) + data = await resp.json() + self.assertFalse(data["success"]) + self.assertIn("too large", data["error"].lower()) + + async def test_body_just_under_cap_is_accepted(self): + payload = {"pad": "x" * (1_000_000 - 20)} + body = json.dumps(payload).encode() + self.assertLessEqual(len(body), 1_000_000) + resp = await self.client.request( + "POST", "/echo", data=body, headers={"Content-Type": "application/json"} + ) + self.assertEqual(resp.status, 200) + + async def test_chunked_upload_without_content_length_is_capped(self): + async def gen(): + for _ in range(10): + yield b'{"pad": "' + b"x" * 20 + b'"}' + + resp = await self.client.request( + "POST", + "/small", + data=gen(), + headers={"Content-Type": "application/json"}, + ) + self.assertEqual(resp.status, 413) + + async def test_declared_oversize_content_length_is_413(self): + body = b'{"pad": "' + b"x" * 200 + b'"}' + resp = await self.client.request( + "POST", "/small", data=body, headers={"Content-Type": "application/json"} + ) + self.assertEqual(resp.status, 413) + + async def test_invalid_json_is_400(self): + for body in (b"{oops", b"\xff\xfe", b""): + resp = await self.client.request( + "POST", + "/echo", + data=body, + headers={"Content-Type": "application/json"}, + ) + self.assertEqual(resp.status, 400, body) + self.assertFalse((await resp.json())["success"]) + + async def test_non_object_json_is_400(self): + for body in (b"[1, 2]", b'"text"', b"3", b"null", b"true"): + resp = await self.client.request( + "POST", + "/echo", + data=body, + headers={"Content-Type": "application/json"}, + ) + self.assertEqual(resp.status, 400, body) + self.assertIn("object", (await resp.json())["error"]) + + def test_helper_is_exposed_on_the_api_object(self): + self.assertIs(PromptManagerAPI._read_json_body, api_module._read_json_body) + + +class TestLongJobLock(unittest.TestCase): + """_acquire_long_job: one job per family at a time, released on exit.""" + + def setUp(self): + self.api = PromptManagerAPI() + + def _run(self, coro): + import asyncio + + return asyncio.run(coro) + + def test_second_acquire_is_refused_until_first_releases(self): + async def scenario(): + first = self.api._acquire_long_job("duplicates") + self.assertIsNotNone(first) + async with first: + self.assertIsNone(self.api._acquire_long_job("duplicates")) + # Other families are independent + other = self.api._acquire_long_job("thumbnails") + self.assertIsNotNone(other) + async with other: + pass + again = self.api._acquire_long_job("duplicates") + self.assertIsNotNone(again) + async with again: + pass + + self._run(scenario()) + + def test_lock_is_released_when_the_job_raises(self): + async def scenario(): + with self.assertRaises(RuntimeError): + async with self.api._acquire_long_job("duplicates"): + raise RuntimeError("boom") + self.assertIsNotNone(self.api._acquire_long_job("duplicates")) + + self._run(scenario()) + + def test_busy_response_is_409_with_envelope(self): + resp = self.api._long_job_busy_response("Duplicate scan") + self.assertEqual(resp.status, 409) + data = json.loads(resp.body) + self.assertFalse(data["success"]) + self.assertEqual(data["error"], "Duplicate scan already running") + + # ── UI and static routes ────────────────────────────────────────────── diff --git a/tests/test_api_logging.py b/tests/test_api_logging.py index 360be5b..3052c6e 100644 --- a/tests/test_api_logging.py +++ b/tests/test_api_logging.py @@ -207,6 +207,18 @@ class TestUpdateLogConfigWhitelist(LoggingAPITestCase): self.assertEqual(resp.status, 400) self.assertEqual(self.manager.update_calls, []) + async def test_oversized_body_is_413(self): + body = json.dumps({"level": "x" * 1_000_100}).encode() + resp = await self.client.request( + "POST", + "/prompt_manager/logs/config", + data=body, + headers={"Content-Type": "application/json"}, + ) + self.assertEqual(resp.status, 413) + self.assertFalse((await resp.json())["success"]) + self.assertEqual(self.manager.update_calls, []) + class TestNoAbsolutePathsInLogResponses(LoggingAPITestCase):