diff --git a/.changeset/clean-oauth-retries.md b/.changeset/clean-oauth-retries.md new file mode 100644 index 00000000000..7ae7815f556 --- /dev/null +++ b/.changeset/clean-oauth-retries.md @@ -0,0 +1,5 @@ +--- +"gradio": patch +--- + +fix: Reset stale OAuth sessions and preserve the retry count across callbacks diff --git a/gradio/oauth.py b/gradio/oauth.py index 0622fcdd4a1..cedfde26af8 100644 --- a/gradio/oauth.py +++ b/gradio/oauth.py @@ -110,10 +110,12 @@ async def oauth_redirect_callback(request: fastapi.Request) -> RedirectResponse: # losing the state. A workaround is to delete the cookie and redirect the user to the login page again. # See https://github.com/lepture/authlib/issues/622 for more details. - # Delete all keys that are related to the OAuth state, just in case - for key in list(request.session.keys()): - if key.startswith("_state_huggingface"): - request.session.pop(key) + # Reset stale OAuth state without logging out an already-authenticated + # user whose callback came from an older login attempt. + session_oauth_info = request.session.get("oauth_info") + request.session.clear() + if session_oauth_info is not None: + request.session["oauth_info"] = session_oauth_info # Parse query params nb_redirects = int(request.query_params.get("_nb_redirects", 0)) @@ -136,11 +138,20 @@ async def oauth_redirect_callback(request: fastapi.Request) -> RedirectResponse: "Gradio is not running in a Space (SPACE_HOST environment variable is not set)." " Cannot redirect to non-iframe view." ) from None + host = host.split(",")[0].strip() host_url = "https://" + host.rstrip("/") - return RedirectResponse(host_url + login_uri) - - # Redirect the user to the login page again - return RedirectResponse(login_uri) + response = RedirectResponse(host_url + login_uri) + else: + # Redirect the user to the login page again + response = RedirectResponse(login_uri) + + if session_oauth_info is None: + # SessionMiddleware cannot expire a cookie whose signature it could + # not decode, so delete it explicitly as well. + response.delete_cookie( + "session", secure=True, httponly=True, samesite="none" + ) + return response # OAuth login worked => store the user info in the session and redirect request.session["oauth_info"] = oauth_info @@ -202,6 +213,10 @@ def _generate_redirect_uri(request: fastapi.Request) -> str: # otherwise => keep query params target = "/?" + urllib.parse.urlencode(request.query_params) + callback_query_params = {"_target_url": target} + if "_nb_redirects" in request.query_params: + callback_query_params["_nb_redirects"] = request.query_params["_nb_redirects"] + # On Spaces, the redirect URI must always be https:///login/callback, # so if a custom domain is used, we need to replace it with the hf.space URL if space_host := os.getenv("SPACE_HOST"): @@ -210,12 +225,12 @@ def _generate_redirect_uri(request: fastapi.Request) -> str: 0 ] # When custom domain is used, SPACE_HOST is a comma-separated list print(f"SPACE_HOST after split: {space_host}") - redirect_uri = f"https://{space_host}/login/callback?{urllib.parse.urlencode({'_target_url': target})}" + redirect_uri = f"https://{space_host}/login/callback?{urllib.parse.urlencode(callback_query_params)}" print(f"Redirect URI: {redirect_uri}") return redirect_uri redirect_uri = request.url_for("oauth_redirect_callback").include_query_params( - _target_url=target + **callback_query_params ) redirect_uri_as_str = str(redirect_uri) return redirect_uri_as_str diff --git a/test/test_routes.py b/test/test_routes.py index 6218a64303f..ec5a79b870a 100644 --- a/test/test_routes.py +++ b/test/test_routes.py @@ -2440,6 +2440,147 @@ def test_json_postprocessing_with_queue_false(connect): class TestOAuthSecurity: + def test_oauth_redirect_preserves_retry_count(self, monkeypatch): + from urllib.parse import parse_qs, urlparse + + from gradio.oauth import _generate_redirect_uri + + monkeypatch.setenv("SPACE_HOST", "space.hf.space") + scope = { + "type": "http", + "method": "GET", + "headers": [], + "query_string": b"_nb_redirects=2&_target_url=%2Fapp", + } + + redirect_uri = _generate_redirect_uri(Request(scope)) + + assert parse_qs(urlparse(redirect_uri).query) == { + "_nb_redirects": ["2"], + "_target_url": ["/app"], + } + + def test_mismatching_oauth_state_clears_stale_session_cookie(self, monkeypatch): + from authlib.integrations.base_client.errors import MismatchingStateError + from fastapi.responses import RedirectResponse + from starlette.middleware.sessions import SessionMiddleware + + from gradio import oauth + + class FakeHuggingFaceOAuth: + async def authorize_redirect(self, request, redirect_uri): + return RedirectResponse(redirect_uri) + + async def authorize_access_token(self, request): + assert request.cookies["session"] == "stale-invalid-cookie" + assert not request.session + raise MismatchingStateError() + + class FakeOAuth: + def __init__(self): + self.huggingface = FakeHuggingFaceOAuth() + + def register(self, **kwargs): + pass + + monkeypatch.setattr("authlib.integrations.starlette_client.OAuth", FakeOAuth) + monkeypatch.setattr(oauth, "OAUTH_CLIENT_ID", "client") + monkeypatch.setattr(oauth, "OAUTH_CLIENT_SECRET", "secret") + monkeypatch.setattr(oauth, "OAUTH_SCOPES", "openid") + monkeypatch.setattr(oauth, "OPENID_PROVIDER_URL", "https://huggingface.co") + + app = FastAPI() + oauth._add_oauth_routes(app) + app.add_middleware( + SessionMiddleware, # type: ignore + secret_key="secret", + https_only=True, + ) + + with TestClient(app, base_url="https://space.hf.space") as client: + client.cookies.set( + "session", "stale-invalid-cookie", domain="space.hf.space" + ) + response = client.get( + "/login/callback?_target_url=/app", follow_redirects=False + ) + + monkeypatch.setenv("SPACE_HOST", "space.hf.space, custom.example") + client.cookies.set( + "session", "stale-invalid-cookie", domain="space.hf.space" + ) + external_response = client.get( + "/login/callback?_nb_redirects=3&_target_url=/app", + follow_redirects=False, + ) + + assert response.headers["location"] == ( + "/login/huggingface?_nb_redirects=1&_target_url=%2Fapp" + ) + assert 'session=""' in response.headers["set-cookie"] + assert "Max-Age=0" in response.headers["set-cookie"] + assert external_response.headers["location"] == ( + "https://space.hf.space/login/huggingface?" + "_nb_redirects=4&_target_url=%2Fapp" + ) + + def test_mismatching_oauth_state_preserves_existing_login(self, monkeypatch): + from authlib.integrations.base_client.errors import MismatchingStateError + from fastapi.responses import RedirectResponse + from starlette.middleware.sessions import SessionMiddleware + + from gradio import oauth + + oauth_info = {"access_token": "valid"} + + class FakeHuggingFaceOAuth: + async def authorize_redirect(self, request, redirect_uri): + return RedirectResponse(redirect_uri) + + async def authorize_access_token(self, request): + assert request.session["oauth_info"] == oauth_info + raise MismatchingStateError() + + class FakeOAuth: + def __init__(self): + self.huggingface = FakeHuggingFaceOAuth() + + def register(self, **kwargs): + pass + + monkeypatch.setattr("authlib.integrations.starlette_client.OAuth", FakeOAuth) + monkeypatch.setattr(oauth, "OAUTH_CLIENT_ID", "client") + monkeypatch.setattr(oauth, "OAUTH_CLIENT_SECRET", "secret") + monkeypatch.setattr(oauth, "OAUTH_SCOPES", "openid") + monkeypatch.setattr(oauth, "OPENID_PROVIDER_URL", "https://huggingface.co") + + app = FastAPI() + + @app.get("/seed-session") + async def seed_session(request: Request): + request.session["oauth_info"] = oauth_info + request.session["_state_huggingface_stale"] = {"state": "stale"} + return {} + + @app.get("/session") + async def get_session(request: Request): + return request.session + + oauth._add_oauth_routes(app) + app.add_middleware( + SessionMiddleware, # type: ignore + secret_key="secret", + https_only=True, + ) + + with TestClient(app, base_url="https://space.hf.space") as client: + client.get("/seed-session") + response = client.get("/login/callback", follow_redirects=False) + session = client.get("/session").json() + + assert "Max-Age=0" not in response.headers["set-cookie"] + assert session == {"oauth_info": oauth_info} + def test_redirect_to_target_blocks_external_urls(self): """_redirect_to_target should strip scheme/host to prevent open redirects.""" from gradio.oauth import _redirect_to_target