Spaces:
Runtime error
Runtime error
File size: 2,917 Bytes
0d86759 e599b38 0d86759 e599b38 0d86759 3014630 0d86759 e599b38 0d86759 e599b38 3014630 0d86759 3014630 0d86759 3014630 0d86759 e599b38 0d86759 |
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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import os
import hashlib
from typing import Any
from authlib.integrations.starlette_client import OAuth
from fastapi import FastAPI
from fastapi.requests import Request
from fastapi.responses import RedirectResponse, HTMLResponse
from starlette.middleware.sessions import SessionMiddleware
from authlib.integrations.base_client import MismatchingStateError
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>")
OPEN_WINDOW_HTML = HTMLResponse("<script>window.open('{url}', '_blank');</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
return CLOSE_WINDOW_HTML if close_tab else RedirectResponse("/")
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
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.
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)
|