From 50c228da975d12e5d750eaf5e5b1fe32fe748643 Mon Sep 17 00:00:00 2001 From: Larry Singleton <166439969+larrysingleton007@users.noreply.github.com> Date: Fri, 14 Aug 2026 03:25:27 -0500 Subject: [PATCH] fix: Reuse IdP-issued client tokens until near expiry On the client_secret branch every outbound RPC built a fresh auth-token factory, manager and OIDCDiscoveryService, paying a discovery GET plus a token POST per call and discarding the token. All three auth interceptors invoke it per RPC, so a batch materialization loop multiplied IdP load by request count and could trip IdP rate limits. Measured before: 5 calls = 5 discovery GETs + 5 token POSTs. After: 1 and 1. Caches IdP tokens in a module-level dict keyed by the token-request identity, since the interceptors build a fresh manager per call and instance state would not survive. Expiry comes from the token's own exp claim, falling back to the token endpoint's expires_in; a token whose expiry is unknowable is not cached, preserving per-call behaviour for opaque tokens. The cache stores true expiry and applies the caller's token_refresh_margin_seconds on read, rather than storing a deadline. The margin is not part of the key, so baking it in let a config with a wider margin reuse a token past its own safety window when another config sharing the same credentials had written the entry. token_refresh_margin_seconds is configurable on OidcClientAuthConfig (default 30, gt=0) rather than hardcoded. Inserting on a miss prunes entries whose stored expiry has passed. Keys are credential identities so the cache is bounded by distinct configs, but a long-lived process rotating credentials would otherwise retain every retired identity. Reuse means a token the IdP revokes mid-life keeps being presented until its own expiry, where fetching per call self-corrected. Adds OidcAuthClientManager.invalidate_token and a transport-agnostic invalidate_auth_token(auth_config), wired into the gRPC interceptor: an UNAUTHENTICATED response drops the cached token so the next call refetches, bounding staleness to the rejected request. The call is not retried, because all four interceptor methods share that path and a stream's request_iterator may already be consumed. 329 permissions tests pass. Signed-off-by: Larry Singleton <166439969+larrysingleton007@users.noreply.github.com> --- .secrets.baseline | 6 +- sdk/python/feast/permissions/auth_model.py | 5 + .../permissions/client/client_auth_token.py | 13 + .../client/grpc_client_auth_interceptor.py | 31 ++- .../oidc_authentication_client_manager.py | 114 +++++++- .../unit/permissions/test_oidc_auth_client.py | 254 +++++++++++++++++- 6 files changed, 414 insertions(+), 9 deletions(-) diff --git a/.secrets.baseline b/.secrets.baseline index 5eee072f6f0..332bb7596f8 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"