stream: principal du middleware ratelimit, fin du repli permissif

stream/auth.py cherchait verify_api_key/resolve_api_key/authenticate_key dans
accounts.security — aucune n'existe — puis acceptait toute chaîne hfmd_live_<32> avec
un principal key:<sha256[:16]> : les lignes livrées étaient débitées à un compteur
Redis orphelin et _tier_for_principal retombait sur free.

* stream_ws lit ws.scope["state"] (principal, principal_kind, tier) posés par
  ratelimit.middleware._websocket ; keyless → JSON AUTH_REQUIRED + 4001 comme avant ;
  aucun contexte alors que HFMD_RATELIMIT est actif → refus 4401 sans handshake.
* authenticate(key) = recherche api_keys.key_hash (sha256(sel+clé)) + statut compte,
  utilisée seulement quand le middleware est désactivé ; plus aucune acceptation par
  forme de clé.
* Tests : principal = key:<id> (≠ hash), clé inconnue bien formée refusée,
  _resolve_principal sans contexte.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Simon-Pierre Boucher committed 18 days ago (Sep 6, 2026) parent 79e5dc3

3 changed files +106 −37

modified hfmarketdata/api/stream/auth.py +35 −33

		`@@ -1,24 +1,22 @@`
1		−"""API-key authentication for the WebSocket (query `api_key` or `Authorization: Bearer`).
	1	+"""Principal of a WebSocket connection.
2	2
3		−Delegates to the accounts module when it is installed (`accounts.security`), looking for one of
4		−`verify_api_key(key) -> principal \| None`, `resolve_api_key(key)` or `authenticate_key(key)`. Without it
5		−(module not deployed yet, tests) a key is accepted when it has the documented shape
6		−`hfmd_live_<32 base62>` and the principal is derived from its salted hash — no plaintext is ever kept.
	3	+The rate-limit middleware (`ratelimit.middleware.RateLimitMiddleware._websocket`) authenticates every `/v1/*`
	4	+socket before the handler runs and leaves the result in `scope["state"]`: `principal` (`key:<id>` or
	5	+`ip:<hash>`), `principal_kind` (`key` \| `keyless`), `tier`, `user_id`, `api_key_id`. The stream handler reads
	6	+those — it never re-authenticates and never derives a principal from the key text (a hashed pseudo-principal
	7	+would bill the delivered rows to a counter nobody reads and resolve to the free tier).
	8	+
	9	+`authenticate(key)` is the fallback for deployments running with `HFMD_RATELIMIT=0` (no middleware): the key is
	10	+looked up in the accounts tables (`sha256(salt + key)` = `api_keys.key_hash`) — an unknown key is refused,
	11	+whatever its shape.
