Added JWKS based token validation check

ravibits · ravibits · commit 379cb7f56fdb · 2025-05-20T16:58:15.000-07:00
diff --git a/examples/servers/simple-auth-remote/mcp_simple_remote_auth/server.py b/examples/servers/simple-auth-remote/mcp_simple_remote_auth/server.py
@@ -1,35 +1,118 @@
 """Simple MCP Server with GitHub OAuth Authentication."""
 
 import logging
-import secrets
-import time
 from typing import Any, Literal
 
 import click
+import jwt
+import requests
 from pydantic import AnyHttpUrl
 from pydantic_settings import BaseSettings, SettingsConfigDict
-from starlette.exceptions import HTTPException
-from starlette.requests import Request
-from starlette.responses import JSONResponse, RedirectResponse, Response
 
-from mcp.server.auth.middleware.auth_context import get_access_token
 from mcp.server.auth.provider import (
     AccessToken,
-    AuthorizationCode,
-    AuthorizationParams,
-    OAuthAuthorizationServerProvider,
-    RefreshToken,
     TokenValidator,
-    construct_redirect_uri,
 )
 from mcp.server.auth.settings import AuthSettings, ClientRegistrationOptions
 from mcp.server.fastmcp.server import FastMCP
-from mcp.shared._httpx_utils import create_mcp_http_client
-from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
+from mcp.shared.auth import ProtectedResourceMetadata
 
 logger = logging.getLogger(__name__)
 
 
+class TokenValidatorJWT(TokenValidator[AccessToken]):
+    def __init__(self, resource_metadata: ProtectedResourceMetadata):
+        self._resource_metadata = resource_metadata
+
+    async def validate_token(self, token: str) -> AccessToken | None:
+        try:
+            return await self.decode_token(token)
+        except Exception as e:
+            logger.error(f"Token validation failed: {e}")
+            return None
+
+    async def _get_jwks_uri(self, auth_server: str) -> str:
+        """Get the JWKS URI from the OIDC or OAuth well-known configuration.
+
+        Args:
+            auth_server: The base URL of the authorization server
+
+        Returns:
+            The JWKS URI
+
+        Raises:
+            ValueError: If the JWKS URI cannot be found in either OIDC or OAuth
+            well-known configurations
+            requests.RequestException: If there's an error fetching the configuration
+        """
+        well_known_paths = [
+            "/.well-known/openid-configuration",  # OIDC well-known
+            "/.well-known/oauth-authorization-server",  # OAuth well-known
+        ]
+
+        last_error = None
+
+        for path in well_known_paths:
+            try:
+                config_url = f"https://{auth_server}{path}"
+                response = requests.get(
+                    config_url,
+                    timeout=10,  # Add timeout to prevent hanging
+                    headers={"Accept": "application/json"},
+                )
+                response.raise_for_status()  # Raise an exception for bad status codes
+                config = response.json()
+
+                # Try to get JWKS URI from the configuration
+                jwks_uri = config.get("jwks_uri")
+                if jwks_uri:
+                    return jwks_uri
+
+            except requests.RequestException as e:
+                last_error = e
+                logger.debug(f"Failed to fetch {path}: {e}")
+                continue
+
+        # If we get here, we couldn't find a valid JWKS URI
+        error_msg = "Could not find jwks_uri in OIDC or OAuth well-known configurations"
+        logger.error(f"{error_msg}. Last error: {last_error}")
+        raise ValueError(error_msg)
+
+    async def decode_token(self, token: str) -> AccessToken | None:
+        try:
+            auth_server = self._resource_metadata.authorization_servers[0]
+            jwks_uri = await self._get_jwks_uri(auth_server)
+            jwks_client = jwt.PyJWKClient(jwks_uri)
+            signing_key = jwks_client.get_signing_key_from_jwt(token)
+
+            # Rest of your decode_token method remains the same
+            payload = jwt.decode(
+                token,
+                key=signing_key.key,
+                algorithms=["RS256"],
+                audience=self._resource_metadata.resource,
+                issuer=f"https://{auth_server}",
+                options={
+                    "verify_signature": True,
+                    "verify_aud": True,
+                    "verify_iss": True,
+                    "verify_exp": True,
+                    "verify_nbf": True,
+                    "verify_iat": True,
+                },
+            )
+
+            return AccessToken(
+                token=token,
+                client_id=payload["client_id"],
+                scopes=payload["scope"].split(" "),
+                expires_at=payload["exp"],
+            )
+        except Exception as e:
+            logger.error(f"Token validation failed: {e}")
+            return None
+
+
 class ServerSettings(BaseSettings):
     """Settings for the simple GitHub MCP server."""
 
