fix: sanitize the download filename extension derived from ?filename=

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)

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()