[v3-1-test] Update authentication to handle JWT token in backend

airflow-core/docs/core-concepts/auth-manager/index.rst

@@ -166,7 +166,7 @@ cookie named ``_token`` before redirecting to the Airflow UI. The Airflow UI wil
     response = RedirectResponse(url="/")
 
     secure = request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback=""))
-    response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure)
+    response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure, httponly=True)
     return response
 
 .. note::

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

@@ -29,7 +29,6 @@
     init_config,
     init_error_handlers,
     init_flask_plugins,
-    init_middlewares,
     init_ui_plugins,
     init_views,
 )
@@ -97,7 +96,6 @@ def create_app(apps: str = "all") -> FastAPI:
         init_ui_plugins(app)
         init_views(app)  # Core views need to be the last routes added - it has a catch all route
         init_error_handlers(app)
-        init_middlewares(app)
 
     init_config(app)
 

airflow-core/src/airflow/api_fastapi/auth/managers/simple/routes/login.py

@@ -94,6 +94,7 @@ def login_all_admins(request: Request) -> RedirectResponse:
         COOKIE_NAME_JWT_TOKEN,
         SimpleAuthManagerLogin.create_token_all_admins(),
         secure=secure,
+        httponly=True,
     )
     return response
 

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

@@ -181,17 +181,6 @@ def init_error_handlers(app: FastAPI) -> None:
         app.add_exception_handler(handler.exception_cls, handler.exception_handler)
 
 
-def init_middlewares(app: FastAPI) -> None:
-    from airflow.configuration import conf
-
-    if "SimpleAuthManager" in conf.get("core", "auth_manager") and conf.getboolean(
-        "core", "simple_auth_manager_all_admins"
-    ):
-        from airflow.api_fastapi.auth.managers.simple.middleware import SimpleAllAdminMiddleware
-
-        app.add_middleware(SimpleAllAdminMiddleware)
-
-
 def init_ui_plugins(app: FastAPI) -> None:
     """Initialize UI plugins."""
     from airflow import plugins_manager

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

@@ -27,6 +27,7 @@
 from pydantic import NonNegativeInt
 
 from airflow.api_fastapi.app import get_auth_manager
