Fix JWT token not refreshed when token expires mid-request

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
@@ -140,25 +141,22 @@ async def dispatch(self, request: Request, call_next):
         response: Response = await call_next(request)
 
         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]
+        token = request.scope.get(_REQUEST_SCOPE_TOKEN_KEY)
+        if token:
             try:
-                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:
+                claims = token.claims.model_dump()
+
+                # Workload tokens are long-lived and meant to survive queue
+                # wait times so avoid refreshing them.
+                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:
                         generator: JWTGenerator = await services.aget(JWTGenerator)
                         refreshed_token = generator.generate(claims)
             except Exception as err:

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

@@ -20,17 +20,76 @@
 from unittest.mock import AsyncMock
 
 import pytest
-from fastapi import FastAPI
+import svcs
+from fastapi import FastAPI, HTTPException, Request, status
+from fastapi.security import HTTPBearer
+from fastapi.testclient import TestClient
 
 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,
+    _jwt_bearer,
+)
+
+
+@pytest.fixture
+def jwt_bearer_client():
+    """Test client that exercises JWTBearer so request.scope is populated for the middleware."""
+    from starlette.routing import Mount
+
+    from airflow.api_fastapi.app import cached_app
+
+    app = cached_app(apps="execution")
+
+    exec_app: FastAPI | None = None
+    for route in app.routes:
+        if isinstance(route, Mount) and route.path == "/execution" and isinstance(route.app, FastAPI):
+            exec_app = route.app
+            break
+    if exec_app is None:
+        raise RuntimeError("Execution API sub-app not found")
+
+    _http_bearer = HTTPBearer(auto_error=False)
+
+    async def mock_jwt_bearer(request: Request):
+        """Drop-in for _jwt_bearer that uses the registered JWTValidator mock and sets scope."""
+        if cached := request.scope.get(_REQUEST_SCOPE_TOKEN_KEY):
+            return cached
+
+        creds = await _http_bearer(request)
+        if not creds:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing auth token")
+
+        async with svcs.Container(request.app.state.svcs_registry) as services:
+            validator: JWTValidator = await services.aget(JWTValidator)
+            try:
+                claims = await validator.avalidated_claims(creds.credentials, {})
+            except Exception:
+                raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid auth token")
+
+        claims.setdefault("scope", "execution")
+        token = TIToken(id=claims["sub"], claims=TIClaims(**claims))
+        request.scope[_REQUEST_SCOPE_TOKEN_KEY] = token
+        return token
+
+    exec_app.dependency_overrides[_jwt_bearer] = mock_jwt_bearer
+
+    with TestClient(app) as c:
+        yield c
+
+    exec_app.dependency_overrides.pop(_jwt_bearer, None)
 
 
 @pytest.fixture
-def exec_app(client):
-    last_route = client.app.routes[-1]
-    assert isinstance(last_route.app, FastAPI)
-    return last_route.app
+def exec_app(jwt_bearer_client):
+    from starlette.routing import Mount
+
+    for route in jwt_bearer_client.app.routes:
+        if isinstance(route, Mount) and route.path == "/execution" and isinstance(route.app, FastAPI):
+            return route.app
+    raise RuntimeError("Execution API sub-app not found")
 
 
 @pytest.mark.parametrize(
@@ -45,7 +104,7 @@ def exec_app(client):
 )
 @pytest.mark.db_test
 def test_expiring_token_is_reissued(
-    client, exec_app: FastAPI, time_machine, age, validity, expect_refreshed_token
+    jwt_bearer_client, exec_app: FastAPI, time_machine, age, validity, expect_refreshed_token
 ):
     moment = 1743451846  # A "random" unix epoch timestamp.
     auth = AsyncMock(spec=JWTValidator)
@@ -61,9 +120,49 @@ def test_expiring_token_is_reissued(
     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
 
-    response = client.get("/execution/variables/key1", headers={"Authorization": "Bearer dummy"})
+    response = jwt_bearer_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
+    # avalidated_claims must be called exactly once — by JWTBearer only, not by the middleware.
+    auth.avalidated_claims.assert_awaited_once_with("dummy", {})
+
+
+@pytest.mark.db_test
+def test_token_expiring_mid_request_is_reissued_without_revalidation(
+    jwt_bearer_client, exec_app: FastAPI, time_machine
+):
+    """Middleware reissues from cached JWTBearer claims without re-validating the token.
+
+    Regression test for the TOCTOU race in JWTReissueMiddleware: a heartbeat arrives with a
+    token that has ~0s left, JWTBearer validates it (still technically valid at that moment),
+    the request starts, and the middleware runs. In the old code the middleware would call
+    avalidated_claims a second time and get ExpiredSignatureError — no Refreshed-API-Token
+    header would be set, and the task would die on the next heartbeat.
+
+    With the fix the middleware reads claims from request.scope (set by JWTBearer) instead of
+    calling avalidated_claims again, so it still issues a fresh token even when the original
+    has since expired.
+    """
+    moment = 1743451846
+    auth = AsyncMock(spec=JWTValidator)
+    auth.avalidated_claims.return_value = {
+        "sub": "edb09971-4e0e-4221-ad3f-800852d38085",
+        "iat": moment,
+        "exp": moment + 600,
+    }
+
+    # Move time to 1 second past the token's expiry. JWTBearer already accepted the token
+    # (mocked); the middleware must still issue a refresh using the cached claims rather than
+    # silently dropping it.
+    time_machine.move_to(moment + 601, tick=False)
+
+    lifespan.registry.register_value(JWTValidator, auth)
+
+    response = jwt_bearer_client.get("/execution/variables/key1", headers={"Authorization": "Bearer dummy"})
+
+    assert "Refreshed-API-Token" in response.headers
+    # avalidated_claims must be called exactly once — by JWTBearer only.
+    auth.avalidated_claims.assert_awaited_once_with("dummy", {})