This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun 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 ea2aae77d1e Keep extra cookies set by a route in JWTRefreshMiddleware
(#75411)
ea2aae77d1e is described below
commit ea2aae77d1e8513c32881485601a86f0a865d2d6
Author: Pierre Jeambrun <[email protected]>
AuthorDate: Wed Oct 28 13:53:09 2026 +0100
Keep extra cookies set by a route in JWTRefreshMiddleware (#75411)
* Keep extra cookies set by a route in JWTRefreshMiddleware
The route-wins rule only covered the token cookie, so a cookie returned
by get_extra_cookies was still appended after the same cookie set by the
route. When the session was invalidated, that cleared the cookie a login
callback had just issued.
The generate_new_user_token test passed one-element tuples to a single
parametrized argument, so the case without a new user never ran.
* Split token and cleanup cookie handling in JWTRefreshMiddleware
* Let route cookies win by sending the middleware cookies first
Browsers key cookies on name and path, so matching the route's cookies by
name alone dropped our clear when the route set the same name on another
path. Sending our headers before the route's lets the later Set-Cookie win
for the same name and path, without parsing them.
* Never send a refreshed token when the route sets its own
Sending the refreshed token first and letting the route's logout cookie
override it still put a freshly minted, unrevoked JWT on the wire, and
clients such as Python's http.cookiejar keep it because they don't apply
Set-Cookie headers in order.
---
.../api_fastapi/auth/middlewares/refresh_token.py | 24 ++++--
.../auth/middlewares/test_refresh_token.py | 92 +++++++++++++++++-----
2 files changed, 90 insertions(+), 26 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 01aa96ffe79..ac175e98474 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
@@ -102,13 +102,20 @@ class JWTRefreshMiddleware:
if message["type"] == "http.response.start":
response_started = True
headers = MutableHeaders(scope=message)
- 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:
+ route_cookies = headers.getlist("set-cookie")
+ route_set_token = any(
+ value.startswith(f"{COOKIE_NAME_JWT_TOKEN}=") for value in
route_cookies
+ )
+ del headers["set-cookie"]
+ # The route owns the token it sets: never send ours, so a
refreshed token can't outlive a logout.
+ cookie_values = (
+ cleanup_cookies
+ if route_set_token or not token_cookie
+ else [token_cookie, *cleanup_cookies]
+ )
+ # Ours go first: for the same cookie name and path the later
Set-Cookie wins, so the
+ # cookies set by the route itself take precedence.
+ for value in [*cookie_values, *route_cookies]:
headers.append("set-cookie", value)
await send(message)
@@ -130,7 +137,8 @@ class JWTRefreshMiddleware:
"""
Get extra cookies from the request context to set in a new token
response.
- Setting a cookie to "" will cause it to be cleared by
build_cookie_headers()
+ Setting a cookie to "" will cause it to be cleared by
build_cookie_headers(). A cookie with the
+ same name and path set by the route itself takes precedence over the
one returned here.
"""
return {}
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 39aa7216b9d..9134ad08bb0 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
@@ -18,6 +18,7 @@
from __future__ import annotations
import threading
+from http.cookies import SimpleCookie
from unittest.mock import MagicMock, patch
import pytest
@@ -37,6 +38,15 @@ from tests_common.test_utils.config import conf_vars
MODULE = "airflow.api_fastapi.auth.middlewares.refresh_token"
+def _get_browser_cookies(response) -> dict[tuple[str, str], str | None]:
+ """Apply the ``Set-Cookie`` headers in order like a browser, keyed on name
and path; ``None`` is cleared."""
+ cookies: dict[tuple[str, str], str | None] = {}
+ for header in response.headers.get_list("set-cookie"):
+ for name, morsel in SimpleCookie(header).items():
+ cookies[(name, morsel["path"])] = None if morsel["max-age"] == "0"
else morsel.value
+ return cookies
+
+
def _build_app(static_dir=None, root_path: str = "") -> tuple[FastAPI, dict]:
"""Return an app wrapped by the middleware and a dict capturing what the
endpoint saw."""
seen: dict = {}
@@ -64,6 +74,19 @@ def _build_app(static_dir=None, root_path: str = "") ->
tuple[FastAPI, dict]:
)
return response
+ @app.get("/auth/login_callback_with_extra_cookie")
+ def login_callback_with_extra_cookie():
+ response = PlainTextResponse("logged in")
+ response.set_cookie(COOKIE_NAME_JWT_TOKEN, "issued_token",
path=refresh_token.get_cookie_path())
+ response.set_cookie("foo", "issued_foo",
path=refresh_token.get_cookie_path())
+ return response
+
+ @app.get("/auth/clear_root_extra_cookie")
+ def clear_root_extra_cookie():
+ response = PlainTextResponse("cleared")
+ response.delete_cookie("foo", path="/")
+ return response
+
@app.get("/api/v2/auth/logout")
def logout():
response = PlainTextResponse("logged out")
@@ -158,9 +181,7 @@ class TestJWTRefreshMiddleware:
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"
- ]
+ assert _get_browser_cookies(response) == {(COOKIE_NAME_JWT_TOKEN,
"/"): "issued_token"}
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)
@@ -169,10 +190,9 @@ class TestJWTRefreshMiddleware:
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]
+ assert _get_browser_cookies(response) == {(COOKIE_NAME_JWT_TOKEN,
"/"): None}
+ # Clients that don't apply Set-Cookie headers in order must not pick
up a token minted during logout
+ assert not any("new_token" in header for header in
response.headers.get_list("set-cookie"))
@patch(f"{MODULE}.get_cookie_path", autospec=True, return_value="/team-a/")
def test_root_cookie_cleared_when_route_sets_token(
@@ -183,10 +203,10 @@ class TestJWTRefreshMiddleware:
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]
+ assert _get_browser_cookies(response) == {
+ (COOKIE_NAME_JWT_TOKEN, "/team-a/"): "issued_token",
+ (COOKIE_NAME_JWT_TOKEN, "/"): None,
+ }
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."""
@@ -317,16 +337,14 @@ class TestJWTRefreshMiddleware:
)
@pytest.mark.asyncio
- @pytest.mark.parametrize("new_user", [(MagicMock(),), (None,)])
- async def test_generate_new_user_token(self, mock_auth_manager, new_user):
+ @pytest.mark.parametrize(
+ ("new_user", "expected_token"), [(MagicMock(spec=BaseUser),
"new_token"), (None, "")]
+ )
+ async def test_generate_new_user_token(self, mock_auth_manager, new_user,
expected_token):
mock_auth_manager.generate_jwt.return_value = "new_token"
app, seen = _build_app()
middleware = JWTRefreshMiddleware(app)
- new_token = await middleware.generate_new_user_token(MagicMock(),
new_user)
- if new_user is None:
- assert new_token == ""
- else:
- assert new_token == "new_token"
+ assert await middleware.generate_new_user_token(MagicMock(), new_user)
== expected_token
@patch(f"{MODULE}.JWTRefreshMiddleware.get_extra_cookies")
def test_refresh_user_with_extra_cookies(
@@ -375,3 +393,41 @@ class TestJWTRefreshMiddleware:
response.headers["set-cookie"]
== '_token=new_token; HttpOnly; Path=/; SameSite=lax, foo="";
HttpOnly; Max-Age=0; Path=/; SameSite=lax'
)
+
+ @patch(f"{MODULE}.JWTRefreshMiddleware.get_extra_cookies")
+ def test_invalid_token_does_not_clear_extra_cookie_set_by_route(
+ self, mock_get_extra_cookies: MagicMock, mock_resolve_user,
mock_auth_manager
+ ):
+ mock_get_extra_cookies.return_value = {"foo": "", "bar": ""}
+ 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_with_extra_cookie"
+ )
+
+ assert _get_browser_cookies(response) == {
+ (COOKIE_NAME_JWT_TOKEN, "/"): "issued_token",
+ ("foo", "/"): "issued_foo",
+ ("bar", "/"): None,
+ }
+
+ @patch(f"{MODULE}.get_cookie_path", autospec=True, return_value="/team-a/")
+ @patch(f"{MODULE}.JWTRefreshMiddleware.get_extra_cookies")
+ def
test_extra_cookie_cleared_on_cookie_path_when_route_clears_it_on_another_path(
+ self, mock_get_extra_cookies: MagicMock, mock_cookie_path,
mock_resolve_user, mock_auth_manager
+ ):
+ mock_get_extra_cookies.return_value = {"foo": ""}
+ 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/clear_root_extra_cookie"
+ )
+
+ assert _get_browser_cookies(response) == {
+ (COOKIE_NAME_JWT_TOKEN, "/team-a/"): None,
+ (COOKIE_NAME_JWT_TOKEN, "/"): None,
+ ("foo", "/team-a/"): None,
+ ("foo", "/"): None,
+ }
