From fdb08fe31f7f6e060a53039bc72f91e7c02568fd Mon Sep 17 00:00:00 2001
From: MingxuanGame <MingxuanGame@outlook.com>
Date: Thu, 28 Aug 2025 16:41:33 +0000
Subject: [PATCH] fix(rate-limit): apply on router to avoid TypeError on ws
 endpoint

---
 app/dependencies/rate_limit.py | 14 ++++++++++++++
 app/router/private/router.py   |  4 +++-
 app/router/v1/router.py        |  3 ++-
 app/router/v2/router.py        |  6 +++---
 main.py                        | 10 +---------
 5 files changed, 23 insertions(+), 14 deletions(-)
 create mode 100644 app/dependencies/rate_limit.py

diff --git a/app/dependencies/rate_limit.py b/app/dependencies/rate_limit.py
new file mode 100644
index 0000000..45e374e
--- /dev/null
+++ b/app/dependencies/rate_limit.py
@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from app.config import settings
+
+from fastapi import Depends
+from fastapi_limiter.depends import RateLimiter
+
+if settings.enable_rate_limit:
+    LIMITERS = [
+        Depends(RateLimiter(times=1200, minutes=1)),
+        Depends(RateLimiter(times=200, seconds=1)),
+    ]
+else:
+    LIMITERS = []
diff --git a/app/router/private/router.py b/app/router/private/router.py
index ca77832..556dbe0 100644
--- a/app/router/private/router.py
+++ b/app/router/private/router.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from app.dependencies.rate_limit import LIMITERS
+
 from fastapi import APIRouter
 
-router = APIRouter(prefix="/api/private")
+router = APIRouter(prefix="/api/private", dependencies=LIMITERS)
diff --git a/app/router/v1/router.py b/app/router/v1/router.py
index ecd9006..f46ea2b 100644
--- a/app/router/v1/router.py
+++ b/app/router/v1/router.py
@@ -3,12 +3,13 @@ from __future__ import annotations
 from datetime import datetime
 from enum import Enum
 
+from app.dependencies.rate_limit import LIMITERS
 from app.dependencies.user import v1_authorize
 
 from fastapi import APIRouter, Depends
 from pydantic import BaseModel, field_serializer
 
-router = APIRouter(prefix="/api/v1", dependencies=[Depends(v1_authorize)], tags=["V1 API"])
+router = APIRouter(prefix="/api/v1", dependencies=[Depends(v1_authorize), *LIMITERS], tags=["V1 API"])
 
 
 class AllStrModel(BaseModel):
diff --git a/app/router/v2/router.py b/app/router/v2/router.py
index e4a6b43..5ef8c3f 100644
--- a/app/router/v2/router.py
+++ b/app/router/v2/router.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
+from app.dependencies.rate_limit import LIMITERS
+
 from fastapi import APIRouter
 
-router = APIRouter(prefix="/api/v2")
-
-# 导入所有子路由模块来注册路由
+router = APIRouter(prefix="/api/v2", dependencies=LIMITERS)
diff --git a/main.py b/main.py
index d5b5fb0..da24190 100644
--- a/main.py
+++ b/main.py
@@ -37,12 +37,11 @@ from app.service.osu_rx_statistics import create_rx_statistics
 from app.service.redis_message_system import redis_message_system
 from app.utils import bg_tasks, utcnow
 
-from fastapi import Depends, FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request
 from fastapi.exceptions import RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi_limiter import FastAPILimiter
-from fastapi_limiter.depends import RateLimiter
 import sentry_sdk
 
 
@@ -152,13 +151,6 @@ app = FastAPI(
     lifespan=lifespan,
     description=desc,
 )
-if settings.enable_rate_limit:
-    app.router.dependencies.extend(
-        [
-            Depends(RateLimiter(times=1200, minutes=1)),
-            Depends(RateLimiter(times=200, seconds=1)),
-        ]
-    )
 
 
 app.include_router(api_v2_router)