mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
fit2cloud-chenywandClaude
bbfc4cc8f5
- TokenMiddleware/ResponseMiddleware/audit logger use scope['path'] instead of request.url.path for whitelist checks: a crafted Host header (e.g. x/api/v1/mcp) was concatenated into Starlette's URL and made protected routes match the /mcp* whitelist entry - validateEmbedded rejects admin accounts and non-type-4 apps, so a self-signed token with account=admin can no longer impersonate the administrator - Add HostValidationMiddleware rejecting malformed Host headers at the outermost layer (defense in depth) - Add regression tests for host validation, whitelist matching and source-level guards Co-Authored-By: Claude <[email protected]>
118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
import json
|
|
from typing import Optional
|
|
|
|
from starlette.exceptions import HTTPException
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.requests import Request
|
|
from starlette.responses import JSONResponse
|
|
|
|
from common.core.config import settings
|
|
from common.utils.utils import SQLBotLogUtil
|
|
|
|
|
|
class ResponseMiddleware(BaseHTTPMiddleware):
|
|
instances = []
|
|
|
|
def __init__(self, app, allow_origins: Optional[list[str]] = None):
|
|
super().__init__(app)
|
|
self.allow_origins = allow_origins or ["'self'"]
|
|
ResponseMiddleware.instances.append(self)
|
|
|
|
def update_allow_origins(self, new_allow_origins: Optional[list[str]] = None):
|
|
if not new_allow_origins:
|
|
return
|
|
self.allow_origins = list(set(self.allow_origins + new_allow_origins))
|
|
|
|
|
|
async def dispatch(self, request, call_next):
|
|
response = await call_next(request)
|
|
|
|
direct_paths = [
|
|
f"{settings.API_V1_STR}/mcp/mcp_question",
|
|
f"{settings.API_V1_STR}/mcp/mcp_assistant",
|
|
f"{settings.CONTEXT_PATH}/openapi.json",
|
|
f"{settings.CONTEXT_PATH}/docs",
|
|
f"{settings.CONTEXT_PATH}/redoc"
|
|
]
|
|
|
|
route = request.scope.get("route")
|
|
# 获取定义的路径模式，例如 '/items/{item_id}'
|
|
path_pattern = '' if not route else route.path_format
|
|
|
|
if (isinstance(response, JSONResponse)
|
|
or (request.scope.get("path") or request.url.path) == f"{settings.CONTEXT_PATH}/openapi.json"
|
|
or path_pattern in direct_paths):
|
|
return response
|
|
if response.status_code != 200:
|
|
return response
|
|
if response.headers.get("content-type") == "application/json":
|
|
try:
|
|
body = b""
|
|
async for chunk in response.body_iterator:
|
|
body += chunk
|
|
|
|
raw_data = json.loads(body.decode())
|
|
|
|
if isinstance(raw_data, dict) and all(k in raw_data for k in ["code", "data", "msg"]):
|
|
return JSONResponse(
|
|
content=raw_data,
|
|
status_code=response.status_code,
|
|
headers={
|
|
k: v for k, v in response.headers.items()
|
|
if k.lower() not in ("content-length", "content-type")
|
|
}
|
|
)
|
|
|
|
wrapped_data = {
|
|
"code": 0,
|
|
"data": raw_data,
|
|
"msg": None
|
|
}
|
|
|
|
return JSONResponse(
|
|
content=wrapped_data,
|
|
status_code=response.status_code,
|
|
headers={
|
|
k: v for k, v in response.headers.items()
|
|
if k.lower() not in ("content-length", "content-type")
|
|
}
|
|
)
|
|
except Exception as e:
|
|
SQLBotLogUtil.error(f"Response processing error: {str(e)}", exc_info=True)
|
|
return JSONResponse(
|
|
status_code=500,
|
|
content=str(e),
|
|
headers={
|
|
k: v for k, v in response.headers.items()
|
|
if k.lower() not in ("content-length", "content-type")
|
|
}
|
|
)
|
|
content_type = response.headers.get("content-type", "")
|
|
static_content_types = ["text/html", "javascript", "typescript", "css"]
|
|
if any(ct in content_type for ct in static_content_types):
|
|
if self.allow_origins:
|
|
frame_ancestors_value = " ".join(self.allow_origins)
|
|
response.headers["Content-Security-Policy"] = f"frame-ancestors {frame_ancestors_value};"
|
|
|
|
return response
|
|
|
|
|
|
class exception_handler():
|
|
@staticmethod
|
|
async def http_exception_handler(request: Request, exc: HTTPException):
|
|
SQLBotLogUtil.error(f"HTTP Exception: {exc.detail}", exc_info=True)
|
|
return JSONResponse(
|
|
status_code=exc.status_code,
|
|
content=exc.detail,
|
|
headers={"Access-Control-Allow-Origin": "*"}
|
|
)
|
|
|
|
@staticmethod
|
|
async def global_exception_handler(request: Request, exc: Exception):
|
|
SQLBotLogUtil.error(f"Unhandled Exception: {str(exc)}", exc_info=True)
|
|
return JSONResponse(
|
|
status_code=500,
|
|
content=str(exc),
|
|
headers={"Access-Control-Allow-Origin": "*"}
|
|
)
|