Skip to content

Commit ca13d4a

Browse files
committed
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.
1 parent a7ad561 commit ca13d4a

12 files changed

Lines changed: 362 additions & 146 deletions

File tree

‎airflow-core/docs/security/jwt_token_authentication.rst‎

Lines changed: 22 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -520,7 +520,7 @@ not yet use this client. Provisioning a token file alone does not switch that co
520520
HTTP or remove its need for database access.
521521

522522
In-memory sessions from workload identity
523-
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
523+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
524524

525525
Processor integrations can use ``WorkloadIdentitySessionTokenProvider`` to obtain and renew
526526
``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
545545
``ES256``, or ``EdDSA``). Use a dedicated audience, distinct from ``[execution_api] jwt_audience``.
546546
The API server verifies the signature, issuer, audience, issued-at time, expiry, and any not-before
547547
claim, using ``[api_auth] jwt_leeway`` for clock skew. Only exact subjects listed in ``grants`` can
548-
obtain a session. Every granted bundle must be configured on the API server, which otherwise refuses
549-
to start. This option is empty by default, disabling the exchange.
548+
obtain a session. If a grant names a Dag bundle that the API server does not have configured, the
549+
API server logs an error at startup and refuses that subject's exchanges with HTTP 403. This option
550+
is empty by default, disabling the exchange.
550551
Apply the same trust configuration to all API-server replicas and restart them after changes.
551552

553+
The option trusts a single issuer. Dag processors whose identities come from different issuers, such
554+
as separate Kubernetes clusters, cannot all obtain sessions from the same API server.
555+
552556
The JWKS endpoint must accept unauthenticated HTTPS GET requests from the Airflow API server.
553-
The fetcher sends no authentication credentials and uses HTTPX's default certificate trust,
554-
including ``SSL_CERT_FILE`` or ``SSL_CERT_DIR`` when configured. These environment settings affect
555-
other HTTPX clients in the API-server process; there is no per-issuer CA or authentication option.
557+
The fetcher sends no authentication credentials and uses HTTPX's default certificate trust.
558+
``SSL_CERT_FILE`` or ``SSL_CERT_DIR``, when set, replace the default CA bundle rather than adding
559+
to it, and they apply to every HTTPX client in the API-server process. A bundle that holds only a
560+
private CA cuts those clients off from public roots, so include the roots they still need. There is
561+
no per-issuer CA or authentication option.
556562
Authenticated JWKS endpoints, including Kubernetes API-server discovery protected by service-account
557563
authentication, are not supported by this implementation. For those clusters, serve the public JWKS
558564
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
586592
active Jobs at the same time.
587593

588594
The processor integration supplies a callable that reads the current platform identity on each
589-
exchange. For example:
595+
exchange. Any exception the callable raises is re-raised as ``WorkloadIdentityUnavailableError``, an
596+
``OSError``, so a failed optional renewal leaves the valid Job credential in use. For example:
590597

591598
.. code-block:: python
592599
@@ -613,8 +620,9 @@ session JWT to disk and does not fall back to a session-token file. ``token_prov
613620
``token_file`` are mutually exclusive. HTTPS is required for the exchange, and redirects are rejected.
614621
The exchange uses the same API TLS configuration and request timeout as the Task SDK client,
615622
including private certificate authorities and client certificates. Registration retries transient
616-
exchange failures using the configured Execution API retry policy; optional early renewal makes
617-
one bounded attempt so it does not delay requests using a valid Job credential.
623+
exchange failures using the configured Execution API retry policy, waiting at least the
624+
``Retry-After`` interval of an HTTP 503 before the next attempt. Optional early renewal makes one
625+
bounded attempt so it does not delay requests using a valid Job credential.
618626

