Merge pull request #496 from vastsa/fix/pr494-security-patches

fix: harden share expiry, download token, quota SQL, and CORS

Author: vastsa

apps/base/quota.py

@@ -8,6 +8,43 @@
 from core.utils import get_now
 
 
+def _detect_sql_dialect() -> str:
+    """识别当前默认连接的 SQL 方言。
+
+    优先使用 Tortoise capabilities.dialect；
+    回退到 db_config 中的 engine 路径字符串。
+    """
+    conn = connections.get("default")
+    dialect = getattr(getattr(conn, "capabilities", None), "dialect", "") or ""
+    if dialect:
+        return dialect.lower()
+
+    try:
+        engine = str(connections.db_config.get("default", {}).get("engine", "")).lower()
+    except Exception:
+        engine = ""
+    if "postgres" in engine or "asyncpg" in engine or "psycopg" in engine:
+        return "postgres"
+    if "mysql" in engine:
+        return "mysql"
+    if "sqlite" in engine:
+        return "sqlite"
+    return "sqlite"
+
+
+def _sql_placeholders(count: int) -> list[str]:
+    """根据数据库方言生成参数占位符（多数据库兼容）。
+
+    SQLite 用 ?，PostgreSQL 用 $1/$2/...，MySQL 用 %s。
+    """
+    dialect = _detect_sql_dialect()
+    if dialect in {"postgres", "postgresql"}:
+        return [f"${i}" for i in range(1, count + 1)]
+    if dialect == "mysql":
+        return ["%s"] * count
+    return ["?"] * count
+
+
 def get_storage_limit() -> int:
     try:
         return max(0, int(getattr(settings, "storageLimit", 0)))
@@ -49,15 +86,16 @@ async def reserve_storage(token: str, size: int, ttl_seconds: int) -> None:
         raise HTTPException(status_code=409, detail="上传容量预留信息不一致")
 
     try:
