Fix Execution API JWT reissue TOCTOU race

airflow-core/src/airflow/api_fastapi/execution_api/app.py

@@ -44,6 +44,7 @@
     get_sig_validation_args,
     get_signing_args,
 )
+from airflow.api_fastapi.execution_api.security import _REQUEST_SCOPE_TOKEN_KEY
 
 if TYPE_CHECKING:
     import httpx
@@ -139,33 +140,33 @@ class JWTReissueMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):
         response: Response = await call_next(request)
 
+        # Reuse claims already validated by JWTBearer (cached on the ASGI scope).
+        # Re-calling avalidated_claims here is a TOCTOU race: the token can expire
+        # between auth and this middleware, raising ExpiredSignatureError which was
+        # swallowed so no Refreshed-API-Token was returned and the next heartbeat 403'd.
+        token = request.scope.get(_REQUEST_SCOPE_TOKEN_KEY)
+        if not token:
+            return response
+
         refreshed_token: str | None = None
-        auth_header = request.headers.get("authorization")
-        if auth_header and auth_header.lower().startswith("bearer "):
-            token = auth_header.split(" ", 1)[1]
-            try:
+        try:
+            claims = token.claims.model_dump()
+
+            # Workload tokens are long-lived and meant to survive queue wait times.
+            if claims.get("scope") == "workload":
+                return response
+
+            now = int(time.time())
+            token_lifetime = int(claims.get("exp", 0)) - int(claims.get("iat", 0))
+            refresh_when_less_than = max(int(token_lifetime * 0.20), 30)
+            valid_left = int(claims.get("exp", 0)) - now
+            if valid_left <= refresh_when_less_than:
                 async with svcs.Container(request.app.state.svcs_registry) as services:
-                    validator: JWTValidator = await services.aget(JWTValidator)
-                    claims = await validator.avalidated_claims(token, {})
-
-                    # Workload tokens are long-lived and meant to survive queue
-                    # wait times so avoid refreshing them. If avalidated_claims
-                    # raises for a workload token, the outer except handles it.
-                    if claims.get("scope") == "workload":
-                        return response
-
-                    now = int(time.time())
-                    token_lifetime = int(claims.get("exp", 0)) - int(claims.get("iat", 0))
-                    refresh_when_less_than = max(int(token_lifetime * 0.20), 30)
-                    valid_left = int(claims.get("exp", 0)) - now
-                    if valid_left <= refresh_when_less_than:
-                        generator: JWTGenerator = await services.aget(JWTGenerator)
-                        refreshed_token = generator.generate(claims)
-            except Exception as err:
-                # Do not block the response if refreshing fails; log a warning for visibility
-                logger.warning(
-                    "JWT reissue middleware failed to refresh token", error=str(err), exc_info=True
-                )
+                    generator: JWTGenerator = await services.aget(JWTGenerator)
+                    refreshed_token = generator.generate(claims)
+        except Exception as err:
+            # Do not block the response if refreshing fails; log a warning for visibility
+            logger.warning("JWT reissue middleware failed to refresh token", error=str(err), exc_info=True)
 
         if refreshed_token:
             response.headers["Refreshed-API-Token"] = refreshed_token

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py

@@ -18,12 +18,15 @@
 from __future__ import annotations
 
 from unittest.mock import AsyncMock
+from uuid import UUID
 
 import pytest
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 
 from airflow.api_fastapi.auth.tokens import JWTValidator
 from airflow.api_fastapi.execution_api.app import lifespan
+from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
+from airflow.api_fastapi.execution_api.security import _REQUEST_SCOPE_TOKEN_KEY, require_auth
 
 
 @pytest.fixture
@@ -33,6 +36,31 @@ def exec_app(client):
     return last_route.app
 
 
