Coverage for node / src / stigmem_node / rate_limit.py: 92%
111 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-06-18 05:34 +0000
1"""Per-principal token-bucket quota middleware — spec §22.4.
3Replaces the legacy sliding-window implementation. Each principal
4(entity_uri resolved from the Bearer token) gets one token-bucket per quota
5dimension. Buckets are stored in the ``quota_buckets`` table and refilled
6lazily on each request.
8Dimension → endpoint mapping (spec §22.4.1):
9 fact_write POST /v1/facts, DELETE /v1/facts/*
10 fact_read GET /v1/facts/*, GET /v1/recall*
11 token_issue POST /v1/federation/capability-tokens
12 admin_action /v1/admin/*
13 audit_export GET /v1/admin/audit*
15Exemptions (same as legacy):
16 /v1/federation/ peer requests (use peer-token auth, not user API keys)
17 Requests without a Bearer token
19Backward-compat settings bridges:
20 STIGMEM_RATE_LIMIT_WRITE_PER_HOUR → fact_write burst capacity (default 100)
21 STIGMEM_RATE_LIMIT_READ_PER_HOUR → fact_read burst capacity (default 500)
22 0 on either setting disables rate limiting entirely.
23"""
25from __future__ import annotations
27import math
28import time
29from typing import Any
31from starlette.middleware.base import BaseHTTPMiddleware
32from starlette.requests import Request
33from starlette.responses import JSONResponse
35from .db import db
36from .settings import settings
38# ---------------------------------------------------------------------------
39# Default quota ceilings (spec §22.4.2)
40# ---------------------------------------------------------------------------
42_SPEC_DEFAULTS: dict[str, tuple[float, float]] = {
43 # dimension: (capacity, rate_per_second)
44 "fact_write": (100.0, 10.0),
45 "fact_read": (500.0, 50.0),
46 "token_issue": (20.0, 1 / 3),
47 "federation_pull": (30.0, 0.5),
48 "admin_action": (10.0, 1 / 6),
49 "subscription_event": (200.0, 20.0),
50 "audit_export": (10_000.0, 167.0),
51}
54def _capacity_for(dimension: str) -> float:
55 """Return burst capacity, honoring legacy per-hour settings for fact_write/read."""
56 if dimension == "fact_write":
57 cap = settings.rate_limit_write_per_hour
58 return float(cap) if cap > 0 else _SPEC_DEFAULTS["fact_write"][0]
59 if dimension == "fact_read": 59 ↛ 62line 59 didn't jump to line 62 because the condition on line 59 was always true
60 cap = settings.rate_limit_read_per_hour
61 return float(cap) if cap > 0 else _SPEC_DEFAULTS["fact_read"][0]
62 return _SPEC_DEFAULTS.get(dimension, (100.0, 1.0))[0]
65def _rate_for(dimension: str) -> float:
66 """Return refill rate (tokens/second)."""
67 if dimension == "fact_write":
68 cap = settings.rate_limit_write_per_hour
69 c = float(cap) if cap > 0 else _SPEC_DEFAULTS["fact_write"][0]
70 return c / 3600.0
71 if dimension == "fact_read": 71 ↛ 75line 71 didn't jump to line 75 because the condition on line 71 was always true
72 cap = settings.rate_limit_read_per_hour
73 c = float(cap) if cap > 0 else _SPEC_DEFAULTS["fact_read"][0]
74 return c / 3600.0
75 return _SPEC_DEFAULTS.get(dimension, (100.0, 1.0))[1]
78# ---------------------------------------------------------------------------
79# Endpoint → dimension routing
80# ---------------------------------------------------------------------------
83def _dimension(path: str, method: str) -> str | None:
84 """Return the quota dimension for this request, or None to skip quota."""
85 m = method.upper()
86 if path.startswith("/v1/admin/audit"): 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true
87 return "audit_export" if m == "GET" else "admin_action"
88 if path.startswith("/v1/admin/"):
89 return "admin_action"
90 if path.startswith("/v1/federation/capability-tokens") and m == "POST": 90 ↛ 91line 90 didn't jump to line 91 because the condition on line 90 was never true
91 return "token_issue"
92 if path.startswith("/v1/recall") or (path.startswith("/v1/facts") and m == "GET"):
93 return "fact_read"
94 if path.startswith("/v1/facts") and m in {"POST", "PUT", "PATCH", "DELETE"}:
95 return "fact_write"
96 # F-AVAIL-3: default coverage — every other authenticated API call still
97 # counts against a quota dimension so no endpoint is left unbounded.
98 if m == "GET":
99 return "fact_read"
100 if m in {"POST", "PUT", "PATCH", "DELETE"}:
101 return "fact_write"
102 return None
105# ---------------------------------------------------------------------------
106# Token-bucket check (SQLite upsert for atomic read-modify-write)
107# ---------------------------------------------------------------------------
110def _check_and_consume(
111 entity_uri: str,
112 tenant_id: str,
113 dimension: str,
114) -> tuple[bool, float]:
115 """Refill and consume one token. Returns (allowed, retry_after_seconds)."""
116 now = time.time()
117 capacity = _capacity_for(dimension)
118 rate = _rate_for(dimension)
120 with db() as conn:
121 conn.execute("BEGIN IMMEDIATE")
122 row = conn.execute(
123 "SELECT tokens, last_refill FROM quota_buckets"
124 " WHERE entity_uri=? AND tenant_id=? AND dimension=?",
125 (entity_uri, tenant_id, dimension),
126 ).fetchone()
128 if row is None:
129 tokens, last_refill = capacity, now
130 else:
131 tokens, last_refill = row["tokens"], row["last_refill"]
132 elapsed = max(0.0, now - last_refill)
133 tokens = min(capacity, tokens + elapsed * rate)
134 last_refill = now
136 if tokens >= 1.0:
137 new_tokens = tokens - 1.0
138 conn.execute(
139 """INSERT INTO quota_buckets (entity_uri, tenant_id, dimension, tokens, last_refill)
140 VALUES (?,?,?,?,?)
141 ON CONFLICT(entity_uri, tenant_id, dimension)
142 DO UPDATE SET tokens=excluded.tokens, last_refill=excluded.last_refill""",
143 (entity_uri, tenant_id, dimension, new_tokens, last_refill),
144 )
145 return True, 0.0
147 # Bucket empty — persist refilled (but not consumed) state
148 conn.execute(
149 """INSERT INTO quota_buckets (entity_uri, tenant_id, dimension, tokens, last_refill)
150 VALUES (?,?,?,?,?)
151 ON CONFLICT(entity_uri, tenant_id, dimension)
152 DO UPDATE SET tokens=excluded.tokens, last_refill=excluded.last_refill""",
153 (entity_uri, tenant_id, dimension, tokens, last_refill),
154 )
155 # retry_after: seconds until one token is earned
156 retry_after = (1.0 - tokens) / rate if rate > 0 else 1.0
157 return False, retry_after
160# ---------------------------------------------------------------------------
161# Identity lookup (lightweight — only entity_uri + tenant_id needed)
162# ---------------------------------------------------------------------------
164_HASH_CACHE: dict[
165 str, tuple[tuple[str, str, str | None], float]
166] = {} # raw-key fingerprint → (result, cached_at)
167_CACHE_TTL = 60.0
170def _lookup_principal(raw_key: str) -> tuple[str, str, str | None] | None:
171 """Return (entity_uri, tenant_id, oidc_sub) for the raw Bearer token, or None."""
172 import hashlib as _hl
174 fingerprint = _hl.sha256(raw_key.encode()).hexdigest()
175 if fingerprint in _HASH_CACHE:
176 result, cached_at = _HASH_CACHE[fingerprint]
177 if time.time() - cached_at < _CACHE_TTL:
178 return result
179 del _HASH_CACHE[fingerprint]
181 from fastapi import HTTPException
183 from .auth import lookup_principal
185 try:
186 principal = lookup_principal(raw_key)
187 except HTTPException:
188 return None
189 if principal is None:
190 return None
191 _HASH_CACHE[fingerprint] = (principal, time.time())
192 return principal
195# ---------------------------------------------------------------------------
196# Middleware
197# ---------------------------------------------------------------------------
200class RateLimitMiddleware(BaseHTTPMiddleware):
201 """Per-principal token-bucket rate limiting (spec §22.4).
203 Federation endpoints (/v1/federation/) are exempt — they use peer-token
204 auth, not user API keys. Requests without a Bearer token are also exempt.
205 Setting rate_limit_write_per_hour=0 AND rate_limit_read_per_hour=0
206 disables quota enforcement entirely (dev/test shortcut).
207 """
209 async def dispatch(self, request: Request, call_next: Any) -> Any:
210 if request.method == "OPTIONS": 210 ↛ 211line 210 didn't jump to line 211 because the condition on line 210 was never true
211 return await call_next(request)
213 if request.url.path.startswith("/v1/federation/"):
214 return await call_next(request)
216 auth_header = request.headers.get("authorization", "")
217 if not auth_header.lower().startswith("bearer "):
218 return await call_next(request)
220 # Global kill-switch: both limits=0 disables enforcement.
221 if settings.rate_limit_write_per_hour == 0 and settings.rate_limit_read_per_hour == 0:
222 return await call_next(request)
224 raw_key = auth_header[7:]
225 principal = _lookup_principal(raw_key)
226 if principal is None:
227 # Unknown/expired key — let auth middleware reject it properly.
228 return await call_next(request)
230 entity_uri, tenant_id, oidc_sub = principal
231 dimension = _dimension(request.url.path, request.method)
232 if dimension is None: 232 ↛ 233line 232 didn't jump to line 233 because the condition on line 232 was never true
233 return await call_next(request)
235 allowed, retry_after = _check_and_consume(entity_uri, tenant_id, dimension)
236 if not allowed:
237 # Write-ahead: emit quota_breach audit event before returning 429.
238 from .observability.audit_event import emit_nofail
240 emit_nofail(
241 "quota_breach",
242 entity_uri=entity_uri,
243 tenant_id=tenant_id,
244 oidc_sub=oidc_sub,
245 detail={
246 "dimension": dimension,
247 "path": request.url.path,
248 "method": request.method,
249 "retry_after": retry_after,
250 },
251 )
252 retry_ceil = math.ceil(retry_after)
253 return JSONResponse(
254 status_code=429,
255 content={
256 "error": "quota_exceeded",
257 "dimension": dimension,
258 "principal": entity_uri,
259 "retry_after": retry_after,
260 },
261 headers={"Retry-After": str(max(retry_ceil, 1))},
262 )
264 return await call_next(request)