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
  best-effort by nature — behind a reverse proxy the client address is read from
  ``X-Forwarded-For``, which a caller can set to anything if the app is exposed
  directly. That is precisely why it is not the only key.

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 quart import request

# Failed sign-ins tolerated per account before it stops answering, and for how long.
# Ten is comfortably above a person mistyping a password and far below anything that
# makes a dictionary worth running.
ACCOUNT_LIMIT = 10
ACCOUNT_WINDOW_S = 15 * 60

# Wider, because one address is legitimately many people: a household, an office
# behind NAT, a phone on carrier-grade NAT.
ADDRESS_LIMIT = 50
ADDRESS_WINDOW_S = 15 * 60

# Registration is scarcer than a sign-in — it creates a row, and on a private
# instance the honest number of accounts anyone needs to make is one.
REGISTER_LIMIT = 5
REGISTER_WINDOW_S = 60 * 60

# Never let the bookkeeping become the denial of service: an attacker rotating a
# forged X-Forwarded-For 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.
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.
    """

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

    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()


sign_in_by_account = SlidingWindow(ACCOUNT_LIMIT, ACCOUNT_WINDOW_S)
sign_in_by_address = SlidingWindow(ADDRESS_LIMIT, ADDRESS_WINDOW_S)
register_by_address = SlidingWindow(REGISTER_LIMIT, REGISTER_WINDOW_S)


def client_address() -> str:
    """The caller's address, as well as it can be known.

    ``X-Forwarded-For`` is a list appended to by each hop, so the leftmost entry is
    the original client — and also the only entry a client can choose for itself.
    It is trusted here anyway, because the alternative behind a reverse proxy is to
    see the proxy's address for every request on earth and rate-limit the entire
    internet as one caller. The account-keyed limit is the one that holds when this
    one is lied to.
    """
    forwarded = request.headers.get("X-Forwarded-For", "")
    if forwarded:
        first = forwarded.split(",")[0].strip()
        if first:
            return first[:64]  # bounded: this becomes a dict key
    return (request.remote_addr or "unknown")[:64]


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()