robertpofuk commented on code in PR #72262:
URL: https://github.com/apache/airflow/pull/72262#discussion_r4016321320
##########
providers/edge3/src/airflow/providers/edge3/worker_api/auth.py:
##########
@@ -30,25 +31,146 @@
InvalidSignatureError,
)
-from airflow.api_fastapi.auth.tokens import JWTValidator
-from airflow.providers.common.compat.sdk import conf
+from airflow.api_fastapi.auth.tokens import JWKS, JWTValidator
+from airflow.providers.common.compat.sdk import AirflowConfigException, conf
+
+if TYPE_CHECKING:
+ from collections.abc import Callable
log = logging.getLogger(__name__)
+class WorkerTokenAuthorization(TypedDict, total=False):
+ """
+ Result of authorizing an OIDC worker token beyond signature verification.
+
+ Returned by a ``[edge] jwt_verifier`` callable to answer "may this token
act
+ as an edge worker?". ``authorized`` must be ``True`` for the request to
+ proceed; a falsy result (or a raised exception) rejects it.
+ """
+
+ authorized: bool
+
+
+def _default_jwt_verifier(claims: dict) -> WorkerTokenAuthorization:
+ """Authorize any token that passed signature, issuer and audience
verification."""
+ return {"authorized": True}
+
+
+def _trusted_jwks_url() -> str:
+ """Return the configured trusted JWKS URL, or an empty string when
unset."""
+ return conf.get("edge", "trusted_jwks_url", fallback="") or ""
+
+
+def _jwt_algorithms() -> list[str]:
+ """Return the accepted signing algorithms for OIDC worker tokens."""
+ configured = conf.get("edge", "jwt_algorithm", fallback="RS256") or "RS256"
+ return [algorithm.strip() for algorithm in configured.split(",") if
algorithm.strip()]
+
+
+def _jwt_audience() -> str | None:
+ """Return the configured audience, or None to accept only tokens without
an ``aud`` claim."""
+ return conf.get("edge", "jwt_audience", fallback="") or None
+
+
+def _jwt_issuer() -> str | None:
+ """Return the expected issuer, or None to skip issuer verification when
left empty."""
+ return conf.get("edge", "jwt_issuer", fallback="") or None
+
+
+def _jwt_leeway() -> int:
+ """Return the clock-skew leeway (seconds) for OIDC worker tokens."""
+ return conf.getint("edge", "jwt_leeway", fallback=30)
+
+
+def _jwt_verifier() -> Callable[[dict], WorkerTokenAuthorization | None]:
+ """Return the configured worker-authorization callable, or the permissive
default."""
+ return conf.getimport("edge", "jwt_verifier", fallback=None) or
_default_jwt_verifier
+
+
+def _jwt_verifier_configured() -> bool:
+ """Return whether an explicit ``[edge] jwt_verifier`` is set."""
+ return bool(conf.get("edge", "jwt_verifier", fallback=""))
+
+
@cache
-def jwt_validator() -> JWTValidator:
+def _oidc_enabled() -> bool:
+ """
+ Return whether OIDC verification is enabled, decided once and cached.
+
+ The validator is also cached, so the mode must be read from a single place;
+ otherwise a request-time re-read could disagree with the cached validator
and
+ skip the ``method``-claim check for a shared-secret token.
+ """
+ return bool(_trusted_jwks_url())
+
+
+def _shared_secret_validator() -> JWTValidator:
+ """Build a validator for worker tokens signed with the shared ``[api_auth]
jwt_secret``."""
return JWTValidator(
secret_key=conf.get("api_auth", "jwt_secret"),
leeway=conf.getint("api_auth", "jwt_leeway", fallback=30),
audience="api",
)
+def _oidc_validator(jwks_url: str) -> JWTValidator:
+ """
+ Build a validator for worker tokens issued by a trusted OIDC provider.
+
+ Verifies the token signature against the provider JWKS and checks the
+ ``iss`` and (optionally) ``aud`` claims. Used when ``[edge]
trusted_jwks_url``
+ is configured, so workers can authenticate with tokens minted by an
+ external identity provider instead of the shared secret.
+
+ Rejects the configuration when issuer verification is skipped (empty
+ ``jwt_issuer``) without a ``jwt_verifier``: that combination would accept
any
+ token signed by a key in the JWKS. The validator is built lazily on the
first
+ request, so this surfaces as a rejected request (403) rather than a startup
+ failure.
+ """
+ if not _jwt_issuer() and not _jwt_verifier_configured():
+ raise AirflowConfigException(
+ "[edge] jwt_verifier must be set when trusted_jwks_url is
configured "
+ "without jwt_issuer, otherwise any token signed by the JWKS is
accepted."
+ )
+ return JWTValidator(
+ jwks=JWKS(url=jwks_url),
+ issuer=_jwt_issuer(),
+ audience=_jwt_audience(),
+ algorithm=_jwt_algorithms(),
+ required_claims=frozenset({"iat", "exp"}),
+ leeway=_jwt_leeway(),
+ )
+
+
+@cache
+def jwt_validator() -> JWTValidator:
+ if _oidc_enabled():
+ return _oidc_validator(_trusted_jwks_url())
+ return _shared_secret_validator()
+
+
def jwt_validate(authorization: str) -> dict:
return jwt_validator().validated_claims(authorization)
Review Comment:
Validator is now invoked async.
--
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]
