Skip to content

Commit 55f3d19

Browse files
committed
Enforce scoped Dag processor token claims and expiry
Task and processor credentials need distinct validation so task-only routes cannot receive processor identities. Expired provisioning credentials must not leave a registered Job behind, while trusted in-process task identities must remain usable without JWT expiry.
1 parent 142cc9d commit 55f3d19

16 files changed

Lines changed: 676 additions & 210 deletions

File tree

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

Lines changed: 41 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -348,15 +348,17 @@ The ``JWTBearer`` security dependency validates the token once per request:
348348
2. Performs cryptographic signature validation via ``JWTValidator``.
349349
3. Verifies standard claims (``exp``, ``iat``, ``aud`` — ``nbf`` and ``iss`` if configured).
350350
4. Defaults the ``scope`` claim to ``"execution"`` if absent.
351-
5. **Validates the task identity claims against a typed Pydantic schema.** ``TIClaims``
352-
(in ``airflow.api_fastapi.execution_api.datamodels.token``) enforces that ``scope`` is one
353-
of the declared ``TokenScope`` literals (``"execution"`` or ``"workload"``); ``TIToken``
354-
then parses the ``sub`` claim through a ``UUID`` field, which rejects non-UUID values.
355-
A token whose ``scope`` is unknown, or whose ``sub`` is not a valid UUID, is rejected with
356-
``403 Forbidden`` even when the cryptographic signature checks pass. ``TIClaims`` keeps
357-
``extra="allow"`` so auth managers can attach additional, deployment-specific claims
358-
without modifying the core schema; only the security-critical fields are typed.
359-
6. Creates a ``TIToken`` object with the task instance ID and validated claims.
351+
5. **Validates the identity and its scope-specific claims against typed Pydantic schemas.**
352+
``ExecutionToken`` (in ``airflow.api_fastapi.execution_api.datamodels.token``) parses ``sub``
353+
as a UUID. Its claims form a discriminated union selected by ``scope``: ``TIClaims`` for
354+
``"execution"``, ``"workload"``, and ``"callback"``; ``DagProcessorSessionClaims`` for
355+
provisioning; ``DagProcessorClaims`` for a registered Job; and ``DagParseClaims`` for a
356+
parsing attempt. Each processor claim type requires expiry and its scope-specific grants
357+
and identity fields. Unknown scopes, invalid subjects, or missing required claims return
358+
``403 Forbidden`` even when cryptographic validation succeeds. All claim types allow extra
359+
fields so auth managers can attach deployment-specific claims.
360+
6. Constructs the corresponding token object: ``TIToken``, ``DagProcessorSessionToken``,
361+
``DagProcessorToken``, or an ``ExecutionToken`` carrying ``DagParseClaims``.
360362
7. Caches the validated token on the ASGI request scope for the duration of the request.
361363

362364
Route-level enforcement is handled by ``require_auth``:
@@ -367,26 +369,32 @@ Route-level enforcement is handled by ``require_auth``:
367369
``{task_instance_id}`` path parameter, preventing a worker from accessing another task's
368370
endpoints.
369371

372+
The task-only, processor-session, and processor-Job dependencies additionally reject tokens of
373+
the wrong principal type. Endpoints shared by tasks and parsing attempts use the general
374+
``CurrentExecutionToken`` dependency with their declared route scopes.
375+
370376
.. mermaid::
371377

