This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch v3-4-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit c2adc83d51b56d2ce6153f5cb542b15ee9fa1745 Author: Pierre Jeambrun <[email protected]> AuthorDate: Fri Nov 6 10:10:49 2026 +0100 [v3-4-test] Keep extra cookies set by a route in JWTRefreshMiddleware (#75869) * [v3-4-test] Keep release-candidate provider pins until each is released or yanked (#75919) (#75924) Constraints for an Airflow release candidate pin the provider wave being voted on at its rc versions, and commit them onto the shared constraints-X-Y branch. The next regular generation for the vX-Y-test branch resolves without pre-releases, so it saw the older finals as a provider downgrade and failed - the v3-4-test canary has been red since 3.4.0rc1 for exactly that reason, although nothing was yanked. The final release has the opposite gap: it ignored the candidates, so a provider rc that was neither released nor yanked was silently passed over and the released constraints never recorded that decision. It now has to be made, and the run says which providers it is still waiting on. The alert that should have reported the canary failure never reached Slack because the action was not told which API method to call, and the twice-daily CI notification announced "All passing" because it read a run that was still in progress, whose missing conclusion looked like no failures. (cherry picked from commit eeae004bd25e02c696f619f61ab67a3fb671c81c) Generated-by: Claude Opus 5 Co-authored-by: Jarek Potiuk <[email protected]> * Make JWTRefreshMiddleware extensible via BaseAuthManager.get_jwt_refresh_middleware() (#70783) Co-authored-by: Stephen Bracken <email-protected> (cherry picked from commit e5d8cbbe67931dd13a027ed3386ddcc1877800cd) * 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. (cherry picked from commit ea2aae77d1e8513c32881485601a86f0a865d2d6) --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Jarek Potiuk <[email protected]> Co-authored-by: stephen-bracken <[email protected]> Co-authored-by: Shahar Epstein <[email protected]> --- .../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, + }
