wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-14 10:27:17 +08:00
Code Issues Actions 18 Packages Projects Releases Wiki Activity

Remove api creds and add oauth.

Browse Source
This commit is contained in:
Talmaj Marinc 2026-07-12 21:56:15 +02:00
parent ee06f45ff3
commit 71cf5a11f1
22 changed files with 978 additions and 833 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
alembic_db/versions/0006_drop_download_credentials.py Normal file
View File

@ -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))

84
app/model_downloader/api/routes.py
View File

@ -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) -----

15
app/model_downloader/api/schemas_in.py
View File

@ -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",
]

23
app/model_downloader/api/schemas_out.py
View File

@ -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()]

248
app/model_downloader/auth/oauth.py Normal file
View File

@ -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:<port>`` at ``/callback/<provider>``, 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

87
app/model_downloader/auth/providers.py Normal file
View File

@ -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))

42
app/model_downloader/auth/resolver.py Normal file
View File

@ -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

64
app/model_downloader/auth/store.py Normal file
View File

@ -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()

72
app/model_downloader/auth/token_store.py Normal file
View File

@ -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

15
app/model_downloader/constants.py
View File

@ -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"

111
app/model_downloader/credentials/resolver.py
View File

@ -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
)

141
app/model_downloader/credentials/store.py
View File

@ -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()

50
app/model_downloader/database/models.py
View File

@ -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"<DownloadSegment {self.download_id}#{self.idx} "
f"{self.start_offset}-{self.end_offset} done={self.bytes_done}>"
)
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"<HostCredential id={self.id} host={self.host!r} scheme={self.auth_scheme}>"

110
app/model_downloader/database/queries.py
View File

@ -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

11
app/model_downloader/engine/job.py
View File

@ -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}")

14
app/model_downloader/manager.py
View File

@ -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(

20
app/model_downloader/net/http.py
View File

@ -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

10
app/model_downloader/net/probe.py
View File

@ -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."
)

1
app/model_downloader/scheduler.py
View File

@ -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,

181
openapi.yaml
View File

@ -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:

288
tests-unit/model_downloader_test/test_auth.py Normal file
View File

@ -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())

166
tests-unit/model_downloader_test/test_credentials.py
View File

@ -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())