DX-115674: JWKS-based JWT verification to fix 401 token refresh

When the Dremio API returns 401 (expired token), the MCP SDK's
except-Exception catch-all wraps it in a 200 tool error response.
The MCP client never sees HTTP 401 so its OAuth refresh flow is
never triggered.

Fix: add optional `jwks_uri` and `jwks_cache_lifespan` settings.
When configured, verify JWT signatures via JWKS and extract claims
(exp, aud) in verify_token(). BearerAuthBackend rejects expired
tokens with HTTP 401 *before* tool execution and SSE headers are
sent, triggering client refresh.

- JWKSVerifier: async PyJWKClient wrapper (blocking JWKS fetch
  runs in thread pool), configurable cache lifespan, auto-refresh
  on key rotation or fetch errors
- Reconciled _extract_jwt_aud: when JWKS is configured, aud is
  extracted from verified claims; falls back to unverified decode
- Graceful degradation: verification failures return None (don't
  block), letting Dremio handle validation as before
- E2E tests proving the SDK swallows 401 as 200
- 8 unit tests for JWKSVerifier

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: aniket-s-kulkarni

src/dremioai/config/settings.py

@@ -26,7 +26,7 @@
     field_serializer,
     AliasChoices,
 )
-from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings import BaseSettings, SettingsConfigDict, PydanticBaseSettingsSource, YamlConfigSettingsSource
 from typing import (
     Optional,
     Union,
@@ -44,7 +44,7 @@
 from dremioai.config.tools import ToolType
 from enum import auto, StrEnum
 from pathlib import Path
-from yaml import safe_load, add_representer, dump
+from yaml import add_representer, dump
 from functools import reduce
 from operator import ior
 from shutil import which
@@ -253,6 +253,17 @@ class Dremio(FlagAwareModel):
         description="Extract org ID from JWT aud claim for LD context targeting",
     )
     auth_issuer_uri_override: Optional[str] = None
+    jwks_uri: Optional[str] = Field(
+        default=None,
+        description="JWKS endpoint URL for JWT signature verification and expiry checking. "
+        "When set, the MCP server validates token expiry before tool execution "
+        "so expired tokens trigger HTTP 401 and the client's OAuth refresh flow. "
+        "Example: https://your-auth0-tenant.auth0.com/.well-known/jwks.json",
+    )
+    jwks_cache_lifespan: Optional[int] = Field(
+        default=3600,
+        description="How long (seconds) to cache JWKS keys before refetching. Default: 3600 (1 hour).",
+    )
     wlm: Optional[Wlm] = None
     api: Optional[ApiSettings] = Field(default_factory=ApiSettings)
     metrics: Optional[Metrics] = None
@@ -400,6 +411,23 @@ class Settings(FlagAwareMixin, BaseSettings):
         validate_assignment=True,
     )
 
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        return (
+            init_settings,
+            env_settings,
+            dotenv_settings,
+            YamlConfigSettingsSource(settings_cls, yaml_file=_yaml_file),
+            file_secret_settings,
+        )
+
     def model_post_init(self, __context):
         _propagate_flag_prefixes(self, "")
         if self.launchdarkly and self.launchdarkly.sdk_key:
@@ -459,6 +487,10 @@ def collect_flag_keys(model_cls: type, prefix: str = "") -> list[str]:
     return sorted(keys)
 
 
+# Module-level holder so configure() can pass the YAML path to the Settings constructor
+_yaml_file: Path | None = None
+
+
 _settings: ContextVar[Settings] = ContextVar("settings", default=None)
 
 
@@ -498,9 +530,9 @@ def configure(cfg: Union[str, Path] = None, force=False) -> ContextVar[Settings]
         cfg.parent.mkdir(parents=True, exist_ok=True)
         cfg.touch()
 
-    with cfg.open() as f:
-        s = safe_load(f)
-        _settings.set(Settings.model_validate(s if s else {}))
+    global _yaml_file
+    _yaml_file = cfg
+    _settings.set(Settings())
 
     return _settings
 

src/dremioai/servers/jwks_verifier.py

