ephraimbuddy commented on code in PR #74700: URL: https://github.com/apache/airflow/pull/74700#discussion_r4361086528
########## airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py: ########## @@ -0,0 +1,145 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Authenticate deployment identities for Dag processor session provisioning.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Annotated, Any, Literal +from uuid import NAMESPACE_URL, UUID, uuid5 + +import httpx +import jwt +from pydantic import BaseModel, ConfigDict, Field, HttpUrl, field_validator + +from airflow.api_fastapi.auth.tokens import JWKS, JWTValidator +from airflow.api_fastapi.execution_api.datamodels.token import DagBundleGrant +from airflow.configuration import conf +from airflow.dag_processing.bundles.manager import get_configured_bundle_names + +NonemptyString = Annotated[str, Field(min_length=1)] + + +class WorkloadIdentityConfig(BaseModel): + """Trust settings owned by the deployment manager, never supplied by a request.""" + + model_config = ConfigDict(extra="forbid") + + issuer: NonemptyString + jwks_url: HttpUrl + audience: NonemptyString + algorithm: Literal["RS256", "ES256", "EdDSA"] = "RS256" + instance_id_claim: list[NonemptyString] = Field(min_length=1) + grants: dict[NonemptyString, DagBundleGrant] = Field(min_length=1) + + @field_validator("jwks_url") + @classmethod + def require_https(cls, value: HttpUrl) -> HttpUrl: + if value.scheme != "https" or value.username or value.password or value.fragment: + raise ValueError( + "The workload identity JWKS URL must use HTTPS without credentials or a fragment" + ) + return value + + +class WorkloadIdentityNotAllowed(PermissionError): + """The verified identity has no Dag processor grant.""" + + +@dataclass(frozen=True) +class DagProcessorIdentity: + """Verified identity with server-selected grants and a stable, non-secret session ID.""" + + session_id: UUID + bundle_names: frozenset[str] + expires_at: float + + +class _IdentityClaims(BaseModel): + sub: NonemptyString + exp: float = Field(allow_inf_nan=False) + + +class DagProcessorIdentityVerifier: + """Validate external JWTs independently of Airflow's runtime token trust.""" + + def __init__(self, config: WorkloadIdentityConfig, *, client: httpx.AsyncClient): + if config.audience in conf.getlist("execution_api", "jwt_audience"): + raise ValueError("Workload identity must use a dedicated audience") + if unknown := frozenset().union(*config.grants.values()) - get_configured_bundle_names(): Review Comment: Not intended; that blast radius is too wide. A grant naming an unconfigured bundle no longer stops startup. The verifier logs an error naming both settings and refuses only that subject's exchanges with a 403. I went with 403 rather than 503 so processors don't spend their retry budget on an error only a config change can fix. Malformed trust config (the JSON shape, or reusing the runtime audience) still fails startup, since it is wrong from the moment it's written. Covered by `test_unconfigured_bundle_grant_refuses_only_its_subject`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
