Fix OAuth redirect loops caused by stale sessions

.changeset/clean-oauth-retries.md

@@ -0,0 +1,5 @@
+---
+"gradio": patch
+---
+
+fix: Reset stale OAuth sessions and preserve the retry count across callbacks

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://<space_host>/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

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