fix(http): route the remaining bare urlopen() call sites through an http(s)-only opener (BE-3313)

comfy_cli/command/jobs.py

@@ -36,7 +36,7 @@
 from comfy_cli import cancellation, execution_errors, tracking
 from comfy_cli.env_checker import check_comfy_server_running
 from comfy_cli.host_port import resolve_host_port as _resolve_host_port
-from comfy_cli.http import authed_urlopen
+from comfy_cli.http import authed_urlopen, plain_urlopen
 from comfy_cli.output import get_renderer
 from comfy_cli.where import cloud_preflight_or_exit
 
@@ -84,7 +84,7 @@ def _server_or_error(host: str, port: int, *, raise_on_missing: bool = True) ->
 def _http_get_json(url: str, *, timeout: float = 10.0) -> Any:
     req = urllib.request.Request(url)
     try:
-        with urllib.request.urlopen(req, timeout=timeout) as resp:
+        with plain_urlopen(req, timeout=timeout) as resp:
             return json.loads(resp.read())
     except urllib.error.URLError as e:
         raise RuntimeError(f"failed to GET {url}: {e}") from e
@@ -992,7 +992,7 @@ def _local_cancel(prompt_id: str, host: str, port: int) -> None:
     )
     queue_ok = True
     try:
-        with urllib.request.urlopen(queue_req, timeout=10) as resp:
+        with plain_urlopen(queue_req, timeout=10) as resp:
             _ = resp.read()
     except (urllib.error.HTTPError, urllib.error.URLError, OSError):
         # Server refused the delete; common when the prompt isn't in queue.
@@ -1016,7 +1016,7 @@ def _local_cancel(prompt_id: str, host: str, port: int) -> None:
     if prompt_id in running_ids:
         interrupt_req = urllib.request.Request(f"{base}/interrupt", method="POST")
         try:
-            with urllib.request.urlopen(interrupt_req, timeout=10) as resp:
+            with plain_urlopen(interrupt_req, timeout=10) as resp:
                 _ = resp.read()
         except (urllib.error.HTTPError, urllib.error.URLError, OSError):
             interrupt_ok = False

comfy_cli/command/run/execution.py

@@ -38,6 +38,7 @@
 
 from comfy_cli import execution_errors
 from comfy_cli.command.run.loader import _MAX_BODY_PREVIEW, _node_errors_to_list
+from comfy_cli.http import no_redirect_urlopen
 from comfy_cli.output import get_renderer
 from comfy_cli.output import rprint as pprint
 from comfy_cli.workspace_manager import WorkspaceManager
@@ -167,8 +168,12 @@ def queue(self):
         req = request.Request(f"http://{self.host}:{self.port}/prompt", json.dumps(data).encode("utf-8"))
         req.add_header("Comfy-Usage-Source", "comfy-cli")
         try:
-            resp = request.urlopen(req, timeout=self.timeout)
-            raw_body = resp.read()
+            # No-redirect, not ``plain_urlopen``: ``extra_data`` can carry a
+            # Comfy Org credential, so this submit gets the same refuse-a-30x
+            # policy as every other credentialed call rather than leaning on
+            # urllib happening to drop the body when it follows a redirect.
+            with no_redirect_urlopen(req, timeout=self.timeout) as resp:
+                raw_body = resp.read()
         except urllib.error.HTTPError as e:
             body_bytes = e.read()
             body_text = body_bytes.decode("utf-8", errors="replace").strip() if body_bytes else ""

comfy_cli/command/run/preflight.py

@@ -11,11 +11,11 @@
 
 import json
 import urllib.error
-from urllib import request
 
 import typer
 
 from comfy_cli.command.run.loader import _MAX_BODY_PREVIEW
+from comfy_cli.http import plain_urlopen
 from comfy_cli.output import get_renderer
 from comfy_cli.output import rprint as pprint
 