@@ -71,8 +154,18 @@ def create_simple_mcp_server(settings: ServerSettings) -> FastMCP:
         port=settings.port,
         debug=True,
         auth=auth_settings,
-        token_validator=TokenValidator(),
-        protected_resource_metadata={"resource": "asdasd", "authorization_servers": ["https://auth.devramp.ai"], "scopes_supported": ["user"]}
+        token_validator=TokenValidatorJWT(
+            ProtectedResourceMetadata(
+                resource="asdasd",
+                authorization_servers=["https://auth.devramp.ai"],
+                scopes_supported=["user"],
+            )
+        ),
+        protected_resource_metadata={
+            "resource": "asdasd",
+            "authorization_servers": ["https://auth.devramp.ai"],
+            "scopes_supported": ["user"],
+        },
     )
 
     @app.tool()
diff --git a/src/mcp/server/auth/middleware/bearer_auth.py b/src/mcp/server/auth/middleware/bearer_auth.py
@@ -10,7 +10,11 @@
 from starlette.requests import HTTPConnection
 from starlette.types import Receive, Scope, Send
 
-from mcp.server.auth.provider import AccessToken, OAuthAuthorizationServerProvider, TokenValidator
+from mcp.server.auth.provider import (
+    AccessToken,
+    OAuthAuthorizationServerProvider,
+    TokenValidator,
+)
 
 
 class AuthenticatedUser(SimpleUser):
diff --git a/src/mcp/server/auth/provider.py b/src/mcp/server/auth/provider.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Generic, Literal, Protocol, TypeVar, Any
+from typing import Generic, Literal, Protocol, TypeVar
 from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
 
 from pydantic import AnyHttpUrl, BaseModel
@@ -10,14 +10,6 @@
 )
 
 
-# Define type variables
-AccessTokenT = TypeVar('AccessTokenT', bound='AccessToken')
-
-class TokenValidator(Generic[AccessTokenT], BaseModel):
-    async def validate_token(self, token: str) -> AccessTokenT | None:
-        ...
-
-
 class AuthorizationParams(BaseModel):
     state: str | None
     scopes: list[str] | None
@@ -104,6 +96,13 @@ class TokenError(Exception):
 AccessTokenT = TypeVar("AccessTokenT", bound=AccessToken)
 
 
+
+class TokenValidator(BaseModel, Generic[AccessTokenT]):
+    async def validate_token(self, token: str) -> AccessTokenT | None:
+        ...
+
+
+
 class OAuthAuthorizationServerProvider(
     Protocol, Generic[AuthorizationCodeT, RefreshTokenT, AccessTokenT]
 ):
diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py
@@ -5,10 +5,7 @@
 import inspect
 import re
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
-from contextlib import (
-    AbstractAsyncContextManager,
-    asynccontextmanager,
-)
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
 from itertools import chain
 from typing import Any, Generic, Literal
 
@@ -18,9 +15,9 @@
 from pydantic.networks import AnyUrl
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from starlette.applications import Starlette
+from starlette.exceptions import HTTPException
 from starlette.middleware import Middleware
 from starlette.middleware.authentication import AuthenticationMiddleware
-from starlette.exceptions import HTTPException
 from starlette.requests import Request
 from starlette.responses import Response
 from starlette.routing import Mount, Route
@@ -32,12 +29,12 @@
     JWTBearerTokenAuthBackend,
     RequireAuthMiddleware,
 )
