diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index b9fdc1ee04..57932e7cd8 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1564,9 +1564,16 @@ async def upload_diffusion_dataset( uploaded = 0 allowed = _DIFFUSION_DATASET_IMAGE_EXTS | _DIFFUSION_DATASET_TEXT_EXTS for f in files: - filename = Path(f.filename or "").name.strip().replace("\x00", "") + # Normalise to a safe basename. Path.name does not split on a backslash on POSIX, so a + # Windows client that sends a backslash path in the multipart filename would otherwise be + # stored verbatim; fold backslashes to forward slashes first so the true basename is + # taken for both separators. The read/caption/delete endpoints run the stored name through + # _safe_dataset_image_path (rejects "\\" / ".." / path chars), so a name that still holds + # ".." here would list an image the labeling grid can never preview, caption, or delete -- + # reject it now instead of persisting an unmanageable orphan. + filename = Path((f.filename or "").replace("\\", "/")).name.strip().replace("\x00", "") ext = Path(filename).suffix.lower() - if not filename or ext not in allowed: + if not filename or ".." in filename or ext not in allowed: exts = ", ".join(sorted(allowed)) raise HTTPException( status_code = 400, diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index a5277ccda6..8e898a41f4 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -758,6 +758,34 @@ def test_diffusion_dataset_upload_accumulates(client, dataset_roots): assert r.json()["image_count"] == 3 +def test_diffusion_dataset_upload_normalizes_windows_and_rejects_dotdot(client, dataset_roots): + ds_root, _ = dataset_roots + # A Windows client can send a backslash path in the multipart filename; POSIX Path.name does + # not split on backslash, so it must be folded to the true basename, or the stored name holds + # backslashes that _safe_dataset_image_path later rejects -- an image the labeling grid can + # list but never preview/caption/delete (an orphan). + r = client.post( + "/api/train/diffusion/dataset", + data = {"name": "winset"}, + files = [("files", ("C:\\Users\\me\\pics\\cat.png", b"png-bytes", "image/png"))], + ) + assert r.status_code == 200, r.text + assert (ds_root / "winset" / "cat.png").read_bytes() == b"png-bytes" + # It is listed under the clean basename and the per-image endpoints accept it (not an orphan). + recs = client.get("/api/train/diffusion/dataset/winset/images").json()["images"] + assert any(rec["filename"] == "cat.png" for rec in recs) + assert client.get("/api/train/diffusion/dataset/winset/image/cat.png").status_code == 200 + + # A basename that still contains ".." (which _safe_dataset_image_path rejects) is refused at + # upload rather than persisted as an unmanageable entry. + r = client.post( + "/api/train/diffusion/dataset", + data = {"name": "winset"}, + files = [("files", ("a..b.png", b"x", "image/png"))], + ) + assert r.status_code == 400 and "Unsupported file" in r.json()["detail"] + + def test_diffusion_dataset_upload_over_cap_keeps_existing_example( client, dataset_roots, monkeypatch ):