Coverage for app/backend/src/couchers/middleware/ratelimit.py: 98%
88 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-17 00:57 +0000
1"""
2API rate limiting. See docs/rate-limit-design.md.
4Limits are a pair of (scope, dimension). Scopes nest: per-RPC ⊂ per-servicer ⊂ all-API; a single request
5increments a counter at every level. Dimensions are per-IP (keyed by subnet), per-user, and global. A
6request is rejected if any applicable limit is exceeded.
8This is independent of couchers.rate_limits, which is the per-user 24h action limiter.
9"""
11import logging
12import time
13from dataclasses import dataclass
14from functools import cache
15from ipaddress import ip_network
16from typing import TYPE_CHECKING
18import sentry_sdk
19import valkey
21from couchers.config import config
22from couchers.constants import RATE_LIMIT_WINDOW_SECONDS
23from couchers.experimentation import get_global_boolean_value
24from couchers.metrics import (
25 observe_rate_limit_check,
26 observe_rate_limit_duration,
27 observe_rate_limit_store_error,
28 observe_rate_limit_trip,
29)
30from couchers.middleware.descriptor_pool import get_descriptor_pool
31from couchers.middleware.proto_annotations import method_extension, optional_field, service_extension, split_method
32from couchers.proto import annotations_pb2
34if TYPE_CHECKING:
35 from couchers.middleware.interceptors import CouchersHeaders, UserAuthInfo
37logger = logging.getLogger(__name__)
40@dataclass(frozen=True, slots=True)
41class ResolvedLimits:
42 """The fully-resolved per-minute limits for one method, per scope and dimension."""
44 service_name: str
45 rpc: dict[str, int]
46 svc: dict[str, int]
47 api: dict[str, int]
50@cache
51def resolve_method_rate_limits(method: str) -> ResolvedLimits:
52 """
53 Resolve the limits for a method from its proto annotations, falling back to global defaults.
55 per-RPC: method rate_limit.<dim> → service rate_limit_default.<dim> → global rpc default
56 per-servicer: service rate_limit_aggregate.<dim> → global svc default
57 all-API: global api default
58 """
59 dimensions = ("per_ip", "per_user", "global")
60 defaults = {
61 "rpc": {"per_ip": 60, "per_user": 120, "global": 6000},
62 "svc": {"per_ip": 300, "per_user": 600, "global": 20000},
63 "api": {"per_ip": 600, "per_user": 1200, "global": 60000},
64 }
66 def resolve(*values: int | None) -> int:
67 # the last value is always a global default, so there is always one to find
68 return next(value for value in values if value is not None)
70 pool = get_descriptor_pool()
71 service_name, _ = split_method(method)
72 method_rl = method_extension(pool, method, annotations_pb2.rate_limit)
73 service_default = service_extension(pool, service_name, annotations_pb2.rate_limit_default)
74 service_aggregate = service_extension(pool, service_name, annotations_pb2.rate_limit_aggregate)
76 return ResolvedLimits(
77 service_name=service_name,
78 rpc={
79 dim: resolve(optional_field(method_rl, dim), optional_field(service_default, dim), defaults["rpc"][dim])
80 for dim in dimensions
81 },
82 svc={dim: resolve(optional_field(service_aggregate, dim), defaults["svc"][dim]) for dim in dimensions},
83 api=dict(defaults["api"]),
84 )
87def ip_to_key(ip: str, ipv6_prefix: int) -> str:
88 """
89 Mask an IP to its subnet and return a canonical string key.
91 IPv4 is keyed at /32 (the exact address); IPv6 is masked to ipv6_prefix bits (default /64).
92 """
93 network = ip_network(ip, strict=False)
94 prefix = 32 if network.version == 4 else ipv6_prefix
95 return str(ip_network(f"{network.network_address}/{prefix}", strict=False))
98_LUA_SCRIPT = """
99local tripped = {}
100local ttl = tonumber(ARGV[1])
101for i, key in ipairs(KEYS) do
102 local count = redis.call('INCR', key)
103 if count == 1 then
104 redis.call('EXPIRE', key, ttl)
105 end
106 if count > tonumber(ARGV[i + 1]) then
107 tripped[#tripped + 1] = i
108 end
109end
110return tripped
111"""
114class ValkeyCounterStore:
115 def __init__(self, host: str, port: int) -> None:
116 self._client = valkey.Valkey(
117 host=host,
118 port=port,
119 socket_connect_timeout=0.1,
120 socket_timeout=0.1,
121 )
122 self._script = self._client.register_script(_LUA_SCRIPT)
124 def incr_and_check(self, entries: list[tuple[str, int]], ttl: int) -> list[int]:
125 """Increment each key's counter; return the indices of entries whose count now exceeds its limit."""
126 keys = [key for key, _ in entries]
127 limits = [str(limit) for _, limit in entries]
128 result = self._script(keys=keys, args=[str(ttl), *limits])
129 # Lua returns 1-based indices.
130 return [i - 1 for i in result]
133@cache
134def _get_store() -> ValkeyCounterStore | None:
135 """The process-wide counter store, or None when no store is configured and rate limiting is off."""
136 if not config.VALKEY_HOST: 136 ↛ 138line 136 didn't jump to line 138 because the condition on line 136 was always true
137 return None
138 return ValkeyCounterStore(config.VALKEY_HOST, config.VALKEY_PORT)
141def _build_entries(
142 limits: ResolvedLimits, method: str, ip_address: str | None, user_id: int | None, bucket: int
143) -> list[tuple[str, str, str, int]]:
144 """The applicable (scope, dimension, key, limit) tuples for one request, across all scopes and dimensions."""
145 dim_ids: dict[str, str | None] = {
146 "per_ip": ip_to_key(ip_address, config.RATE_LIMIT_IPV6_PREFIX) if ip_address else None,
147 "per_user": str(user_id) if user_id is not None else None,
148 "global": "*",
149 }
150 scopes = (
151 ("rpc", method, limits.rpc),
152 ("svc", limits.service_name, limits.svc),
153 ("api", "*", limits.api),
154 )
155 entries = []
156 for scope, scope_id, scope_limits in scopes:
157 for dim, limit in scope_limits.items():
158 dim_id = dim_ids[dim]
159 if dim_id is None:
160 continue
161 key = f"rl:{scope}:{scope_id}:{dim}:{dim_id}:{bucket}"
162 entries.append((scope, dim, key, limit))
163 return entries
166def should_rate_limit(method: str, headers: CouchersHeaders, auth_info: UserAuthInfo | None) -> bool:
167 """
168 Count this request against every applicable limit and decide whether it should be rejected.
170 True only when a limit tripped and enforcement is on; shadow mode (the default) allows the request, as
171 does having no counter store configured at all, which turns rate limiting off entirely.
172 """
173 if auth_info and auth_info.is_superuser:
174 return False
176 store = _get_store()
177 if store is None:
178 return False
180 start = time.perf_counter_ns()
181 try:
182 entries = _build_entries(
183 resolve_method_rate_limits(method),
184 method,
185 headers.ip_address,
186 auth_info.user_id if auth_info else None,
187 int(time.time() // RATE_LIMIT_WINDOW_SECONDS),
188 )
190 try:
191 # keys carry their window bucket, so a stale key is only ever read within its own window; we let
192 # Valkey expire them a window later as housekeeping
193 tripped_idx = store.incr_and_check(
194 [(key, limit) for _, _, key, limit in entries], 2 * RATE_LIMIT_WINDOW_SECONDS
195 )
196 except Exception as e:
197 # nothing could be counted, so fail open: going dark on the counters shouldn't take the API down
198 observe_rate_limit_store_error(type(e).__name__)
199 observe_rate_limit_check(method, "failed_open")
200 sentry_sdk.set_tag("context", "rate_limiting")
201 sentry_sdk.capture_exception(e)
202 return False
204 if not tripped_idx:
205 observe_rate_limit_check(method, "allowed")
206 return False
208 enforced = get_global_boolean_value("rate_limiting_enabled", False)
209 for i in tripped_idx:
210 scope, dimension, _, _ = entries[i]
211 observe_rate_limit_trip(method, scope, dimension, enforced)
212 observe_rate_limit_check(method, "blocked" if enforced else "shadowed")
213 return enforced
214 finally:
215 observe_rate_limit_duration((time.perf_counter_ns() - start) / 1e9)