diff --git a/comfy_cli/file_utils.py b/comfy_cli/file_utils.py index d3f087c2..5ef049dc 100644 --- a/comfy_cli/file_utils.py +++ b/comfy_cli/file_utils.py @@ -12,6 +12,7 @@ from pathspec import PathSpec from comfy_cli import constants, ui +from comfy_cli.http import DEFAULT_HTTP_TIMEOUT, DOWNLOAD_TIMEOUT class DownloadException(Exception): @@ -93,8 +94,10 @@ def check_unauthorized(url: str, headers: dict | None = None) -> bool: bool: True if the response status code is 401, False otherwise. """ try: - response = requests.get(url, headers=headers, allow_redirects=True, stream=True) - return response.status_code == 401 + with requests.get( + url, headers=headers, allow_redirects=True, stream=True, timeout=DEFAULT_HTTP_TIMEOUT + ) as response: + return response.status_code == 401 except requests.RequestException: # If there's an error making the request, we can't determine if it's unauthorized return False @@ -566,7 +569,7 @@ def should_ignore(rel_path: str) -> bool: def upload_file_to_signed_url(signed_url: str, file_path: str): with open(file_path, "rb") as f: headers = {"Content-Type": "application/zip"} - response = requests.put(signed_url, data=f, headers=headers) + response = requests.put(signed_url, data=f, headers=headers, timeout=DOWNLOAD_TIMEOUT) if response.status_code == 200: print("Upload successful.") diff --git a/comfy_cli/http.py b/comfy_cli/http.py index a894afdb..77882b88 100644 --- a/comfy_cli/http.py +++ b/comfy_cli/http.py @@ -1,8 +1,19 @@ -"""Shared HTTP helpers with an auth-leak-safe redirect policy.""" +"""Shared HTTP helpers: default timeouts and an auth-leak-safe redirect policy.""" import urllib.error import urllib.request +# Default timeout (seconds) for plain, non-streaming API calls. Without an +# explicit timeout ``requests`` blocks forever on a stalled peer; this makes such +# a call fail fast with a typed ``requests.Timeout`` instead of hanging the CLI. +DEFAULT_HTTP_TIMEOUT = 30.0 + +# Timeout for streaming downloads and large uploads, as a (connect, read) tuple. +# ``requests`` applies the read timeout per socket read rather than to the whole +# transfer, so this caps how long we wait to *start* connecting/receiving without +# capping a legitimately long transfer. +DOWNLOAD_TIMEOUT = (10.0, 60.0) + class NoRedirectHandler(urllib.request.HTTPRedirectHandler): """Refuse to follow HTTP redirects. diff --git a/comfy_cli/registry/api.py b/comfy_cli/registry/api.py index eb1c65a7..bcddc749 100644 --- a/comfy_cli/registry/api.py +++ b/comfy_cli/registry/api.py @@ -4,6 +4,8 @@ import requests +from comfy_cli.http import DEFAULT_HTTP_TIMEOUT + # Reduced global imports from comfy_cli.registry from comfy_cli.registry.types import ( License, @@ -79,7 +81,7 @@ def publish_node_version( headers = {"Content-Type": "application/json"} body = request_body - response = requests.post(url, headers=headers, data=json.dumps(body)) + response = requests.post(url, headers=headers, data=json.dumps(body), timeout=DEFAULT_HTTP_TIMEOUT) if response.status_code == 201: data = response.json() @@ -98,7 +100,7 @@ def list_all_nodes(self): list: A list of Node instances. """ url = f"{self.base_url}/nodes" - response = requests.get(url) + response = requests.get(url, timeout=DEFAULT_HTTP_TIMEOUT) if response.status_code == 200: raw_nodes = response.json()["nodes"] return [map_node_to_node_class(node) for node in raw_nodes] @@ -123,7 +125,7 @@ def install_node(self, node_id, version=None): # A stalled/blackholed registry must not hang callers indefinitely. # A Timeout surfaces as a RequestException for callers to catch. - response = requests.get(url, timeout=30) + response = requests.get(url, timeout=DEFAULT_HTTP_TIMEOUT) if response.status_code == 200: # Convert the API response to a NodeVersion object logging.debug(f"RegistryAPI install_node response: {response.json()}") @@ -147,7 +149,7 @@ def get_node(self, node_id): """ url = f"{self.base_url}/nodes/{node_id}" # Same rationale as install_node: a stalled registry must not hang callers. - response = requests.get(url, timeout=30) + response = requests.get(url, timeout=DEFAULT_HTTP_TIMEOUT) if response.status_code == 200: logging.debug(f"RegistryAPI get_node response: {response.json()}") return map_node_to_node_class(response.json()) diff --git a/comfy_cli/standalone.py b/comfy_cli/standalone.py index 9c89bde5..b1544e2f 100644 --- a/comfy_cli/standalone.py +++ b/comfy_cli/standalone.py @@ -7,6 +7,7 @@ import requests from comfy_cli.constants import DEFAULT_STANDALONE_PYTHON_MINOR_VERSION, OS, PROC +from comfy_cli.http import DEFAULT_HTTP_TIMEOUT from comfy_cli.typing import PathLike from comfy_cli.utils import create_tarball, download_url, extract_tarball, get_os, get_proc from comfy_cli.uv import DependencyCompiler @@ -35,7 +36,7 @@ def _resolve_python_version(asset_url_prefix: str, minor_version: str) -> str: the available patch version for the requested minor series (e.g. "3.12" -> "3.12.13"). """ sha256sums_url = f"{asset_url_prefix.rstrip('/')}/SHA256SUMS" - response = requests.get(sha256sums_url) + response = requests.get(sha256sums_url, timeout=DEFAULT_HTTP_TIMEOUT) response.raise_for_status() pattern = re.compile(rf"cpython-({re.escape(minor_version)}\.\d+)\+") @@ -73,7 +74,7 @@ def download_standalone_python( if tag == "latest": # try to fetch json with info about latest release - response = requests.get(_latest_release_json_url) + response = requests.get(_latest_release_json_url, timeout=DEFAULT_HTTP_TIMEOUT) if response.status_code != 200: response.raise_for_status() raise RuntimeError(f"Request to {_latest_release_json_url} returned status code {response.status_code}") diff --git a/comfy_cli/utils.py b/comfy_cli/utils.py index e5c582e2..60f736c9 100644 --- a/comfy_cli/utils.py +++ b/comfy_cli/utils.py @@ -19,6 +19,7 @@ from rich.table import Table from comfy_cli.constants import DEFAULT_COMFY_WORKSPACE, OS, PROC +from comfy_cli.http import DOWNLOAD_TIMEOUT # Use the output shim so prints go to stderr (not stdout) in JSON mode, # preserving the one-envelope-on-stdout contract. @@ -121,21 +122,21 @@ def download_url( cwd = Path(cwd).expanduser().resolve() fpath = cwd / fname - response = requests.get(url, stream=True, allow_redirects=allow_redirects) - if response.status_code != 200: - response.raise_for_status() # Will only raise for 4xx codes, so... - raise RuntimeError(f"Request to {url} returned status code {response.status_code}") - - response.raw.read = functools.partial(response.raw.read, decode_content=True) # Decompress if needed - with fpath.open("wb") as f: - if show_progress: - fsize = int(response.headers.get("Content-Length", 0)) - desc = f"downloading {fname}..." + ("(Unknown total file size)" if fsize == 0 else "") - - with progress.wrap_file(cast(BinaryIO, response.raw), total=fsize, description=desc) as response_raw: - shutil.copyfileobj(response_raw, f) - else: - shutil.copyfileobj(response.raw, f) + with requests.get(url, stream=True, allow_redirects=allow_redirects, timeout=DOWNLOAD_TIMEOUT) as response: + if response.status_code != 200: + response.raise_for_status() # Will only raise for 4xx codes, so... + raise RuntimeError(f"Request to {url} returned status code {response.status_code}") + + response.raw.read = functools.partial(response.raw.read, decode_content=True) # Decompress if needed + with fpath.open("wb") as f: + if show_progress: + fsize = int(response.headers.get("Content-Length", 0)) + desc = f"downloading {fname}..." + ("(Unknown total file size)" if fsize == 0 else "") + + with progress.wrap_file(cast(BinaryIO, response.raw), total=fsize, description=desc) as response_raw: + shutil.copyfileobj(response_raw, f) + else: + shutil.copyfileobj(response.raw, f) return fpath diff --git a/tests/comfy_cli/registry/test_api.py b/tests/comfy_cli/registry/test_api.py index b8095b66..ef77ce45 100644 --- a/tests/comfy_cli/registry/test_api.py +++ b/tests/comfy_cli/registry/test_api.py @@ -4,6 +4,7 @@ import unittest from unittest.mock import MagicMock, patch +from comfy_cli.http import DEFAULT_HTTP_TIMEOUT from comfy_cli.registry import PyProjectConfig from comfy_cli.registry.api import RegistryAPI from comfy_cli.registry.types import ComfyConfig, License, ProjectConfig, URLs @@ -201,6 +202,54 @@ def test_install_node_failure(self, mock_get): self.registry_api.install_node("node1") self.assertIn("Failed to install node", str(context.exception)) + @patch("requests.post") + def test_publish_node_version_passes_timeout(self, mock_post): + """Registry calls must set a timeout so a stalled peer can't hang the CLI.""" + mock_response = MagicMock() + mock_response.status_code = 201 + mock_response.json.return_value = { + "node_version": { + "id": "test_node", + "version": "0.1.0", + "changelog": "", + "dependencies": [], + "deprecated": False, + "downloadUrl": "https://example.com/download", + }, + "signedUrl": "https://example.com/signed", + } + mock_post.return_value = mock_response + + self.registry_api.publish_node_version(self.node_config, self.token) + self.assertEqual(mock_post.call_args.kwargs["timeout"], DEFAULT_HTTP_TIMEOUT) + + @patch("requests.get") + def test_list_all_nodes_passes_timeout(self, mock_get): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"nodes": []} + mock_get.return_value = mock_response + + self.registry_api.list_all_nodes() + self.assertEqual(mock_get.call_args.kwargs["timeout"], DEFAULT_HTTP_TIMEOUT) + + @patch("requests.get") + def test_install_node_passes_timeout(self, mock_get): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "node1", + "version": "1.0.0", + "changelog": "", + "dependencies": [], + "deprecated": False, + "downloadUrl": "https://example.com/download1", + } + mock_get.return_value = mock_response + + self.registry_api.install_node("node1") + self.assertEqual(mock_get.call_args.kwargs["timeout"], DEFAULT_HTTP_TIMEOUT) + @patch("requests.get") def test_get_node_success(self, mock_get): mock_response = MagicMock() @@ -233,3 +282,17 @@ def test_get_node_failure(self, mock_get): with self.assertRaises(Exception) as context: self.registry_api.get_node("node1") self.assertIn("Failed to retrieve node", str(context.exception)) + + @patch("requests.get") + def test_get_node_passes_timeout(self, mock_get): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "node1", + "name": "Node One", + "description": "A node", + } + mock_get.return_value = mock_response + + self.registry_api.get_node("node1") + self.assertEqual(mock_get.call_args.kwargs["timeout"], DEFAULT_HTTP_TIMEOUT) diff --git a/tests/comfy_cli/test_http.py b/tests/comfy_cli/test_http.py index d5edf866..1c241f0d 100644 --- a/tests/comfy_cli/test_http.py +++ b/tests/comfy_cli/test_http.py @@ -8,6 +8,8 @@ import comfy_cli.http as http_mod from comfy_cli.http import ( + DEFAULT_HTTP_TIMEOUT, + DOWNLOAD_TIMEOUT, NoRedirectHandler, authed_urlopen, build_authed_request, @@ -23,6 +25,20 @@ def _target(*, api_key=None, auth_token=None, is_cloud=True): return types.SimpleNamespace(api_key=api_key, auth_token=auth_token, is_cloud=is_cloud) +def test_default_http_timeout_is_a_positive_scalar(): + assert isinstance(DEFAULT_HTTP_TIMEOUT, (int, float)) + assert DEFAULT_HTTP_TIMEOUT > 0 + + +def test_download_timeout_is_a_connect_read_tuple(): + # A (connect, read) tuple so connect failures fail fast while a legitimately + # long transfer is not capped (requests applies read timeout per socket read). + assert isinstance(DOWNLOAD_TIMEOUT, tuple) + assert len(DOWNLOAD_TIMEOUT) == 2 + connect, read = DOWNLOAD_TIMEOUT + assert connect > 0 and read > 0 + + def _call(handler, method_name, code=302): req = urllib.request.Request("https://example.com/thing") headers = http.client.HTTPMessage() diff --git a/tests/comfy_cli/test_standalone.py b/tests/comfy_cli/test_standalone.py index 3f63e11b..5f0cb1bf 100644 --- a/tests/comfy_cli/test_standalone.py +++ b/tests/comfy_cli/test_standalone.py @@ -5,6 +5,7 @@ import pytest import requests +from comfy_cli.http import DEFAULT_HTTP_TIMEOUT from comfy_cli.standalone import ( _latest_release_json_url, _resolve_python_version, @@ -79,7 +80,8 @@ def test_picks_highest_patch(self, mock_get): def test_url_construction(self, mock_get): mock_get.return_value = _mock_response(SAMPLE_SHA256SUMS) _resolve_python_version("https://example.com/release/", "3.12") - mock_get.assert_called_once_with("https://example.com/release/SHA256SUMS") + # A timeout must always be passed so a stalled peer can't hang the CLI. + mock_get.assert_called_once_with("https://example.com/release/SHA256SUMS", timeout=DEFAULT_HTTP_TIMEOUT) @patch("comfy_cli.standalone.requests.get") def test_no_false_match_across_minor(self, mock_get): @@ -120,6 +122,17 @@ def test_full_version_skips_resolution(self, mock_get, mock_download): # Should have fetched only latest-release.json, not SHA256SUMS assert mock_get.call_count == 1 + @patch("comfy_cli.standalone.download_url") + @patch("comfy_cli.standalone.requests.get") + def test_latest_release_fetch_passes_timeout(self, mock_get, mock_download): + """The latest-release.json fetch must set a timeout so a stalled peer can't hang.""" + mock_get.return_value = _mock_response('{"tag": "20260310", "asset_url_prefix": "https://example.com/release"}') + mock_download.return_value = "python.tar.gz" + + download_standalone_python(platform="linux", proc="x86_64", version="3.12.13") + + assert mock_get.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + _require_network = pytest.mark.skipif( os.getenv("TEST_NETWORK", "false").lower() != "true", diff --git a/tests/comfy_cli/test_utils.py b/tests/comfy_cli/test_utils.py index d16c01a9..784b0df4 100644 --- a/tests/comfy_cli/test_utils.py +++ b/tests/comfy_cli/test_utils.py @@ -1,6 +1,9 @@ import io from unittest.mock import MagicMock, patch +import pytest + +from comfy_cli.http import DOWNLOAD_TIMEOUT from comfy_cli.utils import create_tarball, download_url, extract_tarball @@ -17,20 +20,46 @@ def read(self, amt=-1, decode_content=False): return super().read(amt) +def _mock_streaming_response(mock_get, content): + """Wire a mock response that ``download_url`` can use as a context manager.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Length": str(len(content))} + mock_response.raw = _FakeRaw(content) + mock_get.return_value.__enter__.return_value = mock_response + return mock_response + + class TestDownloadUrl: @patch("comfy_cli.utils.requests.get") def test_writes_file(self, mock_get, tmp_path): content = b"file contents here" - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"Content-Length": str(len(content))} - mock_response.raw = _FakeRaw(content) - mock_get.return_value = mock_response + _mock_streaming_response(mock_get, content) result = download_url("http://example.com/f.bin", "f.bin", cwd=tmp_path, show_progress=False) assert result == tmp_path / "f.bin" assert (tmp_path / "f.bin").read_bytes() == content + @patch("comfy_cli.utils.requests.get") + def test_passes_download_timeout(self, mock_get, tmp_path): + """A streaming download must set a (connect, read) timeout so a stalled peer can't hang.""" + _mock_streaming_response(mock_get, b"x") + + download_url("http://example.com/f.bin", "f.bin", cwd=tmp_path, show_progress=False) + assert mock_get.call_args.kwargs["timeout"] == DOWNLOAD_TIMEOUT + + @patch("comfy_cli.utils.requests.get") + def test_releases_connection_on_error_status(self, mock_get, tmp_path): + """A non-200 must still release the streamed connection rather than leak it until GC.""" + mock_response = _mock_streaming_response(mock_get, b"") + mock_response.status_code = 500 + mock_response.raise_for_status.return_value = None + + with pytest.raises(RuntimeError): + download_url("http://example.com/f.bin", "f.bin", cwd=tmp_path, show_progress=False) + + mock_get.return_value.__exit__.assert_called_once() + class TestTarballRoundTrip: def test_create_and_extract(self, tmp_path, monkeypatch): diff --git a/tests/test_file_utils_network.py b/tests/test_file_utils_network.py index 2fe1347c..5d1b2b01 100644 --- a/tests/test_file_utils_network.py +++ b/tests/test_file_utils_network.py @@ -17,6 +17,7 @@ guess_status_code_reason, upload_file_to_signed_url, ) +from comfy_cli.http import DEFAULT_HTTP_TIMEOUT, DOWNLOAD_TIMEOUT def test_guess_status_code_reason_401_with_json(): @@ -47,20 +48,42 @@ def test_guess_status_code_reason_unknown(): assert "Unknown error occurred (status code: 500)" in result +def _mock_probe_response(mock_get, status_code): + """Wire a mock response that ``check_unauthorized`` can use as a context manager.""" + mock_response = Mock() + mock_response.status_code = status_code + mock_get.return_value.__enter__.return_value = mock_response + return mock_response + + @patch("requests.get") def test_check_unauthorized_true(mock_get): - mock_response = Mock() - mock_response.status_code = 401 - mock_get.return_value = mock_response + _mock_probe_response(mock_get, 401) assert check_unauthorized("http://example.com") is True +@patch("requests.get") +def test_check_unauthorized_passes_timeout(mock_get): + """The unauthorized probe must set a timeout so a stalled peer can't hang the CLI.""" + _mock_probe_response(mock_get, 200) + + check_unauthorized("http://example.com") + assert mock_get.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +@patch("requests.get") +def test_check_unauthorized_closes_response(mock_get): + """The probe reads only the status line, so it must release the streamed socket promptly.""" + _mock_probe_response(mock_get, 401) + + check_unauthorized("http://example.com") + mock_get.return_value.__exit__.assert_called_once() + + @patch("requests.get") def test_check_unauthorized_false(mock_get): - mock_response = Mock() - mock_response.status_code = 200 - mock_get.return_value = mock_response + _mock_probe_response(mock_get, 200) assert check_unauthorized("http://example.com") is False @@ -153,6 +176,7 @@ def test_upload_file_success(mock_put, tmp_path): upload_file_to_signed_url("http://example.com", str(test_file)) mock_put.assert_called_once() + assert mock_put.call_args.kwargs["timeout"] == DOWNLOAD_TIMEOUT @patch("requests.put")