Move keycloak JWT tokens to separate cookies

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

@@ -25,3 +25,11 @@
 CONF_SERVER_URL_KEY = "server_url"
 CONF_REQUESTS_POOL_SIZE_KEY = "requests_pool_size"
 CONF_REQUESTS_RETRIES_KEY = "requests_retries"
+
+# Extra Cookie names
+COOKIE_NAME_ACCESS_TOKEN = "access_token"
+COOKIE_NAME_ID_TOKEN = "_id_token"
+COOKIE_NAME_NAME = "name"
+COOKIE_NAME_OAUTH_STATE = "_oauth_state"
+COOKIE_NAME_REFRESH_TOKEN = "refresh_token"
+COOKIE_NAME_USER_ID = "user_id"

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

@@ -23,11 +23,11 @@
 import warnings
 from base64 import urlsafe_b64decode
 from concurrent.futures import ThreadPoolExecutor
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Annotated, Any
 from urllib.parse import urljoin
 
 import requests
-from fastapi import FastAPI
+from fastapi import Cookie, FastAPI
 from keycloak import KeycloakOpenID
 from keycloak.exceptions import KeycloakPostError
 from requests.adapters import HTTPAdapter
@@ -92,6 +92,11 @@
 )
 
 
+def _get_keycloak_jwt(user: Annotated[KeycloakAuthManagerUser | None, Cookie(default=None)] = None):
+    """Populate Keycloak user from cookies."""
+    return user
+
+
 class KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
     """
     Keycloak auth manager.
@@ -129,19 +134,18 @@ def http_session(self) -> requests.Session:
         return self._http_session
 
     def deserialize_user(self, token: dict[str, Any]) -> KeycloakAuthManagerUser:
-        return KeycloakAuthManagerUser(
-            user_id=token.pop("user_id"),
-            name=token.pop("name"),
-            access_token=token.pop("access_token"),
-            refresh_token=token.pop("refresh_token"),
-        )
+        user = _get_keycloak_jwt()
+        if user is None:
+            raise ValueError("Couldn't deserialise user from Cookies.")
+        if user_id := token.pop("user_id"):
+            if user.get_id() != user_id:
+                raise ValueError("Keycloak user in Cookies does not match Airflow JWT.")
+        return user
 
     def serialize_user(self, user: KeycloakAuthManagerUser) -> dict[str, Any]:
         return {
             "user_id": user.get_id(),
             "name": user.get_name(),
-            "access_token": user.access_token,
-            "refresh_token": user.refresh_token,
         }
 
     def get_url_login(self, **kwargs) -> str:

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

@@ -47,15 +47,20 @@ class AuthManagerRefreshTokenExpiredException(Exception):  # type: ignore[no-red
 
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.providers.common.compat.sdk import conf
+from airflow.providers.keycloak.auth_manager.constants import (
+    COOKIE_NAME_ACCESS_TOKEN,
+    COOKIE_NAME_ID_TOKEN,
+    COOKIE_NAME_NAME,
+    COOKIE_NAME_OAUTH_STATE,
+    COOKIE_NAME_REFRESH_TOKEN,
+    COOKIE_NAME_USER_ID,
+)
 from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import KeycloakAuthManager
 from airflow.providers.keycloak.auth_manager.user import KeycloakAuthManagerUser
 
 log = logging.getLogger(__name__)
 login_router = AirflowRouter(tags=["KeycloakAuthManagerLogin"])
 
-COOKIE_NAME_ID_TOKEN = "_id_token"
-COOKIE_NAME_OAUTH_STATE = "_oauth_state"
-
 
 def _login_callback_url(request: Request) -> str:
     """
@@ -136,6 +141,20 @@ def login_callback(request: Request):
         COOKIE_NAME_ID_TOKEN, tokens["id_token"], path=cookie_path, secure=secure, httponly=True
     )
 
+    response.set_cookie(COOKIE_NAME_USER_ID, userinfo["sub"], path=cookie_path, secure=secure, httponly=True)
+
+    response.set_cookie(
+        COOKIE_NAME_NAME, userinfo["preferred_username"], path=cookie_path, secure=secure, httponly=True
+    )
+
+    response.set_cookie(
+        COOKIE_NAME_ACCESS_TOKEN, tokens["access_token"], path=cookie_path, secure=secure, httponly=True
+    )
+
+    response.set_cookie(
+        COOKIE_NAME_REFRESH_TOKEN, tokens["refresh_token"], path=cookie_path, secure=secure, httponly=True
+    )
+
     return response
 
 
