From 54a65ed7da386c9cdf9d0fbd52f34dc9738bfb43 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B5=85=E5=BF=B5?= Date: Mon, 12 Oct 2026 20:54:28 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20WS=20=E6=AF=8F=20IP=20=E5=B9=B6?= =?UTF-8?q?=E5=8F=91=E8=BF=9E=E6=8E=A5=E4=B8=8A=E9=99=90=EF=BC=88=E9=BB=98?= =?UTF-8?q?=E8=AE=A410=EF=BC=8C(=E7=A7=9F=E6=88=B7,IP)=20=E9=94=AE?= =?UTF-8?q?=EF=BC=8Cfinally=20=E9=87=8A=E6=94=BE=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Code --- app/engine/routes/chat.py | 31 ++++++++--- app/engine/services/conn_limiter.py | 57 ++++++++++++++++++++ app/platform/defaults.py | 3 +- tests/conftest.py | 10 ++++ tests/test_conn_limiter.py | 80 +++++++++++++++++++++++++++++ 5 files changed, 173 insertions(+), 8 deletions(-) create mode 100644 app/engine/services/conn_limiter.py create mode 100644 tests/test_conn_limiter.py diff --git a/app/engine/routes/chat.py b/app/engine/routes/chat.py index 8509c48..62b1a56 100644 --- a/app/engine/routes/chat.py +++ b/app/engine/routes/chat.py @@ -5,6 +5,8 @@ from fastapi import APIRouter, WebSocket, WebSocketDisconnect from app.engine.models import ConversationSource, ConversationStatus from app.engine.routes.runtime import get_runtime +from app.engine.services.conn_limiter import DEFAULT_WS_CONN_LIMIT +from app.engine.services.conn_limiter import limiter as conn_limiter from app.engine.services.conversation import ConversationManager logger = logging.getLogger(__name__) @@ -13,21 +15,32 @@ router = APIRouter() @router.websocket("/ws/chat") async def ws_chat(websocket: WebSocket): - await websocket.accept() + # Task 8: 每 (租户, IP) 并发连接上限——必须在 accept 前拒绝(close 代 + # accept,客户端握手即断)。IP 提取与下方 client_info 同一来源(XFF 优先, + # scope client 兜底),握手期 headers/scope 已就绪、不依赖首条 init 消息; + # 客户端自报的 client_info.ip 仍按原逻辑只进会话记录,不参与限流。 + runtime = get_runtime(websocket) + client_ip = websocket.client.host if websocket.client else "" + forwarded = websocket.headers.get("x-forwarded-for", "").split(",")[0].strip() + if forwarded: + client_ip = forwarded + + limit = runtime.config.data["engine"].get( + "ws_conn_limit", DEFAULT_WS_CONN_LIMIT) + if not conn_limiter.try_acquire(runtime.config.slug, client_ip, limit=limit): + await websocket.close(code=1013, reason="too many connections") + return + session_id = None try: + await websocket.accept() # Task 5b: 租户运行时经 state/默认租户兜底解析(过渡期无中间件); # manager 由装配线从 runtime 造(db 会话 + runtime 引用) - runtime = get_runtime(websocket) async with runtime.session_factory() as db: manager = ConversationManager(db, runtime=runtime) - # Collect client info - client_ip = websocket.client.host if websocket.client else "" - forwarded = websocket.headers.get("x-forwarded-for", "").split(",")[0].strip() - if forwarded: - client_ip = forwarded + # Collect client info(client_ip/forwarded 已在握手期提取,复用) user_agent = websocket.headers.get("user-agent", "") client_info = {"ip": client_ip, "user_agent": user_agent} @@ -107,3 +120,7 @@ async def ws_chat(websocket: WebSocket): logger.info("Client disconnected: %s", session_id) except Exception: logger.exception("WebSocket error for session %s", session_id) + finally: + # Task 8: 正常关闭与异常断开都归还名额——acquire 成功后所有路径 + # 必经此处,不泄漏(accept 失败/循环中抛错同样覆盖) + conn_limiter.release(runtime.config.slug, client_ip) diff --git a/app/engine/services/conn_limiter.py b/app/engine/services/conn_limiter.py new file mode 100644 index 0000000..7413397 --- /dev/null +++ b/app/engine/services/conn_limiter.py @@ -0,0 +1,57 @@ +"""WS 每 IP 并发连接上限(spec §6.1 P2)——进程内 (租户, IP) 计数闸。 + +动机: 平台化后单 IP 可对所有租户域名各开一圈 WS 耗尽单进程资源; +REST 消息接口已有的 rate_limit.py 管不到长连接握手。key 是 +(tenant_slug, ip) 二元组——各租户名额独立,不共享预算。 + +约束与取舍(与 rate_limit.py 同一族): +- 单进程内存态, ``--workers 1`` 约束下准确; 多 worker 各自一份计数, + 实际上限 = limit * workers。 +- 纯同步 dict 操作、方法内无 await——单事件循环下 check-then-set + 原子,无需锁;不得在多线程里使用。 +- 计数归零即删 key, 长期运行不随历史 (租户, IP) 对积累内存; + release 未持有的 key 是安全 no-op(不出负数)。 +- limit 可在 try_acquire 时按租户配置覆盖: 模块单例服务全部租户,而 + ws_conn_limit 是 engine 组的租户级配置——按次传入让各租户差异实时 + 生效(含 registry_reload 后), 又不必维护"每租户一个 limiter 实例" + 的第二层簿记(实例缓存还有配置改了不生效的滞留问题)。 +""" +from typing import Optional + +DEFAULT_WS_CONN_LIMIT = 10 + + +class ConnectionLimiter: + """每 (tenant, ip) 并发计数器: 满则 try_acquire 拒绝, release 归零清键。""" + + def __init__(self, limit: int = DEFAULT_WS_CONN_LIMIT) -> None: + self._limit = limit + self._counts: dict[tuple[str, str], int] = {} + + def try_acquire(self, tenant: str, ip: str, *, + limit: Optional[int] = None) -> bool: + """未满则记账并返回 True; 已满返回 False(不记账)。""" + cap = self._limit if limit is None else limit + key = (tenant, ip) + n = self._counts.get(key, 0) + if n >= cap: + return False + self._counts[key] = n + 1 + return True + + def release(self, tenant: str, ip: str) -> None: + """归还一个名额; 计数归零删 key; 未持有的 key 是安全 no-op。""" + key = (tenant, ip) + n = self._counts.get(key, 0) + if n <= 1: + self._counts.pop(key, None) + else: + self._counts[key] = n - 1 + + +# app 级单例: chat.py 的 WS 路由共用这一份进程内计数(--workers 1 约束) +limiter = ConnectionLimiter() + + +def _reset_for_test() -> None: + limiter._counts.clear() diff --git a/app/platform/defaults.py b/app/platform/defaults.py index 2bc4b8c..76771ab 100644 --- a/app/platform/defaults.py +++ b/app/platform/defaults.py @@ -16,7 +16,8 @@ DEFAULT_CONFIG: dict = { "push": {"enabled": False, "url": "", "api_key": "", "prefix": "", "timeout": 10, "max_attempts": 3, "retry_backoff_base": 1.0}, "ai": {"provider": "glm", "api_key": None, "base_url": None, "model": "glm-4-flash"}, - "engine": {"context_seed_keys": {}, "rate_limit_per_hour": 30}, + "engine": {"context_seed_keys": {}, "rate_limit_per_hour": 30, + "ws_conn_limit": 10}, "misc": {"name_to_region_extras": {}}, } diff --git a/tests/conftest.py b/tests/conftest.py index b5abc32..bdd7ef6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -64,6 +64,16 @@ def _reset_rate_limit(): _reset_for_test() +@pytest.fixture(autouse=True) +def _reset_conn_limiter(): + """Task 8: 同上——WS 连接限流器的模块级单例计数前后复位,跨测试不泄漏 + (TestClient 所有连接同落 ip="testclient",不复位会互相顶名额)。""" + from app.engine.services.conn_limiter import _reset_for_test + _reset_for_test() + yield + _reset_for_test() + + @pytest.fixture def jiancai_runtime(): """直连测试用的 jiancai TenantRuntime(Task 6: manager 的 runtime 必填)。 diff --git a/tests/test_conn_limiter.py b/tests/test_conn_limiter.py new file mode 100644 index 0000000..cff1fe5 --- /dev/null +++ b/tests/test_conn_limiter.py @@ -0,0 +1,80 @@ +# Task 8: WS 每 IP 并发连接上限(spec §6.1 P2)。 +# +# 两层测试: +# - 纯单元: ConnectionLimiter 算法(名额/键隔离/归零清 key/no-op release); +# - 接入: /ws/chat 握手期按 (租户, IP) 限流——第 limit+1 条并发连接被 +# 1013 拒(close 代 accept),连接关闭后名额必须释放。 +import json +from contextlib import ExitStack + +import pytest +from starlette.websockets import WebSocketDisconnect + +from app.engine.services.conn_limiter import ConnectionLimiter + +# ---- 纯单元: 算法 ----------------------------------------------------------- + + +def test_limit_per_ip(): + lim = ConnectionLimiter(limit=2) + assert lim.try_acquire("t1", "1.1.1.1") and lim.try_acquire("t1", "1.1.1.1") + assert not lim.try_acquire("t1", "1.1.1.1") # 第 3 个拒绝 + lim.release("t1", "1.1.1.1") + assert lim.try_acquire("t1", "1.1.1.1") # 释放后可再进 + assert lim.try_acquire("t2", "1.1.1.1") # 其他租户不受影响 + assert lim.try_acquire("t1", "2.2.2.2") # 其他 IP 不受影响 + + +def test_count_zero_cleans_key(): + # 计数归零删 key: 长期运行下内部 dict 不随历史 (租户,IP) 对无限增长 + lim = ConnectionLimiter(limit=3) + lim.try_acquire("t", "1.1.1.1") + lim.try_acquire("t", "1.1.1.1") + lim.release("t", "1.1.1.1") + assert ("t", "1.1.1.1") in lim._counts # 还剩 1 个,键保留 + lim.release("t", "1.1.1.1") + assert lim._counts == {} # 归零即清空 + + +def test_release_unheld_key_is_safe_noop(): + lim = ConnectionLimiter(limit=2) + lim.release("nobody", "9.9.9.9") # 从未 acquire 过 + lim.try_acquire("t", "1.1.1.1") + lim.release("t", "1.1.1.1") + lim.release("t", "1.1.1.1") # 双重 release + assert lim._counts == {} # 不出负计数 + assert lim.try_acquire("t", "1.1.1.1") # 状态仍可用 + + +# ---- 接入: /ws/chat 握手期限流 ---------------------------------------------- +# TestClient 的 WS scope client 恒为 ("testclient", 50000) 且测试不带 XFF, +# 故所有连接同落 ip="testclient";app_client 默认 jiancai 域——同租户同 IP, +# 正好是被限的键。默认 ws_conn_limit=10(defaults.py engine 组)。 + + +def test_ws_over_limit_rejected_with_1013(app_client): + """填满 10 个名额 → 第 11 条握手即被拒(1013, 未 accept)。""" + with ExitStack() as stack: + conns = [ + stack.enter_context(app_client.websocket_connect("/ws/chat")) + for _ in range(10) + ] + # 前 10 条真实可用(走完 init),不是恰好没崩 + conns[0].send_text(json.dumps({"source": "standalone"})) + assert conns[0].receive_json()["type"] == "init" + + with pytest.raises(WebSocketDisconnect) as ei: + with app_client.websocket_connect("/ws/chat"): + pass # pragma: no cover — 未 accept 即 close,进不来 + assert ei.value.code == 1013 + assert ei.value.reason == "too many connections" + + +def test_ws_conn_release_allows_reentry(app_client): + with ExitStack() as stack: + for _ in range(10): + stack.enter_context(app_client.websocket_connect("/ws/chat")) + # ExitStack 退出 = 10 条全部关闭,服务端 finally 已 release + with app_client.websocket_connect("/ws/chat") as ws: + ws.send_text(json.dumps({"source": "standalone"})) + assert ws.receive_json()["type"] == "init"