为 MCP Server 集成 OAuth 2.0 认证：从概念到实战

from pydantic import AnyHttpUrl
from pydantic_settings import BaseSettings, SettingsConfigDict

class ServerSettings(BaseSettings):
    model_config = SettingsConfigDict(env_prefix="MCP_")

    # MCP Server 配置
    host: str = "localhost"
    port: int = 3000
    server_url: AnyHttpUrl = AnyHttpUrl("http://localhost:3000")
    callback_path: str = "http://localhost:3000/callback"

    # Google OAuth 配置
    client_id: str      # 从环境变量 MCP_CLIENT_ID 读取
    client_secret: str  # 从环境变量 MCP_CLIENT_SECRET 读取
    auth_url: str = "https://accounts.google.com/o/oauth2/auth"
    token_url: str = "https://oauth2.googleapis.com/token"
    scope: str = "https://www.googleapis.com/auth/userinfo.email https://www.googleapis.com/auth/userinfo.profile openid"

settings = ServerSettings()

from typing import Dict, Optional
# ... 其他 mcp_server 导入 ...

class GoogleOAuthProvider(OAuthAuthorizationServerProvider):
    def __init__(self, settings: ServerSettings):
        self.settings = settings
        # 用于存储动态注册的客户端信息
        self.clients: Dict[str, OAuthClientInformationFull] = {}
        # 用于存储生成的 MCP 内部授权码
        self.auth_codes: Dict[str, AuthorizationCode] = {}
        # 用于存储从 Google 获取的令牌信息
        self.tokens: Dict[str, Dict] = {}
        # 用于存储 state，防止 CSRF 攻击
        self.state_mapping: Dict[str, Dict[str, str]] = {}

async def register_client(self, client_info: OAuthClientInformationFull):
    # 将 SDK 自动生成的客户端信息保存起来
    self.clients[client_info.client_id] = client_info

async def get_client(self, client_id: str) -> Optional[OAuthClientInformationFull]:
    # 根据 client_id 获取客户端信息
    return self.clients.get(client_id)

import secrets
import urllib.parse

async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str:
    # 生成 state 参数用于后续验证
    state = secrets.token_hex(16)
    self.state_mapping[state] = {
        "redirect_uri": str(params.redirect_uri),
        "client_id": client.client_id,
        "code_challenge": getattr(params, 'code_challenge', None),
    }

    # 构建 Google 的授权 URL
    auth_params = {
        "client_id": self.settings.client_id,
        "redirect_uri": self.settings.callback_path,
        "response_type": "code",
        "scope": self.settings.scope,
        "state": state,
        "access_type": "online",
    }
    auth_url = f"{self.settings.auth_url}?{urllib.parse.urlencode(auth_params)}"
    return auth_url

import time
from fastapi import HTTPException

async def handle_callback(self, code: str, state: str) -> str:
    state_data = self.state_mapping.get(state)
    if not state_data:
        raise HTTPException(status_code=400, detail="Invalid state parameter")

    # 使用从 Google 获取的 code 交换 Google 的 access_token
    token_response = await create_mcp_http_client().post(
        self.settings.token_url,
        data={
            "client_id": self.settings.client_id,
            "client_secret": self.settings.client_secret,
            "code": code,
            "redirect_uri": self.settings.callback_path,
            "grant_type": "authorization_code",
        },
        headers={"Content-Type": "application/x-www-form-urlencoded"},
    )
    token_response.raise_for_status()
    google_data = token_response.json()

    # 生成 MCP 内部授权码
    mcp_code = f"mcp_{secrets.token_hex(16)}"
    self.auth_codes[mcp_code] = AuthorizationCode(
        code=mcp_code,
        client_id=state_data["client_id"],
        redirect_uri=state_data["redirect_uri"],
        # ... 其他参数 ...
    )

    # 存储 Google 令牌信息，用于后续交换
    self.tokens[mcp_code] = google_data
    del self.state_mapping[state]

    # 构建并返回给 MCP 客户端的回调 URI，附上 MCP 内部授权码
    return construct_redirect_uri(state_data["redirect_uri"], code=mcp_code, state=state)

async def exchange_authorization_code(
    self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
) -> AccessToken:
    # 取出之前存储的 Google 令牌信息
    google_data = self.tokens.pop(authorization_code.code)
    
    # 注意：这里为了简化，直接使用了 Google 的 access_token。
    # 严格实现中，应在此处生成 MCP 服务器自己的 JWT 令牌。
    access_token = google_data["access_token"]
    
    # 存储令牌信息，用于后续验证
    self.tokens[access_token] = {
        "client_id": client.client_id,
        "expires_at": time.time() + google_data.get("expires_in", 3600),
    }
    
    return AccessToken(
        token=access_token,
        # ... 其他参数 ...
    )

async def load_access_token(self, token: str) -> Optional[AccessToken]:
    token_data = self.tokens.get(token)
    if not token_data:
        return None
    
    # 检查令牌是否过期
    if token_data.get("expires_at", 0) < time.time():
        del self.tokens[token]
        return None
        
    return AccessToken(
        token=token,
        client_id=token_data["client_id"],
        scopes=self.settings.scope.split(),
        expires_at=int(token_data.get("expires_at", 0)),
    )

from mcp_server.fast_mcp import FastMCP
from fastapi import Request, Response
from fastapi.responses import RedirectResponse

# 实例化 Provider 和 AuthSettings
oauth_provider = GoogleOAuthProvider(settings)
auth_settings = AuthSettings(
    issuer_url=settings.server_url,
    client_registration_options=ClientRegistrationOptions(enabled=True),
    required_scopes=["openid"],
)

# 创建 FastMCP 应用
app = FastMCP(
    name="Google OAuth MCP Server",
    auth_server_provider=oauth_provider,
    host=settings.host,
    port=settings.port,
    auth=auth_settings,
)

# 添加处理 Google 回调的路由
@app.custom_route("/callback", methods=["GET"])
async def callback_handler(request: Request) -> Response:
    code = request.query_params.get("code")
    state = request.query_params.get("state")
    if not code or not state:
        raise HTTPException(status_code=400, detail="Missing code or state")
    
    redirect_uri = await oauth_provider.handle_callback(code, state)
    return RedirectResponse(url=redirect_uri)

# 启动服务
if __name__ == "__main__":
    app.run()