@@ -24,6 +24,11 @@
 # the authoritative signal is the `api_node: true` flag.
 PARTNER_NODE_CATEGORY_PREFIXES = ("partner/",)
 
+# Cap on what we'll pull off the wire for /object_info, success or error. A
+# real schema dump is a few MiB at most; the bound is there so a wedged or
+# hostile server can't stream us out of memory.
+_MAX_OBJECT_INFO_BYTES = 64 * 1024 * 1024
+
 
 def fetch_object_info(host, port, timeout):
     """GET ``/object_info`` from the running ComfyUI server.
@@ -37,10 +42,10 @@ def fetch_object_info(host, port, timeout):
     renderer = get_renderer()
     url = f"http://{host}:{port}/object_info"
     try:
-        with request.urlopen(url, timeout=timeout) as resp:
-            body = resp.read(64 * 1024 * 1024)
+        with plain_urlopen(url, timeout=timeout) as resp:
+            body = resp.read(_MAX_OBJECT_INFO_BYTES)
     except urllib.error.HTTPError as e:
-        body_text = e.read().decode("utf-8", errors="replace").strip()
+        body_text = e.read(_MAX_OBJECT_INFO_BYTES).decode("utf-8", errors="replace").strip()
         renderer.error(
             code="object_info_unavailable",
             message=f"Failed to fetch /object_info (HTTP {e.code})",

comfy_cli/command/templates.py

@@ -28,6 +28,7 @@
 import typer
 
 from comfy_cli import tracking
+from comfy_cli.http import plain_urlopen
 from comfy_cli.output import get_renderer, rprint
 
 app = typer.Typer(no_args_is_help=True, help="Browse the Comfy workflow-template gallery.")
@@ -48,7 +49,7 @@ def _cache_path() -> Path:
 
 def _fetch_gallery(url: str = GALLERY_URL, timeout: float = 15.0) -> bytes:
     req = urllib.request.Request(url, headers={"User-Agent": "comfy-cli"})
-    with urllib.request.urlopen(req, timeout=timeout) as resp:
+    with plain_urlopen(req, timeout=timeout) as resp:
         if resp.status != 200:
             raise RuntimeError(f"gallery fetch failed: HTTP {resp.status}")
         return resp.read()
@@ -410,7 +411,7 @@ def _fetch_template_workflow(name: str, *, timeout: float = 15.0) -> bytes:
     """Pull a single template's workflow JSON from the canonical GitHub raw URL."""
     url = _TEMPLATE_WORKFLOW_URL.format(name=urllib.parse.quote(name, safe=""))
     req = urllib.request.Request(url, headers={"User-Agent": "comfy-cli"})
-    with urllib.request.urlopen(req, timeout=timeout) as resp:
+    with plain_urlopen(req, timeout=timeout) as resp:
         if resp.status != 200:
             raise RuntimeError(f"template workflow fetch failed: HTTP {resp.status}")
         return resp.read()

comfy_cli/http.py

@@ -82,6 +82,39 @@ def build_http_only_opener(*handlers: urllib.request.BaseHandler) -> urllib.requ
 
 _AUTHED_OPENER = build_http_only_opener(NoRedirectHandler())
 
