oauth

competitions/app.py

@@ -16,6 +16,7 @@ from competitions import utils
 from competitions.errors import AuthenticationError
 from competitions.info import CompetitionInfo
 from competitions.leaderboard import Leaderboard
+from competitions.oauth import attach_oauth
 from competitions.runner import JobRunner
 from competitions.submissions import Submissions
 from competitions.text import SUBMISSION_SELECTION_TEXT, SUBMISSION_TEXT
@@ -27,7 +28,7 @@ COMPETITION_ID = os.environ.get("COMPETITION_ID")
 OUTPUT_PATH = os.environ.get("OUTPUT_PATH", "/tmp/model")
 START_DATE = os.environ.get("START_DATE", "2000-12-31")
 DISABLE_PUBLIC_LB = int(os.environ.get("DISABLE_PUBLIC_LB", 0))
-
+USE_OAUTH = int(os.environ.get("USE_OAUTH", 0))
 
 disable_progress_bars()
 
@@ -69,6 +70,9 @@ thread.start()
 
 
 app = FastAPI()
+if USE_OAUTH == 1:
+    attach_oauth(app)
+
 static_path = os.path.join(BASE_DIR, "static")
 app.mount("/static", StaticFiles(directory=static_path), name="static")
 templates_path = os.path.join(BASE_DIR, "templates")
@@ -92,6 +96,11 @@ async def read_form(request: Request):
     return templates.TemplateResponse("index.html", context)
 
 
+@app.get("/use_oauth", response_class=JSONResponse)
+async def use_oauth(request: Request):
+    return {"response": USE_OAUTH}
+
+
 @app.get("/competition_info", response_class=JSONResponse)
 async def get_comp_info(request: Request):
     info = COMP_INFO.competition_desc

competitions/create.py

@@ -119,6 +119,9 @@ def _create_readme(competition_name):
     _readme += "sdk: docker\n"
     _readme += "pinned: false\n"
     _readme += "duplicated_from: autotrain-projects/autotrain-advanced\n"
+    _readme += "hf_oauth: true\n"
+    _readme += "hf_oauth_scopes:\n"
+    _readme += "  - read-repos\n"
     _readme += "---\n"
     _readme = io.BytesIO(_readme.encode())
     return _readme

competitions/oauth.py

