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)