This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 99eff4608f9 Keep the token cookie set by a route in
JWTRefreshMiddleware (#75217)
99eff4608f9 is described below
commit 99eff4608f9a63874bf6f0efda845abeaa91d2d1
Author: Mathieu Monet <[email protected]>
AuthorDate: Tue Oct 20 17:40:30 2026 +0200
Keep the token cookie set by a route in JWTRefreshMiddleware (#75217)
JWTRefreshMiddleware appends its own _token Set-Cookie after the
response headers. When the route itself sets _token, the browser
applies the middleware cookie last, so it overrides the route:
- A login route, e.g. an auth manager callback mounted under /auth,
sets a fresh token while the request still carries an expired or
invalid one. The middleware appends a deletion and the user is
logged out right after logging in.
- /api/v2/auth/logout deletes the cookie, but if refresh_user returned
a refreshed user for that request the middleware appends a newly
minted, unrevoked token and the user stays logged in.
Skip the middleware token cookie when the response already sets
_token. Stale root-path cookies are still cleared.
Co-authored-by: Claude Opus 5.5 <[email protected]>
---
.../api_fastapi/auth/middlewares/refresh_token.py | 23 +++++++---
.../auth/middlewares/test_refresh_token.py | 51 ++++++++++++++++++++++
2 files changed, 69 insertions(+), 5 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py
b/airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py
index 2e97e6c25c1..37186d9be53 100644
--- a/airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py
+++ b/airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py
@@ -50,7 +50,8 @@ class JWTRefreshMiddleware:
1. Extracts JWT token from cookies and build the user from the token
2. Calls ``refresh_user`` method from auth manager with the user
3. If ``refresh_user`` returns a user, generate a JWT token based upon
this user and send it in the
- response as cookie
+ response as cookie, unless the response already sets that cookie itself
(e.g. a login or logout
+ route)
Requests to ``UNAUTHENTICATED_PATHS`` or under ``skip_path_prefixes``
(paths relative to the
root path) are passed through untouched.
@@ -89,7 +90,9 @@ class JWTRefreshMiddleware:
# token used by the auth manager is expired
new_token = ""
- cookie_headers = self._build_cookie_headers(request, new_token) if
new_token is not None else []
+ token_cookie, cleanup_cookies = (
+ self._build_cookie_headers(request, new_token) if new_token is not
None else (None, [])
+ )
response_started = False
async def send_with_cookies(message: Message) -> None:
@@ -97,7 +100,13 @@ class JWTRefreshMiddleware:
if message["type"] == "http.response.start":
response_started = True
headers = MutableHeaders(scope=message)
- for value in cookie_headers:
+ cookie_values = cleanup_cookies
+ # The token set by the route itself wins: appending ours after
it would override it.
+ if token_cookie and not any(
+ value.startswith(f"{COOKIE_NAME_JWT_TOKEN}=") for value in
headers.getlist("set-cookie")
+ ):
+ cookie_values = [token_cookie, *cleanup_cookies]
+ for value in cookie_values:
headers.append("set-cookie", value)
await send(message)
@@ -115,7 +124,8 @@ class JWTRefreshMiddleware:
return path in UNAUTHENTICATED_PATHS or
path.startswith(self.skip_path_prefixes)
@staticmethod
- def _build_cookie_headers(request: Request, new_token: str) -> list[str]:
+ def _build_cookie_headers(request: Request, new_token: str) -> tuple[str,
list[str]]:
+ """Return the ``Set-Cookie`` value of the token and those clearing
stale token cookies."""
cookie_path = get_cookie_path()
secure = request.base_url.scheme == "https" or bool(conf.get("api",
"ssl_cert", fallback=""))
# Only used to render the Set-Cookie header values
@@ -141,7 +151,10 @@ class JWTRefreshMiddleware:
secure=secure,
samesite="lax",
)
- return [value.decode("latin-1") for key, value in response.raw_headers
if key == b"set-cookie"]
+ token_cookie, *cleanup_cookies = (
+ value.decode("latin-1") for key, value in response.raw_headers if
key == b"set-cookie"
+ )
+ return token_cookie, cleanup_cookies
@staticmethod
async def _refresh_user(current_token: str) -> tuple[BaseUser | None,
BaseUser | None]:
diff --git
a/airflow-core/tests/unit/api_fastapi/auth/middlewares/test_refresh_token.py
b/airflow-core/tests/unit/api_fastapi/auth/middlewares/test_refresh_token.py
index 346c1e8dc61..a601a0417d1 100644
--- a/airflow-core/tests/unit/api_fastapi/auth/middlewares/test_refresh_token.py
+++ b/airflow-core/tests/unit/api_fastapi/auth/middlewares/test_refresh_token.py
@@ -28,6 +28,7 @@ from fastapi.testclient import TestClient
from airflow.api_fastapi.auth.managers.base_auth_manager import
COOKIE_NAME_JWT_TOKEN, BaseAuthManager
from airflow.api_fastapi.auth.managers.models.base_user import BaseUser
+from airflow.api_fastapi.auth.middlewares import refresh_token
from airflow.api_fastapi.auth.middlewares.refresh_token import
JWTRefreshMiddleware
from airflow.api_fastapi.core_api.security import
USER_INJECTED_BY_TRUSTED_MIDDLEWARE
@@ -53,6 +54,20 @@ def _build_app(static_dir=None, root_path: str = "") ->
tuple[FastAPI, dict]:
def version():
return PlainTextResponse("3.x")
+ @app.get("/auth/login_callback")
+ def login_callback():
+ response = PlainTextResponse("logged in")
+ response.set_cookie(
+ COOKIE_NAME_JWT_TOKEN, "issued_token",
path=refresh_token.get_cookie_path(), httponly=True
+ )
+ return response
+
+ @app.get("/api/v2/auth/logout")
+ def logout():
+ response = PlainTextResponse("logged out")
+ response.delete_cookie(COOKIE_NAME_JWT_TOKEN,
path=refresh_token.get_cookie_path(), httponly=True)
+ return response
+
@app.get("/api/v2/stream")
def stream():
return StreamingResponse(iter([b"chunk1-", b"chunk2-", b"chunk3"]),
media_type="text/plain")
@@ -136,6 +151,42 @@ class TestJWTRefreshMiddleware:
mock_auth_manager.generate_jwt.assert_called_once_with(refreshed_user)
assert response.headers["set-cookie"] == "_token=new_token; HttpOnly;
Path=/; SameSite=lax"
+ def test_invalid_token_does_not_clear_token_set_by_route(self,
mock_resolve_user, mock_auth_manager):
+ mock_resolve_user.side_effect = HTTPException(status_code=403,
detail="Invalid JWT token")
+ app, _ = _build_app()
+
+ response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"stale_token"}).get("/auth/login_callback")
+
+ assert response.headers.get_list("set-cookie") == [
+ "_token=issued_token; HttpOnly; Path=/; SameSite=lax"
+ ]
+
+ def test_refreshed_token_does_not_override_logout(self, mock_resolve_user,
mock_auth_manager):
+ mock_auth_manager.refresh_user.return_value = MagicMock(spec=BaseUser)
+ mock_auth_manager.generate_jwt.return_value = "new_token"
+ app, _ = _build_app()
+
+ response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"valid_token"}).get("/api/v2/auth/logout")
+
+ set_cookie_headers = response.headers.get_list("set-cookie")
+ assert len(set_cookie_headers) == 1
+ assert "Max-Age=0" in set_cookie_headers[0]
+ assert "new_token" not in set_cookie_headers[0]
+
+ @patch(f"{MODULE}.get_cookie_path", autospec=True, return_value="/team-a/")
+ def test_root_cookie_cleared_when_route_sets_token(
+ self, mock_cookie_path, mock_resolve_user, mock_auth_manager
+ ):
+ mock_resolve_user.side_effect = HTTPException(status_code=403,
detail="Invalid JWT token")
+ app, _ = _build_app()
+
+ response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"stale_token"}).get("/auth/login_callback")
+
+ set_cookie_headers = response.headers.get_list("set-cookie")
+ assert len(set_cookie_headers) == 2
+ assert set_cookie_headers[0] == "_token=issued_token; HttpOnly;
Path=/team-a/; SameSite=lax"
+ assert "Max-Age=0; Path=/;" in set_cookie_headers[1]
+
def test_refresh_user_runs_off_the_event_loop(self, mock_resolve_user,
mock_auth_manager):
"""``refresh_user`` is synchronous and may do blocking I/O, so it must
not block the loop."""
refresh_threads = []
