Spaces:
Running
Running
File size: 2,172 Bytes
ccec886 48d58d8 ccec886 1af2fc7 ccec886 91bbd67 ccec886 16e8620 ccec886 16e8620 ccec886 91bbd67 48d58d8 91bbd67 ccec886 91bbd67 ccec886 48d58d8 4de0319 48d58d8 4de0319 48d58d8 ccec886 4cd3e4f 48d58d8 ccec886 4de0319 ccec886 4de0319 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import os
from authlib.integrations.starlette_client import OAuth
from fastapi import FastAPI
from fastapi.requests import Request
from fastapi.responses import 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,
)
# Hack to close the login/logout page once the user is logged in/out.
# TODO: can it be less hacky?
# CLOSE_WINDOW_HTML = HTMLResponse("<script>window.close();</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 RedirectResponse("/")
async def oauth_redirect_callback(request: Request) -> RedirectResponse:
token = await oauth.huggingface.authorize_access_token(request)
request.session["user"] = token["userinfo"] # TODO: we should store entire token
return RedirectResponse("/")
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)
|