Skip to content

Commit 1b68b46

Browse files
committed
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.
1 parent a73ffae commit 1b68b46

19 files changed

Lines changed: 1665 additions & 37 deletions

File tree

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

Lines changed: 126 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -519,6 +519,125 @@ API without holding its signing key. The standard ``airflow dag-processor`` comm
519519
not yet use this client. Provisioning a token file alone does not switch that command to
520520
HTTP or remove its need for database access.
521521

522+
In-memory sessions from workload identity
523+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
524+
525+
Processor integrations can use ``WorkloadIdentitySessionTokenProvider`` to obtain and renew
526+
``dag_processor_session`` credentials through ``POST /execution/dag-processor/session-token``.
527+
The provider retains these credentials only in memory. It requires a renewable, signed workload
528+
identity from the deployment platform; Airflow does not create that identity or require the
529+
processor to hold an Airflow signing key.
530+
531+
Configure the API server's ``[execution_api] dag_processor_workload_identity`` option with JSON::
532+
533+
{
534+
"issuer": "https://identity.example/",
535+
"jwks_url": "https://identity.example/openid/v1/jwks",
536+
"audience": "airflow-dag-processor-bootstrap",
537+
"algorithm": "RS256",
538+
"instance_id_claim": ["kubernetes.io", "pod", "uid"],
539+
"grants": {
540+
"system:serviceaccount:airflow:processor": ["dags-folder"]
541+
}
542+
}
543+
544+
Use your identity provider's exact issuer, HTTPS JWKS URL, and signing algorithm (``RS256``,
545+
``ES256``, or ``EdDSA``). Use a dedicated audience, distinct from ``[execution_api] jwt_audience``.
546+
The API server verifies the signature, issuer, audience, issued-at time, expiry, and any not-before
547+
claim, using ``[api_auth] jwt_leeway`` for clock skew. Only exact subjects listed in ``grants`` can
548+
obtain a session, and every granted bundle must be configured on the API server. This option is
549+
empty by default, disabling the exchange.
550+
Apply the same trust configuration to all API-server replicas and restart them after changes.
551+
552+
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.
556+
Authenticated JWKS endpoints, including Kubernetes API-server discovery protected by service-account
557+
authentication, are not supported by this implementation. For those clusters, serve the public JWKS
558+
through a suitable HTTPS endpoint before enabling the exchange. See Kubernetes'
559+
`service account issuer discovery documentation <https://kubernetes.io/docs/tasks/configure-pod-container/configure-service-account/#service-account-issuer-discovery>`_.
560+
The processor-to-Airflow TLS settings described below do not configure this outbound JWKS connection.
561+
562+
``instance_id_claim`` is a path through signed claims to a nonempty string identifying one
563+
processor replica. A shared service-account subject alone is insufficient: replicas must have
564+
different instance identities. Run one processor per identity. The session ID stays stable across
565+
identity-token renewal and API-server restarts. Changes to the issuer, audience, subject, instance
566+
claim, instance value, or bundle grants produce a different session. Request bodies and caller
567+
hostnames cannot select the session or its grants.
568+
569+
Changing grants while API-server replicas serve different configurations can repeatedly retire
570+
processor registrations, because those replicas issue different session IDs for the same identity.
571+
Stop the affected processors, finish deploying identical trust configuration to every API-server
572+
replica, and then restart the processors.
573+
574+
For Kubernetes, the example uses the Pod UID in a
575+
`bound service-account token <https://kubernetes.io/docs/reference/access-authn-authz/service-accounts-admin/#bound-service-account-token-volume-mechanism>`_.
576+
Project a token with the dedicated audience and mount the projection directory read-only in the
577+
processor container. Configure the projection's ``expirationSeconds`` and allow Kubernetes to
578+
rotate it. See `serviceAccountToken projected volumes
579+
<https://kubernetes.io/docs/concepts/storage/projected-volumes/#serviceaccounttoken>`_.
580+
581+
A container restart within the same Pod keeps the Pod UID and therefore the session ID.
582+
A replacement process cannot register while the previous Job is still alive. The client raises
583+
``DagProcessorJobAlreadyRunning``; the integration must retry registration with backoff until the
584+
previous Job completes or its heartbeat becomes stale under ``[dag_processor] health_check_threshold``.
585+
The client does not wait automatically. This prevents two processes sharing an identity from owning
586+
active Jobs at the same time.
587+
588+
The processor integration supplies a callable that reads the current platform identity on each
589+
exchange. For example:
590+
591+
.. code-block:: python
592+
593+
from pathlib import Path
594+
595+
from airflow.dag_processing.api_client import DagProcessorAPIClient
596+
from airflow.dag_processing.credentials import WorkloadIdentitySessionTokenProvider
597+
598+
identity_path = Path("/var/run/airflow-identity/token")
599+
provider = WorkloadIdentitySessionTokenProvider(
600+
base_url="https://airflow.example/execution/",
601+
identity_token=identity_path.read_text,
602+
)
603+
processor = DagProcessorAPIClient(
604+
base_url="https://airflow.example/execution/",
605+
token_provider=provider,
606+
hostname="processor-1",
607+
)
608+
609+
Use this client for the integration's Job lifecycle and close it during shutdown; closing it also
610+
closes the provider and clears its session cache. The file in this example contains the platform's
611+
projected identity, not an Airflow session credential. The provider never writes the exchanged
612+
session JWT to disk and does not fall back to a session-token file. ``token_provider`` and
613+
``token_file`` are mutually exclusive. HTTPS is required for the exchange, and redirects are rejected.
614+
The exchange uses the same API TLS configuration and request timeout as the Task SDK client,
615+
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.
618+
619+
Session expiry is capped by both the external identity's remaining lifetime and
620+
``[execution_api] jwt_expiration_time``. Clock-skew leeway never extends this cap, and an identity
621+
with less than one second remaining cannot obtain a session. Renewal rereads the platform identity
622+
and retains the same session and Job. During an identity-provider or API outage, the client can continue using
623+
an already-issued Job credential until it expires; it cannot renew without a valid identity.
624+
Removing a grant stops future exchanges after API-server configuration is reloaded. Credentials
625+
already issued remain subject to their expiry and Job ownership checks.
626+
627+
Verification uses the issuer's JWKS and does not perform online token introspection. In particular,
628+
it does not check whether a Kubernetes Pod still exists. A previously issued identity may remain
629+
usable until expiry. The JWKS cache refreshes hourly and also refreshes when an identity names an
630+
unknown signing key. Fetches, including failed attempts, are limited to one per ten seconds per
631+
API-server process. Publish new keys before signing with them and retain old keys while their tokens
632+
remain valid.
633+
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.
637+
638+
File-based provisioning
639+
~~~~~~~~~~~~~~~~~~~~~~~
640+
522641
On a trusted host with the API signing key, provision a session token for the bundles
523642
the client may access::
524643