+        ph = _sql_placeholders(6)
         affected, _ = await conn.execute_query(
-            """
+            f"""
             INSERT INTO storagereservation (token, size, expires_at)
-            SELECT ?, ?, ?
+            SELECT {ph[0]}, {ph[1]}, {ph[2]}
             WHERE (
                 COALESCE((SELECT SUM(size) FROM filecodes), 0)
-                + COALESCE((SELECT SUM(size) FROM storagereservation WHERE expires_at > ?), 0)
-                + ?
-            ) <= ?
+                + COALESCE((SELECT SUM(size) FROM storagereservation WHERE expires_at > {ph[3]}), 0)
+                + {ph[4]}
+            ) <= {ph[5]}
             """,
             [token, requested_size, expires_at, now, requested_size, limit],
         )

apps/base/utils.py

@@ -19,6 +19,13 @@
 )
 
 
+def validate_expire_style(expire_style: str) -> str:
+    """校验过期方式是否在管理员配置的白名单内。"""
+    if expire_style not in settings.expireStyle:
+        raise HTTPException(status_code=400, detail="过期时间类型错误")
+    return expire_style
+
+
 async def get_file_path_name(file: UploadFile) -> Tuple[str, str, str, str, str]:
     today = await get_now()
     storage_path = settings.storage_path.strip("/")

apps/base/views.py

@@ -28,6 +28,7 @@
     get_file_path_name,
     ip_limit,
     get_chunk_file_path_name,
+    validate_expire_style,
 )
 from core.response import APIResponse
 from core.settings import settings
@@ -152,6 +153,7 @@ async def share_text(
     expire_style: str = Form(default="day"),
     ip: str = Depends(ip_limit["upload"]),
 ):
+    validate_expire_style(expire_style)
     text_size = len(text.encode("utf-8"))
     max_txt_size = 222 * 1024
     if text_size > max_txt_size:
@@ -187,8 +189,7 @@ async def share_file(
 ):
     file_size = await validate_file_size(file, settings.uploadSize)
     validate_file_type(file.filename or "", file.content_type)
-    if expire_style not in settings.expireStyle:
-        raise HTTPException(status_code=400, detail="过期时间类型错误")
+    validate_expire_style(expire_style)
     path, suffix, prefix, uuid_file_name, save_path = await get_file_path_name(file)
     reservation_token = f"file:{uuid.uuid4().hex}"
     await reserve_storage(reservation_token, file_size, ttl_seconds=3600)
@@ -374,7 +375,12 @@ async def select_file(data: SelectFileModel, ip: str = Depends(ip_limit["error"]
 async def download_file(key: str, code: str, ip: str = Depends(ip_limit["error"])):
     file_storage: FileStorageInterface = storages[settings.file_storage]()
     normalized_code = normalize_share_code(code)
-    if await get_select_token(normalized_code) != key:
+    # 同时接受当前窗口与上一窗口 token，避免时间窗边界竞态导致偶发 403
+    valid_keys = {
+        await get_select_token(normalized_code, offset=0),
+        await get_select_token(normalized_code, offset=1),
+    }
+    if key not in valid_keys:
         ip_limit["error"].add_ip(ip)
         raise HTTPException(status_code=403, detail="下载鉴权失败")
     has, file_code = await get_code_file_by_code(normalized_code)
@@ -659,6 +665,7 @@ async def complete_upload(
     chunk_info = await UploadChunk.filter(upload_id=upload_id, chunk_index=-1).first()
     if not chunk_info:
         raise HTTPException(status.HTTP_404_NOT_FOUND, detail="上传会话不存在")
+    validate_expire_style(data.expire_style)
     await reserve_storage(
         f"chunk:{upload_id}",
         chunk_info.file_size,
@@ -775,8 +782,7 @@ async def presign_upload_init(
             403,
             f"文件大小超过限制，最大为 {settings.uploadSize / (1024 * 1024):.2f} MB",
         )
-    if data.expire_style not in settings.expireStyle:
-        raise HTTPException(400, "过期时间类型错误")
+    validate_expire_style(data.expire_style)
 
     upload_id = uuid.uuid4().hex
     reservation_token = f"presign:{upload_id}"

core/storage.py

@@ -25,6 +25,7 @@
 from apps.base.models import FileCodes, UploadChunk
 from core.utils import get_file_url, sanitize_filename
 from fastapi.responses import FileResponse, StreamingResponse
+from starlette.background import BackgroundTask
 
 
 class FileStorageInterface:
@@ -395,7 +396,9 @@ async def stream_generator():
             return StreamingResponse(
                 stream_generator(),
                 media_type="application/octet-stream",
-                headers=headers
+                headers=headers,
+                # 兜底关闭会话：客户端中断时与 generator finally 双保险
+                background=BackgroundTask(session.close),
             )
         except HTTPException:
             raise
@@ -733,7 +736,9 @@ async def stream_generator():
             return StreamingResponse(
                 stream_generator(),
                 media_type="application/octet-stream",
-                headers=headers
+                headers=headers,
+                # 兜底关闭会话：客户端中断时与 generator finally 双保险
+                background=BackgroundTask(session.close),
             )
         except HTTPException:
             raise
@@ -1209,7 +1214,9 @@ async def stream_generator():
             return StreamingResponse(
                 stream_generator(),
                 media_type="application/octet-stream",
-                headers=headers
+                headers=headers,
+                # 兜底关闭会话：客户端中断时与 generator finally 双保险
+                background=BackgroundTask(session.close),
             )
         except aiohttp.ClientError as e:
             raise HTTPException(

core/utils.py

@@ -41,17 +41,22 @@ async def get_now():
     return datetime.datetime.now(datetime.timezone(datetime.timedelta(hours=8)))
 
 
-async def get_select_token(code: str):
+async def get_select_token(code: str, offset: int = 0):
     """
     获取下载token
