diff --git a/comfy_cli/command/transfer.py b/comfy_cli/command/transfer.py index 08e971fa..9232b3f7 100644 --- a/comfy_cli/command/transfer.py +++ b/comfy_cli/command/transfer.py @@ -265,6 +265,20 @@ def _sanitize_item_name(item: str) -> str: return re.sub(r"[^A-Za-z0-9._-]", "_", item) or "item" +def _sanitize_ext(ext: str) -> str: + """A filesystem-safe extension for a download filename. + + An output server controls `?filename=`, so the suffix read off it can + carry control/ANSI bytes that would land in the on-disk name and in the + saved path echoed to the terminal. Accept only a conservative + alphanumeric extension; anything else falls back to the same `.png` + default the callers already use for a suffix-less name. The 16-char cap + is generous enough for real output extensions (`.safetensors`) while + still bounding the junk a hostile name can produce. + """ + return ext if re.fullmatch(r"\.[A-Za-z0-9]{1,16}", ext) else ".png" + + def _collision_safe_path(path: Path) -> Path: """Never overwrite an existing download: ``name.ext`` → ``name.1.ext``, ``name.2.ext``, … (deterministic, first free slot). @@ -621,14 +635,18 @@ def execute_download( # Derive the extension from the source. A bare path has no # `?filename=` query param, so read the real suffix off the on-disk # file rather than mislabeling everything `.png`; real URLs carry the - # name in the query param. + # name in the query param. Neither side is trusted: the query param is + # server-controlled, and the local path is only *parsed* out of an + # `outputs` entry that itself arrives from untrusted metadata (see the + # `is_local_job` note above) — so both suffixes are sanitized before + # they reach the on-disk name. if local_source is not None: - ext = local_source.suffix or ".png" + ext = _sanitize_ext(local_source.suffix) else: parsed = urllib.parse.urlparse(url) qs = urllib.parse.parse_qs(parsed.query) remote_name = qs.get("filename", ["output.png"])[0] - ext = Path(remote_name).suffix or ".png" + ext = _sanitize_ext(Path(remote_name).suffix) node_id, item = annotations[idx] if item is not None: safe_item = _sanitize_item_name(item) diff --git a/tests/comfy_cli/command/test_transfer_download.py b/tests/comfy_cli/command/test_transfer_download.py index 0a395830..e70fd8cf 100644 --- a/tests/comfy_cli/command/test_transfer_download.py +++ b/tests/comfy_cli/command/test_transfer_download.py @@ -12,6 +12,7 @@ import json import os import sys +import urllib.parse from pathlib import Path from unittest.mock import MagicMock, patch @@ -199,6 +200,80 @@ def test_download_data_validates_against_transfer_schema(self, fake_target, tmp_ assert "item" in download_files +def _write_state_with_urls(fake_target, urls: list[str]) -> None: + """State file carrying arbitrary output URLs (no record/item_map), so a + test can drive the `?filename=` the server would have chosen.""" + state = jobs_state.JobState( + prompt_id=PROMPT_ID, + client_id=None, + workflow="/abs/composed.json", + where="cloud", + base_url=fake_target.base_url, + status="completed", + outputs=list(urls), + record=None, + item_map=None, + ) + assert jobs_state.write(state) is not None + + +def _download_with_remote_name(fake_target, tmp_path, capsys, remote_name=None) -> Path: + """Download a single output whose `?filename=` is `remote_name` (omitted + from the query entirely when None); return the saved path.""" + query = {"type": "output"} + if remote_name is not None: + query["filename"] = remote_name + url = f"{fake_target.base_url}/api/view?{urllib.parse.urlencode(query)}" + _write_state_with_urls(fake_target, [url]) + paths, _ = _run_download(fake_target, tmp_path, capsys) + assert len(paths) == 1 + return Path(paths[0]) + + +class TestDownloadExtensionSanitization: + """`?filename=` is server-controlled, so the extension derived from it is + sanitized before it reaches the on-disk name or the echoed path.""" + + def test_ordinary_extension_is_preserved(self, fake_target, tmp_path, capsys): + path = _download_with_remote_name(fake_target, tmp_path, capsys, "out.png") + assert path.name == f"{SHORT_ID}_000.png" + + def test_long_real_extension_is_preserved(self, fake_target, tmp_path, capsys): + # The whitelist is a character class, not a `.png`-only list — a real + # long output extension must survive it. + path = _download_with_remote_name(fake_target, tmp_path, capsys, "out.safetensors") + assert path.name == f"{SHORT_ID}_000.safetensors" + + def test_control_bytes_never_reach_the_filename(self, fake_target, tmp_path, capsys): + path = _download_with_remote_name(fake_target, tmp_path, capsys, "out.png\x1b[31mHACK") + assert path.name == f"{SHORT_ID}_000.png" + assert not any(ord(c) < 32 for c in path.name) + assert path.is_file() + + @pytest.mark.parametrize( + "remote_name", + [ + "out.", # trailing dot → empty suffix + "out", # no extension at all + None, # no `filename` param → the "output.png" default + "out.p*ng", # non-alphanumeric byte in the suffix + "out." + "a" * 17, # absurdly long suffix + ], + ) + def test_unusable_extension_falls_back_to_png(self, fake_target, tmp_path, capsys, remote_name): + path = _download_with_remote_name(fake_target, tmp_path, capsys, remote_name) + assert path.name == f"{SHORT_ID}_000.png" + + def test_traversal_stays_impossible(self, fake_target, tmp_path, capsys): + # Regression guard: only the last component's suffix is ever read, so a + # traversal-shaped name cannot escape `dest` (it never could — this + # pins that the sanitizer did not introduce a way to). + path = _download_with_remote_name(fake_target, tmp_path, capsys, "x.../../../etc/passwd") + dest = (tmp_path / "out").resolve() + assert path.resolve().parent == dest + assert path.name == f"{SHORT_ID}_000.png" + + class TestExtractOutputEntries: """Module-level pure flatten — the record half of Client.extract_outputs, usable without a Target (download joins URLs to it by query params).""" @@ -700,6 +775,23 @@ def test_copied_output_keeps_real_extension_not_png(self, fake_target, tmp_path, assert Path(paths[0]).suffix == ".webp" + def test_copied_output_extension_is_sanitized(self, fake_target, tmp_path, capsys): + # The output reference is a string from the job state's `outputs`, + # which arrives from untrusted metadata (a piped stdin envelope, or a + # cloud `record`) — `_local_source_path` only parses it, never proves + # the name came from a real local render. So a suffix carrying + # control/ANSI bytes must not reach the on-disk name or the echoed + # saved path, exactly as on the remote branch. + src = tmp_path / "output" / "evil.png\x1b[31mHACK" + src.parent.mkdir() + src.write_bytes(b"\x89PNG-local-bytes") + self._write_local_state([str(src)]) + + paths, _ = self._run(tmp_path, capsys, fake_target) + + assert Path(paths[0]).suffix == ".png" + assert "\x1b" not in paths[0] + def test_file_uri_output_is_copied(self, fake_target, tmp_path, capsys): src = tmp_path / "output" / "vid.mp4" src.parent.mkdir()