inkwell/src/thoughtsync/ratelimit.py

"""Throttling for the endpoints that a public deployment leaves exposed.

Only the credential endpoints are rate-limited: login, register, and the native
device-link exchange. Everything else already needs a session or a device token to
reach, so an attacker has to get through one of these three first.

## Why in-process is enough here, and where that stops being true

State lives in module-level dicts, so it is per-process. That is correct for how
this image actually serves — one hypercorn worker (see the Dockerfile, and the
same assumption the trash sweeper documents in app.py). If that ever gains
``--workers N``, each worker would keep its own counters and the effective limit
would multiply by N; the fix then is a shared store (the Postgres connection is
already there), not a bigger number here.

## Two keys, on purpose

Every attempt is counted against BOTH the account being tried and the address it
came from, and either one can refuse it:

- **The account** is the key that matters, and the key that cannot be forged. It
  is what stops credential stuffing against one known email, no matter how many
  addresses the attempts arrive from.
- **The address** bounds the damage from one source spraying many accounts. It is
  read from ``X-Forwarded-For``, counting in from the RIGHT by
  the **Trusted proxy hops** setting so that only entries our own proxies wrote are
  believed — a forged header lands to the left of those and is never selected. It is
  still the weaker of the two keys, because it depends on that setting matching the
  deployment; the account key depends on nothing.

Counting is by failure for the sign-in routes and by attempt for registration: a
correct password should never move someone closer to being locked out, but every
registration is a row in the users table whether it succeeds or not.
"""
from __future__ import annotations

import time
from collections import deque
from collections.abc import Callable

from .settings import live


# Never let the bookkeeping become the denial of service: an attacker rotating an
# address could otherwise mint an unbounded number of buckets. Well above any real
# deployment's distinct-caller count, so a legitimate instance never reaches it; when
# it is reached the oldest buckets are dropped, which at worst forgives some attempts.
#
# Not a setting: it protects the limiter from itself rather than the app from a
# caller, and there is no operator judgment to apply to it.
MAX_BUCKETS = 10_000


class SlidingWindow:
    """Counts events per key over a trailing window.

    Sliding rather than a fixed window because a fixed one lets twice the limit
    through across a boundary — 10 at 14:59 and 10 at 15:00 — which for a login
    limiter is the difference between the number meaning something and not.

    The limit and window are SUPPLIERS, not values, so an admin saving a new number in
    Settings takes effect on the next attempt instead of the next deploy. They are read
    per call, which is a dict lookup — the settings cache never touches the database.
    """

    def __init__(self, limit: Callable[[], int], window_s: Callable[[], float]) -> None:
        self._limit = limit
        self._window_s = window_s
        self._hits: dict[str, deque[float]] = {}

    @property
    def limit(self) -> int:
        return self._limit()

    @property
    def window_s(self) -> float:
        return self._window_s()

    def _prune(self, key: str, now: float) -> deque[float]:
        hits = self._hits.get(key)
        if hits is None:
            hits = deque()
            # Insertion-ordered, so the first key is the least recently created.
            if len(self._hits) >= MAX_BUCKETS:
                self._hits.pop(next(iter(self._hits)), None)
            self._hits[key] = hits
        cutoff = now - self.window_s
        while hits and hits[0] <= cutoff:
            hits.popleft()
        return hits

    def retry_after(self, key: str, now: float | None = None) -> int | None:
        """Seconds until `key` may try again, or None while it is still under the
        limit. Read-only — it does not count as an attempt."""
        now = time.monotonic() if now is None else now
        hits = self._prune(key, now)
        if len(hits) < self.limit:
            return None
        # The window frees up when its OLDEST hit falls out of it.
        return max(1, int(hits[0] + self.window_s - now) + 1)

    def record(self, key: str, now: float | None = None) -> None:
        now = time.monotonic() if now is None else now
        self._prune(key, now).append(now)

    def forget(self, key: str) -> None:
        """Drop a key's history. Used after a successful sign-in, so someone who
        fumbled a password twice and then got it right starts clean rather than
        carrying those two for the next quarter of an hour."""
        self._hits.pop(key, None)

    def clear(self) -> None:
        self._hits.clear()


def _minutes(key: str) -> Callable[[], float]:
    return lambda: float(live(key)) * 60.0


sign_in_by_account = SlidingWindow(
    lambda: live("signin_limit_per_account"), _minutes("signin_window_minutes")
)
sign_in_by_address = SlidingWindow(
    lambda: live("signin_limit_per_address"), _minutes("signin_window_minutes")
)
register_by_address = SlidingWindow(
    lambda: live("register_limit_per_address"), _minutes("register_window_minutes")
)


def reset_all() -> None:
    """Drop every counter. For tests — nothing in the app calls this."""
    for window in (sign_in_by_account, sign_in_by_address, register_by_address):
        window.clear()