372378
flowchart TD
373379
REQ([Incoming request<br/>Authorization: Bearer ...])
374380
REQ --> CACHE{Cached on<br/>request.scope?}
375-
CACHE -->|yes| RET([Return cached TIToken])
381+
CACHE -->|yes| RET([Return cached token])
376382
CACHE -->|no| SIG[JWTValidator:<br/>verify signature]
377383
SIG -->|fail| F1([403 Forbidden])
378384
SIG -->|ok| STD[Verify exp / iat / nbf<br/>aud / iss]
379385
STD -->|fail| F1
380386
STD -->|ok| SCOPE[Default scope to<br/>'execution' if absent]
381-
SCOPE --> SCHEMA[TIClaims:<br/>typed Pydantic schema]
387+
SCOPE --> SCHEMA[ExecutionToken:<br/>UUID identity and<br/>scope-discriminated claims]
382388
SCHEMA -->|ValidationError| F1
383389
SCHEMA -->|ok| TYP{require_auth:<br/>scope in<br/>route.allowed_token_types?}
384390
TYP -->|no| F1
385391
TYP -->|yes| SELF{ti:self scope<br/>declared?}
386-
SELF -->|no| OK([Return TIToken])
392+
SELF -->|no| PRINCIPAL{Typed dependency:<br/>principal type matches?}
387393
SELF -->|yes| MATCH{token.sub ==<br/>task_instance_id?}
388394
MATCH -->|no| F1
389-
MATCH -->|yes| OK
395+
MATCH -->|yes| PRINCIPAL
396+
PRINCIPAL -->|no| F1
397+
PRINCIPAL -->|yes or general dependency| OK([Return validated token])
390398

391399
classDef fail fill:#ffcdd2,stroke:#c62828,stroke-width:2px,color:#000
392400
classDef pass fill:#c8e6c9,stroke:#2e7d32,stroke-width:2px,color:#000
@@ -489,9 +497,10 @@ guidance, and the planned strategic and tactical improvements.
489497
Dag processor HTTP client credentials
490498
-------------------------------------
491499

492-
The core ``DagProcessorAPIClient`` supports authenticated access to the Execution API
493-
without holding its signing key. This is the client and authentication foundation;
494-
``airflow dag-processor`` still uses its existing database-backed lifecycle.
500+
The core ``DagProcessorAPIClient`` lets processor integrations authenticate to the Execution
501+
API without holding its signing key. The standard ``airflow dag-processor`` command does
502+
not yet use this client. Provisioning a token file alone does not switch that command to
503+
HTTP or remove its need for database access.
495504

496505
On a trusted host with the API signing key, provision a session token for the bundles
497506
the client may access::
@@ -503,17 +512,24 @@ The command writes the token atomically with owner-only permissions. Keep the pr
503512
running and mount its directory rather than a single file so token replacement stays
504513
visible. Each processor process needs its own session.
505514

506-
The ``dag_processor_session`` token registers a Job through ``POST /jobs``. The returned
507-
``dag_processor`` token can heartbeat and complete that Job, read Connections and Variables
508-
for a granted bundle, and exchange a parsing credential through
509-
``POST /jobs/{job_id}/parse-token``. Registration retries keep the same registration UUID.
510-
A new process uses a new UUID; it cannot replace a session's Job while that Job is alive.
511-
Ending or replacing the Job invalidates its runtime credentials.
515+
The three credentials have different issuers and permissions:
516+
517+
- ``dag_processor_session`` is issued by trusted provisioning. It can only register a
518+
Job or renew that Job's credential through ``POST /jobs``.
519+
- ``dag_processor`` is returned by Job registration. It can heartbeat and complete that
520+
Job, read Connections and Variables for a granted bundle, and exchange a parsing
521+
credential through ``POST /jobs/{job_id}/parse-token``. It cannot outlive the session
522+
token used to obtain it.
523+
- ``dag_parse`` is returned by the parsing-token exchange for one attempt, bundle, and
524+
relative file location. It permits parse-time requests, including callback-context reads,
525+
but cannot manage Jobs or exchange tokens. It cannot outlive its parent Job credential.
526+
527+
Registration retries keep the same registration UUID. A new process uses a new UUID;
528+
it cannot replace a session's Job while that Job is alive. Ending or replacing the Job
529+
invalidates its runtime credentials.
512530

513531
The client selects the bundle with ``use_bundle`` and a parsing attempt with ``use_parse``.
514-
A ``dag_parse`` token binds the attempt UUID, bundle, and relative file location to the Job.
515-
It permits parse-time requests for that bundle, including callback-context reads, but cannot
516-
manage Jobs or exchange tokens. A request header cannot override the signed bundle.
532+
A request header cannot override the parsing token's signed bundle.
517533
Connection and Variable access uses the bundle's current team.
518534

