fix: harden MCP endpoint after code review
Critical: the MCP SDK auto-enables DNS-rebinding protection when its host parameter is left at the 127.0.0.1 default, hard-coding a loopback-only Host allow-list. Every remote client (the only deployment this feature targets) was refused with HTTP 421 before auth or the tool ran. Now driven by a new MCP_ALLOWED_HOSTS setting, with '*' as an explicit, logged opt-out. Also bounds query/top_k to match AskRequest (top_k is amplified 4x downstream, so an unbounded value was a resource-exhaustion vector), decodes the Authorization header as latin-1 per the ASGI spec instead of raising a 500 on malformed bytes, and returns WWW-Authenticate on 401 per RFC 7235. Moves the psycopg2 import guard into backend/tests/conftest.py: duplicated across four test modules, it only worked because of alphabetical collection order, and any earlier-sorting package would have reintroduced a live connection attempt against the production database. Registers the mcp module in the authoritative backend architecture doc. 84 backend tests pass. Verified against a live server: allowed remote Host returns a valid initialize result, unknown Host returns 421, missing token returns 401 with WWW-Authenticate. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
@@ -8,13 +8,20 @@ dict. No new retrieval, ranking, or LLM orchestration logic lives here.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Annotated
|
||||
|
||||
from mcp.server import MCPServer
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from pydantic import Field
|
||||
from starlette.responses import PlainTextResponse
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
|
||||
from app.config.settings import settings
|
||||
from app.shared.bootstrap import get_agent_conversation_service, get_jwt_handler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Single shared MCPServer instance — analogous to the single shared FastAPI
|
||||
# `app` instance in app/api/main.py. Tools registered via @mcp.tool() below.
|
||||
# Note: the installed mcp SDK (2.0.0) renamed the older "FastMCP" class to
|
||||
@@ -24,12 +31,22 @@ mcp = MCPServer("ai-regulations")
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def search_regulations(query: str, top_k: int = 5) -> dict:
|
||||
def search_regulations(
|
||||
query: Annotated[str, Field(min_length=1, max_length=2000)],
|
||||
top_k: Annotated[int, Field(ge=1, le=20)] = 5,
|
||||
) -> dict:
|
||||
"""Search the compliance knowledge base and return a grounded answer.
|
||||
|
||||
query: Natural-language search question, e.g. "国六排放标准最新要求".
|
||||
top_k: Maximum number of cited sources to return (default 5).
|
||||
top_k: Maximum number of cited sources to return (1-20, default 5).
|
||||
"""
|
||||
# Bounds mirror AskRequest in app/api/models/agent.py so the MCP path cannot
|
||||
# be used to bypass the REST endpoint's limits. They matter more here than
|
||||
# there: KnowledgeRetrievalService amplifies top_k (candidate_k = top_k * 4)
|
||||
# when reranking, and an LLM client can easily hallucinate a huge value.
|
||||
# Declaring them via Annotated puts them in the advertised JSON schema too,
|
||||
# so well-behaved clients never send an out-of-range value in the first place.
|
||||
#
|
||||
# No session_id is passed: this keeps each call stateless (no
|
||||
# ConversationStore reads/writes), matching "search" semantics rather
|
||||
# than multi-turn chat semantics.
|
||||
@@ -61,19 +78,53 @@ class MCPAuthMiddleware:
|
||||
return
|
||||
|
||||
headers = dict(scope["headers"])
|
||||
auth_header = headers.get(b"authorization", b"").decode()
|
||||
# ASGI header values are raw bytes specified as latin-1, not UTF-8;
|
||||
# decoding strictly as UTF-8 would raise on a malformed byte and turn a
|
||||
# bad request into an unhandled 500.
|
||||
auth_header = headers.get(b"authorization", b"").decode("latin-1")
|
||||
token = auth_header.removeprefix("Bearer ").strip()
|
||||
try:
|
||||
get_jwt_handler().decode_token(token)
|
||||
except ValueError as exc:
|
||||
# Reject before the MCP session/protocol layer ever sees the request.
|
||||
response = PlainTextResponse(str(exc), status_code=401)
|
||||
# WWW-Authenticate matches the get_current_user dependency (auth.py)
|
||||
# and is required by RFC 7235 so clients can tell "needs credentials"
|
||||
# apart from a generic failure.
|
||||
response = PlainTextResponse(
|
||||
str(exc), status_code=401, headers={"WWW-Authenticate": "Bearer"}
|
||||
)
|
||||
await response(scope, receive, send)
|
||||
return
|
||||
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
|
||||
def _build_transport_security() -> TransportSecuritySettings:
|
||||
"""Translate the configured MCP host allow-list into SDK transport settings.
|
||||
|
||||
Without this the SDK infers its own allow-list from the bind host, which
|
||||
defaults to 127.0.0.1 and therefore rejects every remote client with HTTP
|
||||
421 — fatal for a remotely deployed backend.
|
||||
"""
|
||||
allowed = [h.strip() for h in settings.mcp_allowed_hosts.split(",") if h.strip()]
|
||||
if "*" in allowed:
|
||||
# Explicit, logged opt-out. Kept as an escape hatch for environments
|
||||
# behind a proxy that rewrites Host unpredictably, but never the default.
|
||||
logger.warning(
|
||||
"MCP DNS-rebinding protection is disabled (mcp_allowed_hosts='*'). "
|
||||
"Set MCP_ALLOWED_HOSTS to the real deployment host(s) instead."
|
||||
)
|
||||
return TransportSecuritySettings(enable_dns_rebinding_protection=False)
|
||||
return TransportSecuritySettings(
|
||||
enable_dns_rebinding_protection=True,
|
||||
allowed_hosts=allowed,
|
||||
# Browser clients send Origin; reuse the already-maintained CORS list so
|
||||
# there is one place to declare trusted web origins. Non-browser MCP
|
||||
# clients send no Origin at all, which the SDK treats as allowed.
|
||||
allowed_origins=[o.strip() for o in settings.cors_allow_origins.split(",") if o.strip()],
|
||||
)
|
||||
|
||||
|
||||
def build_mcp_asgi_app() -> ASGIApp:
|
||||
"""Return the Streamable HTTP ASGI app for the MCP server, auth-guarded.
|
||||
|
||||
@@ -83,6 +134,9 @@ def build_mcp_asgi_app() -> ASGIApp:
|
||||
path to "/", the effective external path would be the confusing "/mcp/mcp"
|
||||
instead of "/mcp".
|
||||
"""
|
||||
asgi_app = mcp.streamable_http_app(streamable_http_path="/")
|
||||
asgi_app = mcp.streamable_http_app(
|
||||
streamable_http_path="/",
|
||||
transport_security=_build_transport_security(),
|
||||
)
|
||||
asgi_app.add_middleware(MCPAuthMiddleware)
|
||||
return asgi_app
|
||||
|
||||
Reference in New Issue
Block a user