diff --git a/.secrets.baseline b/.secrets.baseline index 9a4ae8bfc7d..19fbb12a185 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -1469,14 +1469,14 @@ "filename": "sdk/python/tests/unit/permissions/test_oidc_auth_client.py", "hashed_secret": "e6eae2da3b4a5bf296d0495192788e2772ac5c79", "is_verified": false, - "line_number": 29 + "line_number": 47 }, { "type": "Secret Keyword", "filename": "sdk/python/tests/unit/permissions/test_oidc_auth_client.py", "hashed_secret": "8318df9ecda039deac9868adf1944a29a95c7114", "is_verified": false, - "line_number": 31 + "line_number": 49 } ], "sdk/python/tests/universal/feature_repos/repo_configuration.py": [ @@ -1564,5 +1564,5 @@ } ] }, - "generated_at": "2026-07-31T18:15:46Z" + "generated_at": "2026-08-14T07:27:08Z" } diff --git a/sdk/python/feast/permissions/auth_model.py b/sdk/python/feast/permissions/auth_model.py index 4648a068bce..c6db30f71aa 100644 --- a/sdk/python/feast/permissions/auth_model.py +++ b/sdk/python/feast/permissions/auth_model.py @@ -67,6 +67,11 @@ class OidcClientAuthConfig(OidcAuthConfig): client_secret: Optional[str] = None token: Optional[str] = None token_env_var: Optional[str] = None + # Stop reusing an IdP-issued token this many seconds before it expires, + # so a reused token still has life left when the server validates it. + # Raise it if clients see sporadic 401s from clock skew or slow calls; + # lower it to squeeze more reuse out of short-lived tokens. + token_refresh_margin_seconds: float = Field(default=30, gt=0) @model_validator(mode="after") def _validate_credentials(self): diff --git a/sdk/python/feast/permissions/client/client_auth_token.py b/sdk/python/feast/permissions/client/client_auth_token.py index 68821e3f9c6..300c1c2ab62 100644 --- a/sdk/python/feast/permissions/client/client_auth_token.py +++ b/sdk/python/feast/permissions/client/client_auth_token.py @@ -12,3 +12,16 @@ def get_auth_token(auth_config: AuthConfig) -> str: .get_auth_client_manager() .get_token() ) + + +def invalidate_auth_token(auth_config: AuthConfig) -> bool: + """Drop any cached token for *auth_config*, returning whether one was held. + + Only the OIDC client manager caches, so this is a no-op for the other auth + types. Callers that can observe an authentication failure should use it: a + token the IdP revokes mid-life still looks valid to the client until its + own expiry, and dropping it bounds that to a single rejected request. + """ + manager = AuthenticationClientManagerFactory(auth_config).get_auth_client_manager() + invalidate = getattr(manager, "invalidate_token", None) + return bool(invalidate()) if callable(invalidate) else False diff --git a/sdk/python/feast/permissions/client/grpc_client_auth_interceptor.py b/sdk/python/feast/permissions/client/grpc_client_auth_interceptor.py index 9a6bef2c072..0384a383b11 100644 --- a/sdk/python/feast/permissions/client/grpc_client_auth_interceptor.py +++ b/sdk/python/feast/permissions/client/grpc_client_auth_interceptor.py @@ -5,7 +5,10 @@ from feast.errors import FeastError from feast.permissions.auth.auth_type import AuthType from feast.permissions.auth_model import AuthConfig -from feast.permissions.client.client_auth_token import get_auth_token +from feast.permissions.client.client_auth_token import ( + get_auth_token, + invalidate_auth_token, +) logger = logging.getLogger(__name__) @@ -44,11 +47,37 @@ def _handle_call(self, continuation, client_call_details, request_iterator): client_call_details = self._append_auth_header_metadata(client_call_details) result = continuation(client_call_details, request_iterator) if result.exception() is not None: + self._invalidate_token_if_rejected(result) mapped_error = FeastError.from_error_detail(result.exception().details()) if mapped_error is not None: raise mapped_error return result + def _invalidate_token_if_rejected(self, result) -> None: + """Drop the cached token when the server rejects it as unauthenticated. + + Tokens are reused until near expiry, so one the IdP revoked mid-life + would otherwise keep being presented for the rest of its lifetime. + Dropping it here bounds that to the single request that was rejected; + the next call fetches a fresh token. + + The call is deliberately not retried. All four interceptor methods + share this path, and a stream's ``request_iterator`` may already be + consumed, so retrying here could replay a partially-sent stream. + """ + if self._auth_config.type == AuthType.NONE.value: + return + try: + if result.code() != grpc.StatusCode.UNAUTHENTICATED: + return + except Exception: # pragma: no cover - result without a status code + return + if invalidate_auth_token(self._auth_config): + logger.debug( + "Server rejected the cached auth token; dropped it so the next " + "call fetches a fresh one." + ) + def _append_auth_header_metadata(self, client_call_details): logger.debug( "Intercepted the grpc api method call to inject Authorization header " diff --git a/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py b/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py index 37e613dafc1..6bc6a7ca263 100644 --- a/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py +++ b/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py @@ -1,6 +1,8 @@ import logging import os -from typing import Optional +import threading +import time +from typing import Dict, Optional, Tuple import jwt import requests @@ -13,6 +15,17 @@ SA_TOKEN_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/token" +# IdP-issued tokens keyed by the token-request identity, stored with their +# true expiry. The auth interceptors build a fresh manager for every outbound +# RPC, so instance state would not survive between calls; without this +# cache every RPC pays a discovery GET plus a token POST against the IdP. +# The refresh margin is applied on read, not baked into the stored value, so +# configs sharing IdP credentials but setting different margins can share +# tokens while each still honors its own margin. +# Concurrent misses may fetch in parallel (benign: last write wins). +_token_cache: Dict[Tuple, Tuple[str, float]] = {} +_token_cache_lock = threading.Lock() + class OidcAuthClientManager(AuthenticationClientManager): def __init__(self, auth_config: OidcClientAuthConfig): @@ -67,7 +80,96 @@ def _read_sa_token() -> Optional[str]: return None def _fetch_token_from_idp(self) -> str: - """Obtain an access token via client_credentials or ROPG flow.""" + """Return a cached IdP token, or obtain a fresh one. + + The cache stores each token's true expiry (its ``exp`` claim, falling + back to the token response's ``expires_in``), and this config's + ``token_refresh_margin_seconds`` is applied when reading. Keeping the + margin out of the stored value lets configs that share IdP credentials + but set different margins share tokens while each honors its own + margin. A token whose expiry is unknowable is not cached, preserving + the previous per-call behavior for opaque tokens. + """ + cache_key = self._cache_key() + with _token_cache_lock: + cached = _token_cache.get(cache_key) + if cached is not None: + cached_token, cached_expiry = cached + margin = self.auth_config.token_refresh_margin_seconds + if time.time() < cached_expiry - margin: + return cached_token + + access_token, expires_in = self._request_token_from_idp() + + expiry = self._token_expiry(access_token, expires_in) + if expiry is not None: + now = time.time() + with _token_cache_lock: + # Prune on miss. Entries are keyed by credential identity, so + # the cache is bounded by the number of distinct configs, but + # a long-lived process that rotates credentials would otherwise + # keep every retired identity forever. + for key in [k for k, (_, exp) in _token_cache.items() if exp <= now]: + del _token_cache[key] + _token_cache[cache_key] = (access_token, expiry) + return access_token + + def _cache_key(self) -> Tuple: + """Identity of the token request: same credentials, same token.""" + return ( + self.auth_config.auth_discovery_url, + self.auth_config.client_id, + self.auth_config.client_secret, + self.auth_config.username, + self.auth_config.password, + ) + + def invalidate_token(self) -> bool: + """Drop this config's cached token so the next call refetches. + + Returns whether an entry was actually removed. + + Reuse means a token the IdP revokes mid-life keeps being presented + until its own expiry, where fetching per call self-corrected. Callers + that can observe an authentication failure should invalidate on it, so + the staleness costs one rejected request rather than the remaining + lifetime of the token. + """ + with _token_cache_lock: + return _token_cache.pop(self._cache_key(), None) is not None + + @staticmethod + def _token_expiry( + access_token: str, expires_in: Optional[float] + ) -> Optional[float]: + """Epoch expiry of *access_token*, or ``None`` when it is unknowable. + + Prefers the token's own ``exp`` claim (authoritative); falls back to + the token endpoint's ``expires_in``. + + The refresh margin is deliberately not subtracted here. Storing one + caller's deadline would let another config with a wider margin reuse + the token past its own safety window, since the margin is not part of + the cache key. + """ + exp: Optional[float] = None + try: + claims = jwt.decode(access_token, options={"verify_signature": False}) + claim = claims.get("exp") + if isinstance(claim, (int, float)): + exp = float(claim) + except jwt.exceptions.DecodeError: + pass + if exp is None and isinstance(expires_in, (int, float)): + exp = time.time() + float(expires_in) + return exp + + def _request_token_from_idp(self) -> Tuple[str, Optional[float]]: + """Obtain an access token via client_credentials or ROPG flow. + + Returns the token and the token response's ``expires_in`` (seconds), + when the IdP provides one. + """ if self.auth_config.auth_discovery_url is None: raise ValueError( "auth_discovery_url is required for IDP token fetch " @@ -106,13 +208,17 @@ def _fetch_token_from_idp(self) -> str: ) if token_response.status_code == 200: - access_token = token_response.json()["access_token"] + response_body = token_response.json() + access_token = response_body["access_token"] if not access_token: logger.debug( f"access_token is empty for the client_id=${self.auth_config.client_id}" ) raise RuntimeError("access token is empty") - return access_token + expires_in = response_body.get("expires_in") + if not isinstance(expires_in, (int, float)): + expires_in = None + return access_token, expires_in else: raise RuntimeError( f"""Failed to obtain oidc access token:url=[{token_endpoint}] {token_response.status_code} - {token_response.text}""" diff --git a/sdk/python/tests/unit/permissions/test_oidc_auth_client.py b/sdk/python/tests/unit/permissions/test_oidc_auth_client.py index 3d74eb2a55f..80653774ee6 100644 --- a/sdk/python/tests/unit/permissions/test_oidc_auth_client.py +++ b/sdk/python/tests/unit/permissions/test_oidc_auth_client.py @@ -1,5 +1,9 @@ -from unittest.mock import patch +import time +from unittest.mock import MagicMock, patch +import jwt +import pytest +from pydantic import ValidationError from requests import Session from feast.permissions.auth_model import ( @@ -7,6 +11,11 @@ NoAuthConfig, OidcClientAuthConfig, ) +from feast.permissions.client import oidc_authentication_client_manager +from feast.permissions.client.client_auth_token import ( + get_auth_token, + invalidate_auth_token, +) from feast.permissions.client.http_auth_requests_wrapper import ( AuthenticatedRequestsSession, get_http_auth_requests_session, @@ -21,6 +30,15 @@ MOCKED_TOKEN_VALUE: str = "dummy_token" +@pytest.fixture(autouse=True) +def clear_idp_token_cache(): + """The IdP token cache is module-level state; reset it around every test + so no test inherits (or leaks) a cached token.""" + oidc_authentication_client_manager._token_cache.clear() + yield + oidc_authentication_client_manager._token_cache.clear() + + def _get_dummy_oidc_auth_type() -> OidcClientAuthConfig: oidc_config = OidcClientAuthConfig( auth_discovery_url="http://localhost:8080/realms/master/.well-known/openid-configuration", @@ -61,3 +79,237 @@ def _assert_auth_requests_session( assert auth_req_session.headers["Authorization"] == f"Bearer {expected_token}", ( "Authorization token is incorrect" ) + + +# --------------------------------------------------------------------------- +# IdP token reuse (client_secret / ROPG flow) +# --------------------------------------------------------------------------- + +_DISCOVERY = { + "token_endpoint": "https://idp.example.com/token", + "authorization_endpoint": "https://idp.example.com/auth", + "jwks_uri": "https://idp.example.com/jwks", +} + + +def _jwt_expiring_in(seconds: int, subject: str = "test-subject") -> str: + """A JWT expiring *seconds* from now. *subject* distinguishes otherwise + identical tokens, so a test can assert which one a caller received.""" + return jwt.encode( + {"exp": int(time.time()) + seconds, "sub": subject}, + "test-key", + algorithm="HS256", + ) + + +def _token_response(access_token: str, expires_in=None) -> MagicMock: + response = MagicMock(status_code=200) + body = {"access_token": access_token} + if expires_in is not None: + body["expires_in"] = expires_in + response.json.return_value = body + return response + + +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post") +def test_idp_token_reused_until_expiry(mock_post, mock_discovery): + """One token POST (and one discovery fetch) serves many RPCs: the auth + interceptors call get_auth_token per outbound request, so without reuse + every RPC pays two IdP round trips.""" + mock_discovery.return_value = _DISCOVERY + mock_post.return_value = _token_response(_jwt_expiring_in(3600)) + + config = _get_dummy_oidc_auth_type() + tokens = {get_auth_token(config) for _ in range(5)} + + assert len(tokens) == 1 + assert mock_post.call_count == 1 + assert mock_discovery.call_count == 1 + + +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post") +def test_token_inside_refresh_margin_is_not_reused(mock_post, mock_discovery): + """A token whose remaining lifetime is inside the refresh margin must not + be cached: a reused token must never expire between header injection and + server-side validation.""" + mock_discovery.return_value = _DISCOVERY + mock_post.return_value = _token_response(_jwt_expiring_in(5)) + + config = _get_dummy_oidc_auth_type() + get_auth_token(config) + get_auth_token(config) + + assert mock_post.call_count == 2 + + +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post") +def test_opaque_token_uses_expires_in_fallback(mock_post, mock_discovery): + """A non-JWT token is cached when the token endpoint supplies expires_in.""" + mock_discovery.return_value = _DISCOVERY + mock_post.return_value = _token_response("opaque-token", expires_in=3600) + + config = _get_dummy_oidc_auth_type() + get_auth_token(config) + get_auth_token(config) + + assert mock_post.call_count == 1 + + +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post") +def test_opaque_token_without_expiry_is_not_cached(mock_post, mock_discovery): + """With no exp claim and no expires_in, expiry is unknowable, so the + per-call behavior is preserved rather than risking a stale token.""" + mock_discovery.return_value = _DISCOVERY + mock_post.return_value = _token_response("opaque-token") + + config = _get_dummy_oidc_auth_type() + get_auth_token(config) + get_auth_token(config) + + assert mock_post.call_count == 2 + + +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post") +def test_refresh_margin_is_configurable(mock_post, mock_discovery): + """token_refresh_margin_seconds tunes how early a token stops being + reused: a token 60s from expiry is reusable under the default 30s margin + but not under a 120s one.""" + mock_discovery.return_value = _DISCOVERY + mock_post.return_value = _token_response(_jwt_expiring_in(60)) + + default_margin = _get_dummy_oidc_auth_type() + get_auth_token(default_margin) + get_auth_token(default_margin) + assert mock_post.call_count == 1 + + mock_post.reset_mock() + mock_post.return_value = _token_response(_jwt_expiring_in(60)) + + wide_margin = _get_dummy_oidc_auth_type() + # Distinct credentials give this config its own cache entry, isolating the + # two margins without reaching into module state to clear the cache. + wide_margin.client_id = "wide_margin_client_id" + wide_margin.token_refresh_margin_seconds = 120 + get_auth_token(wide_margin) + get_auth_token(wide_margin) + assert mock_post.call_count == 2 + + +@pytest.mark.parametrize("margin", [0, -1]) +def test_refresh_margin_rejects_non_positive_values(margin): + """A zero or negative margin would let a token be reused right up to (or + past) its expiry, so reject it at config load rather than at request time.""" + with pytest.raises(ValidationError): + OidcClientAuthConfig( + auth_discovery_url="http://localhost:8080/realms/master/.well-known/openid-configuration", + type="oidc", + client_id="dummy_client_id", + client_secret="client_secret", + token_refresh_margin_seconds=margin, + ) + + +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post") +def test_distinct_configs_do_not_share_tokens(mock_post, mock_discovery): + """The cache is keyed by the token-request identity: two clients with + different credentials must never receive each other's tokens.""" + mock_discovery.return_value = _DISCOVERY + token_a = _jwt_expiring_in(3600) + token_b = _jwt_expiring_in(7200) + mock_post.side_effect = [_token_response(token_a), _token_response(token_b)] + + config_a = _get_dummy_oidc_auth_type() + config_b = _get_dummy_oidc_auth_type() + config_b.client_id = "another_client_id" + + assert get_auth_token(config_a) == token_a + assert get_auth_token(config_b) == token_b + assert get_auth_token(config_a) == token_a + assert get_auth_token(config_b) == token_b + assert mock_post.call_count == 2 + + +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post") +def test_shared_credentials_each_honor_their_own_margin(mock_post, mock_discovery): + """Configs sharing IdP credentials but differing in refresh margin share a + cache entry, so each must apply its own margin when reading it. + + Regression: the cache once stored ``exp - margin`` while the margin was + absent from the cache key, so whichever config wrote the entry imposed its + deadline on the other. Here the 30s-margin config caches a token 60s from + expiry; the 120s-margin config must refetch rather than reuse it with only + 60s left. Deliberately does not clear the cache between the two configs — + clearing is what hid this. + """ + mock_discovery.return_value = _DISCOVERY + narrow_token = _jwt_expiring_in(60, subject="narrow") + wide_token = _jwt_expiring_in(60, subject="wide") + mock_post.side_effect = [ + _token_response(narrow_token), + _token_response(wide_token), + ] + + narrow_margin = _get_dummy_oidc_auth_type() + wide_margin = _get_dummy_oidc_auth_type() + wide_margin.token_refresh_margin_seconds = 120 + + assert get_auth_token(narrow_margin) == narrow_token + assert get_auth_token(wide_margin) == wide_token + assert mock_post.call_count == 2 + + # The reverse direction is safe and must stay cheap: the narrow-margin + # config can reuse the entry the wide-margin one just wrote, because 60s + # of remaining life clears its 30s margin. + assert get_auth_token(narrow_margin) == wide_token + assert mock_post.call_count == 2 + + +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post") +def test_invalidate_auth_token_forces_a_refetch(mock_post, mock_discovery): + """A token the IdP revokes mid-life still looks valid to the client until + its own expiry. Invalidating drops it so the next call refetches, bounding + the staleness to the request that was rejected.""" + mock_discovery.return_value = _DISCOVERY + first = _jwt_expiring_in(3600, subject="first") + second = _jwt_expiring_in(3600, subject="second") + mock_post.side_effect = [_token_response(first), _token_response(second)] + + config = _get_dummy_oidc_auth_type() + assert get_auth_token(config) == first + assert get_auth_token(config) == first + assert mock_post.call_count == 1 + + assert invalidate_auth_token(config) is True + assert get_auth_token(config) == second + assert mock_post.call_count == 2 + + # Nothing cached for this config any more, so a second invalidate is a no-op. + invalidate_auth_token(config) + assert invalidate_auth_token(config) is False + + +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") +@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post") +def test_expired_entries_are_pruned_on_miss(mock_post, mock_discovery): + """Entries are keyed by credential identity, so a long-lived process that + rotates credentials would otherwise retain every retired identity.""" + mock_discovery.return_value = _DISCOVERY + cache = oidc_authentication_client_manager._token_cache + stale_key = ("retired-idp", "retired-client", "secret", None, None) + cache[stale_key] = ("stale-token", time.time() - 1) + live_key = ("live-idp", "live-client", "secret", None, None) + cache[live_key] = ("live-token", time.time() + 3600) + + mock_post.return_value = _token_response(_jwt_expiring_in(3600)) + get_auth_token(_get_dummy_oidc_auth_type()) + + assert stale_key not in cache, "expired entry should be pruned on insert" + assert live_key in cache, "unexpired entries must survive the prune"