Fix logout in Fab and Keycloak auth managers

providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/openapi/v2-fab-auth-manager-generated.yaml

@@ -82,6 +82,19 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+  /auth/logout:
+    get:
+      tags:
+      - FabAuthManager
+      summary: Logout
+      description: Generate a new API token.
+      operationId: logout
+      responses:
+        '307':
+          description: Successful Response
+          content:
+            application/json:
+              schema: {}
   /auth/fab/v1/roles:
     post:
       tags:

providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/routes/login.py

@@ -17,7 +17,11 @@
 from __future__ import annotations
 
 from starlette import status
+from starlette.requests import Request  # noqa: TC002
+from starlette.responses import RedirectResponse
 
+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.common.router import AirflowRouter
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.configuration import conf
@@ -52,3 +56,26 @@ def create_token_cli(body: LoginBody) -> LoginResponse:
         return FABAuthManagerLogin.create_token(
             body=body, expiration_time_in_seconds=conf.getint("api_auth", "jwt_cli_expiration_time")
         )
+
+
+@login_router.get(
+    "/logout",
+    status_code=status.HTTP_307_TEMPORARY_REDIRECT,
+)
+def logout(request: Request) -> RedirectResponse:
+    """Generate a new API token."""
+    with get_application_builder():
+        login_url = get_auth_manager().get_url_login()
+        secure = request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback=""))
+        response = RedirectResponse(login_url)
+        response.delete_cookie(
+            key="session",
+            secure=secure,
+            httponly=True,
+        )
+        response.delete_cookie(
+            key=COOKIE_NAME_JWT_TOKEN,
+            secure=secure,
+            httponly=True,
+        )
+        return response

providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py

@@ -569,7 +569,7 @@ def get_url_login(self, **kwargs) -> str:
 
     def get_url_logout(self) -> str | None:
         """Return the logout page url."""