+def _install_scoped_token(
+    exec_app: FastAPI,
+    *,
+    moment: int,
+    validity: int,
+    sub: str = "edb09971-4e0e-4221-ad3f-800852d38085",
+) -> None:
+    """Override require_auth to cache JWT claims on the ASGI scope like JWTBearer does."""
+
+    async def mock_require_auth(request: Request) -> TIToken:
+        claims = TIClaims.model_validate(
+            {
+                "sub": sub,
+                "scope": "execution",
+                "iat": moment,
+                "exp": moment + validity,
+            }
+        )
+        token = TIToken(id=UUID(sub), claims=claims)
+        request.scope[_REQUEST_SCOPE_TOKEN_KEY] = token
+        return token
+
+    exec_app.dependency_overrides[require_auth] = mock_require_auth
+
+
 @pytest.mark.parametrize(
     ("age", "validity", "expect_refreshed_token"),
     (
@@ -48,22 +76,61 @@ def test_expiring_token_is_reissued(
     client, exec_app: FastAPI, time_machine, age, validity, expect_refreshed_token
 ):
     moment = 1743451846  # A "random" unix epoch timestamp.
-    auth = AsyncMock(spec=JWTValidator)
-    auth.avalidated_claims.return_value = {
-        "sub": "edb09971-4e0e-4221-ad3f-800852d38085",
-        "iat": moment,
-        "exp": moment + validity,
-    }
-
+    _install_scoped_token(exec_app, moment=moment, validity=validity)
     time_machine.move_to(moment + age, tick=False)
 
-    # Inject our fake JWTValidator object. Can be over-ridden by tests if they want
-    lifespan.registry.register_value(JWTValidator, auth)
-    # In order to test this we need any endpoint to hit. The easiest one to use is variable get
-
+    # Any authenticated endpoint exercises JWTReissueMiddleware.
     response = client.get("/execution/variables/key1", headers={"Authorization": "Bearer dummy"})
 
     if expect_refreshed_token:
         assert "Refreshed-API-Token" in response.headers
     else:
         assert "Refreshed-API-Token" not in response.headers
+
+
+@pytest.mark.db_test
+def test_token_expiring_mid_request_is_reissued(client, exec_app: FastAPI, time_machine):
+    """JWT reissue must survive tokens that expire during request handling.
+
+    Regression for the TOCTOU race where JWTBearer accepted a near-expiry token,
+    the handler ran, then JWTReissueMiddleware re-validated the bearer string and
+    hit ExpiredSignatureError — swallowing the error and omitting Refreshed-API-Token
+    so the next heartbeat failed with 403.
+    """
+    moment = 1743451846
+    validity = 600
+    sub = "edb09971-4e0e-4221-ad3f-800852d38085"
+
+    # Start just before expiry so auth would still accept the token.
+    time_machine.move_to(moment + validity - 1, tick=False)
+
+    # Route that mimics JWTBearer caching claims, then advances the clock past exp.
+    @exec_app.get("/__test_mid_request_expiry")
+    async def _expire_during_request(request: Request):
+        request.scope[_REQUEST_SCOPE_TOKEN_KEY] = TIToken(
+            id=UUID(sub),
+            claims=TIClaims.model_validate(
+                {
+                    "sub": sub,
+                    "scope": "execution",
+                    "iat": moment,
+                    "exp": moment + validity,
+                }
+            ),
+        )
+        time_machine.move_to(moment + validity + 1, tick=False)
+        return {"ok": True}
+
+    # Ensure avalidated_claims is not needed by middleware (and would fail if called).
+    auth = AsyncMock(spec=JWTValidator)
+    auth.avalidated_claims.side_effect = Exception("middleware must not re-validate JWT")
+    lifespan.registry.register_value(JWTValidator, auth)
+
+    response = client.get(
+        "/execution/__test_mid_request_expiry",
+        headers={"Authorization": "Bearer dummy"},
+    )
+
+    assert response.status_code == 200
+    assert "Refreshed-API-Token" in response.headers
+    auth.avalidated_claims.assert_not_awaited()