finally?

.gitattributes

@@ -33,6 +33,3 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-
-# whl
-*.whl filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,6 +1,11 @@
+import patch_gradio
+
 import gradio as gr
 import gradio.utils
 
+from auth import attach_oauth
+
+
 if gradio.utils.get_space() is not None:
     URL, PORT = "https://wauplin-gradio-oauth-test.hf.space", 7860
 else:
@@ -24,20 +29,41 @@ def show_profile(request: gr.Request) -> str:
     session = getattr(request, "session", None) or getattr(request.request, "session", None)
     if session is None:  # should never happen...
         return "No session attached"
-    if "oauth_info" not in session:
+    if "user" not in session:
         return "Please login first"
-    return TEMPLATE.format(**session["oauth_info"]["userinfo"])
+    return TEMPLATE.format(**session["user"])
+
+
+def js_open(url: str) -> str:
+    # Taken from https://cmgdo.com/external-link-in-gradio-button/
+    return f"function() {{window.location.assign('{url}');}}"
+    # return f"function() {{window.open('{url}', '_blank');}}"
 
 
 with gr.Blocks() as demo:
     with gr.Row():
-        gr.LoginButton()
-        gr.LogoutButton()
+        login_button = gr.Button("Login")
+        login_button.click(None, None, None, _js=js_open(f"{URL}/login/huggingface"))
+
+        logout_button = gr.Button("Logout")
+        logout_button.click(None, None, None, _js=js_open(f"{URL}/logout"))
 
     profile_btn = gr.Button("Show profile")
     output = gr.Markdown()
     profile_btn.click(fn=show_profile, outputs=output)
 
+
+old_create_app = gradio.networking.App.create_app
+
+
+def patched_create_app(*args, **kwargs):
+    app = old_create_app(*args, **kwargs)
+    attach_oauth(app)
+    return app
+
+
+gradio.networking.App.create_app = patched_create_app
+
 print(URL)
 
 demo.queue()

auth.py

@@ -0,0 +1,60 @@
+import os
+import hashlib
+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:
+    print("this one")
+    print(request.session)
+    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=hashlib.sha256(OAUTH_CLIENT_SECRET.encode()).hexdigest())
+    app.get("/login/huggingface")(oauth_login)
+    app.get("/login/callback")(oauth_redirect_callback)
+    app.get("/logout")(oauth_logout)

gradio-3.37.1689678355-py3-none-any.whl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3b4cfad28418833f7fee6dec1a31ede548f30c886bb3254559f15d4662c66b4b
-size 2819959

patch_gradio.py

@@ -0,0 +1,31 @@
+import gradio.networking
+import gradio.queueing
+
+from auth import attach_oauth
+
+
+# 1. Patch => attach auth before starting app
+
+old_create_app = gradio.networking.App.create_app
+
+
+def patched_create_app(*args, **kwargs):
+    app = old_create_app(*args, **kwargs)
+    attach_oauth(app)
+    return app
+
+
+gradio.networking.App.create_app = patched_create_app
+
+# 2. Patch => forward session info when using websockets (i.e. when queue is enabled which is the default on Spaces)
+
+old_get_request_params = gradio.queueing.Queue.get_request_params
+
+
+def new_get_request_params(self, websocket):
+    params = old_get_request_params(self, websocket)
+    params["session"] = websocket.session  # Forward session info to the internal request
+    return params
+
+
+gradio.queueing.Queue.get_request_params = new_get_request_params

requirements.txt

@@ -1 +1,2 @@
-https://huggingface.co/spaces/Wauplin/gradio-oauth-test/resolve/main/gradio-3.37.1689678355-py3-none-any.whl
+authlib==1.2.1
+itsdangerous==2.1.2