From 0700829ee87109a6f5bc70897265df2afc655d01 Mon Sep 17 00:00:00 2001 From: anmolxlight Date: Wed, 22 Jul 2026 10:23:47 +0000 Subject: [PATCH] Fix Execution API JWT reissue TOCTOU race JWTReissueMiddleware re-validated the bearer token after the handler ran. If the token expired mid-request, ExpiredSignatureError was swallowed and no Refreshed-API-Token was returned, so the next heartbeat failed with 403. Reuse claims already cached on the ASGI scope by JWTBearer instead. Closes: #70222 Related: #67939 Signed-off-by: anmolxlight --- .../airflow/api_fastapi/execution_api/app.py | 51 ++++++----- .../versions/head/test_router.py | 91 ++++++++++++++++--- 2 files changed, 105 insertions(+), 37 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 5a87cdb81e04c..010df596791a6 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/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 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..24d6a83cbb05b 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 @@ -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()