diff --git a/apps/base/dependencies.py b/apps/base/dependencies.py index c7748547c..edf31b1fc 100644 --- a/apps/base/dependencies.py +++ b/apps/base/dependencies.py @@ -1,6 +1,7 @@ +import time +from collections import OrderedDict from ipaddress import ip_address, ip_network -from typing import Dict, Iterable, TypedDict -from datetime import datetime, timedelta +from typing import Iterable, TypedDict from fastapi import HTTPException, Request from core.settings import settings @@ -67,41 +68,66 @@ def get_client_ip(request: Request) -> str: class _IPRecord(TypedDict): count: int - time: datetime + last_seen: float + + +# Upper bound for tracked source IPs per limiter. Without it a scan using +# spoofed/rotating addresses grows the dict forever on a small VPS. +DEFAULT_MAX_TRACKED_IPS = 10_000 class IPRateLimit: - def __init__(self, count: int, minutes: int): - self.ips: Dict[str, _IPRecord] = {} + def __init__( + self, + count: int, + minutes: int, + max_ips: int = DEFAULT_MAX_TRACKED_IPS, + ): + # Ordered by last activity (oldest first) so capacity eviction and the + # periodic cleanup are both O(1)/O(expired) instead of a full rebuild. + self.ips: OrderedDict[str, _IPRecord] = OrderedDict() self.count = count self.minutes = minutes + self.max_ips = max(1, max_ips) + + def _window_seconds(self) -> float: + # minutes is mutable at runtime via apps.base.config._sync_ip_limits. + return float(self.minutes) * 60 def check_ip(self, ip: str) -> bool: - if ip in self.ips: - ip_info = self.ips[ip] - if ip_info["count"] >= self.count: - if ip_info["time"] + timedelta(minutes=self.minutes) > datetime.now(): - return False - self.ips.pop(ip) + record = self.ips.get(ip) + if record is None: + return True + if record["count"] >= self.count: + if record["last_seen"] + self._window_seconds() > time.monotonic(): + return False + self.ips.pop(ip, None) return True def add_ip(self, ip: str) -> int: - ip_info = self.ips.setdefault( - ip, {"count": 0, "time": datetime.now()} - ) - ip_info["count"] += 1 - ip_info["time"] = datetime.now() - self.ips[ip] = ip_info - return ip_info["count"] + now = time.monotonic() + record = self.ips.get(ip) + if record is None: + record = {"count": 0, "last_seen": now} + self.ips[ip] = record + else: + self.ips.move_to_end(ip) + record["count"] += 1 + record["last_seen"] = now + while len(self.ips) > self.max_ips: + # Evict the least recently active IP; fail-open for the evicted + # address is preferable to unbounded memory growth. + self.ips.popitem(last=False) + return record["count"] async def remove_expired_ip(self) -> None: - now = datetime.now() - expiration = timedelta(minutes=self.minutes) - self.ips = { - ip: info - for ip, info in self.ips.items() - if info["time"] + expiration >= now - } + now = time.monotonic() + window = self._window_seconds() + while self.ips: + _, record = next(iter(self.ips.items())) + if record["last_seen"] + window >= now: + break + self.ips.popitem(last=False) def __call__(self, request: Request) -> str: ip = get_client_ip(request) diff --git a/apps/base/pages.py b/apps/base/pages.py index 53c18e0ff..697483570 100644 --- a/apps/base/pages.py +++ b/apps/base/pages.py @@ -1,6 +1,9 @@ """Public-facing pages and routes: setup wizard, theme assets, index, robots.txt, and the public config endpoints.""" import html +import re +from functools import lru_cache +from pathlib import Path from fastapi import APIRouter, HTTPException, Request from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, RedirectResponse @@ -83,9 +86,58 @@ def resolve_theme_file(*parts: str): return file_path +# Vite 产物带内容指纹(如 index-BX2s1QqP.js),内容变化时文件名必然变化, +# 可以安全地长缓存;public/ 下的未哈希资源原地升级时名字不变,只给短缓存, +# 过期后靠 FileResponse 自带的 ETag/Last-Modified 走 304。 +_HASHED_ASSET_RE = re.compile( + r"-[A-Za-z0-9_-]{8,}\.(?:js|mjs|css|map|woff2?|ttf|otf|eot|svg|png|jpe?g|gif|webp|avif|ico)$", + re.IGNORECASE, +) +_IMMUTABLE_ASSET_CACHE = "public, max-age=31536000, immutable" +_DEFAULT_ASSET_CACHE = "public, max-age=86400" + + +def _asset_cache_control(asset_path: str) -> str: + if _HASHED_ASSET_RE.search(asset_path): + return _IMMUTABLE_ASSET_CACHE + return _DEFAULT_ASSET_CACHE + + +@lru_cache(maxsize=8) +def _render_index_html( + template_path: str, + inode: int, + mtime_ns: int, + size: int, + title: str, + description: str, + keywords: str, + opacity: str, + background: str, +) -> str: + """Render index.html once per (theme file version, site config) pair. + + The stat fields keep the cache honest when a theme is replaced in place; + config values are part of the key so admin edits are visible immediately + without any cross-module invalidation hook. + """ + return ( + Path(template_path) + .read_text(encoding="utf-8") + .replace("{{title}}", html.escape(title)) + .replace("{{description}}", html.escape(description)) + .replace("{{keywords}}", html.escape(keywords)) + .replace("{{opacity}}", html.escape(opacity)) + .replace("{{background}}", html.escape(background)) + ) + + @router.get("/assets/{asset_path:path}", include_in_schema=False) async def theme_asset(asset_path: str): - return FileResponse(resolve_theme_file("assets", asset_path)) + return FileResponse( + resolve_theme_file("assets", asset_path), + headers={"Cache-Control": _asset_cache_control(asset_path)}, + ) async def not_found_handler(request, exc=None): @@ -115,14 +167,26 @@ async def index(request=None, exc=None): # Site config is admin input (and during the setup window anyone can claim it); # always escape before injecting into the theme template to prevent stored XSS # (mirrors the setup page). + template_path = resolve_theme_file("index.html") + try: + stat_result = template_path.stat() + except OSError: + # resolve_theme_file 与 stat 之间文件被删除/替换时仍按 404 处理, + # 不能退回缓存内容。 + raise HTTPException(status_code=404, detail="资源不存在") + content = _render_index_html( + str(template_path), + stat_result.st_ino, + stat_result.st_mtime_ns, + stat_result.st_size, + str(settings.name), + str(settings.description), + str(settings.keywords), + str(settings.opacity), + str(settings.background), + ) return HTMLResponse( - content=resolve_theme_file("index.html") - .read_text(encoding="utf-8") - .replace("{{title}}", html.escape(str(settings.name))) - .replace("{{description}}", html.escape(str(settings.description))) - .replace("{{keywords}}", html.escape(str(settings.keywords))) - .replace("{{opacity}}", html.escape(str(settings.opacity))) - .replace("{{background}}", html.escape(str(settings.background))), + content=content, media_type="text/html", headers={"Cache-Control": "no-cache"}, ) diff --git a/main.py b/main.py index c4806da15..935cf64f9 100644 --- a/main.py +++ b/main.py @@ -10,7 +10,7 @@ from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from tortoise import Tortoise -from tortoise.contrib.fastapi import register_tortoise +from tortoise.contrib.fastapi import tortoise_exception_handlers from apps.admin.views import admin_api from apps.base.config import ( @@ -29,7 +29,7 @@ ) from apps.base.views import share_api, chunk_api, presign_api from apps.base.mcp import router as mcp_router -from core.database import db_startup_lock, get_db_config, init_db +from core.database import db_startup_lock, init_db from core.errors import StorageError from core.logger import get_log_level_name, is_access_log_enabled, logger from core.settings import settings @@ -89,7 +89,14 @@ async def load_config(): ) -app = FastAPI(lifespan=lifespan, version=APP_VERSION) +# Tortoise 的 ORM 异常映射(DoesNotExist → 404、IntegrityError → 422)。 +# 只注册处理器;连接初始化统一由 lifespan 里的 init_db() 负责, +# 避免 register_tortoise 合并 lifespan 造成的双重初始化。 +app = FastAPI( + lifespan=lifespan, + version=APP_VERSION, + exception_handlers=tortoise_exception_handlers(), +) @app.exception_handler(StorageError) @@ -126,14 +133,6 @@ async def refresh_settings_middleware(request, call_next): allow_headers=["*"], ) -# 使用 register_tortoise 来添加异常处理器 -register_tortoise( - app, - config=get_db_config(), - generate_schemas=False, - add_exception_handlers=True, -) - app.include_router(share_api) app.include_router(chunk_api) app.include_router(presign_api) diff --git a/tests/test_index_cache.py b/tests/test_index_cache.py new file mode 100644 index 000000000..f400110ab --- /dev/null +++ b/tests/test_index_cache.py @@ -0,0 +1,82 @@ +"""index.html render cache must stay correct across config/theme changes.""" +import asyncio +import os +import tempfile +import unittest +from pathlib import Path + +from fastapi import HTTPException + +import apps.base.pages as pages +from core.settings import settings + + +class IndexRenderCacheTests(unittest.TestCase): + def setUp(self): + self._original_user_config = dict(settings.user_config) + self._original_resolve_theme_file = pages.resolve_theme_file + self._tmpdir = tempfile.TemporaryDirectory() + self.template = Path(self._tmpdir.name) / "index.html" + self.template.write_text( + "
rebuilt
", encoding="utf-8" + ) + os.utime( + self.template, + ns=(old_stat.st_atime_ns, old_stat.st_mtime_ns + 1_000_000), + ) + + html = self._render() + self.assertIn("rebuilt
", html) + self.assertEqual(pages._render_index_html.cache_info().misses, 2) + + def test_missing_template_still_404s_instead_of_serving_cache(self): + self._render() + self.template.unlink() + + with self.assertRaises(HTTPException) as ctx: + asyncio.run(pages.index()) + self.assertEqual(ctx.exception.status_code, 404) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_index_template_escaping.py b/tests/test_index_template_escaping.py index 3b6663c55..8f06c32e9 100644 --- a/tests/test_index_template_escaping.py +++ b/tests/test_index_template_escaping.py @@ -18,6 +18,7 @@ class IndexTemplateEscapingTests(unittest.TestCase): def setUp(self): self._original_user_config = dict(settings.user_config) self._original_resolve_theme_file = pages.resolve_theme_file + pages._render_index_html.cache_clear() tmp = tempfile.NamedTemporaryFile( "w", suffix=".html", delete=False, encoding="utf-8" ) @@ -32,6 +33,7 @@ def setUp(self): def tearDown(self): settings.user_config = self._original_user_config pages.resolve_theme_file = self._original_resolve_theme_file + pages._render_index_html.cache_clear() self._template_path.unlink() def _patch_template(self): diff --git a/tests/test_rate_limit_bounds.py b/tests/test_rate_limit_bounds.py new file mode 100644 index 000000000..29843065c --- /dev/null +++ b/tests/test_rate_limit_bounds.py @@ -0,0 +1,81 @@ +"""IPRateLimit is bounded and uses a monotonic clock.""" +import unittest +from unittest import mock + +from apps.base.dependencies import IPRateLimit + + +class FakeClock: + def __init__(self, now: float = 1000.0): + self.now = now + + def __call__(self) -> float: + return self.now + + +class IPRateLimitBoundsTests(unittest.TestCase): + def test_first_add_returns_one(self): + limiter = IPRateLimit(count=3, minutes=1) + + self.assertEqual(limiter.add_ip("203.0.113.1"), 1) + self.assertEqual(limiter.add_ip("203.0.113.1"), 2) + + def test_capacity_evicts_least_recently_active_ip(self): + limiter = IPRateLimit(count=10, minutes=1, max_ips=3) + + for ip in ("a", "b", "c", "d"): + limiter.add_ip(ip) + + self.assertEqual(len(limiter.ips), 3) + self.assertEqual(list(limiter.ips), ["b", "c", "d"]) + + def test_activity_refreshes_recency(self): + limiter = IPRateLimit(count=10, minutes=1, max_ips=2) + + limiter.add_ip("a") + limiter.add_ip("b") + limiter.add_ip("a") + limiter.add_ip("c") + + self.assertEqual(list(limiter.ips), ["a", "c"]) + + def test_window_blocks_then_resets(self): + clock = FakeClock() + with mock.patch("time.monotonic", clock): + limiter = IPRateLimit(count=3, minutes=1) + for _ in range(3): + limiter.add_ip("198.51.100.9") + + self.assertFalse(limiter.check_ip("198.51.100.9")) + clock.now += 30 + self.assertFalse(limiter.check_ip("198.51.100.9")) + clock.now += 31 + self.assertTrue(limiter.check_ip("198.51.100.9")) + self.assertNotIn("198.51.100.9", limiter.ips) + + def test_remove_expired_compacts_old_entries_only(self): + clock = FakeClock() + with mock.patch("time.monotonic", clock): + limiter = IPRateLimit(count=10, minutes=1) + limiter.add_ip("old") + clock.now += 30 + limiter.add_ip("fresh") + clock.now += 31 + + import asyncio + + asyncio.run(limiter.remove_expired_ip()) + + self.assertEqual(list(limiter.ips), ["fresh"]) + + def test_record_count_is_bounded_under_ip_scan(self): + limiter = IPRateLimit(count=10, minutes=1, max_ips=10_000) + + for i in range(50_000): + limiter.add_ip(f"10.0.{i // 256}.{i % 256}") + + self.assertEqual(len(limiter.ips), 10_000) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_startup_single_db_init.py b/tests/test_startup_single_db_init.py new file mode 100644 index 000000000..3abb81416 --- /dev/null +++ b/tests/test_startup_single_db_init.py @@ -0,0 +1,68 @@ +"""One Tortoise init path; ORM exception mappings still registered.""" +import asyncio +import unittest +from pathlib import Path +from unittest import mock + +from tortoise.exceptions import DoesNotExist, IntegrityError + +import main + + +async def _idle_task(): + # Stand-in for the real background loops; the lifespan cancels it on exit. + await asyncio.sleep(3600) + + +class StartupSingleInitTests(unittest.TestCase): + def test_orm_exception_handlers_are_registered(self): + # Behavior pinned by test_admin_write_paths: DoesNotExist -> 404 JSON. + # IntegrityError -> 422 must stay available too. + self.assertIn(DoesNotExist, main.app.exception_handlers) + self.assertIn(IntegrityError, main.app.exception_handlers) + + def test_register_tortoise_is_not_used(self): + # register_tortoise merges another lifespan that calls Tortoise.init + # before init_db(); batch 4 (SQLite PRAGMAs) needs a single init path. + source = Path("main.py").read_text(encoding="utf-8") + self.assertNotIn("import register_tortoise", source) + self.assertNotIn("register_tortoise(", source) + + def test_lifespan_does_not_initialize_tortoise_outside_init_db(self): + async def scenario(): + init_db_calls = [] + + async def fake_init_db(): + init_db_calls.append(True) + + async def fake_load_config(): + return None + + with ( + mock.patch.object(main, "init_db", fake_init_db), + mock.patch.object(main, "load_config", fake_load_config), + mock.patch.object(main, "delete_expire_files", _idle_task), + mock.patch.object(main, "clean_incomplete_uploads", _idle_task), + mock.patch.object( + main, "clean_expired_presign_sessions", _idle_task + ), + mock.patch.object(main.Tortoise, "init") as tortoise_init, + mock.patch.object( + main.Tortoise, "close_connections" + ), + ): + # FastAPI merges the lifespan of every included router, so use + # the installed context rather than main.lifespan directly. + async with main.app.router.lifespan_context(main.app): + pass + + return init_db_calls, tortoise_init + + init_db_calls, tortoise_init = asyncio.run(scenario()) + + self.assertEqual(init_db_calls, [True]) + tortoise_init.assert_not_awaited() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_theme_asset_cache.py b/tests/test_theme_asset_cache.py new file mode 100644 index 000000000..cd164c373 --- /dev/null +++ b/tests/test_theme_asset_cache.py @@ -0,0 +1,54 @@ +"""Theme assets get a safe Cache-Control policy. + +Fingerprinted (Vite) filenames can be cached immutably; public assets whose +names do not change must use a short TTL so in-place upgrades are picked up. +""" +import asyncio +import tempfile +import unittest +from pathlib import Path + +import apps.base.pages as pages + + +class ThemeAssetCacheTests(unittest.TestCase): + def setUp(self): + self._original_resolve_theme_file = pages.resolve_theme_file + self._tmpdir = tempfile.TemporaryDirectory() + self.root = Path(self._tmpdir.name) + (self.root / "assets").mkdir() + pages.resolve_theme_file = lambda *parts: self.root.joinpath(*parts) + + def tearDown(self): + pages.resolve_theme_file = self._original_resolve_theme_file + self._tmpdir.cleanup() + + def _asset_response(self, name: str): + (self.root / "assets" / name).write_bytes(b"x") + return asyncio.run(pages.theme_asset(name)) + + def test_fingerprinted_asset_is_immutable(self): + response = self._asset_response("index-BX2s1QqP.js") + + self.assertIn("max-age=31536000", response.headers["cache-control"]) + self.assertIn("immutable", response.headers["cache-control"]) + + def test_plain_asset_gets_short_cache(self): + response = self._asset_response("logo.png") + + self.assertIn("max-age=86400", response.headers["cache-control"]) + self.assertNotIn("immutable", response.headers["cache-control"]) + + def test_short_hash_is_not_treated_as_fingerprint(self): + self.assertEqual( + pages._asset_cache_control("logo-2024.png"), + pages._DEFAULT_ASSET_CACHE, + ) + self.assertEqual( + pages._asset_cache_control("index-BX2s1QqP.css"), + pages._IMMUTABLE_ASSET_CACHE, + ) + + +if __name__ == "__main__": + unittest.main()