43 lines
1.9 KiB
Python
43 lines
1.9 KiB
Python
|
"""
|
|||
|
|
简单的 ASGI 中间件：为 SSE 入口添加 x-api-key 鉴权。
|
|||
|
|
当前 fastmcp 版本不支持 authenticate 参数时，可通过该中间件实现最小可用的 Header 鉴权。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from typing import Callable, Awaitable
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ApiKeyAuthMiddleware:
|
|||
|
|
def __init__(self, app, header_name: str = "x-api-key", expected_key: str | None = None, sse_path: str = "/sse"):
|
|||
|
|
self.app = app
|
|||
|
|
self.header_name = header_name.lower()
|
|||
|
|
self.expected_key = expected_key or os.environ.get("MCP_API_KEY", "123")
|
|||
|
|
self.sse_path = sse_path
|
|||
|
|
|
|||
|
|
async def __call__(self, scope, receive, send):
|
|||
|
|
# 仅对 HTTP scope 且 SSE 路径做校验，其余直接放行
|
|||
|
|
if scope.get("type") == "http":
|
|||
|
|
path = scope.get("path", "")
|
|||
|
|
if path.startswith(self.sse_path):
|
|||
|
|
# 提取 header（bytes）
|
|||
|
|
headers = {k.decode("latin1").lower(): v.decode("latin1") for k, v in scope.get("headers", [])}
|
|||
|
|
api_key = headers.get(self.header_name)
|
|||
|
|
if api_key != self.expected_key:
|
|||
|
|
body = b"Unauthorized: missing or invalid x-api-key"
|
|||
|
|
await send({
|
|||
|
|
"type": "http.response.start",
|
|||
|
|
"status": 401,
|
|||
|
|
"headers": [(b"content-type", b"text/plain; charset=utf-8")],
|
|||
|
|
})
|
|||
|
|
await send({
|
|||
|
|
"type": "http.response.body",
|
|||
|
|
"body": body,
|
|||
|
|
"more_body": False,
|
|||
|
|
})
|
|||
|
|
return
|
|||
|
|
return await self.app(scope, receive, send)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def with_api_key_auth(app, header_name: str = "x-api-key", expected_key: str | None = None, sse_path: str = "/sse"):
|
|||
|
|
return ApiKeyAuthMiddleware(app, header_name=header_name, expected_key=expected_key, sse_path=sse_path)
|