This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new cf839aca896 Fix JWT token not refreshed when token expires mid-request
(#68499)
cf839aca896 is described below
commit cf839aca896194b83e0c1d4f573b44aab7a61ec8
Author: GayathriSrividya <[email protected]>
AuthorDate: Thu Oct 15 17:41:12 2026 +0530
Fix JWT token not refreshed when token expires mid-request (#68499)
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.
Co-authored-by: Ash Berlin-Taylor <[email protected]>
---
.../src/airflow/api_fastapi/execution_api/app.py | 39 +++++----
.../execution_api/versions/head/test_router.py | 92 +++++++++++++++++++---
2 files changed, 100 insertions(+), 31 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 76093ecd0bd..bf254e32c4e 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py
@@ -41,6 +41,7 @@ from airflow.api_fastapi.auth.tokens import (
get_signing_args,
)
from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI
+from airflow.api_fastapi.execution_api.security import _REQUEST_SCOPE_TOKEN_KEY
if TYPE_CHECKING:
from airflow.models.dagbag import DBDagBag
@@ -140,27 +141,25 @@ class JWTReissueMiddleware(BaseHTTPMiddleware):
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, {})
-
- # Only short-lived execution tokens are renewed here. Any
other type has a
- # lifetime set by its issuer (workload and callback tokens
outlive queue waits,
- # the dag_processor_session token is rotated by
provisioning, and Job and parsing
- # tokens are renewed by re-registering and re-exchanging),
so a new type must
- # not become renewable by default. Tokens without a scope
are legacy execution tokens.
- if claims.get("scope", "execution") != "execution":
- 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()
+
+ # Only short-lived execution tokens are renewed here. Any
other type has a
+ # lifetime set by its issuer (workload and callback tokens
outlive queue waits,
+ # the dag_processor_session token is rotated by provisioning,
and Job and parsing
+ # tokens are renewed by re-registering and re-exchanging), so
a new type must
+ # not become renewable by default. Tokens without a scope are
legacy execution tokens.
+ if claims.get("scope", "execution") != "execution":
+ 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 b00f63c119e..746c8933c0b 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,21 @@ from __future__ import annotations
from unittest.mock import AsyncMock
import pytest
-from fastapi import FastAPI
+import svcs
+from fastapi import Depends, FastAPI
+from fastapi.testclient import TestClient
+from jwt import ExpiredSignatureError
-from airflow.api_fastapi.auth.tokens import JWTValidator
-from airflow.api_fastapi.execution_api.app import lifespan
+from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator
+from airflow.api_fastapi.execution_api.app import JWTReissueMiddleware,
lifespan
+from airflow.api_fastapi.execution_api.security import _jwt_bearer,
require_auth
@pytest.fixture
-def exec_app(client):
- last_route = client.app.routes[-1]
- assert isinstance(last_route.app, FastAPI)
- return last_route.app
+def jwt_bearer_client(client, exec_app, monkeypatch):
+ """Test client that exercises JWTBearer so request.scope is populated for
the middleware."""
+ monkeypatch.delitem(exec_app.dependency_overrides, require_auth)
+ return client
@pytest.mark.parametrize(
@@ -44,9 +48,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
-):
+def test_expiring_token_is_reissued(jwt_bearer_client, time_machine, age,
validity, expect_refreshed_token):
moment = 1743451846 # A "random" unix epoch timestamp.
auth = AsyncMock(spec=JWTValidator)
auth.avalidated_claims.return_value = {
@@ -61,12 +63,80 @@ 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", {})
+
+
[email protected](
+ ("request_age", "expected_status"),
+ [
+ pytest.param(599, 200, id="expires-during-request"),
+ pytest.param(601, 403, id="expired-before-request"),
+ ],
+)
+def test_token_expiring_mid_request_is_reissued_without_revalidation(
+ time_machine, request_age, expected_status
+):
+ """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 1s left, JWTBearer validates it (still valid at that
moment), the handler
+ advances time past expiry, 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. A token that is already expired when the request
arrives must still be
+ rejected without a refresh.
+ """
+ moment = 1743451846
+ secret = "jwt-refresh-test-secret" * 4
+ audience = "test-jwt-refresh"
+ generator = JWTGenerator(secret_key=secret, audience=audience,
valid_for=600)
+ validator = JWTValidator(secret_key=secret, audience=audience, leeway=0)
+ claims = {
+ "sub": "edb09971-4e0e-4221-ad3f-800852d38085",
+ "scope": "execution",
+ "custom": {"key": "value"},
+ }
+ time_machine.move_to(moment, tick=False)
+ token = generator.generate(claims)
+ app = FastAPI()
+ app.add_middleware(JWTReissueMiddleware)
+
+ @app.get("/protected", dependencies=[Depends(_jwt_bearer)])
+ async def protected():
+ time_machine.move_to(moment + 601, tick=False)
+ return {"ok": True}
+
+ with svcs.Registry() as registry:
+ registry.register_value(JWTGenerator, generator)
+ registry.register_value(JWTValidator, validator)
+ app.state.svcs_registry = registry
+ time_machine.move_to(moment + request_age, tick=False)
+
+ with TestClient(app) as client:
+ response = client.get("/protected", headers={"Authorization":
f"Bearer {token}"})
+
+ assert response.status_code == expected_status
+ with pytest.raises(ExpiredSignatureError):
+ validator.validated_claims(token)
+
+ if expected_status == 200:
+ assert response.json() == {"ok": True}
+ refreshed_claims =
validator.validated_claims(response.headers["Refreshed-API-Token"])
+ assert refreshed_claims.items() >= claims.items()
+ assert refreshed_claims["exp"] == moment + 1201
+ else:
+ assert response.json() == {"detail": "Invalid auth token"}
+ assert "Refreshed-API-Token" not in response.headers
@pytest.mark.db_test
