feat: Host 中间件租户路由(纯ASGI覆盖WS/未知404/停用503暂停页/health豁免)

Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
浅念
2026-10-12 20:30:30 +08:00
parent 1824bb1590
commit eaa5bf7e51
5 changed files with 338 additions and 50 deletions
+9 -16
View File
@@ -1,13 +1,9 @@
"""路由取当前租户运行时——Task 5b→7 的过渡兜底。
"""路由取当前租户运行时——纯访问器(Task 7)。
Task 7 的 Host 中间件上线前,没有人为请求挂 ``state.tenant``。本助手的
兜底规则 = state 里有租户用 state 的(中间件灰度接入即生效);没有则用
``app.state.default_runtime``(lifespan 预热的默认租户 jiancai),API 测试
与直连(无 Host 语义)在 5b→7 过渡期不断。
已知过渡限制(记录在案,Task 7 一并解决):默认租户经 registry_reload
热切换后 default_runtime 仍指向旧 runtime——过渡期不修,Task 7 中间件
按请求 Host 重新 resolve 后本模块整体删除。
TenantHostMiddleware(纯 ASGI,http 与 websocket 双覆盖)保证凡到达路由
的请求 ``state.tenant`` 必已就位;唯一例外是平台级豁免路径(/health),
而它不访问租户资源。本模块不再有任何兜底——state 无租户即 raise,把
"中间件没挂上/被绕过"这类装配错误炸在第一现场而不是静默串库。
"""
from typing import TYPE_CHECKING
@@ -20,13 +16,10 @@ if TYPE_CHECKING:
def get_runtime(request_or_ws: Request | WebSocket) -> "TenantRuntime":
"""state.tenant 优先,缺省落 app.state.default_runtime(Task 7 删)。"""
"""读 state.tenant(中间件写入);缺失即装配错误,清晰失败。"""
runtime = getattr(request_or_ws.state, "tenant", None)
if runtime is not None:
return runtime
runtime = getattr(request_or_ws.app.state, "default_runtime", None)
if runtime is None or runtime.session_factory is None:
if runtime is None:
raise RuntimeError(
"默认租户运行时未就绪——lifespan 未运行、未预热 default_runtime,"
"或 registry_reload 已 dispose 旧 runtime(过渡期限制,Task 7 解决)")
"请求 state 里没有租户运行时——TenantHostMiddleware 未覆盖该"
"请求(中间件未挂载,或豁免路径访问了租户资源)")
return runtime
+9 -10
View File
@@ -6,6 +6,7 @@ from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from app.engine.config import STATIC_DIR
from app.platform.host_middleware import TenantHostMiddleware
from app.platform.registry import TenantRegistry
logger = logging.getLogger(__name__)
@@ -14,9 +15,6 @@ logger = logging.getLogger(__name__)
BASE_DIR = Path(__file__).resolve().parents[1]
SEED_DIR = BASE_DIR / "seed_profiles"
TENANTS_DIR = BASE_DIR / "tenants" # .gitignore 已忽略 tenants/*.db
# 过渡兜底(Task 5b→7,中间件上线后随 routes/runtime.py 一并删除):
# Host 中间件落地前请求 state 里没有租户,一律落默认租户 jiancai
DEFAULT_TENANT_HOST = "chat2.agent.hnxinyouhe.com"
@asynccontextmanager
@@ -29,15 +27,10 @@ async def lifespan(app: FastAPI):
app.state.registry = registry
try:
# 预热默认租户(过渡兜底读 app.state.default_runtime)。engine 建在
# 本 lifespan 的事件循环里——连接与 loop 绑定,别处另起 asyncio.run
# 复用同一 engine 会跨 loop 死锁(测试直连造数走 conftest.async_db)。
# 预热/加载保持 fail-fast:默认租户起不来 = 过渡期全站无租户。
app.state.default_runtime = await registry.resolve(DEFAULT_TENANT_HOST)
# 询盘推送(设计 8.2): 清扫崩溃遗留的 stale pending → failed。逐租户、
# 用该租户自己的 session_factory;db 文件不存在的租户跳过——全新库
# 无遗留行,保持 engine 懒建纪律(预热的默认租户文件已建,照常扫)。
# 无遗留行,保持 engine 懒建纪律(租户 engine 由 Host 中间件按请求
# 懒建,启动期只有已有库文件的租户会被扫描)。
# 清扫失败只记日志继续:一个租户的坏库不炸整个平台启动(其余租户
# 照常服务),该租户下轮重启还会重扫。
from app.engine.services.lead_push_service import sweep_stale_pending
@@ -80,6 +73,12 @@ async def static_no_cache(request, call_next):
return resp
# 租户路由中间件(Task 7): 最后 add = 最外层——租户解析先于一切用户中间件。
# 纯 ASGI 类(http+websocket 双覆盖),registry 经 scope["app"] 延迟取,
# 因此可在 lifespan 挂 Registry 之前于模块级注册。
app.add_middleware(TenantHostMiddleware)
@app.get("/")
async def index():
from fastapi.responses import FileResponse
+99
View File
@@ -0,0 +1,99 @@
"""TenantHostMiddleware——按 Host 头解析租户的纯 ASGI 中间件(Task 7)。
为什么必须是纯 ASGI 类而不是 ``@app.middleware("http")``(前序审查实证过
的坑): 那条路走 BaseHTTPMiddleware,只跑 http scope——WebSocket 请求会
整个绕过租户解析落进 routes/runtime.py 的兜底,租户 A 的 WS 流量静默写进
租户 B 的库(串库不崩溃,极难发现)。本类自己分发 http 与 websocket 两种
scope type,两条路都先过租户解析。
行为:
- scope 非 http/websocket(lifespan 等)原样放行;
- path 精确等于 ``/health`` 豁免(部署健康检查 ``curl 127.0.0.1:8018``
/health,Host 是 IP:端口必然未知;精确匹配,``/healthxxx`` 不豁免);
- Host 头 → ``registry.resolve_site``: 未知 → http 404 / ws close(1008);
租户停用 → http 503 中性"服务暂停"页(不暴露租户信息)/ ws close(1013);
- 解析成功把 tenant/site 写进 ``scope["state"]``——Request/WebSocket.state
与 scope["state"] 共享同一 dict(setdefault 保证),routes 的 get_runtime
即读到本中间件写入的 runtime。
Registry 延迟取: ``add_middleware`` 实例化时 lifespan 还没跑、
``app.state.registry`` 未挂,构造期拿不到——从 ``scope["app"].state``
逐请求取(Starlette 在中间件栈外层保证 scope 带 app)。
"""
from starlette.responses import HTMLResponse, PlainTextResponse
from starlette.types import ASGIApp, Receive, Scope, Send
# 归一复用 registry 的私有实现——同一包内单一事实源,中间件与域名索引
# 的建法永不漂移(resolve_site 内部也会再归一次,双保险无害)
from app.platform.registry import (
TenantDisabledError,
UnknownTenantError,
_normalize_host,
)
# 中性"服务暂停"页: 不暴露租户 slug/域名/停用原因(可能是欠费等敏感信息)
PAUSE_PAGE = (
"<!DOCTYPE html><html lang=\"zh-CN\"><head><meta charset=\"utf-8\">"
"<title>服务暂停</title></head>"
"<body style=\"font-family:system-ui,sans-serif;text-align:center;"
"padding:4rem 1rem;color:#333;\">"
"<p style=\"font-size:1.1rem;\">服务暂停维护中,请稍后再试。</p>"
"</body></html>"
)
def _host_from_scope(scope: Scope) -> str:
"""ASGI headers(小写 bytes 键值对)里的 host → str 并归一。"""
for key, value in scope.get("headers") or []:
if key == b"host":
return _normalize_host(value.decode("latin-1"))
return ""
async def _close_ws(receive: Receive, send: Send, code: int, reason: str) -> None:
"""以 close 代 accept 拒绝 WS(ASGI 规范序列: 先收到 websocket.connect,
再回 websocket.close)。reason 中性,不带租户信息。"""
event = await receive()
if event["type"] == "websocket.connect":
await send({"type": "websocket.close", "code": code, "reason": reason})
class TenantHostMiddleware:
"""纯 ASGI 租户路由中间件(设计动机见模块 docstring)。"""
def __init__(self, app: ASGIApp) -> None:
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] not in ("http", "websocket"):
return await self.app(scope, receive, send)
if scope.get("path") == "/health":
return await self.app(scope, receive, send)
registry = getattr(scope["app"].state, "registry", None)
if registry is None:
raise RuntimeError(
"app.state.registry 未挂载——TenantHostMiddleware 需要 "
"lifespan(或测试注入缝)先加载 TenantRegistry")
try:
runtime, site = await registry.resolve_site(_host_from_scope(scope))
except UnknownTenantError:
if scope["type"] == "websocket":
await _close_ws(receive, send, 1008, "unknown host")
else:
await PlainTextResponse("Not Found", status_code=404)(
scope, receive, send)
return
except TenantDisabledError:
if scope["type"] == "websocket":
await _close_ws(receive, send, 1013, "service paused")
else:
await HTMLResponse(PAUSE_PAGE, status_code=503)(
scope, receive, send)
return
state = scope.setdefault("state", {})
state["tenant"] = runtime
state["site"] = site
await self.app(scope, receive, send)
+80 -24
View File
@@ -3,9 +3,11 @@ import os
import shutil
import tempfile
import warnings
from contextlib import contextmanager
from pathlib import Path
import pytest
from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool
@@ -94,33 +96,87 @@ def build_jiancai_runtime():
return _JIANCAI_RUNTIME
@pytest.fixture
def client(tmp_path):
"""TestClient + 每测试独立 Registry(tmp_path 作 data_dir)。
# Task 7: app_client 默认域名——既有测试不带 host 头,默认落 jiancai
# (与 5b 之前 default_runtime 兜底的语义一致,既有断言零改动)
JIANCAI_HOST = "chat2.agent.hnxinyouhe.com"
AUTO_HOST = "chat.agent.hnxinyouhe.com"
- 预置 app.state.registry(测试注入缝): lifespan 见 registry 已挂则
沿用,不再加载仓库 seed_profiles/tenants——每测试拿到全新租户库,
与生产数据互不污染。
- 默认租户(jiancai)的 engine 由 lifespan 在 TestClient 的 portal 事件
循环内预热,后续请求也在同一循环处理(连接/循环绑定安全)。测试
线程需要直接造数时不得用该 engine——asyncio.run 另起循环,跨循环
复用常驻连接会死锁;直连断言走 async_db,或经真实路由(WS/REST)。
- 退出: lifespan finally 里 registry.dispose_all()(Windows 文件锁,
不关则 tmp 目录删不掉);随后置空 app.state.registry(关残留窗——
不让已 dispose 的测试 Registry 被后续无夹具的 lifespan 误用)并
清理 tmp 目录(删不掉告警留痕,不静默吞)。
def build_app(seed_dir, data_dir):
"""app 工厂(Task 7): 建 Registry → 预置 app.state.registry → 返回 app。
Host 中间件已在 app.main 模块级挂好(add_middleware,经 scope["app"]
延迟取 registry),工厂只负责把 Registry 放进 app.state——lifespan 见
已挂则沿用(测试注入缝),不再加载仓库 seed。调用方用完必须复位
app.state.registry = None(见 app_test_client 的 finally)。
"""
from fastapi.testclient import TestClient
registry = TenantRegistry.load_from_dir(SEED_DIR, data_dir=tmp_path)
registry = TenantRegistry.load_from_dir(seed_dir, data_dir=data_dir)
app.state.registry = registry
return app
def _rmtree_warn(path):
try:
with TestClient(app) as c:
shutil.rmtree(path)
except OSError as exc:
warnings.warn(f"测试夹具清理 tmp 目录失败(文件锁?): {exc}")
shutil.rmtree(path, ignore_errors=True)
class TenantTestClient(TestClient):
"""WS 连接默认 Host 取 base_url 域名。
starlette 的 websocket_connect 硬编码 ``ws://testserver`` 拼绝对 URL,
base_url 对 WS 不生效——不补这层,既有 WS 测试的 Host 全落 testserver
(未知域名,被中间件 1008 拒)。显式传 headers={"host": ...} 仍可覆盖
(切租户用,见 test_host_routing)。"""
def __init__(self, *args, ws_default_host: str = "", **kwargs):
super().__init__(*args, **kwargs)
self._ws_default_host = ws_default_host
def websocket_connect(self, url, subprotocols=None, **kwargs):
headers = dict(kwargs.pop("headers", None) or {})
if self._ws_default_host:
headers.setdefault("host", self._ws_default_host)
kwargs["headers"] = headers
return super().websocket_connect(url, subprotocols=subprotocols, **kwargs)
@contextmanager
def app_test_client(app_instance, base_url, cleanup_root):
"""TestClient 包装:统一退出纪律——lifespan finally 已 dispose_all
(Windows 文件锁),再置空 app.state.registry(关残留窗:不让已 dispose
的测试 Registry 被后续无夹具的 lifespan 误用),最后清 tmp 目录
(删不掉告警留痕,不静默吞)。"""
try:
with TenantTestClient(
app_instance, base_url=base_url,
ws_default_host=base_url.removeprefix("http://")) as c:
yield c
finally:
app.state.registry = None
try:
shutil.rmtree(tmp_path)
except OSError as exc:
warnings.warn(f"client 夹具清理 tmp 目录失败(文件锁?): {exc}")
shutil.rmtree(tmp_path, ignore_errors=True)
app_instance.state.registry = None
_rmtree_warn(cleanup_root)
@pytest.fixture
def app_client(tmp_path):
"""TestClient(base_url=jiancai 域名) + 每测试独立 Registry(tmp data_dir)。
- Host 中间件按请求 Host 逐请求 resolve_site:engine 懒建在 TestClient
的 portal 事件循环里(连接/循环绑定安全)。测试线程需要直接造数时
不得用该 engine——asyncio.run 另起循环,跨循环复用常驻连接会死锁;
直连断言走 async_db,或经真实路由(WS/REST)。
- 切换租户: 请求级 headers={"host": ...} 覆盖(WS 同理),不用另开
TestClient——同一 app 上两个 client 各跑一次 lifespan 会共享
registry 而 engine 各绑一个循环,跨循环死锁。
"""
build_app(SEED_DIR, tmp_path)
with app_test_client(app, f"http://{JIANCAI_HOST}", tmp_path) as c:
yield c
@pytest.fixture
def client(app_client):
"""既有测试的 client 夹具(Task 7 起基于 app_client,仅换名)。"""
yield app_client
+141
View File
@@ -0,0 +1,141 @@
# Task 7: Host 中间件租户路由——纯 ASGI(http+websocket 双覆盖)。
#
# 本文件的存在意义(前序审查实证过的坑): 若中间件走了 BaseHTTPMiddleware
# (@app.middleware("http")),WS 请求会整个绕过租户解析落 get_runtime 兜底,
# auto 域名的 WS 流量静默写进 jiancai 库(串库不崩溃,极难发现)。
# test_ws_routes_to_correct_tenant_db 用双库断言钉死这条红线。
import json
import shutil
import pytest
from starlette.websockets import WebSocketDisconnect
from tests.conftest import (
AUTO_HOST,
JIANCAI_HOST,
SEED_DIR,
app_test_client,
build_app,
)
NOPE_HOST = "nope.example.com"
# ---- 平台级豁免 -------------------------------------------------------------
def test_known_host_serves_health(app_client):
# /health 在任何 Host 下可达(部署健康检查 curl 127.0.0.1:8018,
# Host 是 IP:端口,对 registry 必然未知)
assert app_client.get("/health").status_code == 200
r = app_client.get("/health", headers={"host": NOPE_HOST})
assert r.status_code == 200
def test_health_exemption_is_exact_match(app_client):
# 精确 == "/health",/healthxxx 不豁免:未知 Host 由中间件(而非路由)
# 拒绝——text/plain 的 PlainTextResponse 区别于路由 404 的 JSON body
r = app_client.get("/healthxxx", headers={"host": NOPE_HOST})
assert r.status_code == 404
assert r.headers["content-type"].startswith("text/plain")
# ---- 未知 Host --------------------------------------------------------------
def test_unknown_host_gets_404(app_client):
r = app_client.get("/api/chat/session", headers={"host": NOPE_HOST})
assert r.status_code == 404
def test_unknown_host_ws_rejected_with_1008(app_client):
# 探针实证: 服务器未 accept 直接 close 时,TestClient 在 __enter__ 即抛
# WebSocketDisconnect,code/reason 保留——这就是断言方式
with pytest.raises(WebSocketDisconnect) as ei:
with app_client.websocket_connect("/ws/chat", headers={"host": NOPE_HOST}):
pass # pragma: no cover — 未 accept 即 close,进不来
assert ei.value.code == 1008
# ---- 停用租户 ---------------------------------------------------------------
@pytest.fixture
def disabled_client(tmp_path):
"""jiancai 停用版 app_client: 拷 seed 到 tmp、改 status 再加载。
加载前改盘(而非 registry_reload): 无需跨事件循环碰 registry 锁,
最简;reload 语义本身已被 Task 5a 的测试钉过,此处不重复走。"""
seed = tmp_path / "seed"
shutil.copytree(SEED_DIR, seed)
tenant_json = seed / "jiancai" / "tenant.json"
data = json.loads(tenant_json.read_text(encoding="utf-8"))
data["status"] = "disabled"
tenant_json.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
a = build_app(seed, tmp_path / "data")
with app_test_client(a, f"http://{JIANCAI_HOST}", tmp_path) as c:
yield c
def test_disabled_tenant_gets_pause_page(disabled_client):
r = disabled_client.get("/api/chat/session")
assert r.status_code == 503
assert "暂停" in r.text
# 中性页: 不暴露租户 slug/域名
assert "jiancai" not in r.text.lower()
assert JIANCAI_HOST not in r.text
def test_disabled_tenant_ws_rejected_with_1013(disabled_client):
with pytest.raises(WebSocketDisconnect) as ei:
with disabled_client.websocket_connect("/ws/chat"):
pass # pragma: no cover
assert ei.value.code == 1013
# ---- 租户正确性(防串库红线) --------------------------------------------------
def test_default_host_routes_to_jiancai(app_client):
# app_client 不带 host 头 → base_url 的 jiancai 域(既有全量测试的语义)
r = app_client.post("/api/chat/session", json={})
assert r.status_code == 200
def test_ws_routes_to_correct_tenant_db(app_client, tmp_path):
"""WS 经 auto 域名连接 → 会话必须落在 auto 租户库。
三层断言:
1. 文件级——此刻只碰过 auto(WS 用显式 host;lifespan 清扫因无 db
文件而跳过),auto.db 已懒建而 jiancai.db 从未创建;
2. REST resume 经 auto 域名命中同一 sid(is_new=False 且回原 sid);
3. 同一 sid 经 jiancai 域名(base_url 默认)不可见——resume miss 落
is_new=True(新建了另一个会话),证明 jiancai 库里没有它。
(经 /api/chat/session 的 resume 探测而非 /messages: 后者只放行
app-source 会话,WS 造的 standalone 会话两库都会 404,区分不了;
resume 探测不触发 AI 调用。)"""
with app_client.websocket_connect(
"/ws/chat", headers={"host": AUTO_HOST}) as ws:
ws.send_text(json.dumps({"source": "standalone"}))
init = ws.receive_json()
assert init["type"] == "init"
sid = init["session_id"]
# 1. 文件级(jiancai 还没被任何请求 resolve 过)
assert (tmp_path / "auto.db").exists()
assert not (tmp_path / "jiancai.db").exists()
# 2. auto 库命中
r_auto = app_client.post(
"/api/chat/session", json={"session_id": sid},
headers={"host": AUTO_HOST})
assert r_auto.status_code == 200
assert r_auto.json()["is_new"] is False
assert r_auto.json()["session_id"] == sid
# 3. jiancai 库没有(miss → 新建,is_new=True、sid 不同)
r_jc = app_client.post("/api/chat/session", json={"session_id": sid})
assert r_jc.status_code == 200
assert r_jc.json()["is_new"] is True
assert r_jc.json()["session_id"] != sid