This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0aa1b9a3330 Add OIDC worker token verification to Edge3 (#72262)
0aa1b9a3330 is described below
commit 0aa1b9a333065e9b7f52cf502f89209bf76ba8ab
Author: Robert Pofuk RBI <[email protected]>
AuthorDate: Thu Sep 24 13:46:52 2026 +0200
Add OIDC worker token verification to Edge3 (#72262)
* Add OIDC worker token verification to Edge3
Edge workers could only authenticate with tokens signed by the shared
[api_auth] jwt_secret, forcing every deployment to distribute that secret
to each worker. Deployments that already run a central OpenID Connect
provider want their workers to authenticate with provider-minted tokens
instead, so the secret never has to leave the API server.
Because an externally issued token proves only who signed it, a trusted
signature alone is not authorization to act as a worker: when issuer
verification is skipped the configuration is rejected unless an explicit
jwt_verifier decides which identities are allowed. Verification stays
opt-in behind [edge] trusted_jwks_url, leaving the shared-secret path
untouched for existing deployments.
* Await OIDC JWKS validation and pass request context to edge worker
verifier
Edge worker token validation ran through the synchronous validated_claims
path, which does network I/O across per-request event loops once a JWKS URL
is configured. Switch the authorization entry points to await the async
validator like the rest of the codebase, resolve the jwt_verifier import
once
when the OIDC authenticator is built so a bad path fails fast, and hand the
verifier a WorkerTokenContext carrying the request method so the input can
grow without a breaking signature change while jwt_verifier is unreleased.
---
.../src/airflow/api_fastapi/auth/tokens.py | 3 +-
providers/edge3/provider.yaml | 80 +++++
.../airflow/providers/edge3/get_provider_info.py | 42 +++
.../src/airflow/providers/edge3/worker_api/auth.py | 195 ++++++++++-
.../edge3/tests/unit/edge3/worker_api/test_auth.py | 358 +++++++++++++++++++--
5 files changed, 643 insertions(+), 35 deletions(-)
diff --git a/airflow-core/src/airflow/api_fastapi/auth/tokens.py
b/airflow-core/src/airflow/api_fastapi/auth/tokens.py
index 4ba1943c577..91518d5c72a 100644
--- a/airflow-core/src/airflow/api_fastapi/auth/tokens.py
+++ b/airflow-core/src/airflow/api_fastapi/auth/tokens.py
@@ -279,7 +279,8 @@ class JWTValidator:
)
# By default, we just validate these
required_claims: frozenset[str] = frozenset({"exp", "iat", "nbf"})
- audience: str | Sequence[str]
+ # ``None`` is passed straight to ``jwt.decode`` and accepts only tokens
without an ``aud`` claim.
+ audience: str | Sequence[str] | None
algorithm: list[str] = attrs.field(
factory=_conf_list_factory("api_auth", "jwt_algorithm",
fallback="GUESS"), converter=_to_list
)
diff --git a/providers/edge3/provider.yaml b/providers/edge3/provider.yaml
index 79abba5aa3a..39334a8315a 100644
--- a/providers/edge3/provider.yaml
+++ b/providers/edge3/provider.yaml
@@ -108,6 +108,86 @@ config:
type: string
example: https://airflow.hosting.org/edge_worker/v1/rpcapi
default: ~
+ trusted_jwks_url:
+ description: |
+ When set, edge worker tokens are verified against this JSON Web Key
Set
+ (JWKS) URL of a trusted OpenID Connect provider, using asymmetric
+ signatures, instead of the shared ``[api_auth] jwt_secret``. This
lets
+ edge workers authenticate with tokens minted by an external identity
+ provider.
+
+ When empty (the default) the shared-secret (symmetric) verification
is
+ used, so existing deployments are unaffected.
+
+ Security note: enabling this disables the per-request ``method``
claim
+ check that the shared-secret path uses to scope a token to a single
+ endpoint (externally issued tokens do not carry it). A valid
signature
+ then only proves who signed the token, not that the identity may act
as
+ an edge worker. If the issuer serves other workloads, configure
+ ``jwt_issuer`` and a ``jwt_verifier`` to restrict which identities
are
+ accepted; otherwise any validly signed token from the issuer is
accepted.
+ version_added: 4.4.0
+ type: string
+ example: https://idp.example.com/oauth2/v1/keys
+ default: ~
+ jwt_issuer:
+ description: |
+ Expected ``iss`` claim of edge worker tokens when
``trusted_jwks_url`` is
+ set. Leaving it empty skips issuer verification, so a token from any
issuer
+ whose key is in the JWKS is accepted; in that case ``jwt_verifier``
must be
+ configured to decide which identities may act as edge workers
(requests are
+ rejected otherwise). Set it to bind tokens to your provider. Ignored
when
+ ``trusted_jwks_url`` is empty.
+ version_added: 4.4.0
+ type: string
+ example: https://idp.example.com
+ default: ~
+ jwt_audience:
+ description: |
+ Expected ``aud`` claim of edge worker tokens when
``trusted_jwks_url`` is
+ set. When set, the token must carry a matching ``aud``. When empty,
only
+ tokens that carry no ``aud`` claim are accepted (a token that does
carry
+ an ``aud`` is rejected). Ignored when ``trusted_jwks_url`` is empty.
+ version_added: 4.4.0
+ type: string
+ example: api
+ default: ~
+ jwt_algorithm:
+ description: |
+ Comma-separated list of JWS signing algorithms accepted for edge
worker
+ tokens when ``trusted_jwks_url`` is set. Set it to match the
algorithm the
+ trusted provider signs with. Ignored when ``trusted_jwks_url`` is
empty.
+ version_added: 4.4.0
+ type: string
+ example: RS256,RS512
+ default: RS256
+ jwt_leeway:
+ description: |
+ Number of seconds of leeway allowed when validating the expiry and
issued-at
+ claims of edge worker tokens issued by a trusted OIDC provider, to
account for
+ clock skew between the identity provider and the API server. Ignored
when
+ ``trusted_jwks_url`` is empty.
+ version_added: 4.4.0
+ type: integer
+ example: ~
+ default: "30"
+ jwt_verifier:
+ description: |
+ Import path of a callable that authorizes edge worker tokens issued
by a
+ trusted OIDC provider, run after signature, issuer and audience
checks pass.
+ It receives the validated claims (``dict``) and a request context
object
+ carrying the request ``method``, and returns a mapping with an
+ ``authorized`` boolean (e.g. ``{"authorized": True}``); a falsy
result or a
+ raised exception rejects the request. Use it to restrict which
identities
+ (for example by ``sub``, ``client_id`` or ``scope``) may act as edge
workers.
+ Required when ``trusted_jwks_url`` is set but ``jwt_issuer`` is
empty, since
+ issuer verification is then skipped. When empty (and ``jwt_issuer``
is set)
+ any token that passes signature, issuer and audience verification is
+ authorized. Ignored when ``trusted_jwks_url`` is empty.
+ version_added: 4.4.0
+ type: string
+ example: my_company.edge_auth.verify_worker_token
+ default: ~
job_poll_interval:
description: |
Edge Worker currently polls for new jobs via HTTP. This parameter
defines the number
diff --git a/providers/edge3/src/airflow/providers/edge3/get_provider_info.py
b/providers/edge3/src/airflow/providers/edge3/get_provider_info.py
index 86375fece35..58488f5c327 100644
--- a/providers/edge3/src/airflow/providers/edge3/get_provider_info.py
+++ b/providers/edge3/src/airflow/providers/edge3/get_provider_info.py
@@ -53,6 +53,48 @@ def get_provider_info():
"example":
"https://airflow.hosting.org/edge_worker/v1/rpcapi";,
"default": None,
},
+ "trusted_jwks_url": {
+ "description": "When set, edge worker tokens are
verified against this JSON Web Key Set\n(JWKS) URL of a trusted OpenID Connect
provider, using asymmetric\nsignatures, instead of the shared ``[api_auth]
jwt_secret``. This lets\nedge workers authenticate with tokens minted by an
external identity\nprovider.\n\nWhen empty (the default) the shared-secret
(symmetric) verification is\nused, so existing deployments are
unaffected.\n\nSecurity note: enabling this disable [...]
+ "version_added": "4.4.0",
+ "type": "string",
+ "example": "https://idp.example.com/oauth2/v1/keys";,
+ "default": None,
+ },
+ "jwt_issuer": {
+ "description": "Expected ``iss`` claim of edge worker
tokens when ``trusted_jwks_url`` is\nset. Leaving it empty skips issuer
verification, so a token from any issuer\nwhose key is in the JWKS is accepted;
in that case ``jwt_verifier`` must be\nconfigured to decide which identities
may act as edge workers (requests are\nrejected otherwise). Set it to bind
tokens to your provider. Ignored when\n``trusted_jwks_url`` is empty.\n",
+ "version_added": "4.4.0",
+ "type": "string",
+ "example": "https://idp.example.com";,
+ "default": None,
+ },
+ "jwt_audience": {
+ "description": "Expected ``aud`` claim of edge worker
tokens when ``trusted_jwks_url`` is\nset. When set, the token must carry a
matching ``aud``. When empty, only\ntokens that carry no ``aud`` claim are
accepted (a token that does carry\nan ``aud`` is rejected). Ignored when
``trusted_jwks_url`` is empty.\n",
+ "version_added": "4.4.0",
+ "type": "string",
+ "example": "api",
+ "default": None,
+ },
+ "jwt_algorithm": {
+ "description": "Comma-separated list of JWS signing
algorithms accepted for edge worker\ntokens when ``trusted_jwks_url`` is set.
Set it to match the algorithm the\ntrusted provider signs with. Ignored when
``trusted_jwks_url`` is empty.\n",
+ "version_added": "4.4.0",
+ "type": "string",
+ "example": "RS256,RS512",
+ "default": "RS256",
+ },
+ "jwt_leeway": {
+ "description": "Number of seconds of leeway allowed
when validating the expiry and issued-at\nclaims of edge worker tokens issued
by a trusted OIDC provider, to account for\nclock skew between the identity
provider and the API server. Ignored when\n``trusted_jwks_url`` is empty.\n",
+ "version_added": "4.4.0",
+ "type": "integer",
+ "example": None,
+ "default": "30",
+ },
+ "jwt_verifier": {
+ "description": 'Import path of a callable that
authorizes edge worker tokens issued by a\ntrusted OIDC provider, run after
signature, issuer and audience checks pass.\nIt receives the validated claims
(``dict``) and a request context object\ncarrying the request ``method``, and
returns a mapping with an\n``authorized`` boolean (e.g. ``{"authorized":
True}``); a falsy result or a\nraised exception rejects the request. Use it to
restrict which identities\n(for examp [...]
+ "version_added": "4.4.0",
+ "type": "string",
+ "example": "my_company.edge_auth.verify_worker_token",
+ "default": None,
+ },
"job_poll_interval": {
"description": "Edge Worker currently polls for new
jobs via HTTP. This parameter defines the number\nof seconds it should sleep
between polls for new jobs.\nJob polling only happens if the Edge Worker seeks
for new work. Not if busy.\n",
"version_added": None,
diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/auth.py
b/providers/edge3/src/airflow/providers/edge3/worker_api/auth.py
index 090f65bfc92..a2b6fc64495 100644
--- a/providers/edge3/src/airflow/providers/edge3/worker_api/auth.py
+++ b/providers/edge3/src/airflow/providers/edge3/worker_api/auth.py
@@ -17,7 +17,9 @@
from __future__ import annotations
import logging
+from dataclasses import dataclass
from functools import cache
+from typing import TYPE_CHECKING, TypedDict
from uuid import uuid4
from fastapi import Header, HTTPException, Request, status
@@ -30,14 +32,94 @@ from jwt import (
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
+
+
+@dataclass(frozen=True)
+class WorkerTokenContext:
+ """Request context passed to a ``[edge] jwt_verifier`` alongside the token
claims."""
+
+ method: str
+ """The requested Edge worker API path after the ``/edge_worker/v1/``
prefix."""
+
+
+def _default_jwt_verifier(claims: dict, context: WorkerTokenContext) ->
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 _resolve_jwt_verifier() -> Callable[[dict, WorkerTokenContext],
WorkerTokenAuthorization | None]:
+ """
+ Import the configured worker-authorization callable, or the permissive
default.
+
+ Resolved once when the OIDC authenticator is built (not per request).
+ """
+ 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),
@@ -45,8 +127,74 @@ def jwt_validator() -> JWTValidator:
)
-def jwt_validate(authorization: str) -> dict:
- return jwt_validator().validated_claims(authorization)
+@dataclass(frozen=True)
+class _OidcAuthenticator:
+ """A JWKS-backed validator paired with its worker-authorization
callable."""
+
+ validator: JWTValidator
+ verifier: Callable[[dict, WorkerTokenContext], WorkerTokenAuthorization |
None]
+
+
+def _build_oidc_authenticator(jwks_url: str) -> _OidcAuthenticator:
+ """
+ Build the validator and worker-authorization callable for 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.
+
+ Raises when issuer verification is skipped (empty ``jwt_issuer``) without a
+ ``jwt_verifier``, since that combination would accept any token signed by a
+ key in the JWKS.
+ """
+ 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."
+ )
+ validator = JWTValidator(
+ jwks=JWKS(url=jwks_url),
+ issuer=_jwt_issuer(),
+ audience=_jwt_audience(),
+ algorithm=_jwt_algorithms(),
+ required_claims=frozenset({"iat", "exp"}),
+ leeway=_jwt_leeway(),
+ )
+ return _OidcAuthenticator(validator=validator,
verifier=_resolve_jwt_verifier())
+
+
+@cache
+def _oidc_authenticator() -> _OidcAuthenticator:
+ """Return the cached OIDC validator and verifier for the configured JWKS
URL."""
+ return _build_oidc_authenticator(_trusted_jwks_url())
+
+
+@cache
+def jwt_validator() -> JWTValidator:
+ if _oidc_enabled():
+ return _oidc_authenticator().validator
+ return _shared_secret_validator()
+
+
+async def jwt_validate(authorization: str) -> dict:
+ return await jwt_validator().avalidated_claims(authorization)
+
+
+def _check_worker_authorization(method: str, payload: dict) -> None:
+ """
+ Verify the token identity is allowed to act as an edge worker in OIDC mode.
+
+ A valid signature proves who signed the token, not that the identity may
act
+ as a worker. The configured ``[edge] jwt_verifier`` answers the second
+ question; a falsy result rejects the request.
+ """
+ if not _oidc_enabled():
+ return
+
+ result = _oidc_authenticator().verifier(payload,
WorkerTokenContext(method=method))
+ if not result or not result.get("authorized"):
+ _forbidden_response("Token is not authorized to act as an edge
worker.")
def _forbidden_response(message: str):
@@ -59,17 +207,32 @@ def _forbidden_response(message: str):
)
-def jwt_token_authorization(method: str, authorization: str):
+def _check_method_claim(method: str, payload: dict) -> None:
+ """
+ Verify the signed ``method`` claim for shared-secret tokens.
+
+ Tokens minted by the Edge API carry the request ``method`` they are valid
+ for. Tokens issued by an external OIDC provider do not, so the check is
+ skipped when OIDC verification is enabled.
+ """
+ if _oidc_enabled():
+ return
+
+ signed_method = payload.get("method")
+ if not signed_method or signed_method != method:
+ _forbidden_response(
+ "Invalid method in token authorization. "
+ f"signed method='{signed_method}' "
+ f"called method='{method}'",
+ )
+
+
+async def jwt_token_authorization(method: str, authorization: str):
"""Check if the JWT token is correct."""
try:
- payload = jwt_validate(authorization)
- signed_method = payload.get("method")
- if not signed_method or signed_method != method:
- _forbidden_response(
- "Invalid method in token authorization. "
- f"signed method='{signed_method}' "
- f"called method='{method}'",
- )
+ payload = await jwt_validate(authorization)
+ _check_method_claim(method, payload)
+ _check_worker_authorization(method, payload)
except BadSignature:
_forbidden_response("Bad Signature. Please use only the tokens
provided by the API.")
except InvalidAudienceError:
@@ -92,11 +255,11 @@ def jwt_token_authorization(method: str, authorization:
str):
_forbidden_response("Unable to authenticate API via token.")
-def jwt_token_authorization_rest(
+async def jwt_token_authorization_rest(
request: Request, authorization: str = Header(description="JWT
Authorization Token")
):
"""Check if the JWT token is correct for REST API requests."""
PREFIX = "/edge_worker/v1/"
path = request.url.path
method_path = path[path.find(PREFIX) + len(PREFIX) :] if PREFIX in path
else path
- jwt_token_authorization(method_path, authorization)
+ await jwt_token_authorization(method_path, authorization)
diff --git a/providers/edge3/tests/unit/edge3/worker_api/test_auth.py
b/providers/edge3/tests/unit/edge3/worker_api/test_auth.py
index 3d4263be855..ed7ce8e7d02 100644
--- a/providers/edge3/tests/unit/edge3/worker_api/test_auth.py
+++ b/providers/edge3/tests/unit/edge3/worker_api/test_auth.py
@@ -23,7 +23,9 @@ import pytest
from fastapi import HTTPException, Request
from itsdangerous import BadSignature
-from airflow.api_fastapi.auth.tokens import JWTGenerator
+from airflow.api_fastapi.auth.tokens import JWKS, JWTGenerator,
generate_private_key
+from airflow.providers.common.compat.sdk import AirflowConfigException
+from airflow.providers.edge3.worker_api import auth
from airflow.providers.edge3.worker_api.auth import (
jwt_token_authorization,
jwt_token_authorization_rest,
@@ -33,6 +35,11 @@ from airflow.providers.edge3.worker_api.auth import (
from tests_common.test_utils.config import conf_vars
JWT_SECRET = "test-jwt-secret"
+OIDC_JWKS_URL = "https://idp.example.com/keys";
+OIDC_ISSUER = "https://idp.example.com";
+OIDC_KID = "kid1"
+
+pytestmark = [pytest.mark.asyncio]
def _token(method: str | None = "test.method", secret: str = JWT_SECRET) ->
str:
@@ -45,28 +52,32 @@ def _token(method: str | None = "test.method", secret: str
= JWT_SECRET) -> str:
@pytest.fixture(autouse=True)
def _reset_jwt_validator_cache():
- # jwt_validator() is cached: make sure config overrides in one test can
never leak
- # a stale validator into another.
+ # jwt_validator(), _oidc_enabled() and _oidc_authenticator() are cached:
make sure
+ # config overrides in one test can never leak a stale validator or mode
into another.
jwt_validator.cache_clear()
+ auth._oidc_enabled.cache_clear()
+ auth._oidc_authenticator.cache_clear()
yield
jwt_validator.cache_clear()
+ auth._oidc_enabled.cache_clear()
+ auth._oidc_authenticator.cache_clear()
class TestJwtTokenAuthorization:
- @conf_vars({("api_auth", "jwt_secret"): JWT_SECRET, ("api_auth",
"jwt_leeway"): "5"})
- def test_matching_method_claim_is_authorized(self):
- jwt_token_authorization("test.method", _token("test.method"))
+ async def test_matching_method_claim_is_authorized(self):
+ with conf_vars({("api_auth", "jwt_secret"): JWT_SECRET, ("api_auth",
"jwt_leeway"): "5"}):
+ await jwt_token_authorization("test.method", _token("test.method"))
- @conf_vars({("api_auth", "jwt_secret"): JWT_SECRET, ("api_auth",
"jwt_leeway"): "5"})
- def test_missing_method_claim_is_forbidden(self):
- with pytest.raises(HTTPException) as exc_info:
- jwt_token_authorization("test.method", _token(method=None))
+ async def test_missing_method_claim_is_forbidden(self):
+ with conf_vars({("api_auth", "jwt_secret"): JWT_SECRET, ("api_auth",
"jwt_leeway"): "5"}):
+ with pytest.raises(HTTPException) as exc_info:
+ await jwt_token_authorization("test.method",
_token(method=None))
assert exc_info.value.status_code == 403
- @conf_vars({("api_auth", "jwt_secret"): JWT_SECRET, ("api_auth",
"jwt_leeway"): "5"})
- def test_mismatched_method_claim_is_forbidden(self):
- with pytest.raises(HTTPException) as exc_info:
- jwt_token_authorization("test.method", _token("other.method"))
+ async def test_mismatched_method_claim_is_forbidden(self):
+ with conf_vars({("api_auth", "jwt_secret"): JWT_SECRET, ("api_auth",
"jwt_leeway"): "5"}):
+ with pytest.raises(HTTPException) as exc_info:
+ await jwt_token_authorization("test.method",
_token("other.method"))
assert exc_info.value.status_code == 403
@@ -90,11 +101,11 @@ class TestJwtTokenAuthorizationForbiddenResponse:
],
)
@mock.patch("airflow.providers.edge3.worker_api.auth.jwt_validate",
autospec=True)
- def test_each_handled_failure_is_forbidden_and_anonymized(self,
mock_jwt_validate, error):
+ async def test_each_handled_failure_is_forbidden_and_anonymized(self,
mock_jwt_validate, error):
mock_jwt_validate.side_effect = error
with pytest.raises(HTTPException) as exc_info:
- jwt_token_authorization("test.method", "some-token")
+ await jwt_token_authorization("test.method", "some-token")
assert exc_info.value.status_code == 403
assert "error_id=" in exc_info.value.detail
@@ -141,12 +152,323 @@ class TestJwtTokenAuthorizationRest:
],
)
@mock.patch("airflow.providers.edge3.worker_api.auth.jwt_token_authorization",
autospec=True)
- def test_strips_edge_worker_v1_prefix_and_falls_back_to_full_path(
+ async def test_strips_edge_worker_v1_prefix_and_falls_back_to_full_path(
self, mock_jwt_token_authorization, path, expected_method
):
request = mock.MagicMock(spec=Request)
request.url.path = path
- jwt_token_authorization_rest(request, authorization="some-token")
+ await 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"):
"airflow.providers.edge3.worker_api.auth._default_jwt_verifier",
+ }
+ )
+ 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
+
+ @conf_vars(
+ {
+ ("edge", "trusted_jwks_url"): OIDC_JWKS_URL,
+ ("edge", "jwt_issuer"): OIDC_ISSUER,
+ ("edge", "jwt_verifier"): "no.such.module.verify",
+ }
+ )
+ def test_unimportable_verifier_fails_closed_at_construction(self):
+ """A typo'd ``jwt_verifier`` path fails when the authenticator is
built, not per request."""
+ with pytest.raises(AirflowConfigException):
+ auth.jwt_validator()
+
+
+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", {})
+
+
[email protected](scope="session")
+def oidc_private_key():
+ return generate_private_key()
+
+
[email protected]
+def in_memory_jwks(oidc_private_key):
+ # Replace the network-fetching JWKS(url=...) the OIDC validator builds
with an
+ # in-memory keyset, so signed tokens can be driven through the real
validator.
+ jwks = JWKS.from_private_key((oidc_private_key, OIDC_KID))
+ with mock.patch.object(auth, "JWKS", return_value=jwks):
+ yield
+
+
+def _oidc_token(oidc_private_key, *, audience: str = "", issuer: str =
OIDC_ISSUER, sub: str = "") -> str:
+ generator = JWTGenerator(
+ private_key=oidc_private_key,
+ kid=OIDC_KID,
+ valid_for=300,
+ audience=audience,
+ issuer=issuer,
+ algorithm="RS256",
+ )
+ return generator.generate(extras={"sub": sub} if sub else {})
+
+
+def verify_only_known_sub(claims: dict, context: auth.WorkerTokenContext) ->
auth.WorkerTokenAuthorization:
+ """Test ``jwt_verifier``: authorize only a known service-account
``sub``."""
+ return {"authorized": claims.get("sub") == "dc-service-account"}
+
+
+VERIFY_ONLY_KNOWN_SUB = f"{__name__}.verify_only_known_sub"
+
+
+class TestOidcAudienceVerification:
+ """
+ Pin the real audience semantics by driving signed tokens through the OIDC
validator.
+
+ ``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``.
+ """
+
+ async def test_empty_audience_accepts_aud_less_token(self,
oidc_private_key, in_memory_jwks):
+ with conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL, ("edge",
"jwt_issuer"): OIDC_ISSUER}):
+ await jwt_token_authorization("worker/register",
_oidc_token(oidc_private_key, audience=""))
+
+ async def test_empty_audience_rejects_token_with_aud(self,
oidc_private_key, in_memory_jwks):
+ with conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL, ("edge",
"jwt_issuer"): OIDC_ISSUER}):
+ with pytest.raises(HTTPException) as exc_info:
+ await jwt_token_authorization(
+ "worker/register", _oidc_token(oidc_private_key,
audience="api")
+ )
+ assert exc_info.value.status_code == 403
+
+ async def test_configured_audience_requires_match(self, oidc_private_key,
in_memory_jwks):
+ with conf_vars(
+ {
+ ("edge", "trusted_jwks_url"): OIDC_JWKS_URL,
+ ("edge", "jwt_issuer"): OIDC_ISSUER,
+ ("edge", "jwt_audience"): "api",
+ }
+ ):
+ await jwt_token_authorization("worker/register",
_oidc_token(oidc_private_key, audience="api"))
+ with pytest.raises(HTTPException) as exc_info:
+ await jwt_token_authorization(
+ "worker/register", _oidc_token(oidc_private_key,
audience="other")
+ )
+ assert exc_info.value.status_code == 403
+
+
+class TestOidcTokenValidation:
+ """Drive signed tokens through the real ``_oidc_validator()`` and JWKS
path."""
+
+ async def test_trusted_issuer_but_unrelated_identity_is_rejected(self,
oidc_private_key, in_memory_jwks):
+ with conf_vars(
+ {
+ ("edge", "trusted_jwks_url"): OIDC_JWKS_URL,
+ ("edge", "jwt_issuer"): OIDC_ISSUER,
+ ("edge", "jwt_verifier"): VERIFY_ONLY_KNOWN_SUB,
+ }
+ ):
+ token = _oidc_token(oidc_private_key, issuer=OIDC_ISSUER,
sub="someone-else")
+ with pytest.raises(HTTPException) as exc_info:
+ await jwt_token_authorization("worker/register", token)
+ assert exc_info.value.status_code == 403
+
+ async def test_authorized_identity_from_trusted_issuer_is_accepted(
+ self, oidc_private_key, in_memory_jwks
+ ):
+ with conf_vars(
+ {
+ ("edge", "trusted_jwks_url"): OIDC_JWKS_URL,
+ ("edge", "jwt_issuer"): OIDC_ISSUER,
+ ("edge", "jwt_verifier"): VERIFY_ONLY_KNOWN_SUB,
+ }
+ ):
+ token = _oidc_token(oidc_private_key, issuer=OIDC_ISSUER,
sub="dc-service-account")
+ await jwt_token_authorization("worker/register", token)
+
+ async def test_token_from_untrusted_issuer_is_rejected(self,
oidc_private_key, in_memory_jwks):
+ with conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL, ("edge",
"jwt_issuer"): OIDC_ISSUER}):
+ token = _oidc_token(oidc_private_key,
issuer="https://evil.example.com";)
+ with pytest.raises(HTTPException) as exc_info:
+ await jwt_token_authorization("worker/register", token)
+ assert exc_info.value.status_code == 403
+
+
+class TestWorkerAuthorization:
+ """
+ ``_check_worker_authorization`` gates OIDC tokens on the ``jwt_verifier``.
+
+ A valid signature is not authorization: the verifier answers whether the
+ identity may act as an edge worker.
+ """
+
+ def test_no_check_in_shared_secret_mode(self):
+ """The verifier is not consulted when OIDC is disabled."""
+ with mock.patch.object(auth, "_oidc_authenticator") as authenticator:
+ auth._check_worker_authorization("worker/register", {"sub":
"anyone"})
+ authenticator.assert_not_called()
+
+ @conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL, ("edge",
"jwt_verifier"): VERIFY_ONLY_KNOWN_SUB})
+ def test_configured_verifier_rejects_unauthorized_identity(self):
+ """A token from the trusted issuer but an unrelated identity is
rejected."""
+ with pytest.raises(HTTPException) as exc_info:
+ auth._check_worker_authorization("worker/register", {"sub":
"someone-else"})
+ assert exc_info.value.status_code == 403
+
+ @conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL, ("edge",
"jwt_verifier"): VERIFY_ONLY_KNOWN_SUB})
+ def test_configured_verifier_authorizes_known_identity(self):
+ """A token whose identity the verifier accepts passes without
raising."""
+ auth._check_worker_authorization("worker/register", {"sub":
"dc-service-account"})
+
+ @mock.patch.object(
+ auth, "jwt_validate", new_callable=mock.AsyncMock,
return_value={"sub": "dc-service-account"}
+ )
+ @mock.patch.object(auth, "_oidc_authenticator")
+ async def test_verifier_raising_is_forbidden(self, mock_authenticator,
_mock_validate):
+ """A verifier that raises is caught by the entry point and collapsed
to a 403."""
+ mock_authenticator.return_value.verifier =
mock.Mock(side_effect=ValueError("nope"))
+ with conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL, ("edge",
"jwt_issuer"): OIDC_ISSUER}):
+ with pytest.raises(HTTPException) as exc_info:
+ await auth.jwt_token_authorization("worker/register",
"some-token")
+ assert exc_info.value.status_code == 403
+
+ def test_default_verifier_returns_authorized(self):
+ assert auth._default_jwt_verifier(
+ {"sub": "anyone"},
auth.WorkerTokenContext(method="worker/register")
+ ) == {"authorized": True}
+
+ @conf_vars({("edge", "trusted_jwks_url"): OIDC_JWKS_URL, ("edge",
"jwt_verifier"): VERIFY_ONLY_KNOWN_SUB})
+ @mock.patch.object(auth, "_oidc_authenticator")
+ def test_verifier_receives_request_context(self, mock_authenticator):
+ """The verifier is handed the request method via
``WorkerTokenContext``."""
+ verifier = mock.Mock(return_value={"authorized": True})
+ mock_authenticator.return_value.verifier = verifier
+ auth._check_worker_authorization("jobs/fetch/worker1", {"sub":
"dc-service-account"})
+ verifier.assert_called_once_with(
+ {"sub": "dc-service-account"},
auth.WorkerTokenContext(method="jobs/fetch/worker1")
+ )
