import os from authlib.integrations.starlette_client import OAuth from fastapi import FastAPI from fastapi.requests import Request from fastapi.responses import HTMLResponse 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, ) # Hack to close the login/logout page once the user is logged in/out. # TODO: can it be less hacky? CLOSE_WINDOW_HTML = HTMLResponse("") 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) -> HTMLResponse: request.session.pop("user", None) return CLOSE_WINDOW_HTML async def oauth_redirect_callback(request: Request) -> HTMLResponse: token = await oauth.huggingface.authorize_access_token(request) request.session["user"] = token["userinfo"] # TODO: we should store entire token return CLOSE_WINDOW_HTML def attach_oauth(app: FastAPI) -> None: app.add_middleware(SessionMiddleware, secret_key="session-secret-key") # TODO: make this is secret key app.get("/login/huggingface")(oauth_login) app.get("/login/callback")(oauth_redirect_callback) app.get("/logout")(oauth_logout)