python_Sklyarenko/Lab1/app/rate_limit.py - master - ivanstrike/tasker

python_Sklyarenko/Lab1/app/rate_limit.py
51 строка2 KB
ivanstrike
Lab3
13 дек 2025, 11:15
13 дек 2025, 11:1515c0fa5
import time
import os
from fastapi import Request, HTTPException
from typing import Tuple

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

# In-memory simplistic store: { key: (window_start_ts, count) }
# For demo purposes only.
_store: dict[str, Tuple[float, int]] = {}

def rate_limit_key(request: Request, user_id: int | None) -> str:
    if user_id is not None:
        return f"user:{user_id}"
    # fallback: IP
    ip = request.client.host if request.client else "unknown"
    return f"ip:{ip}"

async def enforce_rate_limit(request: Request, user_id: int | None):
    key = rate_limit_key(request, user_id)
    now = time.time()
    window_start, count = _store.get(key, (now, 0))

    # reset window if expired
    if now - window_start >= RATE_LIMIT_WINDOW_SECONDS:
        window_start, count = now, 0

    if count >= RATE_LIMIT_REQUESTS:
        retry_after = RATE_LIMIT_WINDOW_SECONDS - int(now - window_start)
        # set headers via state for later injection
        request.state.rate_limit = {
            "remaining": 0,
            "retry_after": max(1, retry_after)
        }
        raise HTTPException(status_code=429, detail="Rate limit exceeded")

    count += 1
    _store[key] = (window_start, count)
    remaining = max(0, RATE_LIMIT_REQUESTS - count)
    request.state.rate_limit = {
        "remaining": remaining,
        "retry_after": RATE_LIMIT_WINDOW_SECONDS - int(now - window_start)
    }

def set_rate_limit_headers(request: Request, response):
    data = getattr(request.state, "rate_limit", None)
    if data:
        response.headers["X-Limit-Remaining"] = str(data["remaining"])
        # Only include Retry-After when limit is exceeded or useful
        response.headers["Retry-After"] = str(max(0, data["retry_after"]))