auth.py

import hashlib
import os

from authlib.integrations.base_client import MismatchingStateError
from authlib.integrations.starlette_client import OAuth
from fastapi import FastAPI
from fastapi.requests import Request
from fastapi.responses import HTMLResponse, RedirectResponse
from starlette.middleware.sessions import SessionMiddleware


OAUTH_CLIENT_ID = os.environ.get("OAUTH_CLIENT_ID")
OAUTH_CLIENT_SECRET = os.environ.get("OAUTH_CLIENT_SECRET")
OAUTH_SCOPES = os.environ.get("OAUTH_SCOPES")
OPENID_PROVIDER_URL = os.environ.get("OPENID_PROVIDER_URL")

for value in (OAUTH_CLIENT_ID, OAUTH_CLIENT_SECRET, OAUTH_SCOPES, OPENID_PROVIDER_URL):
    if value is None:
        raise ValueError("Missing environment variable")

USER_INFO_URL = OPENID_PROVIDER_URL + "/oauth/userinfo"
METADATA_URL = OPENID_PROVIDER_URL + "/.well-known/openid-configuration"


oauth = OAuth()
oauth.register(
    name="huggingface",
    client_id=OAUTH_CLIENT_ID,
    client_secret=OAUTH_CLIENT_SECRET,
    client_kwargs={"scope": OAUTH_SCOPES},
    server_metadata_url=METADATA_URL,
)

# Close the login/logout page once the user is logged in/out.
CLOSE_WINDOW_HTML = HTMLResponse("<script>window.close();</script>")

# If no third-party cookie, open a new tab to login/logout + redirect to the gradio app on this tab.
OPEN_WINDOW_HTML = HTMLResponse("<script>window.open('{url}', '_blank'); window.location.replace('/');</script>")


async def oauth_login(request: Request):
    redirect_uri = str(request.url_for("oauth_redirect_callback"))
    if ".hf.space" in redirect_uri:  # In Space, FastAPI redirect as http but we want https
        redirect_uri = redirect_uri.replace("http://", "https://")
    return await oauth.huggingface.authorize_redirect(request, redirect_uri)


async def oauth_logout(request: Request) -> RedirectResponse:
    request.session.pop("user", None)
    return CLOSE_WINDOW_HTML


async def oauth_redirect_callback(request: Request) -> RedirectResponse:
    try:
        token = await oauth.huggingface.authorize_access_token(request)
        request.session["user"] = token["userinfo"]  # TODO: we should store the entire token
        print("New user: ", token["userinfo"]["name"])
        close_tab = True
    except MismatchingStateError:
        # Third-party cookies are most likely forbidden meaning the session will not be set inside the Space iframe.
        # To counterpart this, we redirect the user to use the Space url outside of the iframe.
        print("Mismatch error: open in new window")
        close_tab = False

    return CLOSE_WINDOW_HTML if close_tab else OPEN_WINDOW_HTML.format(url=request.url_for("oauth_login"))


def attach_oauth(app: FastAPI) -> None:
    app.add_middleware(
        SessionMiddleware,
        secret_key="000" + hashlib.sha256(OAUTH_CLIENT_SECRET.encode()).hexdigest(),
        same_site="none",
        https_only=True,
    )
    app.get("/login/huggingface")(oauth_login)
    app.get("/login/callback")(oauth_redirect_callback)
    app.get("/logout")(oauth_logout)