potiuk commented on code in PR #72262:
URL: https://github.com/apache/airflow/pull/72262#discussion_r4014102509
##########
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]:
Review Comment:
Nit: this resolves the import path on every request, while everything around
it is cached.
`conf.getimport` is a config read plus an `import_string` per call, and
`_check_worker_authorization` runs on every worker request — job polling and
heartbeats included. `_oidc_enabled()` and `jwt_validator()` are both
`@cache`d; this is the one piece of OIDC config re-read each time.
It also leaves a gap in the fail-closed check. `_oidc_validator` only asks
whether the option is non-empty, so a typo'd path passes construction and then
fails per-request — `test_empty_issuer_with_verifier_is_allowed` configures
`my_company.edge_auth.verify_worker_token`, which doesn't exist, and the test
passes. The operator sees a generic 403 with a fresh error id on every request
rather than one clear error when the validator is built. (It does fail *closed*
— `getimport` raises `AirflowConfigException` and the catch-all in
`jwt_token_authorization` turns it into a 403 — so this is about diagnosability
and cost, not safety.)
Resolving the callable inside `_oidc_validator` and storing it next to the
validator fixes both: one import at construction time, and a bad path surfaces
immediately with `getimport`'s own message naming the key and section.
##########
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:
Nit, but worth doing before this meets a real IdP: OIDC mode makes this sync
`validated_claims()` do network I/O, and every other JWKS caller in the tree
uses the async variant.
`validated_claims()` is `async_to_sync(avalidated_claims)`. On the
shared-secret path that's free — `_get_validation_key` returns the secret
without awaiting anything. With `jwks=` it now awaits `httpx.AsyncClient.get`
against the provider, on the `JWKS` object held by the `@cache`d validator,
from whatever event loop `async_to_sync` creates. Starlette runs sync
dependencies through `anyio.to_thread`, which asgiref doesn't recognise as its
own `sync_to_async` context, so it builds a fresh loop per call rather than
reusing one. A single long-lived `httpx.AsyncClient` used across different
event loops is the usual way to get "Event loop is closed" / "attached to a
different loop".
It won't show up in testing, and it won't show up on day one either.
`_should_fetch_jwks()` is true on first use and then roughly hourly, so the
first fetch succeeds on whichever loop got there first and the problem only
appears at a later refresh. And it fails quietly: `_fetch_remote_jwks` catches
`Exception`, logs, and returns `None`, so `fetch_jwks` leaves `_jwks` at the
stale keyset. Workers keep authenticating on cached keys until the IdP rotates,
and then every worker starts getting 403s from `get_key(kid)` raising
`KeyError`, with the actual cause sitting hours earlier in the log as "Failed
to fetch remote JWKS".
The rest of the codebase avoids this by awaiting:
-
`airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py:160`
- `airflow-core/src/airflow/api_fastapi/execution_api/security.py:126`
- `airflow-core/src/airflow/api_fastapi/execution_api/app.py:149`
- `airflow-core/src/airflow/utils/serve_logs/log_server.py:67`
`jwt_token_authorization_rest` is only ever used as a FastAPI `Depends(...)`
(`routes/worker.py`, `routes/jobs.py`, `routes/logs.py`), and
`jwt_token_authorization` has no caller outside this module — so making both
`async def` and awaiting `avalidated_claims` here stays contained to this file.
##########
providers/edge3/tests/unit/edge3/worker_api/test_auth.py:
##########
@@ -150,3 +155,220 @@ def
test_strips_edge_worker_v1_prefix_and_falls_back_to_full_path(
jwt_token_authorization_rest(request, authorization="some-token")
mock_jwt_token_authorization.assert_called_once_with(expected_method,
"some-token")
+
+
+class TestOidcConfig:
+ """Config readers translate raw ``[edge]`` options into validator
inputs."""
+
+ @pytest.mark.parametrize(
+ ("configured", "expected"),
+ [
+ pytest.param(None, ["RS256"], id="default-rs256"),
+ pytest.param("RS512", ["RS512"], id="single-algorithm"),
+ pytest.param("RS256, RS512", ["RS256", "RS512"],
id="comma-list-is-split-and-stripped"),
+ ],
+ )
+ def test_oidc_algorithms_parsing(self, configured, expected):
+ """Unset falls back to RS256; a configured value is split on commas
and stripped."""
+ overrides = {} if configured is None else {("edge", "jwt_algorithm"):
configured}
+ with conf_vars(overrides):
+ assert auth._jwt_algorithms() == expected
+
+ @pytest.mark.parametrize(
+ ("configured", "expected"),
+ [
+ pytest.param(None, None,
id="empty-is-none-accepts-only-aud-less-tokens"),
+ pytest.param("api", "api", id="configured-value-passed-through"),
+ ],
+ )
+ def test_oidc_audience_parsing(self, configured, expected):
+ """An empty audience becomes ``None`` (accept only aud-less tokens); a
set value is forwarded."""
+ overrides = {} if configured is None else {("edge", "jwt_audience"):
configured}
+ with conf_vars(overrides):
+ assert auth._jwt_audience() == expected
+
+ @pytest.mark.parametrize(
+ ("configured", "expected"),
+ [
+ pytest.param(None, None, id="empty-is-none-skips-issuer-check"),
+ pytest.param(
+ "https://idp.example.com";, "https://idp.example.com";,
id="configured-passed-through"
+ ),
+ ],
+ )
+ def test_oidc_issuer_parsing(self, configured, expected):
+ """An empty issuer becomes ``None`` (skip issuer check); a set value
is forwarded verbatim."""
+ overrides = {} if configured is None else {("edge", "jwt_issuer"):
configured}
+ with conf_vars(overrides):
+ assert auth._jwt_issuer() == expected
+
+ @pytest.mark.parametrize(
+ ("configured", "expected"),
+ [
+ pytest.param(None, 30, id="default-30"),
+ pytest.param("90", 90, id="configured-value-parsed-as-int"),
+ ],
+ )
+ def test_oidc_leeway_parsing(self, configured, expected):
+ """Unset falls back to 30 seconds; a configured value is read as an
integer."""
+ overrides = {} if configured is None else {("edge", "jwt_leeway"):
configured}
+ with conf_vars(overrides):
+ assert auth._jwt_leeway() == expected
+
+
+class TestJwtValidatorSelection:
+ """``jwt_validator`` picks shared-secret vs OIDC based on
``trusted_jwks_url``."""
+
+ @conf_vars({("api_auth", "jwt_secret"): "secret"})
+ def test_uses_shared_secret_validator_when_oidc_jwks_url_unset(self):
+ """Default path is unchanged: a shared-secret validator with no JWKS
is built."""
+ validator = auth.jwt_validator()
+
+ assert validator.secret_key == "secret"
+
+ assert validator.jwks is None
+
+ @conf_vars(
+ {
+ ("edge", "trusted_jwks_url"): OIDC_JWKS_URL,
+ ("edge", "jwt_issuer"): "https://idp.example.com";,
+ ("edge", "jwt_algorithm"): "RS512",
+ ("edge", "jwt_leeway"): "90",
+ }
+ )
+ def test_uses_oidc_validator_when_jwks_url_set(self):
+ """Setting ``trusted_jwks_url`` builds a JWKS-backed validator wired
from ``[edge]`` config."""
+ validator = auth.jwt_validator()
+
+ assert validator.jwks is not None
+
+ assert validator.algorithm == ["RS512"]
+
+ assert validator.issuer == "https://idp.example.com";
+
+ assert validator.audience is None
+
+ assert validator.leeway == 90
+
+ @conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL})
+ def test_empty_issuer_without_verifier_fails_closed(self):
+ """Skipping issuer verification without a ``jwt_verifier`` is a hard
config error."""
+ with pytest.raises(AirflowConfigException, match="jwt_verifier"):
+ auth.jwt_validator()
+
+ @conf_vars(
+ {
+ ("edge", "trusted_jwks_url"): OIDC_JWKS_URL,
+ ("edge", "jwt_verifier"):
"my_company.edge_auth.verify_worker_token",
+ }
+ )
+ def test_empty_issuer_with_verifier_is_allowed(self):
+ """A configured ``jwt_verifier`` permits skipping issuer
verification."""
+ validator = auth.jwt_validator()
+
+ assert validator.jwks is not None
+
+ assert validator.issuer is None
+
+
+class TestMethodClaimCheck:
+ """``_check_method_claim`` enforces the signed ``method`` only for
shared-secret tokens."""
+
+ @conf_vars({("api_auth", "jwt_secret"): "secret"})
+ def test_shared_secret_rejects_mismatched_method(self):
+ """A shared-secret token minted for another endpoint is forbidden
(403)."""
+ with pytest.raises(HTTPException) as exc_info:
+ auth._check_method_claim("worker/register", {"method":
"worker/other"})
+
+ assert exc_info.value.status_code == 403
+
+ @conf_vars({("api_auth", "jwt_secret"): "secret"})
+ def test_shared_secret_accepts_matching_method(self):
+ """A shared-secret token whose ``method`` matches the request passes
without raising."""
+ auth._check_method_claim("worker/register", {"method":
"worker/register"})
+
+ @conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL})
+ def test_oidc_skips_method_claim(self):
+ """OIDC tokens carry no ``method`` claim, so the check is skipped
rather than 403."""
+ auth._check_method_claim("worker/register", {})
+
+
+def _signed_token(audience: str = "") -> str:
+ """Mint a real HS512 token, omitting ``aud`` when passed empty."""
+ generator = JWTGenerator(secret_key=JWT_SECRET, valid_for=300,
audience=audience)
+ return generator.generate()
+
+
+class TestOidcAudienceVerification:
+ """
+ Pin the real audience semantics against a signed token.
+
+ ``jwt_audience`` empty means ``audience=None``, which PyJWT accepts only
for
+ tokens that carry no ``aud`` claim and rejects for tokens that do. A
configured
+ value requires a matching ``aud``.
+ """
+
+ def test_empty_audience_accepts_aud_less_token(self):
+ with conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL}):
+ validator = JWTValidator(secret_key=JWT_SECRET,
audience=auth._jwt_audience(), leeway=5)
Review Comment:
Nit on what this class covers: it pins PyJWT's audience handling rather than
the OIDC validator this PR builds.
The validator here is hand-constructed with `secret_key=JWT_SECRET` — an
HS512 shared secret — so `_oidc_validator()` never runs, and the only PR code
in the path is `auth._jwt_audience()`, which
`TestOidcConfig.test_oidc_audience_parsing` already covers. The
`conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL})` wrapper is inert:
nothing inside the `with` block reads it, since `_jwt_audience()` doesn't
depend on the mode. That's the part I'd change regardless — it reads as though
OIDC mode is under test, which invites the next person to trust it for more
than it covers.
Taken together with `TestWorkerAuthorization` mocking `_jwt_verifier` (so
`conf.getimport` never runs), nothing in the suite drives a token through the
JWKS path: key resolution by `kid`, the `algorithm` restriction, and issuer
binding on a real signed token are all uncovered.
There's a working recipe in the tree at
`airflow-core/tests/unit/api_fastapi/auth/test_tokens.py:205-208` —
`JWKS.from_private_key((key, "kid1"))` builds an in-memory keyset, so a test
can mint a token with `JWTGenerator(private_key=...)` and drive
`jwt_token_authorization` end to end with no network. That would let these
three audience cases assert against the real `_oidc_validator()` instead of a
stand-in.
Being straight about the limits: `from_private_key` sets `url=os.devnull`,
so such a test would *not* catch the event-loop issue on the other thread — it
closes the semantics gap, not that one.
--
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]