@@ -532,9 +651,12 @@ token replacement stays visible. Each processor process needs its own session. R
532651
provisioner with the same token file and bundles keeps that session, so running processors keep
533652
their Jobs.
534653

654+
Credential permissions and renewal
655+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
656+
535657
The three credentials have different issuers and permissions:
536658

537-
- ``dag_processor_session`` is issued by trusted provisioning. It can only register a
659+
- ``dag_processor_session`` is issued by the identity exchange or trusted file provisioning. It can only register a
538660
Job or renew that Job's credential through ``POST /jobs``.
539661
- ``dag_processor`` is returned by Job registration. It can heartbeat and complete that
540662
Job, read Connections and Variables for a granted bundle, and exchange a parsing
@@ -625,6 +747,9 @@ All JWT-related configuration parameters:
625747
* - ``[execution_api] jwt_audience``
626748
- ``urn:airflow.apache.org:task``
627749
- Audience claim for Execution API tokens.
750+
* - ``[execution_api] dag_processor_workload_identity``
751+
- Empty (disabled)
752+
- JSON trust settings and subject-to-bundle grants for workload identity session provisioning.
628753

629754
.. important::
630755

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -537,6 +537,12 @@ potentially still executes with direct database access in the Dag File Processor
537537
and callback credentials, or ``DagProcessorSessionClaims``, ``DagProcessorClaims``, or
538538
``DagParseClaims`` for processor credentials. Unknown scopes, invalid subjects, and missing
539539
required scope-specific claims are rejected before any route handler runs.
540+
* **Workload identity provisioning.** The optional Dag processor session exchange validates
541+
deployment-issued JWTs against separately configured issuer, audience, and JWKS settings.
542+
Bundle grants come from the API server configuration. Exchanged sessions stay in the client's
543+
memory, but workload identity and runtime credentials remain accessible to code with access
544+
to that process. This does not add isolation between the processor and user code; see
545+
:doc:`jwt_token_authentication` for configuration and expiry limitations.
540546

541547
**Token signing key might be a shared secret**
542548
In symmetric key mode (``[api_auth] jwt_secret``), the same secret key is used to both generate and

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

Lines changed: 65 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,14 @@
1616
# under the License.
1717
from __future__ import annotations
1818

19+
import asyncio
1920
import json
2021
import os
2122
import time
2223
import uuid
2324
from base64 import urlsafe_b64encode
2425
from collections.abc import Callable, Sequence
26+
from contextlib import nullcontext
2527
from datetime import datetime
2628
from typing import TYPE_CHECKING, Any, Literal, overload
2729

