64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
"""In-memory rate limiting for the endpoints worth brute-forcing.
|
|
|
|
State is a plain dict, which is correct for the single-process uvicorn this
|
|
project ships (`server-worker.cmd`). Running multiple workers would give each
|
|
its own counters — move to a shared store before doing that.
|
|
|
|
Keys are the TCP peer address, never `X-Forwarded-For`: a spoofable header
|
|
would let a client mint a fresh bucket per request and defeat the limit. The
|
|
trade-off is that behind a reverse proxy every caller shares one bucket, which
|
|
is fine at this project's team scale.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from collections import defaultdict, deque
|
|
|
|
from fastapi import HTTPException, Request
|
|
|
|
|
|
class SlidingWindow:
|
|
def __init__(self, limit: int, window_seconds: float):
|
|
self.limit = limit
|
|
self.window = window_seconds
|
|
self._hits: dict[str, deque[float]] = defaultdict(deque)
|
|
self._lock = threading.Lock()
|
|
|
|
def check(self, key: str) -> None:
|
|
"""Record an attempt for `key`, or raise 429 when it is over budget."""
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
q = self._hits[key]
|
|
while q and now - q[0] > self.window:
|
|
q.popleft()
|
|
if len(q) >= self.limit:
|
|
retry = max(1, int(self.window - (now - q[0])) + 1)
|
|
raise HTTPException(
|
|
429,
|
|
f"too many attempts, retry in {retry}s",
|
|
headers={"Retry-After": str(retry)},
|
|
)
|
|
q.append(now)
|
|
|
|
def clear(self, key: str) -> None:
|
|
"""Forget a key's history — used after a successful attempt so a user
|
|
who fat-fingered the code is not locked out once they get it right."""
|
|
with self._lock:
|
|
self._hits.pop(key, None)
|
|
|
|
def reset(self) -> None:
|
|
with self._lock:
|
|
self._hits.clear()
|
|
|
|
|
|
def peer(request: Request) -> str:
|
|
return request.client.host if request.client else "unknown"
|
|
|
|
|
|
def limit_invite_attempts(request: Request, limiter: SlidingWindow) -> str:
|
|
"""Charge one attempt to the caller's address; return the bucket key."""
|
|
key = peer(request)
|
|
limiter.check(key)
|
|
return key
|