uxoxo/eb2ab

1"""2Rate limiting middleware for REST API using sliding window algorithm.3"""4 5import os6import time7from collections import defaultdict, deque8from typing import Dict, Deque, Tuple9from fastapi import HTTPException, Request, status10 11 12# Rate limit storage: {identifier: deque of timestamps}13_rate_limit_storage: Dict[str, Deque[float]] = defaultdict(deque)14 15 16def get_rate_limit_config() -> Tuple[int, int]:17    """18    Get rate limit configuration from environment variables.19 20    Returns:21        Tuple of (max_requests, time_window_seconds)22    """23    max_requests = int(os.environ.get("MAX_REQUESTS_PER_HOUR", "100"))24    time_window = 3600  # 1 hour in seconds25    return max_requests, time_window26 27 28def get_rate_limit_identifier(request: Request, api_key: str = None) -> str:29    """30    Generate identifier for rate limiting.31 32    Uses API key if authenticated, otherwise uses IP address.33 34    Args:35        request: FastAPI request object36        api_key: Validated API key (if authenticated)37 38    Returns:39        Identifier string for rate limiting40    """41    if api_key and api_key != "no-key-configured":42        # Use API key for authenticated requests43        return f"key:{api_key}"44    else:45        # Use IP address for unauthenticated requests46        client_ip = request.client.host if request.client else "unknown"47        return f"ip:{client_ip}"48 49 50def check_rate_limit(identifier: str) -> bool:51    """52    Check if request is within rate limit using sliding window algorithm.53 54    Args:55        identifier: Unique identifier for rate limiting56 57    Returns:58        True if within limit, False if exceeded59    """60    max_requests, time_window = get_rate_limit_config()61    current_time = time.time()62 63    # Get request timestamps for this identifier64    timestamps = _rate_limit_storage[identifier]65 66    # Remove timestamps outside the time window67    while timestamps and current_time - timestamps[0] > time_window:68        timestamps.popleft()69 70    # Check if limit exceeded71    if len(timestamps) >= max_requests:72        return False73 74    # Add current request timestamp75    timestamps.append(current_time)76 77    return True78 79 80def get_remaining_requests(identifier: str) -> Tuple[int, int]:81    """82    Get remaining requests and reset time for identifier.83 84    Args:85        identifier: Unique identifier for rate limiting86 87    Returns:88        Tuple of (remaining_requests, seconds_until_reset)89    """90    max_requests, time_window = get_rate_limit_config()91    current_time = time.time()92 93    timestamps = _rate_limit_storage[identifier]94 95    # Remove old timestamps96    while timestamps and current_time - timestamps[0] > time_window:97        timestamps.popleft()98 99    remaining = max_requests - len(timestamps)100 101    # Calculate reset time (when oldest request expires)102    if timestamps:103        seconds_until_reset = int(time_window - (current_time - timestamps[0]))104    else:105        seconds_until_reset = int(time_window)106 107    return max(0, remaining), seconds_until_reset108 109 110def enforce_rate_limit(request: Request, api_key: str = None):111    """112    Enforce rate limit for request. Raises HTTPException if exceeded.113 114    Args:115        request: FastAPI request object116        api_key: Validated API key (if authenticated)117 118    Raises:119        HTTPException: 429 Too Many Requests if limit exceeded120    """121    identifier = get_rate_limit_identifier(request, api_key)122 123    if not check_rate_limit(identifier):124        remaining, reset_time = get_remaining_requests(identifier)125        max_requests, _ = get_rate_limit_config()126 127        raise HTTPException(128            status_code=status.HTTP_429_TOO_MANY_REQUESTS,129            detail=f"Rate limit exceeded. Maximum {max_requests} requests per hour.",130            headers={131                "X-RateLimit-Limit": str(max_requests),132                "X-RateLimit-Remaining": "0",133                "X-RateLimit-Reset": str(int(time.time() + reset_time)),134                "Retry-After": str(reset_time)135            }136        )137 138 139def add_rate_limit_headers(identifier: str) -> Dict[str, str]:140    """141    Generate rate limit headers for response.142 143    Args:144        identifier: Rate limit identifier145 146    Returns:147        Dictionary of rate limit headers148    """149    max_requests, _ = get_rate_limit_config()150    remaining, reset_time = get_remaining_requests(identifier)151 152    return {153        "X-RateLimit-Limit": str(max_requests),154        "X-RateLimit-Remaining": str(remaining),155        "X-RateLimit-Reset": str(int(time.time() + reset_time))156    }157 158 159def cleanup_old_entries():160    """161    Cleanup old entries from rate limit storage.162    Should be called periodically to prevent memory growth.163    """164    _, time_window = get_rate_limit_config()165    current_time = time.time()166 167    identifiers_to_remove = []168 169    for identifier, timestamps in _rate_limit_storage.items():170        # Remove old timestamps171        while timestamps and current_time - timestamps[0] > time_window:172            timestamps.popleft()173 174        # Mark empty identifiers for removal175        if not timestamps:176            identifiers_to_remove.append(identifier)177 178    # Remove empty identifiers179    for identifier in identifiers_to_remove:180        del _rate_limit_storage[identifier]181