Update UI authentication

Author: vincbeck

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/auth/managers/base_auth_manager.py

@@ -141,12 +141,15 @@ def get_url_logout(self) -> str | None:
         """
         return None
 
-    def get_url_refresh(self) -> str | None:
+    def refresh_user(self, *, user: T) -> T | None:
         """
-        Return the URL to refresh the authentication token.
+        Refresh the user if needed.
 
-        This is used to refresh the authentication token when it expires.
-        The default implementation returns None, which means that the auth manager does not support refresh token.
+        By default, does nothing. Some auth managers might need to refresh the user to, for instance,
+        refresh some tokens that are needed to communicate with a service/tool.
+
+        This method is called by every single request, it must be lightweight otherwise the overall API
+        server latency will increase.
         """
         return None
 

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
 

…stapi/auth/managers/simple/middleware.py …api_fastapi/auth/middlewares/__init__.pyairflow-core/src/airflow/api_fastapi/auth/managers/simple/middleware.py renamed to airflow-core/src/airflow/api_fastapi/auth/middlewares/__init__.py airflow-core/src/airflow/api_fastapi/auth/managers/simple/middleware.py renamed to airflow-core/src/airflow/api_fastapi/auth/middlewares/__init__.py

@@ -1,3 +1,4 @@
+#
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
@@ -14,19 +15,3 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
-from __future__ import annotations
-
-from fastapi import Request
-from starlette.middleware.base import BaseHTTPMiddleware
-
-from airflow.api_fastapi.auth.managers.simple.services.login import SimpleAuthManagerLogin
-
-
-class SimpleAllAdminMiddleware(BaseHTTPMiddleware):
-    """Middleware that automatically generates and includes auth header for simple auth manager."""
-
-    async def dispatch(self, request: Request, call_next):
-        token = SimpleAuthManagerLogin.create_token_all_admins()
-        request.scope["headers"].append((b"authorization", f"Bearer {token}".encode()))
-        return await call_next(request)

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

@@ -0,0 +1,76 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from fastapi import Request
+from jwt import ExpiredSignatureError, InvalidTokenError
+from starlette.middleware.base import BaseHTTPMiddleware
+
+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.configuration import conf
+
+
+class JWTRefreshMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware to handle JWT token refresh.
+
+    This middleware:
+    1. Extracts JWT token from cookies and build the user from the token
+    2. Calls ``refresh_user`` method from auth manager with the user
+    3. If ``refresh_user`` returns a user, generate a JWT token based upon this user and send it in the
+       response as cookie
+    """
+
+    async def dispatch(self, request: Request, call_next):
+        new_user = None
+        current_token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+        if current_token:
+            new_user = await self._refresh_user(current_token)
+            if new_user:
+                request.state.user = new_user
+
+        response = await call_next(request)
+
+        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)
+            secure = bool(conf.get("api", "ssl_cert", fallback=""))
+            response.set_cookie(
+                COOKIE_NAME_JWT_TOKEN,
+                new_token,
+                httponly=True,
+                secure=secure,
+                samesite="lax",
+            )
+
+        return response
+
+    @staticmethod
+    async def _refresh_user(current_token: str) -> BaseUser | None:
+        try:
+            # If the token is expired or not valid, then we stop the refresh token flow
+            # This will be caught and handled by the route downstream
+            user = await get_auth_manager().get_user_from_token(current_token)
+        except ExpiredSignatureError:
+            return None
+        except InvalidTokenError:
+            return None
+
+        return get_auth_manager().refresh_user(user=user)

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

@@ -182,14 +182,9 @@ def init_error_handlers(app: FastAPI) -> None:
 
 
 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
+    from airflow.api_fastapi.auth.middlewares.refresh_token import JWTRefreshMiddleware
 
-        app.add_middleware(SimpleAllAdminMiddleware)
+    app.add_middleware(JWTRefreshMiddleware)
 
 
 def init_ui_plugins(app: FastAPI) -> None:

airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml

@@ -8440,40 +8440,6 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
-  /api/v2/auth/refresh:
-    get:
-      tags:
-      - Login
-      summary: Refresh
-      description: Refresh the authentication token.
-      operationId: refresh
-      parameters:
-      - name: next
-        in: query
-        required: false
-        schema:
-          anyOf:
-          - type: string
-          - type: 'null'
-          title: Next
-      responses:
-        '200':
-          description: Successful Response
-          content:
-            application/json:
-              schema: {}
-        '307':
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/HTTPExceptionResponse'
-          description: Temporary Redirect
-        '422':
-          description: Validation Error
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/HTTPValidationError'
 components:
   schemas:
     AppBuilderMenuItemResponse:

airflow-core/src/airflow/api_fastapi/core_api/routes/public/auth.py

@@ -22,7 +22,6 @@
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.core_api.security import is_safe_url
-from airflow.configuration import conf
 
 auth_router = AirflowRouter(tags=["Login"], prefix="/auth")
 
@@ -56,23 +55,3 @@ def logout(request: Request, next: None | str = None) -> RedirectResponse:
         logout_url = request.app.state.auth_manager.get_url_login()
 
     return RedirectResponse(logout_url)
-
-
-@auth_router.get(
-    "/refresh",
-    responses=create_openapi_http_exception_doc([status.HTTP_307_TEMPORARY_REDIRECT]),
-)
-def refresh(request: Request, next: None | str = None) -> RedirectResponse:
-    """Refresh the authentication token."""
-    refresh_url = request.app.state.auth_manager.get_url_refresh()
-
-    if not refresh_url:
-        return RedirectResponse(f"{conf.get('api', 'base_url', fallback='/')}auth/logout")
-
-    if next and not is_safe_url(next, request=request):
-        raise HTTPException(status_code=400, detail="Invalid or unsafe next URL")
-
-    if next:
-        refresh_url += f"?next={next}"
-
-    return RedirectResponse(refresh_url)

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/openapi-gen/queries/common.ts

@@ -762,12 +762,6 @@ export const useLoginServiceLogoutKey = "LoginServiceLogout";
 export const UseLoginServiceLogoutKeyFn = ({ next }: {
   next?: string;
 } = {}, queryKey?: Array<unknown>) => [useLoginServiceLogoutKey, ...(queryKey ?? [{ next }])];
-export type LoginServiceRefreshDefaultResponse = Awaited<ReturnType<typeof LoginService.refresh>>;
-export type LoginServiceRefreshQueryResult<TData = LoginServiceRefreshDefaultResponse, TError = unknown> = UseQueryResult<TData, TError>;
-export const useLoginServiceRefreshKey = "LoginServiceRefresh";
-export const UseLoginServiceRefreshKeyFn = ({ next }: {
-  next?: string;
-} = {}, queryKey?: Array<unknown>) => [useLoginServiceRefreshKey, ...(queryKey ?? [{ next }])];
 export type AuthLinksServiceGetAuthMenusDefaultResponse = Awaited<ReturnType<typeof AuthLinksService.getAuthMenus>>;
 export type AuthLinksServiceGetAuthMenusQueryResult<TData = AuthLinksServiceGetAuthMenusDefaultResponse, TError = unknown> = UseQueryResult<TData, TError>;
 export const useAuthLinksServiceGetAuthMenusKey = "AuthLinksServiceGetAuthMenus";