@@ -45,6 +47,7 @@
4547
__all__ = [
4648
"InvalidClaimError",
4749
"JWKS",
50+
"JWKSUnavailableError",
4851
"JWTGenerator",
4952
"JWTValidator",
5053
"generate_private_key",
@@ -99,9 +102,19 @@ def _guess_best_algorithm(key: AllowedPrivateKeys):
99102
raise ValueError(f"Unknown key object {type(key)}")
100103

101104

105+
class JWKSUnavailableError(RuntimeError):
106+
"""The required signing key could not be loaded because a JWKS fetch failed."""
107+
108+
102109
@attrs.define(repr=False)
103110
class JWKS:
104-
"""A class to fetch and sync a set of JSON Web Keys."""
111+
"""
112+
Fetch and cache a set of JSON Web Keys.
113+
114+
With ``raise_on_fetch_error``, missing keys raise ``JWKSUnavailableError`` after a failed
115+
fetch. Cached matching keys remain usable, and a successfully fetched set with no matching key
116+
still raises ``KeyError``.
117+
"""
105118

106119
url: str
107120
fetched_at: float = 0
@@ -113,6 +126,10 @@ class JWKS:
113126
refresh_jwks: bool = True
114127
refresh_interval_secs: int = 3600
115128
refresh_retry_interval_secs: int = 10
129+
refresh_on_unknown_kid: bool = False
130+
raise_on_fetch_error: bool = False
131+
_fetch_lock: asyncio.Lock = attrs.field(factory=asyncio.Lock, init=False)
132+
_fetch_failed: bool = attrs.field(default=False, init=False)
116133

117134
def __repr__(self) -> str:
118135
return f"JWKS(url={self.url}, fetched_at={self.fetched_at})"
@@ -128,19 +145,30 @@ def from_private_key(cls, *keys: AllowedPrivateKeys | tuple[AllowedPrivateKeys,
128145
obj._jwks = jwt.PyJWKSet(keyset)
129146
return obj
130147

131-
async def fetch_jwks(self) -> None:
132-
if not self._should_fetch_jwks():
133-
return
134-
if self.url.startswith("http"):
135-
data = await self._fetch_remote_jwks()
136-
else:
137-
data = self._fetch_local_jwks()
138-
139-
if not data:
140-
return
141-
142-
self._jwks = jwt.PyJWKSet.from_dict(data)
143-
log.debug("Fetched JWKS", url=self.url, keys=len(self._jwks.keys))
148+
async def fetch_jwks(self, *, force: bool = False) -> None:
149+
lock = self._fetch_lock if self.refresh_on_unknown_kid else nullcontext()
150+
async with lock:
151+
if not self._should_fetch_jwks(force=force):
152+
return
153+
self._fetch_failed = True
154+
if self.url.startswith("http"):
155+
data = await self._fetch_remote_jwks()
156+
else:
157+
data = self._fetch_local_jwks()
158+
159+
if not data:
160+
return
161+
162+
try:
163+
keyset = jwt.PyJWKSet.from_dict(data)
164+
except (jwt.PyJWKError, jwt.PyJWKSetError, KeyError, TypeError, ValueError):
165+
if not self.raise_on_fetch_error:
166+
raise
167+
log.warning("Unable to load JWKS response", url=self.url, exc_info=True)
168+
return
169+
self._jwks = keyset
170+
self._fetch_failed = False
171+
log.debug("Fetched JWKS", url=self.url, keys=len(self._jwks.keys))
144172

145173
async def _fetch_remote_jwks(self) -> dict[str, Any] | None:
146174
try:
@@ -151,7 +179,7 @@ async def _fetch_remote_jwks(self) -> dict[str, Any] | None:
151179
)
152180
if TYPE_CHECKING:
153181
assert self.url
154-
self.last_fetch_attempt_at = int(time.monotonic())
182+
self.last_fetch_attempt_at = time.monotonic()
155183
response = await self.client.get(self.url)
156184
response.raise_for_status()
157185
self.fetched_at = int(time.monotonic())
@@ -172,7 +200,7 @@ def _fetch_local_jwks(self) -> dict[str, Any] | None:
172200
log.exception("Failed to read local JWKS", url=self.url)
173201
return None
174202

175-
def _should_fetch_jwks(self) -> bool:
203+
def _should_fetch_jwks(self, *, force: bool = False) -> bool:
176204
"""
177205
Check if we need to fetch the JWKS based on the last fetch time and the refresh interval.
178206
@@ -188,8 +216,16 @@ def _should_fetch_jwks(self) -> bool:
188216
# For remote fetches we check if the JWKS is not loaded (fetched_at = 0) or if the last fetch was more than
189217
# refresh_interval_secs ago and the last fetch attempt was more than refresh_retry_interval_secs ago
190218
now = time.monotonic()
219+
# Unknown keys are untrusted input; bound fetches even for an empty cache or a failed refresh.
220+
if (
221+
self.refresh_on_unknown_kid
222+
and self.last_fetch_attempt_at
223+
and now - self.last_fetch_attempt_at < self.refresh_retry_interval_secs
224+
):
225+
return False
191226
return self.refresh_jwks and (
192-
not self._jwks
227+
force
228+
or not self._jwks
193229
or (
194230
self.fetched_at == 0
195231
or (
@@ -204,9 +240,19 @@ async def get_key(self, kid: str) -> jwt.PyJWK:
204240
await self.fetch_jwks()
205241

206242
if self._jwks:
207-
return self._jwks[kid]
243+
try:
244+
return self._jwks[kid]
245+
except KeyError:
246+
if self.refresh_on_unknown_kid:
247+
await self.fetch_jwks(force=True)
248+
try:
249+
return self._jwks[kid]
250+
except KeyError:
251+
pass
252+
253+
if self.raise_on_fetch_error and self._fetch_failed:
254+
raise JWKSUnavailableError("Unable to load the required signing key")
208255

209-
# It didn't load!
210256
raise KeyError(f"Key ID {kid} not found in keyset")
211257

212258
def status(self):

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
get_signing_args,
4242
)
4343
from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI
44+
from airflow.api_fastapi.execution_api.workload_identity import create_dag_processor_identity_verifier
4445

4546
if TYPE_CHECKING:
4647
from airflow.models.dagbag import DBDagBag
@@ -105,7 +106,9 @@ async def lifespan(app: FastAPI, registry: svcs.Registry):
105106
# Create an app scoped validator, so that we don't have to fetch it every time
106107
registry.register_value(JWTValidator, _jwt_validator(), ping=JWTValidator.status)
107108

108-
yield
109+
async with create_dag_processor_identity_verifier() as verifier:
110+
app.state.dag_processor_identity_verifier = verifier
111+
yield
109112

110113

111114
class CorrelationIdMiddleware(BaseHTTPMiddleware):
@@ -139,6 +142,9 @@ class JWTReissueMiddleware(BaseHTTPMiddleware):
139142
async def dispatch(self, request: Request, call_next):
140143
response: Response = await call_next(request)
141144

145+
if getattr(request.state, "skip_execution_token_reissue", False):
146+
return response
147+
142148
refreshed_token: str | None = None
143149
auth_header = request.headers.get("authorization")
144150
if auth_header and auth_header.lower().startswith("bearer "):

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -43,10 +43,10 @@ class IssuedToken(NamedTuple):
4343

4444

4545
class ExpiredDagProcessorToken(ValueError):
46-
"""The parent credential expired before a Job or parsing credential could be issued."""
46+
"""The parent credential expired before a session, Job, or parsing credential could be issued."""
4747

4848

49-
def _get_lifetime(generator: JWTGenerator, parent_expiry: float, message: str) -> int:
49+
def get_token_lifetime(generator: JWTGenerator, parent_expiry: float, message: str) -> int:
5050
"""Return whole seconds a child token may live: the signer's lifetime, capped by its parent's expiry."""
5151
lifetime = math.floor(min(generator.valid_for, parent_expiry - timezone.utcnow().timestamp()))
5252
if lifetime < 1:
@@ -77,7 +77,7 @@ def generate_dag_processor_token(
7777
session_expiry: float,
7878
) -> IssuedToken:
7979
"""Issue a Job credential that cannot outlive its provisioned session credential."""
80-
lifetime = _get_lifetime(generator, session_expiry, "Session credential has expired")
80+
lifetime = get_token_lifetime(generator, session_expiry, "Session credential has expired")
8181
token = generator.generate(
8282
extras={
8383
"sub": str(session_id),
@@ -101,7 +101,7 @@ def generate_dag_parse_token(
101101
processor_expiry: float,
102102
) -> IssuedToken:
103103
"""Issue a parsing credential that cannot outlive its parent Job credential."""
104-
lifetime = _get_lifetime(generator, processor_expiry, "Processor credential has expired")
104+
lifetime = get_token_lifetime(generator, processor_expiry, "Processor credential has expired")
105105
token = generator.generate(
106106
extras={
107107
"sub": str(attempt_id),

‎airflow-core/src/airflow/api_fastapi/execution_api/datamodels/job.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,13 @@ class JobRegisterResponse(StrictBaseModel):
6060
expires_in: int = Field(gt=0, description="Seconds until the token expires.")
6161

6262

63+
class DagProcessorSessionTokenResponse(StrictBaseModel):
64+
"""A session credential issued after workload identity verification."""
65+
66+
token: str = Field(min_length=1, repr=False)
67+
expires_in: int = Field(gt=0, description="Seconds until the session credential expires.")
68+
69+
6370
class JobHeartbeatResponse(StrictBaseModel):
6471
"""Current state of the Job; ``restarting`` asks the processor to stop."""
6572

0 commit comments

Comments
 (0)