@@ -0,0 +1,119 @@
+#
+#  Copyright (C) 2017-2025 Dremio Corporation
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+"""
+JWKS-based JWT token verifier for the MCP server.
+
+Verifies JWT signatures and extracts claims (``exp``, ``aud``, etc.)
+so that expired tokens are rejected with HTTP 401 *before* any tool
+execution, triggering the MCP client's OAuth token refresh flow.
+
+The JWKS keyset is cached and refreshed automatically on cache miss
+or verification error (e.g. key rotation).
+"""
+
+import asyncio
+from dataclasses import dataclass
+from typing import Optional
+
+import jwt as pyjwt
+from jwt import PyJWKClient, PyJWKClientError, ExpiredSignatureError
+from dremioai import log
+
+logger = log.logger(__name__)
+
+_DEFAULT_JWKS_CACHE_LIFESPAN = 3600  # 1 hour in seconds
+
+
+@dataclass
+class VerifiedClaims:
+    """Subset of JWT claims extracted after signature verification."""
+    exp: Optional[int] = None
+    aud: Optional[str] = None
+
+
+class JWKSVerifier:
+    """Verify JWT tokens using a remote JWKS endpoint.
+
+    Uses ``PyJWKClient`` with built-in caching (``lifespan`` seconds).
+    On verification failure due to a key-related error the cache is
+    invalidated and verification is retried once with fresh keys.
+    """
+
+    def __init__(self, jwks_uri: str, lifespan: int = _DEFAULT_JWKS_CACHE_LIFESPAN):
+        self._jwks_uri = jwks_uri
+        self._lifespan = lifespan
+        self._client = PyJWKClient(
+            jwks_uri,
+            cache_jwk_set=True,
+            lifespan=lifespan,
+        )
+        logger.info(f"JWKS verifier initialised with uri={jwks_uri}, cache={lifespan}s")
+
+    async def verify(self, token: str) -> Optional[VerifiedClaims]:
+        """Verify *token* and return its claims.
+
+        Returns ``None`` when verification cannot be performed — the
+        token will still be forwarded to Dremio for real validation on
+        the tool call.
+
+        Returns ``VerifiedClaims(exp=0)`` for expired tokens so that
+        ``BearerAuthBackend`` rejects them with HTTP 401.
+
+        PyJWKClient.get_signing_key_from_jwt() makes blocking HTTP calls
+        to fetch JWKS on cache miss, so we run it in a thread pool to
+        avoid blocking the event loop.
+        """
+        loop = asyncio.get_running_loop()
+        try:
+            return await loop.run_in_executor(None, self._verify, token)
+        except (pyjwt.InvalidKeyError, PyJWKClientError, KeyError):
+            logger.info("JWKS cache miss or fetch error, refreshing and retrying")
+            try:
+                self._client = PyJWKClient(
+                    self._jwks_uri,
+                    cache_jwk_set=True,
+                    lifespan=self._lifespan,
+                )
+                return await loop.run_in_executor(None, self._verify, token)
+            except Exception:
+                logger.warning("JWKS verification failed after cache refresh", exc_info=True)
+                return None
+        except ExpiredSignatureError:
+            logger.debug("Token expired")
+            return VerifiedClaims(exp=0)
+        except Exception:
+            logger.debug("JWT verification failed, skipping enforcement", exc_info=True)
+            return None
+
+    def _verify(self, token: str) -> VerifiedClaims:
+        signing_key = self._client.get_signing_key_from_jwt(token)
+        claims = pyjwt.decode(
+            token,
+            signing_key.key,
+            algorithms=["RS256"],
+            options={
+                "verify_aud": False,
+                "verify_iss": False,
+                "verify_exp": True,
+            },
+        )
+        aud = claims.get("aud")
+        if isinstance(aud, list):
+            aud = aud[0] if aud else None
+        return VerifiedClaims(
+            exp=claims.get("exp"),
+            aud=aud,
+        )

src/dremioai/servers/mcp.py

@@ -108,12 +108,22 @@ class Transports(StrEnum):
 class FastMCPServerWithAuthToken(FastMCP):
     class DelegatingTokenVerifier(TokenVerifier):
 
+        def __init__(self):
+            self._jwks_verifier = None
+            dremio = settings.instance().dremio
+            jwks_uri = dremio.get("jwks_uri")
+            if jwks_uri:
+                from dremioai.servers.jwks_verifier import JWKSVerifier
+                lifespan = dremio.get("jwks_cache_lifespan") or 3600
+                self._jwks_verifier = JWKSVerifier(jwks_uri, lifespan=lifespan)
+
         @staticmethod
         def _extract_jwt_aud(token: str) -> str | None:
             """Extract the aud claim from a JWT without signature verification.
 
-            The token is always forwarded to Dremio for real validation;
-            we only read the audience (org ID) for LD context targeting.
+            Fallback for when JWKS is not configured. The token is always
+            forwarded to Dremio for real validation; we only read the
+            audience (org ID) for LD context targeting.
             """
             try:
                 import jwt
@@ -125,13 +135,25 @@ def _extract_jwt_aud(token: str) -> str | None:
 
         async def verify_token(self, token: str) -> AccessToken | None:
             if token:
-                if settings.instance().dremio.get("extract_org_id_from_jwt"):
-                    if org_id := self._extract_jwt_aud(token):
-                        FeatureFlagManager.set_org_id(org_id)
+                expires_at = None
+                org_id = None
+
+                if self._jwks_verifier:
+                    verified = await self._jwks_verifier.verify(token)
+                    if verified:
+                        expires_at = verified.exp
+                        org_id = verified.aud
+                elif settings.instance().dremio.get("extract_org_id_from_jwt"):
+                    org_id = self._extract_jwt_aud(token)
+
+                if org_id:
+                    FeatureFlagManager.set_org_id(org_id)
+
                 return AccessToken(
-                    token=token,  # Include the token itself
+                    token=token,
                     client_id="unused-client",
                     scopes=["read"],
+                    expires_at=expires_at,
                 )
             else:
                 log.logger("verify_token").info(f"Token not provided")