619627
Session expiry is capped by both the external identity's remaining lifetime and
620628
``[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
631639
API-server process. Publish new keys before signing with them and retain old keys while their tokens
632640
remain valid.
633641

634-
If a required signing key cannot be fetched, the exchange returns HTTP 503 with ``Retry-After``,
635-
allowing registration to retry. A cached matching key remains usable during an outage. Invalid
636-
identities and keys absent from a successfully fetched keyset return HTTP 401.
642+
If a required signing key cannot be fetched, the exchange returns HTTP 503 with ``Retry-After``, and
643+
registration waits at least that long before retrying. A cached matching key remains usable during
644+
an outage. Invalid identities and keys absent from a successfully fetched key set return HTTP 401.
637645

638646
File-based provisioning
639-
~~~~~~~~~~~~~~~~~~~~~~~
647+
^^^^^^^^^^^^^^^^^^^^^^^
640648

641649
On a trusted host with the API signing key, provision a session token for the bundles
642650
the client may access::
@@ -652,7 +660,7 @@ provisioner with the same token file and bundles keeps that session, so running
652660
their Jobs.
653661

654662
Credential permissions and renewal
655-
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
663+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
656664

657665
The three credentials have different issuers and permissions:
658666

‎airflow-core/src/airflow/api_fastapi/auth/tokens.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,9 @@ async def fetch_jwks(self, *, force: bool = False) -> None:
163163
return
164164

165165
try:
166+
# PyJWKSet.from_dict raises AttributeError for a JSON array or scalar.
167+
if not isinstance(data, dict):
168+
raise TypeError(f"JWKS must be a JSON object, not {type(data).__name__}")
166169
keyset = jwt.PyJWKSet.from_dict(data)
167170
except (jwt.PyJWKError, jwt.PyJWKSetError, KeyError, TypeError, ValueError):
168171
if not self.raise_on_fetch_error:

‎airflow-core/src/airflow/api_fastapi/execution_api/app.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -112,12 +112,12 @@ async def lifespan(app: FastAPI, registry: svcs.Registry):
112112

113113
# InProcessExecutionAPI stubs out the identity verifier: don't create one in that case.
114114
if DagProcessorIdentityVerifier in registry:
115-
app.state.dag_processor_identity_verifier = None
116115
yield
117116
return
118117

119118
async with create_dag_processor_identity_verifier() as verifier:
120-
app.state.dag_processor_identity_verifier = verifier
119+
# None when the exchange is disabled; the route answers 404 then.
120+
registry.register_value(DagProcessorIdentityVerifier, verifier)
121121
yield
122122

123123

‎airflow-core/src/airflow/api_fastapi/execution_api/in_process.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ def app(self):
116116
# Auth dependency overrides bypass validation, so the stub skips validator setup in lifespan().
117117
registry.register_value(JWTValidator, None)
118118
# No external Dag processor can reach this app, so it must not load trust settings meant for
119-
# API servers, whose Dag bundle checks can fail in other components.
119+
# API servers, which can fail to load in components configured differently.
120120
registry.register_value(DagProcessorIdentityVerifier, None)
121121
_configure_in_process_auth(self._app, triggerer_id=self.triggerer_id)
122122

‎airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py‎

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import jwt
2525
import svcs
2626
from cadwyn import VersionedAPIRouter
27-
from fastapi import Depends, HTTPException, Request, Response, status
27+
from fastapi import Depends, HTTPException, Response, status
2828
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
2929
from pydantic import ValidationError
3030

@@ -48,10 +48,10 @@
4848

4949

5050
async def require_workload_identity(
51-
request: Request,
5251
credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(identity_bearer)],
52+
services: svcs.Container = DepContainer,
5353
) -> DagProcessorIdentity:
54-
verifier: DagProcessorIdentityVerifier | None = request.app.state.dag_processor_identity_verifier
54+
verifier: DagProcessorIdentityVerifier | None = await services.aget(DagProcessorIdentityVerifier)
5555
if verifier is None:
5656
raise HTTPException(status.HTTP_404_NOT_FOUND, "Workload identity provisioning is not enabled")
5757
if credentials is None:
@@ -66,22 +66,20 @@ async def require_workload_identity(
6666
"Workload identity signing keys are unavailable",
6767
headers={"Retry-After": str(verifier.jwks.refresh_retry_interval_secs)},
6868
) from None
69-
except (jwt.PyJWTError, ValidationError, TypeError, OverflowError):
69+
except (jwt.PyJWTError, ValidationError):
7070
raise HTTPException(
7171
status.HTTP_401_UNAUTHORIZED, "Invalid workload identity", headers={"WWW-Authenticate": "Bearer"}
7272
) from None
73-
except WorkloadIdentityNotAllowed:
74-
raise HTTPException(
75-
status.HTTP_403_FORBIDDEN, "Workload identity is not granted any Dag bundles"
76-
) from None
73+
except WorkloadIdentityNotAllowed as error:
74+
raise HTTPException(status.HTTP_403_FORBIDDEN, str(error)) from None
7775

