diff --git a/backend/oasst_backend/api/deps.py b/backend/oasst_backend/api/deps.py index fe26f0a6..b347da48 100644 --- a/backend/oasst_backend/api/deps.py +++ b/backend/oasst_backend/api/deps.py @@ -3,8 +3,9 @@ from secrets import token_hex from typing import Generator from uuid import UUID -from fastapi import Depends, Security +from fastapi import Depends, Request, Response, Security from fastapi.security.api_key import APIKey, APIKeyHeader, APIKeyQuery +from fastapi_limiter.depends import RateLimiter from loguru import logger from oasst_backend.config import settings from oasst_backend.database import engine @@ -84,3 +85,60 @@ def get_trusted_api_client( http_status_code=HTTPStatus.FORBIDDEN, ) return client + + +class UserRateLimiter(RateLimiter): + def __init__( + self, times: int = 100, milliseconds: int = 0, seconds: int = 0, minutes: int = 1, hours: int = 0 + ) -> None: + async def identifier(request: Request) -> str: + """Identify a request based on api_key and user.id""" + api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key") + user = (await request.json()).get("user") + return f"{api_key}:{user.get('id')}" + + super().__init__(times, milliseconds, seconds, minutes, hours, identifier) + + async def __call__(self, request: Request, response: Response, api_key: str = Depends(get_api_key)) -> None: + # Skip if rate limiting is disabled + if not settings.RATE_LIMIT: + return + + # Attempt to retrieve api_key and user information + api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key") + user = (await request.json()).get("user") + + # Skip when api_key and user information are not available + # (such that it will be handled by `APIClientRateLimiter`) + if not api_key or not user or not user.get("id"): + return + + return await super().__call__(request, response) + + +class APIClientRateLimiter(RateLimiter): + def __init__( + self, times: int = 10_000, milliseconds: int = 0, seconds: int = 0, minutes: int = 1, hours: int = 0 + ) -> None: + async def identifier(request: Request) -> str: + """Identify a request based on api_key and user.id""" + api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key") + return f"{api_key}" + + super().__init__(times, milliseconds, seconds, minutes, hours, identifier) + + async def __call__(self, request: Request, response: Response, api_key: str = Depends(get_api_key)) -> None: + # Skip if rate limiting is disabled + if not settings.RATE_LIMIT: + return + + # Attempt to retrieve api_key and user information + api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key") + user = (await request.json()).get("user") + + # Skip if user information is available + # (such that it will be handled by `UserRateLimiter`) + if not api_key or user: + return + + return await super().__call__(request, response) diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py index aaa4a8c1..6571a5a9 100644 --- a/backend/oasst_backend/api/v1/tasks.py +++ b/backend/oasst_backend/api/v1/tasks.py @@ -127,7 +127,14 @@ def generate_task( return task, message_tree_id, parent_message_id -@router.post("/", response_model=protocol_schema.AnyTask) # work with Union once more types are added +@router.post( + "/", + response_model=protocol_schema.AnyTask, + dependencies=[ + Depends(deps.UserRateLimiter(times=100, minutes=5)), + Depends(deps.APIClientRateLimiter(times=10_000, minutes=1)), + ], +) # work with Union once more types are added def request_task( *, db: Session = Depends(deps.get_db),