Keycloak auth manager: keep API tokens usable with the Authorization

airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py

@@ -191,6 +191,22 @@ def generate_jwt(
             self.serialize_user(user)
         )
 
+    def generate_api_jwt(
+        self, user: T, *, expiration_time_in_seconds: int = conf.getint("api_auth", "jwt_expiration_time")
+    ) -> str:
+        """
+        Return the JWT token for a client that authenticates with the ``Authorization`` header.
+
+        Such a client sends no cookies, so an auth manager that keeps part of its state in
+        cookies has to put that state in the token's claims instead for the request to be
+        authorized. Auth managers whose tokens are already self-contained need not override
+        this.
+
+        :param user: the user to generate the token for
+        :param expiration_time_in_seconds: expiration time in seconds of the token
+        """
+        return self.generate_jwt(user, expiration_time_in_seconds=expiration_time_in_seconds)
+
     @abstractmethod
     def get_url_login(self, **kwargs) -> str:
         """Return the login page url."""

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

@@ -149,7 +149,10 @@ def deserialize_user(self, token: dict[str, Any]) -> KeycloakAuthManagerUser:
 
     def serialize_user(self, user: KeycloakAuthManagerUser) -> dict[str, Any]:
         if AIRFLOW_V_3_3_PLUS:
-            # Omit Keycloak JWTs from claims, they are stored in separate cookies
+            # Omit Keycloak JWTs from claims, they are stored in separate cookies.
+            # That keeps the browser session cookie under the 4096 byte limit. Tokens
+            # minted for API clients are never stored in a cookie and keep the JWTs in
+            # their claims instead -- see ``generate_api_jwt``.
             return {
                 "user_id": user.get_id(),
                 "name": user.get_name(),
@@ -161,6 +164,31 @@ def serialize_user(self, user: KeycloakAuthManagerUser) -> dict[str, Any]:
             "refresh_token": user.refresh_token,
         }
 
+    def generate_api_jwt(
+        self,
+        user: KeycloakAuthManagerUser,
+        *,
+        expiration_time_in_seconds: int = conf.getint("api_auth", "jwt_expiration_time"),
+    ) -> str:
+        """
+        Return a JWT for a client that authenticates with the ``Authorization`` header.
+
+        Such a client sends no cookies, so the Keycloak tokens have to travel in the
+        claims for the request to be authorized.
+
+        :param user: the user to generate the token for
+        :param expiration_time_in_seconds: expiration time in seconds of the token
+        """
+        return self._get_token_signer(expiration_time_in_seconds=expiration_time_in_seconds).generate(
+            {
+                # Build on serialize_user so an API token cannot silently miss a claim
+                # that the browser flow gained; only the Keycloak JWTs differ.
+                **self.serialize_user(user),
+                "access_token": user.access_token,
+                "refresh_token": user.refresh_token,
+            }
+        )
+
     async def get_user_from_token(
         self, token: str, access_token: str | None = None, refresh_token: str | None = None
     ):