-    :param code:
+    :param code: 取件码
+    :param offset: 时间窗口偏移（0=当前窗口，1=上一个窗口）。
+        用于兼容窗口边界竞态：用户在某窗口末尾获取的 token，
+        请求到达服务器时可能已进入下一窗口。
     :return:
     """
     token = getattr(settings, "jwt_secret", "")
     if not token:
         raise RuntimeError("应用签名密钥未初始化")
+    # 每个窗口约 1000 秒；offset 允许校验上一窗口，避免边界竞态
+    time_factor = int(time.time() / 1000) - max(0, int(offset))
     return hashlib.sha256(
-        f"{code}{int(time.time() / 1000)}000{token}".encode()
+        f"{code}{time_factor}000{token}".encode()
     ).hexdigest()
 
 

main.py

@@ -767,7 +767,9 @@ async def refresh_settings_middleware(request, call_next):
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_credentials=True,
+    # 前端使用 Bearer Token，不依赖 Cookie credentials。
+    # allow_origins=["*"] 与 allow_credentials=True 组合不符合 CORS 规范。
+    allow_credentials=False,
     allow_methods=["*"],
     allow_headers=["*"],
 )

tests/test_pr494_security_patches.py

@@ -0,0 +1,141 @@
+import asyncio
+import hashlib
+import time
+import unittest
+from unittest.mock import patch
+
+from fastapi import HTTPException
+from tortoise import Tortoise
+
+from apps.base import views
+from apps.base.models import FileCodes
+from apps.base.quota import _detect_sql_dialect, _sql_placeholders, reserve_storage
+from apps.base.utils import validate_expire_style
+from core.settings import settings
+from core.utils import get_select_token
+
+
+class FakeStorage:
+    async def get_file_url(self, file_code):
+        return f"https://example.invalid/{file_code.code}"
+
+    async def get_file_response(self, file_code):
+        return {"downloaded": file_code.code}
+
+
+class SecurityPatchTests(unittest.TestCase):
+    def test_validate_expire_style_rejects_unknown_mode(self):
+        original = dict(settings.user_config)
+        try:
+            settings.expireStyle = ["day", "count"]
+            self.assertEqual(validate_expire_style("day"), "day")
+            with self.assertRaises(HTTPException) as ctx:
+                validate_expire_style("forever")
+            self.assertEqual(ctx.exception.status_code, 400)
+        finally:
+            settings.user_config = original
+
+    def test_sql_placeholders_follow_dialect(self):
+        asyncio.run(self._assert_sql_placeholders())
+
+    async def _assert_sql_placeholders(self):
+        await Tortoise.init(
+            config={
+                "connections": {
+                    "default": {
+                        "engine": "tortoise.backends.sqlite",
+                        "credentials": {"file_path": ":memory:"},
+                    }
+                },
+                "apps": {
+                    "models": {
+                        "models": ["apps.base.models"],
+                        "default_connection": "default",
+                    }
+                },
+                "use_tz": False,
+                "timezone": "Asia/Shanghai",
+            }
+        )
+        try:
+            self.assertEqual(_detect_sql_dialect(), "sqlite")
+            self.assertEqual(_sql_placeholders(3), ["?", "?", "?"])
+
+            with patch(
+                "apps.base.quota._detect_sql_dialect", return_value="postgres"
+            ):
+                self.assertEqual(_sql_placeholders(3), ["$1", "$2", "$3"])
+            with patch("apps.base.quota._detect_sql_dialect", return_value="mysql"):
+                self.assertEqual(_sql_placeholders(2), ["%s", "%s"])
+
+            # sqlite 下 reserve_storage 仍可正常工作
+            settings.storageLimit = 100
+            await Tortoise.generate_schemas()
+            await reserve_storage("patch-token", 10, 300)
+        finally:
+            settings.storageLimit = 0
+            await Tortoise.close_connections()
+
+    def test_download_token_accepts_previous_window(self):
+        asyncio.run(self._assert_download_token_window())
+
+    async def _assert_download_token_window(self):
+        original = dict(settings.user_config)
+        settings.jwt_secret = "test-secret-key-for-token-window"
+        now = 2_000_500  # 落在窗口 2000 内
+        code = "ABCDE"
+        with patch("core.utils.time.time", return_value=now):
+            current = await get_select_token(code, offset=0)
+            previous = await get_select_token(code, offset=1)
+
+        expected_current = hashlib.sha256(
+            f"{code}{int(now / 1000)}000{settings.jwt_secret}".encode()
+        ).hexdigest()
+        expected_previous = hashlib.sha256(
+            f"{code}{int(now / 1000) - 1}000{settings.jwt_secret}".encode()
+        ).hexdigest()
+        self.assertEqual(current, expected_current)
+        self.assertEqual(previous, expected_previous)
+
+        await Tortoise.init(
+            config={
+                "connections": {
+                    "default": {
+                        "engine": "tortoise.backends.sqlite",
+                        "credentials": {"file_path": ":memory:"},
+                    }
+                },
+                "apps": {
+                    "models": {
+                        "models": ["apps.base.models"],
+                        "default_connection": "default",
+                    }
+                },
+                "use_tz": False,
+                "timezone": "Asia/Shanghai",
+            }
+        )
+        await Tortoise.generate_schemas()
+        try:
+            await FileCodes.create(
+                code=code,
+                prefix="demo",
+                suffix=".txt",
+                expired_count=-1,
+                size=1,
+            )
+            # 模拟请求进入下一窗口，但仍携带上一窗口 token
+            next_window = now + 1000
+            with patch("core.utils.time.time", return_value=next_window):
+                with patch.dict(views.storages, {"local": FakeStorage}):
+                    result = await views.download_file(
+                        key=current, code=code, ip="127.0.0.1"
+                    )
+            self.assertEqual(result, {"downloaded": code})
+        finally:
+            settings.user_config = original
+            await Tortoise.close_connections()
+
+
+if __name__ == "__main__":
+    unittest.main()