From 7c5b63c5ba81bb4fb06ae3b41f3e739216e89ec8 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 08:28:29 -0700 Subject: [PATCH 1/2] fix(transfer): sanitize download extension from untrusted ?filename= param (BE-3326) The download extension was taken straight from the untrusted `filename` query param via `Path(remote_name).suffix` and embedded into the on-disk name and echoed path. Unlike the `item` token, `ext` was never sanitized, so a malicious server returning `?filename=out.png` plus control/ANSI bytes could inject into the terminal when the path is printed in human mode. Add `_sanitize_ext` next to `_sanitize_item_name` and apply it to `ext` in both the local-source and query-param branches, whitelisting to a safe extension charset. Directory traversal remains impossible (Path.suffix drops path components). Tests pin both the control-byte scrub and the preserved no-traversal behavior. --- comfy_cli/command/transfer.py | 18 ++- .../command/test_transfer_download.py | 103 ++++++++++++++++++ 2 files changed, 119 insertions(+), 2 deletions(-) diff --git a/comfy_cli/command/transfer.py b/comfy_cli/command/transfer.py index 08e971fa..92bdf0f9 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/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; the caller keeps its ``.png`` + fallback for a now-empty result. Directory traversal is already impossible + because ``Path(...).suffix`` drops path components before we get here.""" + return re.sub(r"[^A-Za-z0-9._-]", "", ext) + + 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 +637,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) diff --git a/tests/comfy_cli/command/test_transfer_download.py b/tests/comfy_cli/command/test_transfer_download.py index 0a395830..7301c2d3 100644 --- a/tests/comfy_cli/command/test_transfer_download.py +++ b/tests/comfy_cli/command/test_transfer_download.py @@ -923,3 +923,106 @@ 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) + # The extension stays a clean `.png`-family suffix (leading dot + safe chars). + assert name.startswith(f"{SHORT_ID}_000.png"), 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) + assert name.startswith(f"{SHORT_ID}_000.png"), name + assert Path(paths[0]).is_file() From ff17cd9699bc9003bdcef01b0f1cc7d9e198fbef Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 08:45:36 -0700 Subject: [PATCH 2/2] fix(transfer): harden _sanitize_ext per cursor-review (BE-3326) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address three cursor-review findings on the download-extension sanitizer: - Collapse an extensionless suffix (`.💥`, `.日本語`, lone control byte, or bare `.`) to "" so the caller's `or ".png"` fallback applies instead of a truthy bare-dot result that writes `_000.` — invalid/normalized on Windows and desyncing the reported path from the on-disk name. - Cap the whitelisted extension length (_MAX_EXT_LEN=20) so a hostile `?filename=out.` can't push local_name past NAME_MAX and raise OSError(ENAMETOOLONG) outside the NDJSON error contract. - Assert the exact sanitized name in the control-byte tests (not a `.png` prefix) so a regression leaking control bytes after `.png` is caught; add direct unit tests for the empty-collapse and length-cap behaviors. Co-Authored-By: Claude Opus 4.8 --- comfy_cli/command/transfer.py | 22 ++++++++-- .../command/test_transfer_download.py | 41 +++++++++++++++++-- 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/comfy_cli/command/transfer.py b/comfy_cli/command/transfer.py index 92bdf0f9..be096c6c 100644 --- a/comfy_cli/command/transfer.py +++ b/comfy_cli/command/transfer.py @@ -265,6 +265,9 @@ 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. @@ -273,10 +276,21 @@ def _sanitize_ext(ext: str) -> str: 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; the caller keeps its ``.png`` - fallback for a now-empty result. Directory traversal is already impossible - because ``Path(...).suffix`` drops path components before we get here.""" - return re.sub(r"[^A-Za-z0-9._-]", "", ext) + extension charset, dropping everything else, and cap the length so a + ``?filename=out.`` 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. ``.💥``, + ``.日本語``, ``.``) carries no real extension — return ``""`` so the + caller's ``or ".png"`` fallback applies instead of writing a name with a + bare trailing dot (``_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: diff --git a/tests/comfy_cli/command/test_transfer_download.py b/tests/comfy_cli/command/test_transfer_download.py index 7301c2d3..468be53c 100644 --- a/tests/comfy_cli/command/test_transfer_download.py +++ b/tests/comfy_cli/command/test_transfer_download.py @@ -977,8 +977,11 @@ def test_query_param_control_bytes_stripped_from_ext(self, fake_target, tmp_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) - # The extension stays a clean `.png`-family suffix (leading dot + safe chars). - assert name.startswith(f"{SHORT_ID}_000.png"), name + # 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): @@ -1024,5 +1027,37 @@ def test_local_source_control_bytes_stripped_from_ext(self, fake_target, tmp_pat assert len(paths) == 1 name = Path(paths[0]).name assert not self._has_control_bytes(name), repr(name) - assert name.startswith(f"{SHORT_ID}_000.png"), 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 `_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.` 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)