From c5a52c11b7b9fd34d445ecb37545d8f77c496e95 Mon Sep 17 00:00:00 2001 From: Gayathri Srividya Rajavarapu Date: Fri, 5 Jun 2026 10:37:37 +0530 Subject: [PATCH 1/5] Fix JWT token not refreshed when token expires mid-request MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit closes: #67939 Long-running tasks (>10 min with default 600s JWT lifetime) could die with repeated 403 "Signature has expired" errors despite the refresh mechanism being in place. The failure was a TOCTOU race: Before: JWTReissueMiddleware called avalidated_claims a second time after the handler completed. If the token expired in the milliseconds between JWTBearer's validation and the middleware's re-validation, jwt raised ExpiredSignatureError, the except block swallowed it, and no Refreshed-API-Token header was set. The client's next heartbeat (30 s later) arrived with a now-fully-expired token, got 403, and after MAX_FAILED_HEARTBEATS the supervisor killed the task. After: JWTBearer already validates and caches the token on request.scope. JWTReissueMiddleware reads claims from request.scope instead of calling avalidated_claims again — eliminating the race entirely. Even if the token crosses its expiry boundary during request processing, valid_left will be <= refresh_when_less_than and a fresh token is still issued. Changes: - app.py: read TIToken from request.scope[_REQUEST_SCOPE_TOKEN_KEY] (cached by JWTBearer) instead of re-parsing the Authorization header and calling avalidated_claims a second time. - test_router.py: assert avalidated_claims is called exactly once per request; add test_just_expired_token_is_reissued_within_grace_period which moves time past the token's expiry and asserts that the middleware still sets Refreshed-API-Token. --- .../airflow/api_fastapi/execution_api/app.py | 34 +++++++++--------- .../versions/head/test_router.py | 36 +++++++++++++++++++ 2 files changed, 52 insertions(+), 18 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index 88c7d23fff26a..549eae19871d1 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -42,6 +42,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 @@ -135,25 +136,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 = {"sub": str(token.id), **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: diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py index 3a3e669f7e3a6..f0cbcc4e64f1f 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py @@ -67,3 +67,39 @@ def test_expiring_token_is_reissued( assert "Refreshed-API-Token" in response.headers else: assert "Refreshed-API-Token" not in response.headers + auth.avalidated_claims.assert_awaited_once_with("dummy", {}) + + +@pytest.mark.db_test +def test_just_expired_token_is_reissued_within_grace_period(client, exec_app: FastAPI, time_machine): + """Token that expires mid-request is still reissued by the middleware. + + 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 completes, 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 + re-validating, 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 — simulates the middleware running after + # the token boundary. JWTBearer already accepted the token (mocked); the middleware must + # still issue a refresh rather than silently dropping it. + time_machine.move_to(moment + 601, tick=False) + + lifespan.registry.register_value(JWTValidator, auth) + + response = 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", {}) From fd960947f19eee22d7ffc9b1972dfc63db98eceb Mon Sep 17 00:00:00 2001 From: Gayathri Srividya Rajavarapu Date: Fri, 5 Jun 2026 10:41:45 +0530 Subject: [PATCH 2/5] Rename test: remove 'grace period' wording Rename test_just_expired_token_is_reissued_within_grace_period to test_token_expiring_mid_request_is_reissued_without_revalidation to make clear the fix avoids re-validation, not that it adds any leeway. --- .../execution_api/versions/head/test_router.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py index f0cbcc4e64f1f..4e54c75aceebd 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py @@ -71,8 +71,8 @@ def test_expiring_token_is_reissued( @pytest.mark.db_test -def test_just_expired_token_is_reissued_within_grace_period(client, exec_app: FastAPI, time_machine): - """Token that expires mid-request is still reissued by the middleware. +def test_token_expiring_mid_request_is_reissued_without_revalidation(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), @@ -81,7 +81,8 @@ def test_just_expired_token_is_reissued_within_grace_period(client, exec_app: Fa 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 - re-validating, so it still issues a fresh token even when the original has since expired. + calling avalidated_claims again, so it still issues a fresh token even when the original + has since expired. """ moment = 1743451846 auth = AsyncMock(spec=JWTValidator) @@ -91,9 +92,9 @@ def test_just_expired_token_is_reissued_within_grace_period(client, exec_app: Fa "exp": moment + 600, } - # Move time to 1 second past the token's expiry — simulates the middleware running after - # the token boundary. JWTBearer already accepted the token (mocked); the middleware must - # still issue a refresh rather than silently dropping it. + # 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) From 27f90065952d7ac9539ff946d8cd6f65b6c7c6b4 Mon Sep 17 00:00:00 2001 From: Gayathri Srividya Rajavarapu Date: Fri, 5 Jun 2026 18:53:11 +0530 Subject: [PATCH 3/5] Fix test fixture: override _jwt_bearer instead of require_auth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The conftest client fixture overrides require_auth, which causes FastAPI to skip _jwt_bearer entirely — so request.scope[_REQUEST_SCOPE_TOKEN_KEY] is never set and the mock JWTValidator is never called. Replace the client fixture in test_router.py with one that overrides _jwt_bearer directly. mock_jwt_bearer calls the registered JWTValidator mock, stores the resulting TIToken in request.scope, and lets the real require_auth run — matching the production code path that JWTReissueMiddleware depends on. --- .../versions/head/test_router.py | 68 +++++++++++++++++-- 1 file changed, 64 insertions(+), 4 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py index 4e54c75aceebd..1eca9dc6878a0 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py +++ b/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 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 + from starlette.routing import Mount + + for route in 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( @@ -67,6 +126,7 @@ def test_expiring_token_is_reissued( 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", {}) From e90a001b5a2865e82de778f9477baa3f55e392fa Mon Sep 17 00:00:00 2001 From: GayathriSrividya Date: Fri, 12 Jun 2026 10:58:42 +0530 Subject: [PATCH 4/5] Update airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py Co-authored-by: Ash Berlin-Taylor --- .../unit/api_fastapi/execution_api/versions/head/test_router.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py index 1eca9dc6878a0..e322d8c9c39e6 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py @@ -136,7 +136,7 @@ def test_token_expiring_mid_request_is_reissued_without_revalidation(client, exe 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 completes, and the middleware runs. In the old code the middleware would call + 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. From 30b0bb522689cb0b7c7de7c6bb3d86475efcf24b Mon Sep 17 00:00:00 2001 From: Gayathri Srividya Rajavarapu Date: Sat, 13 Jun 2026 10:43:40 +0530 Subject: [PATCH 5/5] Address PR #68054 maintainer review comments --- .../src/airflow/api_fastapi/execution_api/app.py | 2 +- .../execution_api/versions/head/test_router.py | 16 +++++++++------- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index 549eae19871d1..6b9fc600ef306 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -139,7 +139,7 @@ async def dispatch(self, request: Request, call_next): token = request.scope.get(_REQUEST_SCOPE_TOKEN_KEY) if token: try: - claims = {"sub": str(token.id), **token.claims.model_dump()} + claims = token.claims.model_dump() # Workload tokens are long-lived and meant to survive queue # wait times so avoid refreshing them. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py index e322d8c9c39e6..662eb29c1049a 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py @@ -35,7 +35,7 @@ @pytest.fixture -def client(): +def jwt_bearer_client(): """Test client that exercises JWTBearer so request.scope is populated for the middleware.""" from starlette.routing import Mount @@ -83,10 +83,10 @@ async def mock_jwt_bearer(request: Request): @pytest.fixture -def exec_app(client): +def exec_app(jwt_bearer_client): from starlette.routing import Mount - for route in client.app.routes: + 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") @@ -104,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) @@ -120,7 +120,7 @@ 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 @@ -131,7 +131,9 @@ def test_expiring_token_is_reissued( @pytest.mark.db_test -def test_token_expiring_mid_request_is_reissued_without_revalidation(client, exec_app: FastAPI, time_machine): +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 @@ -159,7 +161,7 @@ def test_token_expiring_mid_request_is_reissued_without_revalidation(client, exe lifespan.registry.register_value(JWTValidator, auth) - response = client.get("/execution/variables/key1", headers={"Authorization": "Bearer dummy"}) + 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.