diff --git a/airflow-core/src/airflow/api_fastapi/app.py b/airflow-core/src/airflow/api_fastapi/app.py index 4925537302433..3c2c94114a503 100644 --- a/airflow-core/src/airflow/api_fastapi/app.py +++ b/airflow-core/src/airflow/api_fastapi/app.py @@ -45,6 +45,8 @@ from airflow.utils.providers_configuration_loader import providers_configuration_loaded if TYPE_CHECKING: + from fastapi import Request + from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager API_BASE_URL = conf.get("api", "base_url", fallback="") @@ -65,6 +67,22 @@ def get_cookie_path() -> str: return API_ROOT_PATH or "/" +def request_cookie_is_secure(request: Request) -> bool: + """ + Return whether a FastAPI request cookie should be tagged as HTTP secure. + + :param request: FastAPI Request object + + usage: + ```python + secure = request_cookie_is_secure(request) + + response.set_cookie("mycookie", "myvalue", path=get_cookie_path(), secure=secure, httponly=True) + ``` + """ + return request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback="")) + + # Fast API apps mounted under these prefixes are not allowed RESERVED_URL_PREFIXES = ["/api/v2", "/ui", "/execution", "/auth", "/pluginsv2"] diff --git a/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py b/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py index 94b96ba5af27b..8b2302393a023 100644 --- a/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py +++ b/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py @@ -185,6 +185,20 @@ def get_fastapi_middlewares(self) -> list[tuple[_MiddlewareFactory[Any], dict[st """ return [] + def get_jwt_refresh_middleware(self) -> list[tuple[_MiddlewareFactory[Any], dict[str, Any]]]: + """ + Return the JWTRefreshMiddleware to refresh the Airflow JWT token. + + Middlewares added by this method will be initialised before middleware from + ``get_fastapi_middlewares()``. To disable JWT refresh return an empty list. + To add more middleware, simply add more tuples here in the same format as + ``get_fastapi_middlewares()``. + """ + from airflow.api_fastapi.auth.middlewares.refresh_token import JWTRefreshMiddleware + from airflow.api_fastapi.core_api.app import _STATIC_URL_PATH + + return [(JWTRefreshMiddleware, {"skip_path_prefixes": (f"{_STATIC_URL_PATH}/",)})] + def generate_jwt( self, user: T, *, expiration_time_in_seconds: int = conf.getint("api_auth", "jwt_expiration_time") ) -> str: 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 37186d9be53d7..01aa96ffe7948 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 @@ -24,19 +24,19 @@ from starlette.concurrency import run_in_threadpool from starlette.datastructures import MutableHeaders -from airflow.api_fastapi.app import get_auth_manager, get_cookie_path +from airflow.api_fastapi.app import get_auth_manager, get_cookie_path, request_cookie_is_secure from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN from airflow.api_fastapi.auth.managers.exceptions import AuthManagerRefreshTokenExpiredException -from airflow.api_fastapi.auth.managers.models.base_user import BaseUser from airflow.api_fastapi.core_api.security import ( USER_INJECTED_BY_TRUSTED_MIDDLEWARE, resolve_user_from_token, ) -from airflow.configuration import conf if TYPE_CHECKING: from starlette.types import ASGIApp, Message, Receive, Scope, Send + from airflow.api_fastapi.auth.managers.models.base_user import BaseUser + # Unauthenticated endpoints that never read the user, so resolving it there would only cost a # JWT decode and a revocation lookup on every health check and UI boot. UNAUTHENTICATED_PATHS = frozenset(["/api/v2/monitor/health", "/api/v2/version"]) @@ -72,27 +72,29 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: await self.app(scope, receive, send) return - new_token = None + new_user = None + session_invalidated = False try: - new_user, current_user = await self._refresh_user(current_token) + new_user, current_user = await self.refresh_user(request) if user := (new_user or current_user): # Stamp the trust sentinel alongside the user so `get_user()` # can distinguish this trusted assignment from a stray write # by unrelated middleware. request.state.user = user request.state.user_authenticated_via = USER_INJECTED_BY_TRUSTED_MIDDLEWARE - if new_user: - # If we created a new user, serialize it and set it as a cookie - new_token = get_auth_manager().generate_jwt(new_user) except (HTTPException, AuthManagerRefreshTokenExpiredException): # Receive a HTTPException when the Airflow token is expired # Receive a AuthManagerRefreshTokenExpiredException when the potential underlying refresh # token used by the auth manager is expired - new_token = "" + session_invalidated = True - token_cookie, cleanup_cookies = ( - self._build_cookie_headers(request, new_token) if new_token is not None else (None, []) - ) + if session_invalidated or new_user is not None: + extra_cookies = self.get_extra_cookies(request, new_user) + token_cookie, cleanup_cookies = self.build_cookie_headers( + request, await self.generate_new_user_token(request, new_user), extra_cookies + ) + else: + token_cookie, cleanup_cookies = (None, []) response_started = False async def send_with_cookies(message: Message) -> None: @@ -124,10 +126,34 @@ def _is_skipped(self, scope: Scope) -> bool: return path in UNAUTHENTICATED_PATHS or path.startswith(self.skip_path_prefixes) @staticmethod - 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.""" + def get_extra_cookies(request: Request, new_user: BaseUser | None) -> dict[str, str]: + """ + 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() + """ + return {} + + @staticmethod + async def generate_new_user_token(request: Request, new_user: BaseUser | None) -> str: + """Generate a new Airflow JWT token from the user model.""" + return get_auth_manager().generate_jwt(new_user) if new_user else "" + + @staticmethod + def build_cookie_headers( + request: Request, new_token: str, extra_cookies: dict[str, str] + ) -> tuple[str, list[str]]: + """ + Return the ``Set-Cookie`` value of the token and those clearing stale token cookies. + + Any cookies with a nil value ("") will be cleared with Max-Age=0 + + :param request: FastAPI Request + :param new_token: New Airflow JWT token + :param extra_cookies: Additional cookie values to set + """ cookie_path = get_cookie_path() - secure = request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback="")) + secure = request_cookie_is_secure(request) # Only used to render the Set-Cookie header values response = Response() response.set_cookie( @@ -151,13 +177,32 @@ def _build_cookie_headers(request: Request, new_token: str) -> tuple[str, list[s secure=secure, samesite="lax", ) + # Add any additional cookies required + for cookie, value in extra_cookies.items(): + response.set_cookie( + cookie, + value, + path=cookie_path, + httponly=True, + secure=secure, + samesite="lax", + max_age=0 if value == "" else None, + ) 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]: + async def refresh_user(request: Request) -> tuple[BaseUser | None, BaseUser | None]: + """ + Refresh the logged in user using the current JWT Token. + + If the user is not authenticated, return ``None, None`` + + :param request: FastAPI Request + """ + current_token = request.cookies.get(COOKIE_NAME_JWT_TOKEN) user = await resolve_user_from_token(current_token) # Auth managers may do blocking I/O here (e.g. calling an identity provider) return await run_in_threadpool(get_auth_manager().refresh_user, user=user), user diff --git a/airflow-core/src/airflow/api_fastapi/core_api/app.py b/airflow-core/src/airflow/api_fastapi/core_api/app.py index 7b2985746e11b..836ba40561077 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/app.py @@ -304,11 +304,13 @@ def init_config(app: FastAPI) -> None: def init_middlewares(app: FastAPI) -> None: from airflow.api_fastapi.app import get_auth_manager - from airflow.api_fastapi.auth.middlewares.refresh_token import JWTRefreshMiddleware - app.add_middleware(JWTRefreshMiddleware, skip_path_prefixes=(f"{_STATIC_URL_PATH}/",)) + auth_manager = get_auth_manager() - for middleware_cls, middleware_kwargs in get_auth_manager().get_fastapi_middlewares(): + for middleware_cls, middleware_kwargs in ( + *auth_manager.get_jwt_refresh_middleware(), + *auth_manager.get_fastapi_middlewares(), + ): app.add_middleware(middleware_cls, **middleware_kwargs) # GZipMiddleware must be inside HttpMetricsMiddleware and HttpAccessLogMiddleware so that both diff --git a/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py b/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py index 18a9ceeecd2c8..e433c1223da7f 100644 --- a/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py +++ b/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py @@ -36,8 +36,10 @@ TeamDetails, VariableDetails, ) +from airflow.api_fastapi.auth.middlewares.refresh_token import JWTRefreshMiddleware from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator from airflow.api_fastapi.common.types import MenuItem +from airflow.api_fastapi.core_api.app import _STATIC_URL_PATH from airflow.exceptions import RemovedInAirflow4Warning from airflow.models.pool import Pool from airflow.models.team import Team @@ -932,3 +934,11 @@ def test_is_authorized_hitl_task( user = BaseAuthManagerUserTest(name=user_id) result = auth_manager.is_authorized_hitl_task(assigned_users=assigned_users, user=user) assert result == expected + + def test_get_jwt_refresh_middleware(self, auth_manager): + assert auth_manager.get_jwt_refresh_middleware() == [ + (JWTRefreshMiddleware, {"skip_path_prefixes": (f"{_STATIC_URL_PATH}/",)}) + ] + + def test_get_fastapi_middlewares(self, auth_manager): + assert auth_manager.get_fastapi_middlewares() == [] 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 a601a0417d17a..39aa7216b9d38 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 @@ -32,6 +32,8 @@ from airflow.api_fastapi.auth.middlewares.refresh_token import JWTRefreshMiddleware from airflow.api_fastapi.core_api.security import USER_INJECTED_BY_TRUSTED_MIDDLEWARE +from tests_common.test_utils.config import conf_vars + MODULE = "airflow.api_fastapi.auth.middlewares.refresh_token" @@ -97,9 +99,8 @@ def mock_resolve_user(): @pytest.fixture(autouse=True) def no_ssl_cert(): - with patch(f"{MODULE}.conf", autospec=True) as mock_conf: - mock_conf.get.return_value = "" - yield mock_conf + with conf_vars({("api", "ssl_cert"): ""}): + yield class TestJWTRefreshMiddleware: @@ -272,18 +273,18 @@ def test_static_files_served_unchanged( ], ) def test_cookie_secure_flag( - self, mock_resolve_user, mock_auth_manager, no_ssl_cert, scheme, ssl_cert, expected_secure + self, mock_resolve_user, mock_auth_manager, scheme, ssl_cert, expected_secure ): """The cookie Secure flag must follow the request scheme as well as the local ssl_cert.""" mock_auth_manager.refresh_user.return_value = MagicMock(spec=BaseUser) mock_auth_manager.generate_jwt.return_value = "new_token" - no_ssl_cert.get.return_value = ssl_cert - app, _ = _build_app() - client = TestClient( - app, base_url=f"{scheme}://testserver", cookies={COOKIE_NAME_JWT_TOKEN: "valid_token"} - ) + with conf_vars({("api", "ssl_cert"): ssl_cert}): + app, _ = _build_app() + client = TestClient( + app, base_url=f"{scheme}://testserver", cookies={COOKIE_NAME_JWT_TOKEN: "valid_token"} + ) - response = client.get("/api/v2/dags") + response = client.get("/api/v2/dags") assert ("Secure" in response.headers["set-cookie"]) is expected_secure @@ -314,3 +315,63 @@ def test_cookie_uses_subpath_and_clears_root_cookie( assert any( "Path=/" in h and "Path=/team-a/" not in h and "Max-Age=0" in h for h in set_cookie_headers ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("new_user", [(MagicMock(),), (None,)]) + async def test_generate_new_user_token(self, mock_auth_manager, new_user): + 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" + + @patch(f"{MODULE}.JWTRefreshMiddleware.get_extra_cookies") + def test_refresh_user_with_extra_cookies( + self, mock_get_extra_cookies: MagicMock, mock_resolve_user, mock_auth_manager + ): + mock_get_extra_cookies.return_value = {"foo": "bar"} + refreshed_user = MagicMock(spec=BaseUser) + mock_auth_manager.refresh_user.return_value = refreshed_user + mock_auth_manager.generate_jwt.return_value = "new_token" + app, seen = _build_app() + + response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN: "valid_token"}).get("/api/v2/dags") + + # JWTRefreshMiddleware must stamp the trust sentinel so `get_user()` accepts + # the cached user. Without the marker, the auth dependency falls through + # to a fresh JWT validation, defeating the refresh. + assert seen["user"] is refreshed_user + assert seen["authenticated_via"] is USER_INJECTED_BY_TRUSTED_MIDDLEWARE + mock_auth_manager.refresh_user.assert_called_once_with(user=mock_resolve_user.return_value) + mock_auth_manager.generate_jwt.assert_called_once_with(refreshed_user) + assert ( + response.headers["set-cookie"] + == "_token=new_token; HttpOnly; Path=/; SameSite=lax, foo=bar; HttpOnly; Path=/; SameSite=lax" + ) + + @patch(f"{MODULE}.JWTRefreshMiddleware.get_extra_cookies") + def test_refresh_user_clear_extra_cookies( + self, mock_get_extra_cookies: MagicMock, mock_resolve_user, mock_auth_manager + ): + mock_get_extra_cookies.return_value = {"foo": ""} + refreshed_user = MagicMock(spec=BaseUser) + mock_auth_manager.refresh_user.return_value = refreshed_user + mock_auth_manager.generate_jwt.return_value = "new_token" + app, seen = _build_app() + + response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN: "valid_token"}).get("/api/v2/dags") + + # JWTRefreshMiddleware must stamp the trust sentinel so `get_user()` accepts + # the cached user. Without the marker, the auth dependency falls through + # to a fresh JWT validation, defeating the refresh. + assert seen["user"] is refreshed_user + assert seen["authenticated_via"] is USER_INJECTED_BY_TRUSTED_MIDDLEWARE + mock_auth_manager.refresh_user.assert_called_once_with(user=mock_resolve_user.return_value) + mock_auth_manager.generate_jwt.assert_called_once_with(refreshed_user) + assert ( + response.headers["set-cookie"] + == '_token=new_token; HttpOnly; Path=/; SameSite=lax, foo=""; HttpOnly; Max-Age=0; Path=/; SameSite=lax' + ) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_auth.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_auth.py index fc4e84621e420..6fd0143388fb2 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_auth.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_auth.py @@ -367,7 +367,7 @@ def test_logout_clears_cookie_when_logout_url_redirects(self, logout_client): # The refresh middleware is stubbed out so that only the logout route can clear the cookie. with ( patch( - "airflow.api_fastapi.auth.middlewares.refresh_token.JWTRefreshMiddleware._refresh_user", + "airflow.api_fastapi.auth.middlewares.refresh_token.JWTRefreshMiddleware.refresh_user", new=AsyncMock(return_value=(None, None)), ), patch.object(auth_manager, "get_url_logout", return_value="http://external/logout"), @@ -389,7 +389,7 @@ def test_logout_clears_cookie_when_revocation_cannot_be_recorded(self, logout_cl logout_client.cookies.set(COOKIE_NAME_JWT_TOKEN, token_str) with ( patch( - "airflow.api_fastapi.auth.middlewares.refresh_token.JWTRefreshMiddleware._refresh_user", + "airflow.api_fastapi.auth.middlewares.refresh_token.JWTRefreshMiddleware.refresh_user", new=AsyncMock(return_value=(None, None)), ), patch.object(auth_manager, "get_url_logout", return_value="http://external/logout"), diff --git a/airflow-core/tests/unit/api_fastapi/core_api/test_app.py b/airflow-core/tests/unit/api_fastapi/core_api/test_app.py index 3d26d63263966..ae76598776081 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/test_app.py @@ -20,10 +20,12 @@ import inspect import json import typing +from unittest.mock import MagicMock, Mock, patch import pytest from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware +from fastapi.middleware.gzip import GZipMiddleware from fastapi.params import Depends as DependsClass from fastapi.responses import StreamingResponse from fastapi.templating import Jinja2Templates @@ -31,9 +33,11 @@ from starlette.testclient import TestClient from airflow.api_fastapi.app import create_app +from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager from airflow.api_fastapi.core_api.app import ( base_react_app_bundle_urls, init_config, + init_middlewares, init_ui_translation_views, ) from airflow.api_fastapi.core_api.routes.public import authenticated_router @@ -384,3 +388,28 @@ def test_index_escapes_a_bundle_url(self): html = self._render("index.html", ['/x.js" onload="alert(1)']) assert 'onload="alert(1)"' not in html + + +class TestMiddlewareInitialisation: + @patch("airflow.api_fastapi.app.get_auth_manager") + def test_jwt_refresh_middleware_init(self, mock_get_auth_manager): + app = FastAPI() + auth_manager = MagicMock(spec=BaseAuthManager) + mock_jwt_middleware = Mock() + auth_manager.get_jwt_refresh_middleware.return_value = [(mock_jwt_middleware, {})] + auth_manager.get_fastapi_middlewares.return_value = [] + mock_get_auth_manager.return_value = auth_manager + init_middlewares(app) + auth_manager.get_jwt_refresh_middleware.assert_called_once() + assert [m.cls for m in app.user_middleware] == [GZipMiddleware, mock_jwt_middleware] + + @patch("airflow.api_fastapi.app.get_auth_manager") + def test_no_jwt_refresh_middleware_init(self, mock_get_auth_manager): + app = FastAPI() + auth_manager = MagicMock(spec=BaseAuthManager) + auth_manager.get_jwt_refresh_middleware.return_value = [] + auth_manager.get_fastapi_middlewares.return_value = [] + mock_get_auth_manager.return_value = auth_manager + init_middlewares(app) + auth_manager.get_jwt_refresh_middleware.assert_called_once() + assert [m.cls for m in app.user_middleware] == [GZipMiddleware] diff --git a/airflow-core/tests/unit/api_fastapi/test_app.py b/airflow-core/tests/unit/api_fastapi/test_app.py index c5d5e33222f93..149179ebd626f 100644 --- a/airflow-core/tests/unit/api_fastapi/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/test_app.py @@ -345,6 +345,22 @@ def test_nested_subpath(self): assert app_module.get_cookie_path() == "/org/team-a/airflow/" +@pytest.mark.parametrize( + ("scheme", "ssl_cert", "exp"), + [ + ("http", "", False), + ("https", "", True), + ("http", "/foo/bar/cert.crt", True), + ("https", "/foo/bar/cert/crt", True), + ], +) +def test_request_cookie_is_secure(scheme, ssl_cert, exp): + request = mock.Mock() + request.base_url.scheme = scheme + with conf_vars({("api", "ssl_cert"): ssl_cert}): + assert app_module.request_cookie_is_secure(request) == exp + + def test_create_auth_manager_thread_safety(): """Concurrent calls to create_auth_manager must return the same singleton instance.""" call_count = 0