+# The uncredentialed fetches — the template gallery on raw.githubusercontent.com
+# and the REST calls against a local ``http://{host}:{port}`` ComfyUI server.
+# These reached for ``urllib.request.urlopen()``, i.e. the global default
+# opener, which also speaks ``file://``, ``ftp://`` and ``data:``. Nothing here
+# attaches a credential header, so unlike ``_AUTHED_OPENER`` there is no
+# redirect-replay exposure and ``HTTPRedirectHandler`` is installed explicitly
+# to keep the redirect-following those call sites have always had. What the
+# pinning buys is that a URL which stops being trusted — a gallery URL that
+# becomes configurable, say — still can't be steered into a local-file read.
+_PLAIN_OPENER = build_http_only_opener(urllib.request.HTTPRedirectHandler())
+
+
+def plain_urlopen(url, *, timeout: float = 30.0):
+    """Open an uncredentialed request via the http(s)-only shared opener.
+
+    ``url`` is a full URL or a prepared ``Request``. Redirects are followed, as
+    they were when these call sites used the global default opener.
+    """
+    return _PLAIN_OPENER.open(url, timeout=timeout)
+
+
+def no_redirect_urlopen(url, *, timeout: float = 30.0):
+    """Open a prepared credential-bearing ``Request`` without following redirects.
+
+    ``authed_urlopen`` covers the common case where the credential rides a
+    header we attach ourselves. This is the escape hatch for a request whose
+    credential the caller has already placed somewhere we can't build — the
+    ``/prompt`` submit carries ``api_key_comfy_org`` inside its JSON body — and
+    which therefore wants the same no-redirect policy without the header
+    mechanics.
+    """
+    return _AUTHED_OPENER.open(url, timeout=timeout)
+
 
 def build_authed_request(
     url: str,

tests/comfy_cli/command/test_run.py

@@ -18,6 +18,7 @@
     execute,
     fetch_object_info,
     is_ui_workflow,
+    preflight,
 )
 
 
@@ -116,7 +117,7 @@ class TestFetchObjectInfo:
     def test_returns_parsed_json_on_success(self):
         payload = {"KSampler": {"input": {}, "output_node": False}}
         with patch(
-            "comfy_cli.command.run.request.urlopen",
+            "comfy_cli.http._PLAIN_OPENER.open",
             return_value=_ok_response(json.dumps(payload).encode()),
         ) as mock_open:
             result = fetch_object_info("127.0.0.1", 8188, timeout=30)
@@ -125,7 +126,7 @@ def test_returns_parsed_json_on_success(self):
 
     def test_http_error_exits_cleanly(self):
         with patch(
-            "comfy_cli.command.run.request.urlopen",
+            "comfy_cli.http._PLAIN_OPENER.open",
             side_effect=_make_http_error(500, b"server exploded"),
         ):
             with pytest.raises(typer.Exit) as exc_info:
@@ -134,28 +135,38 @@ def test_http_error_exits_cleanly(self):
 
     def test_network_error_exits_cleanly(self):
         with patch(
-            "comfy_cli.command.run.request.urlopen",
+            "comfy_cli.http._PLAIN_OPENER.open",
             side_effect=urllib.error.URLError("Connection refused"),
         ):
             with pytest.raises(typer.Exit) as exc_info:
                 fetch_object_info("127.0.0.1", 8188, timeout=30)
             assert exc_info.value.exit_code == 1
 
     def test_timeout_exits_cleanly(self):
-        with patch("comfy_cli.command.run.request.urlopen", side_effect=TimeoutError("timed out")):
+        with patch("comfy_cli.http._PLAIN_OPENER.open", side_effect=TimeoutError("timed out")):
             with pytest.raises(typer.Exit) as exc_info:
                 fetch_object_info("127.0.0.1", 8188, timeout=5)
             assert exc_info.value.exit_code == 1
 
     def test_invalid_json_exits_cleanly(self):
         with patch(
-            "comfy_cli.command.run.request.urlopen",
+            "comfy_cli.http._PLAIN_OPENER.open",
             return_value=_ok_response(b"<html>not json</html>"),
         ):
             with pytest.raises(typer.Exit) as exc_info:
                 fetch_object_info("127.0.0.1", 8188, timeout=30)
             assert exc_info.value.exit_code == 1
 
+    def test_error_body_read_is_capped(self):
+        """The success path bounds the read; the error path must too, or a
+        hostile server just has to return a 500 to stream us out of memory."""
+        err = _make_http_error(500, b"boom")
+        with patch.object(err, "read", wraps=err.read) as err_read:
+            with patch("comfy_cli.http._PLAIN_OPENER.open", side_effect=err):
+                with pytest.raises(typer.Exit):
+                    fetch_object_info("127.0.0.1", 8188, timeout=30)
+        assert err_read.call_args.args[0] == preflight._MAX_OBJECT_INFO_BYTES
+
 
 class TestWorkflowExecutionAuth:
     """X-API-Key is the credential the ComfyUI server forwards to Partner Nodes."""