@@ -0,0 +1,240 @@
+"""OAuth support for AutoTrain.
+Taken from: https://github.com/gradio-app/gradio/blob/main/gradio/oauth.py
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import typing
+import urllib.parse
+import warnings
+from dataclasses import dataclass, field
+
+import fastapi
+from authlib.integrations.starlette_client import OAuth
+from fastapi.responses import RedirectResponse
+from huggingface_hub import whoami
+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")
+
+
+def attach_oauth(app: fastapi.FastAPI):
+    # Add `/login/huggingface`, `/login/callback` and `/logout` routes to enable OAuth in the Gradio app.
+    # If the app is running in a Space, OAuth is enabled normally. Otherwise, we mock the "real" routes to make the
+    # user log in with a fake user profile - without any calls to hf.co.
+    if os.environ.get("SPACE_ID") is not None and os.environ.get("HF_TOKEN") is None:
+        _add_oauth_routes(app)
+    else:
+        _add_mocked_oauth_routes(app)
+
+    # Session Middleware requires a secret key to sign the cookies. Let's use a hash
+    # of the OAuth secret key to make it unique to the Space + updated in case OAuth
+    # config gets updated.
+    session_secret = (OAUTH_CLIENT_SECRET or "") + "-v3"
+    # ^ if we change the session cookie format in the future, we can bump the version of the session secret to make
+    #   sure cookies are invalidated. Otherwise some users with an old cookie format might get a HTTP 500 error.
+    app.add_middleware(
+        SessionMiddleware,
+        secret_key=hashlib.sha256(session_secret.encode()).hexdigest(),
+        same_site="none",
+        https_only=True,
+    )
+
+
+def _add_oauth_routes(app: fastapi.FastAPI) -> None:
+    """Add OAuth routes to the FastAPI app (login, callback handler and logout)."""
+    # 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."""
+        # Define target (where to redirect after login)
+        redirect_uri = _generate_redirect_uri(request)
+        return await oauth.huggingface.authorize_redirect(request, redirect_uri)  # type: ignore
+
+    @app.get("/login/callback")
+    async def oauth_redirect_callback(request: fastapi.Request) -> RedirectResponse:
+        """Endpoint that handles the OAuth callback."""
+        oauth_info = await oauth.huggingface.authorize_access_token(request)  # type: ignore
+        request.session["oauth_info"] = oauth_info
+        return _redirect_to_target(request)
+
+    @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_info", None)
+        return _redirect_to_target(request)
+
+
+def _add_mocked_oauth_routes(app: fastapi.FastAPI) -> None:
+    """Add fake oauth routes if Gradio is run locally and OAuth is enabled.
+    Instead of authenticating with HF, a mocked user profile is added to the session.
+    """
+    warnings.warn(
+        "AutoTrain does not support OAuth features outside of a Space environment. To help"
+        " you debug your app locally, the login and logout buttons are mocked with your"
+        " profile. To make it work, your machine must be logged in to Huggingface."
+    )
+    mocked_oauth_info = _get_mocked_oauth_info()
+
+    # Define OAuth routes
+    @app.get("/login/huggingface")
+    async def oauth_login(request: fastapi.Request):  # noqa: ARG001
+        """Fake endpoint that redirects to HF OAuth page."""
+        # Define target (where to redirect after login)
+        redirect_uri = _generate_redirect_uri(request)
+        return RedirectResponse("/login/callback?" + urllib.parse.urlencode({"_target_url": redirect_uri}))
+
+    @app.get("/login/callback")
+    async def oauth_redirect_callback(request: fastapi.Request) -> RedirectResponse:
+        """Endpoint that handles the OAuth callback."""
+        request.session["oauth_info"] = mocked_oauth_info
+        return _redirect_to_target(request)
+
+    @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_info", None)
+        logout_url = str(request.url).replace("/logout", "/")  # preserve query params
+        return RedirectResponse(url=logout_url)
+
+
+def _generate_redirect_uri(request: fastapi.Request) -> str:
+    if "_target_url" in request.query_params:
+        # if `_target_url` already in query params => respect it
+        target = request.query_params["_target_url"]
+    else:
+        # otherwise => keep query params
+        target = "/?" + urllib.parse.urlencode(request.query_params)
+
+    redirect_uri = request.url_for("oauth_redirect_callback").include_query_params(_target_url=target)
+    redirect_uri_as_str = str(redirect_uri)
+    if redirect_uri.netloc.endswith(".hf.space"):
+        # In Space, FastAPI redirect as http but we want https
+        redirect_uri_as_str = redirect_uri_as_str.replace("http://", "https://")
+    return redirect_uri_as_str
+
+
+def _redirect_to_target(request: fastapi.Request, default_target: str = "/") -> RedirectResponse:
+    target = request.query_params.get("_target_url", default_target)
+    return RedirectResponse(target)
+
+
+@dataclass
+class OAuthProfile(typing.Dict):  # inherit from Dict for backward compatibility
+    """
+    A OAuthProfile object that can be used to inject the profile of a user in a
+    function. If a function expects `OAuthProfile` or `Optional[OAuthProfile]` as input,
+    the value will be injected from the FastAPI session if the user is logged in. If the
+    user is not logged in and the function expects `OAuthProfile`, an error will be
+    raised.
+
+    Attributes:
+        name (str): The name of the user (e.g. 'abhishek').
+        username (str): The username of the user (e.g. 'abhishek')
+        profile (str): The profile URL of the user (e.g. 'https://huggingface.co/abhishek').
+        picture (str): The profile picture URL of the user.
+    """
+
+    name: str = field(init=False)
+    username: str = field(init=False)
+    profile: str = field(init=False)
+    picture: str = field(init=False)
+
+    def __init__(self, data: dict):  # hack to make OAuthProfile backward compatible
+        self.update(data)
+        self.name = self["name"]
+        self.username = self["preferred_username"]
+        self.profile = self["profile"]
+        self.picture = self["picture"]
+
+
+@dataclass
+class OAuthToken:
+    """
+    A Gradio OAuthToken object that can be used to inject the access token of a user in a
+    function. If a function expects `OAuthToken` or `Optional[OAuthToken]` as input,
+    the value will be injected from the FastAPI session if the user is logged in. If the
+    user is not logged in and the function expects `OAuthToken`, an error will be
+    raised.
+
+    Attributes:
+        token (str): The access token of the user.
+        scope (str): The scope of the access token.
+        expires_at (int): The expiration timestamp of the access token.
+    """
+
+    token: str
+    scope: str
+    expires_at: int
+
+
+def _get_mocked_oauth_info() -> typing.Dict:
+    token = os.environ.get("HF_TOKEN")
+    if token is None:
+        raise ValueError(
+            "Your machine must be logged in to HF to debug AutoTrain locally. Please "
+            "set `HF_TOKEN` as environment variable "
+            "with one of your access token. You can generate a new token in your "
+            "settings page (https://huggingface.co/settings/tokens)."
+        )
+
+    user = whoami(token=token)
+    if user["type"] != "user":
+        raise ValueError(
+            "Your machine is not logged in with a personal account. Please use a "
+            "personal access token. You can generate a new token in your settings page"
+            " (https://huggingface.co/settings/tokens)."
+        )
+
+    return {
+        "access_token": token,
+        "token_type": "bearer",
+        "expires_in": 3600,
+        "id_token": "AAAAAAAAAAAAAAAAAAAAAAAAAA",
+        "scope": "openid profile",
+        "expires_at": 1691676444,
+        "userinfo": {
+            "sub": "11111111111111111111111",
+            "name": user["fullname"],
+            "preferred_username": user["name"],
+            "profile": f"https://huggingface.co/{user['name']}",
+            "picture": user["avatarUrl"],
+            "website": "",
+            "aud": "00000000-0000-0000-0000-000000000000",
+            "auth_time": 1691672844,
+            "nonce": "aaaaaaaaaaaaaaaaaaa",
+            "iat": 1691672844,
+            "exp": 1691676444,
+            "iss": "https://huggingface.co",
+        },
+    }

