ai-kefu-admin/app/platform/host_middleware.py

"""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`` 不豁免);
- path 前缀 ``/api/platform/`` 豁免(Phase 2 Task 3,平台后台不依赖租户
  域名解析;安全模型详见 ``__call__`` 内注释——豁免≠放行);
- path 等于 ``/admin`` 或前缀 ``/admin/`` 豁免(Phase 4 Task 1,生产
  admin SPA 静态面,同为平台级;``/admin-evil`` 不豁免);
- 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 中性,不带租户信息。

    客户端在握手期间放弃时收到的首个事件是 websocket.disconnect 而非
    connect——对端已撤,close 送达不了,直接 return。receive 本身抛异常
    (传输层断)则向上传播,由外层 ServerErrorMiddleware 记日志:拒绝路径
    无任何已建立状态可留,无需也不应再兜一层。"""
    event = await receive()
    if event["type"] == "websocket.disconnect":
        return
    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)

        # 平台级豁免(先于租户解析;判断在 scope type 分发之后,http 与
        # websocket 两类 scope 同构生效——当前 /api/platform 前缀没有 WS
        # 端点,保持与 /health 同构以防将来有):
        # - /health 精确匹配(见模块 docstring);
        # - /api/platform/ 前缀(Phase 2 Task 3);
        # - /admin 与 /admin/ 前缀(Phase 4 Task 1,生产 admin SPA 静态
        #   面;hash 路由下直连的都是 /admin/ 一级,深链防御同构保留)。
        #
        # 安全模型——豁免≠放行:本豁免只跳过"租户解析"这一层,不打开
        # 任何受保护端点:/api/platform/* 的认证由路由依赖守护
        # (platform_user = get_current_user + enforce_password_changed,
        # 未带有效 token 一律 401/409)。/admin 只是静态资源(公开的前端
        # 产物),真正的数据面仍在 /api/platform/* 认证后。未知 Host 下
        # admin 面可达是 **有意设计**:生产的 admin 域名与租户域名在反代
        # 层区分(见 spec §4 边界规则与部署文档),应用层不按 Host 收紧
        # admin 面。
        # 前缀精确到斜杠:startswith("/api/platform/") 与 startswith(
        # "/admin/"),/api/platform-evil、/admin-evil 不豁免(仍走租户
        # 解析,测试 test_prefix_match_precise / test_admin_exempt_from_
        # tenant_middleware 钉住)。
        #
        # 已知观察项(本任务仅注记不改实现):反代部署形态下登录限流键
        # 取 request.client.host 不解析 XFF——所有请求同源(反代 IP),
        # 键退化为按用户名(fail-safe:只更严不会更松,单账号爆破面
        # 不变);可信代理头解析归部署批次统一处理。
        path = scope.get("path") or ""
        if (path == "/health"
                or path.startswith("/api/platform/")
                or path == "/admin"
                or path.startswith("/admin/")):
            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)