Spaces:
Sleeping
Sleeping
File size: 4,849 Bytes
6b8f7db |
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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import logging
import os
import fastapi
import gradio
from fastapi.responses import RedirectResponse
from gradio.oauth import MOCKED_OAUTH_TOKEN
from modules.presets import i18n
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")
def _add_oauth_routes(app: fastapi.FastAPI) -> None:
"""Add OAuth routes to the FastAPI app (login, callback handler and logout)."""
try:
from authlib.integrations.starlette_client import OAuth
except ImportError as e:
raise ImportError(
"Cannot initialize OAuth to due a missing library. Please run `pip install gradio[oauth]` or add "
"`gradio[oauth]` to your requirements.txt file in order to install the required dependencies."
) from e
# Check environment variables
msg = (
"OAuth is required but {} environment variable is not set. Make sure you've enabled OAuth in your Space by"
" setting `hf_oauth: true` in the Space metadata."
)
if OAUTH_CLIENT_ID is None:
raise ValueError(msg.format("OAUTH_CLIENT_ID"))
if OAUTH_CLIENT_SECRET is None:
raise ValueError(msg.format("OAUTH_CLIENT_SECRET"))
if OAUTH_SCOPES is None:
raise ValueError(msg.format("OAUTH_SCOPES"))
if OPENID_PROVIDER_URL is None:
raise ValueError(msg.format("OPENID_PROVIDER_URL"))
# Register OAuth server
oauth = OAuth()
oauth.register(
name="huggingface",
client_id=OAUTH_CLIENT_ID,
client_secret=OAUTH_CLIENT_SECRET,
client_kwargs={"scope": OAUTH_SCOPES},
server_metadata_url=OPENID_PROVIDER_URL + "/.well-known/openid-configuration",
)
# Define OAuth routes
@app.get("/login/huggingface")
async def oauth_login(request: fastapi.Request):
"""Endpoint that redirects to HF OAuth page."""
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)
@app.get("/login/callback")
async def oauth_redirect_callback(request: fastapi.Request) -> RedirectResponse:
"""Endpoint that handles the OAuth callback."""
token = await oauth.huggingface.authorize_access_token(request)
request.session["oauth_profile"] = token["userinfo"]
request.session["oauth_token"] = token
return RedirectResponse("/")
@app.get("/logout")
async def oauth_logout(request: fastapi.Request) -> RedirectResponse:
"""Endpoint that logs out the user (e.g. delete cookie session)."""
request.session.pop("oauth_profile", None)
request.session.pop("oauth_token", None)
# 清除cookie并跳转到首页
response = RedirectResponse(url="/", status_code=302)
response.delete_cookie(key=f"access-token")
response.delete_cookie(key=f"access-token-unsecure")
return response
def _add_mocked_oauth_routes(app: fastapi.FastAPI) -> None:
"""Add fake oauth routes if Gradio is run locally and OAuth is enabled.
Clicking on a gr.LoginButton will have the same behavior as in a Space (i.e. gets redirected in a new tab) but
instead of authenticating with HF, a mocked user profile is added to the session.
"""
# Define OAuth routes
@app.get("/login/huggingface")
async def oauth_login(request: fastapi.Request):
"""Fake endpoint that redirects to HF OAuth page."""
return RedirectResponse("/login/callback")
@app.get("/login/callback")
async def oauth_redirect_callback(request: fastapi.Request) -> RedirectResponse:
"""Endpoint that handles the OAuth callback."""
request.session["oauth_profile"] = MOCKED_OAUTH_TOKEN["userinfo"]
request.session["oauth_token"] = MOCKED_OAUTH_TOKEN
return RedirectResponse("/")
@app.get("/logout")
async def oauth_logout(request: fastapi.Request) -> RedirectResponse:
"""Endpoint that logs out the user (e.g. delete cookie session)."""
request.session.pop("oauth_profile", None)
request.session.pop("oauth_token", None)
# 清除cookie并跳转到首页
response = RedirectResponse(url="/", status_code=302)
response.delete_cookie(key=f"access-token")
response.delete_cookie(key=f"access-token-unsecure")
return response
def reg_patch():
gradio.oauth._add_mocked_oauth_routes = _add_mocked_oauth_routes
gradio.oauth._add_oauth_routes = _add_oauth_routes
logging.info(i18n("覆盖gradio.oauth /logout路由"))
|