From f54e0ea808c5ae48f763278d2f50251ef12d447e Mon Sep 17 00:00:00 2001 From: Seonghwan Lee Date: Wed, 12 Aug 2026 22:23:20 +0900 Subject: [PATCH] Recover the revoked-token check from a stale DB connection The JWT revocation check added in 3.2.0 (RevokedToken.is_revoked in BaseAuthManager.get_user_from_token) is the first DB access in the auth path and runs on the shared scoped session. A prior request (e.g. FAB's deserialize_user) can leave that session bound to a connection the database later drops on idle timeout (MySQL error 4031, "disconnected ... because of inactivity"). The connection is never re-checked-out, so pool_pre_ping / pool_recycle cannot detect it, and the first authenticated request after an idle period fails with HTTP 500. deserialize_user already recovers from this class of failure (#62919), but the revocation check runs earlier in get_user_from_token and had no recovery, so the 500 simply returned one step sooner. Apply the same discard-and-retry: on SQLAlchemyError, remove the poisoned scoped session and retry the check once on a fresh connection. Related to #71395 --- .../auth/managers/base_auth_manager.py | 34 +++++++++- .../auth/managers/test_base_auth_manager.py | 68 +++++++++++++++++++ 2 files changed, 101 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py b/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py index f780097cea484..2bb00a691970b 100644 --- a/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py +++ b/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py @@ -22,13 +22,16 @@ import warnings from abc import ABCMeta, abstractmethod from collections import defaultdict +from contextlib import suppress from enum import Enum from functools import cache, cached_property from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar from jwt import InvalidTokenError from sqlalchemy import select +from sqlalchemy.exc import SQLAlchemyError +from airflow import settings from airflow.api_fastapi.auth.managers.models.base_user import BaseUser from airflow.api_fastapi.auth.managers.models.resource_details import ( ConnectionDetails, @@ -161,7 +164,7 @@ async def get_user_from_token(self, token: str) -> BaseUser: log.error("JWT token is not valid: %s", e) raise e - if (jti := payload.get("jti")) and RevokedToken.is_revoked(jti): + if (jti := payload.get("jti")) and self._is_token_revoked(jti): raise InvalidTokenError("Token has been revoked") try: @@ -170,6 +173,35 @@ async def get_user_from_token(self, token: str) -> BaseUser: log.error("Couldn't deserialize user from token, JWT token is not valid: %s", e) raise InvalidTokenError(str(e)) + @staticmethod + def _is_token_revoked(jti: str) -> bool: + """ + Return whether the token ``jti`` has been revoked, tolerating a stale DB connection. + + This is the first database access in the auth path and runs on the shared scoped + session. A prior request (for example FAB's ``deserialize_user``) can leave that session + bound to a connection the database later drops on idle timeout (MySQL error 4031, + "disconnected ... because of inactivity"). The connection is never re-checked-out, so + ``pool_pre_ping`` / ``pool_recycle`` cannot detect it and the first request after an idle + period fails with a 500. On any ``SQLAlchemyError`` we discard the poisoned scoped session + and retry once on a fresh connection, mirroring the recovery + ``FabAuthManager.deserialize_user`` gained in #62919. See #71395. + """ + try: + return RevokedToken.is_revoked(jti) + except SQLAlchemyError: + log.warning( + "Revoked-token check failed on a stale DB session; discarding the scoped " + "session and retrying once on a fresh connection.", + exc_info=True, + ) + # settings.Session is Optional and only set once the DB is configured; guard so the + # discard is a no-op (rather than a crash) if it is missing. + if (session_registry := settings.Session) is not None: + with suppress(Exception): + session_registry.remove() + return RevokedToken.is_revoked(jti) + def get_fastapi_middlewares(self) -> list[tuple[_MiddlewareFactory[Any], dict[str, Any]]]: """ Return middlewares the auth manager wants registered on the main FastAPI app. diff --git a/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py b/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py index d7f2286aa1508..db14d8e1ba5f7 100644 --- a/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py +++ b/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py @@ -22,6 +22,7 @@ import pytest from jwt import InvalidTokenError +from sqlalchemy.exc import OperationalError from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager, T from airflow.api_fastapi.auth.managers.models.base_user import BaseUser @@ -333,6 +334,73 @@ async def test_get_user_from_token_revoked( mock_is_revoked.assert_called_once_with("some-jti") + @staticmethod + def _stale_connection_error() -> OperationalError: + """Mimic MySQL error 4031 raised when a pooled connection was dropped on idle timeout.""" + return OperationalError( + statement="SELECT EXISTS (SELECT 1 FROM revoked_token WHERE jti = %s)", + params=("some-jti",), + orig=Exception("(4031, 'The client was disconnected by the server because of inactivity.')"), + ) + + @patch("airflow.api_fastapi.auth.managers.base_auth_manager.settings.Session") + @patch("airflow.models.revoked_token.RevokedToken.is_revoked") + @patch( + "airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager._get_token_validator", + autospec=True, + ) + @patch.object(EmptyAuthManager, "deserialize_user") + @pytest.mark.asyncio + async def test_get_user_from_token_recovers_from_stale_session_on_revoked_check( + self, mock_deserialize_user, mock__get_token_validator, mock_is_revoked, mock_session, auth_manager + ): + """ + The revoked-token check must survive a stale pooled DB connection: on SQLAlchemyError it + discards the scoped session and retries once, so the first request after an idle period + succeeds instead of returning 500. Regression test for issue #71395. + """ + token = "token" + payload = {"jti": "some-jti"} + user = BaseAuthManagerUserTest(name="test") + signer = AsyncMock(spec=JWTValidator) + signer.avalidated_claims.return_value = payload + mock__get_token_validator.return_value = signer + mock_deserialize_user.return_value = user + # First check hits the poisoned connection, the retry runs on a fresh one. + mock_is_revoked.side_effect = [self._stale_connection_error(), False] + + result = await auth_manager.get_user_from_token(token) + + assert result == user + assert mock_is_revoked.call_count == 2 + # The poisoned scoped session was discarded before the retry. + mock_session.remove.assert_called_once_with() + mock_deserialize_user.assert_called_once_with(payload) + + @patch("airflow.api_fastapi.auth.managers.base_auth_manager.settings.Session") + @patch("airflow.models.revoked_token.RevokedToken.is_revoked") + @patch( + "airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager._get_token_validator", + autospec=True, + ) + @pytest.mark.asyncio + async def test_get_user_from_token_revoked_after_stale_session_retry( + self, mock__get_token_validator, mock_is_revoked, mock_session, auth_manager + ): + """A revocation confirmed by the retry is still honored after recovering the session.""" + token = "token" + payload = {"jti": "some-jti"} + signer = AsyncMock(spec=JWTValidator) + signer.avalidated_claims.return_value = payload + mock__get_token_validator.return_value = signer + mock_is_revoked.side_effect = [self._stale_connection_error(), True] + + with pytest.raises(InvalidTokenError, match="Token has been revoked"): + await auth_manager.get_user_from_token(token) + + assert mock_is_revoked.call_count == 2 + mock_session.remove.assert_called_once_with() + @patch( "airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager._get_token_validator", autospec=True,