7	12
8	13	Author: Simon-Pierre Boucher <contact@spboucher.ai>
9	14	"""
10	15	from __future__ import annotations
11	16
12		−import hashlib
13		−import importlib
14		−import re
	17	+from collections.abc import Mapping
15	18	from typing import Any
16	19
17		−from core.config import settings
18		−
19		−KEY_RE = re.compile(r"^hfmd_(live\|test)_[0-9A-Za-z]{32}$")
20		−_ACCOUNTS_FUNCS = ("verify_api_key", "resolve_api_key", "authenticate_key")
21		−
22	20
23	21	def extract_key(query_key: str \| None, authorization: str \| None) -> str \| None:
24	22	if query_key:
		`@@ -32,29 +30,33 @@ def extract_key(query_key: str \| None, authorization: str \| None) -> str \| None:`
32	30	return None
33	31
34	32
35		−def _principal_from_hash(key: str) -> str:
36		− h = hashlib.sha256((settings.key_hash_salt + key).encode()).hexdigest()
37		− return f"key:{h[:16]}"
	33	+def principal_from_scope(scope: Mapping[str, Any]) -> tuple[str \| None, str \| None, str \| None]:
	34	+ """(principal, kind, tier) left by the rate-limit middleware, or (None, None, None) when it did not run."""
	35	+ state = scope.get("state") or {}
	36	+ p = state.get("principal")
	37	+ return (str(p) if p else None), state.get("principal_kind"), state.get("tier")
38	38
39	39
40	40	def authenticate(key: str \| None) -> str \| None:
41		− """Return the principal (`key:<id>`) for a valid key, else None."""
	41	+ """`key:<id>` for an active key of a non-disabled account (SQLite lookup through the accounts module), else None."""
42	42	if not key:
43	43	return None
44	44	try:
45		− sec: Any = importlib.import_module("accounts.security")
46		− for name in _ACCOUNTS_FUNCS:
47		− fn = getattr(sec, name, None)
48		− if callable(fn):
49		− res = fn(key)
50		− if not res:
51		− return None
52		− if isinstance(res, str):
53		− return res if res.startswith("key:") else f"key:{res}"
54		− kid = getattr(res, "id", None) or (res.get("id") if isinstance(res, dict) else None)
55		− return f"key:{kid}" if kid is not None else _principal_from_hash(key)
56		− except ModuleNotFoundError:
57		− pass
58		− if KEY_RE.match(key):
59		− return _principal_from_hash(key)
60		− return None
	45	+ from sqlalchemy import select
	46	+
	47	+ from accounts.models import ApiKey, User
	48	+ from accounts.security import hash_key, looks_like_key
	49	+ from core.db import session
	50	+ except ModuleNotFoundError: # accounts module not deployed: nobody can be authenticated
	51	+ return None
	52	+ if not looks_like_key(key):
	53	+ return None
	54	+ with session() as s:
	55	+ row = s.execute(select(ApiKey.id, ApiKey.status, User.status.label("user_status"))
	56	+ .join(User, User.id == ApiKey.user_id).where(ApiKey.key_hash == hash_key(key))).first()
	57	+ if row is None or row.status != "active" or row.user_status == "disabled":
	58	+ return None
	59	+ return f"key:{int(row.id)}"
	60	+
	61	+
	62	+__all__ = ["extract_key", "principal_from_scope", "authenticate"]

modified hfmarketdata/api/stream/routes.py +20 −4

		`@@ -25,11 +25,12 @@ from typing import Any`
25	25	from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
26	26	from starlette.websockets import WebSocketState
27	27
	28	+from core.config import settings
28	29	from core.errors import CODES, docs_link
29	30	from core.responses import json_response
30	31
31	32	from . import accounting, broker
32		−from .auth import authenticate, extract_key
	33	+from .auth import authenticate, extract_key, principal_from_scope
33	34
34	35	log = logging.getLogger("hfmarketdata.stream")
35	36
		`@@ -39,7 +40,7 @@ HEARTBEAT_SECONDS = 20.0`
39	40	POLL_SECONDS = 0.25
40	41	MAX_CONNECTIONS_PER_KEY = 5
41	42	CONN_TTL_SECONDS = 3600
42		−CLOSE_AUTH, CLOSE_LIMIT, CLOSE_PROTOCOL = 4001, 4029, 4400
	43	+CLOSE_AUTH, CLOSE_LIMIT, CLOSE_PROTOCOL, CLOSE_UNAUTHENTICATED = 4001, 4029, 4400, 4401
43	44	CHANNELS = ("filings",)
44	45
45	46
		`@@ -90,16 +91,31 @@ def stream_info():`
90	91	"asyncapi": settings.public_url + "/docs/asyncapi.yaml"})
91	92
92	93
	94	+def _resolve_principal(ws: WebSocket, api_key: str \| None) -> tuple[str \| None, str \| None]:
	95	+ """(principal, tier) of the socket: what the rate-limit middleware stored in `scope["state"]`, or — only when
	96	+ the middleware is disabled (`HFMD_RATELIMIT=0`) — a database lookup of the key. Never a hash of the key."""
	97	+ principal, kind, tier = principal_from_scope(ws.scope)
	98	+ if principal is not None:
	99	+ return (principal if kind == "key" and principal.startswith("key:") else None), tier
	100	+ if settings.ratelimit_enabled:
	101	+ return None, None # the middleware should have run: refuse rather than guess
	102	+ return authenticate(extract_key(api_key, ws.headers.get("authorization"))), None
	103	+
	104	+
93	105	@router.websocket("/v1/stream")
94	106	async def stream_ws(ws: WebSocket, api_key: str \| None = Query(None)):
95		− key = extract_key(api_key, ws.headers.get("authorization"))
96		− principal = authenticate(key)
	107	+ if settings.ratelimit_enabled and principal_from_scope(ws.scope)[0] is None:
	108	+ # no authentication context at all (middleware not installed): deny the handshake
	109	+ await ws.close(code=CLOSE_UNAUTHENTICATED, reason="AUTH_REQUIRED")
	110	+ return
	111	+ principal, tier = _resolve_principal(ws, api_key)
97	112	await ws.accept()
98	113	if principal is None:
99	114	await ws.send_json(_err("AUTH_REQUIRED", "A valid API key is required for the stream: pass ?api_key=… or "
100	115	"Authorization: Bearer …. Create a free key on the dashboard."))
101	116	await ws.close(code=CLOSE_AUTH, reason="AUTH_REQUIRED")
102	117	return
	118	+ log.debug("stream connection %s tier=%s", principal, tier)
103	119	r = broker.get_redis()
104	120	conn_key = f"stream:conns:{principal}"
105	121	n = int(r.incr(conn_key))

modified tests/test_stream.py +51 −0

		`@@ -39,10 +39,61 @@ def test_bad_key_is_refused(client):`
39	39	assert e.code in (4401, 4001, 1008)
40	40
41	41
	42	+def _key_id(raw: str) -> int:
	43	+ from sqlalchemy import select
	44	+
	45	+ from accounts.models import ApiKey
	46	+ from accounts.security import hash_key
	47	+ from core.db import session
	48	+ with session() as s:
	49	+ return int(s.scalar(select(ApiKey.id).where(ApiKey.key_hash == hash_key(raw))))
	50	+
	51	+
	52	+def test_principal_is_the_key_id_not_a_hash(keys):
	53	+ """The stream must bill the same `key:<id>` principal as the HTTP limiter (a hashed pseudo-principal fed an
	54	+ orphan counter and resolved to the free tier)."""
	55	+ import hashlib
	56	+
	57	+ from core.config import settings
	58	+ from stream.auth import authenticate, principal_from_scope
	59	+ kid = _key_id(keys[0])
	60	+ assert authenticate(keys[0]) == f"key:{kid}"
	61	+ fake = "key:" + hashlib.sha256((settings.key_hash_salt + keys[0]).encode()).hexdigest()[:16]
	62	+ assert authenticate(keys[0]) != fake
	63	+ # well-formed but unknown keys are refused, whatever their shape
	64	+ assert authenticate("hfmd_live_" + "a" * 32) is None and authenticate("not-a-key") is None and authenticate(None) is None
	65	+ assert principal_from_scope({"state": {"principal": "key:7", "principal_kind": "key", "tier": "free"}}) == ("key:7", "key", "free")
	66	+ assert principal_from_scope({}) == (None, None, None)
	67	+
	68	+
	69	+def test_well_formed_unknown_key_is_refused(client):
	70	+ with pytest.raises(WebSocketDisconnect) as e:
	71	+ with client.websocket_connect("/v1/stream?api_key=hfmd_live_" + "Z" * 32) as ws:
	72	+ ws.receive_json()
	73	+ assert e.value.code in (4401, 4001)
	74	+
	75	+
	76	+def test_handler_refuses_without_middleware_context(keys):
	77	+ """`_resolve_principal`: middleware state wins; no state while the limiter is enabled → refuse, never guess."""
	78	+ import types
	79	+
	80	+ from core.config import settings
	81	+ from stream.routes import _resolve_principal
	82	+ ws = types.SimpleNamespace(scope={"state": {"principal": "key:42", "principal_kind": "key", "tier": "high_usage"}},
	83	+ headers={"authorization": f"Bearer {keys[0]}"})
	84	+ assert _resolve_principal(ws, None) == ("key:42", "high_usage")
	85	+ ws.scope = {"state": {"principal": "ip:abc", "principal_kind": "keyless", "tier": "keyless"}}
	86	+ assert _resolve_principal(ws, keys[0]) == (None, "keyless")
	87	+ ws.scope = {}
	88	+ expected = (None, None) if settings.ratelimit_enabled else (f"key:{_key_id(keys[0])}", None)
	89	+ assert _resolve_principal(ws, keys[0]) == expected
	90	+
	91	+
42	92	def test_subscribe_receive_and_accounting(keys, client, fundamentals_data):
43	93	from stream import accounting, broker
44	94	from stream.auth import authenticate
45	95	principal = authenticate(keys[0])
	96	+ assert principal == f"key:{_key_id(keys[0])}"
46	97	before = accounting.rows_charged(principal)
47	98	with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:
48	99	hello = ws.receive_json()
49	100