spb/hfmarketdata

1"""Security hardening of the rate-limit layer: X-Forwarded-For handling (A1), per-endpoint auth throttles and2fail-closed behaviour (A5/A14), bounded key cache + cross-worker invalidation (A12/A20), key-lookup throttle."""3from __future__ import annotations45from dataclasses import replace67from starlette.datastructures import Headers89JSON = {"Content-Type": "application/json"}101112def scope_with(client_host: str | None, **headers) -> tuple[dict, Headers]:13    raw = [(k.lower().encode(), v.encode()) for k, v in headers.items()]14    scope = {"type": "http", "headers": raw, "client": (client_host, 1234) if client_host else None}15    return scope, Headers(scope=scope)161718def test_client_ip_takes_the_hop_appended_by_the_proxy(app):19    """A1: the FIRST hop is attacker-controlled; ngrok appends the real peer at the END."""20    from ratelimit import middleware as mw21    scope, h = scope_with("127.0.0.1", **{"X-Forwarded-For": "6.6.6.6, 198.51.100.23"})22    assert mw.client_ip(scope, h) == "198.51.100.23"23    scope, h = scope_with("10.0.0.2", **{"X-Forwarded-For": "198.51.100.23"})24    assert mw.client_ip(scope, h) == "198.51.100.23"25    # two trusted proxies → second from the end26    scope, h = scope_with("10.0.0.2", **{"X-Forwarded-For": "6.6.6.6, 198.51.100.23, 10.0.0.9"})27    assert mw.client_ip(scope, h) == "10.0.0.9"28    old = mw.settings29    mw.settings = replace(old, trusted_proxy_hops=2)30    try:31        assert mw.client_ip(scope, h) == "198.51.100.23"32        # hops=0 → the header is ignored entirely33        mw.settings = replace(old, trusted_proxy_hops=0)34        assert mw.client_ip(scope, h) == "10.0.0.2"35    finally:36        mw.settings = old37    # direct peer is a public address (no proxy in front) → X-Forwarded-For is not trusted38    scope, h = scope_with("8.8.8.8", **{"X-Forwarded-For": "6.6.6.6"})       # (203.0.113.x is "private" for ipaddress)39    assert mw.client_ip(scope, h) == "8.8.8.8"40    # X-Real-IP only when trusted41    scope, h = scope_with("127.0.0.1", **{"X-Real-IP": "198.51.100.77"})42    assert mw.client_ip(scope, h) == "198.51.100.77"43    scope, h = scope_with("8.8.8.8", **{"X-Real-IP": "198.51.100.77"})44    assert mw.client_ip(scope, h) == "8.8.8.8"45    scope, h = scope_with(None)46    assert mw.client_ip(scope, h) == "unknown"474849def test_spoofed_first_hop_cannot_escape_the_keyless_quota(client):50    """A1 (live reproduction): varying the first hop while ngrok appends the same peer = one principal."""51    for i in range(30):52        r = client.get("/v1/status", headers={"X-Forwarded-For": f"1.2.3.{i}, 198.51.100.200"})53        assert r.status_code == 200, i54    r = client.get("/v1/status", headers={"X-Forwarded-For": "9.9.9.9, 198.51.100.200"})55    assert r.status_code == 42956    # …and the auth throttle cannot be bypassed either57    for i in range(20):58        r = client.post("/v1/auth/login", json={"email": "nobody@example.com", "password": "x"},59                        headers={**JSON, "X-Forwarded-For": f"1.2.3.{i}, 198.51.100.201"})60        assert r.status_code == 401, i61    r = client.post("/v1/auth/login", json={"email": "nobody@example.com", "password": "x"},62                    headers={**JSON, "X-Forwarded-For": "7.7.7.7, 198.51.100.201"})63    assert r.status_code == 429646566def test_auth_throttle_is_per_endpoint(client):67    """A5/A14: signup 5/h, forgot 5/h, login 20/h, verify 30/h — separate buckets per IP."""68    h = {**JSON, "X-Forwarded-For": "198.51.100.10"}69    for i in range(5):70        assert client.post("/v1/auth/signup", json={"email": "bad", "name": "x", "password": "long-enough-password"}, headers=h).status_code == 400, i71    r = client.post("/v1/auth/signup", json={"email": "bad", "name": "x", "password": "long-enough-password"}, headers=h)72    assert r.status_code == 429 and r.json()["error"]["details"]["endpoint"] == "signup" and "Retry-After" in r.headers73    # the login bucket of the same IP is untouched74    assert client.post("/v1/auth/login", json={"email": "nobody@example.com", "password": "x"}, headers=h).status_code == 40175    for i in range(5):76        assert client.post("/v1/auth/forgot", json={"email": f"n{i}@example.com"}, headers=h).status_code == 20277    assert client.post("/v1/auth/forgot", json={"email": "n9@example.com"}, headers=h).status_code == 42978    for i in range(30):79        assert client.post("/v1/auth/verify", json={"token": "nope"}, headers=h).status_code == 400, i80    assert client.post("/v1/auth/verify", json={"token": "nope"}, headers=h).status_code == 42981    # logout is never throttled82    for _ in range(40):83        assert client.post("/v1/auth/logout", json={}, headers=h).status_code == 200848586def test_auth_fails_closed_when_redis_is_down_but_data_fails_open(client, monkeypatch):87    """A5: no Redis → login/signup/forgot answer 429 (Retry-After 30); data endpoints keep working."""88    import redis as redis_lib8990    from ratelimit import redis_limiter as rl9192    class Broken:93        def __getattr__(self, name):94            def _raise(*a, **k):95                raise redis_lib.exceptions.ConnectionError("down")96            return _raise9798    monkeypatch.setattr(rl, "_client", Broken())99    monkeypatch.setattr(rl, "_sha", None)100    monkeypatch.setattr(rl, "_down_until", 0.0)101    h = {**JSON, "X-Forwarded-For": "198.51.100.11"}102    for path, body in (("/v1/auth/login", {"email": "a@example.com", "password": "x"}),103                       ("/v1/auth/signup", {"email": "a@example.com", "name": "x", "password": "long-enough-password"}),104                       ("/v1/auth/forgot", {"email": "a@example.com"})):105        r = client.post(path, json=body, headers=h)106        assert r.status_code == 429 and r.headers["Retry-After"] == "30", path107    assert client.post("/v1/auth/logout", json={}, headers=h).status_code == 200108    assert client.post("/v1/auth/verify", json={"token": "nope"}, headers=h).status_code == 400   # not in the fail-closed set109    r = client.get("/v1/status", headers={"X-Forwarded-For": "198.51.100.11"})110    assert r.status_code == 200 and "X-RateLimit-Limit-Requests" not in r.headers111112113def test_key_cache_is_bounded_and_invalidated_across_workers(client, make_user, monkeypatch):114    """A12/A20: LRU 10 000 with TTL; a Redis generation bump makes every worker miss."""115    from ratelimit import middleware as mw116    from ratelimit import redis_limiter as rl117    lru = mw._LRU(maxsize=3, ttl_s=60)118    for i in range(5):119        lru.put(f"k{i}", i, generation=1)120    assert len(lru) == 3 and lru.get("k0", 1) == (False, None) and lru.get("k4", 1) == (True, 4)121    assert lru.get("k4", 2) == (False, None)            # generation changed → miss122    lru.put("neg", None, generation=None)123    assert lru.get("neg", 7) == (True, None)             # negatives are cached, None generation = not tracked124    assert mw._key_cache.maxsize == 10_000 and mw._key_cache.ttl_s == 60125126    _, key, _, _ = make_user()127    calls = {"n": 0}128    real = mw._lookup_key129130    def counting(h, ip_hash=None):131        calls["n"] += 1132        return real(h, ip_hash)133134    monkeypatch.setattr(mw, "_lookup_key", counting)135    for _ in range(3):136        assert client.get("/v1/status", headers={"Authorization": f"Bearer {key}"}).status_code == 200137    assert calls["n"] == 1                               # cached138    rl.bump_keys_version()                               # what another worker does after a revoke / tier change139    assert client.get("/v1/status", headers={"Authorization": f"Bearer {key}"}).status_code == 200140    assert calls["n"] == 2141    # unknown keys are cached negatively: 3 requests, 1 lookup142    unknown = "hfmd_live_" + "z" * 32143    for _ in range(3):144        assert client.get("/v1/status", headers={"Authorization": f"Bearer {unknown}"}).status_code == 401145    assert calls["n"] == 3146147148def test_unknown_key_lookups_are_throttled_before_sqlite(client, monkeypatch):149    """A12: brute-forcing keys hits a per-IP throttle (120/min) before the database."""150    from ratelimit import middleware as mw151    from ratelimit.tiers import KEY_LOOKUPS_PER_MINUTE_PER_IP152    calls = {"n": 0}153    real = mw._lookup_key154155    def counting(h, ip_hash=None):156        calls["n"] += 1157        return real(h, ip_hash)158159    monkeypatch.setattr(mw, "_lookup_key", counting)160    h = {"X-Forwarded-For": "198.51.100.12"}161    for i in range(KEY_LOOKUPS_PER_MINUTE_PER_IP):162        r = client.get("/v1/status", headers={**h, "Authorization": f"Bearer hfmd_live_{i:032d}"})163        assert r.status_code == 401, i164    r = client.get("/v1/status", headers={**h, "Authorization": "Bearer hfmd_live_" + "q" * 32})165    assert r.status_code == 429 and "Retry-After" in r.headers166    assert calls["n"] == KEY_LOOKUPS_PER_MINUTE_PER_IP   # the 121st never reached SQLite167    # another IP is unaffected168    assert client.get("/v1/status", headers={"X-Forwarded-For": "198.51.100.13", "Authorization": "Bearer hfmd_live_" + "q" * 32}).status_code == 401169170171def test_new_error_codes_are_documented(client):172    spec = client.get("/openapi.json").json()173    codes = spec["components"]["schemas"]["Error"]["properties"]["error"]["properties"]["code"]["enum"]174    assert {"ACCOUNT_LOCKED", "SESSION_REQUIRED", "LAST_ADMIN"} <= set(codes)175    login = spec["paths"]["/v1/auth/login"]["post"]176    assert "423" in login["responses"] and "ACCOUNT_LOCKED" in login["responses"]["423"]["description"]177    assert "SESSION_REQUIRED" in spec["paths"]["/v1/me/keys"]["post"]["responses"]["403"]["description"]178    assert spec["paths"]["/v1/auth/verify"]["post"]["summary"] and spec["paths"]["/v1/auth/verify"]["get"]["responses"]["302"]179    for p in ("/v1/me/password", "/v1/me/email", "/v1/me/sessions/revoke-all", "/v1/me/limits", "/v1/me/usage.csv", "/v1/admin/users/{user_id}"):180        assert p in spec["paths"], p181    assert "delete" in spec["paths"]["/v1/me"] and "delete" in spec["paths"]["/v1/admin/users/{user_id}"]182