服务暂停维护中，请稍后再试。
" + "" +) + + +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) diff --git a/tests/conftest.py b/tests/conftest.py index 547252d..c935c52 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/test_host_routing.py b/tests/test_host_routing.py new file mode 100644 index 0000000..5da828c --- /dev/null +++ b/tests/test_host_routing.py @@ -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