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..1ec9dd11579f7 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 @@ -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: 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..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 @@ -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", {})