diff --git a/app/engine/services/rate_limit.py b/app/engine/services/rate_limit.py index 5f68eea..7e697c2 100644 --- a/app/engine/services/rate_limit.py +++ b/app/engine/services/rate_limit.py @@ -39,5 +39,21 @@ def check( return True +def sweep(now: float | None = None) -> int: + """Sweep expired buckets (spec 2026-10-19 §7 memory cap): drop keys whose + timestamps have ALL left the window (boundary inclusive, matching check()'s + keep predicate ``t > now - WINDOW``); return the number removed. + + Companion to check()'s per-key lazy pruning: keys never revisited (client + churning tenant/ip/device_id) would otherwise accumulate forever; called + by main.lifespan's periodic task.""" + now = time.time() if now is None else now + dead = [k for k, ts in _buckets.items() + if not ts or ts[-1] <= now - WINDOW] + for k in dead: + del _buckets[k] + return len(dead) + + def _reset_for_test() -> None: _buckets.clear() diff --git a/app/main.py b/app/main.py index 55db38c..d600553 100644 --- a/app/main.py +++ b/app/main.py @@ -1,5 +1,6 @@ -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, suppress from pathlib import Path +import asyncio import logging import os @@ -12,6 +13,8 @@ from starlette.responses import FileResponse, PlainTextResponse from starlette.types import Receive, Scope, Send from app.engine.config import STATIC_DIR, settings +from app.engine.services import rate_limit +from app.platform.auth import login_limiter from app.platform.db import ( Tenant, bootstrap_superuser, @@ -141,6 +144,10 @@ async def lifespan(app: FastAPI): registry = TenantRegistry.load_from_dir(SEED_DIR, data_dir=TENANTS_DIR) app.state.registry = registry + # 周期清扫任务句柄:启动成功(yield 前)才创建,finally 里按非 None 取消 + # ——启动半途失败时跳过,不给清理段引入未定义名。 + _sweep_task: asyncio.Task | None = None + try: await ensure_platform_schema(platform_engine) # 平台设置快照(Phase 4 Task 5A):schema 就绪后全表读入内存—— @@ -196,11 +203,26 @@ async def lifespan(app: FastAPI): "启动清扫失败,跳过该租户(slug=%s, db=%s)", rt.config.slug, rt.db_path) + # 铁律(2026-10-19)进程内面: 登录/聊天限流器过期键周期清扫(spec + # §7——两处内存态字典只增不删,不被回访的键永生;语义见各 sweep + # 注记)。300s 量级足够:清速远高于键积压速度,首步 sleep 不占启动。 + async def _limiter_sweep_loop() -> None: + while True: + await asyncio.sleep(300) + login_limiter.sweep() + rate_limit.sweep() + + _sweep_task = asyncio.create_task(_limiter_sweep_loop()) + yield finally: # 幂等于"更少的引擎":启动半途失败也把已建好的 engine 关掉 # (Windows 文件锁: 不关则库文件删不掉)。platform_engine 含预置的 # 测试引擎——dispose 后 engine 对象仍可重连,测试侧重置引用即可。 + if _sweep_task is not None: + _sweep_task.cancel() + with suppress(asyncio.CancelledError): + await _sweep_task # 等 cancel 落地:不留 pending task 给循环收尾 await registry.dispose_all() await platform_engine.dispose() diff --git a/app/platform/auth.py b/app/platform/auth.py index 89f9d79..c398123 100644 --- a/app/platform/auth.py +++ b/app/platform/auth.py @@ -213,6 +213,19 @@ class LoginLimiter: else: self._fails.pop(key, None) + def sweep(self, now: float | None = None) -> int: + """清扫过期键(spec 2026-10-19 §7——内存上限铁律):删「全部时间戳 + 均已出窗」的键,返回清除数。与 _prune 的单键懒清理互补——allow/ + record_fail 只顺手清被回访的键,失败一次即消失的键(扫号换 IP/ + 用户名)永不被回访,靠本方法全表兜底;由 main.lifespan 周期任务 + 调用。时间戳按记录序追加,ts[-1] 即最新——最新出窗则全键出窗。""" + now = time.time() if now is None else now + dead = [k for k, ts in self._fails.items() + if not ts or ts[-1] <= now - self.window_seconds] + for k in dead: + del self._fails[k] + return len(dead) + # 模块级单例: 路由与依赖共用同一份进程内存态(同 rate_limit._buckets 模式) token_store = TokenStore() diff --git a/tests/test_limiter_sweep.py b/tests/test_limiter_sweep.py new file mode 100644 index 0000000..8fd2e1b --- /dev/null +++ b/tests/test_limiter_sweep.py @@ -0,0 +1,95 @@ +"""limiter 周期清扫(spec 2026-10-19 §7——日志/内存上限铁律的进程内面)。 + +两处限流内存态只增不删的兜底: +- app.platform.auth.LoginLimiter._fails: (username, ip) → list[float]—— + allow/record_fail 只顺手 _prune 被访问的键;失败一次即消失的键(扫号 + 换 IP/换用户名)永不被回访 → 永久滞留。 +- app.engine.services.rate_limit._buckets: (tenant, ip, device_id) → + list[float]——check() 同款单键懒清,键随客户端轮换只增不减。 +sweep() 全表清「全部时间戳均已出窗」的键,由 main.lifespan 的周期任务 +(300s sleep loop)调用。时间推进手法照 test_auth 既有 lockout 用例: +直改内部时间戳,免真实等待。 +""" +import asyncio +import time + +from tests.conftest import ( + JIANCAI_HOST, + SEED_DIR, + app_test_client, + build_app, +) + + +class TestLoginLimiterSweep: + def test_清过期保活跃(self): + from app.platform.auth import LoginLimiter + limiter = LoginLimiter(window_seconds=900.0, max_fails=5) + limiter.record_fail("u", "1.1.1.1") + limiter.record_fail("u", "2.2.2.2") + limiter._fails[("u", "1.1.1.1")] = [0.0] # 远古 → 待清 + n = limiter.sweep(now=time.time()) + assert n == 1 + assert ("u", "1.1.1.1") not in limiter._fails + assert ("u", "2.2.2.2") in limiter._fails + + def test_全活跃时清零(self): + from app.platform.auth import LoginLimiter + limiter = LoginLimiter(window_seconds=900.0) + limiter.record_fail("u", "1.1.1.1") + assert limiter.sweep(now=time.time()) == 0 + assert ("u", "1.1.1.1") in limiter._fails + + def test_窗沿时间戳算过期_空键也清(self): + # ts[-1] <= now - window 与 _prune 的保留判据 t > now - window 同沿 + from app.platform.auth import LoginLimiter + limiter = LoginLimiter(window_seconds=100.0) + now = time.time() + limiter._fails[("u", "1.1.1.1")] = [now - 100.0] # 恰在窗沿 → 出窗 + limiter._fails[("u", "2.2.2.2")] = [] # 空值键 → 待清 + assert limiter.sweep(now=now) == 2 + assert not list(limiter._fails) # 全清空 + + +class TestRateLimitSweep: + def test_同款语义(self): + from app.engine.services import rate_limit + try: + now = time.time() + rate_limit._buckets[("a", "1.1.1.1", "")] = [0.0] # 远古 + rate_limit._buckets[("b", "2.2.2.2", "")] = [now] # 活跃 + rate_limit._buckets[("c", "3.3.3.3", "")] = [] # 空值键 + rate_limit._buckets[("d", "4.4.4.4", "")] = [ + now - rate_limit.WINDOW] # 恰在窗沿 + n = rate_limit.sweep(now=now) + assert n == 3 + assert set(rate_limit._buckets) == {("b", "2.2.2.2", "")} + finally: + rate_limit._buckets.clear() # 模块级状态:清现场(conftest 亦兜底) + + +class TestLifespanSweepTask: + def test_注册与取消(self, tmp_path, monkeypatch): + """lifespan 起周期清扫任务,关停时取消。 + + 手法: monkeypatch asyncio.create_task 包一层记录真实 task(行为 + 不变);TestClient with 进出 = lifespan 启停。不真等 300s——loop + 首步即 sleep,任务常 pending;取消断言 cancelled/cancelling 任一 + 成立即可(实现侧 cancel 后 await 落地,常态是 cancelled)。""" + created = [] + real_create_task = asyncio.create_task + + def _rec(coro, **kwargs): + t = real_create_task(coro, **kwargs) + created.append(t) + return t + + monkeypatch.setattr(asyncio, "create_task", _rec) + a = build_app(SEED_DIR, tmp_path) + with app_test_client(a, f"http://{JIANCAI_HOST}", tmp_path): + tasks = [t for t in created + if getattr(t.get_coro(), "__name__", "") + == "_limiter_sweep_loop"] + assert len(tasks) == 1 # 启动期创建且恰好一份 + assert not tasks[0].done() # 300s sleep 中不会自然完成 + assert tasks[0].cancelled() or tasks[0].cancelling() > 0