AIP-84 : Add JWT token revokation for logout invalidation (#47952)

airflow-core/docs/img/airflow_erd.sha256

@@ -1 +1 @@
-203aa3570578ef6e24b0f6725545e3ab830b349a9ca43f8f238ee9588245adc0
+de526a7ff575b0f9deb174bf9e779bceb505d27a867bd2e8a67598de35f37c2f

airflow-core/docs/migrations-ref.rst

@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | Revision ID             | Revises ID       | Airflow Version   | Description                                                  |
 +=========================+==================+===================+==============================================================+
-| ``a5a3e5eb9b8d`` (head) | ``82dbd68e6171`` | ``3.2.0``         | Make external_executor_id TEXT to allow for longer           |
+| ``53ff648b8a26`` (head) | ``a5a3e5eb9b8d`` | ``3.2.0``         | Add revoked_token table.                                     |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``a5a3e5eb9b8d``        | ``82dbd68e6171`` | ``3.2.0``         | Make external_executor_id TEXT to allow for longer           |
 |                         |                  |                   | external_executor_ids.                                       |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``82dbd68e6171``        | ``55297ae24532`` | ``3.2.0``         | Add index to task_reschedule ti_id .                         |

airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py

@@ -45,6 +45,7 @@
 from airflow.configuration import conf
 from airflow.models import Connection, DagModel, Pool, Variable
 from airflow.models.dagbundle import DagBundleModel
+from airflow.models.revoked_token import RevokedToken
 from airflow.models.team import Team, dag_bundle_team_association_table
 from airflow.typing_compat import Unpack
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -132,6 +133,10 @@ def deserialize_user(self, token: dict[str, Any]) -> T:
     def serialize_user(self, user: T) -> dict[str, Any]:
         """Create a subject and extra claims dict from a user object."""
 
+    def revoke_token(self, token: str) -> None:
+        """Revoke a JWT token by persisting its JTI in the database."""
+        self._get_token_validator().revoke_token(token)
+
     async def get_user_from_token(self, token: str) -> BaseUser:
         """Verify the JWT token is valid and create a user object from it if valid."""
         try:
@@ -140,6 +145,9 @@ 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):
+            raise InvalidTokenError("Token has been revoked")
+
         try:
             return self.deserialize_user(payload)
         except (ValueError, KeyError) as e:

airflow-core/src/airflow/api_fastapi/auth/tokens.py

@@ -35,6 +35,7 @@
 from cryptography.hazmat.primitives.serialization import load_pem_private_key
 
 from airflow._shared.timezones import timezone
+from airflow.models.revoked_token import RevokedToken
 
 if TYPE_CHECKING:
     from jwt.algorithms import AllowedKeys, AllowedPrivateKeys
