spb/hfmarketdata

1"""Site principal: the /charts page has no quota (ratelimit/middleware.py `is_site_request`).23Three cumulative criteria (`X-HFMD-Client`, `Sec-Fetch-Site`, same-host `Origin`/`Referer`) on data GETs →4principal `site`: no X-RateLimit-* headers, no 429 for quota reasons, max_rows 200 000, usage under `site`5(never under the visitor's IP), and a per-IP burst guard that is a DoS protection, not a product limit.6"""7from __future__ import annotations89from dataclasses import replace1011import pytest12from fastapi import FastAPI, Request   # module level: `from __future__ import annotations` + a fixture-local import13                                        # would leave `Request` unresolvable in the echo handler (→ query param, 422)1415SITE = {"X-HFMD-Client": "charts", "Sec-Fetch-Site": "same-origin", "Origin": "https://www.hfmarketdata.io"}16RL_HDRS = ("X-RateLimit-Limit-Requests", "X-RateLimit-Remaining-Requests", "X-RateLimit-Limit-Rows",17           "X-RateLimit-Remaining-Rows", "X-RateLimit-Reset")181920def ip(n: int) -> dict:21    return {"X-Forwarded-For": f"172.16.0.1, 10.0.0.{n}"}222324def site(n: int, **over) -> dict:25    return {**ip(n), **SITE, **over}262728def no_rl_headers(r) -> bool:29    return not any(h in r.headers for h in RL_HDRS)303132# ----------------------------------------------------------------------------------------------------33# Full app (synthetic lake)34# ----------------------------------------------------------------------------------------------------3536def test_site_requests_are_exempt_and_leave_the_ip_counters_untouched(client):37    for _ in range(31):                              # keyless would 429 after 3038        r = client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=5", headers=site(101))39        assert r.status_code == 200, r.text40        assert no_rl_headers(r)41        assert r.headers["X-Row-Count"] == "5"42    # the same IP, as a plain keyless client, still has its full hourly budget43    r = client.get("/v1/status", headers=ip(101))44    assert r.status_code == 200 and r.headers["X-RateLimit-Remaining-Requests"] == "29"45    assert r.headers["X-RateLimit-Remaining-Rows"] == "100000"464748@pytest.mark.parametrize("path", ["/v1/stock/tickers?limit=3", "/v1/futures/roots", "/v1/options/quarters",49                                  "/v1/bars/stock/AAPL?timeframe=1min&limit=10"])50def test_every_charts_data_path_is_covered(client, path):51    r = client.get(path, headers=site(102))52    assert r.status_code == 200, r.text53    assert no_rl_headers(r)545556def test_web_client_label_and_referer_fallback_are_accepted(client):57    r = client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=2", headers=site(103, **{"X-HFMD-Client": "web"}))58    assert r.status_code == 200 and no_rl_headers(r)59    hdrs = {**ip(103), "X-HFMD-Client": "charts", "Sec-Fetch-Site": "same-site", "Referer": "https://hfmarketdata.io/charts?s=AAPL"}60    r = client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=2", headers=hdrs)61    assert r.status_code == 200 and no_rl_headers(r)626364@pytest.mark.parametrize("missing", [65    {"Sec-Fetch-Site": None},                                        # no Sec-Fetch-Site → keyless66    {"Sec-Fetch-Site": "cross-site"},67    {"Origin": "https://evil.example"},                              # wrong host68    {"Origin": None},                                                # no Origin, no Referer69    {"X-HFMD-Client": None},70    {"X-HFMD-Client": "curl"},71])72def test_any_missing_criterion_falls_back_to_keyless(client, missing):73    hdrs = {**ip(104), **SITE}74    for k, v in missing.items():75        if v is None:76            hdrs.pop(k)77        else:78            hdrs[k] = v79    r = client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=2", headers=hdrs)80    assert r.status_code == 20081    assert r.headers["X-RateLimit-Limit-Requests"] == "30", dict(r.headers)828384def test_site_headers_on_a_non_data_path_or_a_non_get_are_ignored(client):85    r = client.get("/v1/status", headers=site(105))86    assert r.headers["X-RateLimit-Limit-Requests"] == "30"           # /v1/status is not a charts data path87    r = client.get("/v1/limits", headers=site(105))88    assert r.json()["data"]["principal"]["kind"] == "keyless"899091def test_bearer_key_wins_over_the_site_criteria(client, make_user):92    _, key, _, _ = make_user()93    r = client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=2", headers=site(106, Authorization=f"Bearer {key}"))94    assert r.status_code == 200 and r.headers["X-RateLimit-Limit-Requests"] == "120"959697def test_usage_is_recorded_under_site_not_under_the_ip(client):98    from ratelimit import middleware, usage99    client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=7", headers=site(107))100    live = usage._live_minutes()101    by_principal: dict[str, dict] = {}102    for (principal, _minute), fields in live.items():103        acc = by_principal.setdefault(principal, {})104        for k, v in fields.items():105            acc[k] = acc.get(k, 0) + v106    assert by_principal["site"]["requests"] >= 1 and by_principal["site"]["rows"] >= 7107    assert f"ip:{middleware.hash_ip('10.0.0.107')}" not in by_principal108109110def test_burst_guard_is_a_dedicated_429_without_ratelimit_headers(client, monkeypatch):111    from ratelimit import tiers112    monkeypatch.setattr(tiers, "SITE_BURST_PER_MINUTE_PER_IP", 3)113    for _ in range(3):114        assert client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=1", headers=site(108)).status_code == 200115    r = client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=1", headers=site(108))116    assert r.status_code == 429117    body = r.json()["error"]118    assert body["code"] == "RATE_LIMIT_EXCEEDED" and body["details"]["guard"] == "site_burst"119    assert "not a quota" in body["message"] and "free account" not in body["message"]120    assert int(r.headers["Retry-After"]) >= 1 and no_rl_headers(r)121    # another IP is unaffected, and the guarded IP keeps its keyless budget intact122    assert client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=1", headers=site(109)).status_code == 200123    assert client.get("/v1/status", headers=ip(108)).headers["X-RateLimit-Remaining-Requests"] == "29"124125126def test_burst_guard_fails_open_when_redis_is_down(client, monkeypatch):127    from ratelimit import redis_limiter as rl128    from ratelimit import tiers129    monkeypatch.setattr(tiers, "SITE_BURST_PER_MINUTE_PER_IP", 1)130    monkeypatch.setattr(rl, "throttle", lambda *a, **k: None)131    for _ in range(3):132        r = client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=1", headers=site(110))133        assert r.status_code == 200 and no_rl_headers(r)134135136# ----------------------------------------------------------------------------------------------------137# Minimal app: what the handler sees (max_rows, principal), localhost rule138# ----------------------------------------------------------------------------------------------------139140@pytest.fixture141def mini(app):142    """Tiny FastAPI app with only the rate-limit middleware, echoing request.state on a charts data path."""143    from fastapi.testclient import TestClient144145    from core.responses import json_response146    from ratelimit import middleware147148    m = FastAPI()149150    def echo(request: Request, asset: str, ticker: str):151        st = request.state152        return json_response({"principal": st.principal, "kind": st.principal_kind, "max_rows": st.max_rows,153                              "ratelimit": st.ratelimit})154155    m.add_api_route("/v1/bars/{asset}/{ticker}", echo, methods=["GET"])156    middleware.install(m)157    with TestClient(m) as c:158        yield c159160161def test_site_state_max_rows_and_principal(mini):162    r = mini.get("/v1/bars/stock/AAPL", headers=site(111))163    assert r.status_code == 200164    body = r.json()["data"]165    assert body == {"principal": "site", "kind": "site", "max_rows": 200_000,166                    "ratelimit": {"principal": "site", "kind": "site", "tier": "site", "exempt": True, "max_rows_per_request": 200_000}}167    r = mini.get("/v1/bars/stock/AAPL", headers=ip(111)).json()["data"]168    assert r["kind"] == "keyless" and r["max_rows"] == 5000169170171def test_localhost_origin_only_outside_production(mini, monkeypatch):172    from ratelimit import middleware173    dev = site(112, Origin="http://localhost:5173")174    assert mini.get("/v1/bars/stock/AAPL", headers=dev).json()["data"]["kind"] == "site"          # HFMD_ENV=test175    monkeypatch.setattr(middleware, "settings", replace(middleware.settings, env="production"))176    assert mini.get("/v1/bars/stock/AAPL", headers=dev).json()["data"]["kind"] == "keyless"177    assert mini.get("/v1/bars/stock/AAPL", headers=site(112)).json()["data"]["kind"] == "site"    # public host still fine178179180def test_is_site_request_unit():181    from starlette.datastructures import Headers182183    from ratelimit.middleware import is_site_data_path, is_site_request184    assert is_site_data_path("/v1/bars/stock/AAPL") and is_site_data_path("/v1/fx/tickers") and is_site_data_path("/v1/futures/ES/contracts")185    assert not is_site_data_path("/v1/status") and not is_site_data_path("/v1/fundamentals/screen") and not is_site_data_path("/v1/stock/tickers/x")186    scope = {"method": "GET", "path": "/v1/bars/stock/AAPL", "query_string": b""}187    good = Headers({"x-hfmd-client": "charts", "sec-fetch-site": "same-origin", "origin": "https://www.hfmarketdata.io"})188    assert is_site_request(scope, good)189    assert not is_site_request({**scope, "method": "POST"}, good)190    assert not is_site_request({**scope, "query_string": b"api_key=hfmd_live_x"}, good)191    assert not is_site_request(scope, Headers({"x-hfmd-client": "charts", "sec-fetch-site": "same-origin", "origin": "https://www.hfmarketdata.io.evil.example"}))192