From 390910c1af0bbffdadffcd7865cea84f32d7428f Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 14 Oct 2026 08:12:04 +0100 Subject: [PATCH 1/6] Add OIDC workload identity for Dag processor session credentials Dag processor integrations currently need a trusted provisioner to write and rotate an Airflow session token on disk. Deployments that already provide renewable workload identities should not need a second credential file or an Airflow signing key on the processor. An opt-in identity exchange lets these integrations keep Airflow session credentials in memory while retaining the existing bundle grants and Job ownership model. Stable sessions avoid retiring active Jobs when the platform rotates an identity, and bounded lifetimes keep issued credentials within the authority of that identity. This is a follow-up to the Dag processor API foundation. The existing file-based provisioning path remains supported; wiring the standard Dag processor command to the HTTP client is separate work. --- .../security/jwt_token_authentication.rst | 127 ++++- airflow-core/docs/security/security_model.rst | 6 + .../src/airflow/api_fastapi/auth/tokens.py | 84 ++- .../airflow/api_fastapi/execution_api/app.py | 8 +- .../execution_api/dag_processor_tokens.py | 8 +- .../execution_api/datamodels/job.py | 7 + .../execution_api/routes/__init__.py | 2 + .../execution_api/routes/session_tokens.py | 114 ++++ .../execution_api/versions/__init__.py | 7 +- .../execution_api/versions/v2027_02_28.py | 10 + .../execution_api/workload_identity.py | 139 +++++ .../src/airflow/config_templates/config.yml | 20 + .../src/airflow/dag_processing/api_client.py | 63 +- .../src/airflow/dag_processing/credentials.py | 121 ++++ .../unit/api_fastapi/auth/test_tokens.py | 103 ++++ .../test_token_scope_boundaries.py | 19 +- .../versions/head/test_session_tokens.py | 539 ++++++++++++++++++ .../unit/dag_processing/test_credentials.py | 311 ++++++++++ .../airflow/sdk/api/datamodels/_generated.py | 14 + 19 files changed, 1665 insertions(+), 37 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py create mode 100644 airflow-core/src/airflow/dag_processing/credentials.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py create mode 100644 airflow-core/tests/unit/dag_processing/test_credentials.py diff --git a/airflow-core/docs/security/jwt_token_authentication.rst b/airflow-core/docs/security/jwt_token_authentication.rst index e72f9d3faa6a5..8a2c620389962 100644 --- a/airflow-core/docs/security/jwt_token_authentication.rst +++ b/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, and every granted bundle must be configured on the API server. 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 +other HTTPX clients in the API-server process; there is no per-issuer CA or authentication option. +Authenticated JWKS endpoints, including Kubernetes API-server discovery protected by service-account +authentication, are not supported by this implementation. For those clusters, serve the public JWKS +through a suitable HTTPS endpoint before enabling the exchange. See Kubernetes' +`service account issuer discovery documentation `_. +The processor-to-Airflow TLS settings described below do not configure this outbound JWKS connection. + +``instance_id_claim`` is a path through signed claims to a nonempty string identifying one +processor replica. A shared service-account subject alone is insufficient: replicas must have +different instance identities. Run one processor per identity. The session ID stays stable across +identity-token renewal and API-server restarts. Changes to the issuer, audience, subject, instance +claim, instance value, or bundle grants produce a different session. Request bodies and caller +hostnames cannot select the session or its grants. + +Changing grants while API-server replicas serve different configurations can repeatedly retire +processor registrations, because those replicas issue different session IDs for the same identity. +Stop the affected processors, finish deploying identical trust configuration to every API-server +replica, and then restart the processors. + +For Kubernetes, the example uses the Pod UID in a +`bound service-account token `_. +Project a token with the dedicated audience and mount the projection directory read-only in the +processor container. Configure the projection's ``expirationSeconds`` and allow Kubernetes to +rotate it. See `serviceAccountToken projected volumes +`_. + +A container restart within the same Pod keeps the Pod UID and therefore the session ID. +A replacement process cannot register while the previous Job is still alive. The client raises +``DagProcessorJobAlreadyRunning``; the integration must retry registration with backoff until the +previous Job completes or its heartbeat becomes stale under ``[dag_processor] health_check_threshold``. +The client does not wait automatically. This prevents two processes sharing an identity from owning +active Jobs at the same time. + +The processor integration supplies a callable that reads the current platform identity on each +exchange. For example: + +.. code-block:: python + + from pathlib import Path + + from airflow.dag_processing.api_client import DagProcessorAPIClient + from airflow.dag_processing.credentials import WorkloadIdentitySessionTokenProvider + + identity_path = Path("/var/run/airflow-identity/token") + provider = WorkloadIdentitySessionTokenProvider( + base_url="https://airflow.example/execution/", + identity_token=identity_path.read_text, + ) + processor = DagProcessorAPIClient( + base_url="https://airflow.example/execution/", + token_provider=provider, + hostname="processor-1", + ) + +Use this client for the integration's Job lifecycle and close it during shutdown; closing it also +closes the provider and clears its session cache. The file in this example contains the platform's +projected identity, not an Airflow session credential. The provider never writes the exchanged +session JWT to disk and does not fall back to a session-token file. ``token_provider`` and +``token_file`` are mutually exclusive. HTTPS is required for the exchange, and redirects are rejected. +The exchange uses the same API TLS configuration and request timeout as the Task SDK client, +including private certificate authorities and client certificates. Registration retries transient +exchange failures using the configured Execution API retry policy; optional early renewal makes +one bounded attempt so it does not delay requests using a valid Job credential. + +Session expiry is capped by both the external identity's remaining lifetime and +``[execution_api] jwt_expiration_time``. Clock-skew leeway never extends this cap, and an identity +with less than one second remaining cannot obtain a session. Renewal rereads the platform identity +and retains the same session and Job. During an identity-provider or API outage, the client can continue using +an already-issued Job credential until it expires; it cannot renew without a valid identity. +Removing a grant stops future exchanges after API-server configuration is reloaded. Credentials +already issued remain subject to their expiry and Job ownership checks. + +Verification uses the issuer's JWKS and does not perform online token introspection. In particular, +it does not check whether a Kubernetes Pod still exists. A previously issued identity may remain +usable until expiry. The JWKS cache refreshes hourly and also refreshes when an identity names an +unknown signing key. Fetches, including failed attempts, are limited to one per ten seconds per +API-server process. Publish new keys before signing with them and retain old keys while their tokens +remain valid. + +If a required signing key cannot be fetched, the exchange returns HTTP 503 with ``Retry-After``, +allowing registration to retry. A cached matching key remains usable during an outage. Invalid +identities and keys absent from a successfully fetched keyset return HTTP 401. + +File-based provisioning +~~~~~~~~~~~~~~~~~~~~~~~ + On a trusted host with the API signing key, provision a session token for the bundles the client may access:: @@ -532,9 +651,12 @@ token replacement stays visible. Each processor process needs its own session. R provisioner with the same token file and bundles keeps that session, so running processors keep their Jobs. +Credential permissions and renewal +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + The three credentials have different issuers and permissions: -- ``dag_processor_session`` is issued by trusted provisioning. It can only register a +- ``dag_processor_session`` is issued by the identity exchange or trusted file provisioning. It can only register a Job or renew that Job's credential through ``POST /jobs``. - ``dag_processor`` is returned by Job registration. It can heartbeat and complete that Job, read Connections and Variables for a granted bundle, and exchange a parsing @@ -625,6 +747,9 @@ All JWT-related configuration parameters: * - ``[execution_api] jwt_audience`` - ``urn:airflow.apache.org:task`` - Audience claim for Execution API tokens. + * - ``[execution_api] dag_processor_workload_identity`` + - Empty (disabled) + - JSON trust settings and subject-to-bundle grants for workload identity session provisioning. .. important:: diff --git a/airflow-core/docs/security/security_model.rst b/airflow-core/docs/security/security_model.rst index ec5c340851dfc..277672e9899a5 100644 --- a/airflow-core/docs/security/security_model.rst +++ b/airflow-core/docs/security/security_model.rst @@ -540,6 +540,12 @@ potentially still executes with direct database access in the Dag File Processor and callback credentials, or ``DagProcessorSessionClaims``, ``DagProcessorClaims``, or ``DagParseClaims`` for processor credentials. Unknown scopes, invalid subjects, and missing required scope-specific claims are rejected before any route handler runs. + * **Workload identity provisioning.** The optional Dag processor session exchange validates + deployment-issued JWTs against separately configured issuer, audience, and JWKS settings. + Bundle grants come from the API server configuration. Exchanged sessions stay in the client's + memory, but workload identity and runtime credentials remain accessible to code with access + to that process. This does not add isolation between the processor and user code; see + :doc:`jwt_token_authentication` for configuration and expiry limitations. **Token signing key might be a shared secret** In symmetric key mode (``[api_auth] jwt_secret``), the same secret key is used to both generate and diff --git a/airflow-core/src/airflow/api_fastapi/auth/tokens.py b/airflow-core/src/airflow/api_fastapi/auth/tokens.py index dd6682b53c66b..df759a6dd7b7d 100644 --- a/airflow-core/src/airflow/api_fastapi/auth/tokens.py +++ b/airflow-core/src/airflow/api_fastapi/auth/tokens.py @@ -16,12 +16,14 @@ # under the License. from __future__ import annotations +import asyncio import json import os import time import uuid from base64 import urlsafe_b64encode from collections.abc import Callable, Sequence +from contextlib import nullcontext from datetime import datetime from typing import TYPE_CHECKING, Any, Literal, overload @@ -45,6 +47,7 @@ __all__ = [ "InvalidClaimError", "JWKS", + "JWKSUnavailableError", "JWTGenerator", "JWTValidator", "generate_private_key", @@ -99,9 +102,19 @@ def _guess_best_algorithm(key: AllowedPrivateKeys): raise ValueError(f"Unknown key object {type(key)}") +class JWKSUnavailableError(RuntimeError): + """The required signing key could not be loaded because a JWKS fetch failed.""" + + @attrs.define(repr=False) class JWKS: - """A class to fetch and sync a set of JSON Web Keys.""" + """ + Fetch and cache a set of JSON Web Keys. + + With ``raise_on_fetch_error``, missing keys raise ``JWKSUnavailableError`` after a failed + fetch. Cached matching keys remain usable, and a successfully fetched set with no matching key + still raises ``KeyError``. + """ url: str fetched_at: float = 0 @@ -113,6 +126,10 @@ class JWKS: refresh_jwks: bool = True refresh_interval_secs: int = 3600 refresh_retry_interval_secs: int = 10 + refresh_on_unknown_kid: bool = False + raise_on_fetch_error: bool = False + _fetch_lock: asyncio.Lock = attrs.field(factory=asyncio.Lock, init=False) + _fetch_failed: bool = attrs.field(default=False, init=False) def __repr__(self) -> str: return f"JWKS(url={self.url}, fetched_at={self.fetched_at})" @@ -128,19 +145,30 @@ 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(): - 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)) + async def fetch_jwks(self, *, force: bool = False) -> None: + 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): + if not self.raise_on_fetch_error: + raise + log.warning("Unable to load JWKS response", url=self.url, exc_info=True) + return + self._jwks = keyset + self._fetch_failed = False + log.debug("Fetched JWKS", url=self.url, keys=len(self._jwks.keys)) async def _fetch_remote_jwks(self) -> dict[str, Any] | None: try: @@ -151,7 +179,7 @@ async def _fetch_remote_jwks(self) -> dict[str, Any] | None: ) if TYPE_CHECKING: assert self.url - self.last_fetch_attempt_at = int(time.monotonic()) + self.last_fetch_attempt_at = time.monotonic() response = await self.client.get(self.url) response.raise_for_status() self.fetched_at = int(time.monotonic()) @@ -172,7 +200,7 @@ def _fetch_local_jwks(self) -> dict[str, Any] | None: log.exception("Failed to read local JWKS", url=self.url) return None - def _should_fetch_jwks(self) -> bool: + def _should_fetch_jwks(self, *, force: bool = False) -> bool: """ Check if we need to fetch the JWKS based on the last fetch time and the refresh interval. @@ -188,8 +216,16 @@ def _should_fetch_jwks(self) -> bool: # For remote fetches we check if the JWKS is not loaded (fetched_at = 0) or if the last fetch was more than # refresh_interval_secs ago and the last fetch attempt was more than refresh_retry_interval_secs ago now = time.monotonic() + # Unknown keys are untrusted input; bound fetches even for an empty cache or a failed refresh. + if ( + self.refresh_on_unknown_kid + and self.last_fetch_attempt_at + and now - self.last_fetch_attempt_at < self.refresh_retry_interval_secs + ): + return False return self.refresh_jwks and ( - not self._jwks + force + or not self._jwks or ( self.fetched_at == 0 or ( @@ -204,9 +240,19 @@ async def get_key(self, kid: str) -> jwt.PyJWK: await self.fetch_jwks() if self._jwks: - return self._jwks[kid] + try: + return self._jwks[kid] + except KeyError: + if self.refresh_on_unknown_kid: + await self.fetch_jwks(force=True) + try: + return self._jwks[kid] + except KeyError: + pass + + if self.raise_on_fetch_error and self._fetch_failed: + raise JWKSUnavailableError("Unable to load the required signing key") - # It didn't load! raise KeyError(f"Key ID {kid} not found in keyset") def status(self): diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index 21ed44c6764c0..a5bbec41f2fef 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -43,6 +43,7 @@ from airflow.api_fastapi.common.telemetry import NATIVE_TELEMETRY_DISABLED from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI from airflow.api_fastapi.execution_api.security import _REQUEST_SCOPE_TOKEN_KEY, require_auth +from airflow.api_fastapi.execution_api.workload_identity import create_dag_processor_identity_verifier if TYPE_CHECKING: from airflow.models.dagbag import DBDagBag @@ -107,7 +108,9 @@ 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 + async with create_dag_processor_identity_verifier() as verifier: + app.state.dag_processor_identity_verifier = verifier + yield class CorrelationIdMiddleware(BaseHTTPMiddleware): @@ -141,6 +144,9 @@ class JWTReissueMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): response: Response = await call_next(request) + if getattr(request.state, "skip_execution_token_reissue", False): + return response + refreshed_token: str | None = None token = request.scope.get(_REQUEST_SCOPE_TOKEN_KEY) if token: diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/dag_processor_tokens.py b/airflow-core/src/airflow/api_fastapi/execution_api/dag_processor_tokens.py index 535117d861190..216c16f8975fe 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/dag_processor_tokens.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/dag_processor_tokens.py @@ -43,10 +43,10 @@ class IssuedToken(NamedTuple): class ExpiredDagProcessorToken(ValueError): - """The parent credential expired before a Job or parsing credential could be issued.""" + """The parent credential expired before a session, Job, or parsing credential could be issued.""" -def _get_lifetime(generator: JWTGenerator, parent_expiry: float, message: str) -> int: +def get_token_lifetime(generator: JWTGenerator, parent_expiry: float, message: str) -> int: """Return whole seconds a child token may live: the signer's lifetime, capped by its parent's expiry.""" lifetime = math.floor(min(generator.valid_for, parent_expiry - timezone.utcnow().timestamp())) if lifetime < 1: @@ -77,7 +77,7 @@ def generate_dag_processor_token( session_expiry: float, ) -> IssuedToken: """Issue a Job credential that cannot outlive its provisioned session credential.""" - lifetime = _get_lifetime(generator, session_expiry, "Session credential has expired") + lifetime = get_token_lifetime(generator, session_expiry, "Session credential has expired") token = generator.generate( extras={ "sub": str(session_id), @@ -101,7 +101,7 @@ def generate_dag_parse_token( processor_expiry: float, ) -> IssuedToken: """Issue a parsing credential that cannot outlive its parent Job credential.""" - lifetime = _get_lifetime(generator, processor_expiry, "Processor credential has expired") + lifetime = get_token_lifetime(generator, processor_expiry, "Processor credential has expired") token = generator.generate( extras={ "sub": str(attempt_id), diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/job.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/job.py index df1b54bc6b86c..dbda578711dff 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/job.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/job.py @@ -60,6 +60,13 @@ class JobRegisterResponse(StrictBaseModel): expires_in: int = Field(gt=0, description="Seconds until the token expires.") +class DagProcessorSessionTokenResponse(StrictBaseModel): + """A session credential issued after workload identity verification.""" + + token: str = Field(min_length=1, repr=False) + expires_in: int = Field(gt=0, description="Seconds until the session credential expires.") + + class JobHeartbeatResponse(StrictBaseModel): """Current state of the Job; ``restarting`` asks the processor to stop.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py index 7ecc398a23ef9..f56f4d3b2bd65 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py @@ -31,6 +31,7 @@ health, hitl, jobs, + session_tokens, task_instances, task_reschedules, task_state_store, @@ -45,6 +46,7 @@ # prefix=... raises "Prefix and path cannot be both empty" once FastAPI switched to lazy router # inclusion (>=0.137); see https://github.com/apache/airflow/issues/68562. Don't reintroduce a prefix here. execution_api_router.include_router(health.router, tags=["Health"]) +execution_api_router.include_router(session_tokens.router, tags=["Dag Processor Identity"]) # _Every_ single endpoint under here must be authenticated. Some do further checks on top of these authenticated_router = VersionedAPIRouter(dependencies=[Security(require_auth)]) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py new file mode 100644 index 0000000000000..936935bea8dc6 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py @@ -0,0 +1,114 @@ +# 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, +) +from airflow.dag_processing.bundles.manager import get_configured_bundle_names + +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: + request.state.skip_execution_token_reissue = True + 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 + + +@router.post( + "/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, + "Signing keys or the configured bundle grant are unavailable", + ), + ] + ), +) +def exchange_session_token( + response: Response, + identity: Annotated[DagProcessorIdentity, Depends(require_workload_identity)], + services: svcs.Container = DepContainer, +) -> DagProcessorSessionTokenResponse: + if not identity.bundle_names <= get_configured_bundle_names(): + raise HTTPException( + status.HTTP_503_SERVICE_UNAVAILABLE, "Workload identity bundle grant is unavailable" + ) + 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 + token = generate_dag_processor_session_token( + generator, session_id=identity.session_id, bundle_names=identity.bundle_names, valid_for=lifetime + ) + response.headers["Cache-Control"] = "no-store" + response.headers["Pragma"] = "no-cache" + return DagProcessorSessionTokenResponse(token=token, expires_in=lifetime) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index d66d483f2aacd..6730dfc8d637a 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -70,11 +70,14 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30_xcom_params import ( AddPreviousIterationToXComFilterParams, ) -from airflow.api_fastapi.execution_api.versions.v2027_02_28 import AddDagProcessorJobEndpoints +from airflow.api_fastapi.execution_api.versions.v2027_02_28 import ( + AddDagProcessorJobEndpoints, + AddDagProcessorSessionTokenExchange, +) bundle = VersionBundle( HeadVersion(), - Version("2027-02-28", AddDagProcessorJobEndpoints), + Version("2027-02-28", AddDagProcessorJobEndpoints, AddDagProcessorSessionTokenExchange), Version( "2026-10-30", AddArgBindingsToTIRunContext, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py index b54a50236284a..80230d9aeee9b 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py @@ -31,3 +31,13 @@ class AddDagProcessorJobEndpoints(VersionChange): endpoint("/jobs/{job_id}/complete", ["POST"]).didnt_exist, endpoint("/jobs/{job_id}/parse-token", ["POST"]).didnt_exist, ) + + +class AddDagProcessorSessionTokenExchange(VersionChange): + """Exchange deployment workload identity for a Dag processor session credential.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + endpoint("/dag-processor/session-token", ["POST"]).didnt_exist, + ) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py b/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py new file mode 100644 index 0000000000000..bfb72bffa4399 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py @@ -0,0 +1,139 @@ +# 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 + +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") + 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 + for key in self.config.instance_id_claim: + if not isinstance(instance_id, dict) or key not in instance_id: + raise jwt.InvalidTokenError("Workload identity lacks its configured instance claim") + instance_id = instance_id[key] + if not isinstance(instance_id, str) or not instance_id: + raise jwt.InvalidTokenError("Workload identity instance claim must be a nonempty string") + # Bind to a signed replica identity, not a caller-supplied hostname or a shared service account. + # Including grants retires an existing registration when its authorization changes. + session_id = uuid5( + NAMESPACE_URL, + json.dumps( + [ + "airflow:dag-processor", + self.config.issuer, + self.config.audience, + identity.sub, + self.config.instance_id_claim, + instance_id, + sorted(grants), + ], + separators=(",", ":"), + ), + ) + return DagProcessorIdentity(session_id, grants, identity.exp) + + +@asynccontextmanager +async def create_dag_processor_identity_verifier() -> AsyncIterator[DagProcessorIdentityVerifier | None]: + """Create one verifier per application lifespan and close its JWKS connection pool.""" + raw = conf.get("execution_api", "dag_processor_workload_identity", fallback="") + if not raw: + yield None + return + config = WorkloadIdentityConfig.model_validate_json(raw) + async with httpx.AsyncClient(timeout=5.0, follow_redirects=False) as client: + yield DagProcessorIdentityVerifier(config, client=client) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index b87f27c509364..21f87cb91a5d3 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2424,6 +2424,26 @@ execution_api: The ExecutionAPI also uses a lot of settings from the :ref:`config:api_auth` section. options: + dag_processor_workload_identity: + description: | + JSON trust configuration for exchanging OIDC workload identity for an in-memory + Dag processor session token. Empty disables the exchange. + + Requires issuer, HTTPS jwks_url, a dedicated audience, an instance_id_claim path + (a list of keys into the signed identity claims), and grants mapping exact subjects + to allowed Dag bundle names. The optional algorithm defaults to RS256; ES256 and + EdDSA are also supported. Use a signed, per-replica instance claim: a service + account shared by replicas is not sufficient on its own. + The JWKS endpoint must accept unauthenticated HTTPS requests and have a certificate + trusted by the API server's HTTPX environment. Per-issuer CA and JWKS authentication + settings are not supported. + + Configure this only on API servers. Changes take effect when the API server restarts. + See :doc:`/security/jwt_token_authentication` for deployment examples. + version_added: 3.5.0 + type: string + example: ~ + default: "" jwt_expiration_time: description: | Number in seconds until the JWT used for authentication expires. When the token expires, diff --git a/airflow-core/src/airflow/dag_processing/api_client.py b/airflow-core/src/airflow/dag_processing/api_client.py index 673c1353ae1e7..78e06af366603 100644 --- a/airflow-core/src/airflow/dag_processing/api_client.py +++ b/airflow-core/src/airflow/dag_processing/api_client.py @@ -25,6 +25,7 @@ from dataclasses import dataclass, field from pathlib import Path from time import monotonic +from types import TracebackType from typing import Any from uuid import UUID @@ -43,6 +44,7 @@ TerminalJobState, ) from airflow.api_fastapi.execution_api.versions import bundle +from airflow.dag_processing.credentials import FileSessionTokenProvider, SessionTokenProvider from airflow.sdk.api.client import BearerAuth, Client from airflow.sdk.execution_time.comms import GetConnection, GetVariable, MaskSecret from airflow.sdk.execution_time.request_handlers import handle_get_connection, handle_get_variable @@ -105,7 +107,8 @@ def __init__( self, *, base_url: str, - token_file: str | Path, + token_file: str | Path | None = None, + token_provider: SessionTokenProvider | None = None, hostname: str, unixname: str | None = None, bundle_names: list[str] | None = None, @@ -114,13 +117,19 @@ def __init__( ): if not math.isfinite(token_reload_interval) or token_reload_interval < 0: raise ValueError("token_reload_interval must be finite and nonnegative") + if token_file is not None: + if token_provider is not None: + raise ValueError("Specify exactly one of token_file or token_provider") + token_provider = FileSessionTokenProvider(token_file) + elif token_provider is None: + raise ValueError("Specify exactly one of token_file or token_provider") self._registration = JobRegisterBody( registration_id=uuid7(), hostname=hostname, unixname=unixname, bundle_names=bundle_names, ) - self._token_file = Path(token_file) + self._token_provider = token_provider self._token_reload_interval = token_reload_interval self._session_token: str | None = None self._reload_at = 0.0 @@ -212,16 +221,51 @@ def _update_auth(self, response: httpx2.Response) -> None: # Task-token refresh headers cannot replace a provisioned or Job-bound credential. pass - def _read_session_token(self, *, force: bool = False) -> str: + def _read_session_token( + self, + *, + force: bool = False, + refresh: bool = False, + retry: bool = False, + timeout: httpx.Timeout | None = None, + ) -> str: now = monotonic() if self._session_token is None or force or now >= self._reload_at: - token = self._token_file.read_text().strip() + token = self._token_provider.get_token( + force_refresh=refresh, retry=retry, timeout=timeout + ).strip() if not token: - raise ValueError(f"Dag processor token file is empty: {self._token_file}") + raise ValueError("Dag processor session credential is empty") self._session_token = token self._reload_at = now + self._token_reload_interval return self._session_token + def _close_token_provider(self) -> None: + try: + self._token_provider.close() + finally: + self._session_token = None + self._registered_with = None + self.auth = BearerAuth("") + + def close(self) -> None: + """Close the owned credential provider and HTTP connection pool.""" + try: + self._close_token_provider() + finally: + super().close() + + def __exit__( + self, + exc_type: type[BaseException] | None = None, + exc_value: BaseException | None = None, + traceback: TracebackType | None = None, + ) -> None: + try: + self._close_token_provider() + finally: + super().__exit__(exc_type, exc_value, traceback) + def _check_can_run(self) -> None: if self._completion_state is not None: raise RuntimeError("The Dag processor Job is completing; only completion retries are allowed") @@ -253,7 +297,7 @@ def _register_job(self, *, retry: bool, timeout: httpx2.Timeout | None = None) - raise DagProcessorRegistrationRetired( "The Dag processor registration has ended; restart required" ) - session_token = self._read_session_token(force=True) + session_token = self._read_session_token(force=True, retry=retry, timeout=timeout) started_at = monotonic() body = self._registration.model_dump(mode="json") retried_auth = False @@ -277,7 +321,7 @@ def _register_job(self, *, retry: bool, timeout: httpx2.Timeout | None = None) - ) from error if error.response.status_code not in (401, 403) or retried_auth: raise - rotated = self._read_session_token(force=True) + rotated = self._read_session_token(force=True, refresh=True, retry=retry, timeout=timeout) if rotated == session_token: raise session_token = rotated @@ -312,7 +356,10 @@ def _ensure_job_token(self, *, retry: bool) -> None: if self.restart_required or monotonic() < self._retry_renewal_at: return try: - if self._read_session_token() != self._registered_with or monotonic() >= self._renew_at: + if ( + self._read_session_token(timeout=self._get_bounded_timeout()) != self._registered_with + or monotonic() >= self._renew_at + ): # Do not spend the runtime request's retry budget on an optional early renewal. self._register_job(retry=False, timeout=self._get_bounded_timeout()) except (httpx2.HTTPError, OSError, ValueError, DagProcessorRegistrationRetired) as error: diff --git a/airflow-core/src/airflow/dag_processing/credentials.py b/airflow-core/src/airflow/dag_processing/credentials.py new file mode 100644 index 0000000000000..69527f1287ba7 --- /dev/null +++ b/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 httpx + +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: httpx.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: httpx.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: httpx.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 httpx.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: httpx.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, + timeout=timeout or self._client.timeout, + ) + response.raise_for_status() + issued = DagProcessorSessionTokenResponse.model_validate_json(response.content) + expires_at = started_at + issued.expires_in + if monotonic() >= expires_at: + raise ValueError("Session credential expired during the exchange") + self._token = issued.token + self._renew_at = started_at + issued.expires_in * 0.8 + return self._token + + def close(self) -> None: + self._token = None + self._client.close() diff --git a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py index cc43e0e06e4ca..f6693743c1fa6 100644 --- a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py +++ b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py @@ -17,6 +17,7 @@ from __future__ import annotations +import asyncio import json import pathlib from datetime import datetime, timedelta @@ -33,6 +34,7 @@ from airflow.api_fastapi.auth.tokens import ( JWKS, InvalidClaimError, + JWKSUnavailableError, JWTGenerator, JWTValidator, generate_private_key, @@ -58,6 +60,107 @@ def private_key(request): class TestJWKS: + @pytest.mark.parametrize("refresh_on_unknown_kid", [False, True]) + @patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) + async def test_refresh_unknown_key_before_regular_interval( + self, clock, ed25519_private_key, refresh_on_unknown_kid + ): + keys = [key_to_jwk_dict(ed25519_private_key, "old")] + requests = [] + + def handle(request): + requests.append(request) + return httpx.Response(200, json={"keys": keys}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + jwks = JWKS( + url="https://example.com/keys", client=client, refresh_on_unknown_kid=refresh_on_unknown_kid + ) + await jwks.get_key("old") + keys.append(key_to_jwk_dict(generate_private_key("Ed25519"), "new")) + clock.return_value = 111 + if refresh_on_unknown_kid: + assert (await jwks.get_key("new")).key_id == "new" + assert len(requests) == 2 + else: + with pytest.raises(KeyError): + await jwks.get_key("new") + assert len(requests) == 1 + + @pytest.mark.parametrize("status", [200, 503]) + @pytest.mark.parametrize("raise_on_fetch_error", [False, True]) + @patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) + async def test_concurrent_unknown_keys_share_fetch_and_cooldown( + self, clock, ed25519_private_key, status, raise_on_fetch_error + ): + requests = [] + + async def handle(request): + requests.append(request) + await asyncio.sleep(0) + return httpx.Response(status, json={"keys": [key_to_jwk_dict(ed25519_private_key, "known")]}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + jwks = JWKS( + url="https://example.com/keys", + client=client, + refresh_on_unknown_kid=True, + raise_on_fetch_error=raise_on_fetch_error, + ) + results = await asyncio.gather( + *(jwks.get_key(f"unknown-{index}") for index in range(8)), return_exceptions=True + ) + error = JWKSUnavailableError if status == 503 and raise_on_fetch_error else KeyError + assert all(isinstance(result, error) for result in results) + assert len(requests) == 1 + clock.return_value = 111 + expected_requests = 1 if status == 200 else 2 + status = 200 + assert (await jwks.get_key("known")).key_id == "known" + assert len(requests) == expected_requests + results = await asyncio.gather( + *(jwks.get_key(f"unknown-{index}") for index in range(8)), return_exceptions=True + ) + assert all(isinstance(result, KeyError) for result in results) + assert len(requests) == 2 + + @pytest.mark.parametrize("raise_on_fetch_error", [False, True]) + @patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) + async def test_failed_unknown_key_refresh_keeps_cached_keys_and_rate_limits( + self, clock, ed25519_private_key, raise_on_fetch_error + ): + keys = [key_to_jwk_dict(ed25519_private_key, "old")] + status = 200 + requests = [] + + def handle(request): + requests.append(request) + return httpx.Response(status, json={"keys": keys}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + jwks = JWKS( + url="https://example.com/keys", + client=client, + refresh_on_unknown_kid=True, + raise_on_fetch_error=raise_on_fetch_error, + ) + old = await jwks.get_key("old") + clock.return_value = 111 + status = 503 + error = JWKSUnavailableError if raise_on_fetch_error else KeyError + with pytest.raises(error): + await jwks.get_key("new") + clock.return_value = 112 + with pytest.raises(error): + await jwks.get_key("another") + assert await jwks.get_key("old") is old + assert len(requests) == 2 + clock.return_value = 122 + status = 200 + keys.append(key_to_jwk_dict(generate_private_key("Ed25519"), "new")) + assert (await jwks.get_key("new")).key_id == "new" + assert len(requests) == 3 + @pytest.mark.parametrize("private_key", ["rsa_private_key", "ed25519_private_key"], indirect=True) async def test_fetch_jwks_success(self, private_key): jwk_content = json.dumps({"keys": [key_to_jwk_dict(private_key, "kid")]}) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py index 9535ec9522cbc..1e7e00cedb27d 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py @@ -43,7 +43,12 @@ from fastapi.routing import APIRoute from airflow.api_fastapi.execution_api.routes import execution_api_router -from airflow.api_fastapi.execution_api.security import get_selected_dag_bundle, require_dag_in_granted_bundle +from airflow.api_fastapi.execution_api.routes.session_tokens import require_workload_identity +from airflow.api_fastapi.execution_api.security import ( + get_selected_dag_bundle, + require_auth, + require_dag_in_granted_bundle, +) from airflow.dag_processing.processor import ToManager from tests_common.test_utils.api_fastapi import iter_api_routes @@ -51,6 +56,8 @@ # Routes that intentionally deviate from the default (execution-only) policy. # Any route NOT listed here must accept only {"execution"}. NON_DEFAULT_TOKEN_POLICY: dict[str, set[str]] = { + # Bootstrap authenticates a deployment identity, never an Airflow runtime token. + "POST /dag-processor/session-token": set(), # The /run endpoint exchanges a workload token for a short-lived execution token. "PATCH /task-instances/{task_instance_id}/run": {"execution", "workload"}, # Connection test routes run from a queued worker context (workload-only). @@ -145,7 +152,9 @@ def _get_path_and_query_param_names(dependant: Dependant) -> set[str]: def _all_route_policies() -> dict[str, set[str]]: """Return a map of all API routes and their allowed token types.""" return { - key: set(getattr(_get_original_route(route), "allowed_token_types", {"execution"})) + key: set() + if key == "POST /dag-processor/session-token" + else set(getattr(_get_original_route(route), "allowed_token_types", {"execution"})) for key, route in _get_api_routes().items() } @@ -153,6 +162,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): + route = _get_api_routes()["POST /dag-processor/session-token"] + dependencies = _get_dependency_calls(route.dependant) + assert require_workload_identity in dependencies + assert require_auth not in dependencies + def test_all_default_routes_are_execution_only(self): actual = _all_route_policies() non_default = { diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py new file mode 100644 index 0000000000000..68b57735dee61 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py @@ -0,0 +1,539 @@ +# 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 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"]}, +} + + +@pytest.fixture(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 + + +@pytest.fixture +def identity_algorithm(request): + return getattr(request, "param", "EdDSA") + + +@pytest.fixture +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() + + +@pytest.fixture +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 + + +@pytest.fixture +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) + + +@pytest.fixture +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 httpx.Response(response.status_code, headers=response.headers, content=response.content) + + def create(): + transport = httpx.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}) + + +@pytest.mark.parametrize(("external_lifetime", "expected"), [(60, 60), (600, 300)]) +@pytest.mark.parametrize("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 + + +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"] + + +@pytest.mark.parametrize( + "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" + + +@pytest.mark.parametrize("leeway", [0, 3, 10]) +@pytest.mark.parametrize("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 + + +@pytest.mark.parametrize("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() + + +@pytest.mark.parametrize("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 + + +@pytest.mark.parametrize("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 + + +@mock.patch("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) + + +@pytest.mark.asyncio +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 + + +@pytest.mark.parametrize("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" + + +@mock.patch("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 + + +@mock.patch("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_unknown_bundle_grant_does_not_issue_credentials(client, verifier, identity_token): + assert exchange(client, identity_token()).status_code == 503 + + +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"} + + +@pytest.mark.parametrize( + "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 + + +@mock.patch.object( + FileSessionTokenProvider, "get_token", autospec=True, side_effect=AssertionError("file used") +) +@mock.patch("airflow.dag_processing.api_client.monotonic", autospec=True, return_value=0) +@mock.patch("airflow.dag_processing.credentials.monotonic", autospec=True, return_value=0) +def test_session_renewal_and_job_lifecycle_without_session_file( + credential_clock, + job_clock, + file_provider, + client, + verifier, + exec_app, + identity_token, + session, + time_machine, +): + clear_db_jobs() + clear_db_variables() + exec_app.dependency_overrides.pop(require_auth) + exchange_available = True + + def forward(request): + if request.url.path == URL and not exchange_available: + raise httpx.ConnectError("session exchange unavailable", request=request) + response = client.request( + request.method, request.url.path, headers=request.headers, content=request.content + ) + return httpx.Response(response.status_code, headers=response.headers, content=response.content) + + provider = WorkloadIdentitySessionTokenProvider( + base_url="https://api/execution/", + identity_token=lambda: identity_token(exp=NOW + 600), + transport=httpx.MockTransport(forward), + ) + try: + Variable.set("identity-test-key", "value") + with DagProcessorAPIClient( + base_url="https://api/execution/", + hostname="processor", + token_provider=provider, + transport=httpx.MockTransport(forward), + ) as processor: + job_id = processor.register_job() + first_session = provider.get_token() + assert processor.heartbeat() == JobState.RUNNING + context = DagParseContext( + request=DagParseTokenBody( + attempt_id=uuid7(), bundle_name="dags-folder", relative_fileloc="dag.py" + ) + ) + with processor.use_parse(context): + assert processor.variables.get("identity-test-key").value == "value" + time_machine.shift(240) + credential_clock.return_value = job_clock.return_value = 240 + assert processor.heartbeat() == JobState.RUNNING + assert provider.get_token() != first_session + assert processor.job_id == job_id + assert not processor.restart_required + assert session.scalars(select(Job.id)).all() == [job_id] + + exchange_available = False + time_machine.shift(240) + credential_clock.return_value = job_clock.return_value = 480 + assert processor.heartbeat() == JobState.RUNNING + time_machine.shift(61) + credential_clock.return_value = job_clock.return_value = 541 + with pytest.raises(httpx.ConnectError): + processor.heartbeat() + + exchange_available = True + assert processor.heartbeat() == JobState.RUNNING + assert processor.job_id == job_id + processor.complete_job(TerminalJobState.SUCCESS) + assert session.get(Job, job_id).state == JobState.SUCCESS.value + file_provider.assert_not_called() + finally: + clear_db_jobs() + clear_db_variables() diff --git a/airflow-core/tests/unit/dag_processing/test_credentials.py b/airflow-core/tests/unit/dag_processing/test_credentials.py new file mode 100644 index 0000000000000..e1788bf46437a --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_credentials.py @@ -0,0 +1,311 @@ +# 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 + +from contextlib import nullcontext +from unittest import mock + +import httpx +import pytest +from pydantic import ValidationError +from tenacity import wait_none + +from airflow.api_fastapi.execution_api.versions import bundle +from airflow.dag_processing.api_client import DagProcessorAPIClient +from airflow.dag_processing.credentials import SessionTokenProvider, WorkloadIdentitySessionTokenProvider +from airflow.sdk.api.client import API_RETRIES, API_TIMEOUT, BearerAuth, Client + + +@pytest.fixture(autouse=True) +def no_retry_wait(monkeypatch): + monkeypatch.setattr(Client._request_with_retry.retry, "wait", wait_none()) + + +@pytest.fixture +def clock(): + with mock.patch("airflow.dag_processing.credentials.monotonic", autospec=True, return_value=0) as clock: + yield clock + + +@pytest.fixture +def make_provider(clock): + providers = [] + + def create(*responses, identity=None): + pending = iter(responses) + requests = [] + + def handle(request): + requests.append(request) + response = next(pending) + if isinstance(response, Exception): + raise response + return response(request) if callable(response) else response + + provider = WorkloadIdentitySessionTokenProvider( + base_url="https://api/execution/", + identity_token=identity or mock.Mock(spec=lambda: "", return_value="identity-1\n"), + transport=httpx.MockTransport(handle), + ) + providers.append(provider) + return provider, requests + + yield create + for provider in providers: + provider.close() + + +def issued(token="session-1", expires_in=100): + return httpx.Response(200, json={"token": token, "expires_in": expires_in}) + + +@mock.patch("airflow.sdk.api.client.API_SSL_CA_FILE_PATH", "/private-ca.pem") +@mock.patch("airflow.sdk.api.client.API_SSL_CERT_PATH", "/server.pem") +@mock.patch("airflow.sdk.api.client.API_CLIENT_SSL_CERT", "/client.pem") +@mock.patch("airflow.sdk.api.client.API_CLIENT_SSL_KEY", "/client.key") +@mock.patch.object( + Client, "_get_ssl_context_cached", new_callable=mock.Mock, spec=lambda ca_file, ca_path: None +) +def test_exchange_uses_api_tls_settings(get_context, make_provider, monkeypatch): + captured = {} + original_init = httpx.Client.__init__ + + def capture_init(self, **kwargs): + captured.update(kwargs) + original_init(self, **kwargs) + + monkeypatch.setattr(httpx.Client, "__init__", capture_init) + provider, _ = make_provider(issued()) + + get_context.assert_called_once_with("/private-ca.pem", "/server.pem") + assert captured["verify"] is get_context.return_value + assert captured["cert"] == ("/client.pem", "/client.key") + assert provider._client.timeout == httpx.Timeout(API_TIMEOUT) + + +def test_runtime_refresh_header_cannot_replace_external_identity(make_provider): + response = issued() + response.headers["Refreshed-API-Token"] = "runtime-token" + provider, requests = make_provider(response, issued("session-2")) + assert provider.get_token() == "session-1" + assert provider.get_token(force_refresh=True) == "session-2" + assert [request.headers["Authorization"] for request in requests] == ["Bearer identity-1"] * 2 + + +def test_caches_only_in_memory_and_renews_with_current_identity(make_provider, clock): + identity = mock.Mock(spec=lambda: "", side_effect=["identity-1\n", "identity-2"]) + provider, requests = make_provider(issued(), issued("session-2"), identity=identity) + + assert provider.get_token() == "session-1" + clock.return_value = 79 + assert provider.get_token() == "session-1" + clock.return_value = 80 + assert provider.get_token() == "session-2" + + assert identity.call_count == 2 + assert [r.headers["Authorization"] for r in requests] == ["Bearer identity-1", "Bearer identity-2"] + assert all(r.url.path == "/execution/dag-processor/session-token" for r in requests) + assert all(r.headers["Airflow-API-Version"] == bundle.version_values[0] for r in requests) + assert all(r.content == b"" for r in requests) + provider.close() + assert provider._token is None + assert provider._client.is_closed + + +def test_forced_refresh_and_timeout_are_forwarded(make_provider): + provider, requests = make_provider(issued(), issued("session-2")) + + assert provider.get_token() == "session-1" + assert provider.get_token(force_refresh=True, timeout=httpx.Timeout(0.5)) == "session-2" + assert requests[-1].extensions["timeout"] == httpx.Timeout(0.5).as_dict() + + +@pytest.mark.parametrize("failure", [httpx.Response(401), httpx.Response(503), httpx.ConnectError("offline")]) +def test_failed_renewal_does_not_return_cached_session(make_provider, clock, failure): + provider, requests = make_provider(issued(), failure, issued("session-2")) + assert provider.get_token() == "session-1" + clock.return_value = 100 + + with pytest.raises(httpx.HTTPError): + provider.get_token() + + assert provider.get_token() == "session-2" + assert len(requests) == 3 + + +@pytest.mark.parametrize("payload", [{"token": "", "expires_in": 10}, {"token": "x", "expires_in": 0}, {}]) +def test_invalid_exchange_response_is_not_cached(make_provider, payload): + provider, _ = make_provider(httpx.Response(200, json=payload)) + with pytest.raises(ValidationError): + provider.get_token() + assert provider._token is None + + +def test_response_lifetime_includes_request_duration(make_provider, clock): + def delayed(request): + clock.return_value = 100 + return issued(expires_in=100) + + provider, _ = make_provider(delayed) + with pytest.raises(ValueError, match="expired during the exchange"): + provider.get_token() + assert provider._token is None + + +def test_missing_identity_fails_before_request(make_provider): + provider, requests = make_provider(identity=mock.Mock(spec=lambda: "", return_value="\n")) + with pytest.raises(ValueError, match="credential is empty"): + provider.get_token() + assert not requests + + +def test_redirect_does_not_forward_identity(make_provider): + provider, requests = make_provider(httpx.Response(307, headers={"Location": "https://other/execution/"})) + with pytest.raises(httpx.HTTPStatusError): + provider.get_token() + assert len(requests) == 1 + + +def test_exchange_requires_https(): + with pytest.raises(ValueError, match="HTTPS"): + WorkloadIdentitySessionTokenProvider(base_url="http://api/execution/", identity_token=lambda: "x") + + +@pytest.mark.parametrize("sources", [{}, {"token_file": "unused", "token_provider": mock.sentinel.provider}]) +def test_client_requires_exactly_one_credential_source(sources): + with pytest.raises(ValueError, match="exactly one"): + DagProcessorAPIClient(base_url="https://api/execution/", hostname="processor", **sources) + + +@pytest.mark.parametrize("status", [401, 403]) +def test_client_forces_provider_refresh_after_registration_auth_failure(status): + provider = mock.Mock(spec=SessionTokenProvider) + provider.get_token.side_effect = ["session-1", "session-2"] + requests = [] + + def handle(request): + requests.append(request) + if len(requests) == 1: + return httpx.Response(status) + return httpx.Response(201, json={"job_id": 1, "token": "job-1", "expires_in": 300}) + + with DagProcessorAPIClient( + base_url="https://api/execution/", + hostname="processor", + token_provider=provider, + transport=httpx.MockTransport(handle), + ) as client: + assert client.register_job() == 1 + assert provider.get_token.call_args_list == [ + mock.call(force_refresh=False, retry=True, timeout=None), + mock.call(force_refresh=True, retry=True, timeout=None), + ] + assert requests[0].content == requests[1].content + + provider.close.assert_called_once_with() + assert client._session_token is None + assert client._registered_with is None + + +@pytest.mark.parametrize("retry", [False, True]) +@pytest.mark.parametrize("failure", [httpx.Response(503), httpx.ConnectError("offline")]) +def test_registration_retries_session_exchange_only_when_requested(make_provider, failure, retry): + provider, requests = make_provider(failure, issued()) + registrations = [] + + def register(request): + registrations.append(request) + return httpx.Response(201, json={"job_id": 1, "token": "job-1", "expires_in": 300}) + + with DagProcessorAPIClient( + base_url="https://api/execution/", + hostname="processor", + token_provider=provider, + transport=httpx.MockTransport(register), + ) as client: + if retry: + assert client.register_job(retry=True) == 1 + assert len(requests) == 2 + assert len(registrations) == 1 + else: + with pytest.raises(httpx.HTTPError): + client.register_job(retry=False) + assert len(requests) == 1 + assert not registrations + + +@pytest.mark.parametrize(("status", "attempts"), [(401, 1), (403, 1), (503, API_RETRIES)]) +def test_exchange_retries_are_bounded_and_do_not_retry_auth_failures(make_provider, status, attempts): + provider, requests = make_provider(*(httpx.Response(status) for _ in range(API_RETRIES))) + with pytest.raises(httpx.HTTPStatusError): + provider.get_token(retry=True) + assert len(requests) == attempts + assert provider._token is None + + +@mock.patch("airflow.dag_processing.api_client.monotonic", autospec=True) +def test_optional_renewal_attempts_exchange_once_then_uses_valid_job_token( + client_clock, make_provider, clock +): + client_clock.return_value = 0 + provider, exchanges = make_provider(issued(), httpx.Response(503)) + requests = [] + + def handle(request): + requests.append(request) + if request.url.path.endswith("/jobs"): + return httpx.Response(201, json={"job_id": 1, "token": "job-1", "expires_in": 300}) + return httpx.Response(200, json={"state": "running"}) + + with DagProcessorAPIClient( + base_url="https://api/execution/", + hostname="processor", + token_provider=provider, + transport=httpx.MockTransport(handle), + ) as client: + client.register_job() + client_clock.return_value = clock.return_value = 80 + client.heartbeat() + + assert len(exchanges) == 2 + assert exchanges[-1].extensions["timeout"] == httpx.Timeout(1).as_dict() + assert len(requests) == 2 + assert requests[-1].headers["Authorization"] == "Bearer job-1" + + +@pytest.mark.parametrize("use_context", [False, True]) +@pytest.mark.parametrize("close_fails", [False, True]) +def test_client_close_releases_credentials_and_connections(use_context, close_fails): + provider = mock.Mock(spec=SessionTokenProvider) + if close_fails: + provider.close.side_effect = OSError("cleanup failed") + client = DagProcessorAPIClient( + base_url="https://api/execution/", hostname="processor", token_provider=provider + ) + client._session_token = client._registered_with = "session" + client.auth = BearerAuth("job") + with pytest.raises(OSError, match="cleanup failed") if close_fails else nullcontext(): + if use_context: + with client: + pass + else: + client.close() + provider.close.assert_called_once_with() + assert client._session_token is None + assert client._registered_with is None + assert client.auth.token == "" + assert client.is_closed diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 6bef12629d956..1cd3c6c9e2bf1 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -136,6 +136,20 @@ class DagParseTokenResponse(BaseModel): ] +class DagProcessorSessionTokenResponse(BaseModel): + """ + A session credential issued after workload identity verification. + """ + + model_config = ConfigDict( + extra="forbid", + ) + token: Annotated[str, Field(min_length=1, title="Token")] + expires_in: Annotated[ + int, Field(description="Seconds until the session credential expires.", gt=0, title="Expires In") + ] + + class DagResponse(BaseModel): """ Schema for DAG response. From 5ab7948c58a6944c1e4eba93e9914b23bdaab463 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 22 Oct 2026 15:21:45 +0100 Subject: [PATCH 2/6] Validate Dag processor workload identity grants at API server startup A grant naming a Dag bundle the API server does not have can only be fixed by correcting the configuration and restarting. Reporting it as a retryable 503 on every exchange made each processor spend its retry budget on an error that could never clear, and the cause only surfaced on the processor side. Components that host the in-process Execution API, such as the Dag processor manager and the triggerer, can never serve the exchange. They still loaded its trust settings from shared configuration, so a component with a narrower bundle list than the API server would have failed to start. A token with an unknown signing key forces a JWKS refresh, and anyone can send one. While that refresh waited on a slow identity provider, it also held up exchanges whose keys were already cached. --- .../security/jwt_token_authentication.rst | 6 ++-- .../src/airflow/api_fastapi/auth/tokens.py | 3 ++ .../airflow/api_fastapi/execution_api/app.py | 11 +++++- .../api_fastapi/execution_api/in_process.py | 4 +++ .../execution_api/routes/session_tokens.py | 10 +----- .../execution_api/workload_identity.py | 6 ++++ .../src/airflow/config_templates/config.yml | 3 +- .../unit/api_fastapi/auth/test_tokens.py | 34 +++++++++++++++++++ .../api_fastapi/execution_api/test_app.py | 23 +++++++++++++ .../test_token_scope_boundaries.py | 6 +--- .../versions/head/test_session_tokens.py | 20 +++++++++-- 11 files changed, 105 insertions(+), 21 deletions(-) diff --git a/airflow-core/docs/security/jwt_token_authentication.rst b/airflow-core/docs/security/jwt_token_authentication.rst index 8a2c620389962..f560113733165 100644 --- a/airflow-core/docs/security/jwt_token_authentication.rst +++ b/airflow-core/docs/security/jwt_token_authentication.rst @@ -545,8 +545,8 @@ Use your identity provider's exact issuer, HTTPS JWKS URL, and signing algorithm ``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, and every granted bundle must be configured on the API server. This option is -empty by default, disabling the exchange. +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. @@ -621,7 +621,7 @@ Session expiry is capped by both the external identity's remaining lifetime and with less than one second remaining cannot obtain a session. Renewal rereads the platform identity and retains the same session and Job. During an identity-provider or API outage, the client can continue using an already-issued Job credential until it expires; it cannot renew without a valid identity. -Removing a grant stops future exchanges after API-server configuration is reloaded. Credentials +Removing a grant stops future exchanges once the API servers restart with the new configuration. Credentials already issued remain subject to their expiry and Job ownership checks. Verification uses the issuer's JWKS and does not perform online token introspection. In particular, diff --git a/airflow-core/src/airflow/api_fastapi/auth/tokens.py b/airflow-core/src/airflow/api_fastapi/auth/tokens.py index df759a6dd7b7d..e4071ce10e0e7 100644 --- a/airflow-core/src/airflow/api_fastapi/auth/tokens.py +++ b/airflow-core/src/airflow/api_fastapi/auth/tokens.py @@ -146,6 +146,9 @@ def from_private_key(cls, *keys: AllowedPrivateKeys | tuple[AllowedPrivateKeys, return obj 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 lock = self._fetch_lock if self.refresh_on_unknown_kid else nullcontext() async with lock: if not self._should_fetch_jwks(force=force): diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index a5bbec41f2fef..cb49bd0ea8058 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -43,7 +43,10 @@ from airflow.api_fastapi.common.telemetry import NATIVE_TELEMETRY_DISABLED from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI from airflow.api_fastapi.execution_api.security import _REQUEST_SCOPE_TOKEN_KEY, require_auth -from airflow.api_fastapi.execution_api.workload_identity import create_dag_processor_identity_verifier +from airflow.api_fastapi.execution_api.workload_identity import ( + DagProcessorIdentityVerifier, + create_dag_processor_identity_verifier, +) if TYPE_CHECKING: from airflow.models.dagbag import DBDagBag @@ -108,6 +111,12 @@ 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) + # InProcessExecutionAPI stubs out the identity verifier: don't create one in that case. + if DagProcessorIdentityVerifier in registry: + app.state.dag_processor_identity_verifier = None + yield + return + async with create_dag_processor_identity_verifier() as verifier: app.state.dag_processor_identity_verifier = verifier yield diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py index 3807e3f7cc6bc..a9afd2915eaea 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py @@ -106,6 +106,7 @@ def app(self): from airflow.api_fastapi.auth.tokens import JWTValidator from airflow.api_fastapi.execution_api.app import create_task_execution_api_app, lifespan + from airflow.api_fastapi.execution_api.workload_identity import DagProcessorIdentityVerifier if not self._app: # Keep the stub private: a shared None validator would prevent lifespan() from @@ -114,6 +115,9 @@ def app(self): self._app = create_task_execution_api_app(lifespan=attrs.evolve(lifespan, registry=registry)) # Auth dependency overrides bypass validation, so the stub skips validator setup in lifespan(). registry.register_value(JWTValidator, None) + # No external Dag processor can reach this app, so it must not load trust settings meant for + # API servers, whose Dag bundle checks can fail in other components. + registry.register_value(DagProcessorIdentityVerifier, None) _configure_in_process_auth(self._app, triggerer_id=self.triggerer_id) return self._app diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py index 936935bea8dc6..f5e098cd612c7 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py @@ -42,7 +42,6 @@ DagProcessorIdentityVerifier, WorkloadIdentityNotAllowed, ) -from airflow.dag_processing.bundles.manager import get_configured_bundle_names router = VersionedAPIRouter() identity_bearer = HTTPBearer(auto_error=False, scheme_name="DagProcessorWorkloadIdentity") @@ -85,10 +84,7 @@ async def require_workload_identity( (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, - "Signing keys or the configured bundle grant are unavailable", - ), + (status.HTTP_503_SERVICE_UNAVAILABLE, "The identity provider's signing keys are unavailable"), ] ), ) @@ -97,10 +93,6 @@ def exchange_session_token( identity: Annotated[DagProcessorIdentity, Depends(require_workload_identity)], services: svcs.Container = DepContainer, ) -> DagProcessorSessionTokenResponse: - if not identity.bundle_names <= get_configured_bundle_names(): - raise HTTPException( - status.HTTP_503_SERVICE_UNAVAILABLE, "Workload identity bundle grant is unavailable" - ) generator: JWTGenerator = services.get(JWTGenerator) try: lifetime = get_token_lifetime(generator, identity.expires_at, "Workload identity has expired") diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py b/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py index bfb72bffa4399..138b236bd0175 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py @@ -33,6 +33,7 @@ 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)] @@ -83,6 +84,11 @@ class DagProcessorIdentityVerifier: 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 diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 21f87cb91a5d3..6ca76b40c4a9d 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2431,7 +2431,8 @@ execution_api: Requires issuer, HTTPS jwks_url, a dedicated audience, an instance_id_claim path (a list of keys into the signed identity claims), and grants mapping exact subjects - to allowed Dag bundle names. The optional algorithm defaults to RS256; ES256 and + to allowed Dag bundle names. The API server refuses to start if a grant names a Dag + bundle it does not have configured. The optional algorithm defaults to RS256; ES256 and EdDSA are also supported. Use a signed, per-replica instance claim: a service account shared by replicas is not sufficient on its own. The JWKS endpoint must accept unauthenticated HTTPS requests and have a certificate diff --git a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py index f6693743c1fa6..2d2ca1480e707 100644 --- a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py +++ b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py @@ -161,6 +161,40 @@ def handle(request): assert (await jwks.get_key("new")).key_id == "new" assert len(requests) == 3 + @patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) + async def test_slow_unknown_key_refresh_does_not_delay_cached_keys(self, clock, ed25519_private_key): + fetching = asyncio.Event() + release = asyncio.Event() + requests = [] + + async def handle(request): + requests.append(request) + if len(requests) > 1: + fetching.set() + await release.wait() + return httpx.Response(200, json={"keys": [key_to_jwk_dict(ed25519_private_key, "known")]}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + jwks = JWKS(url="https://example.com/keys", client=client, refresh_on_unknown_kid=True) + known = await jwks.get_key("known") + clock.return_value = 111 + refresh = asyncio.create_task(jwks.get_key("unknown")) + # The patched clock also freezes the event loop's, so check progress instead of using timeouts. + fetch_started = asyncio.create_task(fetching.wait()) + await asyncio.wait({refresh, fetch_started}, return_when=asyncio.FIRST_COMPLETED) + lookup = asyncio.create_task(jwks.get_key("known")) + await asyncio.sleep(0) + try: + assert fetching.is_set() + assert lookup.done() + assert lookup.result() is known + finally: + release.set() + fetch_started.cancel() + with pytest.raises(KeyError): + await refresh + assert len(requests) == 2 + @pytest.mark.parametrize("private_key", ["rsa_private_key", "ed25519_private_key"], indirect=True) async def test_fetch_jwks_success(self, private_key): jwk_content = json.dumps({"keys": [key_to_jwk_dict(private_key, "kid")]}) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py index 5f708651294f3..40393ac59fc6a 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py @@ -18,6 +18,7 @@ import asyncio import gc +import json import threading from contextlib import asynccontextmanager from unittest import mock @@ -216,6 +217,28 @@ def test_in_process_execution_api_runs_without_jwt_secret(): assert response.status_code == 200 +@conf_vars( + { + ("execution_api", "dag_processor_workload_identity"): json.dumps( + { + "issuer": "https://identity.example/", + "jwks_url": "https://identity.example/keys", + "audience": "airflow-dag-processor-bootstrap", + "instance_id_claim": ["sub"], + "grants": {"processor": ["bundle-of-another-component"]}, + } + ) + } +) +def test_in_process_execution_api_ignores_workload_identity_trust(): + """Trust settings shared with API servers must not stop components that host the in-process API.""" + api = InProcessExecutionAPI() + with httpx.Client(transport=api.transport) as client: + response = client.post("http://localhost/dag-processor/session-token") + assert response.status_code == 404 + assert response.json() == {"detail": "Workload identity provisioning is not enabled"} + + def test_in_process_task_xcom_write_uses_its_attempt_id(create_task_instance, session): ti = create_task_instance() session.commit() diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py index 1e7e00cedb27d..3de9c35f1b91d 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py @@ -56,8 +56,6 @@ # Routes that intentionally deviate from the default (execution-only) policy. # Any route NOT listed here must accept only {"execution"}. NON_DEFAULT_TOKEN_POLICY: dict[str, set[str]] = { - # Bootstrap authenticates a deployment identity, never an Airflow runtime token. - "POST /dag-processor/session-token": set(), # The /run endpoint exchanges a workload token for a short-lived execution token. "PATCH /task-instances/{task_instance_id}/run": {"execution", "workload"}, # Connection test routes run from a queued worker context (workload-only). @@ -152,9 +150,7 @@ def _get_path_and_query_param_names(dependant: Dependant) -> set[str]: def _all_route_policies() -> dict[str, set[str]]: """Return a map of all API routes and their allowed token types.""" return { - key: set() - if key == "POST /dag-processor/session-token" - else set(getattr(_get_original_route(route), "allowed_token_types", {"execution"})) + key: set(getattr(_get_original_route(route), "allowed_token_types", {"execution"})) for key, route in _get_api_routes().items() } diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py index 68b57735dee61..39fb687293abd 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py @@ -196,6 +196,20 @@ def test_issues_only_configured_grants_and_caps_lifetime( 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 ): @@ -419,8 +433,10 @@ def test_container_restart_waits_for_previous_job_to_become_stale(make_processor @conf_vars({("dag_processor", "dag_bundle_config_list"): "[]"}) -def test_unknown_bundle_grant_does_not_issue_credentials(client, verifier, identity_token): - assert exchange(client, identity_token()).status_code == 503 +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): From 07c1f61d3d581db896ca3a917962087aad9de8c4 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 22 Oct 2026 15:40:52 +0100 Subject: [PATCH 3/6] Use httpx2 for the Dag processor session token client The Task SDK client moved to httpx2 while this change was in review, and its responses, timeouts and errors do not interoperate with httpx objects. The session token provider is built on that client, so it has to use httpx2 types for its timeouts and error handling to keep working. The API server's JWKS fetcher still uses httpx, as the rest of its token validation does. --- .../src/airflow/dag_processing/api_client.py | 2 +- .../src/airflow/dag_processing/credentials.py | 12 ++-- .../api_fastapi/execution_api/test_app.py | 2 +- .../versions/head/test_session_tokens.py | 15 ++--- .../unit/dag_processing/test_credentials.py | 56 ++++++++++--------- 5 files changed, 45 insertions(+), 42 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/api_client.py b/airflow-core/src/airflow/dag_processing/api_client.py index 78e06af366603..39f214913d721 100644 --- a/airflow-core/src/airflow/dag_processing/api_client.py +++ b/airflow-core/src/airflow/dag_processing/api_client.py @@ -227,7 +227,7 @@ def _read_session_token( force: bool = False, refresh: bool = False, retry: bool = False, - timeout: httpx.Timeout | None = None, + timeout: httpx2.Timeout | None = None, ) -> str: now = monotonic() if self._session_token is None or force or now >= self._reload_at: diff --git a/airflow-core/src/airflow/dag_processing/credentials.py b/airflow-core/src/airflow/dag_processing/credentials.py index 69527f1287ba7..20b7cf9abcf92 100644 --- a/airflow-core/src/airflow/dag_processing/credentials.py +++ b/airflow-core/src/airflow/dag_processing/credentials.py @@ -24,7 +24,7 @@ from time import monotonic from typing import Protocol -import httpx +import httpx2 from airflow.api_fastapi.execution_api.datamodels.job import DagProcessorSessionTokenResponse from airflow.api_fastapi.execution_api.versions import bundle @@ -35,7 +35,7 @@ class SessionTokenProvider(Protocol): """Provide credentials without prescribing where they are obtained.""" def get_token( - self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx.Timeout | None = None + self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx2.Timeout | None = None ) -> str: ... def close(self) -> None: ... @@ -48,7 +48,7 @@ def __init__(self, path: str | Path): self.path = Path(path) def get_token( - self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx.Timeout | None = None + self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx2.Timeout | None = None ) -> str: token = self.path.read_text().strip() if not token: @@ -60,7 +60,7 @@ def close(self) -> None: class _SessionTokenClient(Client): - def _update_auth(self, response: httpx.Response) -> None: + def _update_auth(self, response: httpx2.Response) -> None: # Runtime refresh headers must not replace the external workload identity. pass @@ -81,7 +81,7 @@ def __init__( identity_token: Callable[[], str], **kwargs, ): - if httpx.URL(base_url).scheme != "https": + 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 @@ -89,7 +89,7 @@ def __init__( self._client = _SessionTokenClient(base_url=base_url, token="", follow_redirects=False, **kwargs) def get_token( - self, *, force_refresh: bool = False, retry: bool = False, timeout: httpx.Timeout | None = None + 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 diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py index 40393ac59fc6a..175e9b8993db0 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py @@ -233,7 +233,7 @@ def test_in_process_execution_api_runs_without_jwt_secret(): def test_in_process_execution_api_ignores_workload_identity_trust(): """Trust settings shared with API servers must not stop components that host the in-process API.""" api = InProcessExecutionAPI() - with httpx.Client(transport=api.transport) as client: + with httpx2.Client(transport=api.transport) as client: response = client.post("http://localhost/dag-processor/session-token") assert response.status_code == 404 assert response.json() == {"detail": "Workload identity provisioning is not enabled"} diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py index 39fb687293abd..896d6a3c3e13f 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py @@ -21,6 +21,7 @@ from unittest import mock import httpx +import httpx2 import jwt import pytest from cryptography.hazmat.primitives.asymmetric import ec, rsa @@ -144,10 +145,10 @@ def forward(request): response = client.request( request.method, request.url.path, headers=request.headers, content=request.content ) - return httpx.Response(response.status_code, headers=response.headers, content=response.content) + return httpx2.Response(response.status_code, headers=response.headers, content=response.content) def create(): - transport = httpx.MockTransport(forward) + transport = httpx2.MockTransport(forward) provider = WorkloadIdentitySessionTokenProvider( base_url="https://api/execution/", identity_token=identity_token, transport=transport ) @@ -498,16 +499,16 @@ def test_session_renewal_and_job_lifecycle_without_session_file( def forward(request): if request.url.path == URL and not exchange_available: - raise httpx.ConnectError("session exchange unavailable", request=request) + raise httpx2.ConnectError("session exchange unavailable", request=request) response = client.request( request.method, request.url.path, headers=request.headers, content=request.content ) - return httpx.Response(response.status_code, headers=response.headers, content=response.content) + return httpx2.Response(response.status_code, headers=response.headers, content=response.content) provider = WorkloadIdentitySessionTokenProvider( base_url="https://api/execution/", identity_token=lambda: identity_token(exp=NOW + 600), - transport=httpx.MockTransport(forward), + transport=httpx2.MockTransport(forward), ) try: Variable.set("identity-test-key", "value") @@ -515,7 +516,7 @@ def forward(request): base_url="https://api/execution/", hostname="processor", token_provider=provider, - transport=httpx.MockTransport(forward), + transport=httpx2.MockTransport(forward), ) as processor: job_id = processor.register_job() first_session = provider.get_token() @@ -541,7 +542,7 @@ def forward(request): assert processor.heartbeat() == JobState.RUNNING time_machine.shift(61) credential_clock.return_value = job_clock.return_value = 541 - with pytest.raises(httpx.ConnectError): + with pytest.raises(httpx2.ConnectError): processor.heartbeat() exchange_available = True diff --git a/airflow-core/tests/unit/dag_processing/test_credentials.py b/airflow-core/tests/unit/dag_processing/test_credentials.py index e1788bf46437a..dbc973589d7d2 100644 --- a/airflow-core/tests/unit/dag_processing/test_credentials.py +++ b/airflow-core/tests/unit/dag_processing/test_credentials.py @@ -19,7 +19,7 @@ from contextlib import nullcontext from unittest import mock -import httpx +import httpx2 import pytest from pydantic import ValidationError from tenacity import wait_none @@ -59,7 +59,7 @@ def handle(request): provider = WorkloadIdentitySessionTokenProvider( base_url="https://api/execution/", identity_token=identity or mock.Mock(spec=lambda: "", return_value="identity-1\n"), - transport=httpx.MockTransport(handle), + transport=httpx2.MockTransport(handle), ) providers.append(provider) return provider, requests @@ -70,7 +70,7 @@ def handle(request): def issued(token="session-1", expires_in=100): - return httpx.Response(200, json={"token": token, "expires_in": expires_in}) + return httpx2.Response(200, json={"token": token, "expires_in": expires_in}) @mock.patch("airflow.sdk.api.client.API_SSL_CA_FILE_PATH", "/private-ca.pem") @@ -82,19 +82,19 @@ def issued(token="session-1", expires_in=100): ) def test_exchange_uses_api_tls_settings(get_context, make_provider, monkeypatch): captured = {} - original_init = httpx.Client.__init__ + original_init = httpx2.Client.__init__ def capture_init(self, **kwargs): captured.update(kwargs) original_init(self, **kwargs) - monkeypatch.setattr(httpx.Client, "__init__", capture_init) + monkeypatch.setattr(httpx2.Client, "__init__", capture_init) provider, _ = make_provider(issued()) get_context.assert_called_once_with("/private-ca.pem", "/server.pem") assert captured["verify"] is get_context.return_value assert captured["cert"] == ("/client.pem", "/client.key") - assert provider._client.timeout == httpx.Timeout(API_TIMEOUT) + assert provider._client.timeout == httpx2.Timeout(API_TIMEOUT) def test_runtime_refresh_header_cannot_replace_external_identity(make_provider): @@ -130,17 +130,19 @@ def test_forced_refresh_and_timeout_are_forwarded(make_provider): provider, requests = make_provider(issued(), issued("session-2")) assert provider.get_token() == "session-1" - assert provider.get_token(force_refresh=True, timeout=httpx.Timeout(0.5)) == "session-2" - assert requests[-1].extensions["timeout"] == httpx.Timeout(0.5).as_dict() + assert provider.get_token(force_refresh=True, timeout=httpx2.Timeout(0.5)) == "session-2" + assert requests[-1].extensions["timeout"] == httpx2.Timeout(0.5).as_dict() -@pytest.mark.parametrize("failure", [httpx.Response(401), httpx.Response(503), httpx.ConnectError("offline")]) +@pytest.mark.parametrize( + "failure", [httpx2.Response(401), httpx2.Response(503), httpx2.ConnectError("offline")] +) def test_failed_renewal_does_not_return_cached_session(make_provider, clock, failure): provider, requests = make_provider(issued(), failure, issued("session-2")) assert provider.get_token() == "session-1" clock.return_value = 100 - with pytest.raises(httpx.HTTPError): + with pytest.raises(httpx2.HTTPError): provider.get_token() assert provider.get_token() == "session-2" @@ -149,7 +151,7 @@ def test_failed_renewal_does_not_return_cached_session(make_provider, clock, fai @pytest.mark.parametrize("payload", [{"token": "", "expires_in": 10}, {"token": "x", "expires_in": 0}, {}]) def test_invalid_exchange_response_is_not_cached(make_provider, payload): - provider, _ = make_provider(httpx.Response(200, json=payload)) + provider, _ = make_provider(httpx2.Response(200, json=payload)) with pytest.raises(ValidationError): provider.get_token() assert provider._token is None @@ -174,8 +176,8 @@ def test_missing_identity_fails_before_request(make_provider): def test_redirect_does_not_forward_identity(make_provider): - provider, requests = make_provider(httpx.Response(307, headers={"Location": "https://other/execution/"})) - with pytest.raises(httpx.HTTPStatusError): + provider, requests = make_provider(httpx2.Response(307, headers={"Location": "https://other/execution/"})) + with pytest.raises(httpx2.HTTPStatusError): provider.get_token() assert len(requests) == 1 @@ -200,14 +202,14 @@ def test_client_forces_provider_refresh_after_registration_auth_failure(status): def handle(request): requests.append(request) if len(requests) == 1: - return httpx.Response(status) - return httpx.Response(201, json={"job_id": 1, "token": "job-1", "expires_in": 300}) + return httpx2.Response(status) + return httpx2.Response(201, json={"job_id": 1, "token": "job-1", "expires_in": 300}) with DagProcessorAPIClient( base_url="https://api/execution/", hostname="processor", token_provider=provider, - transport=httpx.MockTransport(handle), + transport=httpx2.MockTransport(handle), ) as client: assert client.register_job() == 1 assert provider.get_token.call_args_list == [ @@ -222,27 +224,27 @@ def handle(request): @pytest.mark.parametrize("retry", [False, True]) -@pytest.mark.parametrize("failure", [httpx.Response(503), httpx.ConnectError("offline")]) +@pytest.mark.parametrize("failure", [httpx2.Response(503), httpx2.ConnectError("offline")]) def test_registration_retries_session_exchange_only_when_requested(make_provider, failure, retry): provider, requests = make_provider(failure, issued()) registrations = [] def register(request): registrations.append(request) - return httpx.Response(201, json={"job_id": 1, "token": "job-1", "expires_in": 300}) + return httpx2.Response(201, json={"job_id": 1, "token": "job-1", "expires_in": 300}) with DagProcessorAPIClient( base_url="https://api/execution/", hostname="processor", token_provider=provider, - transport=httpx.MockTransport(register), + transport=httpx2.MockTransport(register), ) as client: if retry: assert client.register_job(retry=True) == 1 assert len(requests) == 2 assert len(registrations) == 1 else: - with pytest.raises(httpx.HTTPError): + with pytest.raises(httpx2.HTTPError): client.register_job(retry=False) assert len(requests) == 1 assert not registrations @@ -250,8 +252,8 @@ def register(request): @pytest.mark.parametrize(("status", "attempts"), [(401, 1), (403, 1), (503, API_RETRIES)]) def test_exchange_retries_are_bounded_and_do_not_retry_auth_failures(make_provider, status, attempts): - provider, requests = make_provider(*(httpx.Response(status) for _ in range(API_RETRIES))) - with pytest.raises(httpx.HTTPStatusError): + provider, requests = make_provider(*(httpx2.Response(status) for _ in range(API_RETRIES))) + with pytest.raises(httpx2.HTTPStatusError): provider.get_token(retry=True) assert len(requests) == attempts assert provider._token is None @@ -262,27 +264,27 @@ def test_optional_renewal_attempts_exchange_once_then_uses_valid_job_token( client_clock, make_provider, clock ): client_clock.return_value = 0 - provider, exchanges = make_provider(issued(), httpx.Response(503)) + provider, exchanges = make_provider(issued(), httpx2.Response(503)) requests = [] def handle(request): requests.append(request) if request.url.path.endswith("/jobs"): - return httpx.Response(201, json={"job_id": 1, "token": "job-1", "expires_in": 300}) - return httpx.Response(200, json={"state": "running"}) + return httpx2.Response(201, json={"job_id": 1, "token": "job-1", "expires_in": 300}) + return httpx2.Response(200, json={"state": "running"}) with DagProcessorAPIClient( base_url="https://api/execution/", hostname="processor", token_provider=provider, - transport=httpx.MockTransport(handle), + transport=httpx2.MockTransport(handle), ) as client: client.register_job() client_clock.return_value = clock.return_value = 80 client.heartbeat() assert len(exchanges) == 2 - assert exchanges[-1].extensions["timeout"] == httpx.Timeout(1).as_dict() + assert exchanges[-1].extensions["timeout"] == httpx2.Timeout(1).as_dict() assert len(requests) == 2 assert requests[-1].headers["Authorization"] == "Bearer job-1" From 181bf3249468ec120e34f0dc49aae8ce275e684a Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 22 Oct 2026 15:46:56 +0100 Subject: [PATCH 4/6] Drop the token reissue opt-out from the Dag processor session exchange The opt-out stopped the reissue middleware from re-validating the external workload identity as if it were an Airflow token. The middleware now only renews tokens that Airflow's own JWT check has already validated, and the exchange never runs that check, so the opt-out no longer prevents anything. --- airflow-core/src/airflow/api_fastapi/execution_api/app.py | 3 --- .../airflow/api_fastapi/execution_api/routes/session_tokens.py | 1 - 2 files changed, 4 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index cb49bd0ea8058..64f94f0229ca0 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -153,9 +153,6 @@ class JWTReissueMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): response: Response = await call_next(request) - if getattr(request.state, "skip_execution_token_reissue", False): - return response - refreshed_token: str | None = None token = request.scope.get(_REQUEST_SCOPE_TOKEN_KEY) if token: diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py index f5e098cd612c7..643b3d4a41239 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py @@ -51,7 +51,6 @@ async def require_workload_identity( request: Request, credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(identity_bearer)], ) -> DagProcessorIdentity: - request.state.skip_execution_token_reissue = True 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") From 794455e86c1f130405bf2279e100d2a3fe1a7bae Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Fri, 23 Oct 2026 16:09:18 +0100 Subject: [PATCH 5/6] Keep Dag processor identity exchange failures recoverable and contained After a failed cold-cache JWKS fetch the API server waits 10 seconds before fetching again, and the SDK's default backoff retried every registration attempt inside that window more than half the time. A brief identity provider outage therefore failed registration, although the 503 told the client to retry. An identity supplier that raises its own exception type, or a JWKS endpoint that returns something other than a JSON object, turned a recoverable condition into failed requests, even for keys that were already cached. A grant naming an unconfigured Dag bundle stopped the whole API server, taking down the UI and the API that workers heartbeat against, when only the subject holding that grant needs to be refused. Catching type errors at the route also reported verifier bugs as invalid identities, when only older PyJWT releases raise them for malformed time claims. --- .../security/jwt_token_authentication.rst | 36 +-- .../src/airflow/api_fastapi/auth/tokens.py | 3 + .../airflow/api_fastapi/execution_api/app.py | 4 +- .../api_fastapi/execution_api/in_process.py | 2 +- .../execution_api/routes/session_tokens.py | 28 ++- .../execution_api/workload_identity.py | 36 ++- .../src/airflow/config_templates/config.yml | 12 +- .../src/airflow/dag_processing/credentials.py | 52 ++++- .../unit/api_fastapi/auth/test_tokens.py | 24 ++ .../test_token_scope_boundaries.py | 30 ++- .../versions/head/test_session_tokens.py | 209 ++++++++++-------- .../unit/dag_processing/test_credentials.py | 73 +++++- 12 files changed, 363 insertions(+), 146 deletions(-) diff --git a/airflow-core/docs/security/jwt_token_authentication.rst b/airflow-core/docs/security/jwt_token_authentication.rst index f560113733165..53a4d9c4ae4a7 100644 --- a/airflow-core/docs/security/jwt_token_authentication.rst +++ b/airflow-core/docs/security/jwt_token_authentication.rst @@ -520,7 +520,7 @@ not yet use this client. Provisioning a token file alone does not switch that co 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``. @@ -545,14 +545,20 @@ Use your identity provider's exact issuer, HTTPS JWKS URL, and signing algorithm ``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. +obtain a session. If a grant names a Dag bundle that the API server does not have configured, the +API server logs an error at startup and refuses that subject's exchanges with HTTP 403. 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 option trusts a single issuer. Dag processors whose identities come from different issuers, such +as separate Kubernetes clusters, cannot all obtain sessions from the same API server. + 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 -other HTTPX clients in the API-server process; there is no per-issuer CA or authentication option. +The fetcher sends no authentication credentials and uses HTTPX's default certificate trust. +``SSL_CERT_FILE`` or ``SSL_CERT_DIR``, when set, replace the default CA bundle rather than adding +to it, and they apply to every HTTPX client in the API-server process. A bundle that holds only a +private CA cuts those clients off from public roots, so include the roots they still need. There is +no per-issuer CA or authentication option. Authenticated JWKS endpoints, including Kubernetes API-server discovery protected by service-account authentication, are not supported by this implementation. For those clusters, serve the public JWKS through a suitable HTTPS endpoint before enabling the exchange. See Kubernetes' @@ -586,7 +592,8 @@ The client does not wait automatically. This prevents two processes sharing an i active Jobs at the same time. The processor integration supplies a callable that reads the current platform identity on each -exchange. For example: +exchange. Any exception the callable raises is re-raised as ``WorkloadIdentityUnavailableError``, an +``OSError``, so a failed optional renewal leaves the valid Job credential in use. For example: .. code-block:: python @@ -613,8 +620,9 @@ session JWT to disk and does not fall back to a session-token file. ``token_prov ``token_file`` are mutually exclusive. HTTPS is required for the exchange, and redirects are rejected. The exchange uses the same API TLS configuration and request timeout as the Task SDK client, including private certificate authorities and client certificates. Registration retries transient -exchange failures using the configured Execution API retry policy; optional early renewal makes -one bounded attempt so it does not delay requests using a valid Job credential. +exchange failures using the configured Execution API retry policy, waiting at least the +``Retry-After`` interval of an HTTP 503 before the next attempt. Optional early renewal makes one +bounded attempt so it does not delay requests using a valid Job credential. Session expiry is capped by both the external identity's remaining lifetime and ``[execution_api] jwt_expiration_time``. Clock-skew leeway never extends this cap, and an identity @@ -631,12 +639,12 @@ unknown signing key. Fetches, including failed attempts, are limited to one per API-server process. Publish new keys before signing with them and retain old keys while their tokens remain valid. -If a required signing key cannot be fetched, the exchange returns HTTP 503 with ``Retry-After``, -allowing registration to retry. A cached matching key remains usable during an outage. Invalid -identities and keys absent from a successfully fetched keyset return HTTP 401. +If a required signing key cannot be fetched, the exchange returns HTTP 503 with ``Retry-After``, and +registration waits at least that long before retrying. A cached matching key remains usable during +an outage. Invalid identities and keys absent from a successfully fetched key set return HTTP 401. File-based provisioning -~~~~~~~~~~~~~~~~~~~~~~~ +^^^^^^^^^^^^^^^^^^^^^^^ On a trusted host with the API signing key, provision a session token for the bundles the client may access:: @@ -652,7 +660,7 @@ provisioner with the same token file and bundles keeps that session, so running their Jobs. Credential permissions and renewal -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ The three credentials have different issuers and permissions: diff --git a/airflow-core/src/airflow/api_fastapi/auth/tokens.py b/airflow-core/src/airflow/api_fastapi/auth/tokens.py index e4071ce10e0e7..688706a3b6257 100644 --- a/airflow-core/src/airflow/api_fastapi/auth/tokens.py +++ b/airflow-core/src/airflow/api_fastapi/auth/tokens.py @@ -163,6 +163,9 @@ async def fetch_jwks(self, *, force: bool = False) -> None: return try: + # PyJWKSet.from_dict raises AttributeError for a JSON array or scalar. + if not isinstance(data, dict): + raise TypeError(f"JWKS must be a JSON object, not {type(data).__name__}") keyset = jwt.PyJWKSet.from_dict(data) except (jwt.PyJWKError, jwt.PyJWKSetError, KeyError, TypeError, ValueError): if not self.raise_on_fetch_error: diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index 64f94f0229ca0..078b68e3196dd 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -113,12 +113,12 @@ async def lifespan(app: FastAPI, registry: svcs.Registry): # InProcessExecutionAPI stubs out the identity verifier: don't create one in that case. if DagProcessorIdentityVerifier in registry: - app.state.dag_processor_identity_verifier = None yield return async with create_dag_processor_identity_verifier() as verifier: - app.state.dag_processor_identity_verifier = verifier + # None when the exchange is disabled; the route answers 404 then. + registry.register_value(DagProcessorIdentityVerifier, verifier) yield diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py index a9afd2915eaea..6dc56b46939a8 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py @@ -116,7 +116,7 @@ def app(self): # Auth dependency overrides bypass validation, so the stub skips validator setup in lifespan(). registry.register_value(JWTValidator, None) # No external Dag processor can reach this app, so it must not load trust settings meant for - # API servers, whose Dag bundle checks can fail in other components. + # API servers, which can fail to load in components configured differently. registry.register_value(DagProcessorIdentityVerifier, None) _configure_in_process_auth(self._app, triggerer_id=self.triggerer_id) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py index 643b3d4a41239..e61174e4a5468 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py @@ -24,7 +24,7 @@ import jwt import svcs from cadwyn import VersionedAPIRouter -from fastapi import Depends, HTTPException, Request, Response, status +from fastapi import Depends, HTTPException, Response, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from pydantic import ValidationError @@ -48,10 +48,10 @@ async def require_workload_identity( - request: Request, credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(identity_bearer)], + services: svcs.Container = DepContainer, ) -> DagProcessorIdentity: - verifier: DagProcessorIdentityVerifier | None = request.app.state.dag_processor_identity_verifier + verifier: DagProcessorIdentityVerifier | None = await services.aget(DagProcessorIdentityVerifier) if verifier is None: raise HTTPException(status.HTTP_404_NOT_FOUND, "Workload identity provisioning is not enabled") if credentials is None: @@ -66,14 +66,12 @@ async def require_workload_identity( "Workload identity signing keys are unavailable", headers={"Retry-After": str(verifier.jwks.refresh_retry_interval_secs)}, ) from None - except (jwt.PyJWTError, ValidationError, TypeError, OverflowError): + except (jwt.PyJWTError, ValidationError): 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 + except WorkloadIdentityNotAllowed as error: + raise HTTPException(status.HTTP_403_FORBIDDEN, str(error)) from None @router.post( @@ -81,7 +79,7 @@ async def require_workload_identity( 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_403_FORBIDDEN, "The identity has no usable Dag bundle grant"), (status.HTTP_404_NOT_FOUND, "Workload identity provisioning is disabled"), (status.HTTP_503_SERVICE_UNAVAILABLE, "The identity provider's signing keys are unavailable"), ] @@ -92,11 +90,21 @@ def exchange_session_token( identity: Annotated[DagProcessorIdentity, Depends(require_workload_identity)], services: svcs.Container = DepContainer, ) -> DagProcessorSessionTokenResponse: + """ + Exchange an external workload identity for a Dag processor session credential. + + The bearer token is an OIDC identity issued by the deployment platform, not an Airflow token. The + API server picks the Dag bundle grants from its own configuration, and the session cannot outlive + the identity. While the identity provider's signing keys cannot be fetched, the exchange returns + 503 with ``Retry-After``. + """ 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 + raise HTTPException( + status.HTTP_401_UNAUTHORIZED, str(error), headers={"WWW-Authenticate": "Bearer"} + ) from None token = generate_dag_processor_session_token( generator, session_id=identity.session_id, bundle_names=identity.bundle_names, valid_for=lifetime ) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py b/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py index 138b236bd0175..b5eab598c8dc2 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py @@ -23,11 +23,12 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Annotated, Any, Literal +from typing import Annotated, Literal from uuid import NAMESPACE_URL, UUID, uuid5 import httpx import jwt +import structlog from pydantic import BaseModel, ConfigDict, Field, HttpUrl, field_validator from airflow.api_fastapi.auth.tokens import JWKS, JWTValidator @@ -35,6 +36,8 @@ from airflow.configuration import conf from airflow.dag_processing.bundles.manager import get_configured_bundle_names +log = structlog.get_logger(logger_name=__name__) + NonemptyString = Annotated[str, Field(min_length=1)] @@ -84,10 +87,20 @@ class DagProcessorIdentityVerifier: 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( + configured = get_configured_bundle_names() + # Not fatal: a renamed bundle must not stop the whole API server, only the subjects granted it. + self._unconfigured_grants = { + subject: missing + for subject, bundles in config.grants.items() + if (missing := bundles - configured) + } + if self._unconfigured_grants: + log.error( "[execution_api] dag_processor_workload_identity grants Dag bundles missing from " - f"[dag_processor] dag_bundle_config_list: {sorted(unknown)}" + "[dag_processor] dag_bundle_config_list; these subjects cannot obtain a session", + unconfigured_grants={ + subject: sorted(missing) for subject, missing in self._unconfigured_grants.items() + }, ) self.config = config self.jwks = JWKS( @@ -102,11 +115,20 @@ def __init__(self, config: WorkloadIdentityConfig, *, client: httpx.AsyncClient) ) async def authenticate(self, token: str) -> DagProcessorIdentity: - claims = await self.validator.avalidated_claims(token) + try: + claims = await self.validator.avalidated_claims(token) + except (TypeError, OverflowError) as error: + # PyJWT before 2.15 raises these for non-numeric or infinite time claims instead of DecodeError. + log.warning("Rejecting a workload identity with malformed time claims", exc_info=True) + raise jwt.InvalidTokenError("Workload identity has malformed time claims") from error 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 + raise WorkloadIdentityNotAllowed("Workload identity is not granted any Dag bundles") + if missing := self._unconfigured_grants.get(identity.sub): + raise WorkloadIdentityNotAllowed( + f"Workload identity is granted Dag bundles this API server does not have: {sorted(missing)}" + ) + instance_id: object = claims for key in self.config.instance_id_claim: if not isinstance(instance_id, dict) or key not in instance_id: raise jwt.InvalidTokenError("Workload identity lacks its configured instance claim") diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 6ca76b40c4a9d..9d1b46271b6b0 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2431,10 +2431,14 @@ execution_api: Requires issuer, HTTPS jwks_url, a dedicated audience, an instance_id_claim path (a list of keys into the signed identity claims), and grants mapping exact subjects - to allowed Dag bundle names. The API server refuses to start if a grant names a Dag - bundle it does not have configured. The optional algorithm defaults to RS256; ES256 and - EdDSA are also supported. Use a signed, per-replica instance claim: a service - account shared by replicas is not sufficient on its own. + to allowed Dag bundle names. If a grant names a Dag bundle the API server does not have + configured, the API server logs an error at startup and refuses that subject's exchanges. + The optional algorithm defaults to RS256; ES256 and EdDSA are also supported. Use a + signed, per-replica instance claim: a service account shared by replicas is not + sufficient on its own. + + Only one issuer is supported. Dag processors whose identities come from different + issuers, such as separate Kubernetes clusters, cannot all use the same API server. The JWKS endpoint must accept unauthenticated HTTPS requests and have a certificate trusted by the API server's HTTPX environment. Per-issuer CA and JWKS authentication settings are not supported. diff --git a/airflow-core/src/airflow/dag_processing/credentials.py b/airflow-core/src/airflow/dag_processing/credentials.py index 20b7cf9abcf92..e64d0aa46b142 100644 --- a/airflow-core/src/airflow/dag_processing/credentials.py +++ b/airflow-core/src/airflow/dag_processing/credentials.py @@ -25,10 +25,37 @@ from typing import Protocol import httpx2 +from tenacity import RetryCallState, retry, retry_if_exception, stop_after_attempt, wait_random_exponential 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 +from airflow.sdk.api.client import ( # noqa: SDK001 - reuse Execution API TLS and retry settings + API_RETRIES, + API_RETRY_WAIT_MAX, + API_RETRY_WAIT_MIN, + Client, + _log_and_trace_retry, + _should_retry_api_request, +) + +_backoff = wait_random_exponential(min=API_RETRY_WAIT_MIN, max=API_RETRY_WAIT_MAX) + + +class WorkloadIdentityUnavailableError(OSError): + """The identity supplier failed; the client handles it like an unreadable credential file.""" + + +def _wait_for_retry_after(retry_state: RetryCallState) -> float: + """Back off like the SDK client, but never retry a 503 before its ``Retry-After`` elapses.""" + wait = _backoff(retry_state) + error = retry_state.outcome.exception() if retry_state.outcome else None + if not isinstance(error, httpx2.HTTPStatusError) or error.response.status_code != 503: + return wait + try: + retry_after = float(error.response.headers["Retry-After"]) + except (KeyError, ValueError): # Missing, or an HTTP-date rather than seconds. + return wait + return max(wait, min(retry_after, API_RETRY_WAIT_MAX)) class SessionTokenProvider(Protocol): @@ -64,14 +91,28 @@ def _update_auth(self, response: httpx2.Response) -> None: # Runtime refresh headers must not replace the external workload identity. pass + # The server refetches signing keys at most every 10 seconds; the SDK's backoff alone often retries + # every attempt inside that cooldown. + @retry( + retry=retry_if_exception(_should_retry_api_request), + stop=stop_after_attempt(API_RETRIES), + wait=_wait_for_retry_after, + before_sleep=_log_and_trace_retry, + reraise=True, + ) + def _request_with_retry(self, *args, **kwargs): + return self._request_without_retry(*args, **kwargs) + 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. + read a projected service-account token that the platform rotates. Any exception it raises is + re-raised as ``WorkloadIdentityUnavailableError``, so a failed early renewal keeps the valid Job + credential in use. The owning Dag processor client closes this provider; do not share it between + clients. """ def __init__( @@ -93,7 +134,10 @@ def get_token( ) -> str: if not force_refresh and self._token is not None and monotonic() < self._renew_at: return self._token - identity = self._identity_token().strip() + try: + identity = self._identity_token().strip() + except Exception as error: + raise WorkloadIdentityUnavailableError("Unable to obtain the workload identity") from error if not identity: raise ValueError("Workload identity credential is empty") started_at = monotonic() diff --git a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py index 2d2ca1480e707..360b63244e5b3 100644 --- a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py +++ b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py @@ -195,6 +195,30 @@ async def handle(request): await refresh assert len(requests) == 2 + @patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) + async def test_jwks_body_that_is_not_an_object_is_a_failed_fetch(self, clock, ed25519_private_key): + bodies = [[{"kty": "OKP"}]] + + def handle(request): + return httpx.Response(200, json=bodies[-1]) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + jwks = JWKS( + url="https://example.com/keys", + client=client, + refresh_on_unknown_kid=True, + raise_on_fetch_error=True, + ) + with pytest.raises(JWKSUnavailableError): + await jwks.get_key("known") + clock.return_value = 111 + bodies.append({"keys": [key_to_jwk_dict(ed25519_private_key, "known")]}) + known = await jwks.get_key("known") + # Once the hourly refresh is due, a bad body must not stop the cached key from validating. + clock.return_value = 111 + jwks.refresh_interval_secs + 1 + bodies.append([{"kty": "OKP"}]) + assert await jwks.get_key("known") is known + @pytest.mark.parametrize("private_key", ["rsa_private_key", "ed25519_private_key"], indirect=True) async def test_fetch_jwks_success(self, private_key): jwk_content = json.dumps({"keys": [key_to_jwk_dict(private_key, "kid")]}) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py index 3de9c35f1b91d..e30aebdee2b2d 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py @@ -22,16 +22,19 @@ accepted by as few routes as possible. See PR #62582 (foundation) and PR #66989 (two-token mechanism). -This test checks the real API router and enforces two rules: +This test checks the real API router and enforces three rules: 1. Routes not listed in NON_DEFAULT_TOKEN_POLICY accept only ``token:execution``. 2. Routes listed in NON_DEFAULT_TOKEN_POLICY accept exactly the declared types. +3. Only HEALTH_ROUTES and WORKLOAD_IDENTITY_ROUTES skip ``require_auth``, and the latter + authenticate a workload identity instead. Rules 1 and 2 apply to every other route. Maintenance: - New route with default (execution-only) scope: no test change needed. - New route with non-default scope: add it to NON_DEFAULT_TOKEN_POLICY. - Route removed: remove it from NON_DEFAULT_TOKEN_POLICY (if present). - New token type: add it to the relevant entries in NON_DEFAULT_TOKEN_POLICY. +- New route outside ``require_auth``: add it to HEALTH_ROUTES or WORKLOAD_IDENTITY_ROUTES. """ from __future__ import annotations @@ -53,6 +56,12 @@ from tests_common.test_utils.api_fastapi import iter_api_routes +# Routes that take no credential at all. +HEALTH_ROUTES = {"GET /health", "GET /health/ping"} +# Routes that authenticate a deployment workload identity instead of an Airflow runtime token. +WORKLOAD_IDENTITY_ROUTES = {"POST /dag-processor/session-token"} + + # Routes that intentionally deviate from the default (execution-only) policy. # Any route NOT listed here must accept only {"execution"}. NON_DEFAULT_TOKEN_POLICY: dict[str, set[str]] = { @@ -148,21 +157,28 @@ def _get_path_and_query_param_names(dependant: Dependant) -> set[str]: def _all_route_policies() -> dict[str, set[str]]: - """Return a map of all API routes and their allowed token types.""" + """Return a map of the runtime-authenticated API routes and their allowed token types.""" return { key: set(getattr(_get_original_route(route), "allowed_token_types", {"execution"})) for key, route in _get_api_routes().items() + if key not in HEALTH_ROUTES | WORKLOAD_IDENTITY_ROUTES } class TestTokenScopeBoundaries: """Execution API routes must not silently gain or lose token type access.""" - def test_session_bootstrap_requires_deployment_identity(self): - route = _get_api_routes()["POST /dag-processor/session-token"] - dependencies = _get_dependency_calls(route.dependant) - assert require_workload_identity in dependencies - assert require_auth not in dependencies + def test_only_allowlisted_routes_skip_require_auth(self): + unauthenticated = { + key + for key, route in _get_api_routes().items() + if require_auth not in _get_dependency_calls(route.dependant) + } + assert unauthenticated == HEALTH_ROUTES | WORKLOAD_IDENTITY_ROUTES + + @pytest.mark.parametrize("route", sorted(WORKLOAD_IDENTITY_ROUTES)) + def test_workload_identity_routes_authenticate_the_identity(self, route): + assert require_workload_identity in _get_dependency_calls(_get_api_routes()[route].dependant) def test_all_default_routes_are_execution_only(self): actual = _all_route_policies() diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py index 896d6a3c3e13f..fa2f1761e16a5 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py @@ -24,14 +24,14 @@ import httpx2 import jwt import pytest +import svcs 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.auth.tokens import JWTValidator, 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 ( @@ -44,10 +44,9 @@ DagProcessorAPIClient, DagProcessorJobAlreadyRunning, ) -from airflow.dag_processing.credentials import FileSessionTokenProvider, WorkloadIdentitySessionTokenProvider +from airflow.dag_processing.credentials import WorkloadIdentitySessionTokenProvider, _SessionTokenClient 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 @@ -129,25 +128,32 @@ def verifier(client, exec_app, signing_key, identity_algorithm, monkeypatch): jwks_client = httpx.AsyncClient( transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"keys": [public_key]})) ) - verifier = exec_app.state.dag_processor_identity_verifier + verifier = svcs.Container(exec_app.state.svcs_registry).get(DagProcessorIdentityVerifier) monkeypatch.setattr(verifier.validator.jwks, "client", jwks_client) yield verifier client.portal.call(jwks_client.aclose) +def use_verifier(exec_app, verifier): + """Serve the exchange with another verifier until this test's API server lifespan ends.""" + exec_app.state.svcs_registry.register_value(DagProcessorIdentityVerifier, verifier) + + @pytest.fixture 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(exchange_available=lambda: True): + def forward(request): + if request.url.path == URL and not exchange_available(): + raise httpx2.ConnectError("session exchange unavailable", request=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 @@ -212,12 +218,12 @@ def test_issues_only_configured_grants_and_caps_lifetime( } ) def test_renewal_and_api_restart_keep_session_but_replicas_and_grants_do_not( - client, verifier, exec_app, identity_token, monkeypatch + client, verifier, exec_app, identity_token ): 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) + use_verifier(exec_app, 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"] @@ -256,11 +262,11 @@ def test_rejects_invalid_signed_claims(client, verifier, identity_token, overrid @pytest.mark.parametrize("leeway", [0, 3, 10]) @pytest.mark.parametrize("claim", ["iat", "nbf"]) def test_identity_clock_skew_uses_configured_leeway( - client, verifier, exec_app, identity_token, monkeypatch, leeway, claim + client, verifier, exec_app, identity_token, 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) + use_verifier(exec_app, replacement) assert exchange(client, identity_token(**{claim: NOW + leeway})).status_code == 200 assert exchange(client, identity_token(**{claim: NOW + leeway + 1})).status_code == 401 @@ -269,6 +275,7 @@ def test_identity_clock_skew_uses_configured_leeway( 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 response.headers["WWW-Authenticate"] == "Bearer" assert "token" not in response.json() @@ -324,8 +331,8 @@ def test_missing_bearer_is_unauthorized(client, verifier): 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) +def test_disabled_exchange_is_unavailable(client, exec_app, identity_token): + use_verifier(exec_app, None) assert exchange(client, identity_token()).status_code == 404 @@ -360,7 +367,7 @@ async def expire(token): assert exchange(client, identity_token()).status_code == 401 -@pytest.mark.parametrize("failure", ["connection", "503", "invalid-json", "invalid-keyset"]) +@pytest.mark.parametrize("failure", ["connection", "503", "invalid-json", "json-array", "invalid-keyset"]) def test_jwks_outage_is_retryable(client, verifier, identity_token, monkeypatch, failure): async def unavailable(url): if failure == "connection": @@ -370,6 +377,8 @@ async def unavailable(url): return httpx.Response(503, request=request) if failure == "invalid-json": return httpx.Response(200, content=b"unavailable", request=request) + if failure == "json-array": + return httpx.Response(200, json=[{"kty": "OKP"}], request=request) return httpx.Response(200, json={"keys": []}, request=request) monkeypatch.setattr(verifier.jwks.client, "get", unavailable) @@ -405,15 +414,19 @@ async def get_keys(url): request=httpx.Request("GET", url), ) - def advance_retry(retry_state): - clock.return_value += 11 + waits = [] + + def sleep(seconds): + waits.append(seconds) + clock.return_value += seconds 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) + # Keep the default wait policy: its backoff alone often retries inside the server's 10s cooldown. + monkeypatch.setattr(_SessionTokenClient._request_with_retry.retry, "sleep", sleep) processor = make_processor() assert processor.register_job() == processor.job_id assert processor.heartbeat() == JobState.RUNNING + assert waits == [10] assert len(attempts) == 2 @@ -433,11 +446,50 @@ def test_container_restart_waits_for_previous_job_to_become_stale(make_processor 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_unconfigured_bundle_grant_refuses_only_its_subject( + client, verifier, exec_app, identity_token, caplog +): + config = WorkloadIdentityConfig.model_validate( + {**TRUST, "grants": {**TRUST["grants"], "system:serviceaccount:team-b:processor": ["team-b"]}} + ) + use_verifier(exec_app, DagProcessorIdentityVerifier(config, client=verifier.validator.jwks.client)) + assert { + "event": "[execution_api] dag_processor_workload_identity grants Dag bundles missing from " + "[dag_processor] dag_bundle_config_list; these subjects cannot obtain a session", + "log_level": "error", + "unconfigured_grants": {"system:serviceaccount:team-b:processor": ["team-b"]}, + } in caplog + + assert exchange(client, identity_token()).status_code == 200 + response = exchange(client, identity_token(sub="system:serviceaccount:team-b:processor")) + assert response.status_code == 403 + assert response.json() == { + "detail": "Workload identity is granted Dag bundles this API server does not have: ['team-b']" + } + + +@pytest.mark.parametrize("error", [TypeError("exp"), OverflowError("exp")]) +@mock.patch.object(JWTValidator, "avalidated_claims", autospec=True) +def test_malformed_time_claims_from_older_pyjwt_are_invalid( + avalidated_claims, client, verifier, identity_token, error, caplog +): + """PyJWT before 2.15 raises these for list or infinite time claims instead of DecodeError.""" + avalidated_claims.side_effect = error + response = exchange(client, identity_token()) + assert response.status_code == 401 + assert response.json() == {"detail": "Invalid workload identity"} + assert { + "event": "Rejecting a workload identity with malformed time claims", + "log_level": "warning", + } in caplog + + +@mock.patch( + "airflow.api_fastapi.execution_api.workload_identity.uuid5", autospec=True, side_effect=TypeError("bug") +) +def test_verifier_bugs_are_not_reported_as_invalid_identity(uuid5, client, verifier, identity_token): + with pytest.raises(TypeError, match="bug"): + exchange(client, identity_token()) def test_older_version_has_no_exchange(client, verifier, identity_token): @@ -476,81 +528,48 @@ def test_identity_and_runtime_tokens_are_not_interchangeable(client, verifier, e assert response.status_code == 403 -@mock.patch.object( - FileSessionTokenProvider, "get_token", autospec=True, side_effect=AssertionError("file used") -) @mock.patch("airflow.dag_processing.api_client.monotonic", autospec=True, return_value=0) @mock.patch("airflow.dag_processing.credentials.monotonic", autospec=True, return_value=0) -def test_session_renewal_and_job_lifecycle_without_session_file( - credential_clock, - job_clock, - file_provider, - client, - verifier, - exec_app, - identity_token, - session, - time_machine, +def test_session_renewal_and_job_lifecycle( + credential_clock, job_clock, make_processor, session, time_machine ): - clear_db_jobs() clear_db_variables() - exec_app.dependency_overrides.pop(require_auth) exchange_available = True - - def forward(request): - if request.url.path == URL and not exchange_available: - raise httpx2.ConnectError("session exchange unavailable", request=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) - - provider = WorkloadIdentitySessionTokenProvider( - base_url="https://api/execution/", - identity_token=lambda: identity_token(exp=NOW + 600), - transport=httpx2.MockTransport(forward), - ) + processor = make_processor(exchange_available=lambda: exchange_available) + provider = processor._token_provider try: Variable.set("identity-test-key", "value") - with DagProcessorAPIClient( - base_url="https://api/execution/", - hostname="processor", - token_provider=provider, - transport=httpx2.MockTransport(forward), - ) as processor: - job_id = processor.register_job() - first_session = provider.get_token() - assert processor.heartbeat() == JobState.RUNNING - context = DagParseContext( - request=DagParseTokenBody( - attempt_id=uuid7(), bundle_name="dags-folder", relative_fileloc="dag.py" - ) + job_id = processor.register_job() + first_session = provider.get_token() + assert processor.heartbeat() == JobState.RUNNING + context = DagParseContext( + request=DagParseTokenBody( + attempt_id=uuid7(), bundle_name="dags-folder", relative_fileloc="dag.py" ) - with processor.use_parse(context): - assert processor.variables.get("identity-test-key").value == "value" - time_machine.shift(240) - credential_clock.return_value = job_clock.return_value = 240 - assert processor.heartbeat() == JobState.RUNNING - assert provider.get_token() != first_session - assert processor.job_id == job_id - assert not processor.restart_required - assert session.scalars(select(Job.id)).all() == [job_id] - - exchange_available = False - time_machine.shift(240) - credential_clock.return_value = job_clock.return_value = 480 - assert processor.heartbeat() == JobState.RUNNING - time_machine.shift(61) - credential_clock.return_value = job_clock.return_value = 541 - with pytest.raises(httpx2.ConnectError): - processor.heartbeat() - - exchange_available = True - assert processor.heartbeat() == JobState.RUNNING - assert processor.job_id == job_id - processor.complete_job(TerminalJobState.SUCCESS) - assert session.get(Job, job_id).state == JobState.SUCCESS.value - file_provider.assert_not_called() + ) + with processor.use_parse(context): + assert processor.variables.get("identity-test-key").value == "value" + time_machine.shift(240) + credential_clock.return_value = job_clock.return_value = 240 + assert processor.heartbeat() == JobState.RUNNING + assert provider.get_token() != first_session + assert processor.job_id == job_id + assert not processor.restart_required + assert session.scalars(select(Job.id)).all() == [job_id] + + exchange_available = False + time_machine.shift(240) + credential_clock.return_value = job_clock.return_value = 480 + assert processor.heartbeat() == JobState.RUNNING + time_machine.shift(61) + credential_clock.return_value = job_clock.return_value = 541 + with pytest.raises(httpx2.ConnectError): + processor.heartbeat() + + exchange_available = True + assert processor.heartbeat() == JobState.RUNNING + assert processor.job_id == job_id + processor.complete_job(TerminalJobState.SUCCESS) + assert session.get(Job, job_id).state == JobState.SUCCESS.value finally: - clear_db_jobs() clear_db_variables() diff --git a/airflow-core/tests/unit/dag_processing/test_credentials.py b/airflow-core/tests/unit/dag_processing/test_credentials.py index dbc973589d7d2..20a54f7524da2 100644 --- a/airflow-core/tests/unit/dag_processing/test_credentials.py +++ b/airflow-core/tests/unit/dag_processing/test_credentials.py @@ -24,15 +24,34 @@ from pydantic import ValidationError from tenacity import wait_none +from airflow.api_fastapi.execution_api.datamodels.job import JobState from airflow.api_fastapi.execution_api.versions import bundle from airflow.dag_processing.api_client import DagProcessorAPIClient -from airflow.dag_processing.credentials import SessionTokenProvider, WorkloadIdentitySessionTokenProvider -from airflow.sdk.api.client import API_RETRIES, API_TIMEOUT, BearerAuth, Client +from airflow.dag_processing.credentials import ( + SessionTokenProvider, + WorkloadIdentitySessionTokenProvider, + WorkloadIdentityUnavailableError, + _SessionTokenClient, +) +from airflow.sdk.api.client import ( + API_RETRIES, + API_RETRY_WAIT_MAX, + API_RETRY_WAIT_MIN, + API_TIMEOUT, + BearerAuth, + Client, +) + + +class RefreshError(Exception): + """Stands in for google-auth's RefreshError, which subclasses Exception directly.""" @pytest.fixture(autouse=True) def no_retry_wait(monkeypatch): monkeypatch.setattr(Client._request_with_retry.retry, "wait", wait_none()) + # The exchange keeps its own wait policy under test; only the sleeping is skipped. + monkeypatch.setattr(_SessionTokenClient._request_with_retry.retry, "sleep", lambda seconds: None) @pytest.fixture @@ -175,6 +194,28 @@ def test_missing_identity_fails_before_request(make_provider): assert not requests +def test_identity_supplier_errors_are_reported_as_unavailable(make_provider): + error = RefreshError("metadata server unavailable") + provider, requests = make_provider(identity=mock.Mock(spec=lambda: "", side_effect=error)) + with pytest.raises(WorkloadIdentityUnavailableError) as raised: + provider.get_token() + assert raised.value.__cause__ is error + assert not requests + + +@pytest.mark.parametrize( + ("retry_after", "wait"), + [("10", 10), ("3600", API_RETRY_WAIT_MAX), ("Wed, 21 Oct 2026 07:28:00 GMT", API_RETRY_WAIT_MIN)], +) +def test_exchange_retry_waits_for_retry_after(make_provider, monkeypatch, retry_after, wait): + waits = [] + monkeypatch.setattr(_SessionTokenClient._request_with_retry.retry, "sleep", waits.append) + provider, requests = make_provider(httpx2.Response(503, headers={"Retry-After": retry_after}), issued()) + assert provider.get_token(retry=True) == "session-1" + assert waits == [wait] + assert len(requests) == 2 + + def test_redirect_does_not_forward_identity(make_provider): provider, requests = make_provider(httpx2.Response(307, headers={"Location": "https://other/execution/"})) with pytest.raises(httpx2.HTTPStatusError): @@ -289,6 +330,34 @@ def handle(request): assert requests[-1].headers["Authorization"] == "Bearer job-1" +@mock.patch("airflow.dag_processing.api_client.monotonic", autospec=True) +def test_failed_identity_supplier_keeps_using_valid_job_token(client_clock, make_provider, clock): + client_clock.return_value = 0 + identity = mock.Mock(spec=lambda: "", side_effect=["identity-1", RefreshError("metadata server down")]) + provider, exchanges = make_provider(issued(), identity=identity) + requests = [] + + def handle(request): + requests.append(request) + if request.url.path.endswith("/jobs"): + return httpx2.Response(201, json={"job_id": 1, "token": "job-1", "expires_in": 300}) + return httpx2.Response(200, json={"state": "running"}) + + with DagProcessorAPIClient( + base_url="https://api/execution/", + hostname="processor", + token_provider=provider, + transport=httpx2.MockTransport(handle), + ) as client: + client.register_job() + client_clock.return_value = clock.return_value = 80 + assert client.heartbeat() == JobState.RUNNING + + assert identity.call_count == 2 + assert len(exchanges) == 1 + assert requests[-1].headers["Authorization"] == "Bearer job-1" + + @pytest.mark.parametrize("use_context", [False, True]) @pytest.mark.parametrize("close_fails", [False, True]) def test_client_close_releases_credentials_and_connections(use_context, close_fails): From 287d7f9c86ecdc284675250229ed0d2da6aab5f6 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 26 Oct 2026 09:54:15 +0100 Subject: [PATCH 6/6] Fix Dag processor identity exchange edge cases Processors must stop trusting keys withdrawn by their issuer and recover from temporary identity-provider failures without exhausting retries during its cooldown. Authentication must also remain the boundary for incoming trace context. --- .../src/airflow/api_fastapi/auth/tokens.py | 19 ++++--- .../airflow/api_fastapi/execution_api/app.py | 7 ++- .../execution_api/routes/session_tokens.py | 1 + .../execution_api/workload_identity.py | 6 +- .../src/airflow/dag_processing/credentials.py | 5 +- .../unit/api_fastapi/auth/test_tokens.py | 55 ++++++++++++++----- .../api_fastapi/execution_api/test_app.py | 47 ++++++++++++++++ .../versions/head/test_session_tokens.py | 45 ++++++++++----- .../unit/dag_processing/test_credentials.py | 19 +++++-- 9 files changed, 162 insertions(+), 42 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/auth/tokens.py b/airflow-core/src/airflow/api_fastapi/auth/tokens.py index 688706a3b6257..5fd1c74fd4712 100644 --- a/airflow-core/src/airflow/api_fastapi/auth/tokens.py +++ b/airflow-core/src/airflow/api_fastapi/auth/tokens.py @@ -163,10 +163,14 @@ async def fetch_jwks(self, *, force: bool = False) -> None: return try: - # PyJWKSet.from_dict raises AttributeError for a JSON array or scalar. + # Older PyJWT versions assume the body and every key are objects. if not isinstance(data, dict): raise TypeError(f"JWKS must be a JSON object, not {type(data).__name__}") - keyset = jwt.PyJWKSet.from_dict(data) + keys = data.get("keys") + if not isinstance(keys, list) or any(not isinstance(key, dict) for key in keys): + raise TypeError("JWKS keys must be an array of JSON objects") + # PyJWT rejects empty sets, but an issuer can remove all keys with an empty array. + keyset = jwt.PyJWKSet.from_dict(data) if keys else None except (jwt.PyJWKError, jwt.PyJWKSetError, KeyError, TypeError, ValueError): if not self.raise_on_fetch_error: raise @@ -174,7 +178,7 @@ async def fetch_jwks(self, *, force: bool = False) -> None: return self._jwks = keyset self._fetch_failed = False - log.debug("Fetched JWKS", url=self.url, keys=len(self._jwks.keys)) + log.debug("Fetched JWKS", url=self.url, keys=len(keyset.keys) if keyset is not None else 0) async def _fetch_remote_jwks(self) -> dict[str, Any] | None: try: @@ -251,10 +255,11 @@ async def get_key(self, kid: str) -> jwt.PyJWK: except KeyError: if self.refresh_on_unknown_kid: await self.fetch_jwks(force=True) - try: - return self._jwks[kid] - except KeyError: - pass + if self._jwks is not None: + try: + return self._jwks[kid] + except KeyError: + pass if self.raise_on_fetch_error and self._fetch_failed: raise JWKSUnavailableError("Unable to load the required signing key") diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index 078b68e3196dd..6518dc6c0a18c 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -289,6 +289,8 @@ def _inject_trace_context_dep(routes, mode: str) -> None: ``unsafe-always`` mode is handled with an app-level dependency instead, so extraction runs before authentication declared on any enclosing router. """ + from airflow.api_fastapi.execution_api.routes.session_tokens import require_workload_identity + for context in iter_route_contexts(routes): route = context.original_route if not isinstance(route, APIRoute): @@ -300,7 +302,10 @@ def _inject_trace_context_dep(routes, mode: str) -> None: if mode != "only-authenticated": continue # context.dependencies also holds those of the enclosing routers and include_router() calls. - if any(getattr(d, "dependency", None) is require_auth for d in context.dependencies): + if any( + getattr(d, "dependency", None) in (require_auth, require_workload_identity) + for d in context.dependencies + ): route.dependencies.append(Depends(_extract_w3c_trace_context)) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py index e61174e4a5468..10da336dda812 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py @@ -76,6 +76,7 @@ async def require_workload_identity( @router.post( "/dag-processor/session-token", + dependencies=[Depends(require_workload_identity)], responses=create_openapi_http_exception_doc( [ (status.HTTP_401_UNAUTHORIZED, "The workload identity is missing, invalid, or expired"), diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py b/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py index b5eab598c8dc2..73c5079ff463b 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py @@ -26,7 +26,7 @@ from typing import Annotated, Literal from uuid import NAMESPACE_URL, UUID, uuid5 -import httpx +import httpx2 import jwt import structlog from pydantic import BaseModel, ConfigDict, Field, HttpUrl, field_validator @@ -84,7 +84,7 @@ class _IdentityClaims(BaseModel): class DagProcessorIdentityVerifier: """Validate external JWTs independently of Airflow's runtime token trust.""" - def __init__(self, config: WorkloadIdentityConfig, *, client: httpx.AsyncClient): + def __init__(self, config: WorkloadIdentityConfig, *, client: httpx2.AsyncClient): if config.audience in conf.getlist("execution_api", "jwt_audience"): raise ValueError("Workload identity must use a dedicated audience") configured = get_configured_bundle_names() @@ -163,5 +163,5 @@ async def create_dag_processor_identity_verifier() -> AsyncIterator[DagProcessor yield None return config = WorkloadIdentityConfig.model_validate_json(raw) - async with httpx.AsyncClient(timeout=5.0, follow_redirects=False) as client: + async with httpx2.AsyncClient(timeout=5.0, follow_redirects=False) as client: yield DagProcessorIdentityVerifier(config, client=client) diff --git a/airflow-core/src/airflow/dag_processing/credentials.py b/airflow-core/src/airflow/dag_processing/credentials.py index e64d0aa46b142..2b5ff6ff6e270 100644 --- a/airflow-core/src/airflow/dag_processing/credentials.py +++ b/airflow-core/src/airflow/dag_processing/credentials.py @@ -19,6 +19,7 @@ from __future__ import annotations +import math from collections.abc import Callable from pathlib import Path from time import monotonic @@ -55,7 +56,9 @@ def _wait_for_retry_after(retry_state: RetryCallState) -> float: retry_after = float(error.response.headers["Retry-After"]) except (KeyError, ValueError): # Missing, or an HTTP-date rather than seconds. return wait - return max(wait, min(retry_after, API_RETRY_WAIT_MAX)) + if not math.isfinite(retry_after): + return wait + return max(wait, retry_after) class SessionTokenProvider(Protocol): diff --git a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py index 360b63244e5b3..d1f91e5a741b6 100644 --- a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py +++ b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py @@ -70,9 +70,9 @@ async def test_refresh_unknown_key_before_regular_interval( def handle(request): requests.append(request) - return httpx.Response(200, json={"keys": keys}) + return httpx2.Response(200, json={"keys": keys}) - async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: jwks = JWKS( url="https://example.com/keys", client=client, refresh_on_unknown_kid=refresh_on_unknown_kid ) @@ -98,9 +98,9 @@ async def test_concurrent_unknown_keys_share_fetch_and_cooldown( async def handle(request): requests.append(request) await asyncio.sleep(0) - return httpx.Response(status, json={"keys": [key_to_jwk_dict(ed25519_private_key, "known")]}) + return httpx2.Response(status, json={"keys": [key_to_jwk_dict(ed25519_private_key, "known")]}) - async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: jwks = JWKS( url="https://example.com/keys", client=client, @@ -135,9 +135,9 @@ async def test_failed_unknown_key_refresh_keeps_cached_keys_and_rate_limits( def handle(request): requests.append(request) - return httpx.Response(status, json={"keys": keys}) + return httpx2.Response(status, json={"keys": keys}) - async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: jwks = JWKS( url="https://example.com/keys", client=client, @@ -172,9 +172,9 @@ async def handle(request): if len(requests) > 1: fetching.set() await release.wait() - return httpx.Response(200, json={"keys": [key_to_jwk_dict(ed25519_private_key, "known")]}) + return httpx2.Response(200, json={"keys": [key_to_jwk_dict(ed25519_private_key, "known")]}) - async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: jwks = JWKS(url="https://example.com/keys", client=client, refresh_on_unknown_kid=True) known = await jwks.get_key("known") clock.return_value = 111 @@ -195,14 +195,15 @@ async def handle(request): await refresh assert len(requests) == 2 + @pytest.mark.parametrize("body", [[{"kty": "OKP"}], {"keys": [None]}, {"keys": ["invalid"]}]) @patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) - async def test_jwks_body_that_is_not_an_object_is_a_failed_fetch(self, clock, ed25519_private_key): - bodies = [[{"kty": "OKP"}]] + async def test_malformed_jwks_is_a_failed_fetch(self, clock, ed25519_private_key, body): + bodies = [body] def handle(request): - return httpx.Response(200, json=bodies[-1]) + return httpx2.Response(200, json=bodies[-1]) - async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: jwks = JWKS( url="https://example.com/keys", client=client, @@ -216,9 +217,37 @@ def handle(request): known = await jwks.get_key("known") # Once the hourly refresh is due, a bad body must not stop the cached key from validating. clock.return_value = 111 + jwks.refresh_interval_secs + 1 - bodies.append([{"kty": "OKP"}]) + bodies.append(body) assert await jwks.get_key("known") is known + @pytest.mark.parametrize("refresh", ["initial", "scheduled", "unknown-key"]) + @patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) + async def test_empty_jwks_removes_cached_keys(self, clock, ed25519_private_key, refresh): + keys = [] if refresh == "initial" else [key_to_jwk_dict(ed25519_private_key, "known")] + requests = [] + + def handle(request): + requests.append(request) + return httpx2.Response(200, json={"keys": keys}) + + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: + jwks = JWKS( + url="https://example.com/keys", + client=client, + refresh_on_unknown_kid=True, + raise_on_fetch_error=True, + ) + if refresh != "initial": + await jwks.get_key("known") + keys.clear() + clock.return_value += jwks.refresh_interval_secs + 1 if refresh == "scheduled" else 11 + if refresh == "unknown-key": + with pytest.raises(KeyError): + await jwks.get_key("unknown") + with pytest.raises(KeyError): + await jwks.get_key("known") + assert len(requests) == (1 if refresh == "initial" else 2) + @pytest.mark.parametrize("private_key", ["rsa_private_key", "ed25519_private_key"], indirect=True) async def test_fetch_jwks_success(self, private_key): jwk_content = json.dumps({"keys": [key_to_jwk_dict(private_key, "kid")]}) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py index 175e9b8993db0..add0ea3274745 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py @@ -21,10 +21,12 @@ import json import threading from contextlib import asynccontextmanager +from datetime import UTC, datetime from unittest import mock from uuid import UUID import httpx2 +import jwt import pytest from fastapi import APIRouter, Depends, FastAPI, Request, status from fastapi.params import Security as SecurityParam @@ -48,6 +50,10 @@ from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken from airflow.api_fastapi.execution_api.security import require_auth from airflow.api_fastapi.execution_api.versions import bundle +from airflow.api_fastapi.execution_api.workload_identity import ( + DagProcessorIdentity, + DagProcessorIdentityVerifier, +) from airflow.models.xcom import XComModel from airflow.utils.session import create_session_async @@ -486,6 +492,47 @@ def test_only_authenticated_detects_auth_added_by_include_router(self): assert any(_is_trace_context_dep(d) for d in router.routes[0].dependencies) + @pytest.mark.parametrize("mode", ["only-authenticated", "unsafe-always", "never"]) + @pytest.mark.parametrize("valid_identity", [True, False]) + @mock.patch.object(otel_propagate, "extract", autospec=True) + def test_workload_identity_trace_context_extraction(self, extract, mode, valid_identity, time_machine): + now = datetime(2026, 10, 5, 12, 0, tzinfo=UTC) + time_machine.move_to(now, tick=False) + events = [] + verifier = mock.create_autospec(DagProcessorIdentityVerifier, instance=True) + + async def authenticate(token): + events.append("authenticate") + if not valid_identity: + raise jwt.InvalidTokenError("Invalid identity") + return DagProcessorIdentity(UUID(int=1), frozenset({"dags-folder"}), now.timestamp() + 600) + + def extract_context(headers): + events.append("extract") + return otel_context.get_current() + + verifier.authenticate.side_effect = authenticate + extract.side_effect = extract_context + app = self._build_app(mode) + with TestClient(app) as client: + app.state.svcs_registry.register_value(DagProcessorIdentityVerifier, verifier) + response = client.post( + "/dag-processor/session-token", + headers={ + "Authorization": "Bearer workload-identity", + "traceparent": "00-12345678901234567890123456789012-1234567890123456-01", + }, + ) + + assert response.status_code == (200 if valid_identity else 401) + expect_extract = mode == "unsafe-always" or (mode == "only-authenticated" and valid_identity) + assert extract.call_count == int(expect_extract) + assert events[0] == ("extract" if mode == "unsafe-always" else "authenticate") + if expect_extract: + assert extract.call_args.args[0]["traceparent"] == ( + "00-12345678901234567890123456789012-1234567890123456-01" + ) + def test_trace_context_dep_cleans_up_on_route_exception(self): """Verify extract and cleanup run correctly when a route handler raises.""" app = self._build_app("unsafe-always") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py index fa2f1761e16a5..779ef5ec98bb8 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py @@ -20,7 +20,6 @@ from datetime import UTC, datetime from unittest import mock -import httpx import httpx2 import jwt import pytest @@ -125,8 +124,8 @@ def verifier(client, exec_app, signing_key, identity_algorithm, monkeypatch): } 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]})) + jwks_client = httpx2.AsyncClient( + transport=httpx2.MockTransport(lambda request: httpx2.Response(200, json={"keys": [public_key]})) ) verifier = svcs.Container(exec_app.state.svcs_registry).get(DagProcessorIdentityVerifier) monkeypatch.setattr(verifier.validator.jwks, "client", jwks_client) @@ -305,7 +304,7 @@ def test_rotated_signing_key_refreshes_after_cooldown( async def get_keys(url): requests.append(url) - return httpx.Response(200, json={"keys": keys}, request=httpx.Request("GET", url)) + return httpx2.Response(200, json={"keys": keys}, request=httpx2.Request("GET", url)) monkeypatch.setattr(verifier.validator.jwks.client, "get", get_keys) original = decode(exchange(client, identity_token())) @@ -371,15 +370,15 @@ async def expire(token): 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) + raise httpx2.ConnectError("offline") + request = httpx2.Request("GET", url) if failure == "503": - return httpx.Response(503, request=request) + return httpx2.Response(503, request=request) if failure == "invalid-json": - return httpx.Response(200, content=b"unavailable", request=request) + return httpx2.Response(200, content=b"unavailable", request=request) if failure == "json-array": - return httpx.Response(200, json=[{"kty": "OKP"}], request=request) - return httpx.Response(200, json={"keys": []}, request=request) + return httpx2.Response(200, json=[{"kty": "OKP"}], request=request) + return httpx2.Response(200, json={"keys": [None]}, request=request) monkeypatch.setattr(verifier.jwks.client, "get", unavailable) response = exchange(client, identity_token()) @@ -388,12 +387,32 @@ async def unavailable(url): assert response.headers["Retry-After"] == "10" +@pytest.mark.parametrize("cached", [False, True]) +@mock.patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100) +def test_empty_jwks_rejects_identities_with_or_without_cached_keys( + clock, client, verifier, identity_token, monkeypatch, cached +): + if cached: + assert exchange(client, identity_token()).status_code == 200 + clock.return_value = 3701 + + async def get_keys(url): + return httpx2.Response(200, json={"keys": []}, request=httpx2.Request("GET", url)) + + monkeypatch.setattr(verifier.jwks.client, "get", get_keys) + response = exchange(client, identity_token()) + assert response.status_code == 401 + assert response.json() == {"detail": "Invalid workload identity"} + assert response.headers["WWW-Authenticate"] == "Bearer" + assert "Retry-After" not in response.headers + + @mock.patch("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") + raise httpx2.ConnectError("offline") monkeypatch.setattr(verifier.jwks.client, "get", unavailable) clock.return_value = 3701 @@ -408,10 +427,10 @@ def test_registration_recovers_from_jwks_outage(clock, make_processor, verifier, async def get_keys(url): attempts.append(url) - return httpx.Response( + return httpx2.Response( 503 if len(attempts) == 1 else 200, json={"keys": [key_to_jwk_dict(signing_key, kid="key-1")]}, - request=httpx.Request("GET", url), + request=httpx2.Request("GET", url), ) waits = [] diff --git a/airflow-core/tests/unit/dag_processing/test_credentials.py b/airflow-core/tests/unit/dag_processing/test_credentials.py index 20a54f7524da2..5a54cdcdf3a06 100644 --- a/airflow-core/tests/unit/dag_processing/test_credentials.py +++ b/airflow-core/tests/unit/dag_processing/test_credentials.py @@ -22,10 +22,11 @@ import httpx2 import pytest from pydantic import ValidationError -from tenacity import wait_none +from tenacity import wait_none, wait_random_exponential from airflow.api_fastapi.execution_api.datamodels.job import JobState from airflow.api_fastapi.execution_api.versions import bundle +from airflow.dag_processing import credentials from airflow.dag_processing.api_client import DagProcessorAPIClient from airflow.dag_processing.credentials import ( SessionTokenProvider, @@ -204,11 +205,21 @@ def test_identity_supplier_errors_are_reported_as_unavailable(make_provider): @pytest.mark.parametrize( - ("retry_after", "wait"), - [("10", 10), ("3600", API_RETRY_WAIT_MAX), ("Wed, 21 Oct 2026 07:28:00 GMT", API_RETRY_WAIT_MIN)], + ("retry_after", "maximum_wait", "wait"), + [ + ("10", API_RETRY_WAIT_MAX, 10), + ("10", 1, 10), + ("3600", API_RETRY_WAIT_MAX, 3600), + ("inf", API_RETRY_WAIT_MAX, API_RETRY_WAIT_MIN), + ("Wed, 21 Oct 2026 07:28:00 GMT", API_RETRY_WAIT_MAX, API_RETRY_WAIT_MIN), + ], ) -def test_exchange_retry_waits_for_retry_after(make_provider, monkeypatch, retry_after, wait): +def test_exchange_retry_waits_for_retry_after(make_provider, monkeypatch, retry_after, maximum_wait, wait): waits = [] + monkeypatch.setattr(credentials, "API_RETRY_WAIT_MAX", maximum_wait) + monkeypatch.setattr( + credentials, "_backoff", wait_random_exponential(min=API_RETRY_WAIT_MIN, max=maximum_wait) + ) monkeypatch.setattr(_SessionTokenClient._request_with_retry.retry, "sleep", waits.append) provider, requests = make_provider(httpx2.Response(503, headers={"Retry-After": retry_after}), issued()) assert provider.get_token(retry=True) == "session-1"