feat: support ratelimit and audit middleware

Author: Coolgiserz

.gitignore

@@ -16,4 +16,5 @@ wheels/
 # Logs
 *.log
 *.out
-.pytest_cache/
+.pytest_cache/
+volumes

src/news_mcp_server/app.py

@@ -4,14 +4,20 @@
 from .mcp_server import mcp_app
 from .middlewares.auth import SimpleAuthMiddleware
 from .middlewares.monitor import MonitorMiddleware, metrics
-from .middlewares.session import SessionMiddleware
+from .middlewares.rate_limit import RedisRateLimitMiddleware
 from .config.settings import app_settings
-middlewares = [Middleware(SessionMiddleware,
-                          secret_key=app_settings.SESSION_SECREY_KEY,
-                          max_age=3600),
-                Middleware(MonitorMiddleware),
-               Middleware(SimpleAuthMiddleware)
-               ]
+
+middlewares = [
+    # IP 速率限制
+    Middleware(RedisRateLimitMiddleware,
+               redis_url=app_settings.REDIS_URL,
+               max_requests=app_settings.RATE_LIMIT_MAX,
+               window_seconds=app_settings.RATE_LIMIT_WINDOW),
+    # 监控中间件
+    Middleware(MonitorMiddleware),
+    # 简单认证
+    Middleware(SimpleAuthMiddleware),
+]
 
 allow_origins = [
     "http://localhost:8000",

src/news_mcp_server/config/settings.py

@@ -10,7 +10,11 @@ class ApplicationSettings(BaseModel):
     CORS_HEADERS: list = ["*"]
     CORS_ALLOW_CREDENTIALS: bool = True
     API_KEY: str | None = os.getenv("NEWS_MCP_API_KEY")
-    SESSION_SECREY_KEY: str = os.getenv("SESSION_SECREY_KEY")
+    SESSION_SECRET_KEY: str = os.getenv("SESSION_SECRET_KEY")
+    REDIS_URL: str = os.getenv("REDIS_URL", "redis://redis:6379/0")
+    # IP 限流配置
+    RATE_LIMIT_MAX: int = int(os.getenv("RATE_LIMIT_MAX", 100))  # 单个 IP 在时间窗口内最大请求数
+    RATE_LIMIT_WINDOW: int = int(os.getenv("RATE_LIMIT_WINDOW", 60))  # 限流窗口时长（秒）
     TRANSPORT: str = "streamable-http"
 
 

src/news_mcp_server/mcp_server.py

@@ -1,9 +1,11 @@
 from typing import List
 from fastmcp import FastMCP, Context
+from fastmcp.server.http import Middleware
 from pydantic import Field
 import contextlib
 from .services.news_service import NewsService
 from .clients.elastic_client import AsyncElasticClient
+from .middlewares.audit import AuditMiddleware
 import structlog
 logger = structlog.get_logger(__name__)
 
@@ -13,7 +15,10 @@ class NewsMCP(FastMCP):
     pass
 
 def create_http_app(mcp):
-    mcp_app = mcp.http_app("/es-news-mcp")
+    middlewares = [
+        Middleware(AuditMiddleware)
+    ]
+    mcp_app = mcp.http_app("/es-news-mcp", middleware=middlewares)
     return mcp_app
 app_services = {}
 

src/news_mcp_server/middlewares/audit.py

@@ -0,0 +1,35 @@
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+import time
+from ..utils.logger import logger
+
+class AuditMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next) -> Response:
+        method = None
+        params = None
+        # 尝试解析 JSON-RPC 请求体中的方法名和参数
+        if request.method.upper() == "POST":
+            try:
+                body = await request.json()
+                method = body.get("method")
+                params = body.get("params")
+            except Exception:
+                pass
+        # 获取客户端 IP
+        client_ip = None
+        if request.client:
+            client_ip = request.client.host
+        start_time = time.time()
+        response = await call_next(request)
+        duration = time.time() - start_time
+        # 审计记录：工具名、参数、客户端IP、状态码、耗时(ms)
+        logger.info(
+            "mcp_tool_audit",
+            method=method,
+            params=params,
+            client_ip=client_ip,
+            status_code=response.status_code,
+            duration_ms=int(duration * 1000)
+        )
+        return response

