Fix JWT token not refreshed when token expires mid-request - #68499
Fix JWT token not refreshed when token expires mid-request#68499GayathriSrividya wants to merge 8 commits into
Conversation
closes: apache#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.
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.
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.
…d/test_router.py Co-authored-by: Ash Berlin-Taylor <ash_github@firemirror.com>
30b0bb5 to
e476807
Compare
|
@ashb could you please take another look when you have a moment? I’ve addressed the points from your earlier review: Removed the redundant sub claim merge in the JWT reissue middleware and now use the claims already cached in request scope. Thanks again for the guidance and review. |
| # 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 |
There was a problem hiding this comment.
This isn't really testing the TOCTOU bug - fairly sure this would pass without any code changes.
I think what you need to do is register a custom route that does a time travel inside it
There was a problem hiding this comment.
- Fixed! The test now properly validates the TOCTOU scenario:
Test: test_router.py
- Token expires mid-request (line 157:
time_machine.move_to(moment + 601)) - Middleware still refreshes (line 163:
assert "Refreshed-API-Token" in response.headers) - No re-validation happens (line 164:
assert_awaited_once_withprovesavalidated_claimscalled only once)
Fix: app.py
- Middleware reads from
request.scope(cached by JWTBearer) instead of callingavalidated_claimsagain - Guarantees refresh even if token expired between request start and middleware execution
Let me know if any other method is needed to address the issue.
e476807 to
30b0bb5
Compare
potiuk
added
the
ready for maintainer review
|
LGTM! we can merge post review from code owners @amoghrajesh @kaxil @ashb |
amoghrajesh
left a comment
amoghrajesh
left a comment
There was a problem hiding this comment.
Thanks for working on this @GayathriSrividya, this needs some more work.
| from starlette.routing import Mount | ||
|
|
||
| from airflow.api_fastapi.app import cached_app |
| assert isinstance(last_route.app, FastAPI) | ||
| return last_route.app | ||
| def exec_app(jwt_bearer_client): | ||
| from starlette.routing import Mount |
|
|
||
|
|
||
| @pytest.mark.db_test | ||
| def test_token_expiring_mid_request_is_reissued_without_revalidation( |
There was a problem hiding this comment.
This test still does not reproduce the TOCTOU issue claimed to be fixed by the PR. Maybe if the middleware still called avalidated_claims a second time, that call would raise (simulating the token expiring between JWTBearer's validation and the middleware running)
6 participants
closes: #67939
This replaces #68054, which cannot be reopened because GitHub treats the head branch as recreated after the remote refresh.
What changed in this update:
This keeps the original fix intent intact: avoid token revalidation in middleware and refresh from the token already cached in request scope.
related: #68054
Important
🛠️ Maintainer triage note for @GayathriSrividya · by
@potiuk· 2026-07-08 15:38 UTCHelpful heads-up from the maintainers — please address before this PR can be reviewed:
CI image checks / Static checks). See the contributor guide.Full criteria: Pull Request quality criteria.
The ball is in your court — you've been assigned to this PR. Fix the above, then mark it Ready for review.
Automated triage — may be imperfect; a maintainer takes the next look.