fix: harden share rate limit and code generation (#479)

Author: vastsa

apps/base/dependencies.py

@@ -1,7 +1,69 @@
-from typing import Dict, Union
+from ipaddress import ip_address, ip_network
+from typing import Dict, Iterable, Union
 from datetime import datetime, timedelta
 from fastapi import HTTPException, Request
 
+from core.settings import settings
+
+
+def _iter_trusted_proxies() -> Iterable[str]:
+    trusted_proxies = getattr(settings, "trustedProxies", [])
+    if isinstance(trusted_proxies, str):
+        trusted_proxies = [item.strip() for item in trusted_proxies.split(",")]
+    return [item for item in trusted_proxies if item]
+
+
+def _is_trusted_proxy(host: str) -> bool:
+    try:
+        remote_addr = ip_address(host)
+    except ValueError:
+        return False
+
+    for proxy in _iter_trusted_proxies():
+        try:
+            if remote_addr in ip_network(proxy, strict=False):
+                return True
+        except ValueError:
+            continue
+    return False
+
+
+def _get_forwarded_for_ip(header_value: str, fallback_ip: str) -> str:
+    forwarded_chain = [
+        item.strip() for item in header_value.split(",") if item.strip()
+    ]
+    if not forwarded_chain:
+        return fallback_ip
+
+    for candidate in reversed(forwarded_chain):
+        try:
+            ip_address(candidate)
+        except ValueError:
+            return fallback_ip
+        if not _is_trusted_proxy(candidate):
+            return candidate
+    return forwarded_chain[0]
+
+
+def get_client_ip(request: Request) -> str:
+    client_host = request.client.host if request.client else "unknown"
+    if not _is_trusted_proxy(client_host):
+        return client_host
+
+    forwarded_for = request.headers.get("X-Forwarded-For")
+    if forwarded_for:
+        return _get_forwarded_for_ip(forwarded_for, client_host)
+
+    real_ip = request.headers.get("X-Real-IP")
+    if real_ip:
+        try:
+            ip_address(real_ip)
+        except ValueError:
+            return client_host
+        return real_ip
+
+    return client_host
+
 
 class IPRateLimit:
     def __init__(self, count: int, minutes: int):
@@ -35,11 +97,7 @@ async def remove_expired_ip(self) -> None:
         }
 
     def __call__(self, request: Request) -> str:
-        ip = (
-            request.headers.get("X-Real-IP")
-            or request.headers.get("X-Forwarded-For")
-            or request.client.host
-        )
+        ip = get_client_ip(request)
         if not self.check_ip(ip):
             raise HTTPException(status_code=423, detail="请求次数过多，请稍后再试")
         return ip

apps/base/utils.py

@@ -76,8 +76,6 @@ async def get_expire_info(
             expired_at, extra = result
             if expire_style == "count":
                 expired_count = extra
-            elif expire_style == "forever":
-                code = await get_random_code(style="string")
         else:
             expired_at = result
         if expired_at and expired_at - now > max_timedelta:
@@ -91,9 +89,26 @@ async def get_expire_info(
     return expired_at, expired_count, used_count, code
 
 
-async def get_random_code(style: str = "num") -> str:
+def get_code_generate_type() -> str:
+    code_generate_type = getattr(settings, "code_generate_type", "number")
+    if code_generate_type in {"secret", "string"}:
+        return "secret"
+    return "number"
+
+
+async def get_random_code(style: str | None = None) -> str:
+    code_style = style or get_code_generate_type()
+    if code_style == "num":
+        code_style = "number"
+    if code_style == "string":
+        code_style = "secret"
+
     while True:
-        code = await get_random_num() if style == "num" else await get_random_string()
+        code = (
+            await get_random_num()
+            if code_style == "number"
+            else await get_random_string()
+        )
         if not await FileCodes.filter(code=code).exists():
             return str(code)
 

core/settings.py

@@ -43,6 +43,7 @@
     "uploadSize": 1024 * 1024 * 10,
     "allowed_file_types": ["*"],
     "expireStyle": ["day", "hour", "minute", "forever", "count"],
+    "code_generate_type": "number",
     "uploadMinute": 1,
     "enableChunk": 0,
     "webdav_url": "",
@@ -73,6 +74,7 @@
     "serverPort": 12345,
     "showAdminAddr": 0,
     "robotsText": "User-agent: *\nDisallow: /",
+    "trustedProxies": [],
 }
 
 

core/utils.py

@@ -5,8 +5,8 @@
 import datetime
 import hashlib
 import os
-import random
 import re
+import secrets
 import string
 import time
 
@@ -18,7 +18,7 @@ async def get_random_num():
     获取随机数
     :return:
     """
-    return random.randint(10000, 99999)
+    return secrets.randbelow(90000) + 10000
 
 
 r_s = string.ascii_uppercase + string.digits
@@ -29,7 +29,7 @@ async def get_random_string():
     获取随机字符串
     :return:
     """
-    return "".join(random.choice(r_s) for _ in range(5))
+    return "".join(secrets.choice(r_s) for _ in range(5))
 
 
 async def get_now():