From 20ae68bb7dce5506030498093df5c425c848a54b Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 1 Jun 2025 11:13:21 -0400 Subject: [PATCH 1/5] Add basic bearer auth for server and client --- pyproject.toml | 1 + src/fastmcp/client/auth.py | 9 + src/fastmcp/server/auth/bearer.py | 256 ++++++++++++ src/fastmcp/server/auth/providers/__init__.py | 0 src/fastmcp/server/auth/providers/bearer.py | 359 +++++++++++++++++ .../in_memory.py} | 7 +- src/fastmcp/server/dependencies.py | 10 + tests/auth/providers/test_bearer.py | 366 ++++++++++++++++++ tests/auth/test_oauth_client.py | 4 +- uv.lock | 58 +++ 10 files changed, 1067 insertions(+), 3 deletions(-) create mode 100644 src/fastmcp/server/auth/bearer.py create mode 100644 src/fastmcp/server/auth/providers/__init__.py create mode 100644 src/fastmcp/server/auth/providers/bearer.py rename src/fastmcp/server/auth/{in_memory_provider.py => providers/in_memory.py} (98%) create mode 100644 tests/auth/providers/test_bearer.py diff --git a/pyproject.toml b/pyproject.toml index a55e60254..35b538b58 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ dev = [ "ipython>=8.12.3", "pdbpp>=0.10.3", "pre-commit", + "pyinstrument>=5.0.2", "pyright>=1.1.389", "pytest>=8.3.3", "pytest-asyncio>=0.23.5", diff --git a/src/fastmcp/client/auth.py b/src/fastmcp/client/auth.py index 43df9d442..28c2bba84 100644 --- a/src/fastmcp/client/auth.py +++ b/src/fastmcp/client/auth.py @@ -392,3 +392,12 @@ def OAuth( ) return oauth_provider + + +class BearerAuth(httpx.Auth): + def __init__(self, token: str): + self.token = token + + def auth_flow(self, request): + request.headers["Authorization"] = f"Bearer {self.token}" + yield request diff --git a/src/fastmcp/server/auth/bearer.py b/src/fastmcp/server/auth/bearer.py new file mode 100644 index 000000000..729fcf1dc --- /dev/null +++ b/src/fastmcp/server/auth/bearer.py @@ -0,0 +1,256 @@ +""" +Simple JWT Bearer Token validation for hosted MCP servers. + +Uses RS256 (asymmetric) where your control plane signs with a private key +and hosted MCP servers validate with the corresponding public key. + +Example usage: +# Static public key +provider = BearerTokenValidatorProvider( + public_key='''-----BEGIN PUBLIC KEY----- + MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA... + -----END PUBLIC KEY-----''', + issuer="https://auth.yourservice.com" +) + +# Or JWKS URI (recommended for production - allows key rotation) +provider = BearerTokenValidatorProvider( + jwks_uri="https://auth.yourservice.com/.well-known/jwks.json", + issuer="https://auth.yourservice.com" +) +""" + +import time +from typing import Any + +import httpx +from authlib.jose import JsonWebKey, JsonWebToken +from authlib.jose.errors import JoseError +from mcp.server.auth.provider import ( + AccessToken, + AuthorizationCode, + AuthorizationParams, + RefreshToken, +) +from mcp.shared.auth import ( + OAuthClientInformationFull, + OAuthToken, +) + +from fastmcp.server.auth.auth import ( + ClientRegistrationOptions, + OAuthProvider, + RevocationOptions, +) + + +class BearerTokenValidatorProvider(OAuthProvider): + """ + Simple JWT Bearer Token validator for hosted MCP servers. + Uses RS256 asymmetric encryption. Supports either static public key + or JWKS URI for key rotation. + """ + + def __init__( + self, + issuer: str, + public_key: str | None = None, + jwks_uri: str | None = None, + audience: str | None = None, + required_scopes: list[str] | None = None, + ): + """ + Initialize the provider. + + Args: + issuer: Expected issuer claim (your control plane) + public_key: RSA public key in PEM format (for static key) + jwks_uri: URI to fetch keys from (for key rotation) + audience: Expected audience claim (optional) + required_scopes: List of required scopes for access + """ + if not (public_key or jwks_uri): + raise ValueError("Either public_key or jwks_uri must be provided") + if public_key and jwks_uri: + raise ValueError("Provide either public_key or jwks_uri, not both") + + super().__init__( + issuer_url=issuer, + client_registration_options=ClientRegistrationOptions(enabled=False), + revocation_options=RevocationOptions(enabled=False), + required_scopes=required_scopes, + ) + + self.issuer = issuer + self.audience = audience + self.public_key = public_key + self.jwks_uri = jwks_uri + self.jwt = JsonWebToken(["RS256"]) + + # Simple JWKS cache + self._jwks_cache: dict[str, str] = {} + self._jwks_cache_time: float = 0 + self._cache_ttl = 3600 # 1 hour + + async def _get_verification_key(self, token: str) -> str: + """Get the verification key for the token.""" + if self.public_key: + return self.public_key + + # Extract kid from token header for JWKS lookup + try: + import base64 + import json + + header_b64 = token.split(".")[0] + header_b64 += "=" * (4 - len(header_b64) % 4) # Add padding + header = json.loads(base64.urlsafe_b64decode(header_b64)) + kid = header.get("kid") + + if not kid: + raise ValueError("Token missing key ID (kid)") + + return await self._get_jwks_key(kid) + + except Exception as e: + raise ValueError(f"Failed to extract key ID from token: {e}") + + async def _get_jwks_key(self, kid: str) -> str: + """Fetch key from JWKS with simple caching.""" + if not self.jwks_uri: + raise ValueError("JWKS URI not configured") + + current_time = time.time() + + # Check cache + if ( + current_time - self._jwks_cache_time < self._cache_ttl + and kid in self._jwks_cache + ): + return self._jwks_cache[kid] + + # Fetch JWKS + try: + async with httpx.AsyncClient() as client: + response = await client.get(self.jwks_uri) + response.raise_for_status() + jwks_data = response.json() + + # Cache all keys + self._jwks_cache = {} + for key_data in jwks_data.get("keys", []): + key_kid = key_data.get("kid") + if key_kid: + jwk = JsonWebKey.import_key(key_data) + self._jwks_cache[key_kid] = jwk.get_public_key() + + self._jwks_cache_time = current_time + + if kid not in self._jwks_cache: + raise ValueError(f"Key ID '{kid}' not found in JWKS") + + return self._jwks_cache[kid] + + except Exception as e: + raise ValueError(f"Failed to fetch JWKS: {e}") + + async def load_access_token(self, token: str) -> AccessToken | None: + """ + Validates the provided JWT bearer token. + + Args: + token: The JWT token string to validate + + Returns: + AccessToken object if valid, None if invalid or expired + """ + try: + # Get verification key (static or from JWKS) + verification_key = await self._get_verification_key(token) + + # Decode and verify the JWT token + claims = self.jwt.decode(token, verification_key) + + # Validate expiration + exp = claims.get("exp") + if exp and exp < time.time(): + return None + + # Validate issuer + if claims.get("iss") != self.issuer: + return None + + # Validate audience if configured + if self.audience: + aud = claims.get("aud") + if isinstance(aud, list): + if self.audience not in aud: + return None + elif aud != self.audience: + return None + + # Extract claims + client_id = claims.get("sub") or claims.get("client_id") or "unknown" + scopes = self._extract_scopes(claims) + + return AccessToken( + token=token, + client_id=str(client_id), + scopes=scopes, + expires_at=int(exp) if exp else None, + ) + + except JoseError: + return None + except Exception: + return None + + def _extract_scopes(self, claims: dict[str, Any]) -> list[str]: + """Extract scopes from JWT claims.""" + scope_claim = claims.get("scope", "") + if isinstance(scope_claim, str): + return scope_claim.split() + elif isinstance(scope_claim, list): + return scope_claim + return [] + + # --- Unused OAuth server methods --- + async def get_client(self, client_id: str) -> OAuthClientInformationFull | None: + raise NotImplementedError("Client management not supported") + + async def register_client(self, client_info: OAuthClientInformationFull) -> None: + raise NotImplementedError("Client registration not supported") + + async def authorize( + self, client: OAuthClientInformationFull, params: AuthorizationParams + ) -> str: + raise NotImplementedError("Authorization flow not supported") + + async def load_authorization_code( + self, client: OAuthClientInformationFull, authorization_code: str + ) -> AuthorizationCode | None: + raise NotImplementedError("Authorization code flow not supported") + + async def exchange_authorization_code( + self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode + ) -> OAuthToken: + raise NotImplementedError("Authorization code exchange not supported") + + async def load_refresh_token( + self, client: OAuthClientInformationFull, refresh_token: str + ) -> RefreshToken | None: + raise NotImplementedError("Refresh token flow not supported") + + async def exchange_refresh_token( + self, + client: OAuthClientInformationFull, + refresh_token: RefreshToken, + scopes: list[str], + ) -> OAuthToken: + raise NotImplementedError("Refresh token exchange not supported") + + async def revoke_token( + self, + token: AccessToken | RefreshToken, + ) -> None: + raise NotImplementedError("Token revocation not supported") diff --git a/src/fastmcp/server/auth/providers/__init__.py b/src/fastmcp/server/auth/providers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/fastmcp/server/auth/providers/bearer.py b/src/fastmcp/server/auth/providers/bearer.py new file mode 100644 index 000000000..3cedfee7b --- /dev/null +++ b/src/fastmcp/server/auth/providers/bearer.py @@ -0,0 +1,359 @@ +""" +Simple JWT Bearer Token validation for hosted MCP servers. + +Uses RS256 (asymmetric) where your control plane signs with a private key +and hosted MCP servers validate with the corresponding public key. + +Example usage: +# Static public key +provider = BearerTokenValidatorProvider( + public_key='''-----BEGIN PUBLIC KEY----- + MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA... + -----END PUBLIC KEY-----''', + issuer="https://auth.yourservice.com" +) + +# Or JWKS URI (recommended for production - allows key rotation) +provider = BearerTokenValidatorProvider( + jwks_uri="https://auth.yourservice.com/.well-known/jwks.json", + issuer="https://auth.yourservice.com" +) +""" + +import time +from dataclasses import dataclass +from typing import Any + +import httpx +from authlib.jose import JsonWebKey, JsonWebToken +from authlib.jose.errors import JoseError +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from mcp.server.auth.provider import ( + AccessToken, + AuthorizationCode, + AuthorizationParams, + RefreshToken, +) +from mcp.shared.auth import ( + OAuthClientInformationFull, + OAuthToken, +) +from pydantic import SecretStr + +from fastmcp.server.auth.auth import ( + ClientRegistrationOptions, + OAuthProvider, + RevocationOptions, +) + + +@dataclass(frozen=True, kw_only=True, repr=False) +class RSAKeyPair: + private_key: SecretStr + public_key: str + + @classmethod + def generate(cls) -> "RSAKeyPair": + """ + Generate an RSA key pair for testing. + + Returns: + tuple: (private_key_pem, public_key_pem) + """ + # Generate private key + private_key = rsa.generate_private_key( + public_exponent=65537, + key_size=2048, + ) + + # Get public key + public_key = private_key.public_key() + + # Serialize private key to PEM format + private_pem = private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ).decode("utf-8") + + # Serialize public key to PEM format + public_pem = public_key.public_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PublicFormat.SubjectPublicKeyInfo, + ).decode("utf-8") + + return cls( + private_key=SecretStr(private_pem), + public_key=public_pem, + ) + + def create_token( + self, + subject: str = "fastmcp-user", + issuer: str = "https://fastmcp.example.com", + audience: str | None = None, + scopes: list[str] | None = None, + expires_in_seconds: int = 3600, + additional_claims: dict[str, Any] | None = None, + ) -> str: + """ + Generate a test JWT token for testing purposes. + + Args: + private_key_pem: RSA private key in PEM format + subject: Subject claim (usually user ID) + issuer: Issuer claim + audience: Audience claim (optional) + scopes: List of scopes to include + expires_in_seconds: Token expiration time in seconds + additional_claims: Any additional claims to include + + Returns: + Signed JWT token string + """ + jwt = JsonWebToken(["RS256"]) + + now = int(time.time()) + + # Build payload + payload = { + "iss": issuer, + "sub": subject, + "iat": now, + "exp": now + expires_in_seconds, + } + + if audience: + payload["aud"] = audience + + if scopes: + payload["scope"] = " ".join(scopes) + + if additional_claims: + payload.update(additional_claims) + + # Create header + header = {"alg": "RS256"} + + # Sign and return token + token_bytes = jwt.encode( + header, + payload, + key=self.private_key.get_secret_value(), + ) + + return token_bytes.decode("utf-8") + + +class BearerAuthProvider(OAuthProvider): + """ + Simple JWT Bearer Token validator for hosted MCP servers. + Uses RS256 asymmetric encryption. Supports either static public key + or JWKS URI for key rotation. + """ + + def __init__( + self, + issuer: str | None = None, + public_key: str | None = None, + jwks_uri: str | None = None, + audience: str | None = None, + required_scopes: list[str] | None = None, + ): + """ + Initialize the provider. + + Args: + issuer: Expected issuer claim (your control plane) + public_key: RSA public key in PEM format (for static key) + jwks_uri: URI to fetch keys from (for key rotation) + audience: Expected audience claim (optional) + required_scopes: List of required scopes for access + """ + if not (public_key or jwks_uri): + raise ValueError("Either public_key or jwks_uri must be provided") + if public_key and jwks_uri: + raise ValueError("Provide either public_key or jwks_uri, not both") + + super().__init__( + issuer_url=issuer or "http://fastmcp.example.com", + client_registration_options=ClientRegistrationOptions(enabled=False), + revocation_options=RevocationOptions(enabled=False), + required_scopes=required_scopes, + ) + + self.issuer = issuer + self.audience = audience + self.public_key = public_key + self.jwks_uri = jwks_uri + self.jwt = JsonWebToken(["RS256"]) + + # Simple JWKS cache + self._jwks_cache: dict[str, str] = {} + self._jwks_cache_time: float = 0 + self._cache_ttl = 3600 # 1 hour + + async def _get_verification_key(self, token: str) -> str: + """Get the verification key for the token.""" + if self.public_key: + return self.public_key + + # Extract kid from token header for JWKS lookup + try: + import base64 + import json + + header_b64 = token.split(".")[0] + header_b64 += "=" * (4 - len(header_b64) % 4) # Add padding + header = json.loads(base64.urlsafe_b64decode(header_b64)) + kid = header.get("kid") + + if not kid: + raise ValueError("Token missing key ID (kid)") + + return await self._get_jwks_key(kid) + + except Exception as e: + raise ValueError(f"Failed to extract key ID from token: {e}") + + async def _get_jwks_key(self, kid: str) -> str: + """Fetch key from JWKS with simple caching.""" + if not self.jwks_uri: + raise ValueError("JWKS URI not configured") + + current_time = time.time() + + # Check cache + if ( + current_time - self._jwks_cache_time < self._cache_ttl + and kid in self._jwks_cache + ): + return self._jwks_cache[kid] + + # Fetch JWKS + try: + async with httpx.AsyncClient() as client: + response = await client.get(self.jwks_uri) + response.raise_for_status() + jwks_data = response.json() + + # Cache all keys + self._jwks_cache = {} + for key_data in jwks_data.get("keys", []): + key_kid = key_data.get("kid") + if key_kid: + jwk = JsonWebKey.import_key(key_data) + self._jwks_cache[key_kid] = jwk.get_public_key() + + self._jwks_cache_time = current_time + + if kid not in self._jwks_cache: + raise ValueError(f"Key ID '{kid}' not found in JWKS") + + return self._jwks_cache[kid] + + except Exception as e: + raise ValueError(f"Failed to fetch JWKS: {e}") + + async def load_access_token(self, token: str) -> AccessToken | None: + """ + Validates the provided JWT bearer token. + + Args: + token: The JWT token string to validate + + Returns: + AccessToken object if valid, None if invalid or expired + """ + try: + # Get verification key (static or from JWKS) + verification_key = await self._get_verification_key(token) + + # Decode and verify the JWT token + claims = self.jwt.decode(token, verification_key) + + # Validate expiration + exp = claims.get("exp") + if exp and exp < time.time(): + return None + + # Validate issuer + if self.issuer: + if claims.get("iss") != self.issuer: + return None + + # Validate audience if configured + if self.audience: + aud = claims.get("aud") + if isinstance(aud, list): + if self.audience not in aud: + return None + elif aud != self.audience: + return None + + # Extract claims - prefer client_id over sub for OAuth application identification + client_id = claims.get("client_id") or claims.get("sub") or "unknown" + scopes = self._extract_scopes(claims) + + return AccessToken( + token=token, + client_id=str(client_id), + scopes=scopes, + expires_at=int(exp) if exp else None, + ) + + except JoseError: + return None + except Exception: + return None + + def _extract_scopes(self, claims: dict[str, Any]) -> list[str]: + """Extract scopes from JWT claims.""" + scope_claim = claims.get("scope", "") + if isinstance(scope_claim, str): + return scope_claim.split() + elif isinstance(scope_claim, list): + return scope_claim + return [] + + # --- Unused OAuth server methods --- + async def get_client(self, client_id: str) -> OAuthClientInformationFull | None: + raise NotImplementedError("Client management not supported") + + async def register_client(self, client_info: OAuthClientInformationFull) -> None: + raise NotImplementedError("Client registration not supported") + + async def authorize( + self, client: OAuthClientInformationFull, params: AuthorizationParams + ) -> str: + raise NotImplementedError("Authorization flow not supported") + + async def load_authorization_code( + self, client: OAuthClientInformationFull, authorization_code: str + ) -> AuthorizationCode | None: + raise NotImplementedError("Authorization code flow not supported") + + async def exchange_authorization_code( + self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode + ) -> OAuthToken: + raise NotImplementedError("Authorization code exchange not supported") + + async def load_refresh_token( + self, client: OAuthClientInformationFull, refresh_token: str + ) -> RefreshToken | None: + raise NotImplementedError("Refresh token flow not supported") + + async def exchange_refresh_token( + self, + client: OAuthClientInformationFull, + refresh_token: RefreshToken, + scopes: list[str], + ) -> OAuthToken: + raise NotImplementedError("Refresh token exchange not supported") + + async def revoke_token( + self, + token: AccessToken | RefreshToken, + ) -> None: + raise NotImplementedError("Token revocation not supported") diff --git a/src/fastmcp/server/auth/in_memory_provider.py b/src/fastmcp/server/auth/providers/in_memory.py similarity index 98% rename from src/fastmcp/server/auth/in_memory_provider.py rename to src/fastmcp/server/auth/providers/in_memory.py index 59ac0d2ad..6494ef18b 100644 --- a/src/fastmcp/server/auth/in_memory_provider.py +++ b/src/fastmcp/server/auth/providers/in_memory.py @@ -1,3 +1,8 @@ +""" +This is a simple in-memory OAuth provider for testing purposes. +It simulates the OAuth 2.0 flow locally without external calls. +""" + import secrets import time @@ -43,7 +48,7 @@ class InMemoryOAuthProvider(OAuthProvider): required_scopes: list[str] | None = None, ): super().__init__( - issuer_url or "https://example.com", + issuer_url=issuer_url or "http://fastmcp.example.com", service_documentation_url=service_documentation_url, client_registration_options=client_registration_options, revocation_options=revocation_options, diff --git a/src/fastmcp/server/dependencies.py b/src/fastmcp/server/dependencies.py index e2d279dc5..572af5282 100644 --- a/src/fastmcp/server/dependencies.py +++ b/src/fastmcp/server/dependencies.py @@ -2,6 +2,8 @@ from __future__ import annotations from typing import TYPE_CHECKING, ParamSpec, TypeVar +from mcp.server.auth.middleware.auth_context import get_access_token +from mcp.server.auth.provider import AccessToken from starlette.requests import Request if TYPE_CHECKING: @@ -10,6 +12,14 @@ if TYPE_CHECKING: P = ParamSpec("P") R = TypeVar("R") +__all__ = [ + "get_context", + "get_http_request", + "get_http_headers", + "get_access_token", + "AccessToken", +] + # --- Context --- diff --git a/tests/auth/providers/test_bearer.py b/tests/auth/providers/test_bearer.py new file mode 100644 index 000000000..8f81d4339 --- /dev/null +++ b/tests/auth/providers/test_bearer.py @@ -0,0 +1,366 @@ +from collections.abc import Generator + +import httpx +import pytest + +from fastmcp import Client, FastMCP +from fastmcp.client.auth import BearerAuth +from fastmcp.server.auth.providers.bearer import BearerAuthProvider, RSAKeyPair +from fastmcp.utilities.tests import run_server_in_process + + +@pytest.fixture(scope="module") +def rsa_key_pair() -> RSAKeyPair: + return RSAKeyPair.generate() + + +@pytest.fixture(scope="module") +def bearer_token(rsa_key_pair: RSAKeyPair) -> str: + return rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + ) + + +@pytest.fixture +def bearer_provider(rsa_key_pair: RSAKeyPair) -> BearerAuthProvider: + return BearerAuthProvider( + public_key=rsa_key_pair.public_key, + issuer="https://test.example.com", + audience="https://api.example.com", + ) + + +def run_mcp_server(public_key: str, host: str, port: int, **kwargs) -> str: + mcp = FastMCP( + auth=BearerAuthProvider( + issuer="https://test.example.com", + public_key=public_key, + ) + ) + + @mcp.tool() + def add(a: int, b: int) -> int: + return a + b + + mcp.run(host=host, port=port, **kwargs) + + +@pytest.fixture(scope="module") +def mcp_server_url(rsa_key_pair: RSAKeyPair) -> Generator[str]: + with run_server_in_process( + run_mcp_server, public_key=rsa_key_pair.public_key, transport="streamable-http" + ) as url: + yield f"{url}/mcp" + + +class TestRSAKeyPair: + def test_generate_key_pair(self): + """Test RSA key pair generation.""" + key_pair = RSAKeyPair.generate() + + assert key_pair.private_key is not None + assert key_pair.public_key is not None + + # Check that keys are in PEM format + private_pem = key_pair.private_key.get_secret_value() + public_pem = key_pair.public_key.get_secret_value() + + assert "-----BEGIN PRIVATE KEY-----" in private_pem + assert "-----END PRIVATE KEY-----" in private_pem + assert "-----BEGIN PUBLIC KEY-----" in public_pem + assert "-----END PUBLIC KEY-----" in public_pem + + def test_create_basic_token(self, rsa_key_pair: RSAKeyPair): + """Test basic token creation.""" + token = rsa_key_pair.create_token( + subject="test-user", issuer="https://test.example.com" + ) + + assert isinstance(token, str) + assert len(token.split(".")) == 3 # JWT has 3 parts + + def test_create_token_with_scopes(self, rsa_key_pair: RSAKeyPair): + """Test token creation with scopes.""" + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + scopes=["read", "write"], + ) + + assert isinstance(token, str) + # We'll validate the scopes in the BearerToken tests + + +class TestBearerToken: + def test_initialization_with_public_key(self, rsa_key_pair: RSAKeyPair): + """Test provider initialization with public key.""" + provider = BearerAuthProvider( + public_key=rsa_key_pair.public_key, issuer="https://test.example.com" + ) + + assert provider.issuer == "https://test.example.com" + assert provider.public_key is not None + assert provider.jwks_uri is None + + def test_initialization_with_jwks_uri(self): + """Test provider initialization with JWKS URI.""" + provider = BearerAuthProvider( + jwks_uri="https://test.example.com/.well-known/jwks.json", + issuer="https://test.example.com", + ) + + assert provider.issuer == "https://test.example.com" + assert provider.jwks_uri == "https://test.example.com/.well-known/jwks.json" + assert provider.public_key is None + + def test_initialization_requires_key_or_uri(self): + """Test that either public_key or jwks_uri is required.""" + with pytest.raises( + ValueError, match="Either public_key or jwks_uri must be provided" + ): + BearerAuthProvider(issuer="https://test.example.com") + + def test_initialization_rejects_both_key_and_uri(self, rsa_key_pair: RSAKeyPair): + """Test that both public_key and jwks_uri cannot be provided.""" + with pytest.raises( + ValueError, match="Provide either public_key or jwks_uri, not both" + ): + BearerAuthProvider( + public_key=rsa_key_pair.public_key, + jwks_uri="https://test.example.com/.well-known/jwks.json", + issuer="https://test.example.com", + ) + + @pytest.mark.asyncio + async def test_valid_token_validation( + self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider + ): + """Test validation of a valid token.""" + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + scopes=["read", "write"], + ) + + access_token = await bearer_provider.load_access_token(token) + + assert access_token is not None + assert access_token.client_id == "test-user" + assert "read" in access_token.scopes + assert "write" in access_token.scopes + assert access_token.expires_at is not None + + @pytest.mark.asyncio + async def test_expired_token_rejection( + self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider + ): + """Test rejection of expired tokens.""" + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + expires_in_seconds=-3600, # Expired 1 hour ago + ) + + access_token = await bearer_provider.load_access_token(token) + assert access_token is None + + @pytest.mark.asyncio + async def test_invalid_issuer_rejection( + self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider + ): + """Test rejection of tokens with invalid issuer.""" + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://evil.example.com", # Wrong issuer + audience="https://api.example.com", + ) + + access_token = await bearer_provider.load_access_token(token) + assert access_token is None + + @pytest.mark.asyncio + async def test_invalid_audience_rejection( + self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider + ): + """Test rejection of tokens with invalid audience.""" + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://wrong-api.example.com", # Wrong audience + ) + + access_token = await bearer_provider.load_access_token(token) + assert access_token is None + + @pytest.mark.asyncio + async def test_no_issuer_validation_when_none(self, rsa_key_pair: RSAKeyPair): + """Test that issuer validation is skipped when provider has no issuer configured.""" + provider = BearerAuthProvider( + public_key=rsa_key_pair.public_key, + issuer=None, # No issuer validation + ) + + token = rsa_key_pair.create_token( + subject="test-user", issuer="https://any.example.com" + ) + + access_token = await provider.load_access_token(token) + assert access_token is not None + + @pytest.mark.asyncio + async def test_no_audience_validation_when_none(self, rsa_key_pair: RSAKeyPair): + """Test that audience validation is skipped when provider has no audience configured.""" + provider = BearerAuthProvider( + public_key=rsa_key_pair.public_key, + issuer="https://test.example.com", + audience=None, # No audience validation + ) + + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://any-api.example.com", + ) + + access_token = await provider.load_access_token(token) + assert access_token is not None + + @pytest.mark.asyncio + async def test_multiple_audiences_validation(self, rsa_key_pair: RSAKeyPair): + """Test validation with multiple audiences in token.""" + provider = BearerAuthProvider( + public_key=rsa_key_pair.public_key, + issuer="https://test.example.com", + audience="https://api.example.com", + ) + + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + additional_claims={ + "aud": ["https://api.example.com", "https://other-api.example.com"] + }, + ) + + access_token = await provider.load_access_token(token) + assert access_token is not None + + @pytest.mark.asyncio + async def test_scope_extraction_string( + self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider + ): + """Test scope extraction from space-separated string.""" + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + scopes=["read", "write", "admin"], + ) + + access_token = await bearer_provider.load_access_token(token) + + assert access_token is not None + assert set(access_token.scopes) == {"read", "write", "admin"} + + @pytest.mark.asyncio + async def test_scope_extraction_list( + self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider + ): + """Test scope extraction from list format.""" + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + additional_claims={"scope": ["read", "write"]}, # List format + ) + + access_token = await bearer_provider.load_access_token(token) + + assert access_token is not None + assert set(access_token.scopes) == {"read", "write"} + + @pytest.mark.asyncio + async def test_no_scopes( + self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider + ): + """Test token with no scopes.""" + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + # No scopes + ) + + access_token = await bearer_provider.load_access_token(token) + + assert access_token is not None + assert access_token.scopes == [] + + @pytest.mark.asyncio + async def test_malformed_token_rejection(self, bearer_provider: BearerAuthProvider): + """Test rejection of malformed tokens.""" + malformed_tokens = [ + "not.a.jwt", + "too.many.parts.here.invalid", + "invalid-token", + "", + "header.body", # Missing signature + ] + + for token in malformed_tokens: + access_token = await bearer_provider.load_access_token(token) + assert access_token is None + + @pytest.mark.asyncio + async def test_invalid_signature_rejection( + self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider + ): + """Test rejection of tokens with invalid signatures.""" + # Create a token with a different key pair + other_key_pair = RSAKeyPair.generate() + token = other_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + ) + + access_token = await bearer_provider.load_access_token(token) + assert access_token is None + + @pytest.mark.asyncio + async def test_client_id_fallback( + self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider + ): + """Test client_id extraction with fallback logic.""" + # Test with explicit client_id claim + token = rsa_key_pair.create_token( + subject="user123", + issuer="https://test.example.com", + audience="https://api.example.com", + additional_claims={"client_id": "app456"}, + ) + + access_token = await bearer_provider.load_access_token(token) + assert access_token is not None + assert access_token.client_id == "app456" # Should prefer client_id over sub + + +class TestFastMCPBearerAuth: + def test_bearer_auth(self): + mcp = FastMCP( + auth=BearerAuthProvider(issuer="https://test.example.com", public_key="abc") + ) + assert isinstance(mcp.auth, BearerAuthProvider) + + async def test_unauthorized_access(self, mcp_server_url: str): + with pytest.raises(httpx.HTTPStatusError, match="401"): + async with Client(mcp_server_url) as client: + await client.ping() + + async def test_authorized_access(self, mcp_server_url: str, bearer_token): + async with Client(mcp_server_url, auth=BearerAuth(bearer_token)) as client: + await client.ping() diff --git a/tests/auth/test_oauth_client.py b/tests/auth/test_oauth_client.py index 5d668c6dc..71db7fe47 100644 --- a/tests/auth/test_oauth_client.py +++ b/tests/auth/test_oauth_client.py @@ -9,7 +9,7 @@ import fastmcp.client.auth # Import module, not the function directly from fastmcp.client import Client from fastmcp.client.transports import StreamableHttpTransport from fastmcp.server.auth.auth import ClientRegistrationOptions -from fastmcp.server.auth.in_memory_provider import InMemoryOAuthProvider as InMemory +from fastmcp.server.auth.providers.in_memory import InMemoryOAuthProvider from fastmcp.server.server import FastMCP from fastmcp.utilities.tests import run_server_in_process @@ -18,7 +18,7 @@ def fastmcp_server(issuer_url: str): """Create a FastMCP server with OAuth authentication.""" server = FastMCP( "TestServer", - auth=InMemory( + auth=InMemoryOAuthProvider( issuer_url=issuer_url, client_registration_options=ClientRegistrationOptions(enabled=True), ), diff --git a/uv.lock b/uv.lock index ae92d334e..e1cd22726 100644 --- a/uv.lock +++ b/uv.lock @@ -449,6 +449,7 @@ dev = [ { name = "ipython", version = "9.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pdbpp" }, { name = "pre-commit" }, + { name = "pyinstrument" }, { name = "pyright" }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -482,6 +483,7 @@ dev = [ { name = "ipython", specifier = ">=8.12.3" }, { name = "pdbpp", specifier = ">=0.10.3" }, { name = "pre-commit" }, + { name = "pyinstrument", specifier = ">=5.0.2" }, { name = "pyright", specifier = ">=1.1.389" }, { name = "pytest", specifier = ">=8.3.3" }, { name = "pytest-asyncio", specifier = ">=0.23.5" }, @@ -998,6 +1000,62 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8a/0b/9fcc47d19c48b59121088dd6da2488a49d5f72dacf8262e2790a1d2c7d15/pygments-2.19.1-py3-none-any.whl", hash = "sha256:9ea1544ad55cecf4b8242fab6dd35a93bbce657034b0611ee383099054ab6d8c", size = 1225293, upload-time = "2025-01-06T17:26:25.553Z" }, ] +[[package]] +name = "pyinstrument" +version = "5.0.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/d0/665828770e8fcd5c50880dc83f03811f814d6260bc6a8068dca0a520e68a/pyinstrument-5.0.2.tar.gz", hash = "sha256:e466033ead16a48ffa8bedbd633b90d416fa772b3b22f61226882ace0371f5f3", size = 263930, upload-time = "2025-05-24T15:47:13.358Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/25/f64d0be5f574d2df9ddac3e7a381863f92d8ad30170b1a9de0cf805f4318/pyinstrument-5.0.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:1aeaf6b39ad40b3f03bea5fa3a9bd453a92aeb721dde29c1597f842ed9c8566a", size = 129638, upload-time = "2025-05-24T15:45:20.113Z" }, + { url = "https://files.pythonhosted.org/packages/6e/b8/bc6657f91a8d2f7cf58b0993aa4e6cf20e027b53aca65c2464a50738d711/pyinstrument-5.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:d734bd236d00e0e7f950019c689eaba1c9dd15e355867d8926c8b18b6077b221", size = 122220, upload-time = "2025-05-24T15:45:22.4Z" }, + { url = "https://files.pythonhosted.org/packages/63/5f/9a7edf13333015a9ccfd3fcf5c75ea793fbb30b153aebf6c6ace40a607b2/pyinstrument-5.0.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:520208a9b6c3985473aa9c3f30875ae5e78e77a81081df1d8aeb4fd8b4caf197", size = 146928, upload-time = "2025-05-24T15:45:23.802Z" }, + { url = "https://files.pythonhosted.org/packages/f2/f9/f7d7b28c9038f1a570e96c8eea2a9ffeeb3ee9e75cfc74a370554776f1a6/pyinstrument-5.0.2-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:75e115b759288b8d65a0bf31a34a542ae102c58ef407e0614a43e0c39d261875", size = 157136, upload-time = "2025-05-24T15:45:25.629Z" }, + { url = "https://files.pythonhosted.org/packages/db/ee/aa99f275b3c5f0f32ccd37f77cb64e57597a1f26280aec03a50d2158eab7/pyinstrument-5.0.2-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:091f93e6787c485a7ddf670608c00448e858a056677fc25ce349f8e44d6a9e54", size = 144680, upload-time = "2025-05-24T15:45:27.06Z" }, + { url = "https://files.pythonhosted.org/packages/ae/77/cd7300a5e099c4ad971a647ea8fb9bd081482a9e5751479034e206cd1f69/pyinstrument-5.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:28b07971afa2652cb4f2bdcffaef11aefa32b5384c0cfb32acf9955e96dd8df8", size = 145624, upload-time = "2025-05-24T15:45:28.517Z" }, + { url = "https://files.pythonhosted.org/packages/30/59/1957e2ca2277ecc69e247383527df331002e23940d5b0a79fc5f3b870d60/pyinstrument-5.0.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:e1bcb28a21b80eea5986eb5cb3180689b1d489b7c6fddf34e1f4df1f95d467ad", size = 145901, upload-time = "2025-05-24T15:45:30.365Z" }, + { url = "https://files.pythonhosted.org/packages/7e/ae/396ebdf387cde376ac4b70d52f3df07374f2501ac4c09992dadf641cd71f/pyinstrument-5.0.2-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:80d28162070ff40c6d2ac7dc15b933ba20ef49e891a2e650cd2b91d30cd262b2", size = 145355, upload-time = "2025-05-24T15:45:31.859Z" }, + { url = "https://files.pythonhosted.org/packages/48/5b/fec77476a9b4a316861b29f14cd0962871ad5c54c21e41c540f7c18c950c/pyinstrument-5.0.2-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:c75e52a9bf76f084ba074323835cba4927ab3e572adfc96439698b097e523780", size = 145008, upload-time = "2025-05-24T15:45:33.417Z" }, + { url = "https://files.pythonhosted.org/packages/fd/75/dcd391ca2790b32e41bbd49ad33626e85eb1ce00116b273d5e1d99b3e829/pyinstrument-5.0.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:ccefdd7dd938548ada43c95b24c42ec57e258ac7994a5ec7e4cc934fa4f1743b", size = 145396, upload-time = "2025-05-24T15:45:34.915Z" }, + { url = "https://files.pythonhosted.org/packages/ca/5c/64e026ccf2c7908d10882955993e73ac35a1a77426bde2617973deeda07c/pyinstrument-5.0.2-cp310-cp310-win32.whl", hash = "sha256:6b617fb024c244738aa2f6b8c2a25853eac765360ac91062578bbbcc8e22ebfe", size = 123419, upload-time = "2025-05-24T15:45:36.276Z" }, + { url = "https://files.pythonhosted.org/packages/d5/7c/7d221db96d461c7d28897499bdad55a8ae5ded983f60743bdfbf17438c20/pyinstrument-5.0.2-cp310-cp310-win_amd64.whl", hash = "sha256:6788c8f93c1a6e0ad8d0ccde1631d17eca3839945d0fa4d506cf5d4bd7a26b77", size = 124299, upload-time = "2025-05-24T15:45:37.642Z" }, + { url = "https://files.pythonhosted.org/packages/fc/f2/b3f2416740be762fdfb052b63e1d85591682fa1d2ea6ee1b10db774f6350/pyinstrument-5.0.2-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:0eec7a263cc1ccfb101594e13256115366338fee2a156be4172fe5315f71ec45", size = 129386, upload-time = "2025-05-24T15:45:39.429Z" }, + { url = "https://files.pythonhosted.org/packages/6b/fa/a55b0bf911041b51d2a7a0e8a3feef5ed5ddb48ff0943fc667079955c14c/pyinstrument-5.0.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:ddd5effefb470d7f1886dc16467501b866e3b5883cf74773f13179e718b28393", size = 122100, upload-time = "2025-05-24T15:45:41.253Z" }, + { url = "https://files.pythonhosted.org/packages/a5/e1/c42b94c795bc89d5a486ad7ef349fe3b7a8c3a4e730c09b5fa54af616a6b/pyinstrument-5.0.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6e7458a6aa4048c1703354fc8a4a3c8b59d27b1409aafb707cf339d3c0bc794c", size = 145385, upload-time = "2025-05-24T15:45:43.024Z" }, + { url = "https://files.pythonhosted.org/packages/ff/41/b511141cc336ffeac284cce7d121f05802ffea4ab2c19df8869adda49743/pyinstrument-5.0.2-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2373dd699711463011ec14e4918427a777f7ab73b31ae374d960725dbd5d5a28", size = 156093, upload-time = "2025-05-24T15:45:44.755Z" }, + { url = "https://files.pythonhosted.org/packages/b2/7e/4a7bc4f1c60d4886efb7397fd5bdcc7e537d01ec7372824cd834fff967a1/pyinstrument-5.0.2-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:38ef498fbe71c2bbd11247b71e722290da93a367d88a5a8e0f66f6cc764c2b60", size = 143136, upload-time = "2025-05-24T15:45:46.469Z" }, + { url = "https://files.pythonhosted.org/packages/d8/69/0ac06cf609153fc5eb30ccc0071ce300a181f422836ca7ce8cd431ac3ab4/pyinstrument-5.0.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0a58a8a50f0cb3ee1c2e43ffec51bf48f48945e141feed7ccd9194917b97fe5b", size = 144077, upload-time = "2025-05-24T15:45:48.333Z" }, + { url = "https://files.pythonhosted.org/packages/e3/24/12bd82822393f708e5da8f6c0b82def3f0cbe1f4fbd72a082688c583d7fa/pyinstrument-5.0.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:ad2a97c79ecf0e610df292abb5c46d01a4f99778598881d6e918650fa39801b6", size = 144545, upload-time = "2025-05-24T15:45:50.137Z" }, + { url = "https://files.pythonhosted.org/packages/c9/62/40e7511fa46247ca56734d34e2d2eb6b14390c72b155255ecd1b2288d02d/pyinstrument-5.0.2-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:57ec0277042ee198eb749b76a975fe60f006cd51ea0c7ce3054c937577d19315", size = 144010, upload-time = "2025-05-24T15:45:52.256Z" }, + { url = "https://files.pythonhosted.org/packages/82/77/6d40880dc46a6243951ad7cd50a77f26f6ad126b80d803616934efccf539/pyinstrument-5.0.2-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:73d34047266f27acb67218e331288c0241cf0080fe4b87dfad5596236c71abd7", size = 143746, upload-time = "2025-05-24T15:45:53.702Z" }, + { url = "https://files.pythonhosted.org/packages/9b/a2/08b056d2420199dab877c665ed45bb685863dc5b83d31b2c4311430b2bbd/pyinstrument-5.0.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:cfdc23284a8e2f27637b357c226a15d52b96608d9dde187b68dfe33a947f4908", size = 143928, upload-time = "2025-05-24T15:45:55.103Z" }, + { url = "https://files.pythonhosted.org/packages/39/a1/bab336f70cd5f798d7fa21ec92784b99d3b2df0b5c1736a64fdaa4521004/pyinstrument-5.0.2-cp311-cp311-win32.whl", hash = "sha256:3e6fa135aee6af2c608e912d8d07906bbac3c5e564d94f92721831a957297c26", size = 123395, upload-time = "2025-05-24T15:45:56.469Z" }, + { url = "https://files.pythonhosted.org/packages/f2/15/8a7ac268ffe913aa64bb42ad43315dd0fc3ac493d451a50d4431ecb736c2/pyinstrument-5.0.2-cp311-cp311-win_amd64.whl", hash = "sha256:6317df42a98a8074ccd25af5482312ec59a1f27c05dab408eb3c7b2081242733", size = 124198, upload-time = "2025-05-24T15:45:57.814Z" }, + { url = "https://files.pythonhosted.org/packages/95/36/4afdffbc4fd77dd0155c8943101f175e701ba00cb374c5e84e64790a2a32/pyinstrument-5.0.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:d0b680ef269b528d8dcd8151362fba9683b0ac22ffe74cc8161c33b53c65b899", size = 129527, upload-time = "2025-05-24T15:45:59.216Z" }, + { url = "https://files.pythonhosted.org/packages/96/fe/7ea5af73d65f8f22585005f6e2ce1016fb3145a8ecc1ded51f965c2e98cc/pyinstrument-5.0.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1c70b50ec90ae793b74733a6fc992723c6ee27c0fcb7d99848239316ded61189", size = 122068, upload-time = "2025-05-24T15:46:01.05Z" }, + { url = "https://files.pythonhosted.org/packages/3f/d2/cf8f3b8fde3f3b6768f8407c681fb57e7b5a5bf5e7450a9fbec15164987b/pyinstrument-5.0.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3aae5f4f78515009f72393fdb271a15861534a586401383785f823cf8f60aa02", size = 146679, upload-time = "2025-05-24T15:46:02.841Z" }, + { url = "https://files.pythonhosted.org/packages/8d/e2/6c00273778596560c7033cfee34aab07da6009f32c5a4dbcc35b64700e73/pyinstrument-5.0.2-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3aec8bc3d1c064ff849ca3568d6b0a7cfa0162d590a9d4d250c7118d09518b22", size = 157606, upload-time = "2025-05-24T15:46:04.551Z" }, + { url = "https://files.pythonhosted.org/packages/4c/cc/ec099f566e381f8e5db21d9523dd97b3255047813da57481ab3f45436089/pyinstrument-5.0.2-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:28d87fac2bc0fed802b14a26982440f36c85dc53f303530ff7665a6e470315bb", size = 144317, upload-time = "2025-05-24T15:46:05.996Z" }, + { url = "https://files.pythonhosted.org/packages/37/a7/e2e54bf6d996b3c807534dbc4fe270f373660b89871c63965d3f895c285d/pyinstrument-5.0.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3b9caac53c7eda8187ed122d4f7fcc6e3392f04c583d6d70b373351cede2b829", size = 145622, upload-time = "2025-05-24T15:46:07.334Z" }, + { url = "https://files.pythonhosted.org/packages/ef/c6/0b084ddf8d836076e04912ea83ccae0f83bf4897d0168b0fd7684efdc2a4/pyinstrument-5.0.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8124419e8731a7bdbb9f7f885a8956806a4e9ab9dd19294f8a99e74c0bbdd327", size = 145645, upload-time = "2025-05-24T15:46:09.236Z" }, + { url = "https://files.pythonhosted.org/packages/b3/4d/3e542c5986cc30bc86c304492f4696e58dc03d1816d35c5b2cabfac1d01e/pyinstrument-5.0.2-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:9990d9bd05fbb4fa83f24f0a62989b8e0a3ac15ff0fa19b49348c8ef5f9db50a", size = 145619, upload-time = "2025-05-24T15:46:10.643Z" }, + { url = "https://files.pythonhosted.org/packages/a7/a7/1e4664bf5ada1cff56852d10954b1ff5a39dad17b9b98a2f27054a0c0d95/pyinstrument-5.0.2-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:1dc35f3d200866a43d4bc7570799a405f001591c8f19a30eb7a983a717c1e1f7", size = 145049, upload-time = "2025-05-24T15:46:12.019Z" }, + { url = "https://files.pythonhosted.org/packages/fb/59/08a5237c8d1343842ac9ed3c661dce40c450f1750128fd4789ad80539253/pyinstrument-5.0.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a335a40d0ba1fe3658ef1a5ff2fc7a6870905828014645cb19dab5c1de379447", size = 145451, upload-time = "2025-05-24T15:46:13.49Z" }, + { url = "https://files.pythonhosted.org/packages/53/d0/321b5301e36ac1577dbf73cb49769779c41ebf72ba70a3f6f62d34df902b/pyinstrument-5.0.2-cp312-cp312-win32.whl", hash = "sha256:29e565ce85e03d2541330a8174124c1ecdb073d945962a8eb738d3b1c806ac83", size = 123491, upload-time = "2025-05-24T15:46:15.319Z" }, + { url = "https://files.pythonhosted.org/packages/ae/a6/40f05febe6ab0856b4bfa119113d550d868d94a36b501e6b9fd64379b4ba/pyinstrument-5.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:300b0cc453ffe7661d5f3ceb94cdd98996fd9118f5ff1182b5336489c7d4e45c", size = 124277, upload-time = "2025-05-24T15:46:16.693Z" }, + { url = "https://files.pythonhosted.org/packages/03/88/48654e4b8c6853f218e0506e0609060a54559500b3af5ed6ac752ac4d64f/pyinstrument-5.0.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8141a5f78b927a88de46fb2bbb17e710e41d16e161fca99991635ff7196dbd5d", size = 129528, upload-time = "2025-05-24T15:46:18.108Z" }, + { url = "https://files.pythonhosted.org/packages/92/a7/885418b733350f6c2b1d8fcca322a1eee87216a266ac516d7aefd6757ec8/pyinstrument-5.0.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:12a0095ae408dbbdd429501fd4c6a3ab51d1aeff5f31be36cc3eedc8c4870ede", size = 122072, upload-time = "2025-05-24T15:46:19.513Z" }, + { url = "https://files.pythonhosted.org/packages/a4/d5/dd0b323d2949d1a3ee0531ec6cdd66c3c69c13b9a8739aeec929a0b55fd2/pyinstrument-5.0.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eca651d840e8e75ae5330abfc5c90f6ea4af3f78f9f0269231328305a5f9c667", size = 146874, upload-time = "2025-05-24T15:46:21.38Z" }, + { url = "https://files.pythonhosted.org/packages/aa/3b/429572b57c9ae2874e86c48db91ddcd5d619bd798f73d7d2e51b28abb08d/pyinstrument-5.0.2-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:89d6ffc5459b19f1c85d4433bb9bbc8925ec04a8d7caf2694218b1f557555f23", size = 155257, upload-time = "2025-05-24T15:46:22.791Z" }, + { url = "https://files.pythonhosted.org/packages/7a/98/03cd22f68607362fd8d1ba72e6367104a9dc32bd4a0dbafc823c4e366f35/pyinstrument-5.0.2-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4c84845ccc5318072708dc5535b6bedd54494e92a68e282e6b97b53c1db65331", size = 144380, upload-time = "2025-05-24T15:46:24.26Z" }, + { url = "https://files.pythonhosted.org/packages/eb/c4/40d7b4be6c9620c4d9bbe9788eb9bac892f386c9bd40f1937464b2b95c09/pyinstrument-5.0.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6511092384b5729bbbf4b35534120d2969c5fdfd4f39080badedd973676b8725", size = 145794, upload-time = "2025-05-24T15:46:25.751Z" }, + { url = "https://files.pythonhosted.org/packages/05/07/3b2084b78521d5bbbc328ca9527fb54fbf645a5e62f25169b49f7bbb0bc3/pyinstrument-5.0.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:73f08cff7a8d9714be15440046289ab1a70cbc429e09967a3a106ac61538773e", size = 145803, upload-time = "2025-05-24T15:46:27.277Z" }, + { url = "https://files.pythonhosted.org/packages/22/eb/e3ffcc8734e3d9f50b6bb750209c3ad0c4626dcc3754529741499d9f1d5c/pyinstrument-5.0.2-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:3905b510cdab1a8255a23fbdedcba4685245cbf814fd80f5b2005b472161d16e", size = 145763, upload-time = "2025-05-24T15:46:28.656Z" }, + { url = "https://files.pythonhosted.org/packages/c6/34/6b94945a02afced9e486e9a6b20de0edcfec543e4942dea96d745e2148ac/pyinstrument-5.0.2-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:cd693a616166679da529168037c294ff25746c7ae5e8b547811fb25bb26439f5", size = 145208, upload-time = "2025-05-24T15:46:30.125Z" }, + { url = "https://files.pythonhosted.org/packages/99/af/0339bbfe52de9a7df01e5a244a5fec4c228d23b1f422a55318fc6d0b9d91/pyinstrument-5.0.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:83a1659a3bc4123c81fcddfcc86608f37bd6a951da9692766c2251500a77ac06", size = 145591, upload-time = "2025-05-24T15:46:31.556Z" }, + { url = "https://files.pythonhosted.org/packages/c8/f4/76a2c652e203c15cbc7aa3f8341e07d1ea865764b3ed9f9a97b3c4a5eda2/pyinstrument-5.0.2-cp313-cp313-win32.whl", hash = "sha256:386d047db6c043dcc86bac592873234a89eaa258460e1ad8f47a11fcc7b024d5", size = 123490, upload-time = "2025-05-24T15:46:32.951Z" }, + { url = "https://files.pythonhosted.org/packages/e4/63/14f5c6253e8c85c758485c7717f542346a0d4487818afc28721912a1574b/pyinstrument-5.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:971c974c061019fa6177a021882255e639399bc15bf71b0a17979830702ad8d3", size = 124287, upload-time = "2025-05-24T15:46:34.333Z" }, +] + [[package]] name = "pyperclip" version = "1.9.0" From ed94a2ae8977f29cef9193dc7fb59c6c17a2d1c0 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 1 Jun 2025 11:57:44 -0400 Subject: [PATCH 2/5] Update test_bearer.py --- tests/auth/providers/test_bearer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/auth/providers/test_bearer.py b/tests/auth/providers/test_bearer.py index 8f81d4339..480c9a1b9 100644 --- a/tests/auth/providers/test_bearer.py +++ b/tests/auth/providers/test_bearer.py @@ -32,7 +32,7 @@ def bearer_provider(rsa_key_pair: RSAKeyPair) -> BearerAuthProvider: ) -def run_mcp_server(public_key: str, host: str, port: int, **kwargs) -> str: +def run_mcp_server(public_key: str, host: str, port: int, **kwargs) -> None: mcp = FastMCP( auth=BearerAuthProvider( issuer="https://test.example.com", @@ -65,7 +65,7 @@ class TestRSAKeyPair: # Check that keys are in PEM format private_pem = key_pair.private_key.get_secret_value() - public_pem = key_pair.public_key.get_secret_value() + public_pem = key_pair.public_key assert "-----BEGIN PRIVATE KEY-----" in private_pem assert "-----END PRIVATE KEY-----" in private_pem From da9c51e13289f406fb2220caa3d7d746b60d355a Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 1 Jun 2025 12:09:30 -0400 Subject: [PATCH 3/5] Add tests; update default issuer --- src/fastmcp/server/auth/providers/bearer.py | 2 +- tests/auth/providers/test_bearer.py | 54 ++++++++++++++++++--- 2 files changed, 48 insertions(+), 8 deletions(-) diff --git a/src/fastmcp/server/auth/providers/bearer.py b/src/fastmcp/server/auth/providers/bearer.py index 3cedfee7b..ccf25f0bd 100644 --- a/src/fastmcp/server/auth/providers/bearer.py +++ b/src/fastmcp/server/auth/providers/bearer.py @@ -177,7 +177,7 @@ class BearerAuthProvider(OAuthProvider): raise ValueError("Provide either public_key or jwks_uri, not both") super().__init__( - issuer_url=issuer or "http://fastmcp.example.com", + issuer_url=issuer or "https://fastmcp.example.com", client_registration_options=ClientRegistrationOptions(enabled=False), revocation_options=RevocationOptions(enabled=False), required_scopes=required_scopes, diff --git a/tests/auth/providers/test_bearer.py b/tests/auth/providers/test_bearer.py index 480c9a1b9..ccc84fe73 100644 --- a/tests/auth/providers/test_bearer.py +++ b/tests/auth/providers/test_bearer.py @@ -1,4 +1,5 @@ from collections.abc import Generator +from typing import Any import httpx import pytest @@ -32,11 +33,17 @@ def bearer_provider(rsa_key_pair: RSAKeyPair) -> BearerAuthProvider: ) -def run_mcp_server(public_key: str, host: str, port: int, **kwargs) -> None: +def run_mcp_server( + public_key: str, + host: str, + port: int, + auth_kwargs: dict[str, Any] | None = None, + run_kwargs: dict[str, Any] | None = None, +) -> None: mcp = FastMCP( auth=BearerAuthProvider( - issuer="https://test.example.com", public_key=public_key, + **auth_kwargs or {}, ) ) @@ -44,13 +51,15 @@ def run_mcp_server(public_key: str, host: str, port: int, **kwargs) -> None: def add(a: int, b: int) -> int: return a + b - mcp.run(host=host, port=port, **kwargs) + mcp.run(host=host, port=port, **run_kwargs or {}) @pytest.fixture(scope="module") def mcp_server_url(rsa_key_pair: RSAKeyPair) -> Generator[str]: with run_server_in_process( - run_mcp_server, public_key=rsa_key_pair.public_key, transport="streamable-http" + run_mcp_server, + public_key=rsa_key_pair.public_key, + run_kwargs=dict(transport="streamable-http"), ) as url: yield f"{url}/mcp" @@ -75,7 +84,8 @@ class TestRSAKeyPair: def test_create_basic_token(self, rsa_key_pair: RSAKeyPair): """Test basic token creation.""" token = rsa_key_pair.create_token( - subject="test-user", issuer="https://test.example.com" + subject="test-user", + issuer="https://test.example.com", ) assert isinstance(token, str) @@ -359,8 +369,38 @@ class TestFastMCPBearerAuth: async def test_unauthorized_access(self, mcp_server_url: str): with pytest.raises(httpx.HTTPStatusError, match="401"): async with Client(mcp_server_url) as client: - await client.ping() + tools = await client.list_tools() # noqa: F841 + assert "tools" not in locals() async def test_authorized_access(self, mcp_server_url: str, bearer_token): async with Client(mcp_server_url, auth=BearerAuth(bearer_token)) as client: - await client.ping() + tools = await client.list_tools() # noqa: F841 + assert tools + + async def test_invalid_token_raises_401(self, mcp_server_url: str): + with pytest.raises(httpx.HTTPStatusError, match="401"): + async with Client(mcp_server_url, auth=BearerAuth("invalid")) as client: + tools = await client.list_tools() # noqa: F841 + assert "tools" not in locals() + + async def test_expired_token(self, mcp_server_url: str, rsa_key_pair: RSAKeyPair): + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + expires_in_seconds=-3600, + ) + + with pytest.raises(httpx.HTTPStatusError, match="401"): + async with Client(mcp_server_url, auth=BearerAuth(token)) as client: + tools = await client.list_tools() # noqa: F841 + assert "tools" not in locals() + + async def test_token_with_bad_signature(self, mcp_server_url: str): + rsa_key_pair = RSAKeyPair.generate() + token = rsa_key_pair.create_token() + + with pytest.raises(httpx.HTTPStatusError, match="401"): + async with Client(mcp_server_url, auth=BearerAuth(token)) as client: + tools = await client.list_tools() # noqa: F841 + assert "tools" not in locals() From 172f67c6176105cc895590a4c3144d5a560391bc Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 1 Jun 2025 12:44:56 -0400 Subject: [PATCH 4/5] Add jwks tests --- pyproject.toml | 1 + src/fastmcp/server/auth/bearer.py | 256 ------------------- src/fastmcp/server/auth/providers/bearer.py | 75 ++++-- tests/auth/providers/test_bearer.py | 264 ++++++++++++++++++-- uv.lock | 15 ++ 5 files changed, 318 insertions(+), 293 deletions(-) delete mode 100644 src/fastmcp/server/auth/bearer.py diff --git a/pyproject.toml b/pyproject.toml index 35b538b58..a6fce4ad6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ dev = [ "pytest-cov>=6.1.1", "pytest-env>=1.1.5", "pytest-flakefinder", + "pytest-httpx>=0.35.0", "pytest-report>=0.2.1", "pytest-timeout>=2.4.0", "pytest-xdist>=3.6.1", diff --git a/src/fastmcp/server/auth/bearer.py b/src/fastmcp/server/auth/bearer.py deleted file mode 100644 index 729fcf1dc..000000000 --- a/src/fastmcp/server/auth/bearer.py +++ /dev/null @@ -1,256 +0,0 @@ -""" -Simple JWT Bearer Token validation for hosted MCP servers. - -Uses RS256 (asymmetric) where your control plane signs with a private key -and hosted MCP servers validate with the corresponding public key. - -Example usage: -# Static public key -provider = BearerTokenValidatorProvider( - public_key='''-----BEGIN PUBLIC KEY----- - MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA... - -----END PUBLIC KEY-----''', - issuer="https://auth.yourservice.com" -) - -# Or JWKS URI (recommended for production - allows key rotation) -provider = BearerTokenValidatorProvider( - jwks_uri="https://auth.yourservice.com/.well-known/jwks.json", - issuer="https://auth.yourservice.com" -) -""" - -import time -from typing import Any - -import httpx -from authlib.jose import JsonWebKey, JsonWebToken -from authlib.jose.errors import JoseError -from mcp.server.auth.provider import ( - AccessToken, - AuthorizationCode, - AuthorizationParams, - RefreshToken, -) -from mcp.shared.auth import ( - OAuthClientInformationFull, - OAuthToken, -) - -from fastmcp.server.auth.auth import ( - ClientRegistrationOptions, - OAuthProvider, - RevocationOptions, -) - - -class BearerTokenValidatorProvider(OAuthProvider): - """ - Simple JWT Bearer Token validator for hosted MCP servers. - Uses RS256 asymmetric encryption. Supports either static public key - or JWKS URI for key rotation. - """ - - def __init__( - self, - issuer: str, - public_key: str | None = None, - jwks_uri: str | None = None, - audience: str | None = None, - required_scopes: list[str] | None = None, - ): - """ - Initialize the provider. - - Args: - issuer: Expected issuer claim (your control plane) - public_key: RSA public key in PEM format (for static key) - jwks_uri: URI to fetch keys from (for key rotation) - audience: Expected audience claim (optional) - required_scopes: List of required scopes for access - """ - if not (public_key or jwks_uri): - raise ValueError("Either public_key or jwks_uri must be provided") - if public_key and jwks_uri: - raise ValueError("Provide either public_key or jwks_uri, not both") - - super().__init__( - issuer_url=issuer, - client_registration_options=ClientRegistrationOptions(enabled=False), - revocation_options=RevocationOptions(enabled=False), - required_scopes=required_scopes, - ) - - self.issuer = issuer - self.audience = audience - self.public_key = public_key - self.jwks_uri = jwks_uri - self.jwt = JsonWebToken(["RS256"]) - - # Simple JWKS cache - self._jwks_cache: dict[str, str] = {} - self._jwks_cache_time: float = 0 - self._cache_ttl = 3600 # 1 hour - - async def _get_verification_key(self, token: str) -> str: - """Get the verification key for the token.""" - if self.public_key: - return self.public_key - - # Extract kid from token header for JWKS lookup - try: - import base64 - import json - - header_b64 = token.split(".")[0] - header_b64 += "=" * (4 - len(header_b64) % 4) # Add padding - header = json.loads(base64.urlsafe_b64decode(header_b64)) - kid = header.get("kid") - - if not kid: - raise ValueError("Token missing key ID (kid)") - - return await self._get_jwks_key(kid) - - except Exception as e: - raise ValueError(f"Failed to extract key ID from token: {e}") - - async def _get_jwks_key(self, kid: str) -> str: - """Fetch key from JWKS with simple caching.""" - if not self.jwks_uri: - raise ValueError("JWKS URI not configured") - - current_time = time.time() - - # Check cache - if ( - current_time - self._jwks_cache_time < self._cache_ttl - and kid in self._jwks_cache - ): - return self._jwks_cache[kid] - - # Fetch JWKS - try: - async with httpx.AsyncClient() as client: - response = await client.get(self.jwks_uri) - response.raise_for_status() - jwks_data = response.json() - - # Cache all keys - self._jwks_cache = {} - for key_data in jwks_data.get("keys", []): - key_kid = key_data.get("kid") - if key_kid: - jwk = JsonWebKey.import_key(key_data) - self._jwks_cache[key_kid] = jwk.get_public_key() - - self._jwks_cache_time = current_time - - if kid not in self._jwks_cache: - raise ValueError(f"Key ID '{kid}' not found in JWKS") - - return self._jwks_cache[kid] - - except Exception as e: - raise ValueError(f"Failed to fetch JWKS: {e}") - - async def load_access_token(self, token: str) -> AccessToken | None: - """ - Validates the provided JWT bearer token. - - Args: - token: The JWT token string to validate - - Returns: - AccessToken object if valid, None if invalid or expired - """ - try: - # Get verification key (static or from JWKS) - verification_key = await self._get_verification_key(token) - - # Decode and verify the JWT token - claims = self.jwt.decode(token, verification_key) - - # Validate expiration - exp = claims.get("exp") - if exp and exp < time.time(): - return None - - # Validate issuer - if claims.get("iss") != self.issuer: - return None - - # Validate audience if configured - if self.audience: - aud = claims.get("aud") - if isinstance(aud, list): - if self.audience not in aud: - return None - elif aud != self.audience: - return None - - # Extract claims - client_id = claims.get("sub") or claims.get("client_id") or "unknown" - scopes = self._extract_scopes(claims) - - return AccessToken( - token=token, - client_id=str(client_id), - scopes=scopes, - expires_at=int(exp) if exp else None, - ) - - except JoseError: - return None - except Exception: - return None - - def _extract_scopes(self, claims: dict[str, Any]) -> list[str]: - """Extract scopes from JWT claims.""" - scope_claim = claims.get("scope", "") - if isinstance(scope_claim, str): - return scope_claim.split() - elif isinstance(scope_claim, list): - return scope_claim - return [] - - # --- Unused OAuth server methods --- - async def get_client(self, client_id: str) -> OAuthClientInformationFull | None: - raise NotImplementedError("Client management not supported") - - async def register_client(self, client_info: OAuthClientInformationFull) -> None: - raise NotImplementedError("Client registration not supported") - - async def authorize( - self, client: OAuthClientInformationFull, params: AuthorizationParams - ) -> str: - raise NotImplementedError("Authorization flow not supported") - - async def load_authorization_code( - self, client: OAuthClientInformationFull, authorization_code: str - ) -> AuthorizationCode | None: - raise NotImplementedError("Authorization code flow not supported") - - async def exchange_authorization_code( - self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode - ) -> OAuthToken: - raise NotImplementedError("Authorization code exchange not supported") - - async def load_refresh_token( - self, client: OAuthClientInformationFull, refresh_token: str - ) -> RefreshToken | None: - raise NotImplementedError("Refresh token flow not supported") - - async def exchange_refresh_token( - self, - client: OAuthClientInformationFull, - refresh_token: RefreshToken, - scopes: list[str], - ) -> OAuthToken: - raise NotImplementedError("Refresh token exchange not supported") - - async def revoke_token( - self, - token: AccessToken | RefreshToken, - ) -> None: - raise NotImplementedError("Token revocation not supported") diff --git a/src/fastmcp/server/auth/providers/bearer.py b/src/fastmcp/server/auth/providers/bearer.py index ccf25f0bd..419ed255c 100644 --- a/src/fastmcp/server/auth/providers/bearer.py +++ b/src/fastmcp/server/auth/providers/bearer.py @@ -6,7 +6,7 @@ and hosted MCP servers validate with the corresponding public key. Example usage: # Static public key -provider = BearerTokenValidatorProvider( +provider = BearerAuthProvider( public_key='''-----BEGIN PUBLIC KEY----- MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA... -----END PUBLIC KEY-----''', @@ -14,7 +14,7 @@ provider = BearerTokenValidatorProvider( ) # Or JWKS URI (recommended for production - allows key rotation) -provider = BearerTokenValidatorProvider( +provider = Bear( jwks_uri="https://auth.yourservice.com/.well-known/jwks.json", issuer="https://auth.yourservice.com" ) @@ -22,7 +22,7 @@ provider = BearerTokenValidatorProvider( import time from dataclasses import dataclass -from typing import Any +from typing import Any, TypedDict import httpx from authlib.jose import JsonWebKey, JsonWebToken @@ -48,6 +48,25 @@ from fastmcp.server.auth.auth import ( ) +class JWKData(TypedDict, total=False): + """JSON Web Key data structure.""" + + kty: str # Key type (e.g., "RSA") - required + kid: str # Key ID (optional but recommended) + use: str # Usage (e.g., "sig") + alg: str # Algorithm (e.g., "RS256") + n: str # Modulus (for RSA keys) + e: str # Exponent (for RSA keys) + x5c: list[str] # X.509 certificate chain (for JWKs) + x5t: str # X.509 certificate thumbprint (for JWKs) + + +class JWKSData(TypedDict): + """JSON Web Key Set data structure.""" + + keys: list[JWKData] + + @dataclass(frozen=True, kw_only=True, repr=False) class RSAKeyPair: private_key: SecretStr @@ -96,6 +115,7 @@ class RSAKeyPair: scopes: list[str] | None = None, expires_in_seconds: int = 3600, additional_claims: dict[str, Any] | None = None, + kid: str | None = None, ) -> str: """ Generate a test JWT token for testing purposes. @@ -108,6 +128,7 @@ class RSAKeyPair: scopes: List of scopes to include expires_in_seconds: Token expiration time in seconds additional_claims: Any additional claims to include + kid: Key ID for JWKS lookup (optional) Returns: Signed JWT token string @@ -135,6 +156,8 @@ class RSAKeyPair: # Create header header = {"alg": "RS256"} + if kid: + header["kid"] = kid # Sign and return token token_bytes = jwt.encode( @@ -209,27 +232,25 @@ class BearerAuthProvider(OAuthProvider): header = json.loads(base64.urlsafe_b64decode(header_b64)) kid = header.get("kid") - if not kid: - raise ValueError("Token missing key ID (kid)") - return await self._get_jwks_key(kid) except Exception as e: raise ValueError(f"Failed to extract key ID from token: {e}") - async def _get_jwks_key(self, kid: str) -> str: + async def _get_jwks_key(self, kid: str | None) -> str: """Fetch key from JWKS with simple caching.""" if not self.jwks_uri: raise ValueError("JWKS URI not configured") current_time = time.time() - # Check cache - if ( - current_time - self._jwks_cache_time < self._cache_ttl - and kid in self._jwks_cache - ): - return self._jwks_cache[kid] + # Check cache first + if current_time - self._jwks_cache_time < self._cache_ttl: + if kid and kid in self._jwks_cache: + return self._jwks_cache[kid] + elif not kid and len(self._jwks_cache) == 1: + # If no kid but only one key cached, use it + return next(iter(self._jwks_cache.values())) # Fetch JWKS try: @@ -242,16 +263,32 @@ class BearerAuthProvider(OAuthProvider): self._jwks_cache = {} for key_data in jwks_data.get("keys", []): key_kid = key_data.get("kid") + jwk = JsonWebKey.import_key(key_data) + public_key = jwk.get_public_key() + if key_kid: - jwk = JsonWebKey.import_key(key_data) - self._jwks_cache[key_kid] = jwk.get_public_key() + self._jwks_cache[key_kid] = public_key + else: + # Key without kid - use a default identifier + self._jwks_cache["_default"] = public_key self._jwks_cache_time = current_time - if kid not in self._jwks_cache: - raise ValueError(f"Key ID '{kid}' not found in JWKS") - - return self._jwks_cache[kid] + # Select the appropriate key + if kid: + if kid not in self._jwks_cache: + raise ValueError(f"Key ID '{kid}' not found in JWKS") + return self._jwks_cache[kid] + else: + # No kid in token - only allow if there's exactly one key + if len(self._jwks_cache) == 1: + return next(iter(self._jwks_cache.values())) + elif len(self._jwks_cache) > 1: + raise ValueError( + "Multiple keys in JWKS but no key ID (kid) in token" + ) + else: + raise ValueError("No keys found in JWKS") except Exception as e: raise ValueError(f"Failed to fetch JWKS: {e}") diff --git a/tests/auth/providers/test_bearer.py b/tests/auth/providers/test_bearer.py index ccc84fe73..a54f4d416 100644 --- a/tests/auth/providers/test_bearer.py +++ b/tests/auth/providers/test_bearer.py @@ -3,10 +3,15 @@ from typing import Any import httpx import pytest +from pytest_httpx import HTTPXMock from fastmcp import Client, FastMCP from fastmcp.client.auth import BearerAuth -from fastmcp.server.auth.providers.bearer import BearerAuthProvider, RSAKeyPair +from fastmcp.server.auth.providers.bearer import ( + BearerAuthProvider, + JWKSData, + RSAKeyPair, +) from fastmcp.utilities.tests import run_server_in_process @@ -103,6 +108,194 @@ class TestRSAKeyPair: # We'll validate the scopes in the BearerToken tests +class TestBearerTokenJWKS: + """Tests for JWKS URI functionality.""" + + @pytest.fixture + def jwks_provider(self, rsa_key_pair: RSAKeyPair) -> BearerAuthProvider: + """Provider configured with JWKS URI.""" + return BearerAuthProvider( + jwks_uri="https://test.example.com/.well-known/jwks.json", + issuer="https://test.example.com", + audience="https://api.example.com", + ) + + @pytest.fixture + def mock_jwks_data(self, rsa_key_pair: RSAKeyPair) -> JWKSData: + """Create mock JWKS data from RSA key pair.""" + from authlib.jose import JsonWebKey + + # Create JWK from the RSA public key + jwk = JsonWebKey.import_key(rsa_key_pair.public_key) + jwk_data = jwk.as_dict() + jwk_data["kid"] = "test-key-1" + jwk_data["alg"] = "RS256" + + return {"keys": [jwk_data]} + + async def test_jwks_token_validation( + self, + rsa_key_pair: RSAKeyPair, + jwks_provider: BearerAuthProvider, + mock_jwks_data: JWKSData, + httpx_mock: HTTPXMock, + ): + """Test token validation using JWKS URI.""" + httpx_mock.add_response( + url="https://test.example.com/.well-known/jwks.json", + json=mock_jwks_data, + ) + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + ) + + access_token = await jwks_provider.load_access_token(token) + assert access_token is not None + assert access_token.client_id == "test-user" + + async def test_jwks_token_validation_with_invalid_key( + self, + rsa_key_pair: RSAKeyPair, + jwks_provider: BearerAuthProvider, + mock_jwks_data: JWKSData, + httpx_mock: HTTPXMock, + ): + httpx_mock.add_response( + url="https://test.example.com/.well-known/jwks.json", + json=mock_jwks_data, + ) + token = RSAKeyPair.generate().create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + ) + + access_token = await jwks_provider.load_access_token(token) + assert access_token is None + + async def test_jwks_token_validation_with_kid( + self, + rsa_key_pair: RSAKeyPair, + jwks_provider: BearerAuthProvider, + mock_jwks_data: JWKSData, + httpx_mock: HTTPXMock, + ): + mock_jwks_data["keys"][0]["kid"] = "test-key-1" + httpx_mock.add_response( + url="https://test.example.com/.well-known/jwks.json", + json=mock_jwks_data, + ) + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + kid="test-key-1", + ) + + access_token = await jwks_provider.load_access_token(token) + assert access_token is not None + assert access_token.client_id == "test-user" + + async def test_jwks_token_validation_with_kid_and_no_kid_in_token( + self, + rsa_key_pair: RSAKeyPair, + jwks_provider: BearerAuthProvider, + mock_jwks_data: JWKSData, + httpx_mock: HTTPXMock, + ): + mock_jwks_data["keys"][0]["kid"] = "test-key-1" + httpx_mock.add_response( + url="https://test.example.com/.well-known/jwks.json", + json=mock_jwks_data, + ) + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + ) + + access_token = await jwks_provider.load_access_token(token) + assert access_token is not None + assert access_token.client_id == "test-user" + + async def test_jwks_token_validation_with_no_kid_and_kid_in_jwks( + self, + rsa_key_pair: RSAKeyPair, + jwks_provider: BearerAuthProvider, + mock_jwks_data: JWKSData, + httpx_mock: HTTPXMock, + ): + mock_jwks_data["keys"][0]["kid"] = "test-key-1" + httpx_mock.add_response( + url="https://test.example.com/.well-known/jwks.json", + json=mock_jwks_data, + ) + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + ) + + access_token = await jwks_provider.load_access_token(token) + assert access_token is not None + assert access_token.client_id == "test-user" + + async def test_jwks_token_validation_with_kid_mismatch( + self, + rsa_key_pair: RSAKeyPair, + jwks_provider: BearerAuthProvider, + mock_jwks_data: JWKSData, + httpx_mock: HTTPXMock, + ): + mock_jwks_data["keys"][0]["kid"] = "test-key-1" + httpx_mock.add_response( + url="https://test.example.com/.well-known/jwks.json", + json=mock_jwks_data, + ) + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + kid="test-key-2", + ) + + access_token = await jwks_provider.load_access_token(token) + assert access_token is None + + async def test_jwks_token_validation_with_multiple_keys_and_no_kid_in_token( + self, + rsa_key_pair: RSAKeyPair, + jwks_provider: BearerAuthProvider, + mock_jwks_data: JWKSData, + httpx_mock: HTTPXMock, + ): + mock_jwks_data["keys"] = [ + { + "kid": "test-key-1", + "alg": "RS256", + }, + { + "kid": "test-key-2", + "alg": "RS256", + }, + ] + + httpx_mock.add_response( + url="https://test.example.com/.well-known/jwks.json", + json=mock_jwks_data, + ) + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + ) + + access_token = await jwks_provider.load_access_token(token) + assert access_token is None + + class TestBearerToken: def test_initialization_with_public_key(self, rsa_key_pair: RSAKeyPair): """Test provider initialization with public key.""" @@ -143,7 +336,6 @@ class TestBearerToken: issuer="https://test.example.com", ) - @pytest.mark.asyncio async def test_valid_token_validation( self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider ): @@ -163,7 +355,6 @@ class TestBearerToken: assert "write" in access_token.scopes assert access_token.expires_at is not None - @pytest.mark.asyncio async def test_expired_token_rejection( self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider ): @@ -178,7 +369,6 @@ class TestBearerToken: access_token = await bearer_provider.load_access_token(token) assert access_token is None - @pytest.mark.asyncio async def test_invalid_issuer_rejection( self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider ): @@ -192,7 +382,6 @@ class TestBearerToken: access_token = await bearer_provider.load_access_token(token) assert access_token is None - @pytest.mark.asyncio async def test_invalid_audience_rejection( self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider ): @@ -206,7 +395,6 @@ class TestBearerToken: access_token = await bearer_provider.load_access_token(token) assert access_token is None - @pytest.mark.asyncio async def test_no_issuer_validation_when_none(self, rsa_key_pair: RSAKeyPair): """Test that issuer validation is skipped when provider has no issuer configured.""" provider = BearerAuthProvider( @@ -221,7 +409,6 @@ class TestBearerToken: access_token = await provider.load_access_token(token) assert access_token is not None - @pytest.mark.asyncio async def test_no_audience_validation_when_none(self, rsa_key_pair: RSAKeyPair): """Test that audience validation is skipped when provider has no audience configured.""" provider = BearerAuthProvider( @@ -239,7 +426,6 @@ class TestBearerToken: access_token = await provider.load_access_token(token) assert access_token is not None - @pytest.mark.asyncio async def test_multiple_audiences_validation(self, rsa_key_pair: RSAKeyPair): """Test validation with multiple audiences in token.""" provider = BearerAuthProvider( @@ -259,7 +445,6 @@ class TestBearerToken: access_token = await provider.load_access_token(token) assert access_token is not None - @pytest.mark.asyncio async def test_scope_extraction_string( self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider ): @@ -276,7 +461,6 @@ class TestBearerToken: assert access_token is not None assert set(access_token.scopes) == {"read", "write", "admin"} - @pytest.mark.asyncio async def test_scope_extraction_list( self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider ): @@ -293,7 +477,6 @@ class TestBearerToken: assert access_token is not None assert set(access_token.scopes) == {"read", "write"} - @pytest.mark.asyncio async def test_no_scopes( self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider ): @@ -310,7 +493,6 @@ class TestBearerToken: assert access_token is not None assert access_token.scopes == [] - @pytest.mark.asyncio async def test_malformed_token_rejection(self, bearer_provider: BearerAuthProvider): """Test rejection of malformed tokens.""" malformed_tokens = [ @@ -325,7 +507,6 @@ class TestBearerToken: access_token = await bearer_provider.load_access_token(token) assert access_token is None - @pytest.mark.asyncio async def test_invalid_signature_rejection( self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider ): @@ -341,7 +522,6 @@ class TestBearerToken: access_token = await bearer_provider.load_access_token(token) assert access_token is None - @pytest.mark.asyncio async def test_client_id_fallback( self, rsa_key_pair: RSAKeyPair, bearer_provider: BearerAuthProvider ): @@ -367,9 +547,10 @@ class TestFastMCPBearerAuth: assert isinstance(mcp.auth, BearerAuthProvider) async def test_unauthorized_access(self, mcp_server_url: str): - with pytest.raises(httpx.HTTPStatusError, match="401"): + with pytest.raises(httpx.HTTPStatusError) as exc_info: async with Client(mcp_server_url) as client: tools = await client.list_tools() # noqa: F841 + assert exc_info.value.response.status_code == 401 assert "tools" not in locals() async def test_authorized_access(self, mcp_server_url: str, bearer_token): @@ -378,9 +559,10 @@ class TestFastMCPBearerAuth: assert tools async def test_invalid_token_raises_401(self, mcp_server_url: str): - with pytest.raises(httpx.HTTPStatusError, match="401"): + with pytest.raises(httpx.HTTPStatusError) as exc_info: async with Client(mcp_server_url, auth=BearerAuth("invalid")) as client: tools = await client.list_tools() # noqa: F841 + assert exc_info.value.response.status_code == 401 assert "tools" not in locals() async def test_expired_token(self, mcp_server_url: str, rsa_key_pair: RSAKeyPair): @@ -391,16 +573,62 @@ class TestFastMCPBearerAuth: expires_in_seconds=-3600, ) - with pytest.raises(httpx.HTTPStatusError, match="401"): + with pytest.raises(httpx.HTTPStatusError) as exc_info: async with Client(mcp_server_url, auth=BearerAuth(token)) as client: tools = await client.list_tools() # noqa: F841 + assert exc_info.value.response.status_code == 401 assert "tools" not in locals() async def test_token_with_bad_signature(self, mcp_server_url: str): rsa_key_pair = RSAKeyPair.generate() token = rsa_key_pair.create_token() - with pytest.raises(httpx.HTTPStatusError, match="401"): + with pytest.raises(httpx.HTTPStatusError) as exc_info: async with Client(mcp_server_url, auth=BearerAuth(token)) as client: tools = await client.list_tools() # noqa: F841 + assert exc_info.value.response.status_code == 401 assert "tools" not in locals() + + async def test_token_with_insufficient_scopes( + self, mcp_server_url: str, rsa_key_pair: RSAKeyPair + ): + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + scopes=["read"], + ) + + with run_server_in_process( + run_mcp_server, + public_key=rsa_key_pair.public_key, + auth_kwargs=dict(required_scopes=["read", "write"]), + run_kwargs=dict(transport="streamable-http"), + ) as url: + mcp_server_url = f"{url}/mcp" + with pytest.raises(httpx.HTTPStatusError) as exc_info: + async with Client(mcp_server_url, auth=BearerAuth(token)) as client: + tools = await client.list_tools() # noqa: F841 + assert exc_info.value.response.status_code == 403 + assert "tools" not in locals() + + async def test_token_with_sufficient_scopes( + self, mcp_server_url: str, rsa_key_pair: RSAKeyPair + ): + token = rsa_key_pair.create_token( + subject="test-user", + issuer="https://test.example.com", + audience="https://api.example.com", + scopes=["read", "write"], + ) + + with run_server_in_process( + run_mcp_server, + public_key=rsa_key_pair.public_key, + auth_kwargs=dict(required_scopes=["read", "write"]), + run_kwargs=dict(transport="streamable-http"), + ) as url: + mcp_server_url = f"{url}/mcp" + async with Client(mcp_server_url, auth=BearerAuth(token)) as client: + tools = await client.list_tools() + assert tools diff --git a/uv.lock b/uv.lock index e1cd22726..5097b2bcc 100644 --- a/uv.lock +++ b/uv.lock @@ -456,6 +456,7 @@ dev = [ { name = "pytest-cov" }, { name = "pytest-env" }, { name = "pytest-flakefinder" }, + { name = "pytest-httpx" }, { name = "pytest-report" }, { name = "pytest-timeout" }, { name = "pytest-xdist" }, @@ -490,6 +491,7 @@ dev = [ { name = "pytest-cov", specifier = ">=6.1.1" }, { name = "pytest-env", specifier = ">=1.1.5" }, { name = "pytest-flakefinder" }, + { name = "pytest-httpx", specifier = ">=0.35.0" }, { name = "pytest-report", specifier = ">=0.2.1" }, { name = "pytest-timeout", specifier = ">=2.4.0" }, { name = "pytest-xdist", specifier = ">=3.6.1" }, @@ -1160,6 +1162,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/33/8b/06787150d0fd0cbd3a8054262b56f91631c7778c1bc91bf4637e47f909ad/pytest_flakefinder-1.1.0-py2.py3-none-any.whl", hash = "sha256:741e0e8eea427052f5b8c89c2b3c3019a50c39a59ce4df6a305a2c2d9ba2bd13", size = 4644, upload-time = "2022-10-26T18:27:52.128Z" }, ] +[[package]] +name = "pytest-httpx" +version = "0.35.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1f/89/5b12b7b29e3d0af3a4b9c071ee92fa25a9017453731a38f08ba01c280f4c/pytest_httpx-0.35.0.tar.gz", hash = "sha256:d619ad5d2e67734abfbb224c3d9025d64795d4b8711116b1a13f72a251ae511f", size = 54146, upload-time = "2024-11-28T19:16:54.237Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b0/ed/026d467c1853dd83102411a78126b4842618e86c895f93528b0528c7a620/pytest_httpx-0.35.0-py3-none-any.whl", hash = "sha256:ee11a00ffcea94a5cbff47af2114d34c5b231c326902458deed73f9c459fd744", size = 19442, upload-time = "2024-11-28T19:16:52.787Z" }, +] + [[package]] name = "pytest-report" version = "0.2.1" From c6d168ac2072fd9024ad00777b827dd045a54590 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 1 Jun 2025 12:49:33 -0400 Subject: [PATCH 5/5] update typing --- src/fastmcp/server/auth/providers/bearer.py | 2 +- tests/auth/providers/test_bearer.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/fastmcp/server/auth/providers/bearer.py b/src/fastmcp/server/auth/providers/bearer.py index 419ed255c..9c7c34512 100644 --- a/src/fastmcp/server/auth/providers/bearer.py +++ b/src/fastmcp/server/auth/providers/bearer.py @@ -264,7 +264,7 @@ class BearerAuthProvider(OAuthProvider): for key_data in jwks_data.get("keys", []): key_kid = key_data.get("kid") jwk = JsonWebKey.import_key(key_data) - public_key = jwk.get_public_key() + public_key = jwk.get_public_key() # type: ignore if key_kid: self._jwks_cache[key_kid] = public_key diff --git a/tests/auth/providers/test_bearer.py b/tests/auth/providers/test_bearer.py index a54f4d416..ff127f727 100644 --- a/tests/auth/providers/test_bearer.py +++ b/tests/auth/providers/test_bearer.py @@ -9,6 +9,7 @@ from fastmcp import Client, FastMCP from fastmcp.client.auth import BearerAuth from fastmcp.server.auth.providers.bearer import ( BearerAuthProvider, + JWKData, JWKSData, RSAKeyPair, ) @@ -126,8 +127,8 @@ class TestBearerTokenJWKS: from authlib.jose import JsonWebKey # Create JWK from the RSA public key - jwk = JsonWebKey.import_key(rsa_key_pair.public_key) - jwk_data = jwk.as_dict() + jwk = JsonWebKey.import_key(rsa_key_pair.public_key) # type: ignore + jwk_data: JWKData = jwk.as_dict() # type: ignore jwk_data["kid"] = "test-key-1" jwk_data["alg"] = "RS256"