Merge pull request #9 from keycardai/feat/starlette-middleware

Feat/starlette middleware

Author: kamil-keycard

packages/mcp/pyproject.toml

@@ -7,9 +7,11 @@ requires-python = ">=3.10"
 license = { text = "MIT" }
 authors = [{ name = "KeyCard AI", email = "support@keycard.ai" }]
 dependencies = [
-    "keycardai-oauth==0.3.0",
-    "mcp==1.14.0",
+    "keycardai-oauth>=0.4.0,<1.0.0",
+    "mcp>=1.13.1",
     "pydantic>=2.11.7",
+    "httpx>=0.27.2",
+    "starlette>=0.47.3",
 ]
 keywords = ["mcp", "model-context-protocol", "authentication", "authorization", "ai", "llm"]
 classifiers = [

packages/mcp/src/keycardai/mcp/server/auth/provider.py

@@ -2,10 +2,12 @@
 from typing import Any
 
 from mcp.server.auth.provider import AccessToken
+from pydantic import AnyHttpUrl
 
+from keycardai.oauth import Client
 from keycardai.oauth.utils.jwt import (
     get_header,
-    get_verification_key,
+    get_jwks_key,
     parse_jwt_access_token,
 )
 
@@ -22,15 +24,19 @@ def __init__(
         jwks_uri: str | None = None,
         allowed_algorithms: list[str] = None,
         cache_ttl: int = 300,  # 5 minutes default
+        enable_multi_zone: bool = False,
     ):
         """Initialize the KeyCard token verifier.
 
         Args:
-            issuer: Expected token issuer (required)
+            issuer: Expected token issuer (required). When enable_multi_zone=True,
+                   this should be the top-level domain URL that will be used as base
+                   for zone-specific issuer construction.
             required_scopes: Required scopes for token validation
-            jwks_uri: JWKS endpoint URL for key fetching
+            jwks_uri: JWKS endpoint URL for key fetching (deprecated, use issuer)
             allowed_algorithms: JWT algorithms (default RS256)
             cache_ttl: JWKS cache TTL in seconds (default 300 = 5 minutes)
+            enable_multi_zone: Enable multi-zone support where issuer is top-level domain
         """
         if not issuer:
             raise ValueError("Issuer is required for token verification")
@@ -43,30 +49,99 @@ def __init__(
         self.cache_ttl = cache_ttl
 
         self._jwks_cache = JWKSCache(ttl=cache_ttl, max_size=10)
+        self._discovered_jwks_uri: str | None = None
+        self._discovered_jwks_uris: dict[str, str] = {}  # Initialize the cache dict
 
-    async def _get_verification_key(self, token: str) -> JWKSKey:
-        """Get the verification key for the token with caching."""
-        if not self.jwks_uri:
-            raise ValueError("JWKS URI not configured")
+        self.enable_multi_zone = enable_multi_zone
+
+    def _discover_jwks_uri(self, zone_id: str | None = None) -> str:
+        """Discover JWKS URI from issuer lazily.
+
+        Args:
+            zone_id: Zone ID for multi-zone scenarios. When provided with enable_multi_zone=True,
+                    constructs zone-specific issuer URL for discovery.
+        """
+        cache_key = f"{zone_id or 'default'}"
+        cached_uri = self._discovered_jwks_uris.get(cache_key)
+        if cached_uri is not None:
+            return cached_uri
+
+        if self.jwks_uri:
+            self._discovered_jwks_uris[cache_key] = self.jwks_uri
+            return self.jwks_uri
+
+        discovery_issuer = self.issuer
+        if self.enable_multi_zone and zone_id:
+            discovery_issuer = self._create_zone_scoped_url(self.issuer, zone_id)
+
+        try:
+            with Client(discovery_issuer) as client:
+                server_metadata = client.discover_server_metadata()
+                discovered_uri = server_metadata.jwks_uri
 
+            if not discovered_uri:
+                raise ValueError(f"Could not discover JWKS URI from issuer: {discovery_issuer}")
+
+            # Cache the successful discovery
+            self._discovered_jwks_uris[cache_key] = discovered_uri
+            return discovered_uri
+
+        except Exception as e:
+            # Don't cache failures, let them retry
+            raise ValueError(f"Could not discover JWKS URI from issuer {discovery_issuer}: {e}") from e
+
+    def _create_zone_scoped_url(self, base_url: str, zone_id: str) -> str:
+        """Create zone-scoped URL by prepending zone_id to the host."""
+        base_url_obj = AnyHttpUrl(base_url)
+
+        port_part = ""
+        if base_url_obj.port and not (
+            (base_url_obj.scheme == "https" and base_url_obj.port == 443) or
+            (base_url_obj.scheme == "http" and base_url_obj.port == 80)
+        ):
+            port_part = f":{base_url_obj.port}"
+
+        zone_url = f"{base_url_obj.scheme}://{zone_id}.{base_url_obj.host}{port_part}"
+        return zone_url
+
+    def _get_kid_and_algorithm(self, token: str) -> tuple[str, str]:
         header = get_header(token)
         kid = header.get("kid")
         algorithm = header.get("alg")
         if algorithm not in self.allowed_algorithms:
             raise ValueError(f"Unsupported algorithm: {algorithm}")
+        return [kid, algorithm]
+
+    def _get_zone_jwks_uri(self, jwks_uri: str, zone_id: str) -> str:
+        jwks_url = AnyHttpUrl(jwks_uri)
+        jwks_zone_host = jwks_url.host.replace(jwks_url.host, f"{zone_id}.{jwks_url.host}")
+        jwks_url.host = jwks_zone_host
+        return jwks_url.to_string()
+
+    async def _get_verification_key(self, token: str, zone_id: str | None = None) -> JWKSKey:
+        """Get the verification key for the token with caching."""
+        kid, algorithm = self._get_kid_and_algorithm(token)
 
         cached_key = self._jwks_cache.get_key(kid)
         if cached_key is not None:
             return cached_key
 
-        verification_key = await get_verification_key(token, self.jwks_uri)
+        if self.enable_multi_zone and zone_id:
+            jwks_uri = self._discover_jwks_uri(zone_id)
+        else:
+            jwks_uri = self._discover_jwks_uri()
+            if zone_id:
+                jwks_uri = self._get_zone_jwks_uri(jwks_uri, zone_id)
+
+        verification_key = await get_jwks_key(kid, jwks_uri)
 
         self._jwks_cache.set_key(kid, verification_key, algorithm)
         cached_key = self._jwks_cache.get_key(kid)
         if cached_key is None:
             raise ValueError("Failed to cache verification key")
         return cached_key
 
+
     def clear_cache(self) -> None:
         """Clear the JWKS key cache."""
         self._jwks_cache.clear()
@@ -79,37 +154,25 @@ def get_cache_stats(self) -> dict[str, Any]:
         """
         return self._jwks_cache.get_stats()
 
-    async def verify_token(self, token: str) -> AccessToken | None:
-        """Verify a JWT token and return AccessToken if valid.
-
-        Performs JWT verification including:
-        - Parse token into structured JWTAccessToken model internally
-        - Validate token expiration
-        - Validate issuer if configured
-        - Validate required scopes if configured
-        - Convert to AccessToken format for return
-
-        Note: This is a simplified implementation that does not perform
-        cryptographic signature verification. For production use, proper
-        signature verification should be implemented.
-
-        Args:
-            token: JWT token string to verify
-
-        Returns:
-            AccessToken object if valid, None if invalid
-        """
-        try:
-            verification_key = await self._get_verification_key(token)
+    async def verify_token_for_zone(self, token: str, zone_id: str) -> AccessToken | None:
+        """Verify a JWT token for a specific zone and return AccessToken if valid."""
+        key = await self._get_verification_key(token, zone_id)
+        return self._verify_token(token, key, zone_id)
 
+    def _verify_token(self, token: str, key: JWKSKey, zone_id: str | None = None) -> AccessToken | None:
             jwt_access_token = parse_jwt_access_token(
-                token, verification_key.key, verification_key.algorithm
+                token, key.key, key.algorithm
             )
 
             if jwt_access_token.exp < time.time():
                 return None
 
-            if jwt_access_token.iss != self.issuer:
+            # Validate issuer, handling multi-zone scenarios
+            expected_issuer = self.issuer
+            if self.enable_multi_zone and zone_id:
+                expected_issuer = self._create_zone_scoped_url(self.issuer, zone_id)
+
+            if jwt_access_token.iss != expected_issuer:
                 return None
 
             if self.required_scopes:
@@ -133,6 +196,31 @@ async def verify_token(self, token: str) -> AccessToken | None:
                 resource=jwt_access_token.get_custom_claim("resource"),
             )
 
+
+    async def verify_token(self, token: str) -> AccessToken | None:
+        """Verify a JWT token and return AccessToken if valid.
+
+        Performs JWT verification including:
+        - Parse token into structured JWTAccessToken model internally
+        - Validate token expiration
+        - Validate issuer if configured
+        - Validate required scopes if configured
+        - Convert to AccessToken format for return
+
+        Note: This is a simplified implementation that does not perform
+        cryptographic signature verification. For production use, proper
+        signature verification should be implemented.
+
+        Args:
+            token: JWT token string to verify
+
+        Returns:
+            AccessToken object if valid, None if invalid
+        """
+        try:
+            key = await self._get_verification_key(token)
+            return self._verify_token(token, key)
+
         except Exception:
             return None
 

packages/mcp/src/keycardai/mcp/server/auth/verifier.py

@@ -0,0 +1,11 @@
+from .metadata import (
+    InferredProtectedResourceMetadata,
+    authorization_server_metadata,
+    protected_resource_metadata,
+)
+
+__all__ = [
+    "protected_resource_metadata",
+    "authorization_server_metadata",
+    "InferredProtectedResourceMetadata",
+]

packages/mcp/src/keycardai/mcp/server/handlers/__init__.py

@@ -0,0 +1,106 @@
+import json
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import httpx
+from mcp.shared.auth import ProtectedResourceMetadata
+from pydantic import AnyHttpUrl, Field
+from starlette.requests import Request
+from starlette.responses import Response
+
+
+class InferredProtectedResourceMetadata(ProtectedResourceMetadata):
+    """Extended ProtectedResourceMetadata that allows resource to be inferred from request."""
+    resource: AnyHttpUrl | None = Field(default=None)  # Override to make it optional
+
+@dataclass
+class AuthorizationServerMetadata:
+    base_url: str
+
+
+def _is_authorization_server_zone_scoped(authorization_server_urls: AnyHttpUrl) -> bool:
+    if len(authorization_server_urls) != 1:
+        return False
+    return len(authorization_server_urls[0].host.split(".")) == 3
+
+def _get_zone_id_from_path(path: str) -> str | None:
+    path = path.lstrip("/").rstrip("/")
+    zone_id = path.split("/")[0]
+    if zone_id == "" or zone_id == "/":
+        return None
+    return zone_id
+
+def _remove_well_known_prefix(path: str) -> str:
+    prefix = ".well-known/oauth-protected-resource"
+    path = path.lstrip("/").rstrip("/")
+    if path.startswith(prefix):
+        return path[len(prefix):]
+    return path
+
+def _create_zone_scoped_authorization_server_url(zone_id: str, authorization_server_url: AnyHttpUrl) -> AnyHttpUrl:
+    port_part = f":{authorization_server_url.port}" if authorization_server_url.port else ""
+    url = f"{authorization_server_url.scheme}://{zone_id}.{authorization_server_url.host}{port_part}"
+    return AnyHttpUrl(url)
+
+def _strip_zone_id_from_path(zone_id: str, path: str) -> str:
+    path = path.lstrip("/").rstrip("/")
+    if path.startswith(zone_id):
+        return path[len(zone_id):]
+    return path
+
+def _create_resource_url(base_url: str | AnyHttpUrl, path: str) -> AnyHttpUrl:
+    base_url_str = str(base_url).rstrip("/")
+    if path and not path.startswith("/"):
+        path = "/" + path
+    url = f"{base_url_str}{path}".rstrip("/")
+    if url.endswith("://") or (path == "/" and not url.endswith("/")):
+        url += "/"
+    return AnyHttpUrl(url)
+
+def _remove_authorization_server_prefix(path: str) -> str:
+    """Remove the /.well-known/oauth-authorization-server prefix from the path."""
+    auth_server_prefix = "/.well-known/oauth-authorization-server"
+    if path.startswith(auth_server_prefix):
+        return path[len(auth_server_prefix):]
+    return path
+
+def protected_resource_metadata(metadata: InferredProtectedResourceMetadata, enable_multi_zone: bool = False) -> Callable:
+    def wrapper(request: Request) -> Response:
+        # Create a copy of the metadata to avoid mutating the original
+        request_metadata = metadata.model_copy(deep=True)
+        path = _remove_well_known_prefix(request.url.path)
+        if enable_multi_zone or not _is_authorization_server_zone_scoped(request_metadata.authorization_servers):
+            zone_id = _get_zone_id_from_path(path)
+            if zone_id:
+                request_metadata.authorization_servers = [ _create_zone_scoped_authorization_server_url(zone_id, request_metadata.authorization_servers[0]) ]
+
+        resource = _create_resource_url(request.base_url, path)
+        mcp_version = request.headers.get("mcp-protocol-version")
+        request_metadata.resource = resource
+        # TODO: what is the reason for this?
+        if mcp_version == "2025-03-26":
+            json["authorization_servers"] = [ request.base_url ]
+        return Response(content=request_metadata.model_dump_json(exclude_none=True), status_code=200)
+    return wrapper
+
+def authorization_server_metadata(issuer: str, enable_multi_zone: bool = False) -> Callable:
+    def wrapper(request: Request) -> Response:
+        try:
+            actual_issuer = issuer
+            path = _remove_authorization_server_prefix(request.url.path)
+
+            if enable_multi_zone or not _is_authorization_server_zone_scoped([AnyHttpUrl(issuer)]):
+                zone_id = _get_zone_id_from_path(path)
+                if zone_id:
+                    actual_issuer = str(_create_zone_scoped_authorization_server_url(zone_id, AnyHttpUrl(issuer)))
+
+            with httpx.Client() as client:
+                resp = client.get(f"{actual_issuer}/.well-known/oauth-authorization-server")
+                resp.raise_for_status()
+                authorization_server_metadata = resp.json()
+                authorization_server_metadata["authorization_endpoint"] = f"{request.base_url}{authorization_server_metadata['authorization_endpoint']}"
+                return Response(content=json.dumps(authorization_server_metadata), status_code=200)
+        except Exception as e:
+            error_message = {"error": str(e), "type": type(e).__name__}
+            return Response(content=json.dumps(error_message), status_code=500)
+    return wrapper

packages/mcp/src/keycardai/mcp/server/handlers/metadata.py

@@ -0,0 +1,5 @@
+from .bearer import BearerAuthMiddleware
+
+__all__ = [
+    "BearerAuthMiddleware",
+]