diff --git a/build_manager/manager.py b/build_manager/manager.py index 315025f..60bbc2e 100644 --- a/build_manager/manager.py +++ b/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) diff --git a/docker-compose.yml b/docker-compose.yml index 385df2c..6e09c61 100644 --- a/docker-compose.yml +++ b/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: diff --git a/utils/ratelimiter.py b/utils/ratelimiter.py index 19b492d..91daafe 100644 --- a/utils/ratelimiter.py +++ b/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 diff --git a/web/api/v1/builds.py b/web/api/v1/builds.py index d0f0ab9..28a6a18 100644 --- a/web/api/v1/builds.py +++ b/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: diff --git a/web/core/limiter.py b/web/core/limiter.py new file mode 100644 index 0000000..9c4ab98 --- /dev/null +++ b/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 diff --git a/web/main.py b/web/main.py index 4f7adbf..a87d716 100755 --- a/web/main.py +++ b/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( diff --git a/web/requirements.txt b/web/requirements.txt index 789a92e..1862bbd 100644 --- a/web/requirements.txt +++ b/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 diff --git a/web/services/builds.py b/web/services/builds.py index ede6d91..8ebf9a7 100644 --- a/web/services/builds.py +++ b/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