python_Sklyarenko/Lab2/app/limiter.py - master - ivanstrike/tasker

python_Sklyarenko/Lab2/app/limiter.py
48 строк2 KB
ivanstrike
Lab3
13 дек 2025, 11:15
13 дек 2025, 11:1515c0fa5
import os
import time
from typing import Callable, Dict
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse

RATE_LIMIT_REQUESTS = int(os.getenv("RATE_LIMIT_REQUESTS", "60"))
RATE_LIMIT_WINDOW_SECONDS = int(os.getenv("RATE_LIMIT_WINDOW_SECONDS", "60"))

# simple fixed-window limiter per client (by IP or API Key)
_buckets: Dict[str, Dict[str, float]] = {}  # key -> {count, window_start}

def _key_from_request(request: Request) -> str:
    api_key = request.headers.get("X-API-Key")
    if api_key:
        return f"api:{api_key}"
    # fallback to client host
    client = request.client.host if request.client else "unknown"
    return f"ip:{client}"

class RateLimitMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: Callable):
        key = _key_from_request(request)
        now = time.time()
        bucket = _buckets.get(key)
        if not bucket or now - bucket["window_start"] >= RATE_LIMIT_WINDOW_SECONDS:
            bucket = {"count": 0, "window_start": now}
            _buckets[key] = bucket

        remaining = max(0, RATE_LIMIT_REQUESTS - bucket["count"])
        # optimistic header before call_next; will update after
        response: Response

        if remaining <= 0:
            retry_after = int(RATE_LIMIT_WINDOW_SECONDS - (now - bucket["window_start"]))
            response = JSONResponse({"detail": "Rate limit exceeded"}, status_code=429)
            response.headers["Retry-After"] = str(max(0, retry_after))
            response.headers["X-Limit-Remaining"] = "0"
            return response

        bucket["count"] += 1
        response = await call_next(request)

        # set headers after
        remaining_after = max(0, RATE_LIMIT_REQUESTS - bucket["count"])
        response.headers["X-Limit-Remaining"] = str(remaining_after)
        return response