diff --git a/comfy_cli/cloud/oauth.py b/comfy_cli/cloud/oauth.py index f6459fa8..fbde81d0 100644 --- a/comfy_cli/cloud/oauth.py +++ b/comfy_cli/cloud/oauth.py @@ -52,7 +52,7 @@ CLIENT_NAME, get_base_url, ) -from comfy_cli.http import NoRedirectHandler +from comfy_cli.http import NoRedirectHandler, build_http_only_opener # --------------------------------------------------------------------------- # Error types — caller maps these to renderer.error(code=...) codes. @@ -882,7 +882,7 @@ def _assert_https_or_loopback(url: str) -> None: raise _HTTPFail(0, f"refusing plaintext HTTP for OAuth endpoint: {url}") -_OAUTH_OPENER = urllib.request.build_opener(NoRedirectHandler()) +_OAUTH_OPENER = build_http_only_opener(NoRedirectHandler()) def _send_and_parse(req: urllib.request.Request) -> dict: diff --git a/comfy_cli/comfy_client.py b/comfy_cli/comfy_client.py index 02b8c8eb..78cfdb07 100644 --- a/comfy_cli/comfy_client.py +++ b/comfy_cli/comfy_client.py @@ -23,7 +23,7 @@ from dataclasses import dataclass from typing import Any -from comfy_cli.http import NoRedirectHandler +from comfy_cli.http import NoRedirectHandler, build_http_only_opener from comfy_cli.target import Target _LOOPBACK_HOSTS = {"localhost", "127.0.0.1", "::1", "[::1]"} @@ -95,7 +95,7 @@ class Unauthenticated(Exception): """Target needs auth but no valid session is present.""" -_OPENER = urllib.request.build_opener(NoRedirectHandler()) +_OPENER = build_http_only_opener(NoRedirectHandler()) def _assert_safe_url(url: str) -> None: diff --git a/comfy_cli/command/jobs.py b/comfy_cli/command/jobs.py index 9de371de..7fb708f8 100644 --- a/comfy_cli/command/jobs.py +++ b/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 diff --git a/comfy_cli/command/run/execution.py b/comfy_cli/command/run/execution.py index 88a987a0..e7cbc695 100644 --- a/comfy_cli/command/run/execution.py +++ b/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 "" diff --git a/comfy_cli/command/run/preflight.py b/comfy_cli/command/run/preflight.py index 17923c09..5ae7cc49 100644 --- a/comfy_cli/command/run/preflight.py +++ b/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})", diff --git a/comfy_cli/command/templates.py b/comfy_cli/command/templates.py index 523a83bd..4b4badb6 100644 --- a/comfy_cli/command/templates.py +++ b/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() diff --git a/comfy_cli/command/transfer.py b/comfy_cli/command/transfer.py index 3b3420f3..1c49418b 100644 --- a/comfy_cli/command/transfer.py +++ b/comfy_cli/command/transfer.py @@ -27,7 +27,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 NoRedirectHandler, build_http_only_opener from comfy_cli.output import get_renderer from comfy_cli.output import rprint as pprint from comfy_cli.target import resolve_target @@ -105,8 +105,8 @@ def redirect_request(self, req, fp, code, msg, headers, newurl): return new_req -_TRANSFER_OPENER = urllib.request.build_opener(NoRedirectHandler("redirect refused (auth leak prevention)")) -_DOWNLOAD_OPENER = urllib.request.build_opener(_DownloadRedirectHandler()) +_TRANSFER_OPENER = build_http_only_opener(NoRedirectHandler("redirect refused (auth leak prevention)")) +_DOWNLOAD_OPENER = build_http_only_opener(_DownloadRedirectHandler()) # Per-output safety cap, shared by the HTTP download stream and local-output copies. _MAX_DOWNLOAD_BYTES = 10 * 1024 * 1024 * 1024 # 10 GB diff --git a/comfy_cli/cql/engine.py b/comfy_cli/cql/engine.py index 04439c21..09643e13 100644 --- a/comfy_cli/cql/engine.py +++ b/comfy_cli/cql/engine.py @@ -22,7 +22,7 @@ from typing import Any from comfy_cli.cql._net import is_loopback_host -from comfy_cli.http import NoRedirectHandler +from comfy_cli.http import NoRedirectHandler, build_http_only_opener # --------------------------------------------------------------------------- # Types — mirrors nodegraph/types.go @@ -1194,7 +1194,7 @@ def _check_autogrow_required( _MAX_OBJECT_INFO_BYTES = 64 * 1024 * 1024 -_opener = urllib.request.build_opener(NoRedirectHandler()) +_opener = build_http_only_opener(NoRedirectHandler()) class LoadError(Exception): diff --git a/comfy_cli/cql/loader.py b/comfy_cli/cql/loader.py index f432882e..7ed7bac0 100644 --- a/comfy_cli/cql/loader.py +++ b/comfy_cli/cql/loader.py @@ -30,7 +30,7 @@ from comfy_cli.cql._net import is_loopback_host from comfy_cli.cql.errors import CQLRuntimeError -from comfy_cli.http import NoRedirectHandler +from comfy_cli.http import NoRedirectHandler, build_http_only_opener # Cap raw bytes read from disk or the network. Real `object_info` dumps are a # few MB; anything past 256 MiB is almost certainly a wrong path or a hostile @@ -38,7 +38,7 @@ MAX_INPUT_BYTES = 256 * 1024 * 1024 -_LOADER_OPENER = urllib.request.build_opener(NoRedirectHandler()) +_LOADER_OPENER = build_http_only_opener(NoRedirectHandler()) def load_graph( diff --git a/comfy_cli/http.py b/comfy_cli/http.py index cb4c559b..ec3b6206 100644 --- a/comfy_cli/http.py +++ b/comfy_cli/http.py @@ -26,31 +26,94 @@ def http_error_301(self, req, fp, code, msg, headers): 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 +def _http_only_proxy_handler() -> urllib.request.ProxyHandler: + """A ProxyHandler that can only proxy http(s). + + ``ProxyHandler()`` defaults to ``getproxies()`` and registers a + ``_open`` method for *every* entry it finds, so an ``ftp_proxy`` in + the environment would give the opener an ``ftp_open`` — servicing + ``ftp://`` through the proxy and stepping straight past the + ``UnknownHandler`` that ``build_http_only_opener`` relies on. Filtering the + map to http(s) keeps real proxy support (``proxy_bypass``/``no_proxy`` read + the environment directly, not this dict) while leaving non-http schemes + with nowhere to go. + """ + proxies = {scheme: url for scheme, url in urllib.request.getproxies().items() if scheme in ("http", "https")} + return urllib.request.ProxyHandler(proxies) + + +def build_http_only_opener(*handlers: urllib.request.BaseHandler) -> urllib.request.OpenerDirector: + """Build an opener that speaks http(s) and nothing else. + + ``build_opener()`` would also install ``FileHandler``/``FTPHandler``/ + ``DataHandler``. Our call sites build their URLs from a trusted + ``target.base_url``, so that isn't reachable today, but these openers + attach credentials — pinning them to http(s) means a future caller can't + be steered into a ``file://``, ``ftp://`` or ``data:`` fetch. Unknown + schemes fall to ``UnknownHandler``, which raises ``URLError("unknown url type")``. + + ``handlers`` are the caller's own additions (e.g. a redirect policy). As in + ``build_opener``, a caller-supplied handler replaces the default it + subclasses rather than being appended behind it. Note that no redirect + handler is installed unless the caller passes one, so a bare opener + surfaces a 30x as an ``HTTPError`` rather than following it. """ + defaults = [ + (urllib.request.ProxyHandler, _http_only_proxy_handler), + (urllib.request.HTTPHandler, urllib.request.HTTPHandler), + (urllib.request.HTTPDefaultErrorHandler, urllib.request.HTTPDefaultErrorHandler), + (urllib.request.HTTPErrorProcessor, urllib.request.HTTPErrorProcessor), + (urllib.request.UnknownHandler, urllib.request.UnknownHandler), + ] + # urllib.request only defines HTTPSHandler on an SSL-capable build; naming it + # unconditionally would blow up at import time on one without. + if hasattr(urllib.request, "HTTPSHandler"): + defaults.append((urllib.request.HTTPSHandler, urllib.request.HTTPSHandler)) + 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(), - ): + for klass, factory in defaults: + if not any(isinstance(handler, klass) for handler in handlers): + opener.add_handler(factory()) + for handler in handlers: opener.add_handler(handler) return opener -_AUTHED_OPENER = _build_authed_opener() +_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( diff --git a/tests/comfy_cli/command/github/test_pr.py b/tests/comfy_cli/command/github/test_pr.py index 1050be5e..7d2f08bc 100644 --- a/tests/comfy_cli/command/github/test_pr.py +++ b/tests/comfy_cli/command/github/test_pr.py @@ -179,7 +179,7 @@ def test_find_pr_by_branch_error(self, mock_get): @patch("requests.get") def test_find_pr_by_branch_rate_limit(self, mock_get): - """A rate-limited 403 surfaces as GitHubRateLimitError, not a silent "no PR found\"""" + """A rate-limited 403 surfaces as GitHubRateLimitError, not a silent "no PR found\" """ mock_response = Mock() mock_response.status_code = 403 mock_response.headers = {"x-ratelimit-remaining": "0", "x-ratelimit-reset": "1777415867"} diff --git a/tests/comfy_cli/command/test_run.py b/tests/comfy_cli/command/test_run.py index 1b9d8672..60a64cc0 100644 --- a/tests/comfy_cli/command/test_run.py +++ b/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,7 +135,7 @@ 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: @@ -142,20 +143,30 @@ def test_network_error_exits_cleanly(self): 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"not json"), ): 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,8 +187,8 @@ 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) @@ -185,16 +196,16 @@ def test_queue_embeds_api_key_in_extra_data(self, workflow): 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): diff --git a/tests/comfy_cli/command/test_run_json.py b/tests/comfy_cli/command/test_run_json.py index 15f16d14..65af669e 100644 --- a/tests/comfy_cli/command/test_run_json.py +++ b/tests/comfy_cli/command/test_run_json.py @@ -160,10 +160,10 @@ def test_events_are_noop_in_pretty_mode(self, simple_workflow, capsys): def test_every_line_carries_schema_and_type(self, workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = [ @@ -281,10 +281,10 @@ class TestSuccessfulRun: def test_no_wait_emits_prompt_preview_then_queued(self, workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run._spawn_watcher", return_value=True), ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p123"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p123"}).encode() lines, exit_code = _run_execute_capture(workflow_file, capsys, wait=False) assert exit_code == 0 # prompt_preview is always emitted before queued so agents have a @@ -303,10 +303,10 @@ def test_envelope_after_success(self, workflow_file, capsys): """Mocked WS flow → queued + executing/executed/output events + ok envelope.""" with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance @@ -342,12 +342,12 @@ class TestQueueHttpErrors: def _setup_and_run(self, workflow_file, http_response, capsys, status=None, body=b""): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket"), ): if status is None: # Success path mock - mock_open.return_value.read.return_value = http_response + mock_open.return_value.__enter__.return_value.read.return_value = http_response else: mock_open.side_effect = _make_http_error(status, body) return _run_execute_capture(workflow_file, capsys) @@ -407,7 +407,7 @@ def test_200_with_utf16_bom_body_routes_to_invalid_response(self, workflow_file, def test_url_error_routes_to_connection_error(self, workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket"), ): mock_open.side_effect = urllib.error.URLError("refused") @@ -418,7 +418,7 @@ def test_validation_warnings_on_200_with_partial_node_errors(self, workflow_file """200 + non-empty node_errors → `queued` with validation_warnings populated.""" with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): body = json.dumps( @@ -427,7 +427,7 @@ def test_validation_warnings_on_200_with_partial_node_errors(self, workflow_file "node_errors": {"3": {"errors": [{"type": "x", "message": "skipped"}], "class_type": "X"}}, } ).encode() - mock_open.return_value.read.return_value = body + mock_open.return_value.__enter__.return_value.read.return_value = body ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = [ @@ -447,10 +447,10 @@ def test_queued_nodes_manifest_from_workflow(self, workflow_file, capsys, simple """`nodes` lists one entry per workflow node with node_id, class_type, title.""" with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run._spawn_watcher", return_value=True), ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() lines, exit_code = _run_execute_capture(workflow_file, capsys, wait=False) queued = next(e for e in _events(lines) if e["type"] == "queued") assert queued["client_id"] @@ -467,10 +467,10 @@ class TestWebSocketEvents: def _run_with_ws_messages(self, workflow_file, recv_side_effect, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = recv_side_effect @@ -865,10 +865,10 @@ class TestPromptPreviewAlwaysEmitted: def test_api_input_emits_prompt_preview_before_queued(self, workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = [ @@ -885,10 +885,10 @@ def test_ui_input_emits_converted_then_prompt_preview_then_queued(self, ui_workf with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), patch("comfy_cli.command.run.fetch_object_info", return_value=OBJECT_INFO), - patch("comfy_cli.command.run.request.urlopen") as mock_post, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_post, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_post.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_post.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = [ @@ -907,10 +907,10 @@ def test_prompt_preview_excludes_client_id_and_extra_data(self, workflow_file, c # POST envelope's runtime fields (client_id, extra_data with api_key). with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = [ @@ -934,7 +934,7 @@ def test_api_input_emits_prompt_preview_and_envelope_only(self, workflow_file, c with ( patch("comfy_cli.command.run.check_comfy_server_running") as mock_check, patch("comfy_cli.command.run.fetch_object_info") as mock_fetch, - patch("comfy_cli.command.run.request.urlopen") as mock_post, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_post, ): lines, exit_code = _run_execute_capture(workflow_file, capsys, print_prompt=True) assert mock_check.call_count == 0 @@ -953,7 +953,7 @@ def test_ui_input_emits_converted_then_prompt_preview(self, ui_workflow_file, ca with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), patch("comfy_cli.command.run.fetch_object_info", return_value=OBJECT_INFO), - patch("comfy_cli.command.run.request.urlopen") as mock_post, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_post, ): lines, exit_code = _run_execute_capture(ui_workflow_file, capsys, print_prompt=True) assert mock_post.call_count == 0 @@ -969,7 +969,7 @@ def test_ui_input_with_unreachable_object_info_routes_to_connection_error(self, # --print-prompt skips the pre-flight server probe, but UI conversion # still needs /object_info, so an unreachable host surfaces here. with ( - patch("comfy_cli.command.run.request.urlopen", side_effect=urllib.error.URLError("Connection refused")), + patch("comfy_cli.http._PLAIN_OPENER.open", side_effect=urllib.error.URLError("Connection refused")), ): lines, exit_code = _run_execute_capture(ui_workflow_file, capsys, print_prompt=True) assert exit_code == 1 @@ -1025,10 +1025,10 @@ def test_converted_event_for_ui_input(self, ui_workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), patch("comfy_cli.command.run.fetch_object_info", return_value=OBJECT_INFO), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = [ @@ -1099,7 +1099,7 @@ class TestObjectInfoFailures: def test_object_info_unavailable_on_http_error(self, ui_workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._PLAIN_OPENER.open") as mock_open, ): # _make_http_error builds a /prompt URL by default — build the # /object_info HTTPError inline so the test exercises that path. @@ -1121,7 +1121,7 @@ def test_object_info_connection_error_on_urlerror(self, ui_workflow_file, capsys """URLError on /object_info → connection_error (NOT object_info_unavailable).""" with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._PLAIN_OPENER.open") as mock_open, ): mock_open.side_effect = urllib.error.URLError("connection refused") lines, exit_code = _run_execute_capture(ui_workflow_file, capsys) @@ -1135,10 +1135,10 @@ class TestNodeCachedIntegration: def test_execution_cached_event_shape(self, workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = [ @@ -1249,10 +1249,10 @@ class TestVerboseNoOpInJsonMode: def test_verbose_does_not_corrupt_json_stream(self, workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = [ @@ -1323,7 +1323,7 @@ def test_object_info_timeout_routes_to_connection_error(self, capsys): try: with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen", side_effect=TimeoutError("timed out")), + patch("comfy_cli.http._AUTHED_OPENER.open", side_effect=TimeoutError("timed out")), ): lines, exit_code = _run_execute_capture(path, capsys) assert _envelope(lines)["error"]["code"] == "connection_error" @@ -1355,7 +1355,7 @@ def test_object_info_non_json_body_routes_to_object_info_unavailable(self, capsy mock_resp.__exit__ = MagicMock(return_value=False) with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen", return_value=mock_resp), + patch("comfy_cli.http._PLAIN_OPENER.open", return_value=mock_resp), ): lines, exit_code = _run_execute_capture(path, capsys) env = _envelope(lines) @@ -1368,7 +1368,7 @@ def test_queue_timeout_error_routes_to_connection_error(self, workflow_file, cap """queue()'s urlopen TimeoutError → connection_error.""" with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen", side_effect=TimeoutError("post timed out")), + patch("comfy_cli.http._AUTHED_OPENER.open", side_effect=TimeoutError("post timed out")), patch("comfy_cli.command.run.WebSocket"), ): lines, exit_code = _run_execute_capture(workflow_file, capsys) @@ -1378,7 +1378,7 @@ def test_queue_oserror_routes_to_connection_error(self, workflow_file, capsys): """queue()'s urlopen OSError → connection_error.""" with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen", side_effect=OSError("network unreachable")), + patch("comfy_cli.http._AUTHED_OPENER.open", side_effect=OSError("network unreachable")), patch("comfy_cli.command.run.WebSocket"), ): lines, exit_code = _run_execute_capture(workflow_file, capsys) @@ -1467,10 +1467,10 @@ def test_two_consecutive_executing_includes_intermediate(self, workflow_file, ca are still included so consumers see the complete 'what ran' picture.""" with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = [ @@ -1506,10 +1506,10 @@ class TestTimeoutAppliesToConnectAndPost: def test_queue_passes_timeout_to_urlopen(self, workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance # Single executing(node=None) → on_executing returns False → loop exits @@ -1561,10 +1561,10 @@ def test_preflight_probe_passes_timeout(self, workflow_file, capsys): def test_connect_passes_timeout_to_ws_connect(self, workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, patch("comfy_cli.command.run.WebSocket") as MockWs, ): - mock_open.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() ws_instance = MagicMock() MockWs.return_value = ws_instance ws_instance.recv.side_effect = [ @@ -1598,7 +1598,7 @@ class TestNoWaitQueueErrorRegression: def test_no_wait_with_400_emits_prompt_rejected(self, workflow_file, capsys): with ( patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), - patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, ): body = json.dumps( { diff --git a/tests/comfy_cli/jobs/test_jobs.py b/tests/comfy_cli/jobs/test_jobs.py index 48138336..81de3f5a 100644 --- a/tests/comfy_cli/jobs/test_jobs.py +++ b/tests/comfy_cli/jobs/test_jobs.py @@ -363,13 +363,13 @@ 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. + # Both paths now open through a shared opener in ``comfy_cli.http``: the + # local queue/interrupt calls via the redirect-following ``_PLAIN_OPENER``, + # the cloud cancel path via the no-redirect ``_AUTHED_OPENER``. 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._PLAIN_OPENER, "open", _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 97109ff5..0ab0bfaf 100644 --- a/tests/comfy_cli/test_http.py +++ b/tests/comfy_cli/test_http.py @@ -7,7 +7,7 @@ import pytest import comfy_cli.http as http_mod -from comfy_cli.http import NoRedirectHandler, authed_urlopen, build_authed_request +from comfy_cli.http import NoRedirectHandler, authed_urlopen, build_authed_request, no_redirect_urlopen def _target(*, api_key=None, auth_token=None): @@ -136,6 +136,35 @@ def test_migrated_caller_propagates_refused_redirect(): assert exc_info.value.code == 302 +# --------------------------------------------------------------------------- +# no_redirect_urlopen +# --------------------------------------------------------------------------- + + +def test_no_redirect_urlopen_uses_no_redirect_opener(): + """A prepared Request goes through _AUTHED_OPENER untouched — no header is + attached, since this helper's callers carry their credential themselves.""" + sentinel = object() + req = urllib.request.Request("http://127.0.0.1:8188/prompt", data=b"{}", method="POST") + with patch.object(http_mod._AUTHED_OPENER, "open", return_value=sentinel) as opened: + result = no_redirect_urlopen(req, timeout=15) + assert result is sentinel + assert opened.call_args.args[0] is req + assert opened.call_args.kwargs["timeout"] == 15 + + +def test_no_redirect_urlopen_propagates_refused_redirect(): + """The /prompt submit embeds a credential in its body, so a 30x must raise + rather than be followed — we don't lean on urllib dropping the body.""" + err = urllib.error.HTTPError( + "http://127.0.0.1:8188/prompt", 307, "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: + no_redirect_urlopen(urllib.request.Request("http://127.0.0.1:8188/prompt")) + assert exc_info.value.code == 307 + + @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`` diff --git a/tests/comfy_cli/test_http_only_openers.py b/tests/comfy_cli/test_http_only_openers.py new file mode 100644 index 00000000..89fbfadc --- /dev/null +++ b/tests/comfy_cli/test_http_only_openers.py @@ -0,0 +1,149 @@ +"""Scheme-pinning guards for every credential-carrying urllib opener. + +``urllib.request.build_opener()`` implicitly installs ``FileHandler``, +``FTPHandler`` and ``DataHandler``. An opener that attaches ``Authorization``/ +``X-API-Key`` must not speak those schemes: a caller steered into a ``file://`` +or ``ftp://`` URL could read unintended local content through an opener that +carries credentials. These tests mirror ``test_http.py``'s checks on the shared +``_AUTHED_OPENER`` for each of the other openers. +""" + +import urllib.error +import urllib.request + +import pytest + +from comfy_cli import comfy_client, http +from comfy_cli.cloud import oauth +from comfy_cli.command import transfer +from comfy_cli.cql import engine, loader + +NON_HTTP_URLS = ["file:///etc/passwd", "ftp://example.com/x", "data:text/plain,hi"] + +# Every opener the CLI builds, including the ones that deliberately follow +# redirects (_DOWNLOAD_OPENER strips auth headers and re-checks the scheme in +# its own handler; _PLAIN_OPENER carries nothing to leak). All of them should +# be scheme-pinned regardless. +OPENERS = [ + ("http._AUTHED_OPENER", http._AUTHED_OPENER), + ("comfy_client._OPENER", comfy_client._OPENER), + ("cql.engine._opener", engine._opener), + ("cql.loader._LOADER_OPENER", loader._LOADER_OPENER), + ("cloud.oauth._OAUTH_OPENER", oauth._OAUTH_OPENER), + ("transfer._TRANSFER_OPENER", transfer._TRANSFER_OPENER), + ("transfer._DOWNLOAD_OPENER", transfer._DOWNLOAD_OPENER), + ("http._PLAIN_OPENER", http._PLAIN_OPENER), +] +OPENER_IDS = [name for name, _ in OPENERS] + +# _PLAIN_OPENER carries no credentials and must keep following redirects: its +# call sites used urllib's global default opener before, which follows them. +# Pinning the scheme must not silently turn it into a no-redirect opener. +REDIRECT_FOLLOWING = ["http._PLAIN_OPENER"] + + +@pytest.mark.parametrize("opener", [o for _, o in OPENERS], ids=OPENER_IDS) +@pytest.mark.parametrize("url", NON_HTTP_URLS) +def test_opener_refuses_non_http_schemes(opener, url): + """Non-http(s) schemes fall through to UnknownHandler.""" + with pytest.raises(urllib.error.URLError) as exc_info: + opener.open(url) + assert "unknown url type" in str(exc_info.value.reason) + + +@pytest.mark.parametrize("opener", [o for _, o in OPENERS], ids=OPENER_IDS) +def test_opener_handler_set(opener): + """Only http(s)-relevant handlers are installed — no File/FTP/Data.""" + names = {type(h).__name__ for h in opener.handlers} + assert {"HTTPHandler", "HTTPSHandler", "UnknownHandler"} <= names + assert not names & {"FileHandler", "FTPHandler", "DataHandler"} + + +_NO_REDIRECT = [n for n, _ in OPENERS if n != "transfer._DOWNLOAD_OPENER" and n not in REDIRECT_FOLLOWING] + + +@pytest.mark.parametrize( + ("name", "opener"), + [(n, o) for n, o in OPENERS if n in _NO_REDIRECT], + ids=_NO_REDIRECT, +) +def test_credentialed_openers_refuse_redirects(name, opener): + """Every opener that carries credentials keeps its NoRedirectHandler.""" + assert any(isinstance(h, http.NoRedirectHandler) for h in opener.handlers) + + +def test_download_opener_still_follows_redirects(): + """_DOWNLOAD_OPENER's redirect handler is intentional — scheme-pinning it + must not turn it into a no-redirect opener.""" + assert any(isinstance(h, transfer._DownloadRedirectHandler) for h in transfer._DOWNLOAD_OPENER.handlers) + assert not any(isinstance(h, http.NoRedirectHandler) for h in transfer._DOWNLOAD_OPENER.handlers) + + +@pytest.mark.parametrize( + ("name", "opener"), + [(n, o) for n, o in OPENERS if n in REDIRECT_FOLLOWING], + ids=REDIRECT_FOLLOWING, +) +def test_uncredentialed_openers_still_follow_redirects(name, opener): + """These sites used urllib's global default opener, which follows + redirects. They carry no credential, so there is nothing to replay at the + redirect target — pinning the scheme must not silently make a 30x start + raising HTTPError at call sites that transparently followed it before.""" + redirect_handlers = [h for h in opener.handlers if isinstance(h, urllib.request.HTTPRedirectHandler)] + assert redirect_handlers, f"{name} lost its redirect handler" + assert not any(isinstance(h, http.NoRedirectHandler) for h in redirect_handlers) + + +def test_build_http_only_opener_installs_caller_handlers(): + handler = http.NoRedirectHandler("custom message") + opener = http.build_http_only_opener(handler) + assert handler in opener.handlers + + +def _proxies(opener): + return next(h.proxies for h in opener.handlers if isinstance(h, urllib.request.ProxyHandler)) + + +def test_build_http_only_opener_keeps_http_proxy_support(monkeypatch): + """Proxy support is preserved for the schemes we actually speak.""" + monkeypatch.setenv("HTTP_PROXY", "http://proxy.example:3128") + monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:8080") + + proxies = _proxies(http.build_http_only_opener()) + assert proxies == {"http": "http://proxy.example:3128", "https": "http://proxy.example:8080"} + # ...and for the schemes we speak, that is exactly what build_opener resolves. + stdlib = _proxies(urllib.request.build_opener()) + assert proxies == {k: v for k, v in stdlib.items() if k in ("http", "https")} + + +def test_build_http_only_opener_ignores_non_http_proxies(monkeypatch): + """A proxy for a non-http scheme must not smuggle that scheme back in. + + ``ProxyHandler`` registers a ``_open`` per proxy entry, so an + ``ftp_proxy`` would otherwise hand the opener an ``ftp_open`` that wins + dispatch over ``UnknownHandler``. + """ + monkeypatch.setenv("FTP_PROXY", "http://proxy.example:3128") + monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:3128") + + opener = http.build_http_only_opener() + assert "ftp" not in _proxies(opener) + assert "ftp" not in opener.handle_open + with pytest.raises(urllib.error.URLError) as exc_info: + opener.open("ftp://example.com/x") + assert "unknown url type" in str(exc_info.value.reason) + + +def test_build_http_only_opener_lets_caller_override_a_default(): + """A caller-supplied handler replaces the default it subclasses, rather + than being appended behind it where it would never win dispatch.""" + + class PinnedHTTPSHandler(urllib.request.HTTPSHandler): + pass + + handler = PinnedHTTPSHandler() + opener = http.build_http_only_opener(handler) + + https_handlers = [h for h in opener.handlers if isinstance(h, urllib.request.HTTPSHandler)] + assert https_handlers == [handler] + assert opener.handle_open["https"] == [handler]