Fix Comfy Cloud node safety

Amp-Thread-ID: https://ampcode.com/threads/T-019fd9e0-653f-74cb-a744-9d35f9264778
Co-authored-by: Amp <amp@ampcode.com>

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)

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(

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:

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(

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 = [