From 7c9fb69d424b2d7e2b8b44822b33833da68788c0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 09:51:44 +0000 Subject: [PATCH] Scope the diffusion dataset body-cap passthrough to the exact upload route The multipart upload passthrough was a bare /api/train/diffusion/dataset prefix, so its JSON sub-routes (PUT .../{name}/caption/{filename}, POST .../import-example) also bypassed the default JSON body cap and inherited the far larger upload limit. A large caption/import body would then be buffered and parsed up to the upload cap before the route-level length checks ran. Match the upload route by EXACT path instead: the JSON sub-routes fall through to the normal small-JSON cap. Adds an upload_passthrough_exact_paths param to MaxBodyMiddleware, with regression tests that the sub-routes keep the default cap and a large sub-route body is 413'd. --- studio/backend/main.py | 29 +++++++++--- studio/backend/tests/test_middleware.py | 63 ++++++++++++++++++++++++- 2 files changed, 83 insertions(+), 9 deletions(-) diff --git a/studio/backend/main.py b/studio/backend/main.py index c947c5df5d..5a65b62e57 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -750,21 +750,26 @@ _DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = ( # image upload under the protected /api/train prefix. Like /api/datasets/upload it enforces # its own get_upload_limit_bytes() cap, so it must bypass the default body cap here or the # middleware would 413 near-limit batches (and ignore a raised max_upload_size_mb) before the -# handler runs. Its small-JSON sub-routes (import-example, caption) merely inherit the more -# generous cap, which is harmless. -_DIFFUSION_DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/train/diffusion/dataset" +# handler runs. Matched as an EXACT path (not a prefix): its JSON sub-routes +# (PUT .../{name}/caption/{filename}, POST .../import-example) live under the same prefix but +# must keep the normal small-JSON cap, or a large caption/import body would be buffered and +# parsed up to the far larger upload limit before the route-level length checks run. +_DIFFUSION_DATASET_UPLOAD_PATH = "/api/train/diffusion/dataset" _BODY_UPLOAD_PASSTHROUGH_PREFIXES = ( _DATASET_UPLOAD_PASSTHROUGH_PREFIX, _DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX, - _DIFFUSION_DATASET_UPLOAD_PASSTHROUGH_PREFIX, ) +# Passthrough routes matched by EXACT path (the multipart upload only), so sibling JSON +# sub-routes under the same prefix are not swept into the generous upload cap. +_BODY_UPLOAD_PASSTHROUGH_EXACT_PATHS = (_DIFFUSION_DATASET_UPLOAD_PATH,) def _get_upload_passthrough_request_max_bytes(path: str) -> int: if path.startswith(_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX): return upload_request_limit_bytes(UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES) - if path.startswith(_DATASET_UPLOAD_PASSTHROUGH_PREFIX) or path.startswith( - _DIFFUSION_DATASET_UPLOAD_PASSTHROUGH_PREFIX + if ( + path.startswith(_DATASET_UPLOAD_PASSTHROUGH_PREFIX) + or path == _DIFFUSION_DATASET_UPLOAD_PATH ): return upload_request_limit_bytes() return default_request_body_limit_bytes() @@ -814,12 +819,21 @@ class MaxBodyMiddleware: protected_prefixes: tuple, upload_passthrough_prefixes: tuple = (), upload_passthrough_max_bytes_getter = None, + upload_passthrough_exact_paths: tuple = (), ): self.app = app self.max_bytes_getter = max_bytes_getter self.protected_prefixes = protected_prefixes self.upload_passthrough_prefixes = upload_passthrough_prefixes self.upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter + # Passthrough routes matched by exact path, not prefix: an upload route whose prefix + # also covers sibling JSON sub-routes that must keep the normal (small) body cap. + self.upload_passthrough_exact_paths = upload_passthrough_exact_paths + + def _is_upload_passthrough(self, path: str) -> bool: + return path in self.upload_passthrough_exact_paths or any( + path.startswith(p) for p in self.upload_passthrough_prefixes + ) def _upload_passthrough_max_bytes(self, path: str) -> int: if self.upload_passthrough_max_bytes_getter is None: @@ -856,7 +870,7 @@ class MaxBodyMiddleware: declared = None break - if any(path.startswith(p) for p in self.upload_passthrough_prefixes): + if self._is_upload_passthrough(path): upload_max_bytes = self._upload_passthrough_max_bytes(path) if declared is None: await _send_411(send) @@ -913,6 +927,7 @@ app.add_middleware( protected_prefixes = _BODY_PROTECTED_PREFIXES, upload_passthrough_prefixes = _BODY_UPLOAD_PASSTHROUGH_PREFIXES, upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes, + upload_passthrough_exact_paths = _BODY_UPLOAD_PASSTHROUGH_EXACT_PATHS, ) # Tracks in-flight inference requests for idle auto-unload; off -> passthrough. diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 8f793a3be6..b2fd1c3613 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -35,6 +35,7 @@ def _make_protected_app( main_module, upload_passthrough_prefixes: tuple = (), upload_passthrough_max_bytes_getter = None, + upload_passthrough_exact_paths: tuple = (), ): app = FastAPI() app.add_middleware( @@ -43,6 +44,7 @@ def _make_protected_app( protected_prefixes = ("/v1/chat/completions", "/api/settings", "/api/train"), upload_passthrough_prefixes = upload_passthrough_prefixes, upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter, + upload_passthrough_exact_paths = upload_passthrough_exact_paths, ) @app.post("/v1/chat/completions") @@ -169,18 +171,39 @@ class TestMaxBodyMiddleware: # The diffusion dataset upload route lives under the protected /api/train prefix, so it # must be in the REAL passthrough allowlist with the DB-aware + multipart-overhead cap; # otherwise MaxBodyMiddleware would 413 near-limit batches (and ignore a raised - # max_upload_size_mb) before the handler's own get_upload_limit_bytes() check runs. + # max_upload_size_mb) before the handler's own get_upload_limit_bytes() check runs. It is + # matched by EXACT path (not prefix) so its JSON sub-routes keep the normal small cap. from utils.upload_limits import ( default_request_body_limit_bytes, upload_request_limit_bytes, ) path = "/api/train/diffusion/dataset" - assert any(path.startswith(p) for p in main_module._BODY_UPLOAD_PASSTHROUGH_PREFIXES) + assert path in main_module._BODY_UPLOAD_PASSTHROUGH_EXACT_PATHS + assert not any(path.startswith(p) for p in main_module._BODY_UPLOAD_PASSTHROUGH_PREFIXES) cap = main_module._get_upload_passthrough_request_max_bytes(path) assert cap == upload_request_limit_bytes() # DB-aware cap + multipart overhead assert cap > default_request_body_limit_bytes() # not the plain default body cap + def test_diffusion_dataset_json_subroutes_keep_default_cap(self, main_module): + # The exact-path passthrough must NOT sweep in the JSON sub-routes that live under the same + # /api/train/diffusion/dataset prefix (PUT .../{name}/caption/{filename}, + # POST .../import-example). A prefix match would let a large caption/import body bypass the + # default JSON cap and be buffered + parsed up to the far larger upload limit. They must + # get the plain default body cap, and never be treated as upload passthrough. + from utils.upload_limits import default_request_body_limit_bytes + for path in ( + "/api/train/diffusion/dataset/my-set/caption/img.png", + "/api/train/diffusion/dataset/import-example", + ): + assert path not in main_module._BODY_UPLOAD_PASSTHROUGH_EXACT_PATHS, path + assert not any( + path.startswith(p) for p in main_module._BODY_UPLOAD_PASSTHROUGH_PREFIXES + ), path + assert main_module._get_upload_passthrough_request_max_bytes(path) == ( + default_request_body_limit_bytes() + ), path + def test_v1_surface_is_body_protected(self, main_module): # /images/generations is mounted at both /api/inference and /v1; the /v1 alias (and every # other /v1 POST route) must be body-capped via the /v1 blanket prefix, or an unbounded @@ -233,6 +256,42 @@ class TestMaxBodyMiddleware: assert r.status_code == 411 assert "Content-Length" in r.json()["detail"] + def test_exact_path_passthrough_does_not_cover_subroutes(self, main_module): + # The exact-path passthrough lifts the cap for the upload path itself, but a sibling + # sub-path under the same prefix stays on the small protected cap: a large JSON body to + # the sub-route is 413'd (buffered+capped), not waved through at the upload limit. + app = FastAPI() + app.add_middleware( + main_module.MaxBodyMiddleware, + max_bytes_getter = lambda: 128, + protected_prefixes = ("/api/train",), + upload_passthrough_exact_paths = ("/api/train/ds",), + upload_passthrough_max_bytes_getter = lambda path: 10_000, + ) + + @app.post("/api/train/ds") + async def _upload(request: Request): + total = 0 + async for chunk in request.stream(): + total += len(chunk) + return {"ok": True, "total": total} + + @app.post("/api/train/ds/import-example") + async def _import(payload: dict): + return {"ok": True} + + c = TestClient(app) + # The exact upload path takes the large cap: a 512-byte body passes. + r = c.post( + "/api/train/ds", + content = b"x" * 512, + headers = {"content-type": "application/octet-stream"}, + ) + assert r.status_code == 200 and r.json()["total"] == 512 + # The sibling JSON sub-route keeps the 128-byte default cap: a large body is 413'd. + r = c.post("/api/train/ds/import-example", json = {"text": "x" * 5000}) + assert r.status_code == 413 + # SecurityHeadersMiddleware / CSP