kaxil commented on code in PR #74700: URL: https://github.com/apache/airflow/pull/74700#discussion_r4359445825
########## airflow-core/src/airflow/dag_processing/credentials.py: ########## @@ -0,0 +1,121 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Session credential sources for Dag processor API clients.""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path +from time import monotonic +from typing import Protocol + +import httpx2 + +from airflow.api_fastapi.execution_api.datamodels.job import DagProcessorSessionTokenResponse +from airflow.api_fastapi.execution_api.versions import bundle +from airflow.sdk.api.client import Client # noqa: SDK001 - reuse Execution API TLS and retry settings + + +class SessionTokenProvider(Protocol): + """Provide credentials without prescribing where they are obtained.""" + + def get_token( + self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx2.Timeout | None = None + ) -> str: ... + + def close(self) -> None: ... + + +class FileSessionTokenProvider: + """Read credentials written by the existing trusted provisioner.""" + + def __init__(self, path: str | Path): + self.path = Path(path) + + def get_token( + self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx2.Timeout | None = None + ) -> str: + token = self.path.read_text().strip() + if not token: + raise ValueError(f"Dag processor token file is empty: {self.path}") + return token + + def close(self) -> None: + pass + + +class _SessionTokenClient(Client): + def _update_auth(self, response: httpx2.Response) -> None: + # Runtime refresh headers must not replace the external workload identity. + pass + + +class WorkloadIdentitySessionTokenProvider: + """ + Exchange deployment identity for a session JWT, retaining the JWT only in memory. + + The identity supplier must return a current, renewable deployment credential. For example, it may + read a projected service-account token that the platform rotates. The owning Dag processor client + closes this provider; do not share it between clients. + """ + + def __init__( + self, + *, + base_url: str, + identity_token: Callable[[], str], + **kwargs, + ): + if httpx2.URL(base_url).scheme != "https": + raise ValueError("Workload identity exchange requires an HTTPS Execution API URL") + self._identity_token = identity_token + self._token: str | None = None + self._renew_at = 0.0 + self._client = _SessionTokenClient(base_url=base_url, token="", follow_redirects=False, **kwargs) + + def get_token( + self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx2.Timeout | None = None + ) -> str: + if not force_refresh and self._token is not None and monotonic() < self._renew_at: + return self._token + identity = self._identity_token().strip() Review Comment: The optional early renewal in `_ensure_job_token` only catches `httpx2.HTTPError`, `OSError`, `ValueError` and `DagProcessorRegistrationRetired`. A supplier that gets the identity from a metadata server or cloud SDK can raise something else (google-auth's `RefreshError` subclasses plain `Exception`), and then every request fails while the Job token is still valid, because `_retry_renewal_at` is never set. Wrapping the supplier call here in an exception type the client already catches, or documenting what the callable may raise, would keep the renewal best effort. ########## airflow-core/src/airflow/dag_processing/credentials.py: ########## @@ -0,0 +1,121 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Session credential sources for Dag processor API clients.""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path +from time import monotonic +from typing import Protocol + +import httpx2 + +from airflow.api_fastapi.execution_api.datamodels.job import DagProcessorSessionTokenResponse +from airflow.api_fastapi.execution_api.versions import bundle +from airflow.sdk.api.client import Client # noqa: SDK001 - reuse Execution API TLS and retry settings + + +class SessionTokenProvider(Protocol): + """Provide credentials without prescribing where they are obtained.""" + + def get_token( + self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx2.Timeout | None = None + ) -> str: ... + + def close(self) -> None: ... + + +class FileSessionTokenProvider: + """Read credentials written by the existing trusted provisioner.""" + + def __init__(self, path: str | Path): + self.path = Path(path) + + def get_token( + self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx2.Timeout | None = None + ) -> str: + token = self.path.read_text().strip() + if not token: + raise ValueError(f"Dag processor token file is empty: {self.path}") + return token + + def close(self) -> None: + pass + + +class _SessionTokenClient(Client): + def _update_auth(self, response: httpx2.Response) -> None: + # Runtime refresh headers must not replace the external workload identity. + pass + + +class WorkloadIdentitySessionTokenProvider: + """ + Exchange deployment identity for a session JWT, retaining the JWT only in memory. + + The identity supplier must return a current, renewable deployment credential. For example, it may + read a projected service-account token that the platform rotates. The owning Dag processor client + closes this provider; do not share it between clients. + """ + + def __init__( + self, + *, + base_url: str, + identity_token: Callable[[], str], + **kwargs, + ): + if httpx2.URL(base_url).scheme != "https": + raise ValueError("Workload identity exchange requires an HTTPS Execution API URL") + self._identity_token = identity_token + self._token: str | None = None + self._renew_at = 0.0 + self._client = _SessionTokenClient(base_url=base_url, token="", follow_redirects=False, **kwargs) + + def get_token( + self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx2.Timeout | None = None + ) -> str: + if not force_refresh and self._token is not None and monotonic() < self._renew_at: + return self._token + identity = self._identity_token().strip() + if not identity: + raise ValueError("Workload identity credential is empty") + started_at = monotonic() + response = self._client.request( + "POST", + "dag-processor/session-token", + headers={ + "Authorization": f"Bearer {identity}", + "Airflow-API-Version": bundle.version_values[0], + }, + retry=retry, Review Comment: This retries on the SDK policy (`wait_random_exponential(min=1, max=90)`, 5 attempts), which ignores `Retry-After`. After a failed cold-cache JWKS fetch the server won't fetch again for 10s, and the four default waits (1 + U(1,2) + U(1,4) + U(1,8)) add up to under 10s about 57% of the time when I simulated them. In that case all five attempts get the cooldown 503 and `register_job()` fails even if the IdP recovered right after the first failure, while the docs say the 503 lets registration retry. `test_registration_recovers_from_jwks_outage` moves the clock 11s before every retry, so it never sees this. Could the exchange honour `Retry-After` on 503 (keeping `retry=False` for the optional renewal), with a test on the default wait policy? ########## airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py: ########## @@ -0,0 +1,145 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Authenticate deployment identities for Dag processor session provisioning.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Annotated, Any, Literal +from uuid import NAMESPACE_URL, UUID, uuid5 + +import httpx +import jwt +from pydantic import BaseModel, ConfigDict, Field, HttpUrl, field_validator + +from airflow.api_fastapi.auth.tokens import JWKS, JWTValidator +from airflow.api_fastapi.execution_api.datamodels.token import DagBundleGrant +from airflow.configuration import conf +from airflow.dag_processing.bundles.manager import get_configured_bundle_names + +NonemptyString = Annotated[str, Field(min_length=1)] + + +class WorkloadIdentityConfig(BaseModel): Review Comment: This takes a single issuer, and grants are keyed on bare `sub`. Kubernetes issuers are per cluster, so Dag processors in a second cluster (per-team clusters in a multi-team setup, say) can't be trusted by the same API server. Should the top level be a list of issuer entries, each with its own `grants`, before the option ships and the shape becomes something we have to keep? If single issuer is intended, saying so in the option description would help. ########## airflow-core/src/airflow/api_fastapi/auth/tokens.py: ########## @@ -128,19 +145,33 @@ def from_private_key(cls, *keys: AllowedPrivateKeys | tuple[AllowedPrivateKeys, obj._jwks = jwt.PyJWKSet(keyset) return obj - async def fetch_jwks(self) -> None: - if not self._should_fetch_jwks(): + async def fetch_jwks(self, *, force: bool = False) -> None: + # A usable cache skips the lock, so a slow refresh for an unknown key does not delay known keys. + if not force and self._jwks and not self._should_fetch_jwks(): return - if self.url.startswith("http"): - data = await self._fetch_remote_jwks() - else: - data = self._fetch_local_jwks() - - if not data: - return - - self._jwks = jwt.PyJWKSet.from_dict(data) - log.debug("Fetched JWKS", url=self.url, keys=len(self._jwks.keys)) + lock = self._fetch_lock if self.refresh_on_unknown_kid else nullcontext() + async with lock: + if not self._should_fetch_jwks(force=force): + return + self._fetch_failed = True + if self.url.startswith("http"): + data = await self._fetch_remote_jwks() + else: + data = self._fetch_local_jwks() + + if not data: + return + + try: + keyset = jwt.PyJWKSet.from_dict(data) + except (jwt.PyJWKError, jwt.PyJWKSetError, KeyError, TypeError, ValueError): Review Comment: `PyJWKSet.from_dict` starts with `obj.get("keys", [])`, so a JWKS body that is a JSON array or scalar raises `AttributeError` (checked on the pyjwt pinned in uv.lock), and that isn't in this tuple. It comes back as a 500 rather than the 503 with `Retry-After`, and once the hourly refresh is due it also fails tokens whose kid is still cached, which the class docstring says stay usable. Treating a non-dict `data` as a failed fetch before parsing would cover it. ########## airflow-core/docs/security/jwt_token_authentication.rst: ########## @@ -519,6 +519,125 @@ API without holding its signing key. The standard ``airflow dag-processor`` comm not yet use this client. Provisioning a token file alone does not switch that command to HTTP or remove its need for database access. +In-memory sessions from workload identity +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Review Comment: The docs build fails on this page. `~` under the `-` section is read as a skip from level 2 to 4 (the page uses `^` for the third level), and this underline is one character shorter than the title. The same `~` problem is on L639 and L655, and "keyset" on L636 fails the spell check. ########## airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py: ########## @@ -0,0 +1,556 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import json +from datetime import UTC, datetime +from unittest import mock + +import httpx +import httpx2 +import jwt +import pytest +from cryptography.hazmat.primitives.asymmetric import ec, rsa +from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey +from pydantic import ValidationError +from sqlalchemy import select +from tenacity import wait_none +from uuid6 import uuid7 + +from airflow.api_fastapi.auth.tokens import key_to_jwk_dict +from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody, JobState, TerminalJobState +from airflow.api_fastapi.execution_api.security import require_auth +from airflow.api_fastapi.execution_api.workload_identity import ( + DagProcessorIdentityVerifier, + WorkloadIdentityConfig, + create_dag_processor_identity_verifier, +) +from airflow.dag_processing.api_client import ( + DagParseContext, + DagProcessorAPIClient, + DagProcessorJobAlreadyRunning, +) +from airflow.dag_processing.credentials import FileSessionTokenProvider, WorkloadIdentitySessionTokenProvider +from airflow.jobs.job import Job +from airflow.models.variable import Variable +from airflow.sdk.api.client import Client + +from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.db import clear_db_jobs, clear_db_variables + +pytestmark = pytest.mark.db_test + +URL = "/execution/dag-processor/session-token" +NOW = datetime(2026, 10, 5, 12, 0, tzinfo=UTC).timestamp() +SECRET = "session-exchange-runtime-signing-secret-" * 3 +TRUST = { + "issuer": "https://identity.example/";, + "jwks_url": "https://identity.example/keys";, + "audience": "airflow-dag-processor-bootstrap", + "algorithm": "EdDSA", + "instance_id_claim": ["kubernetes.io", "pod", "uid"], + "grants": {"system:serviceaccount:airflow:processor": ["dags-folder"]}, +} + + [email protected](autouse=True) +def configure(time_machine, identity_algorithm): + time_machine.move_to(NOW, tick=False) + with conf_vars( + { + ("execution_api", "dag_processor_workload_identity"): json.dumps( + {**TRUST, "algorithm": identity_algorithm} + ), + ("execution_api", "jwt_expiration_time"): "300", + ("api_auth", "jwt_secret"): SECRET, + ("api_auth", "jwt_leeway"): "10", + } + ): + yield + + [email protected] +def identity_algorithm(request): + return getattr(request, "param", "EdDSA") + + [email protected] +def signing_key(identity_algorithm): + if identity_algorithm == "RS256": + return rsa.generate_private_key(public_exponent=65537, key_size=2048) + if identity_algorithm == "ES256": + return ec.generate_private_key(ec.SECP256R1()) + return Ed25519PrivateKey.generate() + + [email protected] +def identity_token(signing_key, identity_algorithm): + def issue(*, remove=(), key=None, headers=None, **overrides): + claims = { + "iss": TRUST["issuer"], + "aud": TRUST["audience"], + "sub": "system:serviceaccount:airflow:processor", + "iat": NOW, + "exp": NOW + 600, + "kubernetes.io": {"pod": {"uid": "pod-1"}}, + **overrides, + } + for name in remove: + claims.pop(name) + return jwt.encode( + claims, key or signing_key, algorithm=identity_algorithm, headers=headers or {"kid": "key-1"} + ) + + return issue + + [email protected] +def verifier(client, exec_app, signing_key, identity_algorithm, monkeypatch): + if identity_algorithm == "ES256": + public_key = jwt.algorithms.ECAlgorithm.to_jwk(signing_key.public_key(), as_dict=True) | { + "kid": "key-1" + } + else: + public_key = key_to_jwk_dict(signing_key, kid="key-1") + jwks_client = httpx.AsyncClient( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"keys": [public_key]})) + ) + verifier = exec_app.state.dag_processor_identity_verifier + monkeypatch.setattr(verifier.validator.jwks, "client", jwks_client) + yield verifier + client.portal.call(jwks_client.aclose) + + [email protected] +def make_processor(client, exec_app, identity_token, verifier): + exec_app.dependency_overrides.pop(require_auth) + clear_db_jobs() + processors = [] + + def forward(request): + response = client.request( + request.method, request.url.path, headers=request.headers, content=request.content + ) + return httpx2.Response(response.status_code, headers=response.headers, content=response.content) + + def create(): + transport = httpx2.MockTransport(forward) + provider = WorkloadIdentitySessionTokenProvider( + base_url="https://api/execution/";, identity_token=identity_token, transport=transport + ) + processor = DagProcessorAPIClient( + base_url="https://api/execution/";, + hostname="processor", + token_provider=provider, + transport=transport, + ) + processors.append(processor) + return processor + + yield create + for processor in processors: + processor.close() + clear_db_jobs() + + +def exchange(client, token, **kwargs): + return client.post(URL, headers={"Authorization": f"Bearer {token}"}, **kwargs) + + +def decode(response): + assert response.status_code == 200, response.text + return jwt.decode(response.json()["token"], SECRET, algorithms=["HS512"], options={"verify_aud": False}) + + [email protected](("external_lifetime", "expected"), [(60, 60), (600, 300)]) [email protected]("identity_algorithm", ["RS256", "ES256", "EdDSA"], indirect=True) +def test_issues_only_configured_grants_and_caps_lifetime( + client, verifier, identity_token, external_lifetime, expected, caplog +): + response = exchange( + client, + identity_token(exp=NOW + external_lifetime, dag_bundles=["unauthorized"], scope="execution"), + json={"bundle_names": ["unauthorized"], "session_id": str(uuid7()), "hostname": "another-pod"}, + ) + claims = decode(response) + assert claims["scope"] == "dag_processor_session" + assert claims["dag_bundles"] == ["dags-folder"] + assert claims["exp"] == NOW + expected + assert response.json()["expires_in"] == expected + assert response.headers["Cache-Control"] == "no-store" + assert response.headers["Pragma"] == "no-cache" + assert "Refreshed-API-Token" not in response.headers + assert "JWT reissue middleware failed to refresh token" not in caplog + + +@conf_vars( + { + ("dag_processor", "dag_bundle_config_list"): json.dumps( + [ + { + "name": name, + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {}, + } + for name in ("dags-folder", "another-bundle") + ] + ) + } +) +def test_renewal_and_api_restart_keep_session_but_replicas_and_grants_do_not( + client, verifier, exec_app, identity_token, monkeypatch +): + first = decode(exchange(client, identity_token())) + renewed = decode(exchange(client, identity_token(exp=NOW + 900))) + replacement = DagProcessorIdentityVerifier(verifier.config, client=verifier.validator.jwks.client) + monkeypatch.setattr(exec_app.state, "dag_processor_identity_verifier", replacement) + restarted = decode(exchange(client, identity_token())) + replica = decode(exchange(client, identity_token(**{"kubernetes.io": {"pod": {"uid": "pod-2"}}}))) + assert first["sub"] == renewed["sub"] == restarted["sub"] + assert first["jti"] != renewed["jti"] + assert replica["sub"] != first["sub"] + changed = WorkloadIdentityConfig.model_validate( + {**TRUST, "grants": {"system:serviceaccount:airflow:processor": ["another-bundle"]}} + ) + replacement = DagProcessorIdentityVerifier(changed, client=verifier.validator.jwks.client) + changed_identity = client.portal.call(replacement.authenticate, identity_token()) + assert str(changed_identity.session_id) != first["sub"] + + [email protected]( + "overrides", + [ + {"iss": "https://untrusted.example/"}, + {"aud": "urn:airflow.apache.org:task"}, + {"exp": []}, + {"exp": float("inf")}, + {"iat": NOW + 60}, + {"nbf": NOW + 60}, + {"sub": ""}, + {"kubernetes.io": {}}, + {"kubernetes.io": {"pod": {"uid": ""}}}, + {"kubernetes.io": {"pod": {"uid": 123}}}, + ], +) +def test_rejects_invalid_signed_claims(client, verifier, identity_token, overrides): + response = exchange(client, identity_token(**overrides)) + assert response.status_code == 401 + assert response.json() == {"detail": "Invalid workload identity"} + assert response.headers["WWW-Authenticate"] == "Bearer" + + [email protected]("leeway", [0, 3, 10]) [email protected]("claim", ["iat", "nbf"]) +def test_identity_clock_skew_uses_configured_leeway( + client, verifier, exec_app, identity_token, monkeypatch, leeway, claim +): + with conf_vars({("api_auth", "jwt_leeway"): str(leeway)}): + replacement = DagProcessorIdentityVerifier(verifier.config, client=verifier.jwks.client) + monkeypatch.setattr(exec_app.state, "dag_processor_identity_verifier", replacement) + assert exchange(client, identity_token(**{claim: NOW + leeway})).status_code == 200 + assert exchange(client, identity_token(**{claim: NOW + leeway + 1})).status_code == 401 + + [email protected]("remaining", [-11, -1, 0, 0.5]) +def test_clock_leeway_never_issues_session_after_identity_expiry(client, verifier, identity_token, remaining): + response = exchange(client, identity_token(exp=NOW + remaining)) + assert response.status_code == 401 + assert "token" not in response.json() + + [email protected]("claim", ["iss", "aud", "sub", "exp", "iat", "kubernetes.io"]) +def test_requires_identity_claims(client, verifier, identity_token, claim): + assert exchange(client, identity_token(remove=[claim])).status_code == 401 + + [email protected]("invalid", ["malformed", "wrong-key", "unknown-kid", "no-kid", "algorithm"]) +def test_rejects_invalid_signatures(client, verifier, identity_token, invalid): + token = { + "malformed": "not-a-jwt", + "wrong-key": identity_token(key=Ed25519PrivateKey.generate()), + "unknown-kid": identity_token(headers={"kid": "unknown"}), + "no-kid": identity_token(headers={"typ": "JWT"}), + "algorithm": jwt.encode({"sub": "x"}, SECRET, algorithm="HS512", headers={"kid": "key-1"}), + }[invalid] + assert exchange(client, token).status_code == 401 + + [email protected]("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) +def test_rotated_signing_key_refreshes_after_cooldown( + clock, client, verifier, identity_token, signing_key, monkeypatch +): + keys = [key_to_jwk_dict(signing_key, kid="key-1")] + requests = [] + + async def get_keys(url): + requests.append(url) + return httpx.Response(200, json={"keys": keys}, request=httpx.Request("GET", url)) + + monkeypatch.setattr(verifier.validator.jwks.client, "get", get_keys) + original = decode(exchange(client, identity_token())) + rotated_key = Ed25519PrivateKey.generate() + keys.append(key_to_jwk_dict(rotated_key, kid="key-2")) + rotated = identity_token(key=rotated_key, headers={"kid": "key-2"}) + assert exchange(client, rotated).status_code == 401 + assert len(requests) == 1 + + clock.return_value = 111 + assert decode(exchange(client, rotated))["sub"] == original["sub"] + assert len(requests) == 2 + assert exchange(client, identity_token(headers={"kid": "unknown"})).status_code == 401 + assert len(requests) == 2 + + +def test_ungranted_identity_is_forbidden(client, verifier, identity_token): + assert exchange(client, identity_token(sub="another-service-account")).status_code == 403 + + +def test_missing_bearer_is_unauthorized(client, verifier): + client.headers.pop("Authorization") + assert client.post(URL).status_code == 401 + + +def test_disabled_exchange_is_unavailable(client, exec_app, identity_token, monkeypatch): + monkeypatch.setattr(exec_app.state, "dag_processor_identity_verifier", None) + assert exchange(client, identity_token()).status_code == 404 + + +@conf_vars({("execution_api", "dag_processor_workload_identity"): ""}) +def test_empty_configuration_disables_verifier(client): + async def check(): + async with create_dag_processor_identity_verifier() as verifier: + assert verifier is None + + client.portal.call(check) + + [email protected] +async def test_verifier_closes_jwks_connection_pool(): + async with create_dag_processor_identity_verifier() as verifier: + client = verifier.validator.jwks.client + assert not client.is_closed + assert client.is_closed + + +def test_identity_expiring_during_exchange_is_rejected( + client, verifier, identity_token, monkeypatch, time_machine +): + authenticate = verifier.authenticate + + async def expire(token): + identity = await authenticate(token) + time_machine.shift(600) + return identity + + monkeypatch.setattr(verifier, "authenticate", expire) + assert exchange(client, identity_token()).status_code == 401 + + [email protected]("failure", ["connection", "503", "invalid-json", "invalid-keyset"]) +def test_jwks_outage_is_retryable(client, verifier, identity_token, monkeypatch, failure): + async def unavailable(url): + if failure == "connection": + raise httpx.ConnectError("offline") + request = httpx.Request("GET", url) + if failure == "503": + return httpx.Response(503, request=request) + if failure == "invalid-json": + return httpx.Response(200, content=b"unavailable", request=request) + return httpx.Response(200, json={"keys": []}, request=request) + + monkeypatch.setattr(verifier.jwks.client, "get", unavailable) + response = exchange(client, identity_token()) + assert response.status_code == 503 + assert response.json() == {"detail": "Workload identity signing keys are unavailable"} + assert response.headers["Retry-After"] == "10" + + [email protected]("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) +def test_jwks_outage_preserves_cached_key_validation(clock, client, verifier, identity_token, monkeypatch): + assert exchange(client, identity_token()).status_code == 200 + + async def unavailable(url): + raise httpx.ConnectError("offline") + + monkeypatch.setattr(verifier.jwks.client, "get", unavailable) + clock.return_value = 3701 + assert exchange(client, identity_token()).status_code == 200 + assert exchange(client, identity_token(headers={"kid": "new-key"})).status_code == 503 + assert exchange(client, identity_token(key=Ed25519PrivateKey.generate())).status_code == 401 + + [email protected]("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) +def test_registration_recovers_from_jwks_outage(clock, make_processor, verifier, signing_key, monkeypatch): + attempts = [] + + async def get_keys(url): + attempts.append(url) + return httpx.Response( + 503 if len(attempts) == 1 else 200, + json={"keys": [key_to_jwk_dict(signing_key, kid="key-1")]}, + request=httpx.Request("GET", url), + ) + + def advance_retry(retry_state): + clock.return_value += 11 + + monkeypatch.setattr(verifier.jwks.client, "get", get_keys) + monkeypatch.setattr(Client._request_with_retry.retry, "wait", wait_none()) + monkeypatch.setattr(Client._request_with_retry.retry, "before_sleep", advance_retry) + processor = make_processor() + assert processor.register_job() == processor.job_id + assert processor.heartbeat() == JobState.RUNNING + assert len(attempts) == 2 + + +@conf_vars({("dag_processor", "health_check_threshold"): "30"}) +def test_container_restart_waits_for_previous_job_to_become_stale(make_processor, time_machine, session): + previous = make_processor() + previous_job = previous.register_job() + replacement = make_processor() + with pytest.raises(DagProcessorJobAlreadyRunning): + replacement.register_job() + time_machine.shift(31) + replacement_job = replacement.register_job() + assert replacement_job != previous_job + session.expire_all() + assert session.get(Job, previous_job).state == JobState.FAILED.value + assert session.get(Job, previous_job).session_id is None + assert session.get(Job, replacement_job).state == JobState.RUNNING.value + + +@conf_vars({("dag_processor", "dag_bundle_config_list"): "[]"}) +def test_identity_grants_must_name_configured_bundles(client, verifier): + config = WorkloadIdentityConfig.model_validate(TRUST) + with pytest.raises(ValueError, match=r"dag_bundle_config_list: \['dags-folder'\]"): + DagProcessorIdentityVerifier(config, client=verifier.validator.jwks.client) + + +def test_older_version_has_no_exchange(client, verifier, identity_token): + client.headers["Airflow-API-Version"] = "2026-10-30" + response = exchange(client, identity_token()) + assert response.status_code == 404 + assert response.json() == {"detail": "Not Found"} + + [email protected]( + "url", ["http://identity/keys";, "https://user:secret@identity/keys";, "https://identity/keys#fragment";] +) +def test_jwks_url_requires_https_without_embedded_credentials(url): + with pytest.raises(ValidationError): + WorkloadIdentityConfig.model_validate({**TRUST, "jwks_url": url}) + + +def test_identity_audience_must_differ_from_runtime(client, verifier): + config = WorkloadIdentityConfig.model_validate({**TRUST, "audience": "urn:airflow.apache.org:task"}) + with pytest.raises(ValueError, match="dedicated audience"): + DagProcessorIdentityVerifier(config, client=verifier.validator.jwks.client) + + +def test_identity_and_runtime_tokens_are_not_interchangeable(client, verifier, exec_app, identity_token): + exec_app.dependency_overrides.pop(require_auth) + identity = identity_token() + session_token = decode(exchange(client, identity)) + assert session_token["scope"] == "dag_processor_session" + issued_token = exchange(client, identity).json()["token"] + assert exchange(client, issued_token).status_code == 401 + response = client.post( + "/execution/jobs", + headers={"Authorization": f"Bearer {identity}"}, + json={"registration_id": str(uuid7()), "hostname": "processor"}, + ) + assert response.status_code == 403 + + [email protected]( Review Comment: The client gets `token_provider=` and no `token_file`, so `FileSessionTokenProvider` is never constructed, and this patch and the `file_provider.assert_not_called()` at the end can't fail. The `forward`/provider/client setup below also repeats `make_processor`; an `exchange_available` hook on the fixture would let this test reuse it. ########## airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py: ########## @@ -0,0 +1,145 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Authenticate deployment identities for Dag processor session provisioning.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Annotated, Any, Literal +from uuid import NAMESPACE_URL, UUID, uuid5 + +import httpx +import jwt +from pydantic import BaseModel, ConfigDict, Field, HttpUrl, field_validator + +from airflow.api_fastapi.auth.tokens import JWKS, JWTValidator +from airflow.api_fastapi.execution_api.datamodels.token import DagBundleGrant +from airflow.configuration import conf +from airflow.dag_processing.bundles.manager import get_configured_bundle_names + +NonemptyString = Annotated[str, Field(min_length=1)] + + +class WorkloadIdentityConfig(BaseModel): + """Trust settings owned by the deployment manager, never supplied by a request.""" + + model_config = ConfigDict(extra="forbid") + + issuer: NonemptyString + jwks_url: HttpUrl + audience: NonemptyString + algorithm: Literal["RS256", "ES256", "EdDSA"] = "RS256" + instance_id_claim: list[NonemptyString] = Field(min_length=1) + grants: dict[NonemptyString, DagBundleGrant] = Field(min_length=1) + + @field_validator("jwks_url") + @classmethod + def require_https(cls, value: HttpUrl) -> HttpUrl: + if value.scheme != "https" or value.username or value.password or value.fragment: + raise ValueError( + "The workload identity JWKS URL must use HTTPS without credentials or a fragment" + ) + return value + + +class WorkloadIdentityNotAllowed(PermissionError): + """The verified identity has no Dag processor grant.""" + + +@dataclass(frozen=True) +class DagProcessorIdentity: + """Verified identity with server-selected grants and a stable, non-secret session ID.""" + + session_id: UUID + bundle_names: frozenset[str] + expires_at: float + + +class _IdentityClaims(BaseModel): + sub: NonemptyString + exp: float = Field(allow_inf_nan=False) + + +class DagProcessorIdentityVerifier: + """Validate external JWTs independently of Airflow's runtime token trust.""" + + def __init__(self, config: WorkloadIdentityConfig, *, client: httpx.AsyncClient): + if config.audience in conf.getlist("execution_api", "jwt_audience"): + raise ValueError("Workload identity must use a dedicated audience") + if unknown := frozenset().union(*config.grants.values()) - get_configured_bundle_names(): + raise ValueError( + "[execution_api] dag_processor_workload_identity grants Dag bundles missing from " + f"[dag_processor] dag_bundle_config_list: {sorted(unknown)}" + ) + self.config = config + self.jwks = JWKS( + url=str(config.jwks_url), client=client, refresh_on_unknown_kid=True, raise_on_fetch_error=True + ) + self.validator = JWTValidator( + jwks=self.jwks, + issuer=config.issuer, + audience=config.audience, + algorithm=[config.algorithm], + required_claims=frozenset({"iss", "sub", "aud", "exp", "iat"}), + ) + + async def authenticate(self, token: str) -> DagProcessorIdentity: + claims = await self.validator.avalidated_claims(token) + identity = _IdentityClaims.model_validate(claims) + if not (grants := self.config.grants.get(identity.sub)): + raise WorkloadIdentityNotAllowed("Workload identity is not authorized as a Dag processor") + instance_id: Any = claims Review Comment: `object` works here, since the two `isinstance` checks already narrow it, so the `Any` isn't needed. ########## airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py: ########## @@ -141,6 +146,12 @@ def _all_route_policies() -> dict[str, set[str]]: class TestTokenScopeBoundaries: """Execution API routes must not silently gain or lose token type access.""" + def test_session_bootstrap_requires_deployment_identity(self): Review Comment: `_all_route_policies` still lists this route as `{"execution"}`, same as `/health`, so the policy table calls it execution-only when it takes no runtime token at all. Now that there are two deliberately unauthenticated routes, how about an invariant over all routes instead: the set without `require_auth` equals an explicit allowlist, and the non-health entries depend on `require_workload_identity`? That would catch the next route added outside `authenticated_router` by mistake. ########## airflow-core/docs/security/jwt_token_authentication.rst: ########## @@ -519,6 +519,125 @@ API without holding its signing key. The standard ``airflow dag-processor`` comm not yet use this client. Provisioning a token file alone does not switch that command to HTTP or remove its need for database access. +In-memory sessions from workload identity +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Processor integrations can use ``WorkloadIdentitySessionTokenProvider`` to obtain and renew +``dag_processor_session`` credentials through ``POST /execution/dag-processor/session-token``. +The provider retains these credentials only in memory. It requires a renewable, signed workload +identity from the deployment platform; Airflow does not create that identity or require the +processor to hold an Airflow signing key. + +Configure the API server's ``[execution_api] dag_processor_workload_identity`` option with JSON:: + + { + "issuer": "https://identity.example/";, + "jwks_url": "https://identity.example/openid/v1/jwks";, + "audience": "airflow-dag-processor-bootstrap", + "algorithm": "RS256", + "instance_id_claim": ["kubernetes.io", "pod", "uid"], + "grants": { + "system:serviceaccount:airflow:processor": ["dags-folder"] + } + } + +Use your identity provider's exact issuer, HTTPS JWKS URL, and signing algorithm (``RS256``, +``ES256``, or ``EdDSA``). Use a dedicated audience, distinct from ``[execution_api] jwt_audience``. +The API server verifies the signature, issuer, audience, issued-at time, expiry, and any not-before +claim, using ``[api_auth] jwt_leeway`` for clock skew. Only exact subjects listed in ``grants`` can +obtain a session. Every granted bundle must be configured on the API server, which otherwise refuses +to start. This option is empty by default, disabling the exchange. +Apply the same trust configuration to all API-server replicas and restart them after changes. + +The JWKS endpoint must accept unauthenticated HTTPS GET requests from the Airflow API server. +The fetcher sends no authentication credentials and uses HTTPX's default certificate trust, +including ``SSL_CERT_FILE`` or ``SSL_CERT_DIR`` when configured. These environment settings affect Review Comment: Could this say that `SSL_CERT_FILE` and `SSL_CERT_DIR` replace the default CA bundle rather than adding to it? With the httpx pinned in uv.lock the variable becomes the only trust source, so pointing it at just the private CA for the JWKS host cuts every other HTTPX client in the API server off from the public roots. ########## airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py: ########## @@ -0,0 +1,105 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Exchange an authenticated workload identity for a Dag processor session.""" + +from __future__ import annotations + +from typing import Annotated + +import jwt +import svcs +from cadwyn import VersionedAPIRouter +from fastapi import Depends, HTTPException, Request, Response, status +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from pydantic import ValidationError + +from airflow.api_fastapi.auth.tokens import JWKSUnavailableError, JWTGenerator +from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.execution_api.dag_processor_tokens import ( + ExpiredDagProcessorToken, + generate_dag_processor_session_token, + get_token_lifetime, +) +from airflow.api_fastapi.execution_api.datamodels.job import DagProcessorSessionTokenResponse +from airflow.api_fastapi.execution_api.deps import DepContainer +from airflow.api_fastapi.execution_api.workload_identity import ( + DagProcessorIdentity, + DagProcessorIdentityVerifier, + WorkloadIdentityNotAllowed, +) + +router = VersionedAPIRouter() +identity_bearer = HTTPBearer(auto_error=False, scheme_name="DagProcessorWorkloadIdentity") + + +async def require_workload_identity( + request: Request, + credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(identity_bearer)], +) -> DagProcessorIdentity: + verifier: DagProcessorIdentityVerifier | None = request.app.state.dag_processor_identity_verifier + if verifier is None: + raise HTTPException(status.HTTP_404_NOT_FOUND, "Workload identity provisioning is not enabled") + if credentials is None: + raise HTTPException( + status.HTTP_401_UNAUTHORIZED, "Workload identity required", headers={"WWW-Authenticate": "Bearer"} + ) + try: + return await verifier.authenticate(credentials.credentials) + except JWKSUnavailableError: + raise HTTPException( + status.HTTP_503_SERVICE_UNAVAILABLE, + "Workload identity signing keys are unavailable", + headers={"Retry-After": str(verifier.jwks.refresh_retry_interval_secs)}, + ) from None + except (jwt.PyJWTError, ValidationError, TypeError, OverflowError): + raise HTTPException( + status.HTTP_401_UNAUTHORIZED, "Invalid workload identity", headers={"WWW-Authenticate": "Bearer"} + ) from None + except WorkloadIdentityNotAllowed: + raise HTTPException( + status.HTTP_403_FORBIDDEN, "Workload identity is not granted any Dag bundles" + ) from None + + [email protected]( + "/dag-processor/session-token", + responses=create_openapi_http_exception_doc( + [ + (status.HTTP_401_UNAUTHORIZED, "The workload identity is missing, invalid, or expired"), + (status.HTTP_403_FORBIDDEN, "The identity is not authorized as a Dag processor"), + (status.HTTP_404_NOT_FOUND, "Workload identity provisioning is disabled"), + (status.HTTP_503_SERVICE_UNAVAILABLE, "The identity provider's signing keys are unavailable"), + ] + ), +) +def exchange_session_token( + response: Response, + identity: Annotated[DagProcessorIdentity, Depends(require_workload_identity)], + services: svcs.Container = DepContainer, +) -> DagProcessorSessionTokenResponse: + generator: JWTGenerator = services.get(JWTGenerator) + try: + lifetime = get_token_lifetime(generator, identity.expires_at, "Workload identity has expired") + except ExpiredDagProcessorToken as error: + raise HTTPException(status.HTTP_401_UNAUTHORIZED, str(error)) from None Review Comment: This 401 doesn't send `WWW-Authenticate: Bearer`, unlike the two in `require_workload_identity`. ########## airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py: ########## @@ -0,0 +1,105 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Exchange an authenticated workload identity for a Dag processor session.""" + +from __future__ import annotations + +from typing import Annotated + +import jwt +import svcs +from cadwyn import VersionedAPIRouter +from fastapi import Depends, HTTPException, Request, Response, status +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from pydantic import ValidationError + +from airflow.api_fastapi.auth.tokens import JWKSUnavailableError, JWTGenerator +from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.execution_api.dag_processor_tokens import ( + ExpiredDagProcessorToken, + generate_dag_processor_session_token, + get_token_lifetime, +) +from airflow.api_fastapi.execution_api.datamodels.job import DagProcessorSessionTokenResponse +from airflow.api_fastapi.execution_api.deps import DepContainer +from airflow.api_fastapi.execution_api.workload_identity import ( + DagProcessorIdentity, + DagProcessorIdentityVerifier, + WorkloadIdentityNotAllowed, +) + +router = VersionedAPIRouter() +identity_bearer = HTTPBearer(auto_error=False, scheme_name="DagProcessorWorkloadIdentity") + + +async def require_workload_identity( + request: Request, + credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(identity_bearer)], +) -> DagProcessorIdentity: + verifier: DagProcessorIdentityVerifier | None = request.app.state.dag_processor_identity_verifier + if verifier is None: + raise HTTPException(status.HTTP_404_NOT_FOUND, "Workload identity provisioning is not enabled") + if credentials is None: + raise HTTPException( + status.HTTP_401_UNAUTHORIZED, "Workload identity required", headers={"WWW-Authenticate": "Bearer"} + ) + try: + return await verifier.authenticate(credentials.credentials) + except JWKSUnavailableError: + raise HTTPException( + status.HTTP_503_SERVICE_UNAVAILABLE, + "Workload identity signing keys are unavailable", + headers={"Retry-After": str(verifier.jwks.refresh_retry_interval_secs)}, + ) from None + except (jwt.PyJWTError, ValidationError, TypeError, OverflowError): Review Comment: Where would `TypeError` or `OverflowError` come from here? PyJWT already turns bad `exp`/`iat`/`nbf` values into `DecodeError`, and `_IdentityClaims` gives a `ValidationError`. As written, a bug in `authenticate` or the JWKS code would come back as "Invalid workload identity" with no log line, instead of a 500. ########## airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py: ########## @@ -0,0 +1,105 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Exchange an authenticated workload identity for a Dag processor session.""" + +from __future__ import annotations + +from typing import Annotated + +import jwt +import svcs +from cadwyn import VersionedAPIRouter +from fastapi import Depends, HTTPException, Request, Response, status +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from pydantic import ValidationError + +from airflow.api_fastapi.auth.tokens import JWKSUnavailableError, JWTGenerator +from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.execution_api.dag_processor_tokens import ( + ExpiredDagProcessorToken, + generate_dag_processor_session_token, + get_token_lifetime, +) +from airflow.api_fastapi.execution_api.datamodels.job import DagProcessorSessionTokenResponse +from airflow.api_fastapi.execution_api.deps import DepContainer +from airflow.api_fastapi.execution_api.workload_identity import ( + DagProcessorIdentity, + DagProcessorIdentityVerifier, + WorkloadIdentityNotAllowed, +) + +router = VersionedAPIRouter() +identity_bearer = HTTPBearer(auto_error=False, scheme_name="DagProcessorWorkloadIdentity") + + +async def require_workload_identity( + request: Request, + credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(identity_bearer)], +) -> DagProcessorIdentity: + verifier: DagProcessorIdentityVerifier | None = request.app.state.dag_processor_identity_verifier + if verifier is None: + raise HTTPException(status.HTTP_404_NOT_FOUND, "Workload identity provisioning is not enabled") + if credentials is None: + raise HTTPException( + status.HTTP_401_UNAUTHORIZED, "Workload identity required", headers={"WWW-Authenticate": "Bearer"} + ) + try: + return await verifier.authenticate(credentials.credentials) + except JWKSUnavailableError: + raise HTTPException( + status.HTTP_503_SERVICE_UNAVAILABLE, + "Workload identity signing keys are unavailable", + headers={"Retry-After": str(verifier.jwks.refresh_retry_interval_secs)}, + ) from None + except (jwt.PyJWTError, ValidationError, TypeError, OverflowError): + raise HTTPException( + status.HTTP_401_UNAUTHORIZED, "Invalid workload identity", headers={"WWW-Authenticate": "Bearer"} + ) from None + except WorkloadIdentityNotAllowed: + raise HTTPException( + status.HTTP_403_FORBIDDEN, "Workload identity is not granted any Dag bundles" + ) from None + + [email protected]( + "/dag-processor/session-token", + responses=create_openapi_http_exception_doc( + [ + (status.HTTP_401_UNAUTHORIZED, "The workload identity is missing, invalid, or expired"), + (status.HTTP_403_FORBIDDEN, "The identity is not authorized as a Dag processor"), + (status.HTTP_404_NOT_FOUND, "Workload identity provisioning is disabled"), + (status.HTTP_503_SERVICE_UNAVAILABLE, "The identity provider's signing keys are unavailable"), + ] + ), +) +def exchange_session_token( Review Comment: No docstring, so this route has no description in the Execution API OpenAPI spec (the Job routes next to it all have one). A non-Python processor integration would read that to learn it takes an external OIDC identity rather than an Airflow token, and that the 503 carries `Retry-After`. ########## airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py: ########## @@ -0,0 +1,145 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Authenticate deployment identities for Dag processor session provisioning.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Annotated, Any, Literal +from uuid import NAMESPACE_URL, UUID, uuid5 + +import httpx +import jwt +from pydantic import BaseModel, ConfigDict, Field, HttpUrl, field_validator + +from airflow.api_fastapi.auth.tokens import JWKS, JWTValidator +from airflow.api_fastapi.execution_api.datamodels.token import DagBundleGrant +from airflow.configuration import conf +from airflow.dag_processing.bundles.manager import get_configured_bundle_names + +NonemptyString = Annotated[str, Field(min_length=1)] + + +class WorkloadIdentityConfig(BaseModel): + """Trust settings owned by the deployment manager, never supplied by a request.""" + + model_config = ConfigDict(extra="forbid") + + issuer: NonemptyString + jwks_url: HttpUrl + audience: NonemptyString + algorithm: Literal["RS256", "ES256", "EdDSA"] = "RS256" + instance_id_claim: list[NonemptyString] = Field(min_length=1) + grants: dict[NonemptyString, DagBundleGrant] = Field(min_length=1) + + @field_validator("jwks_url") + @classmethod + def require_https(cls, value: HttpUrl) -> HttpUrl: + if value.scheme != "https" or value.username or value.password or value.fragment: + raise ValueError( + "The workload identity JWKS URL must use HTTPS without credentials or a fragment" + ) + return value + + +class WorkloadIdentityNotAllowed(PermissionError): + """The verified identity has no Dag processor grant.""" + + +@dataclass(frozen=True) +class DagProcessorIdentity: + """Verified identity with server-selected grants and a stable, non-secret session ID.""" + + session_id: UUID + bundle_names: frozenset[str] + expires_at: float + + +class _IdentityClaims(BaseModel): + sub: NonemptyString + exp: float = Field(allow_inf_nan=False) + + +class DagProcessorIdentityVerifier: + """Validate external JWTs independently of Airflow's runtime token trust.""" + + def __init__(self, config: WorkloadIdentityConfig, *, client: httpx.AsyncClient): + if config.audience in conf.getlist("execution_api", "jwt_audience"): + raise ValueError("Workload identity must use a dedicated audience") + if unknown := frozenset().union(*config.grants.values()) - get_configured_bundle_names(): Review Comment: Since the core lifespan enters every sub-app lifespan in one `AsyncExitStack`, a grant naming an unconfigured bundle stops the whole API server, not just this exchange. Renaming a bundle and forgetting the grant would take down the UI and the Task Execution API that workers heartbeat against. Is that intended, or would logging the error and returning 503 from the exchange be enough? ########## airflow-core/src/airflow/api_fastapi/execution_api/app.py: ########## @@ -106,7 +110,15 @@ async def lifespan(app: FastAPI, registry: svcs.Registry): # Create an app scoped validator, so that we don't have to fetch it every time registry.register_value(JWTValidator, _jwt_validator(), ping=JWTValidator.status) - yield + # InProcessExecutionAPI stubs out the identity verifier: don't create one in that case. + if DagProcessorIdentityVerifier in registry: Review Comment: The verifier lives on `app.state` and the registry entry is only used as a flag, while `JWTValidator` and `JWTGenerator` are resolved through svcs. Registering the real verifier in the registry and resolving it through `DepContainer` in the route would keep one lookup path. As it is, a later `services.get(DagProcessorIdentityVerifier)` gets `None` in-process and `ServiceNotFoundError` on a real API server. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
