Skip to content

Fix JWT token not refreshed when token expires mid-request#68499

Open
GayathriSrividya wants to merge 8 commits into
apache:mainfrom
GayathriSrividya:fix/jwt-token-refresh-67939-v2
Open

Fix JWT token not refreshed when token expires mid-request#68499
GayathriSrividya wants to merge 8 commits into
apache:mainfrom
GayathriSrividya:fix/jwt-token-refresh-67939-v2

Conversation

@GayathriSrividya

@GayathriSrividya GayathriSrividya commented Jun 13, 2026

Copy link
Copy Markdown
Contributor

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:

  • Removed redundant claim injection in middleware and now use cached claims directly.
  • Renamed the generic test fixture to a JWT-specific fixture name for clarity.
  • Kept the regression test wording update requested in review.

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 UTC

Helpful heads-up from the maintainers — please address before this PR can be reviewed:

  • Pre-commit / static checks failing (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.

Gayathri Srividya Rajavarapu and others added 5 commits June 13, 2026 10:37
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>
@boring-cyborg boring-cyborg Bot added area:API Airflow's REST/HTTP API area:task-sdk labels Jun 13, 2026
@GayathriSrividya
GayathriSrividya force-pushed the fix/jwt-token-refresh-67939-v2 branch from 30b0bb5 to e476807 Compare June 13, 2026 06:16
@GayathriSrividya

Copy link
Copy Markdown
Contributor Author

@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.
Renamed the fixture from client to jwt_bearer_client so its purpose is explicit and less generic.
Kept the test-docstring wording adjustment about request start vs middleware timing.
Behavior-wise, the intent is unchanged: middleware no longer revalidates the bearer token and instead refreshes from the already-validated scope claims.

Thanks again for the guidance and review.

Comment on lines +157 to +166
# 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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ashb

  • 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_with proves avalidated_claims called only once)

Fix: app.py

  • Middleware reads from request.scope (cached by JWTBearer) instead of calling avalidated_claims again
  • 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.

@GayathriSrividya
GayathriSrividya force-pushed the fix/jwt-token-refresh-67939-v2 branch from e476807 to 30b0bb5 Compare June 15, 2026 04:37
@eladkal eladkal added this to the Airflow 3.3.0 milestone Jun 15, 2026
@eladkal eladkal added the type:bug-fix Changelog: Bug Fixes label Jun 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:API Airflow's REST/HTTP API area:task-sdk ready for maintainer review Set after triaging when all criteria pass. type:bug-fix Changelog: Bug Fixes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Task fails to refresh the JWT token with LocalExecutor

5 participants