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


##########
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:
   I see



-- 
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