519535
The client attempts renewal at 80% of token lifetime. Transient early-renewal failures leave

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -532,10 +532,11 @@ potentially still executes with direct database access in the Dag File Processor
532532
* **Tightened deserialization allowlist.** The serializer's class-name allowlist regex is
533533
anchored to require a full-string match, so attacker-controlled values cannot smuggle
534534
untrusted class names by appending an allowed suffix to a non-allowed name.
535-
* **Typed JWT claims schema.** Even after cryptographic validation, claims are run through a
536-
typed Pydantic schema (``TIClaims``) that enforces the ``scope`` literal, and then through
537-
``TIToken`` which parses ``sub`` as a UUID. A token whose ``scope`` is unknown or whose
538-
``sub`` is not a valid UUID is rejected before any route handler runs.
535+
* **Typed JWT claims schemas.** Even after cryptographic validation, ``ExecutionToken`` validates
536+
``sub`` as a UUID and selects a claims schema by ``scope``: ``TIClaims`` for task, workload,
537+
and callback credentials, or ``DagProcessorSessionClaims``, ``DagProcessorClaims``, or
538+
``DagParseClaims`` for processor credentials. Unknown scopes, invalid subjects, and missing
539+
required scope-specific claims are rejected before any route handler runs.
539540

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

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656

5757
__all__ = [
5858
"create_task_execution_api_app",
59+
"create_jwt_generator",
5960
"lifespan",
6061
"CorrelationIdMiddleware",
6162
]
@@ -77,7 +78,8 @@ def _jwt_validator():
7778
return validator
7879

7980

80-
def _jwt_generator():
81+
def create_jwt_generator() -> JWTGenerator:
82+
"""Create a signer using the Execution API's configured key, audience, and lifetime."""
8183
from airflow.configuration import conf
8284

