v1.0.0提交

backend/middleware/rate_limit.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+PerToolBox Server - 限流中间件
+Copyright (C) 2024 Sea Network Technology Studio
+Author: Canglan <admin@sea-studio.top>
+License: AGPL v3
+"""
+
+from fastapi import Request, HTTPException
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+from ..config import settings
+
+limiter = Limiter(key_func=get_remote_address)
+
+def setup_rate_limit(app):
+    if settings.RATE_LIMIT_ENABLED:
+        app.state.limiter = limiter
+        app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+
+def rate_limit(requests: int = None, period: int = None):
+    if not settings.RATE_LIMIT_ENABLED:
+        return lambda func: func
+
+    req = requests or settings.RATE_LIMIT_REQUESTS
+    per = period or settings.RATE_LIMIT_PERIOD
+    return limiter.limit(f"{req}/{per} seconds")