Bind Keycloak cookie tokens to the Airflow session identity (#72207)

* Bind Keycloak cookie tokens to the Airflow session identity

For Airflow 3.3+ the Keycloak access and refresh tokens are no longer carried in
the signed Airflow JWT; they travel in separate _access_token and _refresh_token
cookies. get_user_from_token validated the Airflow JWT and then attached whatever
those cookies contained, without checking that they described the same subject.

A caller could therefore pair their own Airflow session with another subject's
Keycloak token. Every authorization decision goes to Keycloak carrying that
token, so the effective privileges were the token's, while get_id() and
get_name() - used for the session identity, audit records and logging - stayed
those of the Airflow JWT.

The access token's sub is now compared against the user id the signed JWT
established before the token is attached. Both are the Keycloak subject: every
place a KeycloakAuthManagerUser is constructed sets user_id from userinfo[sub],
in the interactive login, the password grant and the client_credentials grant
alike. A token whose payload cannot be read yields no subject and so matches
nothing.

The subject is read without signature verification, which is sufficient here:
the value is only ever compared against an identity the signed Airflow JWT has
already established, a forged token is refused by Keycloak when presented, and a
genuine token belonging to somebody else is what the comparison exists to catch.

The two existing tests passed the literal string "access_token" as a cookie
value; they now build a JWT-shaped token naming the same subject. Adds coverage
for a token naming another subject and for one that cannot be parsed.

* Refuse malformed Keycloak access-token cookies with 403, not 500

A Keycloak access-token cookie whose payload decodes to valid JSON that is
not an object reached the subject lookup as a non-mapping, so reading the
claim raised an error the middleware does not translate. The cookie is
attacker-supplied, so any shape it can take has to end in the same refusal
as a token naming the wrong subject.

* Drop the Keycloak changelog note about token-to-session binding

Every place a session is established sets the Airflow user id from the
Keycloak subject, so the two can only disagree in a request whose cookies
were assembled by hand. No deployment reaches the new refusal by ordinary
use, which leaves the note describing a change nobody observes.

Author: potiuk

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

@@ -29,6 +29,7 @@
 
 import requests
 from fastapi import FastAPI
+from jwt import InvalidTokenError
 from keycloak import KeycloakOpenID
 from keycloak.exceptions import KeycloakPostError
 from requests.adapters import HTTPAdapter
@@ -174,6 +175,14 @@ async def get_user_from_token(
         if not AIRFLOW_V_3_3_PLUS:
             return user
         if access_token:
+            # The Airflow JWT is signed and establishes who the caller is. The Keycloak
+            # tokens arrive in separate cookies that the signature does not cover, so
+            # pairing them unchecked would let a caller combine their own Airflow session
+            # with somebody else's Keycloak token: every authorization decision is then
+            # made for that subject, while the session identity, audit trail and logs
+            # continue to show this one.
+            if self._token_subject(access_token) != user.get_id():
+                raise InvalidTokenError("Keycloak access token does not belong to this Airflow session")
             user.access_token = access_token
             user.refresh_token = refresh_token
             return user
@@ -818,6 +827,29 @@ def _get_headers(access_token):
             "Content-Type": "application/x-www-form-urlencoded",
         }
 
+    @staticmethod
+    def _token_subject(token: str) -> str | None:
+        """
+        Return the ``sub`` claim of a JWT without verifying its signature.
+
+        :meta private:
+
+        The value is only ever compared against an identity the signed Airflow JWT has
+        already established, so it is never trusted on its own. A forged token is
+        rejected by Keycloak when it is presented; a genuine token belonging to somebody
+        else is exactly what this comparison exists to catch. A token that cannot be
+        parsed yields ``None``, which matches no user id.
+
+        :param token: the token
+        """
+        try:
+            payload_b64 = token.split(".")[1] + "=="
+            payload = json.loads(urlsafe_b64decode(payload_b64))
+            subject = payload["sub"]
+        except (IndexError, KeyError, TypeError, ValueError):
+            return None
+        return str(subject) if subject is not None else None
+
     @staticmethod
     def _token_expired(token: str) -> bool:
         """

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

@@ -23,6 +23,7 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
+from jwt import InvalidTokenError
 from keycloak import KeycloakPostError
 
 from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
@@ -132,6 +133,21 @@ def _clear_filter_cache():
     cache_module._pending_requests.clear()
 
 
+def token_with_payload(payload: str) -> str:
+    """Build a JWT-shaped token whose payload segment is the given raw text.
+
+    Only the payload segment is read when the token is bound to the Airflow session
+    identity, so the header and signature are placeholders.
+    """
+    encoded = base64.urlsafe_b64encode(payload.encode()).decode().rstrip("=")
+    return f"header.{encoded}.signature"
+
+
+def keycloak_token(subject: str) -> str:
+    """Build a JWT-shaped Keycloak token carrying ``sub``."""
+    return token_with_payload(json.dumps({"sub": subject}))
+
+
 class TestKeycloakAuthManager:
     @pytest.mark.parametrize(
         ("token_data", "exp"),
@@ -207,11 +223,12 @@ async def test_get_user_from_token_with_keycloak_tokens(self, auth_manager):
                 mock_get_user_from_token,
             ),
         ):
-            user = await auth_manager.get_user_from_token("token", "access_token", "refresh_token")
+            access_token = keycloak_token("user_id")
+            user = await auth_manager.get_user_from_token("token", access_token, "refresh_token")
         mock_get_user_from_token.assert_called_with("token")
         assert user.get_id() == "user_id"
         assert user.get_name() == "name"
-        assert user.access_token == "access_token"
+        assert user.access_token == access_token
         assert user.refresh_token == "refresh_token"
 
     @pytest.mark.skipif(AIRFLOW_V_3_3_PLUS, reason="Testing Old Keycloak JWT flow.")
@@ -270,13 +287,54 @@ async def test_get_user_from_token_keycloak_jwt(self, auth_manager):
                 mock_get_user_from_token,
             ),
         ):
-            user = await auth_manager.get_user_from_token("token", "access_token", "refresh_token")
+            access_token = keycloak_token("user_id")
+            user = await auth_manager.get_user_from_token("token", access_token, "refresh_token")
         mock_get_user_from_token.assert_called_with("token")
         assert user.get_id() == "user_id"
         assert user.get_name() == "name"
-        assert user.access_token == "access_token"
+        assert user.access_token == access_token
         assert user.refresh_token == "refresh_token"
 
+    @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_rejects_another_subjects_token(self, auth_manager):
+        """A Keycloak token naming a different subject must not attach to this session."""
+        mock_get_user_from_token = AsyncMock(
+            return_value=KeycloakAuthManagerUser(
+                user_id="user_id", name="name", access_token="", refresh_token=None
+            )
+        )
+        with (
+            patch.object(BaseAuthManager, "get_user_from_token", mock_get_user_from_token),
+            pytest.raises(InvalidTokenError, match="does not belong to this Airflow session"),
+        ):
+            await auth_manager.get_user_from_token("token", keycloak_token("someone_else"), "refresh_token")
+
+    @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Uses KeycloakJWTMiddleware and separate cookies")
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "access_token",
+        [
+            pytest.param("not-a-jwt", id="no-payload-segment"),
+            pytest.param(token_with_payload("not json"), id="payload-not-json"),
+            pytest.param(token_with_payload("1"), id="payload-not-an-object"),
+            pytest.param(token_with_payload('{"other": "user_id"}'), id="payload-without-sub"),
+            pytest.param(token_with_payload('{"sub": null}'), id="payload-with-null-sub"),
+        ],
+    )
+    async def test_get_user_from_token_rejects_unparsable_token(self, auth_manager, access_token):
+        """A token whose subject cannot be read matches no user and is refused."""
+        mock_get_user_from_token = AsyncMock(
+            return_value=KeycloakAuthManagerUser(
+                user_id="user_id", name="name", access_token="", refresh_token=None
+            )
+        )
+        with (
+            patch.object(BaseAuthManager, "get_user_from_token", mock_get_user_from_token),
+            pytest.raises(InvalidTokenError, match="does not belong to this Airflow session"),
+        ):
+            await auth_manager.get_user_from_token("token", access_token, "refresh_token")
+
     def test_get_url_login(self, auth_manager):
         result = auth_manager.get_url_login()
         assert result == f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/login"