Преглед на файлове

web: use slow-api middleware to handle rate limits

Shiv Tyagi преди 7 месеца
родител
ревизия
16ddf45439
променени са 5 файла, в които са добавени 60 реда и са изтрити 20 реда
  1. 13 15
      web/api/v1/builds.py
  2. 35 0
      web/core/limiter.py
  3. 9 0
      web/main.py
  4. 1 0
      web/requirements.txt
  5. 2 5
      web/services/builds.py

+ 13 - 15
web/api/v1/builds.py

@@ -16,7 +16,7 @@ from schemas import (
     BuildOut,
 )
 from services.builds import get_builds_service, BuildsService
-from utils import RateLimitExceededException
+from core.limiter import limiter
 
 router = APIRouter(prefix="/builds", tags=["builds"])
 
@@ -28,9 +28,19 @@ router = APIRouter(prefix="/builds", tags=["builds"])
     responses={
         400: {"description": "Invalid build configuration"},
         404: {"description": "Vehicle, board, or version not found"},
-        429: {"description": "Rate limit exceeded"}
+        429: {
+            "description": "Rate limit exceeded",
+            "content": {
+                "application/json": {
+                    "example": {
+                        "detail": "Too many requests. Try again after some time."
+                    }
+                }
+            }
+        }
     }
 )
+@limiter.limit("10/hour")
 async def create_build(
     build_request: BuildRequest,
     request: Request,
@@ -52,19 +62,7 @@ async def create_build(
         429: Rate limit exceeded
     """
     try:
-        # Get client IP for rate limiting
-        forwarded_for = request.headers.get('X-Forwarded-For', None)
-        if forwarded_for:
-            client_ip = forwarded_for.split(',')[0].strip()
-        else:
-            client_ip = request.client.host if request.client else "unknown"
-
-        return service.create_build(build_request, client_ip)
-    except RateLimitExceededException as e:
-        raise HTTPException(
-            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
-            detail=str(e)
-        )
+        return service.create_build(build_request)
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:

+ 35 - 0
web/core/limiter.py

@@ -0,0 +1,35 @@
+import logging
+from fastapi import Request
+from fastapi.responses import JSONResponse
+from slowapi.errors import RateLimitExceeded
+from slowapi import Limiter
+from slowapi.util import get_remote_address
+from core.config import get_settings
+
+logger = logging.getLogger(__name__)
+
+settings = get_settings()
+
+# We use the same redis instance which is used to store build metadata
+# and other cached data. To keep that data separate, we use db-1 of the
+# redis instance instead of the default db-0.
+REDIS_DB_NUMBER = 1
+limiter = Limiter(
+    key_func=get_remote_address,
+    storage_uri=f"redis://{settings.redis_host}:{settings.redis_port}/{REDIS_DB_NUMBER}",
+    strategy="fixed-window",
+)
+
+
+def rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded) -> JSONResponse:
+    """
+    Response to send when a rate limit is exception is raised
+    """
+    response = JSONResponse(
+        {"detail": "Too many requests. Try again after some time."},
+        status_code=429
+    )
+    response = request.app.state.limiter._inject_headers(
+        response, request.state.view_rate_limit
+    )
+    return response

+ 9 - 0
web/main.py

@@ -11,12 +11,16 @@ import argparse
 
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
+from slowapi.errors import RateLimitExceeded
+from slowapi.middleware import SlowAPIMiddleware
 
 from api.v1 import router as v1_router
 from ui import router as ui_router
+
 from core.config import get_settings
 from core.startup import initialize_application
 from core.logging_config import setup_logging
+from core.limiter import limiter, rate_limit_exceeded_handler
 
 import ap_git
 import metadata_manager
@@ -90,6 +94,7 @@ async def lifespan(app: FastAPI):
     app.state.build_manager = build_mgr
     app.state.inbuilt_builder = inbuilt_builder
     app.state.inbuilt_builder_thread = inbuilt_builder_thread
+    app.state.limiter = limiter
 
     yield
 
@@ -114,6 +119,10 @@ app = FastAPI(
     lifespan=lifespan,
 )
 
+# SlowAPIMiddleware is used for rate limiting
+app.add_middleware(SlowAPIMiddleware)
+app.add_exception_handler(RateLimitExceeded, rate_limit_exceeded_handler)
+
 # Mount static files
 WEB_ROOT = Path(__file__).resolve().parent
 app.mount(

+ 1 - 0
web/requirements.txt

@@ -8,3 +8,4 @@ dill==0.3.8
 packaging==25.0
 jinja2==3.1.2
 python-multipart==0.0.6
+slowapi==0.1.9

+ 2 - 5
web/services/builds.py

@@ -42,15 +42,13 @@ class BuildsService:
 
     def create_build(
         self,
-        build_request: BuildRequest,
-        client_ip: str
+        build_request: BuildRequest
     ) -> BuildSubmitResponse:
         """
         Create a new build request.
 
         Args:
             build_request: Build configuration
-            client_ip: Client IP address for rate limiting
 
         Returns:
             Simple response with build_id and URL
@@ -149,8 +147,7 @@ class BuildsService:
 
         # Submit build
         build_id = self.manager.submit_build(
-            build_info=build_info,
-            client_ip=client_ip,
+            build_info=build_info
         )
 
         # Return simple submission response