Adjust more things to fit spec

Author: praboud-ant

src/mcp/server/auth/handlers/authorize.py

@@ -39,9 +39,7 @@ class AuthorizationRequest(BaseModel):
         description="Optional scope; if specified, should be "
         "a space-separated list of scope strings",
     )
-
-    class Config:
-        extra = "ignore"
+    
 
 
 def validate_scope(

src/mcp/server/auth/handlers/revoke.py

@@ -4,9 +4,9 @@
 Corresponds to TypeScript file: src/server/auth/handlers/revoke.ts
 """
 
-from typing import Callable
+from typing import Callable, Optional
 
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
 from starlette.requests import Request
 from starlette.responses import Response
 
@@ -17,8 +17,8 @@
     ClientAuthenticator,
     ClientAuthRequest,
 )
-from mcp.server.auth.provider import OAuthServerProvider
-from mcp.shared.auth import OAuthTokenRevocationRequest
+from mcp.server.auth.provider import OAuthServerProvider, OAuthTokenRevocationRequest
+
 
 
 class RevocationRequest(OAuthTokenRevocationRequest, ClientAuthRequest):
@@ -28,18 +28,6 @@ class RevocationRequest(OAuthTokenRevocationRequest, ClientAuthRequest):
 def create_revocation_handler(
     provider: OAuthServerProvider, client_authenticator: ClientAuthenticator
 ) -> Callable:
-    """
-    Create a handler for OAuth 2.0 Token Revocation.
-
-    Corresponds to revocationHandler in src/server/auth/handlers/revoke.ts
-
-    Args:
-        provider: The OAuth server provider
-
-    Returns:
-        A Starlette endpoint handler function
-    """
-
     async def revocation_handler(request: Request) -> Response:
         """
         Handler for the OAuth 2.0 Token Revocation endpoint.

src/mcp/server/auth/handlers/token.py

@@ -6,9 +6,10 @@
 
 import base64
 import hashlib
+import time
 from typing import Annotated, Callable, Literal, Optional, Union
 
-from pydantic import Field, RootModel, ValidationError
+from pydantic import AnyHttpUrl, Field, RootModel, ValidationError
 from starlette.requests import Request
 
 from mcp.server.auth.errors import (
@@ -24,13 +25,19 @@
 
 
 class AuthorizationCodeRequest(ClientAuthRequest):
+    # See https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.3
     grant_type: Literal["authorization_code"]
     code: str = Field(..., description="The authorization code")
+    redirect_uri: AnyHttpUrl | None = Field(
+        ..., description="Must be the same as redirect URI provided in /authorize"
+    )
+    client_id: str
+    # See https://datatracker.ietf.org/doc/html/rfc7636#section-4.5
     code_verifier: str = Field(..., description="PKCE code verifier")
-    # TODO: this should take redirect_uri
 
 
 class RefreshTokenRequest(ClientAuthRequest):
+    # See https://datatracker.ietf.org/doc/html/rfc6749#section-6
     grant_type: Literal["refresh_token"]
     refresh_token: str = Field(..., description="The refresh token")
     scope: Optional[str] = Field(None, description="Optional scope parameter")
@@ -42,7 +49,7 @@ class TokenRequest(RootModel):
         Field(discriminator="grant_type"),
     ]
 
-
+AUTH_CODE_TTL = 300 # seconds
 
 def create_token_handler(
     provider: OAuthServerProvider, client_authenticator: ClientAuthenticator
@@ -65,22 +72,28 @@ async def token_handler(request: Request):
 
         match token_request:
             case AuthorizationCodeRequest():
-                # TODO: verify that the redirect URIs match
-                # see https://datatracker.ietf.org/doc/html/rfc6749#section-10.6
-                # TODO: enforce TTL on the authorization code
-
-                # Verify PKCE code verifier
-                expected_challenge = await provider.challenge_for_authorization_code(
+                auth_code_metadata = await provider.load_authorization_code_metadata(
                     client_info, token_request.code
                 )
-                if expected_challenge is None:
+                if auth_code_metadata is None or auth_code_metadata.client_id != token_request.client_id:
                     raise InvalidRequestError("Invalid authorization code")
 
-                # Calculate challenge from verifier
+                # make auth codes expire after a deadline
+                # see https://datatracker.ietf.org/doc/html/rfc6749#section-10.5
+                expires_at = auth_code_metadata.issued_at + AUTH_CODE_TTL
+                if expires_at < time.time():
+                    raise InvalidRequestError("authorization code has expired")
+
+                # verify redirect_uri doesn't change between /authorize and /tokens
+                # see https://datatracker.ietf.org/doc/html/rfc6749#section-10.6
+                if token_request.redirect_uri != auth_code_metadata.redirect_uri:
+                    raise InvalidRequestError("redirect_uri did not match redirect_uri used when authorization code was created")
+
+                # Verify PKCE code verifier
                 sha256 = hashlib.sha256(token_request.code_verifier.encode()).digest()
-                actual_challenge = base64.urlsafe_b64encode(sha256).decode().rstrip("=")
+                hashed_code_verifier = base64.urlsafe_b64encode(sha256).decode().rstrip("=")
 
-                if actual_challenge != expected_challenge:
+                if hashed_code_verifier != auth_code_metadata.code_challenge:
                     raise InvalidRequestError(
                         "code_verifier does not match the challenge"
                     )

src/mcp/server/auth/provider.py

@@ -4,7 +4,7 @@
 Corresponds to TypeScript file: src/server/auth/provider.ts
 """
 
-from typing import List, Optional, Protocol
+from typing import List, Literal, Optional, Protocol
 
 from pydantic import AnyHttpUrl, BaseModel
 
@@ -28,6 +28,18 @@ class AuthorizationParams(BaseModel):
     code_challenge: str
     redirect_uri: AnyHttpUrl
 
+class AuthorizationCodeMeta(BaseModel):
+    issued_at: float
+    client_id: str
+    code_challenge: str
+    redirect_uri: AnyHttpUrl
+class OAuthTokenRevocationRequest(BaseModel):
+    """
+    # See https://datatracker.ietf.org/doc/html/rfc7009#section-2.1
+    """
+
+    token: str
+    token_type_hint: Optional[Literal["access_token", "refresh_token"]] = None
 
 class OAuthRegisteredClientsStore(Protocol):
     """
@@ -91,11 +103,11 @@ async def create_authorization_code(
         """
         ...
 
-    async def challenge_for_authorization_code(
+    async def load_authorization_code_metadata(
         self, client: OAuthClientInformationFull, authorization_code: str
-    ) -> str | None:
+    ) -> AuthorizationCodeMeta | None:
         """
-        Returns the code_challenge that was used when the indicated authorization began.
+        Loads metadata for the authorization code challenge.
 
         Args:
             client: The client that requested the authorization code.

src/mcp/shared/auth.py

@@ -11,25 +11,21 @@
 
 class OAuthErrorResponse(BaseModel):
     """
-    OAuth 2.1 error response.
-
-    Corresponds to OAuthErrorResponseSchema in src/shared/auth.ts
+    See https://datatracker.ietf.org/doc/html/rfc6749#section-5.2
     """
 
-    error: str
+    error: Literal["invalid_request", "invalid_client", "invalid_grant", "unauthorized_client", "unsupported_grant_type", "invalid_scope"]
     error_description: Optional[str] = None
     error_uri: Optional[AnyHttpUrl] = None
 
 
 class OAuthTokens(BaseModel):
     """
-    OAuth 2.1 token response.
-
-    Corresponds to OAuthTokensSchema in src/shared/auth.ts
+    See https://datatracker.ietf.org/doc/html/rfc6749#section-5.1
     """
 
     access_token: str
-    token_type: str
+    token_type: Literal["bearer"] = "bearer"
     expires_in: Optional[int] = None
     scope: Optional[str] = None
     refresh_token: Optional[str] = None

tests/server/fastmcp/auth/test_auth_integration.py

@@ -19,9 +19,11 @@
 
 from mcp.server.auth.errors import InvalidTokenError
 from mcp.server.auth.provider import (
+    AuthorizationCodeMeta,
     AuthorizationParams,
     OAuthRegisteredClientsStore,
     OAuthServerProvider,
+    OAuthTokenRevocationRequest,
 )
 from mcp.server.auth.router import (
     ClientRegistrationOptions,
@@ -32,7 +34,6 @@
 from mcp.server.fastmcp import FastMCP
 from mcp.shared.auth import (
     OAuthClientInformationFull,
-    OAuthTokenRevocationRequest,
     OAuthTokens,
 )
 from mcp.types import JSONRPCRequest
@@ -74,32 +75,19 @@ async def create_authorization_code(
         code = f"code_{int(time.time())}"
 
         # Store the code for later verification
-        self.auth_codes[code] = {
-            "client_id": client.client_id,
-            "code_challenge": params.code_challenge,
-            "redirect_uri": params.redirect_uri,
-            "expires_at": int(time.time()) + 600,  # 10 minutes
-        }
+        self.auth_codes[code] = AuthorizationCodeMeta(
+            client_id= client.client_id,
+            code_challenge= params.code_challenge,
+            redirect_uri= params.redirect_uri,
+            issued_at= time.time(),
+        )
 
         return code
 
-    async def challenge_for_authorization_code(
+    async def load_authorization_code_metadata(
         self, client: OAuthClientInformationFull, authorization_code: str
-    ) -> str:
-        # Get the stored code info
-        code_info = self.auth_codes.get(authorization_code)
-        if not code_info:
-            raise InvalidTokenError("Invalid authorization code")
-
-        # Check if code is expired
-        if code_info["expires_at"] < int(time.time()):
-            raise InvalidTokenError("Authorization code has expired")
-
-        # Check if the code was issued to this client
-        if code_info["client_id"] != client.client_id:
-            raise InvalidTokenError("Authorization code was not issued to this client")
-
-        return code_info["code_challenge"]
+    ) -> AuthorizationCodeMeta | None:
+        return self.auth_codes.get(authorization_code)
 
     async def exchange_authorization_code(
         self, client: OAuthClientInformationFull, authorization_code: str
@@ -109,14 +97,6 @@ async def exchange_authorization_code(
         if not code_info:
             raise InvalidTokenError("Invalid authorization code")
 
-        # Check if code is expired
-        if code_info["expires_at"] < int(time.time()):
-            raise InvalidTokenError("Authorization code has expired")
-
-        # Check if the code was issued to this client
-        if code_info["client_id"] != client.client_id:
-            raise InvalidTokenError("Authorization code was not issued to this client")
-
         # Generate an access token and refresh token
         access_token = f"access_{secrets.token_hex(32)}"
         refresh_token = f"refresh_{secrets.token_hex(32)}"
@@ -436,6 +416,7 @@ async def test_authorization_flow(
                 "client_secret": client_info["client_secret"],
                 "code": auth_code,
                 "code_verifier": code_verifier,
+                "redirect_uri": "https://client.example.com/callback",
             },
         )
         assert response.status_code == 200
@@ -465,6 +446,7 @@ async def test_authorization_flow(
                 "client_id": client_info["client_id"],
                 "client_secret": client_info["client_secret"],
                 "refresh_token": refresh_token,
+                "redirect_uri": "https://client.example.com/callback",
             },
         )
         assert response.status_code == 200
@@ -585,6 +567,7 @@ def test_tool(x: int) -> str:
                 "client_secret": client_info["client_secret"],
                 "code": auth_code,
                 "code_verifier": code_verifier,
+                "redirect_uri": "https://client.example.com/callback",
             },
         )
         assert response.status_code == 200