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``。本助手的 TenantHostMiddleware(纯 ASGI,http 与 websocket 双覆盖)保证凡到达路由
兜底规则 = state 里有租户用 state 的(中间件灰度接入即生效);没有则用 的请求 ``state.tenant`` 必已就位;唯一例外是平台级豁免路径(/health),
``app.state.default_runtime``(lifespan 预热的默认租户 jiancai),API 测试 而它不访问租户资源。本模块不再有任何兜底——state 无租户即 raise,把
与直连(无 Host 语义)在 5b→7 过渡期不断。 "中间件没挂上/被绕过"这类装配错误炸在第一现场而不是静默串库。
已知过渡限制(记录在案,Task 7 一并解决):默认租户经 registry_reload
热切换后 default_runtime 仍指向旧 runtime——过渡期不修,Task 7 中间件
按请求 Host 重新 resolve 后本模块整体删除。
""" """
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@@ -20,13 +16,10 @@ if TYPE_CHECKING:
def get_runtime(request_or_ws: Request | WebSocket) -> "TenantRuntime": 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) runtime = getattr(request_or_ws.state, "tenant", None)
if runtime is not None: if runtime is None:
return runtime
runtime = getattr(request_or_ws.app.state, "default_runtime", None)
if runtime is None or runtime.session_factory is None:
raise RuntimeError( raise RuntimeError(
"默认租户运行时未就绪——lifespan 未运行、未预热 default_runtime," "请求 state 里没有租户运行时——TenantHostMiddleware 未覆盖该"
"或 registry_reload 已 dispose 旧 runtime(过渡期限制,Task 7 解决)") "请求(中间件未挂载,或豁免路径访问了租户资源)")
return runtime return runtime
+9 -10
View File
@@ -6,6 +6,7 @@ from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles from fastapi.staticfiles import StaticFiles
from app.engine.config import STATIC_DIR from app.engine.config import STATIC_DIR
from app.platform.host_middleware import TenantHostMiddleware
from app.platform.registry import TenantRegistry from app.platform.registry import TenantRegistry
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -14,9 +15,6 @@ logger = logging.getLogger(__name__)
BASE_DIR = Path(__file__).resolve().parents[1] BASE_DIR = Path(__file__).resolve().parents[1]
SEED_DIR = BASE_DIR / "seed_profiles" SEED_DIR = BASE_DIR / "seed_profiles"
TENANTS_DIR = BASE_DIR / "tenants" # .gitignore 已忽略 tenants/*.db 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 @asynccontextmanager
@@ -29,15 +27,10 @@ async def lifespan(app: FastAPI):
app.state.registry = registry app.state.registry = registry
try: 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。逐租户、 # 询盘推送(设计 8.2): 清扫崩溃遗留的 stale pending → failed。逐租户、
# 用该租户自己的 session_factory;db 文件不存在的租户跳过——全新库 # 用该租户自己的 session_factory;db 文件不存在的租户跳过——全新库
# 无遗留行,保持 engine 懒建纪律(预热的默认租户文件已建,照常扫)。 # 无遗留行,保持 engine 懒建纪律(租户 engine 由 Host 中间件按请求
# 懒建,启动期只有已有库文件的租户会被扫描)。
# 清扫失败只记日志继续:一个租户的坏库不炸整个平台启动(其余租户 # 清扫失败只记日志继续:一个租户的坏库不炸整个平台启动(其余租户
# 照常服务),该租户下轮重启还会重扫。 # 照常服务),该租户下轮重启还会重扫。
from app.engine.services.lead_push_service import sweep_stale_pending 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 return resp
# 租户路由中间件(Task 7): 最后 add = 最外层——租户解析先于一切用户中间件。
# 纯 ASGI 类(http+websocket 双覆盖),registry 经 scope["app"] 延迟取,
# 因此可在 lifespan 挂 Registry 之前于模块级注册。
app.add_middleware(TenantHostMiddleware)
@app.get("/") @app.get("/")
async def index(): async def index():
from fastapi.responses import FileResponse 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 shutil
import tempfile import tempfile
import warnings import warnings
from contextlib import contextmanager
from pathlib import Path from pathlib import Path
import pytest import pytest
from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool from sqlalchemy.pool import NullPool
@@ -94,33 +96,87 @@ def build_jiancai_runtime():
return _JIANCAI_RUNTIME return _JIANCAI_RUNTIME
@pytest.fixture # Task 7: app_client 默认域名——既有测试不带 host 头,默认落 jiancai
def client(tmp_path): # (与 5b 之前 default_runtime 兜底的语义一致,既有断言零改动)
"""TestClient + 每测试独立 Registry(tmp_path 作 data_dir)。 JIANCAI_HOST = "chat2.agent.hnxinyouhe.com"
AUTO_HOST = "chat.agent.hnxinyouhe.com"
- 预置 app.state.registry(测试注入缝): lifespan 见 registry 已挂则
沿用,不再加载仓库 seed_profiles/tenants——每测试拿到全新租户库, def build_app(seed_dir, data_dir):
与生产数据互不污染。 """app 工厂(Task 7): 建 Registry → 预置 app.state.registry → 返回 app。
- 默认租户(jiancai)的 engine 由 lifespan 在 TestClient 的 portal 事件
循环内预热,后续请求也在同一循环处理(连接/循环绑定安全)。测试 Host 中间件已在 app.main 模块级挂好(add_middleware,经 scope["app"]
线程需要直接造数时不得用该 engine——asyncio.run 另起循环,跨循环 延迟取 registry),工厂只负责把 Registry 放进 app.state——lifespan 见
复用常驻连接会死锁;直连断言走 async_db,或经真实路由(WS/REST)。 已挂则沿用(测试注入缝),不再加载仓库 seed。调用方用完必须复位
- 退出: lifespan finally 里 registry.dispose_all()(Windows 文件锁, app.state.registry = None(见 app_test_client 的 finally)。
不关则 tmp 目录删不掉);随后置空 app.state.registry(关残留窗——
不让已 dispose 的测试 Registry 被后续无夹具的 lifespan 误用)并
清理 tmp 目录(删不掉告警留痕,不静默吞)。
""" """
from fastapi.testclient import TestClient registry = TenantRegistry.load_from_dir(seed_dir, data_dir=data_dir)
registry = TenantRegistry.load_from_dir(SEED_DIR, data_dir=tmp_path)
app.state.registry = registry app.state.registry = registry
return app
def _rmtree_warn(path):
try: 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 yield c
finally: finally:
app.state.registry = None app_instance.state.registry = None
try: _rmtree_warn(cleanup_root)
shutil.rmtree(tmp_path)
except OSError as exc:
warnings.warn(f"client 夹具清理 tmp 目录失败(文件锁?): {exc}") @pytest.fixture
shutil.rmtree(tmp_path, ignore_errors=True) 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