diff --git a/comfy_cli/file_utils.py b/comfy_cli/file_utils.py index 9f31b8ce4..c140a2a86 100644 --- a/comfy_cli/file_utils.py +++ b/comfy_cli/file_utils.py @@ -16,6 +16,7 @@ from comfy_cli import constants, ui from comfy_cli._safe_exec import BinaryNotFoundError, resolve_required_binary +from comfy_cli.http import DEFAULT_HTTP_TIMEOUT, DOWNLOAD_TIMEOUT from comfy_cli.output.sanitize import sanitize_value logger = logging.getLogger(__name__) @@ -296,8 +297,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 @@ -1038,7 +1041,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 af83766c0..b9598ce63 100644 --- a/comfy_cli/http.py +++ b/comfy_cli/http.py @@ -1,10 +1,21 @@ -"""Shared HTTP helpers with an auth-leak-safe redirect policy.""" +"""Shared HTTP helpers: default timeouts and an auth-leak-safe redirect policy.""" import json import urllib.error import urllib.parse 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) + _LOOPBACK_HOSTS = {"localhost", "127.0.0.1", "::1", "[::1]"} diff --git a/comfy_cli/registry/api.py b/comfy_cli/registry/api.py index c6a4ec927..4ef28c466 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, @@ -133,7 +135,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() @@ -159,7 +161,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] @@ -189,7 +191,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()}") @@ -218,7 +220,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 7bfdc5d10..3dcf1f732 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 e41d21205..a39f13148 100644 --- a/comfy_cli/utils.py +++ b/comfy_cli/utils.py @@ -17,6 +17,7 @@ from rich.table import Table from comfy_cli.constants import DEFAULT_COMFY_WORKSPACE, OS, PROC +from comfy_cli.http import DOWNLOAD_TIMEOUT from comfy_cli.typing import PathLike @@ -106,21 +107,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 824e77ab9..cf54b27b9 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 ( MAX_ERROR_BODY_CHARS, @@ -237,6 +238,54 @@ def test_install_node_failure(self, mock_get): self.assertEqual(context.exception.status, 404) self.assertEqual(context.exception.body, "Not Found") + @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() @@ -270,6 +319,20 @@ def test_get_node_failure(self, mock_get): 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) + class TestSanitizeErrorBody(unittest.TestCase): def test_leaves_ordinary_body_untouched(self): diff --git a/tests/comfy_cli/test_http.py b/tests/comfy_cli/test_http.py index 194e7fd98..e00898f14 100644 --- a/tests/comfy_cli/test_http.py +++ b/tests/comfy_cli/test_http.py @@ -9,6 +9,8 @@ import comfy_cli.http as http_mod from comfy_cli.http import ( + DEFAULT_HTTP_TIMEOUT, + DOWNLOAD_TIMEOUT, MAX_RESPONSE_BYTES, NoRedirectHandler, ResponseTooLarge, @@ -28,6 +30,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 3f63e11bc..5f0cb1bf1 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 d16c01a95..784b0df4b 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 6d4e33803..ce3dc93e0 100644 --- a/tests/test_file_utils_network.py +++ b/tests/test_file_utils_network.py @@ -23,6 +23,7 @@ partial_paths_for, upload_file_to_signed_url, ) +from comfy_cli.http import DEFAULT_HTTP_TIMEOUT, DOWNLOAD_TIMEOUT def test_guess_status_code_reason_401_with_json(): @@ -53,20 +54,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 @@ -159,6 +182,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")