Make JWTRefreshMiddleware extensible via BaseAuthManager.get_jwt_refresh_middleware()

airflow-core/src/airflow/api_fastapi/app.py

@@ -60,6 +60,22 @@ def get_cookie_path() -> str:
     return API_ROOT_PATH or "/"
 
 
+def request_cookie_is_secure(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"]
 

airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py

@@ -180,7 +180,17 @@ def get_fastapi_middlewares(self) -> list[tuple[_MiddlewareFactory[Any], dict[st
         unauthenticated requests when public access is configured) should override this
         method.
         """
-        return []
+        return [self.get_jwt_refresh_middleware()]
+
+    def get_jwt_refresh_middleware(self) -> tuple[_MiddlewareFactory[Any], dict[str, Any]]:
+        """
+        Return the JWTRefreshMiddleware to refresh the Airflow JWT token.
+
+        :important: The JWTRefreshMiddleware should be included in get_fastapi_middlewares()
+        """
+        from airflow.api_fastapi.auth.middlewares.refresh_token import JWTRefreshMiddleware
+
+        return JWTRefreshMiddleware, {}
 
     def generate_jwt(
         self, user: T, *, expiration_time_in_seconds: int = conf.getint("api_auth", "jwt_expiration_time")

airflow-core/src/airflow/api_fastapi/auth/managers/simple/simple_auth_manager.py

@@ -396,11 +396,12 @@ def is_authorized_hitl_task(self, *, assigned_users: set[str], user: SimpleAuthM
 
     def get_fastapi_middlewares(self) -> list[tuple[_MiddlewareFactory[Any], dict[str, Any]]]:
         """Register the all-admins middleware when ``[core] simple_auth_manager_all_admins`` is set."""
+        middleware = super().get_fastapi_middlewares()
         if not conf.getboolean("core", "simple_auth_manager_all_admins"):
-            return []
+            return middleware
         from airflow.api_fastapi.auth.managers.simple.middleware import SimpleAllAdminMiddleware
 
-        return [(SimpleAllAdminMiddleware, {})]
+        return middleware + [(SimpleAllAdminMiddleware, {})]
 
     def get_fastapi_app(self) -> FastAPI | None:
         """

airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py

@@ -17,19 +17,24 @@
 # under the License.
 from __future__ import annotations
 
-from fastapi import HTTPException, Request
+from typing import TYPE_CHECKING
+
+from fastapi import HTTPException
 from fastapi.responses import JSONResponse
 from starlette.middleware.base import BaseHTTPMiddleware
 
-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 fastapi import Request, Response
+
+    from airflow.api_fastapi.auth.managers.models.base_user import BaseUser
 
 
 class JWTRefreshMiddleware(BaseHTTPMiddleware):
@@ -45,58 +50,114 @@ class JWTRefreshMiddleware(BaseHTTPMiddleware):
 
     async def dispatch(self, request: Request, call_next):
         new_token = None
-        current_token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+        current_user = None
+        new_user = None
         try:
-            if current_token is not None:
-                try:
-                    new_user, current_user = await self._refresh_user(current_token)
-                    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 = ""
+            try:
+                new_user, current_user = await self._refresh_user(request)
+            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 = ""
+
+            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
 
             response = await call_next(request)
 
-            if new_token is not None:
+            if new_user or new_token is not None:
+                secure = request_cookie_is_secure(request)
                 cookie_path = get_cookie_path()
-                secure = request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback=""))
-                response.set_cookie(
-                    COOKIE_NAME_JWT_TOKEN,
-                    new_token,
-                    path=cookie_path,
-                    httponly=True,
-                    secure=secure,
-                    samesite="lax",
-                    max_age=0 if new_token == "" else None,
-                )
-                # Clear any stale _token cookie at root path "/".
-                # Older Airflow instances may have set the cookie there;
-                # without this, the root-path cookie keeps being sent on
-                # every request, causing an infinite redirect loop.
-                if cookie_path != "/":
-                    response.delete_cookie(
-                        key=COOKIE_NAME_JWT_TOKEN,
-                        path="/",
+                if new_token == "":
+                    response.set_cookie(
+                        COOKIE_NAME_JWT_TOKEN,
+                        new_token,
+                        path=cookie_path,
                         httponly=True,
                         secure=secure,
                         samesite="lax",
+                        max_age=0,
                     )
+                    if cookie_path != "/":
+                        response.set_cookie(
+                            key=COOKIE_NAME_JWT_TOKEN,
+                            path="/",
+                            httponly=True,
+                            secure=secure,
+                            samesite="lax",
+                            max_age=0,
+                        )
+                else:
+                    response = await self._set_new_token(new_user, secure, response, cookie_path)
+
         except HTTPException as exc:
             # If any HTTPException is raised during user resolution or refresh, return it as response
             return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})
         return response
 
