vincbeck commented on code in PR #51657: URL: https://github.com/apache/airflow/pull/51657#discussion_r2161672045
########## airflow-core/src/airflow/api_fastapi/auth/managers/middleware/refresh_token.py: ########## @@ -0,0 +1,65 @@ +# 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 + +import logging + +from fastapi import Request +from starlette.middleware.base import BaseHTTPMiddleware + +from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_REFRESH_JWT_TOKEN, BaseAuthManager +from airflow.configuration import conf + +log = logging.getLogger(__name__) + + +class RefreshTokenMiddleware(BaseHTTPMiddleware): + """Middleware that automatically generates and includes auth header for simple auth manager.""" + + def __init__(self, app, auth_manager: BaseAuthManager): + super().__init__(app) + self.auth_manager = auth_manager + + async def dispatch(self, request: Request, call_next): + # Extract Authorization header + auth = request.headers.get("authorization") + + if auth and auth.lower().startswith("bearer "): + token_str = auth.split(" ", 1)[1] + user = await self.auth_manager.get_user_from_token(token_str) + user_refresh_token = getattr(user, "refresh_token", None) + user_access_token = getattr(user, "access_token", None) + if user_access_token and user_refresh_token: + if tokens := self.auth_manager.refresh_token(user_refresh_token): + access_token = tokens.get("access_token") + refresh_token = tokens.get("refresh_token") + # Update the Authorization header with the new token + setattr(user, "access_token", access_token) + setattr(user, "refresh_token", refresh_token) + new_token_with_updated_user = self.auth_manager.generate_jwt(user=user) Review Comment: I just realized the signature I gave you for the method `refresh_token` is wrong. Since we cannot assume the user object has properties such as `access_token` or `refresh_token` because it depends on the actual implementation of the auth manager, we should exchange the user and not tokens. I simplify A LOT this implementation ```suggestion user = await self.auth_manager.get_user_from_token(token_str) if refreshed_user := self.auth_manager.refresh_token(user=user): refreshed_token = self.auth_manager.generate_jwt(user=refreshed_user) ``` ########## airflow-core/src/airflow/api_fastapi/auth/managers/middleware/refresh_token.py: ########## @@ -0,0 +1,65 @@ +# 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 + +import logging + +from fastapi import Request +from starlette.middleware.base import BaseHTTPMiddleware + +from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_REFRESH_JWT_TOKEN, BaseAuthManager +from airflow.configuration import conf + +log = logging.getLogger(__name__) + + +class RefreshTokenMiddleware(BaseHTTPMiddleware): + """Middleware that automatically generates and includes auth header for simple auth manager.""" Review Comment: Please update the docstring ########## airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py: ########## @@ -133,6 +134,15 @@ def get_url_logout(self) -> str | None: """ return None + def refresh_token(self, refresh_token: str) -> dict | None: Review Comment: See the other comment thread for more details but the signature should actually be: ```suggestion def refresh_token(self, user: T) -> T | None: ``` ########## airflow-core/src/airflow/api_fastapi/auth/managers/middleware/refresh_token.py: ########## @@ -0,0 +1,65 @@ +# 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 + +import logging + +from fastapi import Request +from starlette.middleware.base import BaseHTTPMiddleware + +from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_REFRESH_JWT_TOKEN, BaseAuthManager +from airflow.configuration import conf + +log = logging.getLogger(__name__) + + +class RefreshTokenMiddleware(BaseHTTPMiddleware): + """Middleware that automatically generates and includes auth header for simple auth manager.""" + + def __init__(self, app, auth_manager: BaseAuthManager): + super().__init__(app) + self.auth_manager = auth_manager + + async def dispatch(self, request: Request, call_next): + # Extract Authorization header + auth = request.headers.get("authorization") + + if auth and auth.lower().startswith("bearer "): + token_str = auth.split(" ", 1)[1] + user = await self.auth_manager.get_user_from_token(token_str) + user_refresh_token = getattr(user, "refresh_token", None) + user_access_token = getattr(user, "access_token", None) + if user_access_token and user_refresh_token: + if tokens := self.auth_manager.refresh_token(user_refresh_token): + access_token = tokens.get("access_token") + refresh_token = tokens.get("refresh_token") + # Update the Authorization header with the new token + setattr(user, "access_token", access_token) + setattr(user, "refresh_token", refresh_token) + new_token_with_updated_user = self.auth_manager.generate_jwt(user=user) + response = await call_next(request) + # Set the new access token in the cookies to update the state of the user in the token + secure = bool(conf.get("api", "ssl_cert", fallback="")) + response.set_cookie( + COOKIE_NAME_REFRESH_JWT_TOKEN, new_token_with_updated_user, secure=secure Review Comment: ```suggestion COOKIE_NAME_JWT_TOKEN, new_token_with_updated_user, secure=secure ``` ########## airflow-core/src/airflow/api_fastapi/app.py: ########## @@ -96,6 +96,7 @@ def create_app(apps: str = "all") -> FastAPI: 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_generic_middlewares(app) Review Comment: I feel it confusing to have `init_middlewares` and `init_generic_middlewares`. I think it would be better to add the new middleware in the existing `init_middlewares` ########## airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py: ########## @@ -133,6 +134,15 @@ def get_url_logout(self) -> str | None: """ return None + def refresh_token(self, refresh_token: str) -> dict | None: + """ + Refresh the JWT token. + + This method is called when the JWT token is about to expire. It should return a new JWT token if the + user is still authenticated, or None if the user is not authenticated anymore. Review Comment: ```suggestion This method is called when the JWT token expired. Returns the user if the token were being refreshed, None otherwise. ``` ########## airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py: ########## @@ -73,6 +73,7 @@ COOKIE_NAME_JWT_TOKEN = "_token" +COOKIE_NAME_REFRESH_JWT_TOKEN = "_refresh_token" Review Comment: No need to use a different cookie name, we should use the existing one `COOKIE_NAME_JWT_TOKEN`. Then we need to modify the UI to override the token when provided -- 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]
