inkwell/tests/test_ratelimit.py

"""The credential throttle. DB-free, like the rest of this suite.

The window is exercised with an injected clock, so nothing sleeps. The routes are
exercised only as far as they get without a database: a throttled request returns
429 before any session is opened.
"""
import time

import pytest

from inkwell import ratelimit
from inkwell.settings import live
from inkwell.app import create_app
from inkwell.ratelimit import SlidingWindow


def window(limit: int, window_s: float) -> SlidingWindow:
    """A fixed-value window. The real ones read their numbers from the settings cache
    so an admin's change applies immediately; these tests are about the counting, not
    about where the numbers come from."""
    return SlidingWindow(lambda: limit, lambda: window_s)


@pytest.fixture(autouse=True)
def _clean_counters():
    ratelimit.reset_all()
    yield
    ratelimit.reset_all()


@pytest.fixture
def app():
    return create_app()


def test_under_the_limit_is_not_blocked():
    w = window(3, 60)
    for i in range(3):
        assert w.retry_after("k", now=i) is None
        w.record("k", now=i)
    assert w.retry_after("k", now=3) is not None


def test_window_slides_rather_than_resetting():
    w = window(2, 60)
    w.record("k", now=0)
    w.record("k", now=30)
    assert w.retry_after("k", now=31) is not None
    # The 0s hit falls out at t=60, which frees exactly one slot — the 30s hit is
    # still inside the window, so this is a slide and not a reset.
    assert w.retry_after("k", now=61) is None
    w.record("k", now=61)
    assert w.retry_after("k", now=62) is not None


def test_retry_after_points_past_the_oldest_hit():
    w = window(1, 100)
    w.record("k", now=10)
    wait = w.retry_after("k", now=40)
    # The hit at t=10 leaves the window at t=110, i.e. 70s away. Rounded up, never
    # under-reported — a client that waits exactly this long must not be refused
    # again.
    assert wait is not None
    assert 70 <= wait <= 72
    assert w.retry_after("k", now=40 + wait) is None


def test_keys_are_counted_separately():
    w = window(1, 60)
    w.record("a", now=0)
    assert w.retry_after("a", now=1) is not None
    assert w.retry_after("b", now=1) is None


def test_forget_clears_one_key():
    w = window(1, 60)
    w.record("a", now=0)
    w.record("b", now=0)
    w.forget("a")
    assert w.retry_after("a", now=1) is None
    assert w.retry_after("b", now=1) is not None


def test_bucket_count_is_bounded(monkeypatch):
    # An attacker rotating a forged X-Forwarded-For must not be able to grow this
    # dict without limit — the limiter cannot become the exhaustion it prevents.
    monkeypatch.setattr(ratelimit, "MAX_BUCKETS", 8)
    w = window(5, 60)
    for i in range(50):
        w.record(f"addr-{i}", now=i)
    assert len(w._hits) <= 8


async def test_login_starts_refusing(app):
    client = app.test_client()
    body = {"email": "someone@example.com", "password": "wrong-password"}
    # Pre-load the account's counter to its limit rather than posting that many
    # times: every real attempt would need a database to reach the password check.
    #
    # On the real clock: hits stamped at t=0..9 would be pruned as stale the moment
    # the route reads `time.monotonic()`.
    now = time.monotonic()
    for _ in range(live("signin_limit_per_account")):
        ratelimit.sign_in_by_account.record("someone@example.com", now=now)
    resp = await client.post("/api/auth/login", json=body)
    assert resp.status_code == 429
    assert resp.headers.get("Retry-After")
    # The refusal says nothing about whether that account exists.
    assert "someone@example.com" not in (await resp.get_data(as_text=True))


async def test_device_login_shares_the_account_counter(app):
    client = app.test_client()
    now = time.monotonic()
    for _ in range(live("signin_limit_per_account")):
        ratelimit.sign_in_by_account.record("someone@example.com", now=now)
    resp = await client.post(
        "/api/auth/device-login",
        json={"email": "someone@example.com", "password": "wrong-password"},
    )
    # Same budget as /login — otherwise guessing just moves to the route that hands
    # out a long-lived bearer token.
    assert resp.status_code == 429


async def test_register_is_throttled_by_address(app):
    client = app.test_client()
    now = time.monotonic()
    for _ in range(live("register_limit_per_address")):
        ratelimit.register_by_address.record("203.0.113.9", now=now)
    resp = await client.post(
        "/api/auth/register",
        json={"email": "new@example.com", "password": "a-long-enough-password"},
        # One entry, so with the default single trusted hop this IS the address the
        # limiter keys on. The forged-prefix cases live in test_proxy.py.
        headers={"X-Forwarded-For": "203.0.113.9"},
    )
    assert resp.status_code == 429