From 71cf5a11f18032a804b367fa3c5386476b456fe8 Mon Sep 17 00:00:00 2001 From: Talmaj Marinc Date: Sun, 12 Jul 2026 21:56:15 +0200 Subject: [PATCH] Remove api creds and add oauth. --- .../0006_drop_download_credentials.py | 58 ++++ app/model_downloader/api/routes.py | 84 ++--- app/model_downloader/api/schemas_in.py | 15 - app/model_downloader/api/schemas_out.py | 23 +- app/model_downloader/auth/oauth.py | 248 +++++++++++++++ app/model_downloader/auth/providers.py | 87 ++++++ app/model_downloader/auth/resolver.py | 42 +++ app/model_downloader/auth/store.py | 64 ++++ app/model_downloader/auth/token_store.py | 72 +++++ app/model_downloader/constants.py | 15 - app/model_downloader/credentials/resolver.py | 111 ------- app/model_downloader/credentials/store.py | 141 --------- app/model_downloader/database/models.py | 50 +-- app/model_downloader/database/queries.py | 110 +------ app/model_downloader/engine/job.py | 11 +- app/model_downloader/manager.py | 14 +- app/model_downloader/net/http.py | 20 +- app/model_downloader/net/probe.py | 10 +- app/model_downloader/scheduler.py | 1 - openapi.yaml | 181 ++++------- tests-unit/model_downloader_test/test_auth.py | 288 ++++++++++++++++++ .../model_downloader_test/test_credentials.py | 166 ---------- 22 files changed, 978 insertions(+), 833 deletions(-) create mode 100644 alembic_db/versions/0006_drop_download_credentials.py create mode 100644 app/model_downloader/auth/oauth.py create mode 100644 app/model_downloader/auth/providers.py create mode 100644 app/model_downloader/auth/resolver.py create mode 100644 app/model_downloader/auth/store.py create mode 100644 app/model_downloader/auth/token_store.py delete mode 100644 app/model_downloader/credentials/resolver.py delete mode 100644 app/model_downloader/credentials/store.py create mode 100644 tests-unit/model_downloader_test/test_auth.py delete mode 100644 tests-unit/model_downloader_test/test_credentials.py diff --git a/alembic_db/versions/0006_drop_download_credentials.py b/alembic_db/versions/0006_drop_download_credentials.py new file mode 100644 index 000000000..9c6a0f1e8 --- /dev/null +++ b/alembic_db/versions/0006_drop_download_credentials.py @@ -0,0 +1,58 @@ +""" +Drop download credential storage. + +Removes ``downloads.credential_id`` and the ``host_credentials`` table: stored +API keys are replaced by env-var keys plus an on-disk OAuth token store, neither +of which lives in the database. + +Revision ID: 0006_drop_download_credentials +Revises: 0005_download_manager +Create Date: 2026-07-09 +""" + +from alembic import op +import sqlalchemy as sa + +revision = "0006_drop_download_credentials" +down_revision = "0005_download_manager" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.batch_alter_table("downloads") as batch_op: + batch_op.drop_column("credential_id") + + op.drop_index("uq_host_credentials_host", table_name="host_credentials") + op.drop_table("host_credentials") + + +def downgrade() -> None: + op.create_table( + "host_credentials", + sa.Column("id", sa.String(length=36), primary_key=True), + sa.Column("host", sa.String(length=255), nullable=False), + sa.Column( + "match_subdomains", + sa.Boolean(), + nullable=False, + server_default=sa.text("false"), + ), + sa.Column("label", sa.String(length=255), nullable=True), + sa.Column( + "auth_scheme", sa.String(length=16), nullable=False, server_default="bearer" + ), + sa.Column("header_name", sa.String(length=255), nullable=True), + sa.Column("query_param", sa.String(length=255), nullable=True), + sa.Column("secret", sa.Text(), nullable=False), + sa.Column("secret_last4", sa.String(length=4), nullable=True), + sa.Column("enabled", sa.Boolean(), nullable=False, server_default=sa.text("true")), + sa.Column("created_at", sa.BigInteger(), nullable=False), + sa.Column("updated_at", sa.BigInteger(), nullable=False), + ) + op.create_index( + "uq_host_credentials_host", "host_credentials", ["host"], unique=True + ) + + with op.batch_alter_table("downloads") as batch_op: + batch_op.add_column(sa.Column("credential_id", sa.String(length=36), nullable=True)) \ No newline at end of file diff --git a/app/model_downloader/api/routes.py b/app/model_downloader/api/routes.py index 0a08edb41..ba74b6615 100644 --- a/app/model_downloader/api/routes.py +++ b/app/model_downloader/api/routes.py @@ -7,10 +7,9 @@ envelope used by ``app/assets/api/routes.py``: GET /api/download POST /api/download/availability POST /api/download/clear - POST /api/download/credentials - GET /api/download/credentials - GET /api/download/credentials/{id} - DELETE /api/download/credentials/{id} + GET /api/download/auth + POST /api/download/auth/{provider}/login + POST /api/download/auth/{provider}/logout GET /api/download/{id} DELETE /api/download/{id} POST /api/download/{id}/pause @@ -18,9 +17,9 @@ envelope used by ``app/assets/api/routes.py``: POST /api/download/{id}/cancel POST /api/download/{id}/priority -Note on ordering: the static ``credentials`` routes are registered before the -dynamic ``/api/download/{id}`` route so a request to ``.../credentials`` is not -captured as ``id == "credentials"``. +Note on ordering: the static ``auth`` routes are registered before the dynamic +``/api/download/{id}`` route so a request to ``.../auth`` is not captured as +``id == "auth"``. """ from __future__ import annotations @@ -31,10 +30,9 @@ from aiohttp import web from pydantic import BaseModel, ValidationError from app.model_downloader.api import schemas_in, schemas_out -from app.model_downloader.credentials.store import ( - CREDENTIAL_STORE, - CredentialValidationError, -) +from app.model_downloader.auth.oauth import LoginInProgress, OAuthNotConfigured +from app.model_downloader.auth.providers import PROVIDERS +from app.model_downloader.auth.store import AUTH_STORE from app.model_downloader.manager import DOWNLOAD_MANAGER, DownloadError ROUTES = web.RouteTableDef() @@ -89,7 +87,6 @@ async def enqueue(request: web.Request) -> web.Response: priority=parsed.priority, expected_sha256=parsed.expected_sha256, allow_any_extension=parsed.allow_any_extension, - credential_id=parsed.credential_id, ) except DownloadError as e: return _from_download_error(e) @@ -115,50 +112,35 @@ async def clear(request: web.Request) -> web.Response: return _ok({"deleted": deleted}) -# ----- credentials (secrets are write-only) — must precede /{id} ----- +# ----- auth (OAuth login + env-key status) — must precede /{id} ----- -@ROUTES.post("/api/download/credentials") -async def upsert_credential(request: web.Request) -> web.Response: - parsed = await _parse(request, schemas_in.CredentialUpsertRequest) - if isinstance(parsed, web.Response): - return parsed +@ROUTES.get("/api/download/auth") +async def auth_status(request: web.Request) -> web.Response: + return _ok({"providers": schemas_out.auth_status()}) + + +@ROUTES.post("/api/download/auth/{provider}/login") +async def auth_login(request: web.Request) -> web.Response: + provider = PROVIDERS.get(request.match_info["provider"]) + if provider is None: + return _error(400, "UNKNOWN_PROVIDER", "No such auth provider.") try: - view = await CREDENTIAL_STORE.upsert( - parsed.host, - parsed.secret, - auth_scheme=parsed.auth_scheme, - header_name=parsed.header_name, - query_param=parsed.query_param, - label=parsed.label, - match_subdomains=parsed.match_subdomains, - enabled=parsed.enabled, - ) - except CredentialValidationError as e: - return _error(400, "INVALID_CREDENTIAL", str(e)) - return _ok(schemas_out.credential_to_dict(view), status=201) + authorize_url = await AUTH_STORE.begin_login(provider) + except OAuthNotConfigured as e: + return _error(400, "OAUTH_NOT_CONFIGURED", str(e)) + except LoginInProgress as e: + return _error(409, "LOGIN_IN_PROGRESS", str(e)) + return _ok({"authorize_url": authorize_url}) -@ROUTES.get("/api/download/credentials") -async def list_credentials(request: web.Request) -> web.Response: - views = await CREDENTIAL_STORE.list() - return _ok({"credentials": [schemas_out.credential_to_dict(v) for v in views]}) - - -@ROUTES.get("/api/download/credentials/{id}") -async def get_credential(request: web.Request) -> web.Response: - view = await CREDENTIAL_STORE.get(request.match_info["id"]) - if view is None: - return _error(404, "NOT_FOUND", "No such credential.") - return _ok(schemas_out.credential_to_dict(view)) - - -@ROUTES.delete("/api/download/credentials/{id}") -async def delete_credential(request: web.Request) -> web.Response: - deleted = await CREDENTIAL_STORE.delete(request.match_info["id"]) - if not deleted: - return _error(404, "NOT_FOUND", "No such credential.") - return _ok({"deleted": True}) +@ROUTES.post("/api/download/auth/{provider}/logout") +async def auth_logout(request: web.Request) -> web.Response: + provider = PROVIDERS.get(request.match_info["provider"]) + if provider is None: + return _error(400, "UNKNOWN_PROVIDER", "No such auth provider.") + AUTH_STORE.clear(provider.name) + return _ok({"logged_out": True}) # ----- single download by id (dynamic; registered last) ----- diff --git a/app/model_downloader/api/schemas_in.py b/app/model_downloader/api/schemas_in.py index daccebb0e..bda8d9a69 100644 --- a/app/model_downloader/api/schemas_in.py +++ b/app/model_downloader/api/schemas_in.py @@ -10,8 +10,6 @@ from typing import Optional from pydantic import BaseModel, Field, field_validator -from app.model_downloader.constants import AUTH_SCHEME_BEARER - class EnqueueRequest(BaseModel): url: str @@ -19,7 +17,6 @@ class EnqueueRequest(BaseModel): priority: int = 0 expected_sha256: Optional[str] = None allow_any_extension: bool = False - credential_id: Optional[str] = None @field_validator("url") @classmethod @@ -42,20 +39,8 @@ class AvailabilityRequest(BaseModel): return {k: url.strip() for k, url in v.items()} -class CredentialUpsertRequest(BaseModel): - host: str - secret: str - auth_scheme: str = AUTH_SCHEME_BEARER - header_name: Optional[str] = None - query_param: Optional[str] = None - label: Optional[str] = None - match_subdomains: bool = False - enabled: bool = True - - __all__ = [ "EnqueueRequest", "PriorityRequest", "AvailabilityRequest", - "CredentialUpsertRequest", ] diff --git a/app/model_downloader/api/schemas_out.py b/app/model_downloader/api/schemas_out.py index 59ace6430..0eef0da81 100644 --- a/app/model_downloader/api/schemas_out.py +++ b/app/model_downloader/api/schemas_out.py @@ -1,26 +1,15 @@ """Response helpers for the download manager API. The download/status read models are plain dicts produced by the manager. This -module only needs to mask credentials for output (the secret is never returned). +module serializes the per-provider auth status (never a token) for the API. """ from __future__ import annotations -from app.model_downloader.credentials.store import CredentialView +from app.model_downloader.auth.providers import PROVIDERS +from app.model_downloader.auth.store import AUTH_STORE -def credential_to_dict(view: CredentialView) -> dict: - """API-safe credential representation — never includes the secret.""" - return { - "id": view.id, - "host": view.host, - "auth_scheme": view.auth_scheme, - "header_name": view.header_name, - "query_param": view.query_param, - "label": view.label, - "match_subdomains": view.match_subdomains, - "enabled": view.enabled, - "secret_last4": view.secret_last4, - "created_at": view.created_at, - "updated_at": view.updated_at, - } +def auth_status() -> list[dict]: + """Per-provider auth status — never includes a token.""" + return [AUTH_STORE.status(p) for p in PROVIDERS.values()] diff --git a/app/model_downloader/auth/oauth.py b/app/model_downloader/auth/oauth.py new file mode 100644 index 000000000..92b464de9 --- /dev/null +++ b/app/model_downloader/auth/oauth.py @@ -0,0 +1,248 @@ +"""Generic OAuth 2.0 PKCE engine + transient loopback callback server. + +The flow, per provider: + +1. :func:`start_login_flow` builds a PKCE challenge, binds a loopback callback + server on ``127.0.0.1:`` at ``/callback/``, and returns the + provider's authorize URL for the user to open. +2. The provider redirects the browser back to the loopback URL with a ``code`` + and the ``state`` we generated. The server validates ``state``, exchanges the + code for a :class:`Token`, hands it to the ``deliver`` sink, and tears down. +3. If no callback arrives within :data:`_LOGIN_TIMEOUT`, the server tears down. + +Only public PKCE clients are supported (no client secret). All outbound calls +go to the provider's own authorize/token endpoints, strictly user-initiated. +""" + +from __future__ import annotations + +import asyncio +import base64 +import hashlib +import logging +import os +import secrets +import time +from typing import Callable +from urllib.parse import urlencode + +from aiohttp import web + +from app.model_downloader.auth.providers import Provider +from app.model_downloader.auth.token_store import Token +from app.model_downloader.net.session import get_session, ssl_context + +# 0 means "let the OS pick a free loopback port" — RFC 8252 requires providers +# to accept any port on a loopback redirect, so each concurrent login can bind +# its own port instead of contending for one fixed port. Pin via env only if a +# custom OAuth app registered a specific loopback port. +CALLBACK_PORT = int(os.environ.get("COMFY_OAUTH_CALLBACK_PORT", "0")) +_LOGIN_TIMEOUT = 300.0 # seconds to wait for the browser callback + +# Token sink: called with the provider name and the exchanged Token. +TokenSink = Callable[[str, Token], None] + + +class OAuthError(Exception): + """A user-facing OAuth failure.""" + + +class OAuthNotConfigured(OAuthError): + """The provider has no public client id configured.""" + + +class LoginInProgress(OAuthError): + """A login flow for this provider is already running.""" + + +def _b64url(data: bytes) -> str: + return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii") + + +def _make_pkce() -> tuple[str, str]: + """Return ``(verifier, challenge)`` for the S256 PKCE method.""" + verifier = _b64url(secrets.token_bytes(32)) + challenge = _b64url(hashlib.sha256(verifier.encode("ascii")).digest()) + return verifier, challenge + + +def build_authorize_url( + provider: Provider, challenge: str, state: str, redirect_uri: str +) -> str: + params = { + "response_type": "code", + "client_id": provider.client_id, + "redirect_uri": redirect_uri, + "scope": provider.scope, + "state": state, + "code_challenge": challenge, + "code_challenge_method": "S256", + } + return f"{provider.authorize_url}?{urlencode(params)}" + + +def _token_from_payload(payload: dict) -> Token: + expires_in = payload.get("expires_in") + expires_at = int(time.time()) + int(expires_in) if expires_in else 0 + return Token( + access_token=payload["access_token"], + refresh_token=payload.get("refresh_token"), + expires_at=expires_at, + token_type=payload.get("token_type", "Bearer"), + scope=payload.get("scope"), + ) + + +async def _post_token(provider: Provider, data: dict) -> Token: + session = await get_session() + resp = await session.post( + provider.token_url, + data=data, + headers={"Accept": "application/json"}, + ssl=ssl_context(), + ) + try: + if resp.status != 200: + body = await resp.text() + raise OAuthError( + f"{provider.name} token endpoint returned HTTP {resp.status}: {body[:200]}" + ) + payload = await resp.json() + finally: + await resp.release() + if "access_token" not in payload: + raise OAuthError(f"{provider.name} token response missing access_token") + return _token_from_payload(payload) + + +async def exchange_code( + provider: Provider, code: str, verifier: str, redirect_uri: str +) -> Token: + return await _post_token( + provider, + { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": redirect_uri, + "client_id": provider.client_id, + "code_verifier": verifier, + }, + ) + + +async def refresh_access_token(provider: Provider, token: Token) -> Token: + if not token.refresh_token: + raise OAuthError(f"{provider.name} token is not refreshable") + refreshed = await _post_token( + provider, + { + "grant_type": "refresh_token", + "refresh_token": token.refresh_token, + "client_id": provider.client_id, + }, + ) + # Some providers omit a new refresh token on refresh; keep the old one. + if refreshed.refresh_token is None: + refreshed.refresh_token = token.refresh_token + return refreshed + + +class _LoginFlow: + """A single in-flight login: owns the loopback server and PKCE state.""" + + def __init__(self, provider: Provider, deliver: TokenSink) -> None: + self.provider = provider + self.deliver = deliver + self.verifier, self.challenge = _make_pkce() + self.state = secrets.token_urlsafe(24) + self.redirect_uri = "" # set once the loopback port is bound + self._runner: web.AppRunner | None = None + self._timeout_handle: asyncio.TimerHandle | None = None + + async def start(self) -> str: + app = web.Application() + app.router.add_get(f"/callback/{self.provider.name}", self._handle_callback) + self._runner = web.AppRunner(app) + await self._runner.setup() + site = web.TCPSite(self._runner, "127.0.0.1", CALLBACK_PORT) + await site.start() + port = site._server.sockets[0].getsockname()[1] + self.redirect_uri = f"http://127.0.0.1:{port}/callback/{self.provider.name}" + loop = asyncio.get_running_loop() + self._timeout_handle = loop.call_later( + _LOGIN_TIMEOUT, lambda: asyncio.ensure_future(self._teardown()) + ) + return build_authorize_url( + self.provider, self.challenge, self.state, self.redirect_uri + ) + + async def _handle_callback(self, request: web.Request) -> web.Response: + error = request.query.get("error") + if error: + asyncio.ensure_future(self._teardown()) + return web.Response( + text=f"Login failed: {error}", content_type="text/plain", status=400 + ) + if request.query.get("state") != self.state: + return web.Response( + text="Login failed: state mismatch.", + content_type="text/plain", + status=400, + ) + code = request.query.get("code") + if not code: + return web.Response( + text="Login failed: no authorization code.", + content_type="text/plain", + status=400, + ) + try: + token = await exchange_code( + self.provider, code, self.verifier, self.redirect_uri + ) + except OAuthError as e: + logging.warning("[model_downloader] %s login failed: %s", self.provider.name, e) + asyncio.ensure_future(self._teardown()) + return web.Response( + text=f"Login failed: {e}", content_type="text/plain", status=502 + ) + self.deliver(self.provider.name, token) + asyncio.ensure_future(self._teardown()) + return web.Response( + text="Login successful. You can close this window and return to ComfyUI.", + content_type="text/plain", + ) + + async def _teardown(self) -> None: + _ACTIVE.pop(self.provider.name, None) + if self._timeout_handle is not None: + self._timeout_handle.cancel() + self._timeout_handle = None + if self._runner is not None: + try: + await self._runner.cleanup() + except Exception: + logging.debug("[model_downloader] callback server cleanup error", exc_info=True) + self._runner = None + + +_ACTIVE: dict[str, _LoginFlow] = {} + + +def login_in_progress(provider_name: str) -> bool: + return provider_name in _ACTIVE + + +async def start_login_flow(provider: Provider, deliver: TokenSink) -> str: + """Begin a login flow and return the authorize URL to open in a browser.""" + if not provider.client_id: + raise OAuthNotConfigured( + f"OAuth app not configured for {provider.name}; set " + f"{provider.client_id_env} or use an env API key." + ) + if provider.name in _ACTIVE: + raise LoginInProgress(f"A login for {provider.name} is already in progress.") + flow = _LoginFlow(provider, deliver) + authorize_url = await flow.start() + _ACTIVE[provider.name] = flow + return authorize_url diff --git a/app/model_downloader/auth/providers.py b/app/model_downloader/auth/providers.py new file mode 100644 index 000000000..bf42b9306 --- /dev/null +++ b/app/model_downloader/auth/providers.py @@ -0,0 +1,87 @@ +"""Provider registry for download authentication. + +A :class:`Provider` describes a hub that can authenticate downloads either from +an environment API key or from an OAuth 2.0 access token. Both HuggingFace and +Civitai are public PKCE clients, so no client secret is ever stored; the public +``client_id`` is a placeholder overridable via env. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from urllib.parse import urlsplit + + +def normalize_host(host: str) -> str: + """Lowercase, strip port, IDNA-encode.""" + if not host: + return "" + host = host.strip() + if "://" in host: # a full URL was pasted — extract just the host + host = urlsplit(host).hostname or "" + host = host.lower() + if host.startswith("[") and "]" in host: # bracketed IPv6 literal + host = host[1 : host.index("]")] + elif host.count(":") == 1: # host:port (not IPv6) + host = host.split(":", 1)[0] + try: + host = host.encode("idna").decode("ascii") + except (UnicodeError, ValueError): + pass + return host + + +@dataclass(frozen=True) +class Provider: + name: str + host: str + authorize_url: str + token_url: str + scope: str + # Env vars to try, in order, for a plain API key. + env_keys: tuple[str, ...] + # Env var overriding the public OAuth client id. + client_id_env: str + # Public PKCE client id. Empty means "not configured" until the env sets it. + default_client_id: str = "" + + @property + def client_id(self) -> str: + return os.environ.get(self.client_id_env, self.default_client_id) or "" + + def env_token(self) -> str | None: + for var in self.env_keys: + token = os.environ.get(var) + if token: + return token + return None + + +PROVIDERS: dict[str, Provider] = { + "huggingface": Provider( + name="huggingface", + host="huggingface.co", + authorize_url="https://huggingface.co/oauth/authorize", + token_url="https://huggingface.co/oauth/token", + scope="openid gated-repos", + env_keys=("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"), + client_id_env="COMFY_HF_OAUTH_CLIENT_ID", + ), + "civitai": Provider( + name="civitai", + host="civitai.com", + authorize_url="https://auth.civitai.com/api/auth/oauth/authorize", + token_url="https://auth.civitai.com/api/auth/oauth/token", + scope="4", # ModelsRead; UserRead is auto-granted + env_keys=("CIVITAI_API_TOKEN", "CIVITAI_API_KEY"), + client_id_env="COMFY_CIVITAI_OAUTH_CLIENT_ID", + ), +} + +_HOST_TO_PROVIDER = {p.host: p for p in PROVIDERS.values()} + + +def provider_for_host(host: str) -> Provider | None: + """Return the provider whose host exactly matches ``host`` (normalized).""" + return _HOST_TO_PROVIDER.get(normalize_host(host)) diff --git a/app/model_downloader/auth/resolver.py b/app/model_downloader/auth/resolver.py new file mode 100644 index 000000000..4caf9e2d0 --- /dev/null +++ b/app/model_downloader/auth/resolver.py @@ -0,0 +1,42 @@ +"""Per-hop auth resolution (https only). + +Recomputed from scratch on every redirect hop: a hop only gets a bearer token +when *its own host* matches a configured provider, so a token bound to +``huggingface.co`` is silently dropped when the request is redirected to a +presigned CDN host — which is exactly what these hubs expect. + +For a matching hop: env API key first, then the provider's OAuth access token +(refreshed if expired), else no auth. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from app.model_downloader.auth.providers import provider_for_host +from app.model_downloader.auth.store import AUTH_STORE + + +@dataclass +class RequestAuth: + """How to modify a single request to carry a bearer token.""" + + headers: dict[str, str] = field(default_factory=dict) + + +async def resolve_auth_for_hop(host: str, scheme: str) -> RequestAuth | None: + """Resolve the bearer token (if any) to attach for one request hop.""" + if scheme.lower() != "https": + return None + provider = provider_for_host(host) + if provider is None: + return None + + token = provider.env_token() + if token: + return RequestAuth(headers={"Authorization": f"Bearer {token}"}) + + access = await AUTH_STORE.get_valid_token(provider) + if access: + return RequestAuth(headers={"Authorization": f"Bearer {access}"}) + return None diff --git a/app/model_downloader/auth/store.py b/app/model_downloader/auth/store.py new file mode 100644 index 000000000..30a8c45c2 --- /dev/null +++ b/app/model_downloader/auth/store.py @@ -0,0 +1,64 @@ +"""In-memory OAuth token cache over the on-disk token store. + +:data:`AUTH_STORE` is the process singleton the resolver and API talk to. It +lazily loads each provider's token from disk, refreshes an expired access token +via its refresh token, and orchestrates the login flow (delegating the loopback +callback server to :mod:`oauth`). +""" + +from __future__ import annotations + +from app.model_downloader.auth import oauth, token_store +from app.model_downloader.auth.providers import Provider +from app.model_downloader.auth.token_store import Token + + +class AuthStore: + def __init__(self) -> None: + # provider name -> Token, or None when known to be absent. A missing key + # means "not yet loaded from disk". + self._cache: dict[str, Token | None] = {} + + def _load(self, name: str) -> Token | None: + if name not in self._cache: + self._cache[name] = token_store.load(name) + return self._cache[name] + + def set_token(self, name: str, token: Token) -> None: + self._cache[name] = token + token_store.save(name, token) + + def clear(self, name: str) -> None: + self._cache[name] = None + token_store.delete(name) + + async def get_valid_token(self, provider: Provider) -> str | None: + """Return a valid access token string for ``provider``, or ``None``. + + Refreshes an expired token when a refresh token is available. + """ + token = self._load(provider.name) + if token is None or not token.access_token: + return None + if token.is_expired(): + if not token.refresh_token: + return None + token = await oauth.refresh_access_token(provider, token) + self.set_token(provider.name, token) + return token.access_token + + async def begin_login(self, provider: Provider) -> str: + """Start a login flow; returns the authorize URL to open in a browser.""" + return await oauth.start_login_flow(provider, self.set_token) + + def status(self, provider: Provider) -> dict: + token = self._load(provider.name) + return { + "provider": provider.name, + "logged_in": token is not None and bool(token.access_token), + "login_in_progress": oauth.login_in_progress(provider.name), + "env_key_present": provider.env_token() is not None, + } + + +AUTH_STORE = AuthStore() diff --git a/app/model_downloader/auth/token_store.py b/app/model_downloader/auth/token_store.py new file mode 100644 index 000000000..11e6f19af --- /dev/null +++ b/app/model_downloader/auth/token_store.py @@ -0,0 +1,72 @@ +"""On-disk OAuth token persistence — one ``0600`` JSON file per provider. + +Tokens live under ``folder_paths.get_system_user_directory("download_auth")``, +never in the SQLite DB. The file is written with ``0600`` so only the owner can +read the refresh/access token at rest. +""" + +from __future__ import annotations + +import json +import os +import time +from dataclasses import asdict, dataclass + +import folder_paths + + +@dataclass +class Token: + access_token: str + refresh_token: str | None = None + # Epoch seconds when the access token expires; 0 means "unknown / no expiry". + expires_at: int = 0 + token_type: str = "Bearer" + scope: str | None = None + + def is_expired(self, skew: int = 60) -> bool: + if not self.expires_at: + return False + return time.time() + skew >= self.expires_at + + +def _auth_dir() -> str: + path = folder_paths.get_system_user_directory("download_auth") + os.makedirs(path, exist_ok=True) + return path + + +def _token_path(provider: str) -> str: + return os.path.join(_auth_dir(), f"{provider}.json") + + +def load(provider: str) -> Token | None: + try: + with open(_token_path(provider), "r", encoding="utf-8") as f: + data = json.load(f) + except FileNotFoundError: + return None + except (ValueError, OSError): + return None + return Token( + access_token=data.get("access_token", ""), + refresh_token=data.get("refresh_token"), + expires_at=int(data.get("expires_at", 0) or 0), + token_type=data.get("token_type", "Bearer"), + scope=data.get("scope"), + ) + + +def save(provider: str, token: Token) -> None: + path = _token_path(provider) + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(asdict(token), f) + os.chmod(path, 0o600) + + +def delete(provider: str) -> None: + try: + os.remove(_token_path(provider)) + except FileNotFoundError: + pass diff --git a/app/model_downloader/constants.py b/app/model_downloader/constants.py index 6430b1288..232f3611f 100644 --- a/app/model_downloader/constants.py +++ b/app/model_downloader/constants.py @@ -11,21 +11,6 @@ stable. The lifecycle is: from __future__ import annotations -# Auth schemes for HostCredential -AUTH_SCHEME_BEARER = "bearer" -AUTH_SCHEME_HEADER = "header" -AUTH_SCHEME_QUERY = "query" -AUTH_SCHEMES = (AUTH_SCHEME_BEARER, AUTH_SCHEME_HEADER, AUTH_SCHEME_QUERY) - -# Hosts for which a bearer token can be sourced from the environment when no -# stored credential matches. Values are the env var names to try, in order. -# Only consulted during auto-resolve for an exact host match over https, so the -# same per-hop boundary rules apply (e.g. the token is dropped on a redirect to -# a CDN host). Kept here so the host->env-var mapping lives in one place. -ENV_TOKEN_HOSTS = { - "huggingface.co": ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"), -} - class DownloadStatus: QUEUED = "queued" diff --git a/app/model_downloader/credentials/resolver.py b/app/model_downloader/credentials/resolver.py deleted file mode 100644 index b85fccc89..000000000 --- a/app/model_downloader/credentials/resolver.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Turn a stored credential into a per-hop request modifier (PRD section 9.4.2). - -The critical rule: a credential is only ever attached when *the current hop's -host* matches a stored credential, and only over https. This is recomputed -from scratch on every redirect hop, so a token bound to ``huggingface.co`` is -silently dropped when the request is redirected to a presigned CDN host — -which is exactly what these hubs expect. -""" - -from __future__ import annotations - -import asyncio -import os -from dataclasses import dataclass, field -from typing import Optional -from urllib.parse import urlencode, urlsplit, urlunsplit - -from app.model_downloader.constants import ( - AUTH_SCHEME_BEARER, - AUTH_SCHEME_HEADER, - AUTH_SCHEME_QUERY, - ENV_TOKEN_HOSTS, -) -from app.model_downloader.credentials.store import normalize_host -from app.model_downloader.database import queries -from app.model_downloader.database.models import HostCredential - - -@dataclass -class RequestAuth: - """How to modify a single request to carry a credential.""" - - headers: dict[str, str] = field(default_factory=dict) - query: dict[str, str] = field(default_factory=dict) - - def apply_to_url(self, url: str) -> str: - if not self.query: - return url - parts = urlsplit(url) - # Append only the credential params, leaving the original query string - # (including any repeated keys and existing encoding) untouched. - creds = urlencode(self.query) - query = f"{parts.query}&{creds}" if parts.query else creds - return urlunsplit(parts._replace(query=query)) - - -def _matches(cred: HostCredential, hop_host: str) -> bool: - cred_host = cred.host - if hop_host == cred_host: - return True - if cred.match_subdomains: - # Label-boundary suffix: api.example.com matches example.com, but - # evil-example.com does NOT. - return hop_host.endswith("." + cred_host) - return False - - -def _build_auth(cred: HostCredential) -> RequestAuth: - if cred.auth_scheme == AUTH_SCHEME_BEARER: - return RequestAuth(headers={"Authorization": f"Bearer {cred.secret}"}) - if cred.auth_scheme == AUTH_SCHEME_HEADER: - name = cred.header_name or "Authorization" - return RequestAuth(headers={name: cred.secret}) - if cred.auth_scheme == AUTH_SCHEME_QUERY and cred.query_param: - return RequestAuth(query={cred.query_param: cred.secret}) - return RequestAuth() - - -def _resolve_sync( - host: str, scheme: str, explicit_credential_id: Optional[str] -) -> Optional[RequestAuth]: - # Never attach a secret over a non-https hop (PRD section 9.4.2). - if scheme.lower() != "https": - return None - hop_host = normalize_host(host) - if not hop_host: - return None - - if explicit_credential_id is not None: - cred = queries.get_credential(explicit_credential_id) - # An explicit credential is still subject to the per-hop host check — - # it is not forced onto a non-matching host. - if cred is None or not cred.enabled or not _matches(cred, hop_host): - return None - return _build_auth(cred) - - # Auto-resolve: exact host first, then any subdomain-matching credential. - cred = queries.get_credential_by_host(hop_host) - if cred is not None and cred.enabled: - return _build_auth(cred) - for sub in queries.list_subdomain_credentials(): - if sub.enabled and _matches(sub, hop_host): - return _build_auth(sub) - - # Env fallback: only for an exact host match, and only after the DB lookups - # miss, so a user-set credential always takes precedence. The token is never - # persisted; it is read fresh from the environment on each hop. - for var in ENV_TOKEN_HOSTS.get(hop_host, ()): - token = os.environ.get(var) - if token: - return RequestAuth(headers={"Authorization": f"Bearer {token}"}) - return None - - -async def resolve_auth_for_hop( - host: str, scheme: str, *, explicit_credential_id: Optional[str] = None -) -> Optional[RequestAuth]: - """Resolve the credential (if any) to attach for one request hop.""" - return await asyncio.to_thread( - _resolve_sync, host, scheme, explicit_credential_id - ) diff --git a/app/model_downloader/credentials/store.py b/app/model_downloader/credentials/store.py deleted file mode 100644 index ac82d24e3..000000000 --- a/app/model_downloader/credentials/store.py +++ /dev/null @@ -1,141 +0,0 @@ -"""The credential store: one API key per host. - -Secrets are write-only over the API — :class:`CredentialView` carries only -masked metadata (``secret_last4`` + scheme + label), never the secret itself. -At-rest protection for v1 is filesystem permissions on the shared DB (the DB -is the trust boundary); encryption-at-rest is a noted future seam. -""" - -from __future__ import annotations - -import asyncio -from dataclasses import dataclass -from typing import Optional -from urllib.parse import urlsplit - -from app.model_downloader.constants import ( - AUTH_SCHEME_BEARER, - AUTH_SCHEME_HEADER, - AUTH_SCHEME_QUERY, - AUTH_SCHEMES, -) -from app.model_downloader.database import queries -from app.model_downloader.database.models import HostCredential - - -def normalize_host(host: str) -> str: - """Lowercase, strip port, IDNA-encode.""" - if not host: - return "" - host = host.strip() - if "://" in host: # a full URL was pasted — extract just the host - host = urlsplit(host).hostname or "" - host = host.lower() - if host.startswith("[") and "]" in host: # bracketed IPv6 literal - host = host[1 : host.index("]")] - elif host.count(":") == 1: # host:port (not IPv6) - host = host.split(":", 1)[0] - try: - host = host.encode("idna").decode("ascii") - except (UnicodeError, ValueError): - pass - return host - - -@dataclass(frozen=True) -class CredentialView: - """Masked, API-safe view of a credential — never includes the secret.""" - - id: str - host: str - auth_scheme: str - header_name: Optional[str] - query_param: Optional[str] - label: Optional[str] - match_subdomains: bool - enabled: bool - secret_last4: Optional[str] - created_at: int - updated_at: int - - -def _to_view(row: HostCredential) -> CredentialView: - return CredentialView( - id=row.id, - host=row.host, - auth_scheme=row.auth_scheme, - header_name=row.header_name, - query_param=row.query_param, - label=row.label, - match_subdomains=row.match_subdomains, - enabled=row.enabled, - secret_last4=row.secret_last4, - created_at=row.created_at, - updated_at=row.updated_at, - ) - - -class CredentialValidationError(ValueError): - """A credential upsert had inconsistent fields.""" - - -class CredentialStore: - """Async facade over the ``host_credentials`` table. - - DB access is synchronous (SQLite) and offloaded via ``asyncio.to_thread``. - """ - - async def upsert( - self, - host: str, - secret: str, - *, - auth_scheme: str = AUTH_SCHEME_BEARER, - header_name: Optional[str] = None, - query_param: Optional[str] = None, - label: Optional[str] = None, - match_subdomains: bool = False, - enabled: bool = True, - ) -> CredentialView: - host = normalize_host(host) - if not host: - raise CredentialValidationError("host is required") - if not secret: - raise CredentialValidationError("secret is required") - if auth_scheme not in AUTH_SCHEMES: - raise CredentialValidationError( - f"auth_scheme must be one of {AUTH_SCHEMES}, got {auth_scheme!r}" - ) - if auth_scheme == AUTH_SCHEME_HEADER and not header_name: - header_name = "Authorization" - if auth_scheme == AUTH_SCHEME_QUERY and not query_param: - raise CredentialValidationError( - "query_param is required when auth_scheme='query'" - ) - values = { - "host": host, - "secret": secret, - "secret_last4": secret[-4:] if len(secret) > 4 else None, - "auth_scheme": auth_scheme, - "header_name": header_name, - "query_param": query_param, - "label": label, - "match_subdomains": match_subdomains, - "enabled": enabled, - } - row = await asyncio.to_thread(queries.upsert_credential, values) - return _to_view(row) - - async def list(self) -> list[CredentialView]: - rows = await asyncio.to_thread(queries.list_credentials) - return [_to_view(r) for r in rows] - - async def get(self, credential_id: str) -> Optional[CredentialView]: - row = await asyncio.to_thread(queries.get_credential, credential_id) - return _to_view(row) if row is not None else None - - async def delete(self, credential_id: str) -> bool: - return await asyncio.to_thread(queries.delete_credential, credential_id) - - -CREDENTIAL_STORE = CredentialStore() diff --git a/app/model_downloader/database/models.py b/app/model_downloader/database/models.py index 546c8ba0c..be6c4d2e5 100644 --- a/app/model_downloader/database/models.py +++ b/app/model_downloader/database/models.py @@ -1,10 +1,9 @@ """SQLAlchemy models for the download manager. -Three tables: +Two tables: - ``downloads`` one row per requested file (job + queue state). - ``download_segments`` per-segment byte progress, for segmented resume. -- ``host_credentials`` one API key per host, reused across downloads. On completion a finished file is registered into the assets catalog; ``downloads`` is kept only as job history. @@ -64,13 +63,6 @@ class Download(Base): # Optional hub-provided checksum to verify against (NOT the dedup key). expected_sha256: Mapped[str | None] = mapped_column(String(64), nullable=True) - # Explicit credential override; otherwise auto-resolved by host. - # RESTRICT keeps a credential from being deleted while a download references it. - credential_id: Mapped[str | None] = mapped_column( - String(36), - ForeignKey("host_credentials.id", ondelete="RESTRICT"), - nullable=True, - ) allow_any_extension: Mapped[bool] = mapped_column( Boolean, nullable=False, default=False ) @@ -91,10 +83,6 @@ class Download(Base): order_by="DownloadSegment.idx", ) - credential: Mapped[HostCredential | None] = relationship( - "HostCredential", back_populates="downloads" - ) - __table_args__ = ( Index("ix_downloads_status", "status"), Index("ix_downloads_priority", "priority"), @@ -135,39 +123,3 @@ class DownloadSegment(Base): f"" ) - - -class HostCredential(Base): - __tablename__ = "host_credentials" - - id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) - # Normalized lowercase hostname, e.g. "civitai.com". - host: Mapped[str] = mapped_column(String(255), nullable=False) - match_subdomains: Mapped[bool] = mapped_column( - Boolean, nullable=False, default=False - ) - label: Mapped[str | None] = mapped_column(String(255), nullable=True) - auth_scheme: Mapped[str] = mapped_column( - String(16), nullable=False, default="bearer" - ) - header_name: Mapped[str | None] = mapped_column(String(255), nullable=True) - query_param: Mapped[str | None] = mapped_column(String(255), nullable=True) - # The API key itself. Write-only over the API; never returned. See PRD 9.4.4. - secret: Mapped[str] = mapped_column(Text, nullable=False) - secret_last4: Mapped[str | None] = mapped_column(String(4), nullable=True) - enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True) - created_at: Mapped[int] = mapped_column(BigInteger, nullable=False, default=_now) - updated_at: Mapped[int] = mapped_column( - BigInteger, nullable=False, default=_now, onupdate=_now - ) - - downloads: Mapped[list[Download]] = relationship( - "Download", back_populates="credential" - ) - - __table_args__ = ( - Index("uq_host_credentials_host", "host", unique=True), - ) - - def __repr__(self) -> str: - return f"" diff --git a/app/model_downloader/database/queries.py b/app/model_downloader/database/queries.py index e4b447841..ebbec4cab 100644 --- a/app/model_downloader/database/queries.py +++ b/app/model_downloader/database/queries.py @@ -11,15 +11,10 @@ import time from typing import Optional from sqlalchemy import delete, select -from sqlalchemy.exc import IntegrityError from app.database.db import create_session from app.model_downloader.constants import DownloadStatus -from app.model_downloader.database.models import ( - Download, - DownloadSegment, - HostCredential, -) +from app.model_downloader.database.models import Download, DownloadSegment # ----- downloads ----- @@ -180,106 +175,3 @@ def reconcile_live_downloads() -> list[Download]: ) session.expunge_all() return resumable - - -# ----- host credentials ----- - - -def get_credential(credential_id: str) -> Optional[HostCredential]: - with create_session() as session: - row = session.get(HostCredential, credential_id) - if row is not None: - session.expunge_all() - return row - - -def get_credential_by_host(host: str) -> Optional[HostCredential]: - with create_session() as session: - row = ( - session.execute( - select(HostCredential).where(HostCredential.host == host).limit(1) - ) - .scalars() - .first() - ) - if row is not None: - session.expunge_all() - return row - - -def list_credentials() -> list[HostCredential]: - with create_session() as session: - rows = list( - session.execute( - select(HostCredential).order_by(HostCredential.host) - ).scalars() - ) - session.expunge_all() - return rows - - -def list_subdomain_credentials() -> list[HostCredential]: - """Credentials that opted into subdomain matching, for suffix checks.""" - with create_session() as session: - rows = list( - session.execute( - select(HostCredential).where(HostCredential.match_subdomains.is_(True)) - ).scalars() - ) - session.expunge_all() - return rows - - -def upsert_credential(values: dict) -> HostCredential: - """Insert or update a credential keyed by ``host``. - - Callers can target the same host concurrently (each runs in its own - short-lived session on a separate connection), so the read-then-write here - can race: two callers both see no existing row and both attempt an insert. - The ``host`` column is uniquely indexed, so the loser's insert raises - ``IntegrityError``. We recover by rolling back and retrying, at which point - the now-committed row is found and updated in place, letting concurrent - calls converge instead of failing or creating duplicates. - """ - host = values["host"] - now = int(time.time()) - last_error: IntegrityError | None = None - for _ in range(2): - with create_session() as session: - row = ( - session.execute( - select(HostCredential).where(HostCredential.host == host).limit(1) - ) - .scalars() - .first() - ) - if row is None: - row = HostCredential(**values) - row.created_at = now - row.updated_at = now - session.add(row) - else: - for key, value in values.items(): - setattr(row, key, value) - row.updated_at = now - try: - session.commit() - except IntegrityError as exc: - session.rollback() - last_error = exc - continue - session.refresh(row) - session.expunge(row) - return row - assert last_error is not None - raise last_error - - -def delete_credential(credential_id: str) -> bool: - with create_session() as session: - row = session.get(HostCredential, credential_id) - if row is None: - return False - session.delete(row) - session.commit() - return True diff --git a/app/model_downloader/engine/job.py b/app/model_downloader/engine/job.py index c0f5e9c2d..fe0794f41 100644 --- a/app/model_downloader/engine/job.py +++ b/app/model_downloader/engine/job.py @@ -105,7 +105,6 @@ class JobSpec: dest_path: str temp_path: str priority: int = 0 - credential_id: Optional[str] = None expected_sha256: Optional[str] = None allow_any_extension: bool = False etag: Optional[str] = None @@ -188,7 +187,7 @@ class DownloadJob: # ----- probe + plan ----- async def _probe_and_plan(self): - pr = await probe(self.spec.url, credential_id=self.spec.credential_id) + pr = await probe(self.spec.url) if not pr.ok: if pr.gated: raise FatalError(gated_error_message(self.spec.url, pr)) @@ -341,7 +340,7 @@ class DownloadJob: if self._etag: headers["If-Range"] = self._etag async with open_validated( - "GET", self.spec.url, credential_id=self.spec.credential_id, headers=headers + "GET", self.spec.url, headers=headers ) as (resp, _final): if resp.status == 200: # Server ignored the range -> remote changed / no resume support. @@ -383,7 +382,7 @@ class DownloadJob: if self._etag: headers["If-Range"] = self._etag async with open_validated( - "GET", self.spec.url, credential_id=self.spec.credential_id, headers=headers + "GET", self.spec.url, headers=headers ) as (resp, _final): if offset > 0 and resp.status == 200: # Resume not honoured -> start over from the beginning. Truncate @@ -432,8 +431,8 @@ class DownloadJob: def _raise_for_status(self, status: int) -> None: if status in (401, 403): raise FatalError( - f"{redact_url(self.spec.url)} returned {status}; add/update an API key for " - f"this host at /api/download/credentials." + f"{redact_url(self.spec.url)} returned {status}; authenticate this " + f"host via /api/download/auth or set its API key env var." ) if status in _RETRYABLE_STATUSES: raise RetryableError(f"HTTP {status}") diff --git a/app/model_downloader/manager.py b/app/model_downloader/manager.py index 46d663235..ae1c98c2b 100644 --- a/app/model_downloader/manager.py +++ b/app/model_downloader/manager.py @@ -75,7 +75,6 @@ class DownloadManager: priority: int = 0, expected_sha256: Optional[str] = None, allow_any_extension: bool = False, - credential_id: Optional[str] = None, ) -> str: # Coarse gate first: host/scheme must be allowlisted, and any extension # present in the URL path must be a known model type. A URL whose path @@ -98,7 +97,7 @@ class DownloadManager: # adopt the real extension from the response, forcing the stored # filename to match. Skipped when the caller opted into any extension. if not allow_any_extension and url_path_extension(url) == "": - resolved_ext = await self._resolve_extension(url, credential_id) + resolved_ext = await self._resolve_extension(url) model_id = paths.apply_extension(model_id, resolved_ext) try: @@ -137,7 +136,6 @@ class DownloadManager: "status": DownloadStatus.QUEUED, "priority": priority, "expected_sha256": expected_sha256, - "credential_id": credential_id, "allow_any_extension": allow_any_extension, }, ) @@ -145,18 +143,16 @@ class DownloadManager: await self._scheduler.pump() return download_id - async def _resolve_extension( - self, url: str, credential_id: Optional[str] - ) -> str: + async def _resolve_extension(self, url: str) -> str: """Follow ``url`` to its final response and return the real extension. Used for allowlisted URLs whose path has no extension (e.g. Civitai download endpoints): the filename lives in the ``Content-Disposition`` header or the post-redirect URL. Raises :class:`DownloadError` when the - URL can't be resolved, needs credentials, or resolves to something that - is not a known model file — so we never persist a bogus destination. + URL can't be resolved, needs authentication, or resolves to something + that is not a known model file — so we never persist a bogus destination. """ - pr = await probe(url, credential_id=credential_id) + pr = await probe(url) if not pr.ok: if pr.gated: raise DownloadError( diff --git a/app/model_downloader/net/http.py b/app/model_downloader/net/http.py index 8112b7a20..b4df0fc29 100644 --- a/app/model_downloader/net/http.py +++ b/app/model_downloader/net/http.py @@ -2,9 +2,9 @@ Automatic redirects are disabled. We follow hops ourselves so that on *every* hop we (a) re-validate scheme + reject credentials-in-URL, -(b) recompute which stored credential — if any — applies to that hop's host, -and (c) let the connector's resolver screen the IP. This is the single place -that attaches credentials, so a token can never ride a redirect to a CDN host. +(b) recompute which auth — if any — applies to that hop's host, and (c) let the +connector's resolver screen the IP. This is the single place that attaches a +token, so it can never ride a redirect to a CDN host. """ from __future__ import annotations @@ -17,7 +17,7 @@ from urllib.parse import unquote, urljoin, urlsplit, urlunsplit import aiohttp -from app.model_downloader.credentials.resolver import resolve_auth_for_hop +from app.model_downloader.auth.resolver import resolve_auth_for_hop from app.model_downloader.net.session import get_session from app.model_downloader.security.ssrf import ( MAX_REDIRECTS, @@ -78,7 +78,6 @@ def filename_from_content_disposition(value: Optional[str]) -> Optional[str]: async def _resolve_final_response( method: str, url: str, - credential_id: Optional[str], base_headers: dict[str, str], timeout: aiohttp.ClientTimeout, ) -> tuple[aiohttp.ClientResponse, str]: @@ -93,18 +92,14 @@ async def _resolve_final_response( while True: check_redirect_hop(current, is_initial_url=(hops == 0)) parts = urlsplit(current) - auth = await resolve_auth_for_hop( - parts.hostname or "", parts.scheme, explicit_credential_id=credential_id - ) + auth = await resolve_auth_for_hop(parts.hostname or "", parts.scheme) req_headers = dict(base_headers) - req_url = current if auth is not None: req_headers.update(auth.headers) - req_url = auth.apply_to_url(current) resp = await session.request( method, - req_url, + current, allow_redirects=False, headers=req_headers, timeout=timeout, @@ -127,7 +122,6 @@ async def open_validated( method: str, url: str, *, - credential_id: Optional[str] = None, headers: Optional[dict[str, str]] = None, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT, ) -> AsyncIterator[tuple[aiohttp.ClientResponse, str]]: @@ -137,7 +131,7 @@ async def open_validated( query string. The response is released automatically on exit. """ resp, final_url = await _resolve_final_response( - method, url, credential_id, dict(headers or {}), timeout + method, url, dict(headers or {}), timeout ) try: yield resp, final_url diff --git a/app/model_downloader/net/probe.py b/app/model_downloader/net/probe.py index eca0c7fbb..d396c0cc9 100644 --- a/app/model_downloader/net/probe.py +++ b/app/model_downloader/net/probe.py @@ -78,13 +78,12 @@ def _filename_from_response( return None -async def probe(url: str, *, credential_id: Optional[str] = None) -> ProbeResult: +async def probe(url: str) -> ProbeResult: """Probe ``url`` and return discovered metadata, failing soft.""" try: async with open_validated( "GET", url, - credential_id=credential_id, headers={"Range": "bytes=0-0", "Accept-Encoding": "identity"}, timeout=_PROBE_TIMEOUT, ) as (resp, final_url): @@ -149,9 +148,10 @@ def gated_error_message(url: str, pr: ProbeResult) -> str: detail += "." return ( f"{redacted} is a gated model — {detail} Request access on the model's " - f"page, add an API key for this host at /api/download/credentials, and retry." + f"page, authenticate this host via /api/download/auth (or set its API " + f"key env var), and retry." ) return ( - f"{redacted} requires authentication. Add an API key for this host at " - f"/api/download/credentials and retry." + f"{redacted} requires authentication. Authenticate this host via " + f"/api/download/auth or set its API key env var, and retry." ) diff --git a/app/model_downloader/scheduler.py b/app/model_downloader/scheduler.py index b41d2cb3b..f61ae184d 100644 --- a/app/model_downloader/scheduler.py +++ b/app/model_downloader/scheduler.py @@ -124,7 +124,6 @@ class Scheduler: dest_path=row.dest_path, temp_path=row.temp_path, priority=row.priority, - credential_id=row.credential_id, expected_sha256=row.expected_sha256, allow_any_extension=row.allow_any_extension, etag=row.etag, diff --git a/openapi.yaml b/openapi.yaml index 1bd0edc43..47b23d804 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -229,10 +229,6 @@ components: default: false description: Permit a non-model file extension (default only allows known model extensions). type: boolean - credential_id: - description: Explicit per-host credential to use; otherwise auto-resolved by host. Still subject to the per-hop host match. - nullable: true - type: string expected_sha256: description: Optional hub-provided SHA256 to verify the completed file against (fail-closed). nullable: true @@ -590,77 +586,26 @@ components: required: - history type: object - HostCredentialUpsert: - description: Request body for upserting a per-host credential. The secret is write-only. + DownloadAuthProvider: + description: Per-provider download-auth status. Never includes a token. properties: - auth_scheme: - default: bearer - description: How the secret is attached to requests. - enum: - - bearer - - header - - query - type: string - enabled: - default: true + env_key_present: + description: Whether an API key for this provider is set via an environment variable. type: boolean - header_name: - description: Header name when auth_scheme=header (defaults to Authorization). - nullable: true - type: string - host: - description: Normalized hostname the key applies to (e.g. "civitai.com"). - type: string - label: - description: User-friendly name for display. - nullable: true - type: string - match_subdomains: - default: false - description: Also match label-boundary subdomains of host (off by default; unsafe for hub CDNs). + logged_in: + description: Whether a stored OAuth token exists for this provider. type: boolean - query_param: - description: Query parameter name when auth_scheme=query. - nullable: true - type: string - secret: - description: The API key. Write-only — never returned by any endpoint. + login_in_progress: + description: Whether an OAuth login flow is currently awaiting the browser callback. + type: boolean + provider: + description: Provider name (e.g. "huggingface", "civitai"). type: string required: - - host - - secret - type: object - HostCredentialView: - description: Masked, API-safe view of a stored credential. Never includes the secret. - properties: - auth_scheme: - type: string - created_at: - type: integer - enabled: - type: boolean - header_name: - nullable: true - type: string - host: - type: string - id: - format: uuid - type: string - label: - nullable: true - type: string - match_subdomains: - type: boolean - query_param: - nullable: true - type: string - secret_last4: - description: Last 4 characters of the secret, for masked display only. - nullable: true - type: string - updated_at: - type: integer + - provider + - logged_in + - login_in_progress + - env_key_present type: object JobCancelResponse: description: Response for POST /api/jobs/{job_id}/cancel. Returned on both fresh cancels and idempotent no-ops. @@ -2582,58 +2527,69 @@ paths: summary: Clear terminal downloads from history tags: - download - /api/download/credentials: + /api/download/auth: get: - description: List stored per-host credentials. Secrets are never returned; only masked metadata (last 4 chars, scheme, label). - operationId: listDownloadCredentials + description: Per-provider download-auth status (OAuth login state + env-key presence). Never returns a token. + operationId: getDownloadAuth responses: "200": content: application/json: schema: properties: - credentials: + providers: items: - $ref: '#/components/schemas/HostCredentialView' + $ref: '#/components/schemas/DownloadAuthProvider' type: array type: object - description: Masked credential list - summary: List host credentials (masked) + description: Per-provider auth status + summary: Get download auth status tags: - download + /api/download/auth/{provider}/login: post: description: | - Upsert (by host) a per-host API key used to authenticate downloads. - The secret is write-only: it is stored once here and never returned by any endpoint. - operationId: upsertDownloadCredential - requestBody: - content: - application/json: - schema: - $ref: '#/components/schemas/HostCredentialUpsert' + Start an OAuth 2.0 PKCE login for a provider. Spawns a transient loopback + callback server and returns the authorize URL to open in a browser. + operationId: loginDownloadAuth + parameters: + - in: path + name: provider + required: true + schema: + type: string responses: - "201": + "200": content: application/json: schema: - $ref: '#/components/schemas/HostCredentialView' - description: Credential stored (masked view returned) + properties: + authorize_url: + type: string + type: object + description: Authorize URL to open in a browser "400": content: application/json: schema: $ref: '#/components/schemas/ErrorResponse' - description: Invalid credential - summary: Upsert a host credential + description: Unknown provider or OAuth app not configured + "409": + content: + application/json: + schema: + $ref: '#/components/schemas/ErrorResponse' + description: A login for this provider is already in progress + summary: Start OAuth login for a provider tags: - download - /api/download/credentials/{id}: - delete: - description: Delete a stored host credential. - operationId: deleteDownloadCredential + /api/download/auth/{provider}/logout: + post: + description: Clear the stored OAuth token (memory + on-disk file) for a provider. + operationId: logoutDownloadAuth parameters: - in: path - name: id + name: provider required: true schema: type: string @@ -2643,42 +2599,17 @@ paths: application/json: schema: properties: - deleted: + logged_out: type: boolean type: object - description: Deleted - "404": + description: Token cleared + "400": content: application/json: schema: $ref: '#/components/schemas/ErrorResponse' - description: No such credential - summary: Delete a host credential - tags: - - download - get: - description: Get a single host credential (masked; never includes the secret). - operationId: getDownloadCredential - parameters: - - in: path - name: id - required: true - schema: - type: string - responses: - "200": - content: - application/json: - schema: - $ref: '#/components/schemas/HostCredentialView' - description: Masked credential - "404": - content: - application/json: - schema: - $ref: '#/components/schemas/ErrorResponse' - description: No such credential - summary: Get a host credential (masked) + description: Unknown provider + summary: Log out a provider tags: - download /api/download/enqueue: diff --git a/tests-unit/model_downloader_test/test_auth.py b/tests-unit/model_downloader_test/test_auth.py new file mode 100644 index 000000000..3a0b69d94 --- /dev/null +++ b/tests-unit/model_downloader_test/test_auth.py @@ -0,0 +1,288 @@ +"""Unit tests for download authentication. + +Covers env-key resolution, OAuth token resolution/refresh/expiry, provider +host matching (including the per-hop drop on a CDN host), and the auth routes. +Async tests are driven via ``asyncio.run`` so no pytest-asyncio plugin is needed. +""" + +from __future__ import annotations + +import asyncio +import json +import time +from urllib.parse import parse_qs, urlsplit + +import pytest +from aiohttp.test_utils import make_mocked_request + +from app.model_downloader.api import routes +from app.model_downloader.auth import oauth, token_store +from app.model_downloader.auth.providers import PROVIDERS, provider_for_host +from app.model_downloader.auth.resolver import resolve_auth_for_hop +from app.model_downloader.auth.store import AUTH_STORE +from app.model_downloader.auth.token_store import Token + +_HF_ENV = ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN") +_CIVITAI_ENV = ("CIVITAI_API_TOKEN", "CIVITAI_API_KEY") + + +@pytest.fixture +def auth_tmp(monkeypatch, tmp_path): + """Isolate the on-disk token store and clear the in-memory cache.""" + d = tmp_path / "download_auth" + d.mkdir() + monkeypatch.setattr(token_store, "_auth_dir", lambda: str(d)) + AUTH_STORE._cache.clear() + yield + AUTH_STORE._cache.clear() + + +def _clear_env(monkeypatch, *names): + for name in names: + monkeypatch.delenv(name, raising=False) + + +# ----- provider host matching ----- + + +def test_provider_for_host(): + assert provider_for_host("HuggingFace.co:443").name == "huggingface" + assert provider_for_host("civitai.com").name == "civitai" + # sibling CDN hosts must not match — this is what drops the token on redirect + assert provider_for_host("cdn-lfs.huggingface.co") is None + assert provider_for_host("cas-bridge.xethub.hf.co") is None + assert provider_for_host("example.com") is None + + +# ----- env-key resolution ----- + + +def test_env_key_resolution_hf(monkeypatch, auth_tmp): + monkeypatch.setenv("HF_TOKEN", "hf_env") + + async def _run(): + auth = await resolve_auth_for_hop("huggingface.co", "https") + assert auth is not None + assert auth.headers["Authorization"] == "Bearer hf_env" + # never over http, never on a CDN redirect host + assert await resolve_auth_for_hop("huggingface.co", "http") is None + assert await resolve_auth_for_hop("cdn-lfs.huggingface.co", "https") is None + + asyncio.run(_run()) + + +def test_env_key_resolution_civitai_secondary_var(monkeypatch, auth_tmp): + _clear_env(monkeypatch, "CIVITAI_API_TOKEN") + monkeypatch.setenv("CIVITAI_API_KEY", "civ_env") + + async def _run(): + auth = await resolve_auth_for_hop("civitai.com", "https") + assert auth is not None + assert auth.headers["Authorization"] == "Bearer civ_env" + + asyncio.run(_run()) + + +def test_env_key_takes_precedence_over_oauth(monkeypatch, auth_tmp): + monkeypatch.setenv("HF_TOKEN", "hf_env") + + async def _run(): + AUTH_STORE.set_token("huggingface", Token(access_token="oauth_acc")) + auth = await resolve_auth_for_hop("huggingface.co", "https") + assert auth.headers["Authorization"] == "Bearer hf_env" + + asyncio.run(_run()) + + +# ----- OAuth token resolution / refresh / expiry ----- + + +def test_oauth_token_resolution(monkeypatch, auth_tmp): + _clear_env(monkeypatch, *_HF_ENV) + + async def _run(): + AUTH_STORE.set_token("huggingface", Token(access_token="acc", expires_at=0)) + auth = await resolve_auth_for_hop("huggingface.co", "https") + assert auth is not None + assert auth.headers["Authorization"] == "Bearer acc" + + asyncio.run(_run()) + + +def test_oauth_refresh_on_expiry(monkeypatch, auth_tmp): + _clear_env(monkeypatch, *_HF_ENV) + + async def fake_refresh(provider, tok): + return Token( + access_token="new_acc", + refresh_token="r2", + expires_at=int(time.time()) + 3600, + ) + + monkeypatch.setattr(oauth, "refresh_access_token", fake_refresh) + + async def _run(): + AUTH_STORE.set_token( + "huggingface", + Token(access_token="old", refresh_token="r1", expires_at=1), + ) + access = await AUTH_STORE.get_valid_token(PROVIDERS["huggingface"]) + assert access == "new_acc" + # the refreshed token is persisted (cache + disk) + assert token_store.load("huggingface").access_token == "new_acc" + + asyncio.run(_run()) + + +def test_oauth_expired_without_refresh_returns_none(monkeypatch, auth_tmp): + _clear_env(monkeypatch, *_HF_ENV) + + async def _run(): + AUTH_STORE.set_token( + "huggingface", + Token(access_token="old", refresh_token=None, expires_at=1), + ) + assert await AUTH_STORE.get_valid_token(PROVIDERS["huggingface"]) is None + assert await resolve_auth_for_hop("huggingface.co", "https") is None + + asyncio.run(_run()) + + +def test_no_auth_when_nothing_configured(monkeypatch, auth_tmp): + _clear_env(monkeypatch, *_HF_ENV, *_CIVITAI_ENV) + + async def _run(): + assert await resolve_auth_for_hop("huggingface.co", "https") is None + assert await resolve_auth_for_hop("example.com", "https") is None + + asyncio.run(_run()) + + +# ----- auth routes ----- + + +def test_auth_status_route(monkeypatch, auth_tmp): + _clear_env(monkeypatch, *_HF_ENV, *_CIVITAI_ENV) + + async def _run(): + resp = await routes.auth_status(make_mocked_request("GET", "/api/download/auth")) + data = json.loads(resp.body) + by_name = {p["provider"]: p for p in data["providers"]} + assert set(by_name) == {"huggingface", "civitai"} + assert by_name["huggingface"]["logged_in"] is False + assert by_name["huggingface"]["env_key_present"] is False + + asyncio.run(_run()) + + +def test_login_unconfigured_returns_400(monkeypatch, auth_tmp): + _clear_env(monkeypatch, "COMFY_HF_OAUTH_CLIENT_ID") + + async def _run(): + req = make_mocked_request( + "POST", "/api/download/auth/huggingface/login", + match_info={"provider": "huggingface"}, + ) + resp = await routes.auth_login(req) + assert resp.status == 400 + assert json.loads(resp.body)["error"]["code"] == "OAUTH_NOT_CONFIGURED" + + asyncio.run(_run()) + + +def test_login_unknown_provider_returns_400(auth_tmp): + async def _run(): + req = make_mocked_request( + "POST", "/api/download/auth/nope/login", + match_info={"provider": "nope"}, + ) + resp = await routes.auth_login(req) + assert resp.status == 400 + assert json.loads(resp.body)["error"]["code"] == "UNKNOWN_PROVIDER" + + asyncio.run(_run()) + + +def test_login_start_and_in_progress(monkeypatch, auth_tmp): + monkeypatch.setenv("COMFY_HF_OAUTH_CLIENT_ID", "test-client") + + async def _run(): + try: + req = make_mocked_request( + "POST", "/api/download/auth/huggingface/login", + match_info={"provider": "huggingface"}, + ) + resp = await routes.auth_login(req) + assert resp.status == 200 + url = json.loads(resp.body)["authorize_url"] + assert url.startswith("https://huggingface.co/oauth/authorize?") + assert "code_challenge=" in url and "code_challenge_method=S256" in url + assert "client_id=test-client" in url + # a second concurrent login is rejected + resp2 = await routes.auth_login( + make_mocked_request( + "POST", "/api/download/auth/huggingface/login", + match_info={"provider": "huggingface"}, + ) + ) + assert resp2.status == 409 + finally: + flow = oauth._ACTIVE.get("huggingface") + if flow is not None: + await flow._teardown() + + asyncio.run(_run()) + + +def _redirect_port(authorize_url: str) -> int: + redirect = parse_qs(urlsplit(authorize_url).query)["redirect_uri"][0] + return urlsplit(redirect).port + + +def test_concurrent_logins_different_providers(monkeypatch, auth_tmp): + """Two providers can log in at once — each binds its own loopback port.""" + monkeypatch.setenv("COMFY_HF_OAUTH_CLIENT_ID", "hf-client") + monkeypatch.setenv("COMFY_CIVITAI_OAUTH_CLIENT_ID", "civ-client") + + async def _run(): + try: + hf = await routes.auth_login( + make_mocked_request( + "POST", "/api/download/auth/huggingface/login", + match_info={"provider": "huggingface"}, + ) + ) + civ = await routes.auth_login( + make_mocked_request( + "POST", "/api/download/auth/civitai/login", + match_info={"provider": "civitai"}, + ) + ) + assert hf.status == 200 and civ.status == 200 + hf_port = _redirect_port(json.loads(hf.body)["authorize_url"]) + civ_port = _redirect_port(json.loads(civ.body)["authorize_url"]) + # distinct, OS-assigned loopback ports (no fixed-port collision) + assert hf_port and civ_port and hf_port != civ_port + finally: + for name in ("huggingface", "civitai"): + flow = oauth._ACTIVE.get(name) + if flow is not None: + await flow._teardown() + + asyncio.run(_run()) + + +def test_logout_route_clears_token(auth_tmp): + async def _run(): + AUTH_STORE.set_token("civitai", Token(access_token="x")) + assert AUTH_STORE.status(PROVIDERS["civitai"])["logged_in"] is True + req = make_mocked_request( + "POST", "/api/download/auth/civitai/logout", + match_info={"provider": "civitai"}, + ) + resp = await routes.auth_logout(req) + assert resp.status == 200 + assert json.loads(resp.body)["logged_out"] is True + assert AUTH_STORE.status(PROVIDERS["civitai"])["logged_in"] is False + + asyncio.run(_run()) diff --git a/tests-unit/model_downloader_test/test_credentials.py b/tests-unit/model_downloader_test/test_credentials.py deleted file mode 100644 index a630ad751..000000000 --- a/tests-unit/model_downloader_test/test_credentials.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Unit tests for the credential store and the per-hop credential resolver. - -Covers the critical rule: a secret is only ever attached when the current -hop's host matches a stored credential, and never over a non-https hop. -""" - -from __future__ import annotations - -import asyncio - -import pytest - -from app.model_downloader.credentials import resolver -from app.model_downloader.credentials.store import ( - CREDENTIAL_STORE, - CredentialValidationError, - normalize_host, -) -from app.model_downloader.database.models import HostCredential - - -# ----- pure host normalization + matching ----- - - -@pytest.mark.parametrize( - "raw,expected", - [ - ("Civitai.com", "civitai.com"), - ("HuggingFace.co:443", "huggingface.co"), - (" Example.COM ", "example.com"), - ], -) -def test_normalize_host(raw, expected): - assert normalize_host(raw) == expected - - -def _cred(**kw) -> HostCredential: - base = dict( - id="x", host="civitai.com", match_subdomains=False, auth_scheme="bearer", - secret="SECRET", enabled=True, - ) - base.update(kw) - return HostCredential(**base) - - -def test_matches_exact_only_by_default(): - c = _cred(host="civitai.com") - assert resolver._matches(c, "civitai.com") is True - assert resolver._matches(c, "api.civitai.com") is False - assert resolver._matches(c, "evil-civitai.com") is False - - -def test_matches_subdomain_label_boundary(): - c = _cred(host="example.com", match_subdomains=True) - assert resolver._matches(c, "api.example.com") is True - assert resolver._matches(c, "example.com") is True - # not a label boundary -> no match - assert resolver._matches(c, "evil-example.com") is False - - -def test_build_auth_shapes(): - assert resolver._build_auth(_cred(auth_scheme="bearer")).headers == { - "Authorization": "Bearer SECRET" - } - assert resolver._build_auth( - _cred(auth_scheme="header", header_name="X-Api-Key") - ).headers == {"X-Api-Key": "SECRET"} - q = resolver._build_auth(_cred(auth_scheme="query", query_param="token")) - assert q.query == {"token": "SECRET"} - assert q.apply_to_url("https://civitai.com/x") == "https://civitai.com/x?token=SECRET" - - -# ----- DB-backed store + resolver ----- - - -def test_store_upsert_is_write_only_and_masked(): - async def _run(): - view = await CREDENTIAL_STORE.upsert("civitai.com", "abcd1234", label="my key") - # The view never carries the secret, only the last 4. - assert not hasattr(view, "secret") - assert view.secret_last4 == "1234" - assert view.host == "civitai.com" - listed = await CREDENTIAL_STORE.list() - assert any(v.host == "civitai.com" for v in listed) - await CREDENTIAL_STORE.delete(view.id) - asyncio.run(_run()) - - -def test_query_scheme_requires_param(): - async def _run(): - with pytest.raises(CredentialValidationError): - await CREDENTIAL_STORE.upsert("civitai.com", "k", auth_scheme="query") - asyncio.run(_run()) - - -def test_resolver_never_crosses_host_boundary(): - async def _run(): - view = await CREDENTIAL_STORE.upsert("huggingface.co", "hf_secret_key") - try: - # matching host over https -> attached - auth = await resolver.resolve_auth_for_hop("huggingface.co", "https") - assert auth is not None - assert auth.headers["Authorization"] == "Bearer hf_secret_key" - # CDN redirect host -> dropped - assert await resolver.resolve_auth_for_hop("cdn-lfs.huggingface.co", "https") is None - # non-https hop -> never attached - assert await resolver.resolve_auth_for_hop("huggingface.co", "http") is None - finally: - await CREDENTIAL_STORE.delete(view.id) - asyncio.run(_run()) - - -# ----- env-based HF token fallback ----- - - -def test_env_token_fallback_attaches_when_no_db_credential(monkeypatch): - monkeypatch.setenv("HF_TOKEN", "env_hf_token") - - async def _run(): - # exact host over https -> env token attached - auth = await resolver.resolve_auth_for_hop("huggingface.co", "https") - assert auth is not None - assert auth.headers["Authorization"] == "Bearer env_hf_token" - # non-https hop -> never attached - assert await resolver.resolve_auth_for_hop("huggingface.co", "http") is None - # CDN redirect host -> dropped (exact-host only) - assert await resolver.resolve_auth_for_hop("cdn-lfs.huggingface.co", "https") is None - asyncio.run(_run()) - - -def test_env_token_secondary_var_is_honored(monkeypatch): - monkeypatch.delenv("HF_TOKEN", raising=False) - monkeypatch.setenv("HUGGING_FACE_HUB_TOKEN", "env_hub_token") - - async def _run(): - auth = await resolver.resolve_auth_for_hop("huggingface.co", "https") - assert auth is not None - assert auth.headers["Authorization"] == "Bearer env_hub_token" - asyncio.run(_run()) - - -def test_db_credential_takes_precedence_over_env(monkeypatch): - monkeypatch.setenv("HF_TOKEN", "env_hf_token") - - async def _run(): - view = await CREDENTIAL_STORE.upsert("huggingface.co", "db_secret_key") - try: - auth = await resolver.resolve_auth_for_hop("huggingface.co", "https") - assert auth is not None - assert auth.headers["Authorization"] == "Bearer db_secret_key" - finally: - await CREDENTIAL_STORE.delete(view.id) - asyncio.run(_run()) - - -def test_env_token_does_not_leak_into_explicit_path(monkeypatch): - monkeypatch.setenv("HF_TOKEN", "env_hf_token") - - async def _run(): - # An explicit credential id that doesn't resolve must stay None; the env - # fallback only applies to the auto-resolve branch. - auth = await resolver.resolve_auth_for_hop( - "huggingface.co", "https", explicit_credential_id="does-not-exist" - ) - assert auth is None - asyncio.run(_run())