@@ -345,6 +346,15 @@ async def avalidated_claims(
 
         return claims
 
+    def revoke_token(self, token: str) -> None:
+        """Validate the token, extract jti and exp, and revoke it in the database."""
+        try:
+            claims = self.validated_claims(token)
+            if (jti := claims.get("jti")) and (exp := claims.get("exp")):
+                RevokedToken.revoke(jti, datetime.fromtimestamp(exp, tz=timezone.utc))
+        except (jwt.InvalidTokenError, Exception):
+            log.warning("Failed to revoke token", exc_info=True)
+
     def status(self):
         if self.jwks:
             self.jwks.status()

airflow-core/src/airflow/api_fastapi/core_api/routes/public/auth.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import structlog
 from fastapi import HTTPException, Request, status
 from fastapi.responses import RedirectResponse
 
@@ -25,6 +26,8 @@
 from airflow.api_fastapi.core_api.security import AuthManagerDep, is_safe_url
 from airflow.configuration import conf
 
+log = structlog.get_logger(logger_name=__name__)
+
 auth_router = AirflowRouter(tags=["Login"], prefix="/auth")
 
 
@@ -55,6 +58,10 @@ def logout(request: Request, auth_manager: AuthManagerDep) -> RedirectResponse:
     if logout_url:
         return RedirectResponse(logout_url)
 
+    # Revoke the current token before deleting the cookie
+    if token_str := request.cookies.get(COOKIE_NAME_JWT_TOKEN):
+        auth_manager.revoke_token(token_str)
+
     secure = request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback=""))
     response = RedirectResponse(auth_manager.get_url_login())
     response.delete_cookie(

airflow-core/src/airflow/migrations/versions/0102_3_2_0_add_revoked_token_table.py

@@ -0,0 +1,54 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add revoked_token table.
+
+Revision ID: 53ff648b8a26
+Revises: a5a3e5eb9b8d
+Create Date: 2026-02-01 00:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+from airflow.utils.sqlalchemy import UtcDateTime
+
+# revision identifiers, used by Alembic.
+revision = "53ff648b8a26"
+down_revision = "a5a3e5eb9b8d"
+branch_labels = None
+depends_on = None
+airflow_version = "3.2.0"
+
+
+def upgrade():
+    """Add revoked_token table."""
+    op.create_table(
+        "revoked_token",
+        sa.Column("jti", sa.String(32), primary_key=True, nullable=False),
+        sa.Column("exp", UtcDateTime, nullable=False, index=True),
+    )
+
+
+def downgrade():
+    """Drop revoked_token table."""
+    op.drop_table("revoked_token")

airflow-core/src/airflow/models/__init__.py

@@ -69,6 +69,7 @@ def import_all_models():
     import airflow.models.dagwarning
     import airflow.models.deadline_alert
     import airflow.models.errors
+    import airflow.models.revoked_token
     import airflow.models.serialized_dag
     import airflow.models.taskinstancehistory
     import airflow.models.tasklog

airflow-core/src/airflow/models/revoked_token.py

@@ -0,0 +1,79 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import time
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, ClassVar
+
+import structlog
+from sqlalchemy import String, delete, exists, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow.configuration import conf
+from airflow.models.base import Base
+from airflow.utils.session import NEW_SESSION, provide_session
+from airflow.utils.sqlalchemy import UtcDateTime
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+log = structlog.get_logger(__name__)
+
+
+class RevokedToken(Base):
+    """Stores revoked JWT token JTIs to support token invalidation on logout."""
+
+    __tablename__ = "revoked_token"
+
+    # Track last cleanup time to avoid running cleanup on every request
+    _last_cleanup_time: ClassVar[float] = 0.0
+
+    jti: Mapped[str] = mapped_column(String(32), primary_key=True)
+    exp: Mapped[datetime] = mapped_column(UtcDateTime, nullable=False, index=True)
+
+    @classmethod
+    @provide_session
+    def revoke(cls, jti: str, exp: datetime, session: Session = NEW_SESSION) -> None:
+        """Add a token JTI to the revoked tokens."""
+        session.merge(cls(jti=jti, exp=exp))
+
+    @classmethod
+    @provide_session
+    def is_revoked(cls, jti: str, session: Session = NEW_SESSION) -> bool:
+        """Check if a token JTI has been revoked."""
+        cls._maybe_cleanup_expired(session)
+        return bool(session.scalar(select(exists().where(cls.jti == jti))))
+
+    @classmethod
+    def _maybe_cleanup_expired(cls, session: Session) -> None:
+        """
+        Periodically clean up expired revoked tokens.
+
+        Cleanup interval is based on jwt_expiration_time config to ensure expired
+        tokens are cleaned up after they're no longer useful. Uses monotonic time
+        to track intervals.
+        """
+        now = time.monotonic()
+        cleanup_interval = conf.getint("api_auth", "jwt_expiration_time", fallback=3600) * 2
+        if now - cls._last_cleanup_time >= cleanup_interval:
+            cls._last_cleanup_time = now
+            try:
+                session.execute(delete(cls).where(cls.exp < datetime.now(tz=timezone.utc)))
+            except Exception:
+                log.exception("Failed to clean up expired revoked tokens")

airflow-core/src/airflow/utils/db.py

@@ -112,7 +112,7 @@ class MappedClassProtocol(Protocol):
     "3.0.0": "29ce7909c52b",
     "3.0.3": "fe199e1abd77",
     "3.1.0": "cc92b33c6709",
-    "3.2.0": "a5a3e5eb9b8d",
+    "3.2.0": "53ff648b8a26",
 }
 
 # Prefix used to identify tables holding data moved during migration.

airflow-core/src/airflow/utils/db_cleanup.py

@@ -171,6 +171,7 @@ def readable_config(self):
         keep_last_group_by=["dag_id"],
     ),
     _TableConfig(table_name="deadline", recency_column_name="deadline_time", dag_id_column_name="dag_id"),
+    _TableConfig(table_name="revoked_token", recency_column_name="exp"),
 ]
 
 # We need to have `fallback="database"` because this is executed at top level code and provider configuration

airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py

