server/oauth.py

"""Sign in with Hugging Face, so inference is billed to the viewer.

A Space with ``hf_oauth: true`` in its README header is given
``OAUTH_CLIENT_ID``, ``OAUTH_CLIENT_SECRET``, ``OAUTH_SCOPES`` and
``OPENID_PROVIDER_URL`` in its environment. Gradio implements this flow for
you; a Docker Space has to do it itself, which is this file.

The scope that matters is ``inference-api``: the access token it yields can
call the router, and the *viewer's* account is billed rather than the Space
owner's. Without it a sign-in tells us who someone is and nothing else, which
is the wrong trade for the amount of machinery involved.

The token is held in a signed, http-only session cookie and never leaves the
server. Nothing here is reachable off a Space, because no client id exists.
"""

from __future__ import annotations

import os
import secrets
from urllib.parse import urlencode

import httpx
from fastapi import APIRouter, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse

router = APIRouter()

# The paths a Space's OAuth app is allowed to redirect back to are fixed by the
# platform. `/auth/callback` is the one this app registers.
CALLBACK_PATH = "/auth/callback"


def provider() -> str:
    return os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co").rstrip("/")


def client_id() -> str | None:
    return os.environ.get("OAUTH_CLIENT_ID")


def configured() -> bool:
    """True when the platform gave us an OAuth app to use."""
    return bool(client_id() and os.environ.get("OAUTH_CLIENT_SECRET"))


def scopes() -> str:
    # `openid profile` identifies the viewer; `inference-api` is what lets us
    # spend their quota instead of ours.
    return os.environ.get("OAUTH_SCOPES") or "openid profile inference-api"


def _redirect_uri(request: Request) -> str:
    space = os.environ.get("SPACE_HOST")
    if space:
        return f"https://{space}{CALLBACK_PATH}"
    return str(request.base_url).rstrip("/") + CALLBACK_PATH


@router.get("/api/auth/login")
async def login(request: Request):
    if not configured():
        return JSONResponse(
            {"error": "Sign-in is only available on the deployed Space."},
            status_code=503,
        )
    # CSRF: the callback must present the state we generated for this session.
    state = secrets.token_urlsafe(24)
    request.session["oauth_state"] = state
    query = urlencode({
        "client_id": client_id(),
        "redirect_uri": _redirect_uri(request),
        "response_type": "code",
        "scope": scopes(),
        "state": state,
        "prompt": "consent",
    })
    return RedirectResponse(f"{provider()}/oauth/authorize?{query}", status_code=302)


@router.get(CALLBACK_PATH)
async def callback(request: Request, code: str | None = None,
                   state: str | None = None, error: str | None = None):
    if error:
        return _close_window(f"Hugging Face returned: {error}")

    expected = request.session.pop("oauth_state", None)
    if not expected:
        # We issued a state and it did not come back, which is a cookie
        # problem rather than a forgery: the session is gone. Say which,
        # because the two have completely different fixes.
        return _close_window(
            "The sign-in session was not returned by your browser. This is "
            "usually third-party cookies being blocked for the embedded "
            "Space \u2014 open the Space in its own tab and try again."
        )
    # Constant time, and both halves required: a callback whose state we did
    # not issue is not ours to trust.
    if not state or not secrets.compare_digest(state, expected):
        return _close_window("Sign-in state did not match. Please try again.")
    if not code:
        return _close_window("Hugging Face did not return an authorization code.")

    try:
        async with httpx.AsyncClient(timeout=20.0) as client:
            tok = await client.post(
                f"{provider()}/oauth/token",
                data={
                    "client_id": client_id(),
                    "client_secret": os.environ.get("OAUTH_CLIENT_SECRET"),
                    "grant_type": "authorization_code",
                    "code": code,
                    "redirect_uri": _redirect_uri(request),
                },
                headers={"Accept": "application/json"},
            )
            if tok.status_code >= 400:
                return _close_window(f"Token exchange failed ({tok.status_code}).")
            payload = tok.json()
            access = payload.get("access_token")
            if not access:
                return _close_window("Hugging Face did not return an access token.")

            who = await client.get(
                f"{provider()}/oauth/userinfo",
                headers={"Authorization": f"Bearer {access}"},
            )
            info = who.json() if who.status_code < 400 else {}
    except httpx.HTTPError as e:
        return _close_window(f"Could not reach Hugging Face: {type(e).__name__}")

    request.session["oauth_token"] = access
    request.session["oauth_user"] = info.get("preferred_username") or info.get("name")
    request.session["oauth_profile"] = {
        "name": info.get("name"),
        "username": info.get("preferred_username"),
        "picture": info.get("picture"),
        # `isPro` and org membership come back on the userinfo for HF.
        "isPro": bool(info.get("isPro")),
        "orgs": [o.get("preferred_username") or o.get("name")
                 for o in (info.get("orgs") or []) if isinstance(o, dict)],
        "grantedScopes": scopes(),
    }
    return _close_window(None)


@router.post("/api/auth/logout")
async def logout(request: Request):
    for key in ("oauth_token", "oauth_user", "oauth_profile", "oauth_state"):
        request.session.pop(key, None)
    return {"ok": True}


def _close_window(error: str | None) -> HTMLResponse:
    """Land the browser back on the app.

    Sign-in is started in the same tab, so the callback simply navigates home;
    an error is carried in the fragment, which never reaches a server log the
    way a query string would.
    """
    target = "/" if not error else "/#auth-error=" + error.replace('"', "")
    return HTMLResponse(
        f'<!doctype html><meta charset="utf-8">'
        f'<meta http-equiv="refresh" content="0;url={target}">'
        f'<body style="background:#1a1c1f;color:#e8e2d6;font:13px ui-monospace,monospace">'
        f'<p style="padding:2rem">{error or "Signed in. Returning…"}</p>'
    )