vincbeck commented on code in PR #70783:
URL: https://github.com/apache/airflow/pull/70783#discussion_r3691076361


##########
airflow-core/src/airflow/api_fastapi/core_api/app.py:
##########
@@ -172,12 +172,15 @@ 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)
+    auth_manager = get_auth_manager()
 
-    for middleware_cls, middleware_kwargs in 
get_auth_manager().get_fastapi_middlewares():
+    jwt_refresh_middleware, jwt_refresh_middleware_kwargs = 
auth_manager.get_jwt_refresh_middleware()
+
+    app.add_middleware(jwt_refresh_middleware, **jwt_refresh_middleware_kwargs)
+
+    for middleware_cls, middleware_kwargs in 
auth_manager.get_fastapi_middlewares():

Review Comment:
   Why having 2 kind of middlewares? Why not adding the JWT refresh middleware 
as part of `get_fastapi_middlewares` by default?



##########
airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py:
##########
@@ -68,35 +69,65 @@ async def dispatch(self, request: Request, call_next):
             response = await call_next(request)
 
             if new_token is not None:
-                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,
+                response = await self._set_new_token(
+                    new_token, new_user, request_cookie_is_secure(request), 
response
                 )
-                # 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="/",
-                        httponly=True,
-                        secure=secure,
-                        samesite="lax",
-                    )
+
         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_token: str,
+        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_token: New JWT Token to set in cookies
+        :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()
+        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="/",
+                httponly=True,
+                secure=secure,
+                samesite="lax",
+            )
+        return response
+
     @staticmethod
-    async def _refresh_user(current_token: str) -> tuple[BaseUser | None, 
BaseUser | None]:
+    async def _refresh_user(current_token: str, request: Request) -> 
tuple[BaseUser | None, BaseUser | None]:

Review Comment:
   Why passing the request as parameter if you do not use it?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to