diff --git a/comfy_cli/command/jobs.py b/comfy_cli/command/jobs.py index 9d6db398..df39eee8 100644 --- a/comfy_cli/command/jobs.py +++ b/comfy_cli/command/jobs.py @@ -1065,10 +1065,10 @@ def _cloud_cancel(prompt_id: str) -> None: # 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}") + from comfy_cli.http import target_auth_headers + + for k, v in target_auth_headers(target).items(): + req.add_header(k, v) try: with urllib.request.urlopen(req, timeout=15) as resp: diff --git a/comfy_cli/command/models/search.py b/comfy_cli/command/models/search.py index 3f8565e9..04e5cd30 100644 --- a/comfy_cli/command/models/search.py +++ b/comfy_cli/command/models/search.py @@ -99,11 +99,11 @@ def _models_path_parts(target) -> tuple[str, ...]: def _authed_request(url: str, target) -> urllib.request.Request: + from comfy_cli.http import target_auth_headers + 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}") + for k, v in target_auth_headers(target).items(): + req.add_header(k, v) return req diff --git a/comfy_cli/command/transfer.py b/comfy_cli/command/transfer.py index 3b3420f3..7899c3f2 100644 --- a/comfy_cli/command/transfer.py +++ b/comfy_cli/command/transfer.py @@ -28,6 +28,7 @@ from comfy_cli import jobs_state from comfy_cli.comfy_client import Client, Unauthenticated, extract_output_entries from comfy_cli.http import NoRedirectHandler +from comfy_cli.http import target_auth_headers as _auth_headers from comfy_cli.output import get_renderer from comfy_cli.output import rprint as pprint from comfy_cli.target import resolve_target @@ -67,17 +68,6 @@ def _default_out_dir() -> str: # --------------------------------------------------------------------------- -def _auth_headers(target: Any) -> dict[str, str]: - """Build auth headers for a target (cloud only).""" - headers: dict[str, str] = {} - if target.is_cloud: - if target.api_key: - headers["X-API-Key"] = target.api_key - elif target.auth_token: - headers["Authorization"] = f"Bearer {target.auth_token}" - return headers - - # Stripped on every download redirect so auth never crosses origins. _AUTH_HEADERS_TO_STRIP = frozenset({"authorization", "x-api-key", "x-comfy-api-key", "cookie"}) _MAX_REDIRECTS = 5 diff --git a/comfy_cli/command/workflow.py b/comfy_cli/command/workflow.py index 383651d7..6289a926 100644 --- a/comfy_cli/command/workflow.py +++ b/comfy_cli/command/workflow.py @@ -607,11 +607,11 @@ def _authed_request( loosely to keep urllib out of the module's top-level imports.""" import urllib.request + from comfy_cli.http import target_auth_headers + 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}") + for k, v in target_auth_headers(target).items(): + req.add_header(k, v) if content_type: req.add_header("Content-Type", content_type) return req diff --git a/comfy_cli/cql/engine.py b/comfy_cli/cql/engine.py index 04439c21..f809b9b9 100644 --- a/comfy_cli/cql/engine.py +++ b/comfy_cli/cql/engine.py @@ -1249,11 +1249,10 @@ def _load_from_target(*, mode: str = "local", host: str | None = None, port: int req.add_header("Accept", "application/json") # Auth headers (cloud only — local has no auth) - if target.is_cloud: - 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}") + from comfy_cli.http import target_auth_headers + + for k, v in target_auth_headers(target).items(): + req.add_header(k, v) try: with _opener.open(req, timeout=30.0) as resp: diff --git a/comfy_cli/http.py b/comfy_cli/http.py index 793307b9..0f0f4a21 100644 --- a/comfy_cli/http.py +++ b/comfy_cli/http.py @@ -24,3 +24,21 @@ 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 target_auth_headers(target) -> dict[str, str]: + """Auth headers for a routing Target — cloud only. + + Local ComfyUI has no auth; refusing to attach credentials to a non-cloud + target means a stray token on a local Target can never leak to a + plaintext server (same defense-in-depth gate as comfy_client.Client). + ``resolve_target`` populates at most one of api_key / auth_token, so the + precedence branch is mechanics, not policy. + """ + headers: dict[str, str] = {} + if target.is_cloud: + if target.api_key: + headers["X-API-Key"] = target.api_key + elif target.auth_token: + headers["Authorization"] = f"Bearer {target.auth_token}" + return headers diff --git a/tests/comfy_cli/test_http.py b/tests/comfy_cli/test_http.py index 7bac0220..43fa04ad 100644 --- a/tests/comfy_cli/test_http.py +++ b/tests/comfy_cli/test_http.py @@ -4,7 +4,8 @@ import pytest -from comfy_cli.http import NoRedirectHandler +from comfy_cli.http import NoRedirectHandler, target_auth_headers +from comfy_cli.target import Target def _call(handler, method_name, code=302): @@ -42,3 +43,31 @@ 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 + + +def test_target_auth_headers_local_attaches_nothing_even_with_creds(): + """The security property this builder exists to enforce: a local Target + NEVER contributes auth headers, even if stray credentials are set on it + (the exact misuse the ``is_cloud`` gate defends against).""" + target = Target( + kind="local", + base_url="http://127.0.0.1:8188", + auth_token="stray", + api_key="stray", + ) + assert target_auth_headers(target) == {} + + +def test_target_auth_headers_cloud_api_key_only(): + target = Target(kind="cloud", base_url="https://cloud.example", api_key="k") + assert target_auth_headers(target) == {"X-API-Key": "k"} + + +def test_target_auth_headers_cloud_auth_token_only(): + target = Target(kind="cloud", base_url="https://cloud.example", auth_token="t") + assert target_auth_headers(target) == {"Authorization": "Bearer t"} + + +def test_target_auth_headers_cloud_both_api_key_wins(): + target = Target(kind="cloud", base_url="https://cloud.example", auth_token="t", api_key="k") + assert target_auth_headers(target) == {"X-API-Key": "k"}