diff --git a/comfy_cli/command/jobs.py b/comfy_cli/command/jobs.py index 9d6db398..9de371de 100644 --- a/comfy_cli/command/jobs.py +++ b/comfy_cli/command/jobs.py @@ -36,6 +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.output import get_renderer from comfy_cli.where import cloud_preflight_or_exit @@ -1064,14 +1065,9 @@ def _cloud_cancel(prompt_id: str) -> None: # escape (e.g. ``../foo`` → ``%2E%2E%2Ffoo``). Cloud rejects bad UUIDs # upstream too; encoding here is defense in depth. url = target.url("jobs", urllib.parse.quote(prompt_id, safe=""), "cancel") - req = urllib.request.Request(url, data=b"", method="POST") - if target.api_key: - req.add_header("X-API-Key", target.api_key) - elif target.auth_token: - req.add_header("Authorization", f"Bearer {target.auth_token}") try: - with urllib.request.urlopen(req, timeout=15) as resp: + with authed_urlopen(url, target, method="POST", data=b"", timeout=15) as resp: body = resp.read() except urllib.error.HTTPError as e: body_text = (e.read() or b"")[:1000].decode("utf-8", "replace") diff --git a/comfy_cli/command/models/search.py b/comfy_cli/command/models/search.py index 3f8565e9..3ea8d4c7 100644 --- a/comfy_cli/command/models/search.py +++ b/comfy_cli/command/models/search.py @@ -28,12 +28,12 @@ import re import urllib.error import urllib.parse -import urllib.request from typing import Annotated, Any, NoReturn import typer from comfy_cli import tracking +from comfy_cli.http import authed_urlopen from comfy_cli.output import get_renderer, rprint app = typer.Typer(no_args_is_help=True, help="Discover models — folders, files, and the cloud asset catalog.") @@ -98,23 +98,13 @@ def _models_path_parts(target) -> tuple[str, ...]: return ("experiment", "models") if target.is_cloud else ("models",) -def _authed_request(url: str, target) -> urllib.request.Request: - req = urllib.request.Request(url) - if target.api_key: - req.add_header("X-API-Key", target.api_key) - elif target.auth_token: - req.add_header("Authorization", f"Bearer {target.auth_token}") - return req - - def _http_get_json(url: str, target, timeout: float = 30.0) -> Any: """Issue an authenticated GET and decode JSON. Raises urllib/JSON errors verbatim. Response body is capped at ``_MAX_RESPONSE_BYTES`` to bound memory use on a misbehaving server. A ``ValueError`` is raised if the cap is exceeded. """ - req = _authed_request(url, target) - with urllib.request.urlopen(req, timeout=timeout) as resp: + with authed_urlopen(url, target, timeout=timeout) as resp: # ``read(N)`` returns up to N bytes; reading N+1 lets us distinguish # "fits exactly" from "exceeds cap" without buffering the whole stream # twice on the happy path. diff --git a/comfy_cli/command/workflow.py b/comfy_cli/command/workflow.py index d953c10e..1c90660e 100644 --- a/comfy_cli/command/workflow.py +++ b/comfy_cli/command/workflow.py @@ -500,13 +500,12 @@ def _userdata_request( """Authed HTTP call to a ComfyUI ``/userdata`` endpoint returning (status, raw_bytes). Raises urllib errors verbatim so callers can map them to envelope codes. - Local ComfyUI needs no auth; ``_authed_request`` is a no-op on the headers - when the Target carries no credential. + Local ComfyUI needs no auth; ``authed_urlopen`` attaches no header when the + Target carries no credential. """ - import urllib.request + from comfy_cli.http import authed_urlopen - req = _authed_request(url, target, method=method, data=data, content_type=content_type) - with urllib.request.urlopen(req, timeout=timeout) as resp: + with authed_urlopen(url, target, method=method, data=data, content_type=content_type, timeout=timeout) as resp: status = resp.status # Read one byte past the cap so we can tell a full body from a truncated one. raw = resp.read(_USERDATA_MAX_BYTES + 1) @@ -585,34 +584,16 @@ def _userdata_file_url(target, key: str, query: dict | None = None) -> str: return url -def _authed_request( - url: str, target, *, method: str = "GET", data: bytes | None = None, content_type: str | None = None -): - """Build an authenticated urllib Request. The return type is annotated - loosely to keep urllib out of the module's top-level imports.""" - import urllib.request - - req = urllib.request.Request(url, data=data, method=method) - if target.api_key: - req.add_header("X-API-Key", target.api_key) - elif target.auth_token: - req.add_header("Authorization", f"Bearer {target.auth_token}") - if content_type: - req.add_header("Content-Type", content_type) - return req - - def _http_request( url: str, target, *, method: str = "GET", body: dict | None = None, timeout: float = 30.0 ) -> tuple[int, dict | None]: """Authed HTTP call returning (status, parsed_json_or_none). Raises urllib errors verbatim so callers can surface the right error code.""" - import urllib.request + from comfy_cli.http import authed_urlopen data = json.dumps(body).encode("utf-8") if body is not None else None ct = "application/json" if data is not None else None - req = _authed_request(url, target, method=method, data=data, content_type=ct) - with urllib.request.urlopen(req, timeout=timeout) as resp: + with authed_urlopen(url, target, method=method, data=data, content_type=ct, timeout=timeout) as resp: status = resp.status raw = resp.read(64 * 1024 * 1024) # 64 MiB cap if not raw: diff --git a/comfy_cli/http.py b/comfy_cli/http.py index 793307b9..cb4c559b 100644 --- a/comfy_cli/http.py +++ b/comfy_cli/http.py @@ -24,3 +24,69 @@ def http_error_301(self, req, fp, code, msg, headers): raise urllib.error.HTTPError(req.full_url, code, self._message, headers, fp) http_error_302 = http_error_303 = http_error_307 = http_error_308 = http_error_301 + + +def _build_authed_opener() -> urllib.request.OpenerDirector: + """Build the credential-carrying opener with http/https handlers only. + + ``build_opener()`` would also install ``FileHandler``/``FTPHandler``. Every + call site builds its URL from a trusted ``target.base_url``, so that isn't + reachable today, but this is the opener that attaches credentials — pinning + it to http(s) means a future caller can't be steered into a ``file://`` or + ``ftp://`` fetch. Unknown schemes fall to ``UnknownHandler``, which raises + ``URLError("unknown url type")``. + """ + opener = urllib.request.OpenerDirector() + for handler in ( + urllib.request.ProxyHandler(), + urllib.request.HTTPHandler(), + urllib.request.HTTPSHandler(), + urllib.request.HTTPDefaultErrorHandler(), + urllib.request.HTTPErrorProcessor(), + NoRedirectHandler(), + urllib.request.UnknownHandler(), + ): + opener.add_handler(handler) + return opener + + +_AUTHED_OPENER = _build_authed_opener() + + +def build_authed_request( + url: str, + target, + *, + method: str = "GET", + data: bytes | None = None, + content_type: str | None = None, +) -> urllib.request.Request: + """Build a urllib Request carrying the target's credential header. + + api_key wins over auth_token; the policy layer (resolve_target) populates + at most one, so this is a mechanic, not a policy choice. No header is + attached for an uncredentialed (local) target. + """ + req = urllib.request.Request(url, data=data, method=method) + if target.api_key: + req.add_header("X-API-Key", target.api_key) + elif target.auth_token: + req.add_header("Authorization", f"Bearer {target.auth_token}") + if content_type: + req.add_header("Content-Type", content_type) + return req + + +def authed_urlopen( + url: str, + target, + *, + method: str = "GET", + data: bytes | None = None, + content_type: str | None = None, + timeout: float = 30.0, +): + """Open an authed request via the no-redirect opener. A 30x raises + HTTPError instead of replaying credentials at the redirect target.""" + req = build_authed_request(url, target, method=method, data=data, content_type=content_type) + return _AUTHED_OPENER.open(req, timeout=timeout) diff --git a/tests/comfy_cli/command/models/test_search.py b/tests/comfy_cli/command/models/test_search.py index 8316896b..bc5c3b13 100644 --- a/tests/comfy_cli/command/models/test_search.py +++ b/tests/comfy_cli/command/models/test_search.py @@ -190,7 +190,12 @@ def _fake(req, timeout=None): return _fake_resp(body) raise AssertionError(f"unexpected URL hit by mock: {url}") - monkeypatch.setattr("urllib.request.urlopen", _fake) + # ``_http_get_json`` now opens through the shared no-redirect opener in + # ``comfy_cli.http``; the fake still receives a ``Request`` object, so the + # route-matching above is unchanged. + import comfy_cli.http as http_mod + + monkeypatch.setattr(http_mod._AUTHED_OPENER, "open", _fake) return calls diff --git a/tests/comfy_cli/command/test_workflow_saved.py b/tests/comfy_cli/command/test_workflow_saved.py index 2328480b..3010639b 100644 --- a/tests/comfy_cli/command/test_workflow_saved.py +++ b/tests/comfy_cli/command/test_workflow_saved.py @@ -125,7 +125,12 @@ def _fake(req, timeout=None): return _fake_resp(json.dumps(payload).encode()) raise AssertionError(f"unexpected URL: {url}") - monkeypatch.setattr("urllib.request.urlopen", _fake) + # Both the local ``/userdata`` and cloud workflow paths now open through the + # shared no-redirect opener in ``comfy_cli.http``; the fake still receives a + # ``Request`` object, so the route-matching above is unchanged. + import comfy_cli.http as http_mod + + monkeypatch.setattr(http_mod._AUTHED_OPENER, "open", _fake) return calls diff --git a/tests/comfy_cli/jobs/test_jobs.py b/tests/comfy_cli/jobs/test_jobs.py index 5ed15898..48138336 100644 --- a/tests/comfy_cli/jobs/test_jobs.py +++ b/tests/comfy_cli/jobs/test_jobs.py @@ -363,7 +363,14 @@ def _fake(req, timeout=None): return _Resp(payload if isinstance(payload, bytes) else json.dumps(payload).encode()) raise AssertionError(f"unexpected URL: {url}") + # Local queue/interrupt paths still use plain ``urllib.request.urlopen``; + # the cloud cancel path now opens through the shared no-redirect opener in + # ``comfy_cli.http``. Both receive a ``Request`` object, so route the same + # fake through both. + import comfy_cli.http as http_mod + monkeypatch.setattr("urllib.request.urlopen", _fake) + monkeypatch.setattr(http_mod._AUTHED_OPENER, "open", _fake) return calls diff --git a/tests/comfy_cli/test_http.py b/tests/comfy_cli/test_http.py index 7bac0220..97109ff5 100644 --- a/tests/comfy_cli/test_http.py +++ b/tests/comfy_cli/test_http.py @@ -1,10 +1,19 @@ import http.client +import types import urllib.error import urllib.request +from unittest.mock import patch import pytest -from comfy_cli.http import NoRedirectHandler +import comfy_cli.http as http_mod +from comfy_cli.http import NoRedirectHandler, authed_urlopen, build_authed_request + + +def _target(*, api_key=None, auth_token=None): + """Minimal stand-in for a resolved Target — build_authed_request only reads + ``.api_key`` and ``.auth_token``.""" + return types.SimpleNamespace(api_key=api_key, auth_token=auth_token) def _call(handler, method_name, code=302): @@ -42,3 +51,104 @@ def test_custom_message_passthrough(): _call(handler, "http_error_302", code=302) assert str(exc_info.value.reason) == "redirect refused (auth leak prevention)" assert exc_info.value.code == 302 + + +# --------------------------------------------------------------------------- +# build_authed_request +# --------------------------------------------------------------------------- + + +def test_api_key_wins_over_auth_token(): + """When both credentials are set, X-API-Key is attached and Authorization is not.""" + req = build_authed_request("https://x/thing", _target(api_key="k", auth_token="t")) + assert req.get_header("X-api-key") == "k" + assert req.get_header("Authorization") is None + + +def test_bearer_when_only_auth_token(): + req = build_authed_request("https://x/thing", _target(auth_token="t")) + assert req.get_header("Authorization") == "Bearer t" + assert req.get_header("X-api-key") is None + + +def test_no_auth_header_when_uncredentialed(): + """A local (uncredentialed) target gets no credential header.""" + req = build_authed_request("https://x/thing", _target()) + assert req.get_header("Authorization") is None + assert req.get_header("X-api-key") is None + + +def test_content_type_passthrough(): + req = build_authed_request("https://x/thing", _target(), content_type="application/json") + assert req.get_header("Content-type") == "application/json" + + +def test_no_content_type_header_by_default(): + req = build_authed_request("https://x/thing", _target()) + assert req.get_header("Content-type") is None + + +def test_method_and_data_passthrough(): + req = build_authed_request("https://x/thing", _target(), method="POST", data=b"payload") + assert req.get_method() == "POST" + assert req.data == b"payload" + + +# --------------------------------------------------------------------------- +# authed_urlopen +# --------------------------------------------------------------------------- + + +def test_authed_urlopen_uses_no_redirect_opener(): + """authed_urlopen opens through the shared _AUTHED_OPENER, passing the + built Request and timeout through.""" + sentinel = object() + with patch.object(http_mod._AUTHED_OPENER, "open", return_value=sentinel) as opened: + result = authed_urlopen("https://x/thing", _target(api_key="k"), method="POST", data=b"d", timeout=15) + assert result is sentinel + req = opened.call_args.args[0] + assert isinstance(req, urllib.request.Request) + assert req.full_url == "https://x/thing" + assert req.get_header("X-api-key") == "k" + assert req.get_method() == "POST" + assert opened.call_args.kwargs["timeout"] == 15 + + +def test_authed_urlopen_propagates_refused_redirect(): + """A 30x refused by the opener surfaces as HTTPError to the caller instead + of following the redirect with credentials attached.""" + err = urllib.error.HTTPError("https://x/thing", 302, "redirect refused", http.client.HTTPMessage(), None) + with patch.object(http_mod._AUTHED_OPENER, "open", side_effect=err): + with pytest.raises(urllib.error.HTTPError) as exc_info: + authed_urlopen("https://x/thing", _target(auth_token="t")) + assert exc_info.value.code == 302 + + +def test_migrated_caller_propagates_refused_redirect(): + """End-to-end: a migrated caller (models search's _http_get_json) surfaces a + refused-redirect HTTPError rather than swallowing or following it.""" + from comfy_cli.command.models.search import _http_get_json + + err = urllib.error.HTTPError("https://x/api/assets", 302, "redirect refused", http.client.HTTPMessage(), None) + with patch.object(http_mod._AUTHED_OPENER, "open", side_effect=err): + with pytest.raises(urllib.error.HTTPError) as exc_info: + _http_get_json("https://x/api/assets", _target(api_key="k")) + assert exc_info.value.code == 302 + + +@pytest.mark.parametrize("url", ["file:///etc/passwd", "ftp://example.com/x", "data:text/plain,hi"]) +def test_authed_opener_refuses_non_http_schemes(url, tmp_path): + """The credential-carrying opener is pinned to http(s). ``build_opener`` + would install File/FTP/Data handlers, so a caller steered into a local-file + or ftp URL could exfiltrate or read unintended content; those schemes must + fall through to UnknownHandler instead.""" + with pytest.raises(urllib.error.URLError) as exc_info: + authed_urlopen(url, _target(api_key="k")) + assert "unknown url type" in str(exc_info.value.reason) + + +def test_authed_opener_handler_set(): + """Only http(s)-relevant handlers are installed — no File/FTP/Data.""" + names = {type(h).__name__ for h in http_mod._AUTHED_OPENER.handlers} + assert {"HTTPHandler", "HTTPSHandler", "NoRedirectHandler"} <= names + assert not names & {"FileHandler", "FTPHandler", "DataHandler"}