potiuk commented on code in PR #72262:
URL: https://github.com/apache/airflow/pull/72262#discussion_r4014286282
##########
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)
+def _check_worker_authorization(payload: dict) -> None:
Review Comment:
Nit, and explicitly *not* a request to validate anything: consider passing
`method` through to the verifier now, unused, so the signature doesn't have to
change later.
You future-proofed the *output* of this hook on purpose — returning a
mapping rather than a bool so `teams` and friends can be added without breaking
callers. The same argument applies to the input, and right now it's one-way:
the verifier receives `payload` only, so it can never see which endpoint the
token is being presented to. `jwt_token_authorization` has `method` in hand two
lines up.
Why that's worth more than it looks: `method` is the request path after
`/edge_worker/v1/`, and the worker routes are
`@worker_router.post("/{worker_name}")` — so the path carries the worker name.
An operator who wants their verifier to ask "does this token's `sub` match the
worker it's addressing?" currently can't. That's the scoping the shared-secret
path gets from its `method` claim and OIDC mode gives up. Handing the verifier
the path doesn't restore it — whether to enforce anything stays the operator's
call, in their own callable — it just makes it possible.
The timing is the actual argument. `jwt_verifier` is new in 4.4.0 and
unreleased, so nothing implements `Callable[[dict], ...]` against it yet. Once
deployments do, widening it is a breaking change for every one of them.
Something like:
```python
def _default_jwt_verifier(claims: dict, method: str) ->
WorkerTokenAuthorization:
return {"authorized": True}
```
with `_check_worker_authorization(method, payload)` forwarding it, costs
nothing today and is awkward to retrofit. Worth a line in `jwt_verifier`'s
`provider.yaml` description too — it currently says the callable "receives the
validated claims (``dict``)".
---
Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting
--
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]
