This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 35d64d2f11a Add jwt auth keycloak provider (#72978)
35d64d2f11a is described below
commit 35d64d2f11a5895b1045671e38bc4536bf439959
Author: Dustin Lactin <[email protected]>
AuthorDate: Mon Sep 14 11:05:46 2026 -0600
Add jwt auth keycloak provider (#72978)
---
.../keycloak/docs/auth-manager/setup/config.rst | 3 +
providers/keycloak/docs/auth-manager/token.rst | 72 ++++++++-
providers/keycloak/provider.yaml | 10 ++
.../providers/keycloak/auth_manager/constants.py | 3 +
.../keycloak/auth_manager/datamodels/token.py | 22 ++-
.../v2-keycloak-auth-manager-generated.yaml | 26 ++++
.../keycloak/auth_manager/services/token.py | 93 ++++++++++-
.../providers/keycloak/get_provider_info.py | 7 +
.../keycloak/auth_manager/datamodels/test_token.py | 30 ++++
.../keycloak/auth_manager/routes/test_token.py | 24 +++
.../keycloak/auth_manager/services/test_token.py | 173 +++++++++++++++++++++
11 files changed, 460 insertions(+), 3 deletions(-)
diff --git a/providers/keycloak/docs/auth-manager/setup/config.rst
b/providers/keycloak/docs/auth-manager/setup/config.rst
index ef3c6ad275d..57806a6ccd2 100644
--- a/providers/keycloak/docs/auth-manager/setup/config.rst
+++ b/providers/keycloak/docs/auth-manager/setup/config.rst
@@ -72,3 +72,6 @@ Optional config options:
- ``server_url``. Keycloak server URL. This server URL is used by the Airflow
API server to communicate with Keycloak.
If the Airflow API server and Keycloak are running in Docker, set
"http://host.docker.internal:<PORT>" (default value).
You do not need to set this configuration option if you are running Keycloak
with Breeze.
+- ``jwt_federated_client_ids``. Comma-separated allow-list of Keycloak client
ids (``azp`` claim) permitted to
+ exchange a Keycloak-issued access token for an Airflow token via the
``urn:ietf:params:oauth:grant-type:jwt-bearer``
+ grant at ``/auth/token``. See :doc:`../token` for details. Unset or empty
denies every caller.
diff --git a/providers/keycloak/docs/auth-manager/token.rst
b/providers/keycloak/docs/auth-manager/token.rst
index c775814a5c1..491797e32cc 100644
--- a/providers/keycloak/docs/auth-manager/token.rst
+++ b/providers/keycloak/docs/auth-manager/token.rst
@@ -29,7 +29,7 @@ Several endpoints exist to create tokens depending on the
authentication method
If a user or service needs to interact with the Airflow public API, they can
create a token using their credentials.
-- ``/auth/token``: Create token using username and password or client
credentials with a ``[config][api_auth]jwt_expiration_time`` expiration time.
+- ``/auth/token``: Create token using username and password, client
credentials, or a Keycloak-issued JWT assertion, with a
``[config][api_auth]jwt_expiration_time`` expiration time.
- ``/auth/token/cli``: Create token for Airflow CLI using username and
password with a ``[config][api_auth]jwt_cli_expiration_time`` expiration time.
@@ -65,3 +65,73 @@ The body can also contain a ``grant_type`` field with value
``password`` but it
If other services need to interact with the Airflow public API, they can
create a token using the client credentials grant flow.
The client must live in the same realm the Auth Manager is configured to use.
Its service account must have the appropriate roles / permissions to access the
Airflow public API.
This process will return a token obtained using client credentials grant flow.
+
+.. code-block:: bash
+
+ ENDPOINT_URL="http://localhost:8080";
+ curl -X 'POST' \
+ "${ENDPOINT_URL}/auth/token" \
+ -H 'Content-Type: application/json' \
+ -d '{
+ "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+ "assertion": "<keycloak_access_token>"
+ }'
+
+If a client already authenticated to Keycloak by some other means -- for
example, a
+"Signed JWT - Federated" client bound to an external OIDC identity provider
such as a
+Kubernetes ServiceAccount issuer or AWS IAM outbound identity federation -- it
can
+exchange the resulting Keycloak access token for an Airflow token without
+re-authenticating to Keycloak through Airflow. Airflow does not obtain the
assertion on
+the client's behalf, it verifies the assertion and calls Keycloak's
``/userinfo``
+endpoint to validate it and retrieve user information.
+
+The assertion must be a valid, unexpired access token issued by this realm,
and its
+``aud`` claim must include Airflow's configured client id. The calling client
(from the
+token's ``azp`` claim) must also appear in the
+``[keycloak_auth_manager]jwt_federated_client_ids`` allow-list -- an ``aud``
match alone
+is not sufficient, since it only proves the token was meant for Airflow, not
that the
+issuing client has been vetted for machine authentication. An unset or empty
allow-list
+denies every caller.
+
+If the assertion fails JWT validation, is not allow-listed, or Keycloak
rejects it at
+``/userinfo``, the endpoint returns ``403 Invalid Keycloak assertion``. This
generic
+response does not reveal which assertion validation check failed.
+
+**Keycloak client requirements:** The federated client that obtains the
assertion (not
+the ``airflow`` client itself) must
+be configured with:
+
+- **Client authentication**: ON (confidential client), using whatever
mechanism suits
+ the caller -- a static client secret, or a "Signed JWT - Federated"
authenticator bound
+ to an external OIDC identity provider (e.g. a Kubernetes ServiceAccount
issuer or AWS
+ IAM outbound identity federation) for machine auth with no static secret.
+- **Service accounts roles**: ON, so the client can obtain its own access
token via the
+ ``client_credentials`` grant.
+- **Default client scopes** must include ``openid``. Keycloak's ``/userinfo``
endpoint
+ rejects tokens whose ``scope`` claim omits ``openid`` with a bare ``403``,
and Airflow
+ calls ``/userinfo`` to build the resulting user -- a token issued with only,
say,
+ ``profile email`` in its ``scope`` claim will fail here even though the JWT
itself is
+ perfectly valid. This is easy to miss since it is a realm-wide default that
some
+ client scope configurations exclude for service accounts.
+- Realm/client roles appropriate for whatever Airflow permissions the client
needs
+ (e.g. ``SuperAdmin``), assigned the same way as any other service-account
client.
+
+On the Airflow side, the client's id (from its token's ``azp`` claim) must be
added to
+``[keycloak_auth_manager]jwt_federated_client_ids``.
+
+The response contains an Airflow-minted JWT, not the Keycloak assertion
itself. Exchange
+the assertion for it first, then use that JWT -- not ``$KEYCLOAK_TOKEN`` -- as
the Bearer
+token for subsequent public API calls:
+
+.. code-block:: bash
+
+ ENDPOINT_URL="http://airflow-api-server:8080";
+
+ AIRFLOW_TOKEN=$(curl -s -X 'POST' \
+ "${ENDPOINT_URL}/auth/token" \
+ -H 'Content-Type: application/json' \
+ -d "{\"grant_type\": \"urn:ietf:params:oauth:grant-type:jwt-bearer\",
\"assertion\": \"${KEYCLOAK_TOKEN}\"}" \
+ | python3 -c 'import json, sys;
print(json.load(sys.stdin)["access_token"])')
+
+ curl -s "${ENDPOINT_URL}/api/v2/dags" \
+ -H "Authorization: Bearer ${AIRFLOW_TOKEN}"
diff --git a/providers/keycloak/provider.yaml b/providers/keycloak/provider.yaml
index 721f0557a77..d6105593bdc 100644
--- a/providers/keycloak/provider.yaml
+++ b/providers/keycloak/provider.yaml
@@ -93,6 +93,16 @@ config:
version_added: 0.0.1
example: ~
default: ~
+ jwt_federated_client_ids:
+ description: |
+ Comma-separated allow-list of Keycloak client ids (``azp`` claim)
permitted to
+ exchange a Keycloak-issued access token for an Airflow token via the
+ ``urn:ietf:params:oauth:grant-type:jwt-bearer`` grant at
``/auth/token``. Unset
+ or empty denies every caller.
+ type: string
+ version_added: 0.10.0
+ example: "team-platform-admin-sa,other-team-sa"
+ default: ~
realm:
description: |
Realm configured in Keycloak associated to Airflow.
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 169042d1f45..262166de5b6 100644
---
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/constants.py
+++
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/constants.py
@@ -27,6 +27,9 @@ CONF_REALM_KEY = "realm"
CONF_SERVER_URL_KEY = "server_url"
CONF_REQUESTS_POOL_SIZE_KEY = "requests_pool_size"
CONF_REQUESTS_RETRIES_KEY = "requests_retries"
+# Comma-separated allow-list of federated Keycloak client ids (azp claim)
permitted
+# to use the jwt-bearer grant. Empty/unset denies all callers.
+CONF_JWT_FEDERATED_CLIENT_IDS_KEY = "jwt_federated_client_ids"
# Extra Cookie names
COOKIE_NAME_ACCESS_TOKEN = "_access_token"
diff --git
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/datamodels/token.py
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/datamodels/token.py
index 6d60e7e4020..f61fdd3615c 100644
---
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/datamodels/token.py
+++
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/datamodels/token.py
@@ -24,6 +24,7 @@ from pydantic import Field, RootModel, model_validator
from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
from airflow.providers.keycloak.auth_manager.services.token import (
create_client_credentials_token,
+ create_jwt_federated_token,
create_token_for,
)
@@ -62,8 +63,27 @@ class TokenClientCredentialsBody(StrictBaseModel):
)
+class TokenJwtFederatedBody(StrictBaseModel):
+ """
+ JWT-bearer grant token serializer for post bodies.
+
+ Accepts a Keycloak access token obtained via any Keycloak-native
authentication
+ method (e.g. a client federated to an external OIDC identity provider)
instead of
+ Airflow re-authenticating to Keycloak itself.
+ """
+
+ grant_type: Literal["urn:ietf:params:oauth:grant-type:jwt-bearer"]
+ assertion: str = Field()
+
+ def create_token(self, expiration_time_in_seconds: int) -> str:
+ """Create token by validating a pre-obtained Keycloak assertion."""
+ return create_jwt_federated_token(
+ self.assertion,
expiration_time_in_seconds=expiration_time_in_seconds
+ )
+
+
TokenUnion = Annotated[
- TokenPasswordBody | TokenClientCredentialsBody,
+ TokenPasswordBody | TokenClientCredentialsBody | TokenJwtFederatedBody,
Field(discriminator="grant_type"),
]
diff --git
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/openapi/v2-keycloak-auth-manager-generated.yaml
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/openapi/v2-keycloak-auth-manager-generated.yaml
index 0c3d4fa8c49..2158cc3a65b 100644
---
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/openapi/v2-keycloak-auth-manager-generated.yaml
+++
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/openapi/v2-keycloak-auth-manager-generated.yaml
@@ -164,6 +164,7 @@ components:
oneOf:
- $ref: '#/components/schemas/TokenPasswordBody'
- $ref: '#/components/schemas/TokenClientCredentialsBody'
+ - $ref: '#/components/schemas/TokenJwtFederatedBody'
title: TokenBody
description: Token request body.
discriminator:
@@ -171,6 +172,7 @@ components:
mapping:
client_credentials: '#/components/schemas/TokenClientCredentialsBody'
password: '#/components/schemas/TokenPasswordBody'
+ urn:ietf:params:oauth:grant-type:jwt-bearer:
'#/components/schemas/TokenJwtFederatedBody'
TokenClientCredentialsBody:
properties:
grant_type:
@@ -191,6 +193,30 @@ components:
- client_secret
title: TokenClientCredentialsBody
description: Client credentials grant token serializer for post bodies.
+ TokenJwtFederatedBody:
+ properties:
+ grant_type:
+ type: string
+ const: urn:ietf:params:oauth:grant-type:jwt-bearer
+ title: Grant Type
+ assertion:
+ type: string
+ title: Assertion
+ additionalProperties: false
+ type: object
+ required:
+ - grant_type
+ - assertion
+ title: TokenJwtFederatedBody
+ description: 'JWT-bearer grant token serializer for post bodies.
+
+
+ Accepts a Keycloak access token obtained via any Keycloak-native
authentication
+
+ method (e.g. a client federated to an external OIDC identity provider)
instead
+ of
+
+ Airflow re-authenticating to Keycloak itself.'
TokenPasswordBody:
properties:
grant_type:
diff --git
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/token.py
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/token.py
index 94da35310f7..6a0566bfce0 100644
---
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/token.py
+++
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/token.py
@@ -18,20 +18,32 @@
from __future__ import annotations
import json
+from functools import cache
+import jwt
from fastapi import HTTPException, status
-from keycloak import KeycloakAuthenticationError
+from jwt import PyJWKClient
+from keycloak import KeycloakAuthenticationError, KeycloakError
from airflow.api_fastapi.app import get_auth_manager
from airflow.providers.common.compat.sdk import conf
from airflow.providers.keycloak.auth_manager.constants import (
CONF_CLIENT_ID_KEY,
+ CONF_JWT_FEDERATED_CLIENT_IDS_KEY,
+ CONF_REALM_KEY,
CONF_SECTION_NAME,
+ CONF_SERVER_URL_KEY,
)
from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import
KeycloakAuthManager
from airflow.providers.keycloak.auth_manager.user import
KeycloakAuthManagerUser
+@cache
+def _get_jwks_client(issuer: str) -> PyJWKClient:
+ """Return the cached JWK client for a Keycloak realm issuer."""
+ return PyJWKClient(f"{issuer}/protocol/openid-connect/certs")
+
+
def create_token_for(
username: str,
password: str,
@@ -61,6 +73,85 @@ def create_token_for(
return get_auth_manager().generate_api_jwt(user,
expiration_time_in_seconds=expiration_time_in_seconds)
+def create_jwt_federated_token(
+ assertion: str,
+ expiration_time_in_seconds: int = conf.getint("api_auth",
"jwt_expiration_time"),
+) -> str:
+ """
+ Create a token from a Keycloak access token obtained outside of Airflow.
+
+ This authentication flow accepts an access token issued by Keycloak
through any
+ Keycloak-native mechanism (e.g. a "Signed JWT - Federated" client bound to
an
+ external OIDC identity provider such as a Kubernetes ServiceAccount
issuer, or AWS
+ IAM outbound identity federation). Airflow does not obtain the token on
the caller's
+ behalf, the caller must obtain it directly from Keycloak's token endpoint.
Airflow
+ verifies the token and calls Keycloak's UserInfo endpoint to validate it
and retrieve
+ user information.
+
+ The token's signature, issuer, and audience are verified against this
realm's JWKS.
+ The ``aud`` claim (a string or a list) must include this Airflow client's
id, which
+ requires an Audience mapper on the federated client's scope in Keycloak.
The calling
+ client (``azp``) must also appear in the ``jwt_federated_client_ids``
allow-list
+ below -- an ``aud`` match alone only proves the token was meant for
Airflow, not
+ that the issuing client has been vetted for machine auth. Any Keycloak
validation
+ error is returned as a generic ``403 Invalid Keycloak assertion`` response.
+ """
+ realm = conf.get(CONF_SECTION_NAME, CONF_REALM_KEY)
+ server_url = conf.get(CONF_SECTION_NAME, CONF_SERVER_URL_KEY)
+ client_id = conf.get(CONF_SECTION_NAME, CONF_CLIENT_ID_KEY)
+ issuer = f"{server_url.rstrip('/')}/realms/{realm}"
+
+ try:
+ jwks_client = _get_jwks_client(issuer)
+ signing_key = jwks_client.get_signing_key_from_jwt(assertion)
+ claims = jwt.decode(
+ assertion,
+ signing_key.key,
+ algorithms=["RS256"],
+ audience=client_id,
+ issuer=issuer,
+ )
+ except jwt.PyJWTError:
+ raise HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail="Invalid Keycloak assertion",
+ )
+
+ allowed_client_ids = {
+ allowed.strip()
+ for allowed in conf.get(CONF_SECTION_NAME,
CONF_JWT_FEDERATED_CLIENT_IDS_KEY, fallback="").split(",")
+ if allowed.strip()
+ }
+ federated_client_id = claims.get("azp") or claims.get("client_id")
+ if not federated_client_id or federated_client_id not in
allowed_client_ids:
+ raise HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail="Invalid Keycloak assertion",
+ )
+
+ # Confirm the assertion is still live (not revoked) and fetch the same
shape of
+ # user info create_client_credentials_token uses, rather than trusting the
JWT's
+ # own claims alone.
+ client = KeycloakAuthManager.get_keycloak_client()
+ try:
+ userinfo_raw: dict | bytes = client.userinfo(assertion)
+ except KeycloakError:
+ raise HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail="Invalid Keycloak assertion",
+ )
+ userinfo: dict = json.loads(userinfo_raw) if isinstance(userinfo_raw,
bytes) else userinfo_raw
+
+ user = KeycloakAuthManagerUser(
+ user_id=userinfo["sub"],
+ name=userinfo.get("preferred_username", userinfo.get("clientId",
"service-account")),
+ access_token=assertion,
+ refresh_token=None,
+ )
+
+ return get_auth_manager().generate_api_jwt(user,
expiration_time_in_seconds=expiration_time_in_seconds)
+
+
def create_client_credentials_token(
client_id: str,
client_secret: str,
diff --git
a/providers/keycloak/src/airflow/providers/keycloak/get_provider_info.py
b/providers/keycloak/src/airflow/providers/keycloak/get_provider_info.py
index d186cf2b868..f4b0ae3a890 100644
--- a/providers/keycloak/src/airflow/providers/keycloak/get_provider_info.py
+++ b/providers/keycloak/src/airflow/providers/keycloak/get_provider_info.py
@@ -71,6 +71,13 @@ def get_provider_info():
"example": None,
"default": None,
},
+ "jwt_federated_client_ids": {
+ "description": "Comma-separated allow-list of Keycloak
client ids (``azp`` claim) permitted to\nexchange a Keycloak-issued access
token for an Airflow token via
the\n``urn:ietf:params:oauth:grant-type:jwt-bearer`` grant at ``/auth/token``.
Unset\nor empty denies every caller.\n",
+ "type": "string",
+ "version_added": "0.10.0",
+ "example": "team-platform-admin-sa,other-team-sa",
+ "default": None,
+ },
"realm": {
"description": "Realm configured in Keycloak
associated to Airflow.\nThis realm define all users, roles and groups used in
Airflow.\n",
"type": "string",
diff --git
a/providers/keycloak/tests/unit/keycloak/auth_manager/datamodels/test_token.py
b/providers/keycloak/tests/unit/keycloak/auth_manager/datamodels/test_token.py
index 378a5d7da55..50c82c4a5bd 100644
---
a/providers/keycloak/tests/unit/keycloak/auth_manager/datamodels/test_token.py
+++
b/providers/keycloak/tests/unit/keycloak/auth_manager/datamodels/test_token.py
@@ -24,6 +24,7 @@ from pydantic import ValidationError
from airflow.providers.keycloak.auth_manager.datamodels.token import (
TokenBody,
TokenClientCredentialsBody,
+ TokenJwtFederatedBody,
TokenPasswordBody,
TokenResponse,
)
@@ -68,6 +69,17 @@ class TestTokenBody:
"client_secret": "client_secret",
},
),
+ (
+ {
+ "grant_type":
"urn:ietf:params:oauth:grant-type:jwt-bearer",
+ "assertion": "assertion",
+ },
+ TokenJwtFederatedBody,
+ {
+ "grant_type":
"urn:ietf:params:oauth:grant-type:jwt-bearer",
+ "assertion": "assertion",
+ },
+ ),
],
)
def test_model_validate_and_dump(self, payload, expected_type,
expected_dump):
@@ -84,6 +96,8 @@ class TestTokenBody:
{"grant_type": "unsupported", "username": "username", "password":
"password"},
{"username": "username", "password": "password", "extra": "value"},
{"grant_type": "client_credentials", "client_secret":
"client_secret"},
+ {"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer"},
+ {"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
"assertion": "a", "extra": "x"},
],
)
def test_rejects_invalid_payload(self, payload):
@@ -118,3 +132,19 @@ class TestTokenClientCredentialsBody:
mock_create_client_credentials_token.assert_called_once_with(
"client_id", "client_secret", expiration_time_in_seconds=60
)
+
+
+class TestTokenJwtFederatedBody:
+ @mock.patch(
+
"airflow.providers.keycloak.auth_manager.datamodels.token.create_jwt_federated_token",
+ autospec=True,
+ )
+ def test_create_token(self, mock_create_jwt_federated_token):
+ mock_create_jwt_federated_token.return_value = "token"
+ body = TokenJwtFederatedBody(
+ grant_type="urn:ietf:params:oauth:grant-type:jwt-bearer",
+ assertion="assertion",
+ )
+
+ assert body.create_token(expiration_time_in_seconds=60) == "token"
+ mock_create_jwt_federated_token.assert_called_once_with("assertion",
expiration_time_in_seconds=60)
diff --git
a/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_token.py
b/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_token.py
index 064288152ff..d7903ca4879 100644
--- a/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_token.py
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_token.py
@@ -92,6 +92,28 @@ class TestTokenRouter:
"client_id", "client_secret", expiration_time_in_seconds=10
)
+ @conf_vars(
+ {
+ ("api_auth", "jwt_expiration_time"): "10",
+ }
+ )
+
@patch("airflow.providers.keycloak.auth_manager.datamodels.token.create_jwt_federated_token")
+ def test_create_token_jwt_bearer_grant(self,
mock_create_jwt_federated_token, client):
+ mock_create_jwt_federated_token.return_value = self.token
+ response = client.post(
+ AUTH_MANAGER_FASTAPI_APP_PREFIX + "/token",
+ json={
+ "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+ "assertion": "a-keycloak-access-token",
+ },
+ )
+
+ assert response.status_code == 201
+ assert response.json() == {"access_token": self.token}
+ mock_create_jwt_federated_token.assert_called_once_with(
+ "a-keycloak-access-token", expiration_time_in_seconds=10
+ )
+
@pytest.mark.parametrize(
"body",
[
@@ -102,6 +124,8 @@ class TestTokenRouter:
{"grant_type": "client_credentials", "username": "username",
"password": "password"},
{"grant_type": "client_credentials", "client_id": "client_id",
"password": "password"},
{"grant_type": "client_credentials", "username": "username",
"client_secret": "client_secret"},
+ {"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer"},
+ {"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
"username": "username"},
],
)
@conf_vars(
diff --git
a/providers/keycloak/tests/unit/keycloak/auth_manager/services/test_token.py
b/providers/keycloak/tests/unit/keycloak/auth_manager/services/test_token.py
index b02bcc1c95a..1a5b8aa5a8a 100644
--- a/providers/keycloak/tests/unit/keycloak/auth_manager/services/test_token.py
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/services/test_token.py
@@ -19,12 +19,15 @@ from __future__ import annotations
from unittest.mock import Mock, patch
import fastapi
+import jwt
import pytest
from keycloak import KeycloakAuthenticationError
+from keycloak.exceptions import KeycloakGetError
from airflow.providers.common.compat.sdk import conf
from airflow.providers.keycloak.auth_manager.services.token import (
create_client_credentials_token,
+ create_jwt_federated_token,
create_token_for,
)
@@ -179,3 +182,173 @@ class TestTokenService:
assert wrong_secret.value.status_code == wrong_client.value.status_code
assert wrong_secret.value.detail == wrong_client.value.detail
+
+
+class TestCreateJwtFederatedToken:
+ token = "token"
+ test_assertion = "assertion"
+ test_access_token = "assertion"
+
+ @conf_vars(
+ {
+ ("api_auth", "jwt_expiration_time"): "10",
+ ("keycloak_auth_manager", "client_id"): "airflow",
+ ("keycloak_auth_manager", "realm"): "airflow",
+ ("keycloak_auth_manager", "server_url"):
"https://keycloak.example.com";,
+ ("keycloak_auth_manager", "jwt_federated_client_ids"):
"team-platform-admin-sa, other-sa",
+ }
+ )
+
@patch("airflow.providers.keycloak.auth_manager.services.token.get_auth_manager")
+
@patch("airflow.providers.keycloak.auth_manager.services.token.KeycloakAuthManager.get_keycloak_client")
+ @patch("airflow.providers.keycloak.auth_manager.services.token.jwt.decode")
+
@patch("airflow.providers.keycloak.auth_manager.services.token._get_jwks_client")
+ def test_create_jwt_federated_token(
+ self, mock_jwks_client_cls, mock_jwt_decode, mock_get_keycloak_client,
mock_get_auth_manager
+ ):
+
mock_jwks_client_cls.return_value.get_signing_key_from_jwt.return_value =
Mock(key="fake-key")
+ mock_jwt_decode.return_value = {
+ "sub": "service-account-sub",
+ "azp": "team-platform-admin-sa",
+ "aud": ["airflow"],
+ }
+ mock_keycloak_client = Mock()
+ mock_keycloak_client.userinfo.return_value = {
+ "sub": "service-account-sub",
+ "preferred_username": "service-account-team-platform-admin-sa",
+ }
+ 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_api_jwt.return_value = self.token
+
+ result = create_jwt_federated_token(assertion=self.test_assertion)
+
+ assert result == self.token
+ mock_jwt_decode.assert_called_once_with(
+ self.test_assertion,
+ "fake-key",
+ algorithms=["RS256"],
+ audience="airflow",
+ issuer="https://keycloak.example.com/realms/airflow";,
+ )
+
mock_keycloak_client.userinfo.assert_called_once_with(self.test_assertion)
+ mock_auth_manager.generate_api_jwt.assert_called_once()
+
+ @conf_vars(
+ {
+ ("api_auth", "jwt_expiration_time"): "10",
+ ("keycloak_auth_manager", "client_id"): "airflow",
+ ("keycloak_auth_manager", "realm"): "airflow",
+ ("keycloak_auth_manager", "server_url"):
"https://keycloak.example.com";,
+ ("keycloak_auth_manager", "jwt_federated_client_ids"):
"team-platform-admin-sa",
+ }
+ )
+
@patch("airflow.providers.keycloak.auth_manager.services.token.KeycloakAuthManager.get_keycloak_client")
+ @patch("airflow.providers.keycloak.auth_manager.services.token.jwt.decode")
+
@patch("airflow.providers.keycloak.auth_manager.services.token._get_jwks_client")
+ def test_create_jwt_federated_token_invalid_signature(
+ self, mock_jwks_client_cls, mock_jwt_decode, mock_get_keycloak_client
+ ):
+
mock_jwks_client_cls.return_value.get_signing_key_from_jwt.return_value =
Mock(key="fake-key")
+ mock_jwt_decode.side_effect = jwt.InvalidSignatureError()
+
+ with pytest.raises(fastapi.exceptions.HTTPException) as exc_info:
+ create_jwt_federated_token(assertion=self.test_assertion)
+
+ assert exc_info.value.status_code == 403
+
+ @conf_vars(
+ {
+ ("api_auth", "jwt_expiration_time"): "10",
+ ("keycloak_auth_manager", "client_id"): "airflow",
+ ("keycloak_auth_manager", "realm"): "airflow",
+ ("keycloak_auth_manager", "server_url"):
"https://keycloak.example.com";,
+ ("keycloak_auth_manager", "jwt_federated_client_ids"): "other-sa",
+ }
+ )
+
@patch("airflow.providers.keycloak.auth_manager.services.token.KeycloakAuthManager.get_keycloak_client")
+ @patch("airflow.providers.keycloak.auth_manager.services.token.jwt.decode")
+
@patch("airflow.providers.keycloak.auth_manager.services.token._get_jwks_client")
+ def test_create_jwt_federated_token_client_not_allowlisted(
+ self, mock_jwks_client_cls, mock_jwt_decode, mock_get_keycloak_client
+ ):
+ """A validly-signed token for an un-allow-listed client must still be
rejected."""
+
mock_jwks_client_cls.return_value.get_signing_key_from_jwt.return_value =
Mock(key="fake-key")
+ mock_jwt_decode.return_value = {
+ "sub": "service-account-sub",
+ "azp": "team-platform-admin-sa",
+ "aud": ["airflow"],
+ }
+
+ with pytest.raises(fastapi.exceptions.HTTPException) as exc_info:
+ create_jwt_federated_token(assertion=self.test_assertion)
+
+ assert exc_info.value.status_code == 403
+ # No userinfo call is made once the client is rejected.
+ mock_get_keycloak_client.assert_not_called()
+
+ @conf_vars(
+ {
+ ("api_auth", "jwt_expiration_time"): "10",
+ ("keycloak_auth_manager", "client_id"): "airflow",
+ ("keycloak_auth_manager", "realm"): "airflow",
+ ("keycloak_auth_manager", "server_url"):
"https://keycloak.example.com";,
+ }
+ )
+
@patch("airflow.providers.keycloak.auth_manager.services.token.KeycloakAuthManager.get_keycloak_client")
+ @patch("airflow.providers.keycloak.auth_manager.services.token.jwt.decode")
+
@patch("airflow.providers.keycloak.auth_manager.services.token._get_jwks_client")
+ def test_create_jwt_federated_token_allowlist_unset_fails_closed(
+ self, mock_jwks_client_cls, mock_jwt_decode, mock_get_keycloak_client
+ ):
+ """An unconfigured jwt_federated_client_ids denies every caller, not
just none."""
+
mock_jwks_client_cls.return_value.get_signing_key_from_jwt.return_value =
Mock(key="fake-key")
+ mock_jwt_decode.return_value = {
+ "sub": "service-account-sub",
+ "azp": "team-platform-admin-sa",
+ "aud": ["airflow"],
+ }
+
+ with pytest.raises(fastapi.exceptions.HTTPException) as exc_info:
+ create_jwt_federated_token(assertion=self.test_assertion)
+
+ assert exc_info.value.status_code == 403
+ mock_get_keycloak_client.assert_not_called()
+
+ @pytest.mark.parametrize(
+ "userinfo_error",
+ [
+ KeycloakAuthenticationError(),
+ KeycloakGetError(response_code=403, response_body=b""),
+ ],
+ )
+ @conf_vars(
+ {
+ ("api_auth", "jwt_expiration_time"): "10",
+ ("keycloak_auth_manager", "client_id"): "airflow",
+ ("keycloak_auth_manager", "realm"): "airflow",
+ ("keycloak_auth_manager", "server_url"):
"https://keycloak.example.com";,
+ ("keycloak_auth_manager", "jwt_federated_client_ids"):
"team-platform-admin-sa",
+ }
+ )
+
@patch("airflow.providers.keycloak.auth_manager.services.token.KeycloakAuthManager.get_keycloak_client")
+ @patch("airflow.providers.keycloak.auth_manager.services.token.jwt.decode")
+
@patch("airflow.providers.keycloak.auth_manager.services.token._get_jwks_client")
+ def test_create_jwt_federated_token_userinfo_rejected(
+ self, mock_jwks_client_cls, mock_jwt_decode, mock_get_keycloak_client,
userinfo_error
+ ):
+ """A Keycloak userinfo rejection returns a generic invalid-assertion
response."""
+
mock_jwks_client_cls.return_value.get_signing_key_from_jwt.return_value =
Mock(key="fake-key")
+ mock_jwt_decode.return_value = {
+ "sub": "service-account-sub",
+ "azp": "team-platform-admin-sa",
+ "aud": ["airflow"],
+ }
+ mock_keycloak_client = Mock()
+ mock_keycloak_client.userinfo.side_effect = userinfo_error
+ mock_get_keycloak_client.return_value = mock_keycloak_client
+
+ with pytest.raises(fastapi.exceptions.HTTPException) as exc_info:
+ create_jwt_federated_token(assertion=self.test_assertion)
+
+ assert exc_info.value.status_code == 403