-        return urljoin(self.apiserver_endpoint, f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/logout/")
+        return urljoin(self.apiserver_endpoint, f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/logout")
 
     def register_views(self) -> None:
         self.security_manager.register_views()

providers/fab/tests/unit/fab/auth_manager/api_fastapi/conftest.py

@@ -28,7 +28,7 @@
 
 @pytest.fixture(scope="module")
 def fab_auth_manager():
-    return FabAuthManager(None)
+    return FabAuthManager()
 
 
 @pytest.fixture(scope="module")

providers/fab/tests/unit/fab/auth_manager/api_fastapi/routes/test_login.py

@@ -21,6 +21,7 @@
 
 import pytest
 
+from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
 from airflow.providers.fab.auth_manager.api_fastapi.datamodels.login import LoginBody, LoginResponse
 
 
@@ -50,3 +51,11 @@ def test_create_token_cli(self, mock_fab_auth_manager_login, test_client):
         )
         assert response.status_code == 201
         assert response.json()["access_token"] == self.dummy_token.access_token
+
+    def test_logout(self, test_client):
+        response = test_client.get("/logout", follow_redirects=False)
+        assert response.status_code == 307
+        assert response.headers["location"] == "/auth/login"
+        cookies = response.headers.get_list("set-cookie")
+        assert any("session=" in c for c in cookies)
+        assert any(f"{COOKIE_NAME_JWT_TOKEN}=" in c for c in cookies)

providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py

@@ -824,7 +824,7 @@ def test_get_url_login(self, auth_manager):
 
     def test_get_url_logout(self, auth_manager):
         result = auth_manager.get_url_logout()
-        assert result == f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/logout/"
+        assert result == f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/logout"
 
     @mock.patch.object(FabAuthManager, "_is_authorized", return_value=True)
     def test_get_extra_menu_items(self, _, auth_manager_with_appbuilder, flask_app):

providers/fab/www-hash.txt

@@ -1 +1 @@
-a5e7be19ee6766a961ef8cba57ebfad5bd550e448c7cbef54307dec7c301446e
+c45890ac6b17386adfb076a4baec17589a26aab189d6257a92138aaa720fbc7f

providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py

@@ -110,6 +110,10 @@ def get_url_login(self, **kwargs) -> str:
         base_url = conf.get("api", "base_url", fallback="/")
         return urljoin(base_url, f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/login")
 
+    def get_url_logout(self) -> str | None:
+        base_url = conf.get("api", "base_url", fallback="/")
+        return urljoin(base_url, f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/logout")
+
     def refresh_user(self, *, user: KeycloakAuthManagerUser) -> KeycloakAuthManagerUser | None:
         if self._token_expired(user.access_token):
             log.debug("Refreshing the token")

providers/keycloak/src/airflow/providers/keycloak/auth_manager/openapi/v2-keycloak-auth-manager-generated.yaml

@@ -32,6 +32,39 @@ paths:
           content:
             application/json:
               schema: {}
+  /auth/logout:
+    get:
+      tags:
+      - KeycloakAuthManagerLogin
+      summary: Logout
+      description: Log out the user from Keycloak.
+      operationId: logout
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema: {}
+      security:
+      - OAuth2PasswordBearer: []
+      - HTTPBearer: []
+  /auth/logout_callback:
+    get:
+      tags:
+      - KeycloakAuthManagerLogin
+      summary: Logout Callback
+      description: 'Complete the log-out.
+
+
+        This callback is redirected by Keycloak after the user has been logged out
+        from Keycloak.'
+      operationId: logout_callback
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema: {}
   /auth/refresh:
     get:
       tags:

providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py

@@ -20,7 +20,7 @@
 import logging
 from typing import Annotated
 
-from fastapi import Depends, HTTPException, Request
+from fastapi import Depends, Request
 from starlette.responses import HTMLResponse, RedirectResponse
 
 from airflow.api_fastapi.app import get_auth_manager
@@ -80,16 +80,46 @@ def login_callback(request: Request):
     return response
 
 
+@login_router.get("/logout")
+def logout(request: Request, user: Annotated[KeycloakAuthManagerUser, Depends(get_user)]):
+    """Log out the user from Keycloak."""
+    client = KeycloakAuthManager.get_keycloak_client()
+    keycloak_config = client.well_known()
+    end_session_endpoint = keycloak_config["end_session_endpoint"]
+
+    # Use the refresh flow to get the id token, it avoids us to save the id token
+    tokens = client.refresh_token(user.refresh_token)
+    post_logout_redirect_uri = request.url_for("logout_callback")
+    logout_url = f"{end_session_endpoint}?post_logout_redirect_uri={post_logout_redirect_uri}&id_token_hint={tokens['id_token']}"
+
+    return RedirectResponse(logout_url)
+
+
+@login_router.get("/logout_callback")
+def logout_callback(request: Request):
+    """
+    Complete the log-out.
+
+    This callback is redirected by Keycloak after the user has been logged out from Keycloak.
+    """
+    login_url = get_auth_manager().get_url_login()
+    secure = request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback=""))
+    response = RedirectResponse(login_url)
+    response.delete_cookie(
+        key=COOKIE_NAME_JWT_TOKEN,
+        secure=secure,
+        httponly=True,
+    )
+    return response
+
+
 @login_router.get("/refresh")
 def refresh(
     request: Request, user: Annotated[KeycloakAuthManagerUser, Depends(get_user)]
 ) -> RedirectResponse:
     """Refresh the token."""
     client = KeycloakAuthManager.get_keycloak_client()
 
-    if not user or not user.refresh_token:
-        raise HTTPException(status_code=400, detail="User is empty or has no refresh token")
-
     tokens = client.refresh_token(user.refresh_token)
     user.refresh_token = tokens["refresh_token"]
     user.access_token = tokens["access_token"]

providers/keycloak/tests/unit/keycloak/auth_manager/routes/conftest.py

@@ -16,7 +16,11 @@
 # under the License.
 from __future__ import annotations
 
+import datetime
+from typing import TYPE_CHECKING
+
 import pytest
+import time_machine
 from fastapi.testclient import TestClient
 
 from airflow.api_fastapi.app import create_app
@@ -26,9 +30,13 @@
     CONF_REALM_KEY,
     CONF_SECTION_NAME,
 )
+from airflow.providers.keycloak.auth_manager.user import KeycloakAuthManagerUser
 
 from tests_common.test_utils.config import conf_vars
 