8385
generator = JWTGenerator(
@@ -99,7 +101,7 @@ async def lifespan(app: FastAPI, registry: svcs.Registry):
99101
# record this here
100102
app.state.svcs_registry = registry
101103

102-
registry.register_factory(JWTGenerator, _jwt_generator)
104+
registry.register_factory(JWTGenerator, create_jwt_generator)
103105

104106
# InProcessExecutionAPI stubs out JWTValidator: don't re-register in that case.
105107
if JWTValidator not in registry:

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

Lines changed: 60 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -15,32 +15,34 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717
"""
18-
Issue Dag processor session tokens.
18+
Issue Dag processor session, Job, and parsing tokens.
1919
20-
Only trusted provisioning runs this: it needs the Execution API signing key, which a Dag processor
21-
must never hold, because the processor runs the code it parses.
20+
Only trusted provisioning and the API server issue these tokens. A Dag processor must never hold
21+
the signing key, because the processor runs the code it parses.
2222
"""
2323

2424
from __future__ import annotations
2525

26-
import contextlib
27-
import os
28-
import tempfile
2926
from typing import TYPE_CHECKING
3027

31-
from airflow.api_fastapi.execution_api.app import _jwt_generator
28+
from airflow._shared.timezones import timezone
3229

3330
if TYPE_CHECKING:
3431
from collections.abc import Collection
35-
from pathlib import Path
3632
from uuid import UUID
3733

34+
from airflow.api_fastapi.auth.tokens import JWTGenerator
35+
36+
37+
class ExpiredDagProcessorToken(ValueError):
38+
"""The parent credential expired before a Job or parsing credential could be issued."""
39+
3840

3941
def generate_dag_processor_session_token(
40-
*, session_id: UUID, bundle_names: Collection[str], valid_for: float
42+
generator: JWTGenerator, *, session_id: UUID, bundle_names: Collection[str], valid_for: float
4143
) -> str:
4244
"""Return a ``dag_processor_session`` token, which a Dag processor exchanges for a Job token."""
43-
return _jwt_generator().generate(
45+
return generator.generate(
4446
extras={
4547
"sub": str(session_id),
4648
"scope": "dag_processor_session",
@@ -50,16 +52,51 @@ def generate_dag_processor_session_token(
5052
)
5153

5254

53-
def write_token_file(path: Path, token: str) -> None:
54-
"""Replace the token file in one step, so a processor rereading it never sees a partial token."""
55-
fd, tmp_path = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.")
56-
try:
57-
with os.fdopen(fd, "w") as tmp:
58-
tmp.write(token)
59-
tmp.flush()
60-
os.fsync(tmp.fileno())
61-
os.replace(tmp_path, path)
62-
except BaseException:
63-
with contextlib.suppress(FileNotFoundError):
64-
os.unlink(tmp_path)
65-
raise
55+
def generate_dag_processor_token(
56+
generator: JWTGenerator,
57+
*,
58+
session_id: UUID,
59+
job_id: int,
60+
bundle_names: Collection[str],
61+
session_expiry: float,
62+
) -> str:
63+
"""Issue a Job credential that cannot outlive its provisioned session credential."""
64+
remaining = session_expiry - timezone.utcnow().timestamp()
65+
if remaining <= 0:
66+
raise ExpiredDagProcessorToken("Session credential has expired")
67+
return generator.generate(
68+
extras={
69+
"sub": str(session_id),
70+
"scope": "dag_processor",
71+
"dag_bundles": sorted(bundle_names),
72+
"job_id": job_id,
73+
},
74+
valid_for=min(generator.valid_for, remaining),
75+
)
76+
77+
78+
def generate_dag_parse_token(
79+
generator: JWTGenerator,
80+
*,
81+
session_id: UUID,
82+
job_id: int,
83+
attempt_id: UUID,
84+
bundle_name: str,
85+
relative_fileloc: str,
86+
processor_expiry: float,
87+
) -> str:
88+
"""Issue a parsing credential that cannot outlive its parent Job credential."""
89+
remaining = processor_expiry - timezone.utcnow().timestamp()
90+
if remaining <= 0:
91+
raise ExpiredDagProcessorToken("Processor credential has expired")
92+
return generator.generate(
93+
extras={
94+
"sub": str(attempt_id),
95+
"scope": "dag_parse",
96+
"session_id": str(session_id),
97+
"job_id": job_id,
98+
"dag_bundles": [bundle_name],
99+
"relative_fileloc": relative_fileloc,
100+
},
101+
valid_for=min(generator.valid_for, remaining),
102+
)

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

Lines changed: 66 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -20,48 +20,70 @@
2020
from typing import Annotated, Literal
2121
from uuid import UUID
2222

23-
from pydantic import ConfigDict, Field, model_validator
23+
from pydantic import ConfigDict, Field
2424

2525
from airflow.api_fastapi.core_api.base import BaseModel
26-
from airflow.typing_compat import Self
2726

28-
TokenScope = Literal[
29-
"execution", "workload", "callback", "dag_processor_session", "dag_processor", "dag_parse"
30-
]
27+
TaskTokenScope = Literal["execution", "workload", "callback"]
28+
DagProcessorSessionScope = Literal["dag_processor_session"]
29+
DagProcessorScope = Literal["dag_processor"]
30+
DagParseScope = Literal["dag_parse"]
31+
TokenScope = Literal[TaskTokenScope, DagProcessorSessionScope, DagProcessorScope, DagParseScope]
3132

3233

33-
class ExecutionClaims(BaseModel):
34+
class _JWTClaims(BaseModel):
3435
"""
3536
Validated JWT claims for an Execution API principal.
3637
37-
JWTValidator validates exp/iat/nbf/aud before these claims are constructed.
38+
JWTValidator checks standard JWT claims before constructing claims for HTTP requests.
3839
Extra claims are allowed for compatibility with existing task tokens.
3940
"""
4041

4142
model_config = ConfigDict(extra="allow")
4243

43-
scope: TokenScope = "execution"
44+
45+
class TIClaims(_JWTClaims):
46+
"""Claims for task execution, workload exchange, or callback execution."""
47+
48+
scope: TaskTokenScope = "execution"
49+
# Trusted in-process callers construct task identities without a JWT or expiry.
4450
exp: float | None = None
45-
dag_bundles: frozenset[Annotated[str, Field(min_length=1)]] | None = None
46-
"""Dag bundles a Dag processor token may act for."""
47-
job_id: int | None = None
48-
"""Job a ``dag_processor`` token was issued for when that Job registered."""
49-
session_id: UUID | None = None
50-
"""Processor session that owns a parsing attempt, whose own identity is the token subject."""
51-
relative_fileloc: str | None = Field(default=None, min_length=1, max_length=2000)
52-
"""Bundle-relative file being parsed; an archive is one file under the current processor model."""
53-
54-
@model_validator(mode="after")
55-
def validate_dag_processor_claims(self) -> Self:
56-
if self.scope in ("dag_processor_session", "dag_processor", "dag_parse") and not self.dag_bundles:
57-
raise ValueError(f"A {self.scope} token must grant at least one Dag bundle")
58-
if self.scope in ("dag_processor", "dag_parse") and self.job_id is None:
59-
raise ValueError(f"A {self.scope} token must name the Job it was issued for")
60-
if self.scope == "dag_parse" and (
61-
len(self.dag_bundles or ()) != 1 or self.session_id is None or self.relative_fileloc is None
62-
):
63-
raise ValueError("A dag_parse token must name one bundle, its file, and its processor session")
64-
return self
51+
52+
53+
DagBundleGrant = Annotated[frozenset[Annotated[str, Field(min_length=1)]], Field(min_length=1)]
54+
55+
56+
class DagProcessorSessionClaims(_JWTClaims):
57+
"""Provisioned bundle grants used only to register a Job or renew its credential."""
58+
59+
scope: DagProcessorSessionScope = "dag_processor_session"
60+
exp: float
61+
dag_bundles: DagBundleGrant
62+
63+
64+
class DagProcessorClaims(_JWTClaims):
65+
"""Claims for managing one registered Job and exchanging its parsing credentials."""
66+
67+
scope: DagProcessorScope = "dag_processor"
68+
exp: float
69+
dag_bundles: DagBundleGrant
70+
job_id: int
71+
72+
73+
class DagParseClaims(_JWTClaims):
74+
"""Claims for one file-parsing attempt within a registered Job's bundle."""
75+
76+
scope: DagParseScope = "dag_parse"
77+
exp: float
78+
dag_bundles: DagBundleGrant = Field(max_length=1)
79+
job_id: int
80+
session_id: UUID
81+
relative_fileloc: str = Field(min_length=1, max_length=2000)
82+
83+
84+
ExecutionClaims = Annotated[
85+
TIClaims | DagProcessorSessionClaims | DagProcessorClaims | DagParseClaims, Field(discriminator="scope")
86+
]
6587

6688

6789
class ExecutionToken(BaseModel):
@@ -71,6 +93,19 @@ class ExecutionToken(BaseModel):
7193
claims: ExecutionClaims
7294

7395

74-
# Preserve imports for task-specific consumers while shared endpoints use the general principal.
75-
TIClaims = ExecutionClaims
76-
TIToken = ExecutionToken
96+
class TIToken(ExecutionToken):
97+
"""Authenticated task, workload, or callback identity."""
98+
99+
claims: TIClaims
100+
101+
102+
class DagProcessorSessionToken(ExecutionToken):
103+
"""Authenticated provisioned Dag processor session."""
104+
105+
claims: DagProcessorSessionClaims
106+
107+
108+
class DagProcessorToken(ExecutionToken):
109+
"""Authenticated Dag processor Job."""
110+
111+
claims: DagProcessorClaims

0 commit comments

Comments
 (0)