src/news_mcp_server/middlewares/auth.py

@@ -1,10 +1,10 @@
 # @Author: Zhu Guowei
 # @Date: 2025/6/17
 # @Function:
-from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.requests import Request
 from starlette.responses import JSONResponse
+from starlette import status
 from structlog import get_logger
 logger = get_logger(__name__)
 import os
@@ -20,24 +20,26 @@ async def dispatch(self, request: Request, call_next):
         # 如果未配置 API_KEY，且允许跳过，则直接放行（便于开发环境）
 
         host = request.headers.get("HOST")
-        logger.info(f"host: {host}, {request.headers}")
+        headers = request.headers
+        auth_header = headers.get("authorization")
+
         for h in ALLOW_HOSTS:
             if h in host:
                 return await call_next(request)
         if not self.api_key:
-            logger.info(f"request.url: {request.url}")
             return await call_next(request)
 
         # 获取 Authorization头
-        auth_header = request.headers.get("authorization")
-        logger.info(f"url: {request.url}, query_params: {request.query_params}")
-        logger.info(f"auth-header: {auth_header}")
+        logger.info(f"simple-auth", host=host, header=auth_header)
+
         if not auth_header or not auth_header.lower().startswith("bearer "):
-            return JSONResponse({"detail": "Bearer Token Not Provided"}, status_code=401)
+            logger.warning(f"simple-auth", host=host, detail="Bearer Token Not Provided", header=auth_header)
+            return JSONResponse({"detail": "Bearer Token Not Provided"}, status_code=status.HTTP_401_UNAUTHORIZED)
 
         token = auth_header[7:].strip()
         if token != self.api_key:
-            return JSONResponse({"detail": "Invalid Token"}, status_code=403)
+            logger.warning(f"simple-auth", host=host, detail="Invalid Token", token=token, header=auth_header)
+            return JSONResponse({"detail": "Invalid Token"}, status_code=status.HTTP_403_FORBIDDEN)
 
         # 认证通过，继续处理请求
         return await call_next(request)

src/news_mcp_server/middlewares/rate_limit.py

@@ -0,0 +1,55 @@
+import time
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette import status
+from redis import asyncio as aioredis
+from ..utils.logger import logger
+
+
+class RedisRateLimitMiddleware(BaseHTTPMiddleware):
+    """
+    基于 Redis 的简单 IP 限流中间件。
+    每个 IP 在固定时间窗口内最多允许 max_requests 次请求。
+    """
+    def __init__(self, app, redis_url: str, max_requests: int, window_seconds: int):
+        super().__init__(app)
+        self.redis_url = redis_url
+        self.max_requests = max_requests
+        self.window = window_seconds
+        self._redis = None
+
+    async def _get_redis(self):
+        if self._redis is None:
+            # 仅在首次调用时创建 Redis 连接
+            self._redis = await aioredis.from_url(
+                self.redis_url, encoding="utf-8", decode_responses=True
+            )
+        return self._redis
+
+    async def dispatch(self, request: Request, call_next):
+        # 获取客户端 IP
+        client_host = request.client.host if request.client else "unknown"
+        logger.debug("rate-limiter", host=client_host)
+        # 计算当前时间窗口
+        now = int(time.time())
+        window_key = now // self.window
+        key = f"ratelimit:{client_host}:{window_key}"
+        redis = await self._get_redis()
+        # 自增计数
+        count = await redis.incr(key)
+        if count == 1:
+            # 设置过期时间为一个窗口长度
+            await redis.expire(key, self.window)
+
+        # 超出限流阈值，返回 429
+        if count > self.max_requests:
+            logger.info("rate-limiter", host=client_host, key=key)
+            return JSONResponse(
+                {"detail": "请求过多，请稍后重试"},
+                status_code=status.HTTP_429_TOO_MANY_REQUESTS
+            )
+
+        # 继续处理请求
+        response = await call_next(request)
+        return response 

