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>
feast-dev · ntkathole · Aug 18, 2026 · Aug 14, 2026 · Aug 18, 2026 · 50c228da975d12e5d750eaf5e5b1fe32fe748643
commit 50c228da975d12e5d750eaf5e5b1fe32fe748643
diff --git a/.secrets.baseline b/.secrets.baseline
@@ -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):

@@ -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
@@ -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 "

@@ -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}"""