fix(transfer): sanitize download extension from untrusted ?filename= param (BE-3326)

comfy_cli/command/transfer.py

@@ -265,6 +265,34 @@ def _sanitize_item_name(item: str) -> str:
     return re.sub(r"[^A-Za-z0-9._-]", "_", item) or "item"
 
 
+_MAX_EXT_LEN = 20  # generous — real extensions are short; caps an untrusted suffix
+
+
+def _sanitize_ext(ext: str) -> str:
+    """A filesystem/terminal-safe download extension.
+
+    The extension can be derived from an untrusted ``?filename=`` query param
+    (a compromised/malicious server), so ``Path(remote_name).suffix`` can carry
+    control/ANSI bytes — e.g. ``out.png\\x1b[31mHACK`` yields ``.png\\x1b[31mHACK``
+    — which would otherwise land in the on-disk name and inject into the
+    terminal when the path is echoed in human mode. Whitelist to a known-safe
+    extension charset, dropping everything else, and cap the length so a
+    ``?filename=out.<thousands of safe chars>`` suffix can't push ``local_name``
+    past NAME_MAX and raise ``OSError(ENAMETOOLONG)`` outside the NDJSON error
+    contract. Directory traversal is already impossible because ``Path(...).suffix``
+    drops path components before we get here.
+
+    A suffix that survives to only dots/dashes/underscores (e.g. ``.💥``,
+    ``.日本語``, ``.<ESC>``) carries no real extension — return ``""`` so the
+    caller's ``or ".png"`` fallback applies instead of writing a name with a
+    bare trailing dot (``<id>_000.``), which is invalid/silently normalized on
+    Windows and desyncs the reported path from the on-disk name."""
+    cleaned = re.sub(r"[^A-Za-z0-9._-]", "", ext)[:_MAX_EXT_LEN]
+    if not any(c.isalnum() for c in cleaned):
+        return ""
+    return cleaned
+
+
 def _collision_safe_path(path: Path) -> Path:
     """Never overwrite an existing download: ``name.ext`` → ``name.1.ext``,
     ``name.2.ext``, … (deterministic, first free slot).
@@ -623,12 +651,12 @@ def execute_download(
         # file rather than mislabeling everything `.png`; real URLs carry the
         # name in the query param.
         if local_source is not None:
-            ext = local_source.suffix or ".png"
+            ext = _sanitize_ext(local_source.suffix) or ".png"
         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) or ".png"
         node_id, item = annotations[idx]
         if item is not None:
             safe_item = _sanitize_item_name(item)

tests/comfy_cli/command/test_transfer_download.py

@@ -923,3 +923,141 @@ def get(self, key):
 
         monkeypatch.setattr(config_manager, "ConfigManager", _FakeCM)
         assert transfer._default_out_dir() == "./outputs"
+
+
+class TestDownloadExtensionSanitized:
+    """The download extension can be derived from an untrusted `?filename=`
+    query param (a compromised/malicious server). Unlike the `item` token, it
+    was never sanitized, so control/ANSI bytes reached the on-disk name and
+    the echoed path — a terminal-injection vector in human mode. `_sanitize_ext`
+    whitelists it to a safe charset while preserving the no-traversal guarantee
+    (BE-3326)."""
+
+    # A real ESC control byte, exactly as a hostile server could return it.
+    _ATTACK_NAME = "out.png\x1b[31mHACK"
+
+    @staticmethod
+    def _has_control_bytes(s: str) -> bool:
+        return any(ord(c) < 0x20 or ord(c) == 0x7F for c in s)
+
+    def _write_cloud_state(self, target: Target, outputs: list[str]) -> None:
+        state = jobs_state.JobState(
+            prompt_id=PROMPT_ID,
+            client_id=None,
+            workflow="/abs/composed.json",
+            where="cloud",
+            base_url=target.base_url,
+            status="completed",
+            outputs=outputs,
+            record=None,
+            item_map=None,
+        )
+        assert jobs_state.write(state) is not None
+
+    def _run(self, fake_target, tmp_path, capsys) -> tuple[list[str], dict]:
+        set_renderer(Renderer(mode=OutputMode.NDJSON, command="download"))
+        with (
+            patch("comfy_cli.command.transfer.resolve_target", return_value=fake_target),
+            patch.object(transfer._DOWNLOAD_OPENER, "open", side_effect=lambda req, timeout=None: _FakeResp()),
+        ):
+            paths = transfer.execute_download(PROMPT_ID, out_dir=str(tmp_path / "out"))
+        lines = [ln for ln in capsys.readouterr().out.splitlines() if ln.strip()]
+        return paths, json.loads(lines[-1])["data"]
+
+    def test_query_param_control_bytes_stripped_from_ext(self, fake_target, tmp_path, capsys):
+        # A hostile server names the output with a control-byte-laden extension.
+        url = f"https://cloud.example.com/api/view?filename={self._ATTACK_NAME}&subfolder=&type=output"
+        self._write_cloud_state(fake_target, [url])
+
+        paths, data = self._run(fake_target, tmp_path, capsys)
+
+        assert len(paths) == 1
+        name = Path(paths[0]).name
+        echoed = data["files"][0]["path"]
+        # No control/ANSI bytes survive into the on-disk name or the echoed path.
+        assert not self._has_control_bytes(name), repr(name)
+        assert not self._has_control_bytes(echoed), repr(echoed)
+        # Exact name: control/ANSI bytes are gone (`\x1b`, `[` stripped) while the
+        # benign alphanumeric payload remnant survives the whitelist. Asserting the
+        # full name (not just a `.png` prefix) catches any regression that leaks
+        # control bytes after `.png`.
+        assert name == f"{SHORT_ID}_000.png31mHACK", name
+        assert Path(paths[0]).is_file()
+
+    def test_query_param_no_directory_traversal(self, fake_target, tmp_path, capsys):
+        # `Path(...).suffix` already drops directory components; confirm a
+        # traversal-shaped filename never escapes the out-dir and degrades to
+        # the `.png` default (no dotted suffix to keep).
+        out_dir = tmp_path / "out"
+        url = "https://cloud.example.com/api/view?filename=../../../etc/passwd&subfolder=&type=output"
+        self._write_cloud_state(fake_target, [url])
+
+        paths, _ = self._run(fake_target, tmp_path, capsys)
+
+        assert len(paths) == 1
+        written = Path(paths[0])
+        assert written.name == f"{SHORT_ID}_000.png"
+        # Nothing was written outside the requested out-dir.
+        assert written.resolve().parent == out_dir.resolve()
+
+    def test_local_source_control_bytes_stripped_from_ext(self, fake_target, tmp_path, capsys):
+        # The local-copy branch derives `ext` from the on-disk source name and
+        # must be sanitized too (BE-3326 applies to both branches).
+        src_dir = tmp_path / "output"
+        src_dir.mkdir()
+        src = src_dir / self._ATTACK_NAME
+        src.write_bytes(b"\x89PNG-local")
+        state = jobs_state.JobState(
+            prompt_id=PROMPT_ID,
+            client_id=None,
+            workflow="/abs/composed.json",
+            where="local",
+            base_url="http://127.0.0.1:8188",
+            status="completed",
+            outputs=[str(src)],
+            record=None,
+            item_map=None,
+        )
+        assert jobs_state.write(state) is not None
+
+        set_renderer(Renderer(mode=OutputMode.NDJSON, command="download"))
+        with patch("comfy_cli.command.transfer.resolve_target", return_value=fake_target):
+            paths = transfer.execute_download(PROMPT_ID, out_dir=str(tmp_path / "out"))
+
+        assert len(paths) == 1
+        name = Path(paths[0]).name
+        assert not self._has_control_bytes(name), repr(name)
+        # Exact name (see query-param twin above): control bytes stripped, benign
+        # payload remnant kept — asserting the full name catches control-byte leaks.
+        assert name == f"{SHORT_ID}_000.png31mHACK", name
+        assert Path(paths[0]).is_file()
+
+    @pytest.mark.parametrize(
+        "suffix",
+        [
+            ".💥",  # emoji-only
+            ".日本語",  # unicode-only
+            ".\x1b",  # lone control byte
+            ".",  # already just a dot
+            "..",  # multiple dots
+            ".-_",  # dots/dashes/underscores, no alnum
+            "",  # empty
+        ],
+    )
+    def test_sanitize_ext_collapses_extensionless_suffix_to_empty(self, suffix):
+        # A suffix with no surviving alphanumeric char carries no real extension:
+        # it must return "" so the caller's `or ".png"` fallback applies rather than
+        # a truthy bare-dot result that bypasses the fallback and writes `<id>_000.`.
+        assert transfer._sanitize_ext(suffix) == ""
+
+    def test_sanitize_ext_keeps_real_extension(self):
+        assert transfer._sanitize_ext(".png") == ".png"
+        assert transfer._sanitize_ext(".7z") == ".7z"
+        assert transfer._sanitize_ext(".tar.gz") == ".tar.gz"
+
+    def test_sanitize_ext_caps_length(self):
+        # A hostile `?filename=out.<thousands of safe chars>` must not yield an
+        # over-long extension that pushes local_name past NAME_MAX.
+        result = transfer._sanitize_ext("." + "a" * 5000)
+        assert len(result) <= transfer._MAX_EXT_LEN
+        assert not self._has_control_bytes(result)