ComfyUI/app/model_downloader/allowlist.py
DoronGenzelHass fdd84d04a0 feat(model_downloader): server-side model download + HuggingFace OAuth subsystem
Self-contained package under app/model_downloader/:
- Allowlist + path-validated downloads (SSRF guard: HF/Civitai/localhost + model extension).
- Streaming worker: writes to <final>.tmp, atomic rename on success, cooperative cancellation with epoch-based session identity, orphan .tmp sweep.
- Unified availability probe with per-URL gated/size caching; is_hf_downloadable recomputed per call so login/license changes surface within one poll.
- HuggingFace OAuth 2.0 PKCE flow with loopback callback server and on-disk (0600) token storage + transparent refresh.
- Pydantic request/response schemas and aiohttp routes under api/.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-22 15:16:59 +03:00

47 lines
1.5 KiB
Python

"""URL allowlist for server-side model fetches.
Mirrors the frontend's ``isModelDownloadable`` allowlist so the two flows
agree on which URLs are eligible for download. Server-side allowlisting is
the primary SSRF defense for this subsystem — workflow JSON is untrusted
input (anyone can hand-craft one), so we never let the server fetch URLs
outside this list.
"""
from urllib.parse import urlparse
# Frontend parity: ``missingModelDownload-*.js`` exports the same triple as
# ``i = [...]`` (Civitai / HuggingFace / localhost).
_ALLOWED_URL_PREFIXES = (
"https://huggingface.co/",
"https://civitai.com/",
"http://localhost:",
"http://127.0.0.1:",
)
# Frontend parity: same set as ``a = [...]`` in the bundle.
_ALLOWED_MODEL_EXTENSIONS = (
".safetensors",
".sft",
".ckpt",
".pth",
".pt",
)
def is_url_allowed(url: str) -> bool:
"""Check whether ``url`` is permitted as a server-side download source.
Returns True only when both:
- the URL starts with one of the allowed prefixes, AND
- the URL's final path segment ends with a known model extension.
Both checks are required to keep arbitrary HTML / API endpoints on
allowlisted hosts (e.g. ``https://huggingface.co/api/...``) off the table.
"""
if not isinstance(url, str) or not url:
return False
if not any(url.startswith(p) for p in _ALLOWED_URL_PREFIXES):
return False
path = urlparse(url).path
return any(path.endswith(ext) for ext in _ALLOWED_MODEL_EXTENSIONS)