fix: enforce streamed request body limits

Author: jiwenxu025-boop

scripts/bazi_engine/api.py

@@ -53,6 +53,7 @@
 }
 _MAX_AI_STREAMS = int(os.getenv("BAZI_MAX_AI_STREAMS", "3"))
 _AI_STREAM_SLOTS = threading.BoundedSemaphore(max(1, _MAX_AI_STREAMS))
+_ALLOW_LEGACY_CHART_GET = os.getenv("BAZI_ALLOW_LEGACY_CHART_GET", "1").lower() in ("1", "true", "yes")
 _FUSION_STREAM_TOTAL_TIMEOUT = float(os.getenv("BAZI_FUSION_STREAM_TOTAL_TIMEOUT", "150"))
 _FUSION_STREAM_IDLE_TIMEOUT = float(os.getenv("BAZI_FUSION_STREAM_IDLE_TIMEOUT", "45"))
 _STREAM_HEARTBEAT_INTERVAL = float(os.getenv("BAZI_STREAM_HEARTBEAT_INTERVAL", "15"))
@@ -63,6 +64,10 @@
 ]
 
 
+class RequestBodyTooLargeError(Exception):
+    """Raised while consuming a chunked request body over the configured limit."""
+
+
 class ChartStreamRequest(BaseModel):
     name: str = Field(default="", max_length=40)
     gender: Literal["男", "女"]
@@ -114,10 +119,31 @@ async def reject_large_request_bodies(request: Request, call_next):
     if request.method in {"POST", "PUT", "PATCH"}:
         content_length = request.headers.get("content-length")
         max_bytes = _LARGE_REQUEST_BODY_LIMITS.get(request.url.path, _MAX_REQUEST_BODY_BYTES)
-        if content_length and int(content_length) > max_bytes:
-            response = JSONResponse({"error": "请求体过大"}, status_code=413)
+        if content_length:
+            try:
+                if int(content_length) > max_bytes:
+                    response = JSONResponse({"error": "请求体过大"}, status_code=413)
+            except ValueError:
+                response = JSONResponse({"error": "无效的请求体长度"}, status_code=400)
+        if response is None:
+            receive = request._receive
+            received_bytes = 0
+
+            async def limited_receive():
+                nonlocal received_bytes
+                message = await receive()
+                if message["type"] == "http.request":
+                    received_bytes += len(message.get("body", b""))
+                    if received_bytes > max_bytes:
+                        raise RequestBodyTooLargeError
+                return message
+
+            request._receive = limited_receive
     if response is None:
-        response = await call_next(request)
+        try:
+            response = await call_next(request)
+        except RequestBodyTooLargeError:
+            response = JSONResponse({"error": "请求体过大"}, status_code=413)
     response.headers.setdefault("X-Content-Type-Options", "nosniff")
     response.headers.setdefault("X-Frame-Options", "DENY")
     response.headers.setdefault("Referrer-Policy", "no-referrer")
@@ -176,6 +202,8 @@ def chart_api(
     hour_confirmed: bool = Query(False, description="出生时辰是否经用户确认"),
     practical: bool = Query(False, description="实用模式：仅返回白话解读，不包含技术推导"),
 ):
+    if not _ALLOW_LEGACY_CHART_GET:
+        return JSONResponse({"error": "该入口已停用，请使用 POST /api/chart/stream"}, status_code=410)
     ln_range = None
     if liunian_from and liunian_to:
         ln_range = (liunian_from, liunian_to)
@@ -527,6 +555,9 @@ async def chart_stream(
       data: {"phase":"done"}                             — 全流程结束
     """
 
+    if not _ALLOW_LEGACY_CHART_GET:
+        return JSONResponse({"error": "该入口已停用，请使用 POST /api/chart/stream"}, status_code=410)
+    logger.info("legacy chart stream GET used")
     payload = ChartStreamRequest(
         name=name,
         gender=gender,

scripts/tests/test_api.py

@@ -12,6 +12,7 @@
 
 import pytest
 from fastapi.testclient import TestClient
+from starlette.requests import Request
 
 import bazi_engine._chart_context as chart_context
 from tests._chat_fixtures import chart_data_with_current_dayun
@@ -55,6 +56,28 @@ def test_api_sets_baseline_security_headers():
     assert response.headers["referrer-policy"] == "no-referrer"
 
 
+def test_request_body_limit_applies_to_chunked_body(monkeypatch):
+    import bazi_engine.api as api_module
+
+    monkeypatch.setattr(api_module, "_MAX_REQUEST_BODY_BYTES", 4)
+    chunks = iter((
+        {"type": "http.request", "body": b"abc", "more_body": True},
+        {"type": "http.request", "body": b"de", "more_body": False},
+    ))
+
+    async def receive():
+        return next(chunks)
+
+    async def call_next(request):
+        await request.body()
+        return api_module.JSONResponse({"ok": True})
+
+    request = Request({"type": "http", "method": "POST", "path": "/api/feedback", "headers": []}, receive)
+    response = asyncio.run(api_module.reject_large_request_bodies(request, call_next))
+
+    assert response.status_code == 413
+
+
 def test_chart_stream_returns_rules_and_done_events(monkeypatch):
     monkeypatch.setenv("BAZI_LLM_REVIEW", "0")
     monkeypatch.setenv("BAZI_AI_ENABLED", "0")
@@ -83,6 +106,16 @@ def test_chart_stream_returns_rules_and_done_events(monkeypatch):
     assert "data: [DONE]" in body
 
 
+def test_legacy_chart_get_can_be_disabled(monkeypatch):
+    import bazi_engine.api as api_module
+    from bazi_engine.api import app
+
+    monkeypatch.setattr(api_module, "_ALLOW_LEGACY_CHART_GET", False)
+    params = {"gender": "男", "year": 2007, "month": 8, "day": 26, "hour": 20}
+    assert TestClient(app).get("/api/chart", params=params).status_code == 410
+    assert TestClient(app).get("/api/chart/stream", params=params).status_code == 410
+
+
 def test_chart_stream_post_returns_rules_and_done_events(monkeypatch):
     monkeypatch.setenv("BAZI_LLM_REVIEW", "0")
     monkeypatch.setenv("BAZI_AI_ENABLED", "0")