Coverage for app/backend/src/tests/test_ratelimit.py: 99%
181 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-17 00:57 +0000
1import os
2from uuid import uuid4
4import grpc
5import pytest
6import valkey
8from couchers import metrics
9from couchers.middleware import ratelimit
10from couchers.middleware.interceptors import CouchersHeaders
11from couchers.proto import api_pb2, auth_pb2
12from tests.fixtures.db import generate_user
13from tests.fixtures.sessions import auth_api_session, real_api_session
15AUTHENTICATE = "/org.couchers.auth.Auth/Authenticate"
16USERNAME_VALID = "/org.couchers.auth.Auth/UsernameValid"
18# Where the Valkey integration tests look for a server; docker-compose.test.yml publishes one on 6545.
19VALKEY_TEST_HOST = os.environ.get("VALKEY_TEST_HOST", "localhost")
20VALKEY_TEST_PORT = int(os.environ.get("VALKEY_TEST_PORT", "6545"))
23@pytest.fixture(autouse=True)
24def _(testconfig):
25 pass
28class InMemoryCounterStore:
29 """A pure-Python fixed-window store mirroring the Valkey one, for hermetic tests."""
31 def __init__(self) -> None:
32 self.counts: dict[str, int] = {}
34 def incr_and_check(self, entries: list[tuple[str, int]], ttl: int) -> list[int]:
35 tripped = []
36 for i, (key, limit) in enumerate(entries):
37 self.counts[key] = self.counts.get(key, 0) + 1
38 if self.counts[key] > limit:
39 tripped.append(i)
40 return tripped
43class AlwaysTripStore:
44 def incr_and_check(self, entries: list[tuple[str, int]], ttl: int) -> list[int]:
45 return list(range(len(entries)))
48class BrokenStore:
49 def incr_and_check(self, entries: list[tuple[str, int]], ttl: int) -> list[int]:
50 raise RuntimeError("valkey down")
53def _use_store(monkeypatch, store):
54 """Point the rate limiter at this counter store (None meaning none configured)."""
55 monkeypatch.setattr(ratelimit, "_get_store", lambda: store)
58@pytest.fixture
59def store(monkeypatch):
60 """Inject an in-memory counter store, bypassing Valkey."""
61 s = InMemoryCounterStore()
62 _use_store(monkeypatch, s)
63 return s
66@pytest.fixture
67def valkey_client():
68 """A raw client against the test Valkey, skipping the test if there isn't one running."""
69 client = valkey.Valkey(host=VALKEY_TEST_HOST, port=VALKEY_TEST_PORT, socket_connect_timeout=1, socket_timeout=1)
70 try:
71 client.ping()
72 except valkey.ConnectionError as e:
73 pytest.skip(
74 f"no Valkey at {VALKEY_TEST_HOST}:{VALKEY_TEST_PORT} ({e}); "
75 f"start one with `docker compose -f docker-compose.test.yml up -d valkey_tests`"
76 )
77 return client
80@pytest.fixture
81def valkey_store(valkey_client):
82 """The real Valkey-backed store, so the Lua script itself is exercised rather than a stand-in."""
83 return ratelimit.ValkeyCounterStore(VALKEY_TEST_HOST, VALKEY_TEST_PORT)
86@pytest.fixture
87def key_prefix():
88 """A prefix unique to this test run, so counters can't collide with a previous run's leftovers."""
89 return f"test:{uuid4().hex}"
92def _limited(method: str, ip: str | None = None) -> bool:
93 """should_rate_limit for an unauthenticated call from this IP; the limiter reads no other header."""
94 headers = CouchersHeaders(
95 token=None,
96 is_api_key=False,
97 ip_address=ip,
98 user_agent=None,
99 client_platform=None,
100 ui_lang=None,
101 user_id_str=None,
102 sofa=None,
103 )
104 return ratelimit.should_rate_limit(method, headers, None)
107def test_ip_to_key_ipv4():
108 assert ratelimit.ip_to_key("1.2.3.4", 64) == "1.2.3.4/32"
111def test_ip_to_key_ipv6_masks_to_prefix():
112 assert ratelimit.ip_to_key("2001:db8::1", 64) == "2001:db8::/64"
113 assert ratelimit.ip_to_key("2001:0db8:0000:0000:dead:beef:0:1", 64) == "2001:db8::/64"
114 assert ratelimit.ip_to_key("2001:db8::1", 64) == ratelimit.ip_to_key("2001:db8::ffff", 64)
117def test_ip_to_key_ipv6_prefix_configurable():
118 assert ratelimit.ip_to_key("2001:db8:abcd:1234::1", 48) == "2001:db8:abcd::/48"
121def test_resolve_method_rate_limits_method_override():
122 limits = ratelimit.resolve_method_rate_limits(AUTHENTICATE)
123 assert limits.rpc == {"per_ip": 10, "per_user": 120, "global": 6000}
126def test_resolve_method_rate_limits_defaults():
127 limits = ratelimit.resolve_method_rate_limits(USERNAME_VALID)
128 assert limits.rpc == {"per_ip": 60, "per_user": 120, "global": 6000}
129 assert limits.svc == {"per_ip": 300, "per_user": 600, "global": 20000}
130 assert limits.api == {"per_ip": 600, "per_user": 1200, "global": 60000}
133def test_no_store_means_rate_limiting_is_off():
134 # this is the one test that goes through the real accessor, so it can't reuse another test's store
135 ratelimit._get_store.cache_clear()
136 assert ratelimit._get_store() is None
137 assert not _limited(AUTHENTICATE, ip="1.2.3.4")
140def test_trips_per_ip(feature_flags, store):
141 feature_flags.set("rate_limiting_enabled", True)
142 # Authenticate per_ip = 10 and every other limit is far higher, so the 11th call from an IP is the first
143 # to trip anything
144 for _ in range(10):
145 assert not _limited(AUTHENTICATE, ip="1.2.3.4")
146 assert _limited(AUTHENTICATE, ip="1.2.3.4")
149def test_per_ip_skipped_without_ip(feature_flags, store):
150 feature_flags.set("rate_limiting_enabled", True)
151 # no IP → the per_ip dimension is not counted, so the per_ip=10 limit can never trip
152 for _ in range(20):
153 assert not _limited(AUTHENTICATE)
156def test_separate_subnets_counted_separately(feature_flags, store):
157 feature_flags.set("rate_limiting_enabled", True)
158 for _ in range(11):
159 _limited(AUTHENTICATE, ip="2001:db8:1::1")
160 assert not _limited(AUTHENTICATE, ip="2001:db8:2::1")
163def test_store_error_fails_open(feature_flags, monkeypatch):
164 feature_flags.set("rate_limiting_enabled", True)
165 _use_store(monkeypatch, BrokenStore())
166 captured = []
167 monkeypatch.setattr("couchers.middleware.ratelimit.sentry_sdk.capture_exception", lambda e: captured.append(e))
168 monkeypatch.setattr("couchers.middleware.ratelimit.sentry_sdk.set_tag", lambda *a, **k: None)
170 assert not _limited(AUTHENTICATE, ip="1.2.3.4")
171 assert len(captured) == 1
174def test_interceptor_superuser_exempt_when_enforcing(db, feature_flags, monkeypatch):
175 feature_flags.set("rate_limiting_enabled", True)
176 _use_store(monkeypatch, AlwaysTripStore())
177 superuser, token = generate_user(is_superuser=True)
179 # real_api_session, not api_session: only the real server runs the interceptor the limiter lives in
180 with real_api_session(token) as api:
181 assert api.Ping(api_pb2.PingReq()).user.user_id == superuser.id
184def test_interceptor_non_superuser_still_blocked_when_enforcing(db, feature_flags, monkeypatch):
185 feature_flags.set("rate_limiting_enabled", True)
186 _use_store(monkeypatch, AlwaysTripStore())
187 _, token = generate_user()
189 with real_api_session(token) as api:
190 with pytest.raises(grpc.RpcError) as e:
191 api.Ping(api_pb2.PingReq())
192 assert e.value.code() == grpc.StatusCode.RESOURCE_EXHAUSTED
195def test_interceptor_no_store_allows(db, feature_flags, monkeypatch):
196 feature_flags.set("rate_limiting_enabled", True)
197 _use_store(monkeypatch, None)
198 with auth_api_session() as (auth_api, _):
199 assert auth_api.UsernameValid(auth_pb2.UsernameValidReq(username="test")).valid
202def test_interceptor_shadow_allows(db, feature_flags, monkeypatch):
203 feature_flags.set("rate_limiting_enabled", False)
204 _use_store(monkeypatch, AlwaysTripStore())
205 with auth_api_session() as (auth_api, _):
206 assert auth_api.UsernameValid(auth_pb2.UsernameValidReq(username="test")).valid
209def test_interceptor_enforce_rejects(db, feature_flags, monkeypatch):
210 feature_flags.set("rate_limiting_enabled", True)
211 _use_store(monkeypatch, AlwaysTripStore())
212 with auth_api_session() as (auth_api, _):
213 with pytest.raises(grpc.RpcError) as e:
214 auth_api.UsernameValid(auth_pb2.UsernameValidReq(username="test"))
215 assert e.value.code() == grpc.StatusCode.RESOURCE_EXHAUSTED
218def test_interceptor_fails_open_when_enforcing(db, feature_flags, monkeypatch):
219 feature_flags.set("rate_limiting_enabled", True)
220 _use_store(monkeypatch, BrokenStore())
221 with auth_api_session() as (auth_api, _):
222 assert auth_api.UsernameValid(auth_pb2.UsernameValidReq(username="test")).valid
225def _metric_value(counter, name: str, **labels: str) -> float:
226 return float(
227 sum(
228 s.value
229 for m in counter.collect()
230 for s in m.samples
231 if s.name == name and all(s.labels.get(k) == v for k, v in labels.items())
232 )
233 )
236def test_interceptor_emits_metrics_on_enforce(db, feature_flags, monkeypatch):
237 feature_flags.set("rate_limiting_enabled", True)
238 _use_store(monkeypatch, AlwaysTripStore())
240 blocked_before = _metric_value(
241 metrics.rate_limit_checks_counter,
242 "couchers_rate_limit_checks_total",
243 method=USERNAME_VALID,
244 decision="blocked",
245 )
246 # no IP/user on this call, so the global dimension trips at every scope
247 trip_before = _metric_value(
248 metrics.rate_limit_trips_counter,
249 "couchers_rate_limit_trips_total",
250 method=USERNAME_VALID,
251 scope="rpc",
252 dimension="global",
253 enforced="true",
254 )
256 with auth_api_session() as (auth_api, _):
257 with pytest.raises(grpc.RpcError):
258 auth_api.UsernameValid(auth_pb2.UsernameValidReq(username="test"))
260 assert (
261 _metric_value(
262 metrics.rate_limit_checks_counter,
263 "couchers_rate_limit_checks_total",
264 method=USERNAME_VALID,
265 decision="blocked",
266 )
267 == blocked_before + 1
268 )
269 assert (
270 _metric_value(
271 metrics.rate_limit_trips_counter,
272 "couchers_rate_limit_trips_total",
273 method=USERNAME_VALID,
274 scope="rpc",
275 dimension="global",
276 enforced="true",
277 )
278 == trip_before + 1
279 )
282# The tests below run the real Lua script against a real Valkey; everything above uses a stand-in store.
285def test_valkey_store_counts_and_trips(valkey_store, key_prefix):
286 key = f"{key_prefix}:counted"
287 for _ in range(3):
288 assert valkey_store.incr_and_check([(key, 3)], 120) == []
289 assert valkey_store.incr_and_check([(key, 3)], 120) == [0]
290 assert valkey_store.incr_and_check([(key, 3)], 120) == [0]
293def test_valkey_store_returns_indices_of_tripped_entries(valkey_store, key_prefix):
294 # a limit of 0 trips on the first increment, a high limit never does; this pins the Lua script's
295 # 1-based indices being translated back to the 0-based positions of the entries passed in
296 entries = [
297 (f"{key_prefix}:high:0", 100),
298 (f"{key_prefix}:zero:1", 0),
299 (f"{key_prefix}:high:2", 100),
300 (f"{key_prefix}:zero:3", 0),
301 ]
302 assert valkey_store.incr_and_check(entries, 120) == [1, 3]
305def test_valkey_store_counts_keys_independently(valkey_store, key_prefix):
306 a, b = f"{key_prefix}:a", f"{key_prefix}:b"
307 for _ in range(5):
308 valkey_store.incr_and_check([(a, 5)], 120)
309 assert valkey_store.incr_and_check([(a, 5), (b, 5)], 120) == [0]
312def test_valkey_store_sets_ttl_on_first_increment(valkey_store, valkey_client, key_prefix):
313 key = f"{key_prefix}:ttl"
314 valkey_store.incr_and_check([(key, 100)], 120)
315 # without a TTL the counter would never reset and the key would leak
316 assert 0 < valkey_client.ttl(key) <= 120
319def test_valkey_store_does_not_extend_ttl_on_later_increments(valkey_store, valkey_client, key_prefix):
320 key = f"{key_prefix}:ttl-once"
321 valkey_store.incr_and_check([(key, 100)], 120)
322 # pull the expiry in, then increment again: the window must not slide, or a sustained flood would keep
323 # renewing its own counter and the fixed window would never roll over
324 valkey_client.expire(key, 5)
325 valkey_store.incr_and_check([(key, 100)], 120)
326 assert 0 < valkey_client.ttl(key) <= 5
329def test_rate_limiting_end_to_end_against_valkey(feature_flags, valkey_store, monkeypatch):
330 feature_flags.set("rate_limiting_enabled", True)
331 _use_store(monkeypatch, valkey_store)
332 # a /64 unique to this run, so the per-IP counters start clean
333 ip = f"2001:db8:{uuid4().hex[:4]}:{uuid4().hex[:4]}::1"
335 # Authenticate annotates per_ip = 10
336 for _ in range(10):
337 assert not _limited(AUTHENTICATE, ip=ip)
338 assert _limited(AUTHENTICATE, ip=ip)
340 other_ip = f"2001:db8:{uuid4().hex[:4]}:{uuid4().hex[:4]}::1"
341 assert not _limited(AUTHENTICATE, ip=other_ip)