+from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
 from airflow.api_fastapi.auth.managers.models.base_user import BaseUser
 from airflow.api_fastapi.auth.managers.models.batch_apis import (
     IsAuthorizedConnectionRequest,
@@ -96,14 +97,22 @@ async def resolve_user_from_token(token_str: str | None) -> BaseUser:
 
 
 async def get_user(
+    request: Request,
     oauth_token: str | None = Depends(oauth2_scheme),
     bearer_credentials: HTTPAuthorizationCredentials | None = Depends(bearer_scheme),
 ) -> BaseUser:
-    token_str = None
+    # A user might have been already built by a middleware, if so, it is stored in `request.state.user`
+    user: BaseUser | None = getattr(request.state, "user", None)
+    if user:
+        return user
+
+    token_str: str | None
     if bearer_credentials and bearer_credentials.scheme.lower() == "bearer":
         token_str = bearer_credentials.credentials
     elif oauth_token:
         token_str = oauth_token
+    else:
+        token_str = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
 
     return await resolve_user_from_token(token_str)
 

airflow-core/src/airflow/ui/src/layouts/Nav/LogoutModal.tsx

@@ -22,7 +22,6 @@ import { useTranslation } from "react-i18next";
 
 import { ConfirmationModal } from "src/components/ConfirmationModal";
 import { getRedirectPath } from "src/utils/links.ts";
-import { TOKEN_STORAGE_KEY } from "src/utils/tokenHandler";
 
 type LogoutModalProps = {
   readonly isOpen: boolean;
@@ -38,7 +37,7 @@ const LogoutModal: React.FC<LogoutModalProps> = ({ isOpen, onClose }) => {
       onConfirm={() => {
         const logoutPath = getRedirectPath("api/v2/auth/logout");
 
-        localStorage.removeItem(TOKEN_STORAGE_KEY);
+        document.cookie = "_token=; Path=/; Max-Age=-99999999;";
         globalThis.location.replace(logoutPath);
       }}
       onOpenChange={onClose}

airflow-core/src/airflow/ui/src/main.tsx

@@ -37,7 +37,6 @@ import { getRedirectPath } from "src/utils/links.ts";
 import i18n from "./i18n/config";
 import { client } from "./queryClient";
 import { system } from "./theme";
-import { clearToken, tokenHandler } from "./utils/tokenHandler";
 
 // Set React, ReactDOM, and ReactJSXRuntime on globalThis to share them with the dynamically imported React plugins.
 // Only one instance of React should be used.
@@ -55,7 +54,6 @@ axios.interceptors.response.use(
       error.response?.status === 401 ||
       (error.response?.status === 403 && error.response.data.detail === "Invalid JWT token")
     ) {
-      clearToken();
       const params = new URLSearchParams();
 
       params.set("next", globalThis.location.href);
@@ -68,8 +66,6 @@ axios.interceptors.response.use(
   },
 );
 
-axios.interceptors.request.use(tokenHandler);
-
 createRoot(document.querySelector("#root") as HTMLDivElement).render(
   <StrictMode>
     <I18nextProvider i18n={i18n}>

airflow-core/tests/unit/api_fastapi/core_api/test_security.py

@@ -23,6 +23,7 @@
 from jwt import ExpiredSignatureError, InvalidTokenError
 
 from airflow.api_fastapi.app import create_app
+from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
 from airflow.api_fastapi.auth.managers.models.resource_details import (
     ConnectionDetails,
     DagAccessEntity,
@@ -35,6 +36,7 @@
 from airflow.api_fastapi.core_api.datamodels.pools import PoolBody
 from airflow.api_fastapi.core_api.datamodels.variables import VariableBody
 from airflow.api_fastapi.core_api.security import (
+    get_user,
     is_safe_url,
     requires_access_connection,
     requires_access_connection_bulk,
@@ -104,6 +106,46 @@ async def test_get_user_expired_token(self, mock_get_auth_manager):
 
         auth_manager.get_user_from_token.assert_called_once_with(token_str)
 
+    @patch("airflow.api_fastapi.core_api.security.resolve_user_from_token")
+    async def test_get_user_with_request_state(self, mock_resolve_user_from_token):
+        user = Mock()
+        request = Mock()
+        request.state.user = user
+
+        result = await get_user(request, None, None)
+
+        assert result == user
+        mock_resolve_user_from_token.assert_not_called()
+
+    @pytest.mark.parametrize(
+        "oauth_token, bearer_credentials_creds, cookies, expected",
+        [
+            ("oauth_token", None, {}, "oauth_token"),
+            (None, "bearer_credentials_creds", {}, "bearer_credentials_creds"),
+            (None, None, {COOKIE_NAME_JWT_TOKEN: "cookie_token"}, "cookie_token"),
+        ],
+    )
+    @patch("airflow.api_fastapi.core_api.security.resolve_user_from_token")
+    async def test_get_user_with_token(
+        self, mock_resolve_user_from_token, oauth_token, bearer_credentials_creds, cookies, expected
+    ):
+        user = Mock()
+        mock_resolve_user_from_token.return_value = user
+
+        request = Mock()
+        request.state.user = None
+        request.cookies = cookies
+        bearer_credentials = None
+        if bearer_credentials_creds:
+            bearer_credentials = Mock()
+            bearer_credentials.scheme = "bearer"
+            bearer_credentials.credentials = bearer_credentials_creds
+
+        result = await get_user(request, oauth_token, bearer_credentials)
+
+        assert result == user
+        mock_resolve_user_from_token.assert_called_once_with(expected)
+
     @pytest.mark.db_test
     @patch("airflow.api_fastapi.core_api.security.get_auth_manager")
     async def test_requires_access_dag_authorized(self, mock_get_auth_manager):

dev/breeze/src/airflow_breeze/files/simple_auth_manager_passwords.json

@@ -1 +1 @@
-{"admin": "admin", "viewer": "viewer"}
+{"admin": "admin", "viewer":  "viewer"}

providers/amazon/src/airflow/providers/amazon/aws/auth_manager/routes/login.py

@@ -35,6 +35,7 @@
 from airflow.providers.amazon.aws.auth_manager.constants import CONF_SAML_METADATA_URL_KEY, CONF_SECTION_NAME
 from airflow.providers.amazon.aws.auth_manager.datamodels.login import LoginResponse
 from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_1_1_PLUS
 
 try:
     from onelogin.saml2.auth import OneLogin_Saml2_Auth
@@ -101,7 +102,12 @@ def login_callback(request: Request):
     if relay_state == "login-redirect":
         response = RedirectResponse(url=url, status_code=303)
         secure = bool(conf.get("api", "ssl_cert", fallback=""))
-        response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure)
+        # In Airflow 3.1.1 authentication changes, front-end no longer handle the token
+        # See https://github.com/apache/airflow/pull/55506
+        if AIRFLOW_V_3_1_1_PLUS:
+            response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure, httponly=True)
+        else:
+            response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure)
         return response
     if relay_state == "login-token":
         return LoginResponse(access_token=token)

providers/amazon/src/airflow/providers/amazon/version_compat.py

@@ -34,6 +34,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
 
 AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
 AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0)
+AIRFLOW_V_3_1_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 1)
 
 if AIRFLOW_V_3_1_PLUS:
     from airflow.sdk import BaseHook

providers/fab/src/airflow/providers/fab/version_compat.py

@@ -33,3 +33,4 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
 
 
 AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0)
+AIRFLOW_V_3_1_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 1)