7876

7977
@router.post(
8078
"/dag-processor/session-token",
8179
responses=create_openapi_http_exception_doc(
8280
[
8381
(status.HTTP_401_UNAUTHORIZED, "The workload identity is missing, invalid, or expired"),
84-
(status.HTTP_403_FORBIDDEN, "The identity is not authorized as a Dag processor"),
82+
(status.HTTP_403_FORBIDDEN, "The identity has no usable Dag bundle grant"),
8583
(status.HTTP_404_NOT_FOUND, "Workload identity provisioning is disabled"),
8684
(status.HTTP_503_SERVICE_UNAVAILABLE, "The identity provider's signing keys are unavailable"),
8785
]
@@ -92,11 +90,21 @@ def exchange_session_token(
9290
identity: Annotated[DagProcessorIdentity, Depends(require_workload_identity)],
9391
services: svcs.Container = DepContainer,
9492
) -> DagProcessorSessionTokenResponse:
93+
"""
94+
Exchange an external workload identity for a Dag processor session credential.
95+
96+
The bearer token is an OIDC identity issued by the deployment platform, not an Airflow token. The
97+
API server picks the Dag bundle grants from its own configuration, and the session cannot outlive
98+
the identity. While the identity provider's signing keys cannot be fetched, the exchange returns
99+
503 with ``Retry-After``.
100+
"""
95101
generator: JWTGenerator = services.get(JWTGenerator)
96102
try:
97103
lifetime = get_token_lifetime(generator, identity.expires_at, "Workload identity has expired")
98104
except ExpiredDagProcessorToken as error:
99-
raise HTTPException(status.HTTP_401_UNAUTHORIZED, str(error)) from None
105+
raise HTTPException(
106+
status.HTTP_401_UNAUTHORIZED, str(error), headers={"WWW-Authenticate": "Bearer"}
107+
) from None
100108
token = generate_dag_processor_session_token(
101109
generator, session_id=identity.session_id, bundle_names=identity.bundle_names, valid_for=lifetime
102110
)

‎airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py‎

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,18 +23,21 @@
2323
from collections.abc import AsyncIterator
2424
from contextlib import asynccontextmanager
2525
from dataclasses import dataclass
26-
from typing import Annotated, Any, Literal
26+
from typing import Annotated, Literal
2727
from uuid import NAMESPACE_URL, UUID, uuid5
2828

2929
import httpx
3030
import jwt
31+
import structlog
3132
from pydantic import BaseModel, ConfigDict, Field, HttpUrl, field_validator
3233

3334
from airflow.api_fastapi.auth.tokens import JWKS, JWTValidator
3435
from airflow.api_fastapi.execution_api.datamodels.token import DagBundleGrant
3536
from airflow.configuration import conf
3637
from airflow.dag_processing.bundles.manager import get_configured_bundle_names
3738

39+
log = structlog.get_logger(logger_name=__name__)
40+
3841
NonemptyString = Annotated[str, Field(min_length=1)]
3942

4043

@@ -84,10 +87,20 @@ class DagProcessorIdentityVerifier:
8487
def __init__(self, config: WorkloadIdentityConfig, *, client: httpx.AsyncClient):
8588
if config.audience in conf.getlist("execution_api", "jwt_audience"):
8689
raise ValueError("Workload identity must use a dedicated audience")
87-
if unknown := frozenset().union(*config.grants.values()) - get_configured_bundle_names():
88-
raise ValueError(
90+
configured = get_configured_bundle_names()
91+
# Not fatal: a renamed bundle must not stop the whole API server, only the subjects granted it.
92+
self._unconfigured_grants = {
93+
subject: missing
94+
for subject, bundles in config.grants.items()
95+
if (missing := bundles - configured)
96+
}
97+
if self._unconfigured_grants:
98+
log.error(
8999
"[execution_api] dag_processor_workload_identity grants Dag bundles missing from "
90-
f"[dag_processor] dag_bundle_config_list: {sorted(unknown)}"
100+
"[dag_processor] dag_bundle_config_list; these subjects cannot obtain a session",
101+
unconfigured_grants={
102+
subject: sorted(missing) for subject, missing in self._unconfigured_grants.items()
103+
},
91104
)
92105
self.config = config
93106
self.jwks = JWKS(
@@ -102,11 +115,20 @@ def __init__(self, config: WorkloadIdentityConfig, *, client: httpx.AsyncClient)
102115
)
103116

104117
async def authenticate(self, token: str) -> DagProcessorIdentity:
105-
claims = await self.validator.avalidated_claims(token)
118+
try:
119+
claims = await self.validator.avalidated_claims(token)
120+
except (TypeError, OverflowError) as error:
121+
# PyJWT before 2.15 raises these for non-numeric or infinite time claims instead of DecodeError.
122+
log.warning("Rejecting a workload identity with malformed time claims", exc_info=True)
123+
raise jwt.InvalidTokenError("Workload identity has malformed time claims") from error
106124
identity = _IdentityClaims.model_validate(claims)
107125
if not (grants := self.config.grants.get(identity.sub)):
108-
raise WorkloadIdentityNotAllowed("Workload identity is not authorized as a Dag processor")
109-
instance_id: Any = claims
126+
raise WorkloadIdentityNotAllowed("Workload identity is not granted any Dag bundles")
127+
if missing := self._unconfigured_grants.get(identity.sub):
128+
raise WorkloadIdentityNotAllowed(
129+
f"Workload identity is granted Dag bundles this API server does not have: {sorted(missing)}"
130+
)
131+
instance_id: object = claims
110132
for key in self.config.instance_id_claim:
111133
if not isinstance(instance_id, dict) or key not in instance_id:
112134
raise jwt.InvalidTokenError("Workload identity lacks its configured instance claim")

‎airflow-core/src/airflow/config_templates/config.yml‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2431,10 +2431,14 @@ execution_api:
24312431
24322432
Requires issuer, HTTPS jwks_url, a dedicated audience, an instance_id_claim path
24332433
(a list of keys into the signed identity claims), and grants mapping exact subjects
2434-
to allowed Dag bundle names. The API server refuses to start if a grant names a Dag
2435-
bundle it does not have configured. The optional algorithm defaults to RS256; ES256 and
2436-
EdDSA are also supported. Use a signed, per-replica instance claim: a service
2437-
account shared by replicas is not sufficient on its own.
2434+
to allowed Dag bundle names. If a grant names a Dag bundle the API server does not have
2435+
configured, the API server logs an error at startup and refuses that subject's exchanges.
2436+
The optional algorithm defaults to RS256; ES256 and EdDSA are also supported. Use a
2437+
signed, per-replica instance claim: a service account shared by replicas is not
2438+
sufficient on its own.
2439+
2440+
Only one issuer is supported. Dag processors whose identities come from different
2441+
issuers, such as separate Kubernetes clusters, cannot all use the same API server.
24382442
The JWKS endpoint must accept unauthenticated HTTPS requests and have a certificate
24392443
trusted by the API server's HTTPX environment. Per-issuer CA and JWKS authentication
24402444
settings are not supported.

‎airflow-core/src/airflow/dag_processing/credentials.py‎

Lines changed: 48 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,37 @@
2525
from typing import Protocol
2626

2727
import httpx2
28+
from tenacity import RetryCallState, retry, retry_if_exception, stop_after_attempt, wait_random_exponential
2829

2930
from airflow.api_fastapi.execution_api.datamodels.job import DagProcessorSessionTokenResponse
3031
from airflow.api_fastapi.execution_api.versions import bundle
31-
from airflow.sdk.api.client import Client # noqa: SDK001 - reuse Execution API TLS and retry settings
32+
from airflow.sdk.api.client import ( # noqa: SDK001 - reuse Execution API TLS and retry settings
33+
API_RETRIES,
34+
API_RETRY_WAIT_MAX,
35+
API_RETRY_WAIT_MIN,
36+
Client,
37+
_log_and_trace_retry,
38+
_should_retry_api_request,
39+
)
40+
41+
_backoff = wait_random_exponential(min=API_RETRY_WAIT_MIN, max=API_RETRY_WAIT_MAX)
42+
43+
44+
class WorkloadIdentityUnavailableError(OSError):
45+
"""The identity supplier failed; the client handles it like an unreadable credential file."""
46+
47+
48+
def _wait_for_retry_after(retry_state: RetryCallState) -> float:
49+
"""Back off like the SDK client, but never retry a 503 before its ``Retry-After`` elapses."""
50+
wait = _backoff(retry_state)
51+
error = retry_state.outcome.exception() if retry_state.outcome else None
52+
if not isinstance(error, httpx2.HTTPStatusError) or error.response.status_code != 503:
53+
return wait
54+
try:
55+
retry_after = float(error.response.headers["Retry-After"])
56+
except (KeyError, ValueError): # Missing, or an HTTP-date rather than seconds.
57+
return wait
58+
return max(wait, min(retry_after, API_RETRY_WAIT_MAX))
3259

3360

3461
class SessionTokenProvider(Protocol):
@@ -64,14 +91,28 @@ def _update_auth(self, response: httpx2.Response) -> None:
6491
# Runtime refresh headers must not replace the external workload identity.
6592
pass
6693

94+
# The server refetches signing keys at most every 10 seconds; the SDK's backoff alone often retries
95+
# every attempt inside that cooldown.
96+
@retry(
97+
retry=retry_if_exception(_should_retry_api_request),
98+
stop=stop_after_attempt(API_RETRIES),
99+
wait=_wait_for_retry_after,
100+
before_sleep=_log_and_trace_retry,
101+
reraise=True,
102+
)
103+
def _request_with_retry(self, *args, **kwargs):
104+
return self._request_without_retry(*args, **kwargs)
105+
67106

68107
class WorkloadIdentitySessionTokenProvider:
69108
"""
70109
Exchange deployment identity for a session JWT, retaining the JWT only in memory.
71110
72111
The identity supplier must return a current, renewable deployment credential. For example, it may
73-
read a projected service-account token that the platform rotates. The owning Dag processor client
74-
closes this provider; do not share it between clients.
112+
read a projected service-account token that the platform rotates. Any exception it raises is
113+
re-raised as ``WorkloadIdentityUnavailableError``, so a failed early renewal keeps the valid Job
114+
credential in use. The owning Dag processor client closes this provider; do not share it between
115+
clients.
75116
"""
76117

77118
def __init__(
@@ -93,7 +134,10 @@ def get_token(
93134
) -> str:
94135
if not force_refresh and self._token is not None and monotonic() < self._renew_at:
95136
return self._token
96-
identity = self._identity_token().strip()
137+
try:
138+
identity = self._identity_token().strip()
139+
except Exception as error:
140+
raise WorkloadIdentityUnavailableError("Unable to obtain the workload identity") from error
97141
if not identity:
98142
raise ValueError("Workload identity credential is empty")
99143
started_at = monotonic()

0 commit comments

Comments
 (0)