+if TYPE_CHECKING:
+    from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import KeycloakAuthManager
+
 
 @pytest.fixture
 def client():
@@ -44,4 +52,21 @@ def client():
             (CONF_SECTION_NAME, "base_url"): "http://host.docker.internal:48080",
         }
     ):
-        yield TestClient(create_app())
+        app = create_app()
+        auth_manager: KeycloakAuthManager = app.state.auth_manager
+        time_very_before = datetime.datetime(2014, 1, 1, 0, 0, 0)
+        time_after = datetime.datetime.now() + datetime.timedelta(days=1)
+        with time_machine.travel(time_very_before, tick=False):
+            token = auth_manager._get_token_signer(
+                expiration_time_in_seconds=(time_after - time_very_before).total_seconds()
+            ).generate(
+                auth_manager.serialize_user(
+                    KeycloakAuthManagerUser(
+                        user_id="user_id",
+                        name="name",
+                        access_token="access_token",
+                        refresh_token="refresh_token",
+                    )
+                ),
+            )
+        yield TestClient(create_app(), headers={"Authorization": f"Bearer {token}"})

providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py

@@ -74,3 +74,41 @@ def test_login_callback(self, mock_get_keycloak_client, mock_get_auth_manager, c
     def test_login_callback_without_code(self, client):
         response = client.get(AUTH_MANAGER_FASTAPI_APP_PREFIX + "/login_callback")
         assert response.status_code == 400
+
+    @patch("airflow.providers.keycloak.auth_manager.routes.login.KeycloakAuthManager.get_keycloak_client")
+    def test_logout(self, mock_get_keycloak_client, client):
+        mock_keycloak_client = Mock()
+        mock_keycloak_client.well_known.return_value = {"end_session_endpoint": "logout_url"}
+        mock_keycloak_client.refresh_token.return_value = {"id_token": "id_token"}
+        mock_get_keycloak_client.return_value = mock_keycloak_client
+        response = client.get(AUTH_MANAGER_FASTAPI_APP_PREFIX + "/logout", follow_redirects=False)
+        assert response.status_code == 307
+        assert "location" in response.headers
+        assert (
+            response.headers["location"]
+            == "logout_url?post_logout_redirect_uri=http://testserver/auth/logout_callback&id_token_hint=id_token"
+        )
+        mock_keycloak_client.refresh_token.assert_called_once_with("refresh_token")
+
+    @patch("airflow.providers.keycloak.auth_manager.routes.login.get_auth_manager")
+    @patch("airflow.providers.keycloak.auth_manager.routes.login.KeycloakAuthManager.get_keycloak_client")
+    def test_refresh_token(self, mock_get_keycloak_client, mock_get_auth_manager, client):
+        mock_keycloak_client = Mock()
+        mock_keycloak_client.refresh_token.return_value = {
+            "access_token": "new_access_token",
+            "refresh_token": "new_refresh_token",
+        }
+        mock_get_keycloak_client.return_value = mock_keycloak_client
+
+        mock_auth_manager = Mock()
+        mock_auth_manager.generate_jwt.return_value = "token"
+        mock_get_auth_manager.return_value = mock_auth_manager
+
+        response = client.get(AUTH_MANAGER_FASTAPI_APP_PREFIX + "/refresh", follow_redirects=False)
+        assert response.status_code == 303
+        assert "location" in response.headers
+        assert response.headers["location"] == "/"
+        assert "_token" in response.cookies
+        assert response.cookies["_token"] == "token"
+        mock_keycloak_client.refresh_token.assert_called_once_with("refresh_token")
+        mock_auth_manager.generate_jwt.assert_called_once()

providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py

@@ -103,6 +103,10 @@ def test_get_url_login(self, auth_manager):
         result = auth_manager.get_url_login()
         assert result == f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/login"
 
+    def test_get_url_logout(self, auth_manager):
+        result = auth_manager.get_url_logout()
+        assert result == f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/logout"
+
     @patch.object(KeycloakAuthManager, "_token_expired")
     def test_refresh_user_not_expired(self, mock_token_expired, auth_manager):
         mock_token_expired.return_value = False