+    @classmethod
+    async def _set_new_token(
+        cls,
+        new_user: BaseUser | None,
+        secure: bool,
+        response: Response,
+        cookie_path: str | None = None,
+    ) -> Response:
+        """
+        Set Cookies in the response based on a new JWT token and a new user model.
+
+        :param new_user: User model for the JWT token
+        :param secure: HTTP secure property for cookies
+        :param response: FastAPI response object to set the cookies on
+        :param cookie_path: Path for cookies in the response
+        """
+        if cookie_path is None:
+            cookie_path = get_cookie_path()
+        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)
+        else:
+            new_token = ""
+        response.set_cookie(
+            COOKIE_NAME_JWT_TOKEN,
+            new_token,
+            path=cookie_path,
+            httponly=True,
+            secure=secure,
+            samesite="lax",
+            max_age=0 if new_token == "" else None,
+        )
+        # Clear any stale _token cookie at root path "/".
+        # Older Airflow instances may have set the cookie there;
+        # without this, the root-path cookie keeps being sent on
+        # every request, causing an infinite redirect loop.
+        if cookie_path != "/":
+            response.set_cookie(
+                key=COOKIE_NAME_JWT_TOKEN,
+                path="/",
+                httponly=True,
+                secure=secure,
+                samesite="lax",
+                max_age=0,
+            )
+        return response
+
     @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)
+        if not current_token:
+            return None, None
         user = await resolve_user_from_token(current_token)
         return get_auth_manager().refresh_user(user=user), user

airflow-core/src/airflow/api_fastapi/core_api/app.py

@@ -172,11 +172,8 @@ 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
     from airflow.api_fastapi.common.http_access_log import HttpAccessLogMiddleware
 
-    app.add_middleware(JWTRefreshMiddleware)
-
     for middleware_cls, middleware_kwargs in get_auth_manager().get_fastapi_middlewares():
         app.add_middleware(middleware_cls, **middleware_kwargs)
 

airflow-core/tests/unit/api_fastapi/auth/managers/simple/test_simple_auth_manager.py

@@ -511,13 +511,16 @@ def test_get_teams(self, auth_manager):
 
     @conf_vars({("core", "simple_auth_manager_all_admins"): "false"})
     def test_get_fastapi_middlewares_disabled(self, auth_manager):
-        assert auth_manager.get_fastapi_middlewares() == []
+        assert auth_manager.get_fastapi_middlewares() == [auth_manager.get_jwt_refresh_middleware()]
 
     @conf_vars({("core", "simple_auth_manager_all_admins"): "true"})
     def test_get_fastapi_middlewares_enabled(self, auth_manager):
         from airflow.api_fastapi.auth.managers.simple.middleware import SimpleAllAdminMiddleware
 
-        assert auth_manager.get_fastapi_middlewares() == [(SimpleAllAdminMiddleware, {})]
+        assert auth_manager.get_fastapi_middlewares() == [
+            auth_manager.get_jwt_refresh_middleware(),
+            (SimpleAllAdminMiddleware, {}),
+        ]
 
     def test_generate_password_uses_expected_alphabet_and_length(self):
         alphabet = set("abcdefghkmnpqrstuvwxyzABCDEFGHKMNPQRSTUVWXYZ23456789")

airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py

@@ -33,6 +33,7 @@
     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.exceptions import RemovedInAirflow4Warning
@@ -817,3 +818,6 @@ 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_fastapi_middlewares(self, auth_manager):
+        assert auth_manager.get_fastapi_middlewares() == [(JWTRefreshMiddleware, {})]