@@ -203,6 +203,29 @@ async def test_get_user_from_token(self, mock_deserialize_user, mock__get_token_
         signer.avalidated_claims.assert_called_once_with(token)
         assert result == user
 
+    @patch(
+        "airflow.models.revoked_token.RevokedToken.is_revoked",
+        return_value=True,
+    )
+    @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(
+        self, mock__get_token_validator, mock_is_revoked, auth_manager
+    ):
+        token = "token"
+        payload = {"jti": "some-jti"}
+        signer = AsyncMock(spec=JWTValidator)
+        signer.avalidated_claims.return_value = payload
+        mock__get_token_validator.return_value = signer
+
+        with pytest.raises(InvalidTokenError, match="Token has been revoked"):
+            await auth_manager.get_user_from_token(token)
+
+        mock_is_revoked.assert_called_once_with("some-jti")
+
     @patch(
         "airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager._get_token_validator",
         autospec=True,
@@ -224,6 +247,19 @@ async def test_get_user_from_token_invalid_token_payload(
         mock_deserialize_user.assert_called_once_with(payload)
         signer.avalidated_claims.assert_called_once_with(token)
 
+    @patch(
+        "airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager._get_token_validator",
+        autospec=True,
+    )
+    def test_revoke_token(self, mock__get_token_validator, auth_manager):
+        token = "token"
+        validator = Mock(spec=JWTValidator)
+        mock__get_token_validator.return_value = validator
+
+        auth_manager.revoke_token(token)
+
+        validator.revoke_token.assert_called_once_with(token)
+
     @patch("airflow.api_fastapi.auth.managers.base_auth_manager.JWTGenerator", autospec=True)
     @patch.object(EmptyAuthManager, "serialize_user")
     def test_generate_jwt_token(self, mock_serialize_user, mock_jwt_generator, auth_manager):

airflow-core/tests/unit/api_fastapi/auth/test_tokens.py

@@ -239,6 +239,88 @@ async def test_jwt_generate_validate_roundtrip_with_jwks(private_key, algorithm,
         assert await validator.avalidated_claims(token)
 
 
+class TestRevokeToken:
+    pytestmark = [pytest.mark.db_test]
+
+    @pytest.fixture(autouse=True)
+    def cleanup_revoked_tokens(self):
+        from tests_common.test_utils.db import clear_db_revoked_tokens
+
+        clear_db_revoked_tokens()
+        yield
+        clear_db_revoked_tokens()
+
+    def test_revoke_token_persists_in_db(self):
+        """Test that revoke_token validates the token and persists the jti in the database."""
+        import time
+
+        from airflow.models.revoked_token import RevokedToken
+
+        now = int(time.time())
+        payload = {
+            "sub": "user",
+            "jti": "revoke-test-jti",
+            "exp": now + 3600,
+            "iat": now,
+            "nbf": now,
+            "aud": "test",
+        }
+        token = jwt.encode(payload, "secret", algorithm="HS256")
+
+        validator = JWTValidator(secret_key="secret", audience="test", algorithm=["HS256"], leeway=0)
+        validator.revoke_token(token)
+
+        assert RevokedToken.is_revoked("revoke-test-jti") is True
+
+    def test_revoke_token_without_jti_does_not_persist(self):
+        """Test that a token without jti does not create a revoked token entry."""
+        import time
+
+        from airflow.models.revoked_token import RevokedToken
+
+        now = int(time.time())
+        payload = {"sub": "user", "exp": now + 3600, "iat": now, "nbf": now, "aud": "test"}
+        token = jwt.encode(payload, "secret", algorithm="HS256")
+
+        validator = JWTValidator(secret_key="secret", audience="test", algorithm=["HS256"], leeway=0)
+        validator.revoke_token(token)
+
+        assert RevokedToken.is_revoked("any-jti") is False
+
+    def test_revoke_token_with_invalid_token_does_not_raise(self):
+        """Test that revoke_token logs a warning instead of raising for an invalid token."""
+        from airflow.models.revoked_token import RevokedToken
+
+        validator = JWTValidator(secret_key="secret", audience="test", algorithm=["HS256"], leeway=0)
+        validator.revoke_token("invalid-token")
+
+        assert RevokedToken.is_revoked("any-jti") is False
+
+    def test_revoke_token_with_db_error_does_not_raise(self):
+        """Test that revoke_token handles database errors gracefully."""
+        import time
+        from unittest.mock import patch
+
+        from sqlalchemy.exc import SQLAlchemyError
+
+        now = int(time.time())
+        payload = {
+            "sub": "user",
+            "jti": "db-error-jti",
+            "exp": now + 3600,
+            "iat": now,
+            "nbf": now,
+            "aud": "test",
+        }
+        token = jwt.encode(payload, "secret", algorithm="HS256")
+
+        validator = JWTValidator(secret_key="secret", audience="test", algorithm=["HS256"], leeway=0)
+        with patch(
+            "airflow.models.revoked_token.RevokedToken.revoke", side_effect=SQLAlchemyError("db down")
+        ):
+            validator.revoke_token(token)
+
+
 @pytest.fixture(scope="session")
 def rsa_private_key():
     return generate_private_key()