From 6df0c1851c26a96c19864a0076bfb7f0dfe6b507 Mon Sep 17 00:00:00 2001 From: Simon Pinfold Date: Thu, 18 Jun 2026 12:21:19 +1200 Subject: [PATCH] Add explicit asset upload subfolder handling Amp-Thread-ID: https://ampcode.com/threads/T-019ecf39-2e6f-747d-ae80-addba6b8e4f5 Co-authored-by: Amp --- app/assets/api/routes.py | 2 ++ app/assets/api/schemas_in.py | 15 +++++++-- app/assets/api/upload.py | 4 +++ app/assets/services/ingest.py | 3 +- app/assets/services/path_utils.py | 20 +++++++++-- server.py | 5 ++- .../assets_test/services/test_path_utils.py | 13 ++++++++ tests-unit/assets_test/test_uploads.py | 33 +++++++++++++++++++ 8 files changed, 89 insertions(+), 6 deletions(-) diff --git a/app/assets/api/routes.py b/app/assets/api/routes.py index bd53552d4..cdcb39c79 100644 --- a/app/assets/api/routes.py +++ b/app/assets/api/routes.py @@ -409,6 +409,7 @@ async def upload_asset(request: web.Request) -> web.Response: "hash": parsed.provided_hash, "mime_type": parsed.provided_mime_type, "preview_id": parsed.provided_preview_id, + "subfolder": parsed.provided_subfolder, } ) except ValidationError as ve: @@ -454,6 +455,7 @@ async def upload_asset(request: web.Request) -> web.Response: expected_hash=spec.hash, mime_type=spec.mime_type, preview_id=spec.preview_id, + subfolder=spec.subfolder, ) except AssetValidationError as e: delete_temp_file_if_exists(parsed.tmp_path) diff --git a/app/assets/api/schemas_in.py b/app/assets/api/schemas_in.py index 33565cad6..e588fd63a 100644 --- a/app/assets/api/schemas_in.py +++ b/app/assets/api/schemas_in.py @@ -47,6 +47,7 @@ class ParsedUpload: provided_hash_exists: bool | None provided_mime_type: str | None = None provided_preview_id: str | None = None + provided_subfolder: str | None = None class ListAssetsQuery(BaseModel): @@ -239,8 +240,9 @@ class TagsRemove(TagsAdd): class UploadAssetSpec(BaseModel): """Upload Asset operation. - - tags: optional list; if provided, first is root ('models'|'input'|'output'); - if root == 'models', second must be a valid category + - tags: labels plus one destination role ('models'|'input'|'output') for new bytes; + if role == 'models', exactly one model_type: tag is required + - subfolder: optional destination subfolder for new bytes - name: display name - user_metadata: arbitrary JSON object (optional) - hash: optional canonical 'blake3:' for validation / fast-path @@ -258,6 +260,7 @@ class UploadAssetSpec(BaseModel): hash: str | None = Field(default=None) mime_type: str | None = Field(default=None) preview_id: str | None = Field(default=None) # references an asset_reference id + subfolder: str | None = Field(default=None, max_length=1024) @field_validator("hash", mode="before") @classmethod @@ -315,6 +318,14 @@ class UploadAssetSpec(BaseModel): norm.append(tnorm) return norm + @field_validator("subfolder", mode="before") + @classmethod + def _parse_subfolder(cls, v): + if v is None: + return None + s = str(v).strip() + return s or None + @field_validator("user_metadata", mode="before") @classmethod def _parse_metadata_json(cls, v): diff --git a/app/assets/api/upload.py b/app/assets/api/upload.py index 13d3d372c..ce4a807a6 100644 --- a/app/assets/api/upload.py +++ b/app/assets/api/upload.py @@ -54,6 +54,7 @@ async def parse_multipart_upload( provided_hash_exists: bool | None = None provided_mime_type: str | None = None provided_preview_id: str | None = None + provided_subfolder: str | None = None file_written = 0 tmp_path: str | None = None @@ -140,6 +141,8 @@ async def parse_multipart_upload( provided_mime_type = ((await field.text()) or "").strip() or None elif fname == "preview_id": provided_preview_id = ((await field.text()) or "").strip() or None + elif fname == "subfolder": + provided_subfolder = ((await field.text()) or "").strip() or None if not file_present and not (provided_hash and provided_hash_exists): raise UploadError( @@ -166,6 +169,7 @@ async def parse_multipart_upload( provided_hash_exists=provided_hash_exists, provided_mime_type=provided_mime_type, provided_preview_id=provided_preview_id, + provided_subfolder=provided_subfolder, ) diff --git a/app/assets/services/ingest.py b/app/assets/services/ingest.py index 8b2021a61..28aac33d5 100644 --- a/app/assets/services/ingest.py +++ b/app/assets/services/ingest.py @@ -463,6 +463,7 @@ def upload_from_temp_path( expected_hash: str | None = None, mime_type: str | None = None, preview_id: str | None = None, + subfolder: str | None = None, ) -> UploadResult: try: digest, _ = hashing.compute_blake3_hash(temp_path) @@ -507,7 +508,7 @@ def upload_from_temp_path( if not tags: raise ValueError("tags are required for new asset uploads") - base_dir, subdirs = resolve_destination_from_tags(tags) + base_dir, subdirs = resolve_destination_from_tags(tags, subfolder=subfolder) dest_dir = os.path.join(base_dir, *subdirs) if subdirs else base_dir os.makedirs(dest_dir, exist_ok=True) diff --git a/app/assets/services/path_utils.py b/app/assets/services/path_utils.py index 0e4656fe7..b6b87ba14 100644 --- a/app/assets/services/path_utils.py +++ b/app/assets/services/path_utils.py @@ -25,11 +25,27 @@ def get_comfy_models_folders() -> list[tuple[str, list[str]]]: return targets -def resolve_destination_from_tags(tags: list[str]) -> tuple[str, list[str]]: +def _validate_subfolder(subfolder: str | None) -> list[str]: + if not subfolder: + return [] + + parts = Path(subfolder).parts + invalid = {"", ".", ".."} + if Path(subfolder).is_absolute() or any(part in invalid for part in parts): + raise ValueError("invalid subfolder path") + if any("/" in part or "\\" in part for part in parts): + raise ValueError("invalid subfolder path") + return list(parts) + + +def resolve_destination_from_tags( + tags: list[str], subfolder: str | None = None +) -> tuple[str, list[str]]: """Validates and maps upload routing tags -> (base_dir, subdirs_for_fs). The request tags are only used to choose the write destination. Extra tags remain labels; they do not become path components or trusted classification. + Explicit subfolder is the only request field that can add path components. """ destination_roles = [t for t in tags if t in {"input", "models", "output"}] if len(destination_roles) != 1: @@ -56,7 +72,7 @@ def resolve_destination_from_tags(tags: list[str]) -> tuple[str, list[str]]: else: base_dir = os.path.abspath(folder_paths.get_output_directory()) - return base_dir, [] + return base_dir, _validate_subfolder(subfolder) def validate_path_within_base(candidate: str, base: str) -> None: diff --git a/server.py b/server.py index 361850f38..faf40e501 100644 --- a/server.py +++ b/server.py @@ -440,7 +440,10 @@ class PromptServer(): if args.enable_assets: try: tag = image_upload_type if image_upload_type in ("input", "output") else "input" - result = register_file_in_place(abs_path=filepath, name=filename, tags=[tag]) + tags = [tag] + if subfolder in {"3d", "pasted", "painter", "threed", "webcam"}: + tags.append(subfolder) + result = register_file_in_place(abs_path=filepath, name=filename, tags=tags) resp["asset"] = { "id": result.ref.id, "name": result.ref.name, diff --git a/tests-unit/assets_test/services/test_path_utils.py b/tests-unit/assets_test/services/test_path_utils.py index 5c4871c1f..029d746d8 100644 --- a/tests-unit/assets_test/services/test_path_utils.py +++ b/tests-unit/assets_test/services/test_path_utils.py @@ -171,6 +171,19 @@ class TestGetAssetCategoryAndRelativePath: class TestResolveDestinationFromTags: + def test_explicit_subfolder_is_path_component(self, fake_dirs): + base_dir, subdirs = resolve_destination_from_tags( + ["input", "unit-tests", "foo"], subfolder="foo/bar" + ) + + assert base_dir == os.path.abspath(fake_dirs["input"]) + assert subdirs == ["foo", "bar"] + + @pytest.mark.parametrize("subfolder", ["../escape", "foo/../bar", "/abs", "foo\\bar"]) + def test_explicit_subfolder_rejects_unsafe_paths(self, fake_dirs, subfolder: str): + with pytest.raises(ValueError, match="invalid subfolder"): + resolve_destination_from_tags(["input", "unit-tests"], subfolder=subfolder) + def test_model_upload_rejects_non_writable_registered_folders(self): with tempfile.TemporaryDirectory() as root: root_path = Path(root) diff --git a/tests-unit/assets_test/test_uploads.py b/tests-unit/assets_test/test_uploads.py index 123be298b..7bdab6499 100644 --- a/tests-unit/assets_test/test_uploads.py +++ b/tests-unit/assets_test/test_uploads.py @@ -366,6 +366,39 @@ def test_upload_extra_tags_are_labels_not_path_components(http: requests.Session assert "model_type:checkpoints" in body["tags"] +def test_upload_subfolder_is_explicit_path_component( + http: requests.Session, api_base: str, comfy_tmp_base_dir: Path +): + files = {"file": ("nested.bin", b"nested" * 64, "application/octet-stream")} + form = { + "tags": json.dumps(["input", "unit-tests", "foo"]), + "subfolder": "foo/bar", + "name": "nested.bin", + } + r = http.post(api_base + "/api/assets", data=form, files=files, timeout=120) + body = r.json() + + assert r.status_code == 201, body + stored_name = get_asset_filename(body["asset_hash"], ".bin") + assert (comfy_tmp_base_dir / "input" / "foo" / "bar" / stored_name).exists() + assert body["file_path"] == f"input/foo/bar/{stored_name}" + assert "foo" in body["tags"] + + +def test_upload_rejects_unsafe_subfolder(http: requests.Session, api_base: str): + files = {"file": ("escape.bin", b"escape" * 64, "application/octet-stream")} + form = { + "tags": json.dumps(["input", "unit-tests"]), + "subfolder": "../escape", + "name": "escape.bin", + } + r = http.post(api_base + "/api/assets", data=form, files=files, timeout=120) + body = r.json() + + assert r.status_code == 400, body + assert body["error"]["code"] == "INVALID_BODY" + + def test_multipart_upload_accepts_system_looking_extra_labels( http: requests.Session, api_base: str ):