src/news_mcp_server/middlewares/redis_session.py

@@ -0,0 +1,68 @@
+import json
+import uuid
+from itsdangerous import TimestampSigner, BadSignature
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+from redis import asyncio as aioredis
+
+
+class RedisSessionMiddleware(BaseHTTPMiddleware):
+    """
+    TODO 基于 Redis 的服务端 Session 中间件。
+    - session 使用 UUID 作为 key 存储在 Redis 中
+    - cookie 存储签名后的 session_id
+    """
+    def __init__(self, app, secret_key: str, redis_url: str, cookie_name: str = "session", max_age: int = 14*24*60*60):
+        super().__init__(app)
+        if not secret_key:
+            raise ValueError("SESSION_SECRET_KEY 未配置")
+        self.signer = TimestampSigner(secret_key)
+        self.redis_url = redis_url
+        self.cookie_name = cookie_name
+        self.max_age = max_age
+        self._redis = None
+
+    async def _get_redis(self):
+        if self._redis is None:
+            self._redis = await aioredis.from_url(self.redis_url, encoding="utf-8", decode_responses=True)
+        return self._redis
+
+    async def dispatch(self, request: Request, call_next):
+        # 获取或生成 session_id
+        session_id = None
+        cookie = request.cookies.get(self.cookie_name)
+        if cookie:
+            try:
+                unsigned = self.signer.unsign(cookie, max_age=self.max_age)
+                session_id = unsigned.decode()
+            except BadSignature:
+                session_id = None
+        if not session_id:
+            session_id = str(uuid.uuid4())
+
+        # 取 Redis 中的数据
+        redis = await self._get_redis()
+        raw = await redis.get(f"session:{session_id}")
+        try:
+            session_data = json.loads(raw) if raw else {}
+        except json.JSONDecodeError:
+            session_data = {}
+        request.scope["session"] = session_data
+
+        # 调用下游
+        response: Response = await call_next(request)
+
+        # 写回 Redis
+        await redis.setex(f"session:{session_id}", self.max_age, json.dumps(request.scope.get('session', {})))
+
+        # 设置 cookie
+        signed = self.signer.sign(session_id.encode()).decode()
+        response.set_cookie(
+            self.cookie_name,
+            signed,
+            max_age=self.max_age,
+            httponly=True,
+            samesite="lax"
+        )
+        return response 

src/news_mcp_server/utils/logger.py

@@ -1,2 +1,30 @@
+import os
+import logging
+from pathlib import Path
 from structlog import get_logger
-logger = get_logger("suwen-news-mcp-server")
+from logging.handlers import TimedRotatingFileHandler
+BASE_DIR = Path(__file__).parent.parent.parent
+
+
+logger = get_logger("suwen-news-mcp-server")
+# === 文件日志处理器 ===
+LOG_DIR = os.getenv("LOG_DIR", os.path.join(BASE_DIR, "logs"))
+os.makedirs(LOG_DIR, exist_ok=True)
+LOG_FILE = os.path.join(LOG_DIR, "es_news_mcp_server.log")
+
+# === 标准库日志根配置 ===
+root_logger = logging.getLogger()
+root_logger.setLevel(logging.INFO)
+
+# 控制台 Handler（保留 JSON 格式）
+console_handler = logging.StreamHandler()
+console_handler.setFormatter(logging.Formatter("%(message)s"))
+
+# 文件 Handler：每日 0 点轮转，保留 7 天
+file_handler = TimedRotatingFileHandler(LOG_FILE, when="midnight", backupCount=7, encoding="utf-8")
+file_handler.setFormatter(logging.Formatter("%(message)s"))
+
+# 仅在首次配置时添加（防止重复）
+if not root_logger.handlers:
+    root_logger.addHandler(console_handler)
+    root_logger.addHandler(file_handler)