From 83dcd5b714373c1f87ca79910ce1359139a6ec4e Mon Sep 17 00:00:00 2001 From: Hunter Senft-Grupp Date: Sat, 8 Aug 2026 17:32:46 +0000 Subject: [PATCH] Fix Comfy Cloud node safety Amp-Thread-ID: https://ampcode.com/threads/T-019fd9e0-653f-74cb-a744-9d35f9264778 Co-authored-by: Amp --- comfy_api_nodes/apis/comfy_cloud.py | 2 +- comfy_api_nodes/nodes_comfy_cloud.py | 85 ++++++++++--------- comfy_api_nodes/util/conversions.py | 12 ++- comfy_api_nodes/util/download_helpers.py | 2 +- .../comfy_api_nodes_test/comfy_cloud_test.py | 61 ++++++++++++- 5 files changed, 117 insertions(+), 45 deletions(-) diff --git a/comfy_api_nodes/apis/comfy_cloud.py b/comfy_api_nodes/apis/comfy_cloud.py index a781d3aef..5449dc5fe 100644 --- a/comfy_api_nodes/apis/comfy_cloud.py +++ b/comfy_api_nodes/apis/comfy_cloud.py @@ -107,7 +107,7 @@ class ComfyCloudGenerateResponse(BaseModel): class ComfyCloudStatusResponse(BaseModel): task_id: str = Field(..., min_length=1) status: str = Field(...) - progress: float | None = Field(None, ge=0, le=100) + progress: float | None = Field(None) output_url: str | None = Field(None) output_urls: dict[str, str] | None = Field(None) error: str | None = Field(None) diff --git a/comfy_api_nodes/nodes_comfy_cloud.py b/comfy_api_nodes/nodes_comfy_cloud.py index 188c1f0cc..cdcd63af9 100644 --- a/comfy_api_nodes/nodes_comfy_cloud.py +++ b/comfy_api_nodes/nodes_comfy_cloud.py @@ -1,7 +1,8 @@ import math +import posixpath import re from typing import ClassVar -from urllib.parse import quote, urlsplit +from urllib.parse import quote, unquote, urlsplit import torch @@ -25,6 +26,7 @@ from comfy_api_nodes.util import ( get_number_of_images, poll_op, sync_op, + sync_op_raw, upload_audio_to_comfyapi, upload_image_to_comfyapi, upload_video_to_comfyapi, @@ -85,8 +87,20 @@ def _with_input_sockets(inputs: list[IO.Input]) -> list[IO.Input]: def _validated_output_url(url: str) -> str: parsed = urlsplit(url) - is_proxy_path = not parsed.scheme and not parsed.netloc and parsed.path.startswith("/proxy/comfy-cloud/") - is_signed_https_url = parsed.scheme == "https" and bool(parsed.netloc) and parsed.username is None + decoded_path = unquote(parsed.path) + is_proxy_path = ( + not parsed.scheme + and not parsed.netloc + and decoded_path.startswith("/proxy/comfy-cloud/") + and posixpath.normpath(decoded_path) == decoded_path + ) + is_signed_https_url = ( + parsed.scheme == "https" + and parsed.hostname == "storage.googleapis.com" + and parsed.port is None + and parsed.username is None + and parsed.password is None + ) if not is_proxy_path and not is_signed_https_url: raise RuntimeError("Comfy Cloud returned an invalid output URL.") return url @@ -110,6 +124,31 @@ def _validate_audio_upload(audio: Input.Audio) -> None: raise ValueError("Decoded audio exceeds the 256 MiB Comfy Cloud limit.") +def _progress(response: ComfyCloudStatusResponse) -> float | None: + if response.progress is None or not math.isfinite(response.progress): + return None + return min(100.0, max(0.0, response.progress)) + + +async def _poll_task(cls: type[IO.ComfyNode], task_id: str) -> ComfyCloudStatusResponse: + polling_endpoint, cancel_endpoint = _task_endpoints(task_id) + try: + return await poll_op( + cls, + polling_endpoint, + response_model=ComfyCloudStatusResponse, + status_extractor=lambda response: response.status, + progress_extractor=_progress, + cancel_endpoint=cancel_endpoint, + ) + except Exception: + try: + await sync_op_raw(cls, cancel_endpoint, max_retries=0) + except Exception: + pass + raise + + def _validate_node_inputs(cls: type[IO.ComfyNode], values: dict) -> dict: validated = dict(values) for input_spec in cls.define_schema().inputs: @@ -220,15 +259,7 @@ class _ComfyCloudWorkflowNode(IO.ComfyNode): inputs=inputs, ), ) - polling_endpoint, cancel_endpoint = _task_endpoints(task.task_id) - result = await poll_op( - cls, - polling_endpoint, - response_model=ComfyCloudStatusResponse, - status_extractor=lambda response: response.status, - progress_extractor=lambda response: response.progress, - cancel_endpoint=cancel_endpoint, - ) + result = await _poll_task(cls, task.task_id) if not result.output_url: raise RuntimeError("Comfy Cloud task completed without an output URL.") @@ -557,15 +588,7 @@ class ComfyCloudSeedVR2ImageUpscaleNode(_ComfyCloudWorkflowNode): async def _run_video_workflow(cls: type[IO.ComfyNode], workflow: ComfyCloudWorkflow, inputs: ComfyCloudWorkflowInputs) -> IO.NodeOutput: task = await sync_op(cls, _GENERATE_ENDPOINT, response_model=ComfyCloudGenerateResponse, data=ComfyCloudGenerateRequest(workflow=workflow, inputs=inputs)) - polling_endpoint, cancel_endpoint = _task_endpoints(task.task_id) - result = await poll_op( - cls, - polling_endpoint, - response_model=ComfyCloudStatusResponse, - status_extractor=lambda response: response.status, - progress_extractor=lambda response: response.progress, - cancel_endpoint=cancel_endpoint, - ) + result = await _poll_task(cls, task.task_id) if not result.output_url: raise RuntimeError("Comfy Cloud task completed without an output URL.") return IO.NodeOutput( @@ -819,15 +842,7 @@ async def _audio_asset(cls: type[IO.ComfyNode], name: str, audio: Input.Audio) - async def _run_audio_workflow(cls: type[IO.ComfyNode], workflow: ComfyCloudWorkflow, inputs: ComfyCloudWorkflowInputs, output_names: tuple[str, ...] = ()) -> IO.NodeOutput: task = await sync_op(cls, _GENERATE_ENDPOINT, response_model=ComfyCloudGenerateResponse, data=ComfyCloudGenerateRequest(workflow=workflow, inputs=inputs)) - polling_endpoint, cancel_endpoint = _task_endpoints(task.task_id) - result = await poll_op( - cls, - polling_endpoint, - response_model=ComfyCloudStatusResponse, - status_extractor=lambda response: response.status, - progress_extractor=lambda response: response.progress, - cancel_endpoint=cancel_endpoint, - ) + result = await _poll_task(cls, task.task_id) if output_names: if not result.output_urls or any(not result.output_urls.get(name) for name in output_names): raise RuntimeError("Comfy Cloud task completed without all named output URLs.") @@ -998,15 +1013,7 @@ class ComfyCloudMelBandRoFormerStemSeparationNode(IO.ComfyNode): async def _run_3d_workflow(cls: type[IO.ComfyNode], workflow: ComfyCloudWorkflow, inputs: ComfyCloudWorkflowInputs, file_format: str) -> IO.NodeOutput: task = await sync_op(cls, _GENERATE_ENDPOINT, response_model=ComfyCloudGenerateResponse, data=ComfyCloudGenerateRequest(workflow=workflow, inputs=inputs)) - polling_endpoint, cancel_endpoint = _task_endpoints(task.task_id) - result = await poll_op( - cls, - polling_endpoint, - response_model=ComfyCloudStatusResponse, - status_extractor=lambda response: response.status, - progress_extractor=lambda response: response.progress, - cancel_endpoint=cancel_endpoint, - ) + result = await _poll_task(cls, task.task_id) if not result.output_url: raise RuntimeError("Comfy Cloud task completed without an output URL.") return IO.NodeOutput( diff --git a/comfy_api_nodes/util/conversions.py b/comfy_api_nodes/util/conversions.py index f46cac3f8..b0b3e5ba2 100644 --- a/comfy_api_nodes/util/conversions.py +++ b/comfy_api_nodes/util/conversions.py @@ -15,6 +15,8 @@ from comfy_api.latest import Input, InputImpl, Types from ._helpers import mimetype_to_extension +_MAX_DECODED_AUDIO_BYTES = 256 * 1024 * 1024 + def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str = "RGBA") -> torch.Tensor: """Converts image data from BytesIO to a torch.Tensor. @@ -568,12 +570,14 @@ def _f32_pcm(wav: torch.Tensor) -> torch.Tensor: raise ValueError(f"Unsupported wav dtype: {wav.dtype}") -def audio_bytes_to_audio_input(audio_bytes: bytes) -> dict: +def audio_bytes_to_audio_input(audio_bytes: bytes | BytesIO) -> dict: """ Decode any common audio container from bytes using PyAV and return a Comfy AUDIO dict: {"waveform": [1, C, T] float32, "sample_rate": int}. """ - with av.open(BytesIO(audio_bytes)) as af: + source = audio_bytes if isinstance(audio_bytes, BytesIO) else BytesIO(audio_bytes) + source.seek(0) + with av.open(source) as af: if not af.streams.audio: raise ValueError("No audio stream found in response.") stream = af.streams.audio[0] @@ -583,6 +587,7 @@ def audio_bytes_to_audio_input(audio_bytes: bytes) -> dict: frames: list[torch.Tensor] = [] n_channels = stream.channels or 1 + decoded_bytes = 0 for frame in af.decode(streams=stream.index): arr = frame.to_ndarray() # shape can be [C, T] or [T, C] or [T] @@ -593,6 +598,9 @@ def audio_bytes_to_audio_input(audio_bytes: bytes) -> dict: buf = buf.transpose(0, 1).contiguous() # [T, C] -> [C, T] elif buf.shape[0] != n_channels: buf = buf.reshape(-1, n_channels).t().contiguous() # fallback to [C, T] + decoded_bytes += buf.numel() * buf.element_size() + if decoded_bytes > _MAX_DECODED_AUDIO_BYTES: + raise ValueError("Decoded audio exceeds the 256 MiB limit.") frames.append(buf) if not frames: diff --git a/comfy_api_nodes/util/download_helpers.py b/comfy_api_nodes/util/download_helpers.py index 3edf66fad..5c3ddd30f 100644 --- a/comfy_api_nodes/util/download_helpers.py +++ b/comfy_api_nodes/util/download_helpers.py @@ -286,7 +286,7 @@ async def download_url_to_audio_input( cls=cls, allow_redirects=allow_redirects, ) - return audio_bytes_to_audio_input(result.getvalue()) + return audio_bytes_to_audio_input(result) async def download_url_as_bytesio( diff --git a/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py b/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py index 4a8efb625..9ccf351b9 100644 --- a/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py +++ b/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py @@ -20,7 +20,7 @@ from comfy_api_nodes.apis.comfy_cloud import ( ComfyCloudWorkflowInputs, ) from comfy_api_nodes import nodes_comfy_cloud -from comfy_api_nodes.util import download_helpers +from comfy_api_nodes.util import conversions, download_helpers @pytest.mark.parametrize( @@ -111,6 +111,25 @@ def test_contract_omits_optional_status_fields(): assert status.model_dump(exclude_none=True) == {"task_id": "task-1", "status": "queued"} +def test_status_progress_is_clamped_for_display(): + assert nodes_comfy_cloud._progress(ComfyCloudStatusResponse(task_id="task-1", status="running", progress=100.5)) == 100 + assert nodes_comfy_cloud._progress(ComfyCloudStatusResponse(task_id="task-1", status="running", progress=-1)) == 0 + + +def test_poll_failure_cancels_submitted_task(monkeypatch): + poll = AsyncMock(side_effect=ValueError("invalid status response")) + cancel = AsyncMock(return_value={"status": "cancellation_requested"}) + monkeypatch.setattr(nodes_comfy_cloud, "poll_op", poll) + monkeypatch.setattr(nodes_comfy_cloud, "sync_op_raw", cancel) + + with pytest.raises(ValueError, match="invalid status response"): + asyncio.run(nodes_comfy_cloud._poll_task(nodes_comfy_cloud.ComfyCloudTextToImageNode, "task/1")) + + cancel.assert_awaited_once() + assert cancel.call_args.args[1].path == "/proxy/comfy-cloud/workflow/tasks/task%2F1/cancel" + assert cancel.call_args.kwargs["max_retries"] == 0 + + @pytest.mark.parametrize("response_model", [ComfyCloudGenerateResponse, ComfyCloudStatusResponse]) @pytest.mark.parametrize("task_id", ["", " "]) def test_contract_rejects_empty_task_ids(response_model, task_id): @@ -130,6 +149,11 @@ def test_contract_rejects_empty_task_ids(response_model, task_id): "//169.254.169.254/latest/meta-data", "/unrelated/path/output.png", "https://user@example.com/output.png", + "/proxy/comfy-cloud/../../v1/users/me", + "/proxy/comfy-cloud/%2e%2e/%2e%2e/v1/users/me", + "https://127.0.0.1/output.png", + "https://169.254.169.254/latest/meta-data", + "https://attacker.example/output.png", ], ) def test_cloud_workflows_reject_untrusted_output_urls(monkeypatch, url): @@ -738,9 +762,42 @@ def test_download_cloud_audio_url_to_audio_input(monkeypatch): assert isinstance(download_call.call_args.kwargs["dest"], BytesIO) assert download_call.call_args.kwargs["timeout"] == 30 assert download_call.call_args.kwargs["max_retries"] == 2 + audio_decode.assert_called_once() + assert isinstance(audio_decode.call_args.args[0], BytesIO) + assert audio_decode.call_args.kwargs == {} assert download_call.call_args.kwargs["cls"] is node assert download_call.call_args.kwargs["allow_redirects"] is True - audio_decode.assert_called_once_with(downloaded) + + +def test_audio_decode_stops_before_exceeding_budget(monkeypatch): + class Frame: + def to_ndarray(self): + return torch.ones(2, 8).numpy() + + stream = type( + "Stream", + (), + {"codec_context": type("Codec", (), {"sample_rate": 48000})(), "channels": 2, "index": 0}, + )() + + class AudioFile: + streams = type("Streams", (), {"audio": [stream]})() + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def decode(self, streams): + yield Frame() + yield Frame() + + monkeypatch.setattr(conversions.av, "open", lambda source: AudioFile()) + monkeypatch.setattr(conversions, "_MAX_DECODED_AUDIO_BYTES", 64) + + with pytest.raises(ValueError, match="Decoded audio exceeds"): + conversions.audio_bytes_to_audio_input(BytesIO(b"encoded")) AUDIO_POC_NODES = [