distinct_server/auth_routes.py

"""The sign-in routes a self-hosted server needs, mounted on Gradio's app.

**Why these exist at all.** On a Space, Hugging Face injects an OAuth app and
Gradio's own `/login/huggingface` and `/login/callback` are real, so
`gr.OAuthProfile` arrives filled in and none of this runs. Off a Space those
same routes are mocked and sign every visitor in as a fake profile, which is
the trap `identity.refuse_mocked_oauth` refuses to start into. These routes are
the other half of that refusal: the way to have real sign-in without a Space,
rather than merely being told you cannot have the fake one.

The flow is `distinct_server.identity.OpenIdProvider`, which follows
https://huggingface.co/docs/hub/en/oauth. Nothing about the protocol is
decided here. What is decided here is where the three per-attempt secrets live
between the redirect and the callback: in the visitor's own signed session
cookie, never in a process-wide dictionary that two visitors could collide in
and never in a URL.

**The session cookie is signed, not encrypted, and that is enough because of
what goes in it.** A subject and a username, which the person already knows
about themselves, plus a state and a nonce which are worthless once used.
There is no token in it: the access token is used once, at the callback, to
ask who this is, and then dropped. A server that kept it would be holding a
credential to somebody's Hugging Face account for no reason.
"""

from __future__ import annotations

import os
import secrets
from typing import Any, Mapping, Optional

from fastapi import Request
from fastapi.responses import RedirectResponse
from starlette.middleware.sessions import SessionMiddleware

from .identity import (
    Identity,
    IdentityError,
    NotConfigured,
    OpenIdProvider,
    PendingSignIn,
    dev_auth_permitted,
    from_environment,
    gradio_oauth_is_real,
)

#: Where the browser goes to start and finish. The callback must match the
#: redirect URI registered with the OAuth app, which is why it is named here
#: once and read from configuration rather than assembled twice.
LOGIN_PATH = "/auth/login"
CALLBACK_PATH = "/auth/callback"
LOGOUT_PATH = "/auth/logout"

#: The keys used inside the session cookie.
SESSION_IDENTITY = "distinct_identity"
SESSION_PENDING = "distinct_pending"

#: What signs the cookie. Generated per process when nothing is configured,
#: which logs everybody out on restart. That is the safe default: the
#: alternative is a predictable key, and a predictable key on a signed session
#: cookie is the same as no sign-in at all.
ENV_SESSION_SECRET = "DISTINCT_SESSION_SECRET"


def session_secret(environ: Optional[Mapping[str, str]] = None) -> str:
    environ = os.environ if environ is None else environ
    configured = environ.get(ENV_SESSION_SECRET)
    if configured:
        return configured
    return secrets.token_urlsafe(32)


def identity_from_session(session: Mapping[str, Any]) -> Optional[Identity]:
    """Read back a signed-in person, or None.

    Rebuilt through `Identity` rather than trusted as a dictionary, so a
    cookie that somehow carried a blank subject is refused here instead of
    becoming a viewer id that matches everything.
    """

    stored = session.get(SESSION_IDENTITY)
    if not isinstance(stored, Mapping):
        return None
    try:
        return Identity(
            provider=str(stored.get("provider", "")),
            subject=str(stored.get("subject", "")),
            username=str(stored.get("username", "")),
            name=str(stored.get("name", "")),
        )
    except IdentityError:
        return None


def attach(app: Any, *, provider: Optional[OpenIdProvider] = None) -> bool:
    """Mount the routes, and say whether there was anything to mount.

    Returns False when sign-in is not configured for this deployment, which is
    the Space case and the development case. It is not an error: on a Space
    Gradio's own routes are real and these would be a second, conflicting way
    in, and in development `DISTINCT_DEV_AUTH` has already said out loud that
    there is no sign-in.
    """

    # Two conditions, and both are asked the same way the rest of the codebase
    # asks them. `SPACE_ID` alone used to stand for "this is a Space", which is
    # not what Gradio means by it, and `DISTINCT_DEV_AUTH` alone used to be
    # enough to suppress these routes even when a real OAuth app was
    # configured: an operator with a working self-hosted login and a leftover
    # development variable got no login at all.
    if gradio_oauth_is_real() or dev_auth_permitted():
        return False
    if provider is None:
        try:
            provider = from_environment()
        except NotConfigured:
            return False

    app.add_middleware(
        SessionMiddleware,
        secret_key=session_secret(),
        same_site="lax",
        https_only=provider.redirect_uri.startswith("https://"),
    )

    # `Request` is imported at module scope on purpose. With
    # `from __future__ import annotations` these annotations are strings, and
    # FastAPI resolves them against this module's globals: a name imported
    # inside this function is not there, and `request: Request` then reads as
    # an undeclared query parameter, so every route answers 422.
    @app.get(LOGIN_PATH)
    async def login(request: Request):  # pragma: no cover - exercised over HTTP
        pending = provider.begin()
        request.session[SESSION_PENDING] = {
            "state": pending.state,
            "nonce": pending.nonce,
            "code_verifier": pending.code_verifier,
        }
        return RedirectResponse(pending.url, status_code=302)

    @app.get(CALLBACK_PATH)
    async def callback(request: Request):  # pragma: no cover - exercised over HTTP
        stored = request.session.pop(SESSION_PENDING, None)
        if not isinstance(stored, Mapping):
            # No attempt in this browser's session. Either the cookie was
            # dropped or somebody sent this link to somebody else, and both
            # are answered the same way: start again from the button.
            return RedirectResponse("/?signin=expired", status_code=302)
        pending = PendingSignIn(
            url="",
            state=str(stored.get("state", "")),
            nonce=str(stored.get("nonce", "")),
            code_verifier=str(stored.get("code_verifier", "")),
        )
        try:
            identity = provider.complete(
                request.query_params.get("code", ""),
                state=request.query_params.get("state", ""),
                pending=pending,
            )
        except IdentityError:
            # Deliberately no detail in the URL. The reason is for the log,
            # not for whoever engineered the failure.
            return RedirectResponse("/?signin=refused", status_code=302)
        request.session[SESSION_IDENTITY] = {
            "provider": identity.provider,
            "subject": identity.subject,
            "username": identity.username,
            "name": identity.name,
        }
        return RedirectResponse("/", status_code=302)

    @app.get(LOGOUT_PATH)
    async def logout(request: Request):  # pragma: no cover - exercised over HTTP
        request.session.clear()
        return RedirectResponse("/", status_code=302)

    return True