Refactor back to authorize()

Author: praboud-ant

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

@@ -19,6 +19,10 @@
 from mcp.server.auth.provider import AuthorizationParams, OAuthServerProvider
 from mcp.shared.auth import OAuthClientInformationFull
 
+import logging
+
+logger = logging.getLogger(__name__)
+
 
 class AuthorizationRequest(BaseModel):
     client_id: str = Field(..., description="The client ID")
@@ -122,28 +126,18 @@ async def authorization_handler(request: Request) -> Response:
         )
 
         try:
-            # Let the provider handle the authorization flow
-            authorization_code = await provider.create_authorization_code(
-                client, auth_params
-            )
+            # Let the provider pick the next URI to redirect to
             response = RedirectResponse(
                 url="", status_code=302, headers={"Cache-Control": "no-store"}
             )
-
-            # Redirect with code
-            parsed_uri = urlparse(str(auth_params.redirect_uri))
-            query_params = [(k, v) for k, vs in parse_qs(parsed_uri.query) for v in vs]
-            query_params.append(("code", authorization_code))
-            if auth_params.state:
-                query_params.append(("state", auth_params.state))
-
-            redirect_url = urlunparse(
-                parsed_uri._replace(query=urlencode(query_params))
+            response.headers["location"] = await provider.authorize(
+                client, auth_params
             )
-            response.headers["location"] = redirect_url
 
             return response
         except Exception as e:
+            logger.exception("error from authorize()", exc_info=e)
+            
             return RedirectResponse(
                 url=create_error_redirect(redirect_uri, e, auth_request.state),
                 status_code=302,

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

@@ -21,7 +21,7 @@
     ClientAuthRequest,
 )
 from mcp.server.auth.provider import OAuthServerProvider
-from mcp.shared.auth import OAuthTokens
+from mcp.shared.auth import TokenErrorResponse, TokenSuccessResponse
 
 
 class AuthorizationCodeRequest(ClientAuthRequest):
@@ -54,53 +54,79 @@ class TokenRequest(RootModel):
 def create_token_handler(
     provider: OAuthServerProvider, client_authenticator: ClientAuthenticator
 ) -> Callable:
+    def response(obj: TokenSuccessResponse | TokenErrorResponse):
+        return PydanticJSONResponse(
+            content=obj,
+            headers={
+                "Cache-Control": "no-store",
+                "Pragma": "no-cache",
+            },
+        )
+
     async def token_handler(request: Request):
         try:
             form_data = await request.form()
             token_request = TokenRequest.model_validate(dict(form_data)).root
-        except ValidationError as e:
-            raise InvalidRequestError(f"Invalid request body: {e}")
+        except ValidationError as validation_error:
+            return response(TokenErrorResponse(
+                error="invalid_request",
+                error_description="\n".join(e['msg'] for e in validation_error.errors())
+
+            ))
         client_info = await client_authenticator(token_request)
 
         if token_request.grant_type not in client_info.grant_types:
-            raise InvalidRequestError(
-                f"Unsupported grant type (supported grant types are "
+            return response(TokenErrorResponse(
+                error="unsupported_grant_type",
+                error_description=f"Unsupported grant type (supported grant types are "
                 f"{client_info.grant_types})"
-            )
+            ))
 
-        tokens: OAuthTokens
+        tokens: TokenSuccessResponse
 
         match token_request:
             case AuthorizationCodeRequest():
-                auth_code_metadata = await provider.load_authorization_code_metadata(
+                auth_code = await provider.load_authorization_code(
                     client_info, token_request.code
                 )
-                if auth_code_metadata is None or auth_code_metadata.client_id != token_request.client_id:
-                    raise InvalidRequestError("Invalid authorization code")
+                if auth_code is None or auth_code.client_id != token_request.client_id:
+                    # if the authoriation code belongs to a different client, pretend it doesn't exist
+                    return response(TokenErrorResponse(
+                        error="invalid_grant",
+                        error_description=f"authorization code does not exist"
+                    ))
 
                 # 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
+                expires_at = auth_code.issued_at + AUTH_CODE_TTL
                 if expires_at < time.time():
-                    raise InvalidRequestError("authorization code has expired")
+                    return response(TokenErrorResponse(
+                        error="invalid_grant",
+                        error_description=f"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")
+                if token_request.redirect_uri != auth_code.redirect_uri:
+                    return response(TokenErrorResponse(
+                        error="invalid_request",
+                        error_description=f"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()
                 hashed_code_verifier = base64.urlsafe_b64encode(sha256).decode().rstrip("=")
 
-                if hashed_code_verifier != auth_code_metadata.code_challenge:
-                    raise InvalidRequestError(
-                        "code_verifier does not match the challenge"
-                    )
+                if hashed_code_verifier != auth_code.code_challenge:
+                    # see https://datatracker.ietf.org/doc/html/rfc7636#section-4.6
+                    return response(TokenErrorResponse(
+                        error="invalid_grant",
+                        error_description=f"incorrect code_verifier"
+                    ))
 
                 # Exchange authorization code for tokens
                 tokens = await provider.exchange_authorization_code(
-                    client_info, token_request.code
+                    client_info, auth_code
                 )
 
             case RefreshTokenRequest():
@@ -112,12 +138,6 @@ async def token_handler(request: Request):
                     client_info, token_request.refresh_token, scopes
                 )
 
-        return PydanticJSONResponse(
-            content=tokens,
-            headers={
-                "Cache-Control": "no-store",
-                "Pragma": "no-cache",
-            },
-        )
+        return response(tokens)
 
     return token_handler

src/mcp/server/auth/provider.py

@@ -5,13 +5,14 @@
 """
 
 from typing import List, Literal, Optional, Protocol
+from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
 
-from pydantic import AnyHttpUrl, BaseModel
+from pydantic import AnyHttpUrl, AnyUrl, BaseModel
 
 from mcp.server.auth.types import AuthInfo
 from mcp.shared.auth import (
     OAuthClientInformationFull,
-    OAuthTokens,
+    TokenSuccessResponse,
 )
 
 
@@ -27,7 +28,9 @@ class AuthorizationParams(BaseModel):
     code_challenge: str
     redirect_uri: AnyHttpUrl
 
-class AuthorizationCodeMeta(BaseModel):
+class AuthorizationCode(BaseModel):
+    code: str
+    scopes: list[str]
     issued_at: float
     client_id: str
     code_challenge: str
@@ -88,12 +91,33 @@ def clients_store(self) -> OAuthRegisteredClientsStore:
         """
         ...
 
-    async def create_authorization_code(
+    async def authorize(
         self, client: OAuthClientInformationFull, params: AuthorizationParams
     ) -> str:
         """
-        Generates and stores an authorization code as part of completing the /authorize
-        OAuth step.
+        Called as part of the /authorize endpoint, and returns a URL that the client
+        will be redirected to.
+        Many MCP implementations will redirect to a third-party provider to perform
+        a second OAuth exchange with that provider. In this sort of setup, the client
+        has an OAuth connection with the MCP server, and the MCP server has an OAuth
+        connection with the 3rd-party provider. At the end of this flow, the client
+        should be redirected to the redirect_uri from params.redirect_uri.
+
+        +--------+     +------------+     +-------------------+
+        |        |     |            |     |                   |
+        | Client | --> | MCP Server | --> | 3rd Party OAuth   |
+        |        |     |            |     | Server            |
+        +--------+     +------------+     +-------------------+
+                            |   ^                  |
+        +------------+      |   |                  |
+        |            |      |   |    Redirect      |
+        |redirect_uri|<-----+   +------------------+
+        |            |
+        +------------+          
+
+        Implementations will need to define another handler on the MCP server return
+        flow to perform the second redirect, and generates and stores an authorization
+        code as part of completing the OAuth authorization step.
 
         Implementations SHOULD generate an authorization code with at least 160 bits of
         entropy,
@@ -102,9 +126,9 @@ async def create_authorization_code(
         """
         ...
 
-    async def load_authorization_code_metadata(
+    async def load_authorization_code(
         self, client: OAuthClientInformationFull, authorization_code: str
-    ) -> AuthorizationCodeMeta | None:
+    ) -> AuthorizationCode | None:
         """
         Loads metadata for the authorization code challenge.
 
@@ -118,8 +142,8 @@ async def load_authorization_code_metadata(
         ...
 
     async def exchange_authorization_code(
-        self, client: OAuthClientInformationFull, authorization_code: str
-    ) -> OAuthTokens:
+        self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
+    ) -> TokenSuccessResponse:
         """
         Exchanges an authorization code for an access token.
 
@@ -137,7 +161,7 @@ async def exchange_refresh_token(
         client: OAuthClientInformationFull,
         refresh_token: str,
         scopes: Optional[List[str]] = None,
-    ) -> OAuthTokens:
+    ) -> TokenSuccessResponse:
         """
         Exchanges a refresh token for an access token.
 
@@ -178,3 +202,15 @@ async def revoke_token(
             request: The token revocation request.
         """
         ...
+
+def construct_redirect_uri(redirect_uri_base: str, authorization_code: AuthorizationCode, state: Optional[str]) -> str:
+    parsed_uri = urlparse(redirect_uri_base)
+    query_params = [(k, v) for k, vs in parse_qs(parsed_uri.query) for v in vs]
+    query_params.append(("code", authorization_code.code))
+    if state:
+        query_params.append(("state", state))
+
+    redirect_uri = urlunparse(
+        parsed_uri._replace(query=urlencode(query_params))
+    )
+    return redirect_uri

src/mcp/shared/auth.py

@@ -9,7 +9,7 @@
 from pydantic import AnyHttpUrl, BaseModel, Field
 
 
-class OAuthErrorResponse(BaseModel):
+class TokenErrorResponse(BaseModel):
     """
     See https://datatracker.ietf.org/doc/html/rfc6749#section-5.2
     """
@@ -19,7 +19,7 @@ class OAuthErrorResponse(BaseModel):
     error_uri: Optional[AnyHttpUrl] = None
 
 
-class OAuthTokens(BaseModel):
+class TokenSuccessResponse(BaseModel):
     """
     See https://datatracker.ietf.org/doc/html/rfc6749#section-5.1
     """