diff --git a/comfy_cli/command/transfer.py b/comfy_cli/command/transfer.py index e32ef7af..4974c075 100644 --- a/comfy_cli/command/transfer.py +++ b/comfy_cli/command/transfer.py @@ -9,11 +9,13 @@ from __future__ import annotations +import http.client import json import mimetypes +import os import re -import shutil import sys +import tempfile import urllib.error import urllib.parse import urllib.request @@ -106,6 +108,24 @@ def redirect_request(self, req, fp, code, msg, headers, newurl): _TRANSFER_OPENER = urllib.request.build_opener(NoRedirectHandler("redirect refused (auth leak prevention)")) _DOWNLOAD_OPENER = urllib.request.build_opener(_DownloadRedirectHandler()) +# Per-output safety cap, shared by the HTTP download stream and local-output copies. +_MAX_DOWNLOAD_BYTES = 10 * 1024 * 1024 * 1024 # 10 GB + +# Per-socket-op (connect / each read) timeout for output downloads: a stalled +# transfer aborts instead of hanging forever, while a steadily-flowing body of +# any size is unaffected. +_DOWNLOAD_TIMEOUT_S = 30 + +# Stream/copy chunk size. 1 MiB keeps syscall volume low on multi-GB outputs +# while still bounding memory and letting the size cap trip promptly. +_DOWNLOAD_CHUNK = 1024 * 1024 + +# The process umask, captured once at import. os has no getter, so reading it +# means the classic set-and-restore dance; doing it here (single-threaded under +# the import lock) avoids a per-download window where the process umask is 0. +_UMASK = os.umask(0) +os.umask(_UMASK) + def _sanitize_multipart_filename(name: str) -> str: """Escape a filename for use in Content-Disposition per RFC 7578. @@ -123,6 +143,83 @@ def _assert_download_url(url: str) -> None: raise ValueError(f"refusing to download from non-HTTP URL: {url}") +def _declared_content_length(resp: Any) -> int | None: + """Parse the response's Content-Length header; None when absent or invalid. + + A misconfigured proxy can fold duplicate headers into a single + ``"123, 123"`` value; accept it only when every part agrees (so the + verified length is unambiguous), otherwise treat the header as absent + rather than letting ``int()`` raise and silently skip verification. + """ + raw = resp.headers.get("Content-Length") + if raw is None: + return None + parts = {p.strip() for p in str(raw).split(",")} + if len(parts) != 1: + return None + try: + value = int(parts.pop()) + except ValueError: + return None + return value if value >= 0 else None + + +def _open_part_file(dst: Path) -> tuple[Any, Path]: + """Exclusively create a random ``..part`` sibling of ``dst``, + returning it open for binary writing plus its path. + + mkstemp's O_EXCL + random name defeat a symlink planted in the out-dir + (a plain ``open("wb")`` would follow it) and keep concurrent downloads + off each other's temp files; its restrictive 0600 mode is widened to the + process umask so the renamed result keeps ``open("wb")``-equivalent + permissions. Returning an already-open file keeps the raw descriptor from + ever crossing back to a caller, and any failure here closes the fd and + removes the temp file so it can't leak a descriptor or orphan a ``.part``. + """ + fd, name = tempfile.mkstemp(dir=str(dst.parent), prefix=dst.name + ".", suffix=".part") + try: + if hasattr(os, "fchmod"): + os.fchmod(fd, 0o666 & ~_UMASK) + return os.fdopen(fd, "wb"), Path(name) + except OSError: + os.close(fd) + Path(name).unlink(missing_ok=True) + raise + + +def _copy_local_output_capped(src: Path, dst: Path) -> None: + """Copy ``src`` to ``dst`` via an exclusive sibling temp file. + + The caller's pre-copy ``stat()`` cap check can be defeated by a source + that grows mid-copy or a pseudo-file that under-reports ``st_size`` — + enforce ``_MAX_DOWNLOAD_BYTES`` on the bytes actually read, and rename + into place so ``dst`` never holds a partial copy. + """ + part_file, part_path = _open_part_file(dst) + try: + total = 0 + # part_file is entered first so that if open(src) raises (a source + # unlinked between the caller's stat() and here), its already-entered + # context still closes the temp fd — the descriptor never leaks. + with part_file as df, open(src, "rb") as sf: + while True: + chunk = sf.read(_DOWNLOAD_CHUNK) + if not chunk: + break + total += len(chunk) + if total > _MAX_DOWNLOAD_BYTES: + raise ValueError(f"local output exceeds {_MAX_DOWNLOAD_BYTES} byte safety limit") + df.write(chunk) + part_path.replace(dst) + part_path = None + finally: + if part_path is not None: + try: + part_path.unlink(missing_ok=True) + except OSError: + pass + + def _local_source_path(url: str) -> Path | None: """Return the on-disk source for a LOCAL output reference, else ``None``. @@ -551,11 +648,10 @@ def execute_download( details={"url": url, "path": str(local_source), "index": idx}, ) raise typer.Exit(code=1) - # Mirror the HTTP branch's 10 GB safety cap so a pathological source + # Mirror the HTTP branch's safety cap so a pathological source # (e.g. an unbounded pseudo-file that still reports as regular) can't # exhaust the disk. stat() follows the symlinked source, matching # what copyfile actually reads. - max_download = 10 * 1024 * 1024 * 1024 # 10 GB safety cap try: source_size = local_source.stat().st_size except OSError as e: @@ -566,26 +662,21 @@ def execute_download( details={"url": url, "path": str(local_source), "index": idx}, ) raise typer.Exit(code=1) - if source_size > max_download: + if source_size > _MAX_DOWNLOAD_BYTES: renderer.error( code="download_failed", - message=f"Local output {idx} exceeds {max_download} byte safety limit", + message=f"Local output {idx} exceeds {_MAX_DOWNLOAD_BYTES} byte safety limit", hint="the source file is too large to copy", details={"url": url, "path": str(local_source), "size": source_size, "index": idx}, ) raise typer.Exit(code=1) # Wrap the copy like the HTTP branch's failure handling: an OSError - # (permission denied, full/read-only dest) or SameFileError (source - # already equals the collision-safe dest) must surface as a - # structured envelope, not an unhandled traceback that breaks - # machine-mode/NDJSON consumers. + # (permission denied, full/read-only dest) or the mid-copy size cap + # must surface as a structured envelope, not an unhandled traceback + # that breaks machine-mode/NDJSON consumers. try: - shutil.copyfile(local_source, local_path) - except shutil.SameFileError: - # Source and dest are the same file — nothing to copy; the file - # is already at the destination path. - pass - except OSError as e: + _copy_local_output_capped(local_source, local_path) + except (OSError, ValueError) as e: renderer.error( code="download_failed", message=f"Failed to copy local output {idx}: {e}", @@ -609,19 +700,74 @@ def execute_download( for hdr, val in auth_hdrs.items(): req.add_header(hdr, val) - max_download = 10 * 1024 * 1024 * 1024 # 10 GB safety cap + # Stream into an exclusively-created temp file and rename it into + # place only once the body is complete and verified, so + # `local_path` never holds a partial file no matter how the + # transfer dies. + part_path: Path | None = None try: - with _DOWNLOAD_OPENER.open(req) as resp: + with _DOWNLOAD_OPENER.open(req, timeout=_DOWNLOAD_TIMEOUT_S) as resp: + expected = _declared_content_length(resp) + if expected is not None and expected > _MAX_DOWNLOAD_BYTES: + renderer.error( + code="download_failed", + message=f"Output {idx} declares {expected} bytes, over the {_MAX_DOWNLOAD_BYTES} byte safety limit", + hint="the output is too large to download", + details={"url": url, "index": idx, "declared_bytes": expected}, + ) + raise typer.Exit(code=1) + part_file, part_path = _open_part_file(local_path) total = 0 - with open(local_path, "wb") as fp: + with part_file as fp: while True: - chunk = resp.read(65536) + chunk = resp.read(_DOWNLOAD_CHUNK) if not chunk: break total += len(chunk) - if total > max_download: - raise ValueError(f"download exceeds {max_download} byte safety limit") + if expected is not None and total > expected: + # http.client clips plain Content-Length bodies, + # but ignores Content-Length when the response + # is chunked — that pairing could stream far + # past the declared size before the post-loop + # check fires. + renderer.error( + code="download_failed", + message=( + f"Download of output {idx} exceeds its declared " + f"Content-Length of {expected} bytes" + ), + hint="the server sent more data than it declared", + details={ + "url": url, + "index": idx, + "declared_bytes": expected, + "received_bytes": total, + }, + ) + raise typer.Exit(code=1) + if total > _MAX_DOWNLOAD_BYTES: + renderer.error( + code="download_failed", + message=f"Download of output {idx} exceeds {_MAX_DOWNLOAD_BYTES} byte safety limit", + hint="the output is too large to download", + details={"url": url, "index": idx, "received_bytes": total}, + ) + raise typer.Exit(code=1) fp.write(chunk) + # http.client returns EOF instead of raising IncompleteRead when + # a Content-Length body is cut short and read in chunks, so a + # dropped connection otherwise looks like a completed download — + # verify the byte count explicitly. + if expected is not None and total != expected: + renderer.error( + code="download_failed", + message=f"Download of output {idx} truncated: received {total} of {expected} bytes", + hint="the connection dropped mid-transfer; retry the download", + details={"url": url, "index": idx, "declared_bytes": expected, "received_bytes": total}, + ) + raise typer.Exit(code=1) + part_path.replace(local_path) + part_path = None except urllib.error.HTTPError as e: renderer.error( code="download_failed", @@ -630,6 +776,35 @@ def execute_download( details={"status": e.code, "url": url, "index": idx}, ) raise typer.Exit(code=1) + except (OSError, http.client.HTTPException) as e: + # Everything that isn't an HTTP status lands here: URLError + # (refused/DNS/TLS — a subclass of OSError), a socket timeout + # or reset mid-read, filesystem errors from the temp-file + # create/write/rename, and a truncated *chunked* body — which + # raises http.client.IncompleteRead (an HTTPException, not an + # OSError) rather than the silent EOF a Content-Length body + # gives. Emit the envelope instead of a traceback so + # machine-mode consumers keep their contract. + reason = getattr(e, "reason", None) or e + # A bare TimeoutError()/IncompleteRead can stringify to "", + # which would emit a reason-less envelope — fall back to the + # exception's type name so the cause is always diagnosable. + reason_text = str(reason) or type(e).__name__ + renderer.error( + code="download_failed", + message=f"Failed to download output {idx}: {reason_text}", + hint="check that the server is reachable and the out-dir is writable", + details={"url": url, "index": idx, "reason": reason_text}, + ) + raise typer.Exit(code=1) + finally: + # Cleared after the success rename; on every failure path this + # removes the partial download. + if part_path is not None: + try: + part_path.unlink(missing_ok=True) + except OSError: + pass file_size = local_path.stat().st_size entry: dict[str, Any] = { diff --git a/tests/comfy_cli/command/test_transfer_download.py b/tests/comfy_cli/command/test_transfer_download.py index f5b24a43..0a395830 100644 --- a/tests/comfy_cli/command/test_transfer_download.py +++ b/tests/comfy_cli/command/test_transfer_download.py @@ -10,8 +10,10 @@ from __future__ import annotations import json +import os +import sys from pathlib import Path -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -89,13 +91,29 @@ def _write_state(target: Target, *, record=None, item_map=None) -> list[str]: class _FakeResp: - """Context-manager response yielding one chunk then EOF.""" - - def __init__(self, data: bytes = b"\x89PNG-fake"): - self._chunks = [data] + """Context-manager stand-in for http.client.HTTPResponse. ``read(n)`` + honours ``n`` and advances through the body like the real object, so the + streaming loop and its running byte total are faithfully exercised. + Declares a matching Content-Length by default; pass an int to lie, None to + omit. ``chunk_size`` caps each read below ``n`` to force multi-read + streaming (and to model a chunked body that over-delivers past a small + declared Content-Length, which a real clipped Content-Length body cannot).""" + + def __init__(self, data: bytes = b"\x89PNG-fake", *, content_length="auto", chunk_size=None): + self._buf = data + self._pos = 0 + self._chunk_size = chunk_size + self.reads = 0 + if content_length == "auto": + content_length = len(data) + self.headers = {} if content_length is None else {"Content-Length": str(content_length)} def read(self, n: int) -> bytes: - return self._chunks.pop(0) if self._chunks else b"" + self.reads += 1 + take = n if self._chunk_size is None else min(n, self._chunk_size) + chunk = self._buf[self._pos : self._pos + take] + self._pos += len(chunk) + return chunk def __enter__(self): return self @@ -109,7 +127,7 @@ def _run_download(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: _FakeResp()), + 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()] @@ -333,7 +351,7 @@ class TestMachineModeStdoutPurity: def _download(self, fake_target, tmp_path): with ( patch("comfy_cli.command.transfer.resolve_target", return_value=fake_target), - patch.object(transfer._DOWNLOAD_OPENER, "open", side_effect=lambda req: _FakeResp()), + patch.object(transfer._DOWNLOAD_OPENER, "open", side_effect=lambda req, timeout=None: _FakeResp()), ): transfer.execute_download(PROMPT_ID, out_dir=str(tmp_path / "out")) @@ -378,6 +396,249 @@ def test_ndjson_mode_emits_no_human_progress_line(self, fake_target, tmp_path, c assert "downloaded" not in captured.err +class TestDownloadIntegrity: + """A transfer that dies mid-body or oversteps the size cap must fail with a + structured envelope and leave neither the final file nor a `.part` partial + behind; completed downloads land atomically.""" + + def _failing_download(self, fake_target, tmp_path, capsys, resp=None, raises=None) -> dict: + import typer + + def _open(req, **kw): + if raises is not None: + raise raises + return resp + + set_renderer(Renderer(mode=OutputMode.JSON, command="download")) + _write_state(fake_target) + with ( + patch("comfy_cli.command.transfer.resolve_target", return_value=fake_target), + patch.object(transfer._DOWNLOAD_OPENER, "open", side_effect=_open), + ): + with pytest.raises(typer.Exit) as excinfo: + transfer.execute_download(PROMPT_ID, out_dir=str(tmp_path / "out")) + assert excinfo.value.exit_code == 1 + lines = [ln for ln in capsys.readouterr().out.splitlines() if ln.strip()] + envelope = json.loads(lines[-1]) + assert envelope["ok"] is False + return envelope["error"] + + def test_truncated_body_errors_and_leaves_nothing(self, fake_target, tmp_path, capsys): + err = self._failing_download(fake_target, tmp_path, capsys, _FakeResp(b"x" * 400, content_length=1000)) + assert err["code"] == "download_failed" + # Every size failure reports the server-declared length under the same + # key so a machine consumer never has to guess which spelling is present. + assert err["details"]["declared_bytes"] == 1000 + assert err["details"]["received_bytes"] == 400 + assert list((tmp_path / "out").iterdir()) == [] + + def test_declared_size_over_cap_refused_before_body_read(self, fake_target, tmp_path, capsys, monkeypatch): + monkeypatch.setattr(transfer, "_MAX_DOWNLOAD_BYTES", 100) + resp = _FakeResp(b"x" * 40, content_length=500) + err = self._failing_download(fake_target, tmp_path, capsys, resp) + assert err["code"] == "download_failed" + assert err["details"]["declared_bytes"] == 500 + assert resp.reads == 0 + assert list((tmp_path / "out").iterdir()) == [] + + def test_body_exceeding_declared_length_aborts(self, fake_target, tmp_path, capsys): + # Chunked over-delivery is the real trigger: a body streamed in small + # chunks past a declared Content-Length of 100 (a plain Content-Length + # body would be clipped and never over-deliver). The abort must fire + # mid-stream, not only after the loop ends. + resp = _FakeResp(b"x" * 400, content_length=100, chunk_size=32) + err = self._failing_download(fake_target, tmp_path, capsys, resp) + assert err["code"] == "download_failed" + assert err["details"]["declared_bytes"] == 100 + assert err["details"]["received_bytes"] > 100 + assert resp.reads > 1 # streamed across multiple reads before aborting + assert list((tmp_path / "out").iterdir()) == [] + + def test_lengthless_body_over_cap_errors_with_envelope(self, fake_target, tmp_path, capsys, monkeypatch): + monkeypatch.setattr(transfer, "_MAX_DOWNLOAD_BYTES", 100) + resp = _FakeResp(b"x" * 400, content_length=None, chunk_size=32) + err = self._failing_download(fake_target, tmp_path, capsys, resp) + assert err["code"] == "download_failed" + assert err["details"]["received_bytes"] > 100 + assert resp.reads > 1 # cap tripped while accumulating, not on one read + assert list((tmp_path / "out").iterdir()) == [] + + def test_multichunk_body_reassembled_correctly(self, fake_target, tmp_path, capsys): + # A large body delivered across many reads must land byte-exact, so the + # running total and the temp-file writes stay in step with the stream. + body = bytes(range(256)) * 40 # 10240 bytes + set_renderer(Renderer(mode=OutputMode.NDJSON, command="download")) + _write_state(fake_target) + 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(body, chunk_size=100), + ), + ): + paths = transfer.execute_download(PROMPT_ID, out_dir=str(tmp_path / "out")) + assert all(Path(p).read_bytes() == body for p in paths) + + def test_lengthless_body_downloads_without_verification(self, fake_target, tmp_path, capsys): + set_renderer(Renderer(mode=OutputMode.NDJSON, command="download")) + _write_state(fake_target) + 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(content_length=None) + ), + ): + paths = transfer.execute_download(PROMPT_ID, out_dir=str(tmp_path / "out")) + assert len(paths) == 4 + assert all(Path(p).read_bytes() == b"\x89PNG-fake" for p in paths) + + def test_completed_download_leaves_no_part_file(self, fake_target, tmp_path, capsys): + _write_state(fake_target) + paths, _ = _run_download(fake_target, tmp_path, capsys) + assert paths + assert list((tmp_path / "out").glob("*.part")) == [] + + def test_preexisting_part_sibling_survives_success(self, fake_target, tmp_path, capsys): + out = tmp_path / "out" + out.mkdir(parents=True) + bystander = out / f"{SHORT_ID}_000.png.part" + bystander.write_bytes(b"user data, not ours") + _write_state(fake_target) + paths, _ = _run_download(fake_target, tmp_path, capsys) + assert len(paths) == 4 + assert bystander.read_bytes() == b"user data, not ours" + + def test_preexisting_part_sibling_survives_failure(self, fake_target, tmp_path, capsys): + out = tmp_path / "out" + out.mkdir(parents=True) + bystander = out / f"{SHORT_ID}_000.png.part" + bystander.write_bytes(b"user data, not ours") + err = self._failing_download(fake_target, tmp_path, capsys, _FakeResp(b"x" * 40, content_length=100)) + assert err["code"] == "download_failed" + assert bystander.read_bytes() == b"user data, not ours" + + @pytest.mark.skipif(sys.platform == "win32", reason="POSIX permission bits") + def test_downloaded_file_keeps_umask_permissions(self, fake_target, tmp_path, capsys): + _write_state(fake_target) + paths, _ = _run_download(fake_target, tmp_path, capsys) + umask = os.umask(0) + os.umask(umask) + assert (Path(paths[0]).stat().st_mode & 0o777) == 0o666 & ~umask + + def test_connection_failure_emits_envelope_not_traceback(self, fake_target, tmp_path, capsys): + import urllib.error + + err = self._failing_download( + fake_target, tmp_path, capsys, raises=urllib.error.URLError("[Errno 111] Connection refused") + ) + assert err["code"] == "download_failed" + assert "refused" in err["details"]["reason"] + assert list((tmp_path / "out").iterdir()) == [] + + def test_midstream_reset_emits_envelope_and_cleans_temp(self, fake_target, tmp_path, capsys): + class _ResettingResp(_FakeResp): + def read(self, n: int) -> bytes: + chunk = super().read(n) + if not chunk: + raise ConnectionResetError("Connection reset by peer") + return chunk + + err = self._failing_download(fake_target, tmp_path, capsys, _ResettingResp(b"x" * 100, content_length=None)) + assert err["code"] == "download_failed" + assert "reset" in err["details"]["reason"] + assert list((tmp_path / "out").iterdir()) == [] + + def test_chunked_truncation_incompleteread_emits_envelope(self, fake_target, tmp_path, capsys): + # A truncated chunked body raises http.client.IncompleteRead — an + # HTTPException, NOT an OSError — which used to escape both except arms + # as a raw traceback and break the machine-mode envelope contract. + import http.client + + class _IncompleteResp(_FakeResp): + def read(self, n: int) -> bytes: + chunk = super().read(n) + if not chunk: + raise http.client.IncompleteRead(b"", 500) + return chunk + + err = self._failing_download(fake_target, tmp_path, capsys, _IncompleteResp(b"y" * 100, content_length=None)) + assert err["code"] == "download_failed" + assert list((tmp_path / "out").iterdir()) == [] + + def test_rename_failure_emits_envelope_and_cleans_temp(self, fake_target, tmp_path, capsys): + with patch("pathlib.Path.replace", side_effect=OSError("Permission denied")): + err = self._failing_download(fake_target, tmp_path, capsys, _FakeResp()) + assert err["code"] == "download_failed" + assert "Permission denied" in err["details"]["reason"] + assert list((tmp_path / "out").iterdir()) == [] + + def test_download_passes_socket_timeout(self, fake_target, tmp_path, capsys): + set_renderer(Renderer(mode=OutputMode.NDJSON, command="download")) + _write_state(fake_target) + opener = MagicMock(side_effect=lambda req, timeout=None: _FakeResp()) + with ( + patch("comfy_cli.command.transfer.resolve_target", return_value=fake_target), + patch.object(transfer._DOWNLOAD_OPENER, "open", opener), + ): + transfer.execute_download(PROMPT_ID, out_dir=str(tmp_path / "out")) + assert opener.call_args.kwargs["timeout"] == transfer._DOWNLOAD_TIMEOUT_S + + +class TestDeclaredContentLength: + """Content-Length parsing must never raise (a ValueError would silently + disable size verification): unparseable or ambiguous headers degrade to + None, a folded duplicate that agrees is accepted.""" + + class _Resp: + def __init__(self, value): + self.headers = {} if value is None else {"Content-Length": value} + + @pytest.mark.parametrize( + "value,expected", + [ + ("1000", 1000), + ("0", 0), + (None, None), + ("1000, 1000", 1000), # folded duplicate that agrees + ("1000, 2000", None), # folded duplicate that disagrees → ambiguous + ("not-a-number", None), + ("-5", None), + ("", None), + ], + ) + def test_parses_or_degrades_to_none(self, value, expected): + assert transfer._declared_content_length(self._Resp(value)) == expected + + +class TestPartFileHelper: + """`_open_part_file` owns the raw fd end-to-end: any failure closes it and + removes the temp file, and a copy whose source vanishes must not leak the + descriptor either (the fd is wrapped before the source is opened).""" + + def _open_fd_count(self): + return len(os.listdir("/proc/self/fd")) + + @pytest.mark.skipif(not hasattr(os, "fchmod"), reason="os.fchmod is POSIX-only") + def test_fchmod_failure_closes_fd_and_removes_temp(self, tmp_path, monkeypatch): + monkeypatch.setattr(os, "fchmod", lambda *a: (_ for _ in ()).throw(OSError("ENOTSUP"))) + with pytest.raises(OSError, match="ENOTSUP"): + transfer._open_part_file(tmp_path / "out.png") + assert list(tmp_path.glob("*.part")) == [] + + @pytest.mark.skipif(not sys.platform.startswith("linux"), reason="/proc/self/fd fd accounting") + def test_missing_source_copy_leaks_no_fd_and_no_temp(self, tmp_path): + dst = tmp_path / "dst.bin" + missing = tmp_path / "gone.bin" + before = self._open_fd_count() + for _ in range(5): + with pytest.raises(FileNotFoundError): + transfer._copy_local_output_capped(missing, dst) + assert self._open_fd_count() == before + assert list(tmp_path.glob("*.part")) == [] + assert not dst.exists() + + class TestLocalOutputCopy: """A `comfy run --where local --wait` job records bare on-disk output paths (execution.format_image_path returns an absolute path for loopback @@ -450,6 +711,37 @@ def test_file_uri_output_is_copied(self, fake_target, tmp_path, capsys): assert Path(paths[0]).suffix == ".mp4" assert Path(paths[0]).read_bytes() == b"fake-mp4" + def test_copy_cap_enforced_during_read(self, tmp_path, monkeypatch): + monkeypatch.setattr(transfer, "_MAX_DOWNLOAD_BYTES", 100) + src = tmp_path / "big.bin" + src.write_bytes(b"x" * 300) + dst = tmp_path / "out" / "big.bin" + dst.parent.mkdir() + with pytest.raises(ValueError, match="safety limit"): + transfer._copy_local_output_capped(src, dst) + assert list(dst.parent.iterdir()) == [] + + def test_copy_cap_breach_surfaces_as_envelope(self, fake_target, tmp_path, capsys, monkeypatch): + import typer + + src = tmp_path / "output" / "grew.png" + src.parent.mkdir() + src.write_bytes(b"\x89PNG-local") + self._write_local_state([str(src)]) + + def _cap_breach(src_, dst_): + raise ValueError("local output exceeds 100 byte safety limit") + + monkeypatch.setattr(transfer, "_copy_local_output_capped", _cap_breach) + set_renderer(Renderer(mode=OutputMode.JSON, command="download")) + with patch("comfy_cli.command.transfer.resolve_target", return_value=fake_target): + with pytest.raises(typer.Exit) as excinfo: + transfer.execute_download(PROMPT_ID, out_dir=str(tmp_path / "out")) + assert excinfo.value.exit_code == 1 + envelope = json.loads([ln for ln in capsys.readouterr().out.splitlines() if ln.strip()][-1]) + assert envelope["error"]["code"] == "download_failed" + assert "safety limit" in envelope["error"]["message"] + def test_missing_local_source_errors_cleanly(self, fake_target, tmp_path, capsys): self._write_local_state([str(tmp_path / "output" / "gone.png")]) set_renderer(Renderer(mode=OutputMode.JSON, command="download"))