competitions/templates/index.html

@@ -244,6 +244,31 @@
         });
 
     </script>
+    <script>
+        function makeApiRequest(url, callback) {
+            var xhr = new XMLHttpRequest();
+            xhr.open("GET", url, true);
+            xhr.onreadystatechange = function () {
+                if (xhr.readyState === 4 && xhr.status === 200) {
+                    var response = JSON.parse(xhr.responseText);
+                    callback(response.response);
+                }
+            };
+            xhr.send();
+        }
+
+        function checkOAuth() {
+            var url = "/use_oauth";
+            makeApiRequest(url, function (response) {
+                if (response === 1) {
+                    document.getElementById("userToken").style.display = "none";
+                    document.getElementById("loginButton").style.display = "block";
+                }
+            });
+        }
+
+        window.onload = checkOAuth;
+    </script>
 </head>
 
 <body class="flex h-screen">
@@ -355,13 +380,18 @@
                         <span class="flex-1 ms-3 whitespace-nowrap">Team</span>
                     </a>
                 </li> -->
-                <li>
+                <li id="userToken">
                     <label for="user_token" class="text-xs font-medium">Hugging Face <a
                             href="https://huggingface.co/settings/tokens" target="_blank">Token</a> (read-only)
                     </label>
                     <input type="password" name="user_token" id="user_token"
                         class="mt-1 block w-full border border-gray-300 px-3 py-1.5 bg-white rounded-md shadow-sm focus:outline-none focus:ring-indigo-500 focus:border-indigo-500">
                 </li>
+                <li id="loginButton" style="display: none;">
+                    <a href="/login/huggingface"
+                        class="bg-blue-500 hover:bg-blue-700 text-white font-bold py-2 px-4 rounded">Login with Hugging
+                        Face</a>
+                </li>
             </ul>
             <footer>
                 <div class="w-full mx-auto max-w-screen-xl p-4 md:flex md:items-center md:justify-between">