@@ -186,8 +214,13 @@ async def get_user_from_token(
             user.access_token = access_token
             user.refresh_token = refresh_token
             return user
-        # Skip refreshing JWT if Keycloak JWTs are not included.
-        return None
+        # No cookie-supplied tokens. A token minted for an API client carries the
+        # Keycloak JWTs in its own claims, so the user is already complete -- and unlike
+        # the cookie path those claims are covered by the Airflow JWT signature, so they
+        # need no separate subject check. A browser token does not carry them and can
+        # only be completed by KeycloakJWTMiddleware from the cookies; without them
+        # there is nothing to authorize against.
+        return user if user.access_token else None
 
     def get_url_login(self, **kwargs) -> str:
         base_url = conf.get("api", "base_url", fallback="/")

providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/token.py

@@ -58,7 +58,7 @@ def create_token_for(
         refresh_token=tokens["refresh_token"],
     )
 
-    return get_auth_manager().generate_jwt(user, expiration_time_in_seconds=expiration_time_in_seconds)
+    return get_auth_manager().generate_api_jwt(user, expiration_time_in_seconds=expiration_time_in_seconds)
 
 
 def create_client_credentials_token(
@@ -118,4 +118,4 @@ def create_client_credentials_token(
         ),  # client_credentials may not return refresh_token (RFC6749 section 4.4.3)
     )
 
-    return get_auth_manager().generate_jwt(user, expiration_time_in_seconds=expiration_time_in_seconds)
+    return get_auth_manager().generate_api_jwt(user, expiration_time_in_seconds=expiration_time_in_seconds)

providers/keycloak/tests/unit/keycloak/auth_manager/services/test_token.py

@@ -55,9 +55,13 @@ def test_create_token(self, mock_get_keycloak_client, mock_get_auth_manager):
         mock_get_keycloak_client.return_value = mock_keycloak_client
         mock_auth_manager = Mock()
         mock_get_auth_manager.return_value = mock_auth_manager
-        mock_auth_manager.generate_jwt.return_value = self.token
+        mock_auth_manager.generate_api_jwt.return_value = self.token
 
         assert create_token_for(username=self.test_username, password=self.test_password) == self.token
+        # API tokens must be minted with the Keycloak JWTs in their claims, since a
+        # header-authenticated client sends no cookies to supply them.
+        mock_auth_manager.generate_api_jwt.assert_called_once()
+        mock_auth_manager.generate_jwt.assert_not_called()
         mock_keycloak_client.token.assert_called_once_with(self.test_username, self.test_password)
         mock_keycloak_client.userinfo.assert_called_once_with(self.test_access_token)
 
@@ -104,7 +108,7 @@ def test_create_token_client_credentials(self, mock_get_keycloak_client, mock_ge
         mock_get_keycloak_client.return_value = mock_keycloak_client
         mock_auth_manager = Mock()
         mock_get_auth_manager.return_value = mock_auth_manager
-        mock_auth_manager.generate_jwt.return_value = self.token
+        mock_auth_manager.generate_api_jwt.return_value = self.token
 
         result = create_client_credentials_token(client_id=test_client_id, client_secret=test_client_secret)
 

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

@@ -255,6 +255,61 @@ async def test_get_user_from_token(self, auth_manager):
         assert user.access_token == "access_token"
         assert user.refresh_token == "refresh_token"
 
+    @pytest.mark.asyncio
+    async def test_api_token_authenticates_without_cookies(self, auth_manager):
+        """A token from ``POST /auth/token`` must authorize an ``Authorization`` header request.
+
+        Such a client sends no cookies, so ``get_user_from_token`` is called with the
+        token alone and the user has to be reconstructible from its claims.
+        """
+        access_token = keycloak_token("user_id")
+        user = KeycloakAuthManagerUser(
+            user_id="user_id", name="name", access_token=access_token, refresh_token="refresh_token"
+        )
+
+        minted_claims: dict = {}
+
+        class _CapturingSigner:
+            def generate(self, claims):
+                minted_claims.update(claims)
+                return "token"
+
+        with patch.object(KeycloakAuthManager, "_get_token_signer", Mock(return_value=_CapturingSigner())):
+            auth_manager.generate_api_jwt(user)
+
+        mock_token_validator = Mock()
+        mock_token_validator.avalidated_claims = AsyncMock(return_value=minted_claims)
+        with patch.object(
+            KeycloakAuthManager, "_get_token_validator", Mock(return_value=mock_token_validator)
+        ):
+            resolved = await auth_manager.get_user_from_token("token")
+
+        assert resolved is not None
+        assert resolved.get_id() == "user_id"
+        assert resolved.access_token == access_token
+        assert resolved.refresh_token == "refresh_token"
+
+    @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Browser tokens omit the Keycloak JWTs")
+    @pytest.mark.asyncio
+    async def test_browser_token_without_cookies_is_not_authenticated(self, auth_manager):
+        """A browser token carries no Keycloak JWTs, so it cannot stand in for the cookies."""
+        browser_claims = auth_manager.serialize_user(
+            KeycloakAuthManagerUser(
+                user_id="user_id",
+                name="name",
+                access_token=keycloak_token("user_id"),
+                refresh_token="refresh_token",
+            )
+        )
+        assert "access_token" not in browser_claims
+
+        mock_token_validator = Mock()
+        mock_token_validator.avalidated_claims = AsyncMock(return_value=browser_claims)
+        with patch.object(
+            KeycloakAuthManager, "_get_token_validator", Mock(return_value=mock_token_validator)
+        ):
+            assert await auth_manager.get_user_from_token("token") is None
+
     @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Uses KeycloakJWTMiddleware and separate cookies")
     @pytest.mark.asyncio
     async def test_get_user_from_token_keycloak_jwts_missing(self, auth_manager):