From 48db08ae1a2d92ead186c7639d2f45db19e32afc Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Tue, 14 Jul 2026 01:33:19 +0200 Subject: [PATCH] Harden network requests, downloads, and path handling - Add core/network_utils.py with URL validation, public-IP enforcement, redirect validation, and safe request_public_url/validate_public_http_url helpers - Replace direct requests calls in downloader/type_utils/sources with request_public_url so redirects are validated before follow-up - Replace netloc string checks with host_matches_domain for CivitAI, CivArchive, and HuggingFace host matching - Add HuggingFace Xet download backend with progress reporting and cancellation support via hf_xet/huggingface_hub - Harden aria2 installer: reject symlinks/special files in archives, restrict custom aria2c paths to the managed install, and update UI copy - Normalize paths with realpath and enforce configured model roots in local hash/metadata routes, plus remove caller-selected metadata_path injection in sidecar writing - Sanitize download/network errors to strip signed query strings before logging or showing them to the user --- __init__.py | 124 +++++--- core/aria2_installer.py | 19 +- core/downloader.py | 393 ++++++++++++++++++++++++-- core/network_utils.py | 163 +++++++++++ core/sources/civarchive.py | 13 +- core/sources/huggingface.py | 5 +- core/type_utils.py | 21 +- web/resolver/views/options_methods.js | 4 +- 8 files changed, 646 insertions(+), 96 deletions(-) create mode 100644 core/network_utils.py diff --git a/__init__.py b/__init__.py index e7adc0c..15eca05 100644 --- a/__init__.py +++ b/__init__.py @@ -265,11 +265,16 @@ class ModelResolverExtension: read_json_safe, write_json_atomic, ) + from .core.network_utils import ( + UnsafeUrlError, + host_matches_domain, + request_public_url, + validate_public_http_url, + ) from .core.type_utils import ( first_non_empty, to_int, to_bool, - MODEL_EXTENSIONS, format_size_bytes, normalize_sha256, extract_sha256_from_metadata, @@ -305,9 +310,11 @@ class ModelResolverExtension: stop_aria2_daemon, get_download_directory, get_metadata_sidecar_path, + get_safe_metadata_sidecar_path, normalize_download_category, write_lora_manager_metadata, sanitize_download_filename, + is_allowed_model_download_filename, ) from .core.aria2_installer import Aria2InstallError, install_aria2_engine from .core.sources.popular import ( @@ -733,8 +740,17 @@ class ModelResolverExtension: {"error": "path is required"}, status=400 ) + normalized_path = os.path.realpath( + os.path.abspath(os.path.normpath(str(path))) + ) + if not is_path_in_configured_model_roots(normalized_path): + return web.json_response( + {"error": "path is outside configured model directories"}, + status=403, + ) + return web.json_response( - get_local_model_hash_metadata(path, model=model) + get_local_model_hash_metadata(normalized_path, model=model) ) @routes.post("/model_resolver/workflow-model-hashes") @@ -905,7 +921,9 @@ class ModelResolverExtension: {"error": "path is required"}, status=400 ) - normalized_path = os.path.abspath(os.path.normpath(target_path)) + normalized_path = os.path.realpath( + os.path.abspath(os.path.normpath(target_path)) + ) if not os.path.exists(normalized_path): return web.json_response( {"error": "path does not exist"}, status=404 @@ -942,7 +960,9 @@ class ModelResolverExtension: if not file_path: return "", "file_path is required" - normalized_path = _os.path.abspath(_os.path.normpath(file_path)) + normalized_path = _os.path.realpath( + _os.path.abspath(_os.path.normpath(file_path)) + ) if not _os.path.exists(normalized_path) or not _os.path.isfile(normalized_path): return "", "file does not exist" if not is_path_in_configured_model_roots(normalized_path): @@ -951,25 +971,12 @@ class ModelResolverExtension: def write_calculated_hash_metadata( normalized_path, - metadata_path, sha256, sha256_source="file", ): import os as _os - model_dir = _os.path.abspath(_os.path.dirname(normalized_path)) - resolved_metadata_path = "" - if metadata_path: - candidate_path = _os.path.abspath(_os.path.normpath(metadata_path)) - try: - same_dir = _os.path.commonpath([model_dir, candidate_path]) == model_dir - except ValueError: - same_dir = False - if same_dir: - resolved_metadata_path = candidate_path - - if not resolved_metadata_path: - resolved_metadata_path = get_metadata_sidecar_path(normalized_path) + resolved_metadata_path = get_safe_metadata_sidecar_path(normalized_path) metadata_updated = False try: @@ -1128,7 +1135,6 @@ class ModelResolverExtension: async def calculate_file_hash_route(request): """Calculate SHA256 for a local model and persist it to sidecar metadata.""" data = await request.json() - metadata_path = data.get("metadata_path") or "" normalized_path, error = resolve_hash_file_request(data) if error == "file_path is required": @@ -1148,7 +1154,6 @@ class ModelResolverExtension: ) resolved_metadata_path, metadata_updated = write_calculated_hash_metadata( normalized_path, - metadata_path, sha256, sha256_source, ) @@ -1172,7 +1177,6 @@ class ModelResolverExtension: import uuid data = await request.json() - metadata_path = data.get("metadata_path") or "" normalized_path, error = resolve_hash_file_request(data) if error == "file_path is required": return web.json_response( @@ -1212,7 +1216,6 @@ class ModelResolverExtension: ) resolved_metadata_path, metadata_updated = write_calculated_hash_metadata( normalized_path, - metadata_path, sha256, sha256_source, ) @@ -1958,6 +1961,29 @@ class ModelResolverExtension: if not file_path: file_path = find_local_file_path(filename, category) + if file_path: + file_path = _os.path.realpath( + _os.path.abspath(_os.path.normpath(str(file_path))) + ) + if not is_path_in_configured_model_roots(file_path): + if resolved_path: + return web.json_response( + { + "error": ( + "resolved_path is outside configured model directories" + ) + }, + status=403, + ) + file_path = None + elif not _os.path.isfile(file_path): + if resolved_path: + return web.json_response( + {"error": "resolved_path is not an existing model file"}, + status=404, + ) + file_path = None + if file_path and _os.path.exists(file_path): file_location = _os.path.dirname(file_path).replace("\\", "/") if file_location and not file_location.endswith("/"): @@ -2208,20 +2234,15 @@ class ModelResolverExtension: return urls def remote_download_url_is_alive(url): - try: - import requests - except Exception: - return True - headers = { "User-Agent": "ComfyUI-Model-Resolver/1.0", "Accept": "*/*", } try: - response = requests.head( + response, _final_url, _final_headers = request_public_url( + "HEAD", url, headers=headers, - allow_redirects=True, timeout=8, ) try: @@ -2235,10 +2256,10 @@ class ModelResolverExtension: pass try: - response = requests.get( + response, _final_url, _final_headers = request_public_url( + "GET", url, headers={**headers, "Range": "bytes=0-0"}, - allow_redirects=True, stream=True, timeout=8, ) @@ -3113,15 +3134,16 @@ class ModelResolverExtension: if raw_url.startswith("hf://"): normalized_url = raw_url else: - from urllib.parse import urlparse - - parsed_input = urlparse(raw_url) - if parsed_input.scheme not in {"http", "https"}: + try: + normalized_url = await asyncio.to_thread( + validate_public_http_url, + raw_url, + ) + except UnsafeUrlError as exc: return web.json_response( - {"error": "Only http, https, and hf:// URLs are supported"}, + {"error": str(exc)}, status=400, ) - normalized_url = raw_url category = data.get("category") or "" expected_filename = ( @@ -3225,9 +3247,9 @@ class ModelResolverExtension: from urllib.parse import urlparse parsed_direct = urlparse(normalized_url) - direct_host = parsed_direct.netloc.lower() + direct_host = parsed_direct.hostname if ( - direct_host.endswith("civarchive.com") + host_matches_domain(direct_host, "civarchive.com") and "/api/download/" in parsed_direct.path ): source = "civarchive" @@ -4549,6 +4571,14 @@ class ModelResolverExtension: return web.json_response( {"error": "URL is required"}, status=400 ) + try: + url = await asyncio.to_thread(validate_public_http_url, url) + except UnsafeUrlError as exc: + return web.json_response({"error": str(exc)}, status=400) + + from urllib.parse import urlparse as _download_urlparse + + download_host = _download_urlparse(url).hostname if not filename: # Extract filename from URL @@ -4562,14 +4592,22 @@ class ModelResolverExtension: return web.json_response( {"error": "Could not determine filename"}, status=400 ) + if not is_allowed_model_download_filename(filename): + return web.json_response( + {"error": "Unsupported model file extension"}, status=400 + ) # Build headers if needed headers = {} - if "huggingface.co" in url: + if host_matches_domain(download_host, "huggingface.co"): hf_token = data.get("hf_token", "") if hf_token: headers["Authorization"] = f"Bearer {hf_token}" - elif "civitai.com" in url or "civitai.red" in url: + elif host_matches_domain( + download_host, + "civitai.com", + "civitai.red", + ): civitai_key = data.get("civitai_key", "") if civitai_key and "token=" not in url: url += ( @@ -4590,9 +4628,9 @@ class ModelResolverExtension: return to_int(value) inferred_source = "" - if "civitai.com" in url: + if host_matches_domain(download_host, "civitai.com", "civitai.red"): inferred_source = "civitai" - elif "huggingface.co" in url: + elif host_matches_domain(download_host, "huggingface.co"): inferred_source = "huggingface" download_metadata.setdefault("filename", filename) diff --git a/core/aria2_installer.py b/core/aria2_installer.py index 593cfa5..9eb493a 100644 --- a/core/aria2_installer.py +++ b/core/aria2_installer.py @@ -169,7 +169,7 @@ def _select_release_asset(release: Dict[str, Any], tokens: Dict[str, str]) -> Di asset_names = ", ".join(_asset_name(asset) for asset in assets if _asset_name(asset)) raise Aria2InstallError( f"No prebuilt aria2 desktop binary for {tokens['label']} was found in {release_name}. " - "Install aria2 with your system package manager or configure aria2c_path manually. " + "Install aria2 with your system package manager so aria2c is available on PATH. " f"Available assets: {asset_names or 'none'}" ) @@ -184,7 +184,7 @@ def _fetch_latest_release() -> Dict[str, Any]: except requests.RequestException as exc: raise Aria2InstallError( "Could not contact GitHub to check the latest aria2 release. " - "Check your internet connection, firewall/proxy settings, or configure aria2c_path manually. " + "Check your internet connection or firewall/proxy settings. " f"Details: {exc}" ) from exc data = response.json() @@ -206,6 +206,11 @@ def _safe_extract_zip(archive_path: Path, destination: Path) -> None: destination_abs = destination.resolve() with zipfile.ZipFile(archive_path) as archive: for member in archive.infolist(): + unix_mode = member.external_attr >> 16 + if unix_mode and stat.S_ISLNK(unix_mode): + raise Aria2InstallError( + f"Refusing symbolic link in aria2 archive: {member.filename}" + ) target = destination_abs / member.filename target_abs = target.resolve() if not is_path_within(str(target_abs), str(destination_abs)) and target_abs != destination_abs: @@ -217,6 +222,14 @@ def _safe_extract_tar(archive_path: Path, destination: Path) -> None: destination_abs = destination.resolve() with tarfile.open(archive_path) as archive: for member in archive.getmembers(): + if member.issym() or member.islnk(): + raise Aria2InstallError( + f"Refusing link in aria2 archive: {member.name}" + ) + if not (member.isfile() or member.isdir()): + raise Aria2InstallError( + f"Refusing special file in aria2 archive: {member.name}" + ) target = destination_abs / member.name target_abs = target.resolve() if not is_path_within(str(target_abs), str(destination_abs)) and target_abs != destination_abs: @@ -356,7 +369,7 @@ def _download_file(url: str, destination: Path) -> int: except requests.RequestException as exc: raise Aria2InstallError( "Could not download the aria2 release asset from GitHub. " - "Check your internet connection, firewall/proxy settings, or configure aria2c_path manually. " + "Check your internet connection or firewall/proxy settings. " f"Details: {exc}" ) from exc return total diff --git a/core/downloader.py b/core/downloader.py index dce92b2..4a918d6 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -26,7 +26,12 @@ log = create_module_logger(__name__) from .resolver import invalidate_local_hash_match_cache, normalize_sha256 from .scanner import invalidate_model_files_cache from .path_utils import is_path_within, get_path_identity, write_json_atomic, read_json_safe, get_comfy_root_path, calculate_file_sha256, get_filename_from_path -from .type_utils import as_dict, as_list, first_non_empty, format_size_bytes as format_bytes, DEFAULT_BROWSER_USER_AGENT, extract_response_file_size, normalize_category_to_model_type, get_category_folder_keys +from .network_utils import ( + host_matches_domain, + request_public_url, + validate_public_http_url, +) +from .type_utils import as_dict, as_list, first_non_empty, format_size_bytes as format_bytes, DEFAULT_BROWSER_USER_AGENT, extract_response_file_size, normalize_category_to_model_type, get_category_folder_keys, MODEL_EXTENSIONS try: @@ -55,6 +60,7 @@ CLI_LOG_INTERVAL = 5 # Log progress to CLI every N seconds ARIA2_RPC_TIMEOUT = (2, 5) # local JSON-RPC should respond quickly ARIA2_IDLE_STOP_SECONDS = 5 * 60 DOWNLOAD_USER_AGENT = DEFAULT_BROWSER_USER_AGENT +MANAGED_ARIA2_ROOT = Path(__file__).resolve().parents[1] / "tools" / "aria2" from .settings import CATEGORY_MAP, load_settings, normalize_download_backend, normalize_download_category, normalize_relative_subfolder @@ -88,6 +94,7 @@ SENSITIVE_QUERY_KEYS = { } _INVALID_DOWNLOAD_FILENAME_RE = re.compile(r'[<>:"/\\|?*\x00-\x1f]+') +_HTTP_URL_IN_TEXT_RE = re.compile(r"https?://[^\s]+", re.IGNORECASE) def sanitize_download_filename(filename: Any) -> str: @@ -100,6 +107,12 @@ def sanitize_download_filename(filename: Any) -> str: return text +def is_allowed_model_download_filename(filename: Any) -> bool: + """Return True only for model file extensions supported by the resolver.""" + safe_name = sanitize_download_filename(filename) + return bool(safe_name and os.path.splitext(safe_name)[1].lower() in MODEL_EXTENSIONS) + + def _is_sensitive_metadata_key(key: Any) -> bool: key_text = str(key or "").strip().lower() return ( @@ -129,6 +142,23 @@ def _strip_sensitive_url_params(value: str) -> str: return urlunparse(parsed._replace(query=urlencode(filtered, doseq=True))) +def _sanitize_download_error(value: Any) -> str: + """Remove signed query strings and credentials from errors shown or logged.""" + text = str(value or "") + + def redact_url(match: re.Match) -> str: + raw_url = match.group(0) + try: + parsed = urlparse(raw_url) + if parsed.query: + return urlunparse(parsed._replace(query="", fragment="")) + except Exception: + pass + return raw_url + + return _HTTP_URL_IN_TEXT_RE.sub(redact_url, text) + + def _clean_http_header_value(value: Any) -> str: return str(value or "").replace("\r", "").replace("\n", "").strip() @@ -162,8 +192,8 @@ def build_download_headers( _set_header_default(request_headers, "Accept", "*/*") _set_header_default(request_headers, "Accept-Encoding", "identity") - host = urlparse(str(url or "")).netloc.lower() - if host in {"civitai.com", "www.civitai.com", "civitai.red", "www.civitai.red"}: + host = urlparse(str(url or "")).hostname + if host_matches_domain(host, "civitai.com", "civitai.red"): _set_header_default(request_headers, "Referer", "https://civitai.com/") _set_header_default(request_headers, "Origin", "https://civitai.com") @@ -172,9 +202,9 @@ def build_download_headers( def _is_civitai_api_download_url(url: str) -> bool: parsed = urlparse(str(url or "")) - host = parsed.netloc.lower() + host = parsed.hostname return ( - host in {"civitai.com", "www.civitai.com", "civitai.red", "www.civitai.red"} + host_matches_domain(host, "civitai.com", "civitai.red") and parsed.path.startswith("/api/download/") ) @@ -184,14 +214,15 @@ def _resolve_civitai_download_url_for_aria2( headers: Optional[Dict[str, str]] = None, ) -> str: """Resolve CivitAI's API download redirect before handing the URL to aria2.""" - if not _is_civitai_api_download_url(url): - return url + validated_url = validate_public_http_url(url) + if not _is_civitai_api_download_url(validated_url): + return validated_url response = None try: response = requests.get( - url, - headers=build_download_headers(url, headers), + validated_url, + headers=build_download_headers(validated_url, headers), allow_redirects=False, stream=True, timeout=20, @@ -203,12 +234,12 @@ def _resolve_civitai_download_url_for_aria2( if not location: return url - resolved_url = urljoin(url, location.strip()) + resolved_url = validate_public_http_url(urljoin(validated_url, location.strip())) log.debug("Resolved CivitAI download redirect for aria2") return resolved_url except Exception as exc: log.warning(f"Could not pre-resolve CivitAI download URL for aria2: {exc}") - return url + return validated_url finally: if response is not None: try: @@ -217,6 +248,28 @@ def _resolve_civitai_download_url_for_aria2( pass +def _resolve_download_url_for_aria2( + url: str, + headers: Optional[Dict[str, str]] = None, +) -> tuple[str, Dict[str, str]]: + """Preflight an aria2 URL and validate every redirect before RPC handoff.""" + request_headers = build_download_headers(url, headers) + response = None + try: + response, resolved_url, resolved_headers = request_public_url( + "GET", + url, + headers=request_headers, + timeout=20, + stream=True, + ) + response.raise_for_status() + return resolved_url, resolved_headers + finally: + if response is not None: + response.close() + + def _json_safe_metadata(value: Any, depth: int = 0) -> Any: if depth > 10: return str(value) @@ -272,6 +325,24 @@ def get_metadata_sidecar_path(file_path: str) -> str: return f"{base_path}.metadata.json" +def get_safe_metadata_sidecar_path(file_path: str) -> str: + """Return the canonical sidecar path without allowing caller-selected targets.""" + raw_model_path = str(file_path or "").strip() + if not raw_model_path: + raise ValueError("A model path is required") + model_path = os.path.realpath(os.path.abspath(raw_model_path)) + model_dir = os.path.realpath(os.path.dirname(model_path)) + metadata_path = os.path.realpath(get_metadata_sidecar_path(model_path)) + if ( + not model_dir + or os.path.dirname(metadata_path) != model_dir + or metadata_path == model_path + or not is_path_within(metadata_path, model_dir) + ): + raise ValueError("Metadata path is outside the model directory") + return metadata_path + + def _resolve_lora_manager_model_type(category: str, source_type: Any = "") -> str: res = normalize_category_to_model_type(category) if res in ("checkpoint", "diffusion_model", "embedding"): @@ -778,16 +849,34 @@ def _resolve_aria2c_executable(settings: Optional[Dict[str, Any]] = None) -> str active_settings = settings if isinstance(settings, dict) else load_settings() configured = str(active_settings.get("aria2c_path") or "").strip() candidate = os.path.expandvars(os.path.expanduser(configured or "aria2c")) + expected_names = {"aria2c", "aria2c.exe"} + candidate_name = os.path.basename(candidate).lower() + has_path_component = bool( + os.path.isabs(candidate) + or os.path.dirname(candidate) + or "/" in candidate + or "\\" in candidate + ) - resolved = shutil.which(candidate) - if resolved: - return resolved + if has_path_component: + if ( + candidate_name in expected_names + and os.path.isfile(candidate) + and is_path_within(candidate, MANAGED_ARIA2_ROOT) + ): + return os.path.realpath(os.path.abspath(candidate)) + raise Aria2Error( + "Custom aria2c paths are restricted to the managed Model Resolver install. " + "Use the built-in aria2 installer or place aria2c on PATH." + ) - if configured and os.path.isfile(candidate): - return candidate + if candidate_name in expected_names: + resolved = shutil.which(candidate) + if resolved and os.path.basename(resolved).lower() in expected_names: + return os.path.realpath(os.path.abspath(resolved)) raise Aria2Error( - "aria2c executable was not found. Install aria2 or configure aria2c_path." + "aria2c executable was not found. Use the built-in installer or place aria2c on PATH." ) @@ -1150,6 +1239,230 @@ def _delete_partial_download_files(dest_path: str) -> None: log.warning(f"Could not delete incomplete download file {path}: {exc}") +class _HuggingFaceXetDownloadCancelled(Exception): + """Raised by the Xet progress adapter when a download is cancelled.""" + + +class _HuggingFaceXetProgressAdapter: + """Forward hf_xet byte progress to the resolver's download state.""" + + def __init__(self, download_id: str, total_size: int, start_time: float) -> None: + self.download_id = download_id + self.total_size = max(0, int(total_size or 0)) + self.downloaded = 0 + self.last_update = start_time + + def update(self, byte_delta: Any) -> None: + if self.download_id in cancelled_downloads: + raise _HuggingFaceXetDownloadCancelled("Download cancelled") + + try: + delta = max(0, int(float(byte_delta or 0))) + except (TypeError, ValueError): + delta = 0 + + now = time.time() + elapsed = max(now - self.last_update, 0.001) + self.last_update = now + self.downloaded += delta + speed = int(delta / elapsed) if delta else 0 + progress = ( + min(100, int((self.downloaded / self.total_size) * 100)) + if self.total_size > 0 + else 0 + ) + + with download_lock: + state = download_progress.get(self.download_id) + if state is not None: + state.update( + { + "status": "downloading", + "progress": progress, + "downloaded": self.downloaded, + "total_size": self.total_size, + "speed": speed, + "download_backend": "huggingface_xet", + } + ) + + +def _download_huggingface_xet( + url: str, + dest_path: str, + download_id: str, + headers: Optional[Dict[str, str]] = None, + metadata: Optional[Dict[str, Any]] = None, + category: str = "", +) -> Optional[Dict[str, Any]]: + """Download Hugging Face Xet files with the official hf_xet transport.""" + validated_url = validate_public_http_url(url) + parsed_url = urlparse(validated_url) + if not ( + host_matches_domain(parsed_url.hostname, "huggingface.co") + and "/resolve/" in parsed_url.path + ): + return None + + try: + __import__("hf_xet") + from huggingface_hub.file_download import get_hf_file_metadata, xet_get + except ImportError: + return None + + request_headers = build_download_headers(validated_url, headers) + try: + file_metadata = get_hf_file_metadata( + validated_url, + headers=request_headers, + timeout=20, + ) + except Exception as exc: + log.debug( + "Hugging Face Xet metadata probe failed; using HTTP fallback: " + f"{type(exc).__name__}" + ) + return None + + xet_file_data = getattr(file_metadata, "xet_file_data", None) + if xet_file_data is None: + return None + + try: + expected_size = max(0, int(getattr(file_metadata, "size", 0) or 0)) + except (TypeError, ValueError): + expected_size = 0 + + result = { + "success": False, + "download_id": download_id, + "path": dest_path, + "error": None, + "size": 0, + } + start_time = time.time() + filename = get_filename_from_path(dest_path) + partial_path = f"{dest_path}.xet-part" + progress_adapter = _HuggingFaceXetProgressAdapter( + download_id, + expected_size, + start_time, + ) + + with download_lock: + download_progress[download_id] = { + "status": "starting", + "progress": 0, + "total_size": expected_size, + "downloaded": 0, + "filename": filename, + "path": dest_path, + "directory": os.path.dirname(dest_path), + "url": validated_url, + "error": None, + "speed": 0, + "start_time": start_time, + "download_backend": "huggingface_xet", + } + + cancelled_downloads.discard(download_id) + try: + os.makedirs(os.path.dirname(dest_path), exist_ok=True) + if os.path.exists(partial_path): + os.remove(partial_path) + + log.info(f"Starting Hugging Face Xet download: {filename}") + xet_get( + incomplete_path=Path(partial_path), + xet_file_data=xet_file_data, + headers=request_headers, + expected_size=expected_size or None, + displayed_filename=filename, + _tqdm_bar=progress_adapter, + ) + + if download_id in cancelled_downloads: + raise _HuggingFaceXetDownloadCancelled("Download cancelled") + + size = os.path.getsize(partial_path) + if expected_size and size != expected_size: + raise IOError( + f"Downloaded size mismatch: expected {expected_size}, received {size}" + ) + os.replace(partial_path, dest_path) + + metadata_path = write_lora_manager_metadata( + dest_path, + metadata or {}, + category, + validated_url, + ) + with download_lock: + state = download_progress.get(download_id) + if state is not None: + state.update( + { + "status": "completed", + "progress": 100, + "downloaded": size, + "total_size": expected_size or size, + "speed": 0, + } + ) + if metadata_path: + state["metadata_path"] = metadata_path + + result.update( + { + "success": True, + "size": size, + "metadata_path": metadata_path, + } + ) + elapsed = time.time() - start_time + avg_speed = size / elapsed if elapsed > 0 else 0 + log.info(f"✓ Hugging Face Xet download complete: {filename}") + log.info( + f"Size: {format_bytes(size)}, Time: {elapsed:.1f}s, " + f"Avg speed: {format_bytes(int(avg_speed))}/s" + ) + invalidate_model_files_cache() + invalidate_local_hash_match_cache() + return result + except Exception as exc: + was_cancelled = ( + isinstance(exc, _HuggingFaceXetDownloadCancelled) + or download_id in cancelled_downloads + ) + try: + if os.path.exists(partial_path): + os.remove(partial_path) + except Exception as cleanup_exc: + log.warning(f"Could not delete incomplete Xet file: {cleanup_exc}") + + error_msg = ( + "Download cancelled" if was_cancelled else _sanitize_download_error(exc) + ) + with download_lock: + state = download_progress.get(download_id) + if state is not None: + state.update( + { + "status": "cancelled" if was_cancelled else "error", + "error": error_msg, + "speed": 0, + } + ) + cancelled_downloads.discard(download_id) + result["error"] = error_msg + if was_cancelled: + log.info(f"Hugging Face Xet download cancelled: {filename}") + else: + log.error(f"✗ Hugging Face Xet download failed: {filename}") + log.error(f"Error: {error_msg}") + return result + + def download_file_with_aria2( url: str, dest_path: str, @@ -1189,13 +1502,7 @@ def download_file_with_aria2( try: os.makedirs(os.path.dirname(dest_path), exist_ok=True) _ensure_aria2_daemon(settings) - aria2_url = _resolve_civitai_download_url_for_aria2(url, headers) - aria2_headers = ( - None - if aria2_url != url and _is_civitai_api_download_url(url) - else headers - ) - request_headers = build_download_headers(aria2_url, aria2_headers) + aria2_url, request_headers = _resolve_download_url_for_aria2(url, headers) options: Dict[str, Any] = { "dir": os.path.dirname(dest_path), @@ -1350,7 +1657,7 @@ def download_file_with_aria2( time.sleep(0.5) except Exception as exc: - error_msg = str(exc) + error_msg = _sanitize_download_error(exc) with download_lock: if download_id in download_progress: download_progress[download_id]["status"] = "error" @@ -1396,6 +1703,17 @@ def download_file( """ global download_progress, cancelled_downloads + xet_result = _download_huggingface_xet( + url, + dest_path, + download_id, + headers=headers, + metadata=metadata, + category=category, + ) + if xet_result is not None: + return xet_result + if _download_backend_from_settings() == "aria2": return download_file_with_aria2( url, @@ -1447,11 +1765,12 @@ def download_file( # Verbose logging - what model and from where filename = get_filename_from_path(dest_path) + source_host = urlparse(url).hostname source = ( "HuggingFace" - if "huggingface.co" in url + if host_matches_domain(source_host, "huggingface.co") else "CivitAI" - if "civitai.com" in url + if host_matches_domain(source_host, "civitai.com", "civitai.red") else "URL" ) log.info(f"Starting download: {filename}") @@ -1460,8 +1779,16 @@ def download_file( # Start download request_headers = build_download_headers(url, headers) - response = requests.get(url, headers=request_headers, stream=True, timeout=30) + response, final_url, _final_headers = request_public_url( + "GET", + url, + headers=request_headers, + stream=True, + timeout=30, + ) response.raise_for_status() + if final_url != url: + log.debug("Validated download redirect target") # Get total size total_size = extract_response_file_size(response) or 0 @@ -1607,7 +1934,7 @@ def download_file( invalidate_local_hash_match_cache() except requests.exceptions.RequestException as e: - error_msg = str(e) + error_msg = _sanitize_download_error(e) # Check for specific HTTP errors if hasattr(e, "response") and e.response is not None: status_code = e.response.status_code @@ -1640,7 +1967,7 @@ def download_file( pass except Exception as e: - error_msg = str(e) + error_msg = _sanitize_download_error(e) with download_lock: download_progress[download_id]["status"] = "error" download_progress[download_id]["error"] = error_msg @@ -1690,6 +2017,12 @@ def download_model( "download_id": download_id, "error": "Invalid filename", } + if not is_allowed_model_download_filename(filename): + return { + "success": False, + "download_id": download_id, + "error": "Unsupported model file extension", + } subfolder = normalize_relative_subfolder(subfolder) # Get destination directory diff --git a/core/network_utils.py b/core/network_utils.py new file mode 100644 index 0000000..cd683ab --- /dev/null +++ b/core/network_utils.py @@ -0,0 +1,163 @@ +"""Network security helpers for outbound HTTP requests.""" + +from __future__ import annotations + +import ipaddress +import socket +from typing import Any, Dict, Optional, Tuple +from urllib.parse import urljoin, urlparse + +import requests + + +REDIRECT_STATUS_CODES = {301, 302, 303, 307, 308} +SENSITIVE_REDIRECT_HEADERS = {"authorization", "cookie", "proxy-authorization"} + + +class UnsafeUrlError(ValueError): + """Raised when an outbound URL could access a non-public network target.""" + + +def host_matches_domain(host: Any, *domains: str) -> bool: + """Return True for an exact domain or one of its real DNS subdomains.""" + normalized_host = str(host or "").strip().lower().rstrip(".") + if not normalized_host: + return False + for domain in domains: + normalized_domain = str(domain or "").strip().lower().rstrip(".") + if normalized_domain and ( + normalized_host == normalized_domain + or normalized_host.endswith(f".{normalized_domain}") + ): + return True + return False + + +def _is_public_ip(address: str) -> bool: + try: + parsed = ipaddress.ip_address(str(address).split("%", 1)[0]) + except ValueError: + return False + if isinstance(parsed, ipaddress.IPv6Address) and parsed.ipv4_mapped: + parsed = parsed.ipv4_mapped + return bool(parsed.is_global) + + +def validate_public_http_url(url: Any) -> str: + """Validate an HTTP(S) URL and reject local, private, or reserved targets.""" + text = str(url or "").strip() + if not text or len(text) > 8192: + raise UnsafeUrlError("A valid download URL is required") + + try: + parsed = urlparse(text) + port = parsed.port + except ValueError as exc: + raise UnsafeUrlError("The download URL is malformed") from exc + + if parsed.scheme.lower() not in {"http", "https"}: + raise UnsafeUrlError("Only http and https download URLs are supported") + if parsed.username is not None or parsed.password is not None: + raise UnsafeUrlError("Credentials embedded in download URLs are not allowed") + + host = str(parsed.hostname or "").strip().lower().rstrip(".") + if not host: + raise UnsafeUrlError("The download URL must include a hostname") + if host_matches_domain(host, "localhost"): + raise UnsafeUrlError("Localhost download URLs are not allowed") + + try: + direct_ip = ipaddress.ip_address(host.split("%", 1)[0]) + except ValueError: + direct_ip = None + + if direct_ip is not None: + if not _is_public_ip(str(direct_ip)): + raise UnsafeUrlError("Private, local, and reserved download addresses are not allowed") + return text + + try: + ascii_host = host.encode("idna").decode("ascii") + addresses = { + item[4][0] + for item in socket.getaddrinfo( + ascii_host, + port or (443 if parsed.scheme.lower() == "https" else 80), + type=socket.SOCK_STREAM, + ) + if item and len(item) > 4 and item[4] + } + except (OSError, UnicodeError) as exc: + raise UnsafeUrlError("The download hostname could not be resolved") from exc + + if not addresses: + raise UnsafeUrlError("The download hostname did not resolve to an address") + if any(not _is_public_ip(address) for address in addresses): + raise UnsafeUrlError("Private, local, and reserved download addresses are not allowed") + return text + + +def _url_origin(url: str) -> Tuple[str, str, int]: + parsed = urlparse(url) + scheme = parsed.scheme.lower() + port = parsed.port or (443 if scheme == "https" else 80) + return scheme, str(parsed.hostname or "").lower().rstrip("."), port + + +def _redirect_headers(headers: Dict[str, str], source_url: str, target_url: str) -> Dict[str, str]: + if _url_origin(source_url) == _url_origin(target_url): + return dict(headers) + return { + key: value + for key, value in headers.items() + if str(key).lower() not in SENSITIVE_REDIRECT_HEADERS + } + + +def request_public_url( + method: str, + url: Any, + *, + headers: Optional[Dict[str, str]] = None, + timeout: Any = 30, + stream: bool = True, + max_redirects: int = 5, +) -> Tuple[requests.Response, str, Dict[str, str]]: + """Perform a GET/HEAD request while validating every redirect target.""" + current_url = validate_public_http_url(url) + current_headers = dict(headers or {}) + request_method = str(method or "GET").upper() + if request_method not in {"GET", "HEAD"}: + raise ValueError("request_public_url supports only GET and HEAD") + + for redirect_count in range(max_redirects + 1): + request_func = requests.head if request_method == "HEAD" else requests.get + response = request_func( + current_url, + headers=current_headers, + allow_redirects=False, + stream=stream, + timeout=timeout, + ) + if response.status_code not in REDIRECT_STATUS_CODES: + return response, current_url, current_headers + + location = response.headers.get("Location") or response.headers.get("location") + if not location: + return response, current_url, current_headers + if redirect_count >= max_redirects: + response.close() + raise UnsafeUrlError("The download URL has too many redirects") + + candidate_url = urljoin(current_url, str(location).strip()) + try: + next_url = validate_public_http_url(candidate_url) + except UnsafeUrlError: + response.close() + raise + next_headers = _redirect_headers(current_headers, current_url, next_url) + response.close() + current_url = next_url + current_headers = next_headers + + raise UnsafeUrlError("The download URL has too many redirects") diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 45ece65..123a869 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -14,6 +14,7 @@ from urllib.parse import parse_qs, unquote, urljoin, urlparse import requests +from ..network_utils import host_matches_domain from ..matcher import ( calculate_archived_model_confidence, calculate_model_title_confidence, @@ -481,7 +482,7 @@ def parse_civarchive_url(url: str) -> Optional[Dict[str, Any]]: return None parsed = urlparse(urljoin(CIVARCHIVE_BASE_URL, url)) - if parsed.netloc and "civarchive.com" not in parsed.netloc: + if parsed.hostname and not host_matches_domain(parsed.hostname, "civarchive.com"): return None sha_match = re.search(r"/sha256/([a-fA-F0-9]{64})", parsed.path) @@ -573,18 +574,18 @@ def _prepare_size_probe_url(url: Any) -> Optional[str]: return None parsed = urlparse(normalized) - host = parsed.netloc.lower() + host = parsed.hostname path = parsed.path or "" - if host.endswith("huggingface.co") and "/blob/" in path: + if host_matches_domain(host, "huggingface.co") and "/blob/" in path: normalized = normalized.replace("/blob/", "/resolve/", 1) parsed = urlparse(normalized) - host = parsed.netloc.lower() + host = parsed.hostname path = parsed.path or "" - if host.endswith("civarchive.com") and not path.startswith("/api/download/"): + if host_matches_domain(host, "civarchive.com") and not path.startswith("/api/download/"): return None - if (host.endswith("civitai.com") or host.endswith("civitai.red")) and not path.startswith("/api/download/"): + if host_matches_domain(host, "civitai.com", "civitai.red") and not path.startswith("/api/download/"): return None return normalized diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 17b79d0..0d1e019 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -18,6 +18,7 @@ from ..log_system import create_module_logger log = create_module_logger(__name__) from ..path_utils import write_json_atomic, METADATA_DIR, read_json_safe +from ..network_utils import host_matches_domain from ..type_utils import check_credential_http, parse_size_header as _parse_size_header, fetch_remote_file_size, fetch_remote_file_size_cached, clear_remote_size_cache, extract_file_size from ..matcher import build_filename_search_queries, clean_filename_for_search @@ -302,7 +303,7 @@ def parse_huggingface_url(url: str) -> Optional[Dict[str, str]]: return None parsed = urlparse(url) - if "huggingface.co" not in parsed.netloc: + if not host_matches_domain(parsed.hostname, "huggingface.co"): return None match = re.match(r"^/([^/]+/[^/]+)/(resolve|blob)/([^/]+)/(.+)$", parsed.path) @@ -332,7 +333,7 @@ def _normalize_huggingface_size_probe_url(url: str) -> Optional[str]: return None value = url.strip() parsed = urlparse(value) - if "huggingface.co" not in parsed.netloc: + if not host_matches_domain(parsed.hostname, "huggingface.co"): return None if "/blob/" in parsed.path: value = value.replace("/blob/", "/resolve/", 1) diff --git a/core/type_utils.py b/core/type_utils.py index 4be4460..1caee5e 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -559,15 +559,15 @@ def fetch_remote_file_size( """ Check a remote file size using HEAD and, when needed, GET Range=bytes=0-0. """ - import requests + from .network_utils import request_public_url request_headers = {**(headers or {}), "Accept-Encoding": "identity"} size = None try: - response = requests.head( + response, _final_url, _final_headers = request_public_url( + "HEAD", url, headers=request_headers, - allow_redirects=True, timeout=timeout, ) try: @@ -580,10 +580,10 @@ def fetch_remote_file_size( if not size: try: - response = requests.get( + response, _final_url, _final_headers = request_public_url( + "GET", url, headers={**request_headers, "Range": "bytes=0-0"}, - allow_redirects=True, stream=True, timeout=timeout, ) @@ -604,26 +604,27 @@ def looks_like_model_file(url: str, expected_filename: str = "") -> bool: """ import os from urllib.parse import unquote, urlparse + from .network_utils import host_matches_domain text = str(url or "").strip() if not text.startswith(("http://", "https://", "hf://")): return False try: parsed = urlparse(text) - host = parsed.netloc.lower() + host = parsed.hostname path = unquote(parsed.path or "") except Exception: host = "" path = text - if host.endswith("huggingface.co") and path.startswith("/spaces/"): + if host_matches_domain(host, "huggingface.co") and path.startswith("/spaces/"): return False # Civitai/CivArchive api/download url prefix check if ( - host.endswith("civitai.com") - or host.endswith("civitai.red") - or host.endswith("civarchive.com") + host_matches_domain(host, "civitai.com") + or host_matches_domain(host, "civitai.red") + or host_matches_domain(host, "civarchive.com") ) and "/api/download/" in path: return True diff --git a/web/resolver/views/options_methods.js b/web/resolver/views/options_methods.js index b14dac0..e082143 100644 --- a/web/resolver/views/options_methods.js +++ b/web/resolver/views/options_methods.js @@ -294,7 +294,7 @@ export const optionsMethods = {
- aria2c path ? + aria2c path ?
@@ -2269,7 +2269,7 @@ export const optionsMethods = { if (data?.available) { setAria2Status(''); } else { - setAria2Status(data?.error || 'aria2c was not found. Install aria2 or configure the aria2c path.', 'is-invalid'); + setAria2Status(data?.error || 'aria2c was not found. Use the built-in installer or add aria2c to PATH.', 'is-invalid'); } } catch (error) { setAria2Status(error.message || 'aria2 check failed.', 'is-invalid');