diff --git a/docs/getting-started/components/authz_manager.md b/docs/getting-started/components/authz_manager.md index abc97594e61..12fdf2f39e5 100644 --- a/docs/getting-started/components/authz_manager.md +++ b/docs/getting-started/components/authz_manager.md @@ -122,6 +122,19 @@ A token whose `aud` (or `iss`) claim does not match is rejected at authenticatio Set these to the values your IdP puts **in the token itself**, which are not always the ones in the discovery document. For example, Microsoft Entra ID commonly issues v1.0 tokens (`iss: https://sts.windows.net//`, `aud: api://`) even when `auth_discovery_url` points at the v2.0 endpoint. That setup keeps working with these options unset, or set to the v1.0 values — but copying the v2.0 issuer from the discovery document would reject every v1.0 token. {% endhint %} +To validate token signatures the server fetches the provider's JWKS document and caches it, refetching when the cache expires or when a token presents an unknown key id. Two options tune that behavior: + +```yaml +auth: + type: oidc + client_id: _CLIENT_ID_ + auth_discovery_url: https://login.example.com/.well-known/openid-configuration + jwks_cache_lifespan_seconds: 300 # default; how long the fetched key set is reused + jwks_request_timeout_seconds: 10 # default; network timeout for the JWKS fetch +``` + +`jwks_cache_lifespan_seconds` also bounds how long a key the provider has **revoked** continues to validate tokens, so lower it if your provider rotates or revokes aggressively; each reduction costs proportionally more JWKS fetches. Key rotations that introduce a new key id are picked up immediately regardless of this setting, because an unknown key id triggers a refetch. `jwks_request_timeout_seconds` bounds how long an unresponsive provider can block request serving. Both must be greater than zero. + #### Client-Side Configuration The client supports multiple token source modes. The SDK resolves tokens in the following priority order: diff --git a/pixi.lock b/pixi.lock index b663a50e679..ec4c79ff7c3 100644 --- a/pixi.lock +++ b/pixi.lock @@ -2353,7 +2353,7 @@ packages: - prometheus-client>=0.20.0,<0.25.0 - psutil - bigtree>=0.19.2 - - pyjwt + - pyjwt>=2.13.0 - aerospike>=19.0.0,<20.0.0 ; extra == 'aerospike' - boto3>=1.38.27 ; extra == 'aws' - fsspec>=2024.1.0 ; extra == 'aws' diff --git a/pyproject.toml b/pyproject.toml index e4623b51509..e2441730106 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,9 @@ dependencies = [ "prometheus_client>=0.20.0,<0.25.0", "psutil", "bigtree>=0.19.2", - "pyjwt", + # >=2.13 for PyJWKClient's JWK-set cache surviving transient fetch + # failures; the OIDC token parser reuses one client and relies on it. + "pyjwt>=2.13.0", ] [project.optional-dependencies] diff --git a/sdk/python/feast/permissions/auth/oidc_token_parser.py b/sdk/python/feast/permissions/auth/oidc_token_parser.py index c02ec2b81ef..e2b8aeb79cf 100644 --- a/sdk/python/feast/permissions/auth/oidc_token_parser.py +++ b/sdk/python/feast/permissions/auth/oidc_token_parser.py @@ -41,6 +41,41 @@ def __init__(self, auth_config: OidcAuthConfig): ca_cert_path=self._auth_config.ca_cert_path, ) self._k8s_auth_api = None + self._jwks_client: Optional[PyJWKClient] = None # Initialize it lazily. + + def _get_jwks_client(self) -> PyJWKClient: + """Lazily build and cache a parser-lifetime ``PyJWKClient``. + + A per-request client starts with a cold JWK-set cache, forcing a + full HTTPS fetch of the JWKS document on every authenticated + request. Reusing one client lets PyJWT cache the JWK set for + ``jwks_cache_lifespan_seconds``, which also bounds two staleness + windows: a key the IdP has removed keeps validating, and a + rotation that reuses an existing ``kid`` keeps failing, for at + most that long. Rotations that introduce a new ``kid`` recover + immediately (``PyJWKClient.get_signing_key`` refreshes and + retries once on a cache miss). + """ + if self._jwks_client is None: + ssl_ctx = ssl.create_default_context() + if not self._auth_config.verify_ssl: + ssl_ctx.check_hostname = False + ssl_ctx.verify_mode = ssl.CERT_NONE + elif self._auth_config.ca_cert_path and os.path.exists( + self._auth_config.ca_cert_path + ): + ssl_ctx.load_verify_locations(self._auth_config.ca_cert_path) + self._jwks_client = PyJWKClient( + self.oidc_discovery_service.get_jwks_url(), + headers={"User-agent": "custom-user-agent"}, + ssl_context=ssl_ctx, + # Explicit so upgrades cannot silently change the staleness + # window documented above, and so a hung IdP bounds how long + # a fetch can block the serving path. + lifespan=self._auth_config.jwks_cache_lifespan_seconds, + timeout=self._auth_config.jwks_request_timeout_seconds, + ) + return self._jwks_client async def _validate_token(self, access_token: str): """ @@ -125,21 +160,7 @@ def _decode_token(self, access_token: str) -> dict: metadata (e.g. Entra ID v1.0 tokens validated against a v2.0 discovery document). """ - optional_custom_headers = {"User-agent": "custom-user-agent"} - ssl_ctx = ssl.create_default_context() - if not self._auth_config.verify_ssl: - ssl_ctx.check_hostname = False - ssl_ctx.verify_mode = ssl.CERT_NONE - elif self._auth_config.ca_cert_path and os.path.exists( - self._auth_config.ca_cert_path - ): - ssl_ctx.load_verify_locations(self._auth_config.ca_cert_path) - jwks_client = PyJWKClient( - self.oidc_discovery_service.get_jwks_url(), - headers=optional_custom_headers, - ssl_context=ssl_ctx, - ) - signing_key = jwks_client.get_signing_key_from_jwt(access_token) + signing_key = self._get_jwks_client().get_signing_key_from_jwt(access_token) expected_audience = self._auth_config.audience expected_issuer = self._auth_config.issuer return jwt.decode( diff --git a/sdk/python/feast/permissions/auth_model.py b/sdk/python/feast/permissions/auth_model.py index 03d65c5973d..4648a068bce 100644 --- a/sdk/python/feast/permissions/auth_model.py +++ b/sdk/python/feast/permissions/auth_model.py @@ -2,7 +2,7 @@ from typing import Literal, Optional, Tuple -from pydantic import ConfigDict, model_validator +from pydantic import ConfigDict, Field, model_validator from feast.repo_config import FeastConfigBaseModel @@ -47,6 +47,15 @@ class OidcAuthConfig(AuthConfig): # against a v2.0 discovery URL). audience: Optional[str] = None issuer: Optional[str] = None + # How long the fetched JWK set is reused before the server refetches it. + # This also bounds how long a key the IdP has revoked keeps validating + # tokens, so lower it if your provider rotates or revokes aggressively; + # every reduction costs a corresponding increase in JWKS fetches. + jwks_cache_lifespan_seconds: int = Field(default=300, gt=0) + # Network timeout for the JWKS fetch. This fetch happens inline on the + # request path, so an unresponsive IdP blocks serving for at most this + # long. + jwks_request_timeout_seconds: float = Field(default=10, gt=0) class OidcClientAuthConfig(OidcAuthConfig): diff --git a/sdk/python/tests/unit/permissions/auth/test_token_parser.py b/sdk/python/tests/unit/permissions/auth/test_token_parser.py index 8f0c82367d5..a5056393b60 100644 --- a/sdk/python/tests/unit/permissions/auth/test_token_parser.py +++ b/sdk/python/tests/unit/permissions/auth/test_token_parser.py @@ -1,5 +1,6 @@ import asyncio import os +import ssl import time from unittest import mock from unittest.mock import MagicMock, patch @@ -7,6 +8,7 @@ import assertpy import jwt import pytest +from pydantic import ValidationError from starlette.authentication import ( AuthenticationError, ) @@ -472,6 +474,177 @@ async def mock_oath2(self, request): assertpy.assert_that(user.has_matching_role(["updater"])).is_false() +# --------------------------------------------------------------------------- +# JWKS client lifecycle (one lazy client per parser) +# --------------------------------------------------------------------------- + + +@patch( + "feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__" +) +@patch("feast.permissions.auth.oidc_token_parser.jwt.decode") +@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient") +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +def test_oidc_jwks_client_is_lazy_and_reused_per_parser( + mock_discovery_data, + mock_jwks_client_cls, + mock_jwt, + mock_oauth2, + oidc_config, + discovery_data, +): + """One JWKS client per parser: built on the first request (not at + construction, which would move a blocking discovery fetch into server + startup), reused across requests, and scoped to the parser instance.""" + mock_discovery_data.return_value = discovery_data + mock_jwt.return_value = {"preferred_username": "my-name"} + + token_parser = OidcTokenParser(auth_config=oidc_config) + assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(0) + + for _ in range(3): + asyncio.run( + token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc") + ) + assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(1) + + # A second parser must not see the first parser's client: each parser + # verifies against the JWKS of its own configured provider. + other_parser = OidcTokenParser(auth_config=oidc_config) + asyncio.run(other_parser.user_details_from_access_token(access_token="aaa-bbb-ccc")) + assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(2) + + +@pytest.mark.parametrize("verify_ssl", [True, False]) +@patch( + "feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__" +) +@patch("feast.permissions.auth.oidc_token_parser.jwt.decode") +@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient") +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +def test_oidc_jwks_client_ssl_context_follows_config( + mock_discovery_data, + mock_jwks_client_cls, + mock_jwt, + mock_oauth2, + verify_ssl, + discovery_data, +): + """The client is built with the discovery JWKS URL and an SSL context + matching verify_ssl: default configs must keep certificate verification + on, and verify_ssl=False must be the only way to turn it off.""" + mock_discovery_data.return_value = discovery_data + mock_jwt.return_value = {"preferred_username": "my-name"} + + token_parser = OidcTokenParser(auth_config=_oidc_config_with(verify_ssl=verify_ssl)) + asyncio.run(token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc")) + + call = mock_jwks_client_cls.call_args + assertpy.assert_that(call.args[0]).is_equal_to(discovery_data["jwks_uri"]) + ssl_ctx = call.kwargs["ssl_context"] + if verify_ssl: + assertpy.assert_that(ssl_ctx.verify_mode).is_equal_to(ssl.CERT_REQUIRED) + assertpy.assert_that(ssl_ctx.check_hostname).is_true() + else: + assertpy.assert_that(ssl_ctx.verify_mode).is_equal_to(ssl.CERT_NONE) + assertpy.assert_that(ssl_ctx.check_hostname).is_false() + + +@pytest.mark.parametrize( + "overrides,expected_lifespan,expected_timeout", + [ + ({}, 300, 10), + ( + { + "jwks_cache_lifespan_seconds": 60, + "jwks_request_timeout_seconds": 2.5, + }, + 60, + 2.5, + ), + ], +) +@patch( + "feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__" +) +@patch("feast.permissions.auth.oidc_token_parser.jwt.decode") +@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient") +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +def test_oidc_jwks_client_cache_and_timeout_follow_config( + mock_discovery_data, + mock_jwks_client_cls, + mock_jwt, + mock_oauth2, + overrides, + expected_lifespan, + expected_timeout, + discovery_data, +): + """The JWK-set cache lifespan and the fetch timeout are operator-tunable: + the lifespan bounds how long a revoked key keeps validating, and the + timeout bounds how long an unresponsive IdP blocks the serving path.""" + mock_discovery_data.return_value = discovery_data + mock_jwt.return_value = {"preferred_username": "my-name"} + + token_parser = OidcTokenParser(auth_config=_oidc_config_with(**overrides)) + asyncio.run(token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc")) + + kwargs = mock_jwks_client_cls.call_args.kwargs + assertpy.assert_that(kwargs["lifespan"]).is_equal_to(expected_lifespan) + assertpy.assert_that(kwargs["timeout"]).is_equal_to(expected_timeout) + + +@pytest.mark.parametrize( + "overrides", + [ + {"jwks_cache_lifespan_seconds": 0}, + {"jwks_cache_lifespan_seconds": -1}, + {"jwks_request_timeout_seconds": 0}, + {"jwks_request_timeout_seconds": -1}, + ], +) +def test_oidc_jwks_tunables_reject_non_positive_values(overrides): + """A non-positive lifespan would expire the cache immediately, silently + restoring a JWKS fetch per request; a non-positive timeout is equally + meaningless. Reject both at config load rather than at serving time.""" + with pytest.raises(ValidationError): + _oidc_config_with(**overrides) + + +@patch( + "feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__" +) +@patch("feast.permissions.auth.oidc_token_parser.jwt.decode") +@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient") +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +def test_oidc_jwks_client_construction_failure_is_retried( + mock_discovery_data, + mock_jwks_client_cls, + mock_jwt, + mock_oauth2, + oidc_config, + discovery_data, +): + """A failed first construction must leave the parser able to retry on + the next request: the parser is a process singleton, so caching a failed + or half-built client would wedge authentication until restart.""" + mock_discovery_data.return_value = discovery_data + mock_jwt.return_value = {"preferred_username": "my-name"} + mock_jwks_client_cls.side_effect = [RuntimeError("IdP unreachable"), MagicMock()] + + token_parser = OidcTokenParser(auth_config=oidc_config) + with pytest.raises(RuntimeError): + asyncio.run( + token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc") + ) + + user = asyncio.run( + token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc") + ) + assertpy.assert_that(user.username).is_equal_to("my-name") + assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(2) + + # --------------------------------------------------------------------------- # Optional audience / issuer verification (opt-in via OidcAuthConfig) # ---------------------------------------------------------------------------