diff --git a/src/mcp/server/auth/handlers/token.py b/src/mcp/server/auth/handlers/token.py index 00df115573..ac0d25dd55 100644 --- a/src/mcp/server/auth/handlers/token.py +++ b/src/mcp/server/auth/handlers/token.py @@ -1,5 +1,6 @@ import base64 import hashlib +import inspect import time from dataclasses import dataclass from typing import Annotated, Any, Literal @@ -80,6 +81,47 @@ class TokenErrorResponse(BaseModel): TokenSuccessResponse = OAuthToken +def _accepts_resource(func: Any) -> bool: + try: + sig = inspect.signature(func) + return "resource" in sig.parameters or any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() + ) + except (ValueError, TypeError): + return True + + +async def _exchange_authorization_code( + provider: OAuthAuthorizationServerProvider[Any, Any, Any], + client_info: Any, + auth_code: Any, + resource: str | None, +) -> OAuthToken: + if _accepts_resource(provider.exchange_authorization_code): + try: + return await provider.exchange_authorization_code(client_info, auth_code, resource=resource) + except TypeError as te: + if "resource" not in str(te): + raise + return await provider.exchange_authorization_code(client_info, auth_code) + + +async def _exchange_refresh_token( + provider: OAuthAuthorizationServerProvider[Any, Any, Any], + client_info: Any, + refresh_token: Any, + scopes: list[str], + resource: str | None, +) -> OAuthToken: + if _accepts_resource(provider.exchange_refresh_token): + try: + return await provider.exchange_refresh_token(client_info, refresh_token, scopes, resource=resource) + except TypeError as te: + if "resource" not in str(te): + raise + return await provider.exchange_refresh_token(client_info, refresh_token, scopes) + + @dataclass class TokenHandler: provider: OAuthAuthorizationServerProvider[Any, Any, Any] @@ -119,8 +161,11 @@ async def handle(self, request: Request): try: form_data = await request.form() - # TODO(Marcelo): Can someone check if this `dict()` wrapper is necessary? - token_request = token_request_adapter.validate_python(dict(form_data)) + # Parse token request from body form data, falling back to query parameters + # for parameters like RFC 8707 `resource` if provided in query string. + request_data = dict(request.query_params) + request_data.update(dict(form_data)) + token_request = token_request_adapter.validate_python(request_data) except ValidationError as validation_error: return self.response( TokenErrorResponse( @@ -197,7 +242,9 @@ async def handle(self, request: Request): try: # Exchange authorization code for tokens - tokens = await self.provider.exchange_authorization_code(client_info, auth_code) + tokens = await _exchange_authorization_code( + self.provider, client_info, auth_code, token_request.resource + ) except TokenError as e: return self.response(TokenErrorResponse(error=e.error, error_description=e.error_description)) @@ -235,7 +282,9 @@ async def handle(self, request: Request): try: # Exchange refresh token for new tokens - tokens = await self.provider.exchange_refresh_token(client_info, refresh_token, scopes) + tokens = await _exchange_refresh_token( + self.provider, client_info, refresh_token, scopes, token_request.resource + ) except TokenError as e: return self.response(TokenErrorResponse(error=e.error, error_description=e.error_description)) diff --git a/src/mcp/server/auth/provider.py b/src/mcp/server/auth/provider.py index bc8ba3e518..9066545d65 100644 --- a/src/mcp/server/auth/provider.py +++ b/src/mcp/server/auth/provider.py @@ -228,13 +228,17 @@ async def load_authorization_code( ... async def exchange_authorization_code( - self, client: OAuthClientInformationFull, authorization_code: AuthorizationCodeT + self, + client: OAuthClientInformationFull, + authorization_code: AuthorizationCodeT, + resource: str | None = None, ) -> OAuthToken: """Exchanges an authorization code for an access token and refresh token. Args: client: The client exchanging the authorization code. authorization_code: The authorization code to exchange. + resource: Optional RFC 8707 resource indicator requested by the client in the token request. Returns: The OAuth token, containing access and refresh tokens. @@ -261,6 +265,7 @@ async def exchange_refresh_token( client: OAuthClientInformationFull, refresh_token: RefreshTokenT, scopes: list[str], + resource: str | None = None, ) -> OAuthToken: """Exchanges a refresh token for an access token and refresh token. @@ -270,6 +275,7 @@ async def exchange_refresh_token( client: The client exchanging the refresh token. refresh_token: The refresh token to exchange. scopes: Optional scopes to request with the new access token. + resource: Optional RFC 8707 resource indicator requested by the client in the token request. Returns: The OAuth token, containing access and refresh tokens. diff --git a/tests/interaction/auth/_provider.py b/tests/interaction/auth/_provider.py index 4e422947b4..9c8fae515a 100644 --- a/tests/interaction/auth/_provider.py +++ b/tests/interaction/auth/_provider.py @@ -79,6 +79,8 @@ def __init__( # The most recent jwt-bearer request the SDK handler passed to exchange_identity_assertion, # for tests to assert what the client sent (None until the first exchange). self.last_assertion_params: IdentityAssertionParams | None = None + self.last_auth_code_resource: str | None = None + self.last_refresh_resource: str | None = None def _next_expires_in(self) -> int: self._tokens_issued += 1 @@ -145,19 +147,24 @@ async def load_authorization_code( return self.codes.get(authorization_code) async def exchange_authorization_code( - self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode + self, + client: OAuthClientInformationFull, + authorization_code: AuthorizationCode, + resource: str | None = None, ) -> OAuthToken: """Mint an access token and a refresh token for a valid authorization code, then consume the code.""" assert client.client_id is not None + self.last_auth_code_resource = resource + target_resource = resource if resource is not None else authorization_code.resource access = self.mint_access_token( - client_id=client.client_id, scopes=authorization_code.scopes, resource=authorization_code.resource + client_id=client.client_id, scopes=authorization_code.scopes, resource=target_resource ) refresh = f"refresh_{secrets.token_hex(16)}" self.refresh_tokens[refresh] = RefreshToken( token=refresh, client_id=client.client_id, scopes=authorization_code.scopes, - resource=authorization_code.resource, + resource=target_resource, ) del self.codes[authorization_code.code] return OAuthToken( @@ -177,17 +184,23 @@ async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_t return self.refresh_tokens.get(refresh_token) async def exchange_refresh_token( - self, client: OAuthClientInformationFull, refresh_token: RefreshToken, scopes: list[str] + self, + client: OAuthClientInformationFull, + refresh_token: RefreshToken, + scopes: list[str], + resource: str | None = None, ) -> OAuthToken: """Mint a new access token and rotate the refresh token, consuming the old one.""" assert client.client_id is not None + self.last_refresh_resource = resource if self._fail_next_refresh: self._fail_next_refresh = False raise TokenError(error="invalid_grant", error_description="refresh denied by harness") - access = self.mint_access_token(client_id=client.client_id, scopes=scopes, resource=refresh_token.resource) + target_resource = resource if resource is not None else refresh_token.resource + access = self.mint_access_token(client_id=client.client_id, scopes=scopes, resource=target_resource) new_refresh = f"refresh_{secrets.token_hex(16)}" self.refresh_tokens[new_refresh] = RefreshToken( - token=new_refresh, client_id=client.client_id, scopes=scopes, resource=refresh_token.resource + token=new_refresh, client_id=client.client_id, scopes=scopes, resource=target_resource ) del self.refresh_tokens[refresh_token.token] return OAuthToken( diff --git a/tests/server/auth/test_token_resource.py b/tests/server/auth/test_token_resource.py new file mode 100644 index 0000000000..0e4d649837 --- /dev/null +++ b/tests/server/auth/test_token_resource.py @@ -0,0 +1,361 @@ +"""Tests for RFC 8707 resource indicator handling in token exchange grants.""" + +import base64 +import hashlib +import time + +import httpx2 +import pytest +from httpx2 import ASGITransport +from pydantic import AnyHttpUrl, AnyUrl +from starlette.applications import Starlette + +from mcp.server.auth.provider import ( + AccessToken, + AuthorizationCode, + AuthorizationParams, + IdentityAssertionParams, + OAuthAuthorizationServerProvider, + RefreshToken, + TokenError, +) +from mcp.server.auth.routes import create_auth_routes +from mcp.server.auth.settings import ClientRegistrationOptions, RevocationOptions +from mcp.shared.auth import OAuthClientInformationFull, OAuthToken + +pytestmark = pytest.mark.anyio + +CLIENT_ID = "test-client" +CLIENT_SECRET = "test-secret" +REDIRECT_URI = "https://client.example.com/callback" +CODE_VERIFIER = "test-verifier-string-long-enough-for-pkce-verification-12345" +sha256 = hashlib.sha256(CODE_VERIFIER.encode()).digest() +CODE_CHALLENGE = base64.urlsafe_b64encode(sha256).decode().rstrip("=") + + +class ResourceAwareProvider(OAuthAuthorizationServerProvider[AuthorizationCode, RefreshToken, AccessToken]): + def __init__(self) -> None: + self.clients: dict[str, OAuthClientInformationFull] = { + CLIENT_ID: OAuthClientInformationFull( + client_id=CLIENT_ID, + client_secret=CLIENT_SECRET, + redirect_uris=[AnyUrl(REDIRECT_URI)], + grant_types=["authorization_code", "refresh_token"], + token_endpoint_auth_method="client_secret_post", + scope="mcp", + ) + } + self.codes: dict[str, AuthorizationCode] = {} + self.refresh_tokens: dict[str, RefreshToken] = {} + self.access_tokens: dict[str, AccessToken] = {} + self.last_exchange_auth_code_resource: str | None = None + self.last_exchange_refresh_resource: str | None = None + + async def get_client(self, client_id: str) -> OAuthClientInformationFull | None: + return self.clients.get(client_id) + + async def register_client(self, client_info: OAuthClientInformationFull) -> None: + pass + + async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str: + raise NotImplementedError + + async def load_authorization_code( + self, client: OAuthClientInformationFull, authorization_code: str + ) -> AuthorizationCode | None: + return self.codes.get(authorization_code) + + async def exchange_authorization_code( + self, + client: OAuthClientInformationFull, + authorization_code: AuthorizationCode, + resource: str | None = None, + ) -> OAuthToken: + self.last_exchange_auth_code_resource = resource + if resource == "https://forbidden.example.com": + raise TokenError(error="invalid_target", error_description="forbidden target resource") + return OAuthToken(access_token="access-123", token_type="Bearer", expires_in=3600) + + async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_token: str) -> RefreshToken | None: + return self.refresh_tokens.get(refresh_token) + + async def exchange_refresh_token( + self, + client: OAuthClientInformationFull, + refresh_token: RefreshToken, + scopes: list[str], + resource: str | None = None, + ) -> OAuthToken: + self.last_exchange_refresh_resource = resource + if resource == "https://forbidden.example.com": + raise TokenError(error="invalid_target", error_description="forbidden target resource") + return OAuthToken(access_token="access-refreshed", token_type="Bearer", expires_in=3600) + + async def load_access_token(self, token: str) -> AccessToken | None: + return None + + async def revoke_token(self, token: AccessToken | RefreshToken) -> None: + pass + + async def exchange_identity_assertion( + self, client: OAuthClientInformationFull, params: IdentityAssertionParams + ) -> OAuthToken: + raise NotImplementedError + + +class LegacyProvider(OAuthAuthorizationServerProvider[AuthorizationCode, RefreshToken, AccessToken]): + """Provider without resource parameter on exchange methods for backwards compatibility testing.""" + + def __init__(self) -> None: + self.clients: dict[str, OAuthClientInformationFull] = { + CLIENT_ID: OAuthClientInformationFull( + client_id=CLIENT_ID, + client_secret=CLIENT_SECRET, + redirect_uris=[AnyUrl(REDIRECT_URI)], + grant_types=["authorization_code", "refresh_token"], + token_endpoint_auth_method="client_secret_post", + scope="mcp", + ) + } + self.codes: dict[str, AuthorizationCode] = {} + self.refresh_tokens: dict[str, RefreshToken] = {} + + async def get_client(self, client_id: str) -> OAuthClientInformationFull | None: + return self.clients.get(client_id) + + async def register_client(self, client_info: OAuthClientInformationFull) -> None: + pass + + async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str: + raise NotImplementedError + + async def load_authorization_code( + self, client: OAuthClientInformationFull, authorization_code: str + ) -> AuthorizationCode | None: + return self.codes.get(authorization_code) + + async def exchange_authorization_code( + self, + client: OAuthClientInformationFull, + authorization_code: AuthorizationCode, + ) -> OAuthToken: + return OAuthToken(access_token="legacy-access-123", token_type="Bearer", expires_in=3600) + + async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_token: str) -> RefreshToken | None: + return self.refresh_tokens.get(refresh_token) + + async def exchange_refresh_token( + self, + client: OAuthClientInformationFull, + refresh_token: RefreshToken, + scopes: list[str], + ) -> OAuthToken: + return OAuthToken(access_token="legacy-refreshed", token_type="Bearer", expires_in=3600) + + async def load_access_token(self, token: str) -> AccessToken | None: + return None + + async def revoke_token(self, token: AccessToken | RefreshToken) -> None: + pass + + async def exchange_identity_assertion( + self, client: OAuthClientInformationFull, params: IdentityAssertionParams + ) -> OAuthToken: + raise NotImplementedError + + +@pytest.fixture +def provider() -> ResourceAwareProvider: + prov = ResourceAwareProvider() + prov.codes["code-123"] = AuthorizationCode( + code="code-123", + client_id=CLIENT_ID, + scopes=["mcp"], + expires_at=time.time() + 300, + code_challenge=CODE_CHALLENGE, + redirect_uri=AnyUrl(REDIRECT_URI), + redirect_uri_provided_explicitly=True, + resource="https://original.example.com", + ) + prov.refresh_tokens["refresh-123"] = RefreshToken( + token="refresh-123", + client_id=CLIENT_ID, + scopes=["mcp"], + expires_at=int(time.time()) + 3600, + resource="https://original.example.com", + ) + return prov + + +@pytest.fixture +def legacy_provider() -> LegacyProvider: + prov = LegacyProvider() + prov.codes["code-legacy"] = AuthorizationCode( + code="code-legacy", + client_id=CLIENT_ID, + scopes=["mcp"], + expires_at=time.time() + 300, + code_challenge=CODE_CHALLENGE, + redirect_uri=AnyUrl(REDIRECT_URI), + redirect_uri_provided_explicitly=True, + resource="https://original.example.com", + ) + prov.refresh_tokens["refresh-legacy"] = RefreshToken( + token="refresh-legacy", + client_id=CLIENT_ID, + scopes=["mcp"], + expires_at=int(time.time()) + 3600, + resource="https://original.example.com", + ) + return prov + + +@pytest.fixture +def app(provider: ResourceAwareProvider) -> Starlette: + routes = create_auth_routes( + provider, + issuer_url=AnyHttpUrl("https://auth.example.com"), + client_registration_options=ClientRegistrationOptions(enabled=True, valid_scopes=["mcp"]), + revocation_options=RevocationOptions(enabled=False), + ) + return Starlette(routes=routes) + + +@pytest.fixture +def legacy_app(legacy_provider: LegacyProvider) -> Starlette: + routes = create_auth_routes( + legacy_provider, + issuer_url=AnyHttpUrl("https://auth.example.com"), + client_registration_options=ClientRegistrationOptions(enabled=True, valid_scopes=["mcp"]), + revocation_options=RevocationOptions(enabled=False), + ) + return Starlette(routes=routes) + + +@pytest.fixture +async def client(app: Starlette): + transport = ASGITransport(app=app) + async with httpx2.AsyncClient(transport=transport, base_url="https://auth.example.com") as http: + yield http + + +@pytest.fixture +async def legacy_client(legacy_app: Starlette): + transport = ASGITransport(app=legacy_app) + async with httpx2.AsyncClient(transport=transport, base_url="https://auth.example.com") as http: + yield http + + +async def test_authorization_code_preserves_resource_in_body( + client: httpx2.AsyncClient, provider: ResourceAwareProvider +): + response = await client.post( + "/token", + data={ + "grant_type": "authorization_code", + "client_id": CLIENT_ID, + "client_secret": CLIENT_SECRET, + "code": "code-123", + "code_verifier": CODE_VERIFIER, + "redirect_uri": REDIRECT_URI, + "resource": "https://requested.example.com", + }, + ) + assert response.status_code == 200 + assert provider.last_exchange_auth_code_resource == "https://requested.example.com" + + +async def test_authorization_code_preserves_resource_in_query( + client: httpx2.AsyncClient, provider: ResourceAwareProvider +): + response = await client.post( + "/token?resource=https://requested-query.example.com", + data={ + "grant_type": "authorization_code", + "client_id": CLIENT_ID, + "client_secret": CLIENT_SECRET, + "code": "code-123", + "code_verifier": CODE_VERIFIER, + "redirect_uri": REDIRECT_URI, + }, + ) + assert response.status_code == 200 + assert provider.last_exchange_auth_code_resource == "https://requested-query.example.com" + + +async def test_refresh_token_preserves_resource_in_body(client: httpx2.AsyncClient, provider: ResourceAwareProvider): + response = await client.post( + "/token", + data={ + "grant_type": "refresh_token", + "client_id": CLIENT_ID, + "client_secret": CLIENT_SECRET, + "refresh_token": "refresh-123", + "resource": "https://refreshed.example.com", + }, + ) + assert response.status_code == 200 + assert provider.last_exchange_refresh_resource == "https://refreshed.example.com" + + +async def test_refresh_token_preserves_resource_in_query(client: httpx2.AsyncClient, provider: ResourceAwareProvider): + response = await client.post( + "/token?resource=https://refreshed-query.example.com", + data={ + "grant_type": "refresh_token", + "client_id": CLIENT_ID, + "client_secret": CLIENT_SECRET, + "refresh_token": "refresh-123", + }, + ) + assert response.status_code == 200 + assert provider.last_exchange_refresh_resource == "https://refreshed-query.example.com" + + +async def test_provider_rejects_mismatched_resource_with_invalid_target( + client: httpx2.AsyncClient, provider: ResourceAwareProvider +): + response = await client.post( + "/token", + data={ + "grant_type": "refresh_token", + "client_id": CLIENT_ID, + "client_secret": CLIENT_SECRET, + "refresh_token": "refresh-123", + "resource": "https://forbidden.example.com", + }, + ) + assert response.status_code == 400 + data = response.json() + assert data["error"] == "invalid_target" + assert data["error_description"] == "forbidden target resource" + + +async def test_legacy_provider_without_resource_parameter_supported( + legacy_client: httpx2.AsyncClient, legacy_provider: LegacyProvider +): + response = await legacy_client.post( + "/token", + data={ + "grant_type": "authorization_code", + "client_id": CLIENT_ID, + "client_secret": CLIENT_SECRET, + "code": "code-legacy", + "code_verifier": CODE_VERIFIER, + "redirect_uri": REDIRECT_URI, + "resource": "https://requested.example.com", + }, + ) + assert response.status_code == 200 + + response = await legacy_client.post( + "/token", + data={ + "grant_type": "refresh_token", + "client_id": CLIENT_ID, + "client_secret": CLIENT_SECRET, + "refresh_token": "refresh-legacy", + "resource": "https://refreshed.example.com", + }, + ) + assert response.status_code == 200