ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Scope the diffusion dataset body-cap passthrough to the exact upload route

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-07 09:51:44 +00:00
commit 7c9fb69d42
2 changed files with 83 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

29
studio/backend/main.py
View file

@ -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.

63
studio/backend/tests/test_middleware.py
View file

@ -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