Use slow-api middleware for rate limiting

build_manager/manager.py

@@ -2,7 +2,6 @@
 import redis
 import dill
 from enum import Enum
-from utils import RateLimiter
 import logging
 import hashlib
 from metadata_manager import RemoteInfo
@@ -136,14 +135,6 @@ def __init__(self,
         self.__task_queue = redis_task_queue_name
         self.__outdir = outdir
 
-        # Initialide an IP-based rate limiter.
-        # Allow 10 builds per hour per client
-        self.__ip_rate_limiter = RateLimiter(
-            redis_host=redis_host,
-            redis_port=redis_port,
-            time_window_sec=3600,
-            allowed_requests=10
-        )
         self.__build_entry_prefix = "buildmeta-"
         self.logger = logging.getLogger(__name__)
         self.logger.info(
@@ -218,21 +209,17 @@ def __generate_build_id(self, build_info: BuildInfo) -> str:
         return bid
 
     def submit_build(self,
-                     build_info: BuildInfo,
-                     client_ip: str) -> str:
+                     build_info: BuildInfo) -> str:
         """
         Submit a new build request, generate a build ID, and queue the
         build for processing.
 
         Parameters:
             build_info (BuildInfo): The build information.
-            client_ip (str): The IP address of the client submitting the
-            build request.
 
         Returns:
             str: The generated build ID for the submitted build.
         """
-        self.__ip_rate_limiter.count(client_ip)
         build_id = self.__generate_build_id(build_info)
         self.__insert_build_info(build_id=build_id, build_info=build_info)
         self.__queue_build(build_id=build_id)

docker-compose.yml

@@ -22,6 +22,7 @@ services:
       CBS_REMOTES_RELOAD_TOKEN: ${CBS_REMOTES_RELOAD_TOKEN}
       PYTHONPATH: /app
       CBS_BUILD_TIMEOUT_SEC: ${CBS_BUILD_TIMEOUT_SEC:-900}
+      FORWARDED_ALLOW_IPS: ${FORWARDED_ALLOW_IPS:-*}
     volumes:
       - ./base:/base:rw
     depends_on:

utils/ratelimiter.py

@@ -1,3 +1,6 @@
+# NOTE: This module is not currently used. The application uses slowapi
+# rate limiting middleware instead. See web/core/limiter.py.
+
 import redis
 import logging
 

web/api/v1/builds.py

@@ -16,7 +16,7 @@
     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 @@
     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:

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

web/main.py

@@ -11,12 +11,16 @@
 
 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 @@ async def lifespan(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(

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

web/services/builds.py

@@ -42,15 +42,13 @@ def __init__(
 
     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 @@ def create_build(
 
         # Submit build
         build_id = self.manager.submit_build(
-            build_info=build_info,
-            client_ip=client_ip,
+            build_info=build_info
         )
 
         # Return simple submission response