`services/common/rate_limit.py`

54 lines
"""In-memory sliding-window rate limiter (stdlib).

Replaces the search API's `require_rate_limit_STUB` with a real control. This is the
single-process reference; production keys the same algorithm off Redis (atomic INCR + TTL,
or a sorted-set sliding window) so limits hold across many API instances.

The clock is INJECTABLE so time-based behavior is tested deterministically (no sleeps).

SECURITY: the limiter key MUST be derived from the TRUSTED-proxy client IP (+ optional
session), never a raw client-supplied X-Forwarded-For header (that is spoofable and would
let an attacker mint unlimited buckets). The caller is responsible for passing a trusted key.
"""
from __future__ import annotations

import time
from collections import defaultdict, deque
from typing import Callable


class SlidingWindowRateLimiter:
    def __init__(self, max_requests: int, window_seconds: float,
                 clock: Callable[[], float] = time.monotonic) -> None:
        if max_requests < 1 or window_seconds <= 0:
            raise ValueError("max_requests>=1 and window_seconds>0 required")
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        self._clock = clock
        self._hits: dict[str, deque[float]] = defaultdict(deque)

    def allow(self, key: str) -> bool:
        """Record a request for `key`; return False if it exceeds the window budget."""
        now = self._clock()
        cutoff = now - self.window_seconds
        bucket = self._hits[key]
        while bucket and bucket[0] <= cutoff:
            bucket.popleft()
        if len(bucket) >= self.max_requests:
            return False
        bucket.append(now)
        return True

    def retry_after(self, key: str) -> float:
        """Seconds until the oldest in-window hit for `key` ages out (for a 429 header)."""
        bucket = self._hits.get(key)
        if not bucket:
            return 0.0
        return max(0.0, self.window_seconds - (self._clock() - bucket[0]))

    def reset(self, key: str | None = None) -> None:
        if key is None:
            self._hits.clear()
        else:
            self._hits.pop(key, None)