diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/constants.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/constants.py index 605ae62c62d01..85749f2650049 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/constants.py +++ b/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" diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py index fd7bfaf899078..9755d5742793f 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py +++ b/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: diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py index 58386b9aada48..95758f6122fd7 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py +++ b/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 diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/user.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/user.py index c2ec7d59c8883..75ee2219a67b9 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/user.py +++ b/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 diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py b/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py index ca4bda4f9e57f..58698292bc941 100644 --- a/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py +++ b/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):