@@ -186,4 +205,23 @@ def logout_callback(request: Request):
         secure=secure,
         httponly=True,
     )
+    response.delete_cookie(key=COOKIE_NAME_USER_ID, path=cookie_path, secure=secure, httponly=True)
+    response.delete_cookie(
+        key=COOKIE_NAME_NAME,
+        path=cookie_path,
+        secure=secure,
+        httponly=True,
+    )
+    response.delete_cookie(
+        key=COOKIE_NAME_ACCESS_TOKEN,
+        path=cookie_path,
+        secure=secure,
+        httponly=True,
+    )
+    response.delete_cookie(
+        key=COOKIE_NAME_REFRESH_TOKEN,
+        path=cookie_path,
+        secure=secure,
+        httponly=True,
+    )
     return response

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

@@ -16,17 +16,18 @@
 # under the License.
 from __future__ import annotations
 
+from pydantic import BaseModel
+
 from airflow.api_fastapi.auth.managers.models.base_user import BaseUser
 
 
-class KeycloakAuthManagerUser(BaseUser):
+class KeycloakAuthManagerUser(BaseModel, BaseUser):
     """User model for users managed by Keycloak auth manager."""
 
-    def __init__(self, *, user_id: str, name: str, access_token: str, refresh_token: str | None) -> None:
-        self.user_id = user_id
-        self.name = name
-        self.access_token = access_token
-        self.refresh_token = refresh_token
+    user_id: str
+    name: str
+    access_token: str
+    refresh_token: str | None
 
     def get_id(self) -> str:
         return self.user_id

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

@@ -122,20 +122,34 @@ def _clear_filter_cache():
 
 
 class TestKeycloakAuthManager:
-    def test_deserialize_user(self, auth_manager):
-        result = auth_manager.deserialize_user(
-            {
-                "user_id": "user_id",
-                "name": "name",
-                "access_token": "access_token",
-                "refresh_token": "refresh_token",
-            }
+    @patch("airflow.providers.keycloak.auth_manager.keycloak_auth_manager._get_keycloak_jwt")
+    def test_deserialize_user(self, mock_get_keycloak_jwt, auth_manager):
+        mock_get_keycloak_jwt.return_value = KeycloakAuthManagerUser(
+            user_id="user_id", name="name", access_token="access_token", refresh_token="refresh_token"
         )
+        result = auth_manager.deserialize_user({"user_id": "user_id", "name": "name"})
         assert result.user_id == "user_id"
         assert result.name == "name"
         assert result.access_token == "access_token"
         assert result.refresh_token == "refresh_token"
 
+    @patch("airflow.providers.keycloak.auth_manager.keycloak_auth_manager._get_keycloak_jwt")
+    def test_deserialize_user_missing(self, mock_get_keycloak_jwt, auth_manager):
+        mock_get_keycloak_jwt.return_value = None
+        with pytest.raises(ValueError, match="Couldn't deserialise user from Cookies."):
+            auth_manager.deserialize_user({"user_id": "user_id", "name": "name"})
+
+    @patch("airflow.providers.keycloak.auth_manager.keycloak_auth_manager._get_keycloak_jwt")
+    def test_deserialize_user_doesnt_match(self, mock_get_keycloak_jwt, auth_manager):
+        mock_get_keycloak_jwt.return_value = KeycloakAuthManagerUser(
+            user_id="user_2",
+            name="name",
+            access_token="access_token",
+            refresh_token="refresh_token",
+        )
+        with pytest.raises(ValueError, match="Keycloak user in Cookies does not match Airflow JWT."):
+            auth_manager.deserialize_user({"user_id": "user_id", "name": "name"})
+
     def test_serialize_user(self, auth_manager):
         result = auth_manager.serialize_user(
             KeycloakAuthManagerUser(
@@ -145,8 +159,6 @@ def test_serialize_user(self, auth_manager):
         assert result == {
             "user_id": "user_id",
             "name": "name",
-            "access_token": "access_token",
-            "refresh_token": "refresh_token",
         }
 
     def test_get_url_login(self, auth_manager):