-from mcp.server.auth.provider import AccessToken, TokenValidator
-from mcp.server.auth.provider import OAuthAuthorizationServerProvider
-from mcp.shared.auth import ProtectedResourceMetadata
-from mcp.server.auth.settings import (
-    AuthSettings,
+from mcp.server.auth.provider import (
+    AccessToken,
+    OAuthAuthorizationServerProvider,
+    TokenValidator,
 )
+from mcp.server.auth.settings import AuthSettings
 from mcp.server.fastmcp.exceptions import ResourceError
 from mcp.server.fastmcp.prompts import Prompt, PromptManager
 from mcp.server.fastmcp.resources import FunctionResource, Resource, ResourceManager
@@ -53,6 +50,7 @@
 from mcp.server.stdio import stdio_server
 from mcp.server.streamable_http import EventStore
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
+from mcp.shared.auth import ProtectedResourceMetadata
 from mcp.shared.context import LifespanContextT, RequestContext
 from mcp.types import (
     AnyFunction,
@@ -143,8 +141,9 @@ def __init__(
         name: str | None = None,
         instructions: str | None = None,
         auth_server_details: dict[str, Any] | None = None,
-        auth_server_provider: OAuthAuthorizationServerProvider[Any, Any, Any]
-        | None = None,
+        auth_server_provider: (
+            OAuthAuthorizationServerProvider[Any, Any, Any] | None
+        ) = None,
         protected_resource_metadata: dict[str, Any] | None = None,
         event_store: EventStore | None = None,
         token_validator: TokenValidator[AccessToken] | None = None,
@@ -154,9 +153,11 @@ def __init__(
         self._auth_server_details = auth_server_details
         self._protected_resource_metadata = None
         if protected_resource_metadata:
-            self._protected_resource_metadata = ProtectedResourceMetadata(**protected_resource_metadata)
+            self._protected_resource_metadata = ProtectedResourceMetadata(
+                **protected_resource_metadata
+            )
         self._token_validator = token_validator
-       
+
         self._mcp_server = MCPServer(
             name=name or "FastMCP",
             instructions=instructions,
@@ -176,7 +177,9 @@ def __init__(
             warn_on_duplicate_prompts=self.settings.warn_on_duplicate_prompts
         )
         # don't do this check if protected_resource_metadata is not None
-        if (self.settings.auth is not None) != (auth_server_provider is not None) and self._protected_resource_metadata is None:
+        if (self.settings.auth is not None) != (
+            auth_server_provider is not None
+        ) and self._protected_resource_metadata is None:
             # TODO: after we support separate authorization servers (see
             # https://github.com/modelcontextprotocol/modelcontextprotocol/pull/284)
             # we should validate that if auth is enabled, we have either an
@@ -228,8 +231,13 @@ def session_manager(self) -> StreamableHTTPSessionManager:
     async def _serve_protected_resource_metadata(self, request: Request) -> Response:
         """Serve the OAuth protected resource metadata."""
         if not self._protected_resource_metadata:
-            raise HTTPException(status_code=404, detail="Protected resource metadata not configured")
-        return Response(self._protected_resource_metadata.model_dump_json(), media_type="application/json")
+            raise HTTPException(
+                status_code=404, detail="Protected resource metadata not configured"
+            )
+        return Response(
+            self._protected_resource_metadata.model_dump_json(),
+            media_type="application/json",
+        )
 
     def run(
         self,
@@ -706,11 +714,10 @@ async def handle_sse(scope: Scope, receive: Receive, send: Send):
 
         # Create routes
         routes: list[Route | Mount] = []
-        
+
         middleware: list[Middleware] = []
         required_scopes = []
 
-
         # Add auth endpoints if auth provider is configured
         if self._auth_server_provider:
             assert self.settings.auth
@@ -808,21 +815,20 @@ async def handle_streamable_http(
         routes: list[Route | Mount] = []
         middleware: list[Middleware] = []
         required_scopes = []
-        print("Protected resource metadata: ", self._protected_resource_metadata)
         if self._protected_resource_metadata and self._token_validator:
-            print("Adding protected resource metadata route")
-            # only add the well-known route if the protected resource metadata is configured
+            # only add the well-known route if the protected resource metadata is
+            # configured
             routes.append(
                 Route(
                     "/.well-known/oauth-protected-resource",
                     self._serve_protected_resource_metadata,
-                    methods=["GET"]
+                    methods=["GET"],
                 )
             )
-            # by default assuming that this would be a JWT Bearer Token; 
-            # Make this also optional somehow; may be as part of the protected resource metadata, 
-            # take a class for validting the token
-            middleware= [
+            # by default assuming that this would be a JWT Bearer Token;
+            # Make this also optional somehow; may be as part of the protected resource
+            # metadata, take a class for validating the token
+            middleware = [
                 Middleware(
                     AuthenticationMiddleware,
                     backend=JWTBearerTokenAuthBackend(
diff --git a/src/mcp/shared/auth.py b/src/mcp/shared/auth.py
@@ -2,6 +2,7 @@
 
 from pydantic import AnyHttpUrl, BaseModel, Field
 
+
 class ProtectedResourceMetadata(BaseModel):
     # create a pydantic model with required params as resource, authorization_servers
     resource: str