@@ -176,25 +187,25 @@ def _make_exec(self, workflow, api_key=None):
 
     def test_queue_embeds_api_key_in_extra_data(self, workflow):
         ex = self._make_exec(workflow, api_key="sk-secret")
-        with patch("comfy_cli.command.run.request.urlopen") as mock_open:
-            mock_open.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
+        with patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open:
+            mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
             ex.queue()
         req = mock_open.call_args[0][0]
         body = json.loads(req.data)
         assert body["extra_data"] == {"comfy_usage_source": "comfy-cli", "api_key_comfy_org": "sk-secret"}
 
     def test_queue_does_not_send_x_api_key_header(self, workflow):
         ex = self._make_exec(workflow, api_key="sk-secret")
-        with patch("comfy_cli.command.run.request.urlopen") as mock_open:
-            mock_open.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
+        with patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open:
+            mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
             ex.queue()
         req = mock_open.call_args[0][0]
         assert req.get_header("X-api-key") is None
 
     def test_queue_omits_api_key_when_not_set(self, workflow):
         ex = self._make_exec(workflow)
-        with patch("comfy_cli.command.run.request.urlopen") as mock_open:
-            mock_open.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
+        with patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open:
+            mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
             ex.queue()
         req = mock_open.call_args[0][0]
         body = json.loads(req.data)
@@ -206,12 +217,48 @@ def test_queue_omits_api_key_when_not_set(self, workflow):
 
     def test_queue_sends_usage_source_header(self, workflow):
         ex = self._make_exec(workflow)
-        with patch("comfy_cli.command.run.request.urlopen") as mock_open:
-            mock_open.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
+        with patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open:
+            mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
             ex.queue()
         req = mock_open.call_args[0][0]
         assert req.get_header("Comfy-usage-source") == "comfy-cli"
 
+    def test_queue_submits_through_the_no_redirect_opener(self, workflow):
+        """The api_key rides the request body, so the submit must go through the
+        opener that refuses a 30x rather than the redirect-following one."""
+        ex = self._make_exec(workflow, api_key="sk-secret")
+        with patch("comfy_cli.http._PLAIN_OPENER.open") as plain:
+            with patch("comfy_cli.http._AUTHED_OPENER.open") as authed:
+                authed.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
+                ex.queue()
+        assert authed.call_count == 1
+        assert plain.call_count == 0
+
+    def test_queue_surfaces_a_refused_redirect_as_an_error(self, workflow):
+        """A 30x on /prompt is a misconfiguration or an attack, not something to
+        follow with a credential in the body — it exits rather than resubmitting."""
+        ex = self._make_exec(workflow, api_key="sk-secret")
+        redirect = urllib.error.HTTPError(
+            url="http://127.0.0.1:8188/prompt",
+            code=307,
+            msg="redirect refused",
+            hdrs=None,
+            fp=io.BytesIO(b"redirect refused"),
+        )
+        with patch("comfy_cli.http._AUTHED_OPENER.open", side_effect=redirect):
+            with pytest.raises(typer.Exit) as exc_info:
+                ex.queue()
+        assert exc_info.value.exit_code == 1
+
+    def test_queue_closes_the_response(self, workflow):
+        """The submit reads inside a `with`, so the connection doesn't linger
+        until GC while the run moves on to the websocket."""
+        ex = self._make_exec(workflow)
+        with patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open:
+            mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
+            ex.queue()
+        assert mock_open.return_value.__exit__.called
+
 
 class TestWatchExecution:
     def test_successful_execution(self, mock_execution):