feat(core): add resilient fixed-window throttling
This commit is contained in:
21
docs/THROTTLING.md
Normal file
21
docs/THROTTLING.md
Normal file
@@ -0,0 +1,21 @@
|
||||
# Shared fixed-window throttling
|
||||
|
||||
Core exposes `govoplan_core.core.throttling` for security-sensitive endpoints
|
||||
that need bounded fixed-window counters. Subjects are SHA-256 hashed before they
|
||||
become store keys. A configured Redis instance provides atomic counters shared
|
||||
across API workers; development, a missing Redis configuration, and temporary
|
||||
Redis outages use a bounded process-local fallback. When Redis fails, local
|
||||
attempts are still mirrored so losing the distributed store does not reset the
|
||||
active worker's protection window.
|
||||
|
||||
Callers define one or more `ThrottleDimension` values with a controlled
|
||||
namespace, a subject and a positive limit. They must call `check` before an
|
||||
expensive verifier, `record` after a failed attempt, and may `reset` the relevant
|
||||
dimension after successful verification. A blocked decision includes a
|
||||
`retry_after_seconds` value suitable for an HTTP `Retry-After` header.
|
||||
|
||||
The first consumer is Scheduling's anonymous participation password challenge.
|
||||
Its namespace is `poll-participation-password`; its subject combines tenant,
|
||||
scheduling request and Poll's non-secret invitation-token fingerprint. Access's
|
||||
login throttle predates this primitive and should be migrated onto it in a
|
||||
separate compatibility-preserving slice.
|
||||
349
src/govoplan_core/core/throttling.py
Normal file
349
src/govoplan_core/core/throttling.py
Normal file
@@ -0,0 +1,349 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol
|
||||
|
||||
from redis import Redis
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_NAMESPACE_PATTERN = re.compile(r"^[a-z0-9][a-z0-9._-]{0,99}$")
|
||||
_REDIS_INCREMENT_SCRIPT = """
|
||||
local count = redis.call('INCR', KEYS[1])
|
||||
if count == 1 then
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[1])
|
||||
end
|
||||
local ttl = redis.call('TTL', KEYS[1])
|
||||
return {count, ttl}
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class FixedWindowBucket:
|
||||
count: int = 0
|
||||
retry_after_seconds: int = 0
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ThrottleDimension:
|
||||
"""One independently limited dimension without putting its value in keys."""
|
||||
|
||||
namespace: str
|
||||
subject: str
|
||||
limit: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ThrottleDecision:
|
||||
allowed: bool
|
||||
retry_after_seconds: int = 0
|
||||
|
||||
|
||||
class FixedWindowStore(Protocol):
|
||||
def read(self, key: str) -> FixedWindowBucket: ...
|
||||
|
||||
def increment(self, key: str, *, window_seconds: int) -> FixedWindowBucket: ...
|
||||
|
||||
def delete(self, key: str) -> None: ...
|
||||
|
||||
|
||||
class InMemoryFixedWindowStore:
|
||||
"""Bounded process-local store for development and distributed-store loss."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
max_entries: int = 10_000,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> None:
|
||||
self._entries: dict[str, tuple[int, float]] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._max_entries = max(2, max_entries)
|
||||
self._clock = clock
|
||||
|
||||
def read(self, key: str) -> FixedWindowBucket:
|
||||
now = self._clock()
|
||||
with self._lock:
|
||||
entry = self._active_entry(key, now=now)
|
||||
if entry is None:
|
||||
return FixedWindowBucket()
|
||||
count, expires_at = entry
|
||||
return FixedWindowBucket(
|
||||
count=count,
|
||||
retry_after_seconds=max(1, int(expires_at - now + 0.999)),
|
||||
)
|
||||
|
||||
def increment(self, key: str, *, window_seconds: int) -> FixedWindowBucket:
|
||||
now = self._clock()
|
||||
with self._lock:
|
||||
entry = self._active_entry(key, now=now)
|
||||
if entry is None:
|
||||
self._make_room(now=now, incoming_key=key)
|
||||
count = 1
|
||||
expires_at = now + window_seconds
|
||||
else:
|
||||
count = entry[0] + 1
|
||||
expires_at = entry[1]
|
||||
self._entries[key] = (count, expires_at)
|
||||
return FixedWindowBucket(
|
||||
count=count,
|
||||
retry_after_seconds=max(1, int(expires_at - now + 0.999)),
|
||||
)
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
with self._lock:
|
||||
self._entries.pop(key, None)
|
||||
|
||||
def _active_entry(self, key: str, *, now: float) -> tuple[int, float] | None:
|
||||
entry = self._entries.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
if entry[1] <= now:
|
||||
self._entries.pop(key, None)
|
||||
return None
|
||||
return entry
|
||||
|
||||
def _make_room(self, *, now: float, incoming_key: str) -> None:
|
||||
if incoming_key in self._entries or len(self._entries) < self._max_entries:
|
||||
return
|
||||
for key, (_count, expires_at) in tuple(self._entries.items()):
|
||||
if expires_at <= now:
|
||||
self._entries.pop(key, None)
|
||||
while len(self._entries) >= self._max_entries:
|
||||
self._entries.pop(next(iter(self._entries)))
|
||||
|
||||
|
||||
class RedisFixedWindowStore:
|
||||
"""Atomic Redis fixed-window counters shared across API workers."""
|
||||
|
||||
def __init__(self, redis_url: str) -> None:
|
||||
self._client = Redis.from_url(
|
||||
redis_url,
|
||||
decode_responses=True,
|
||||
socket_connect_timeout=0.25,
|
||||
socket_timeout=0.25,
|
||||
health_check_interval=30,
|
||||
)
|
||||
|
||||
def read(self, key: str) -> FixedWindowBucket:
|
||||
pipeline = self._client.pipeline(transaction=False)
|
||||
pipeline.get(key)
|
||||
pipeline.ttl(key)
|
||||
raw_count, raw_ttl = pipeline.execute()
|
||||
return FixedWindowBucket(
|
||||
count=int(raw_count or 0),
|
||||
retry_after_seconds=max(0, int(raw_ttl or 0)),
|
||||
)
|
||||
|
||||
def increment(self, key: str, *, window_seconds: int) -> FixedWindowBucket:
|
||||
result = self._client.eval(
|
||||
_REDIS_INCREMENT_SCRIPT,
|
||||
1,
|
||||
key,
|
||||
window_seconds,
|
||||
)
|
||||
if not isinstance(result, (list, tuple)) or len(result) != 2:
|
||||
raise RedisError("Unexpected fixed-window throttle response from Redis")
|
||||
return FixedWindowBucket(
|
||||
count=int(result[0]),
|
||||
retry_after_seconds=max(1, int(result[1])),
|
||||
)
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
self._client.delete(key)
|
||||
|
||||
|
||||
class ResilientFixedWindowStore:
|
||||
"""Mirror locally, prefer Redis, and fall back safely during an outage."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
primary: FixedWindowStore | None,
|
||||
fallback: FixedWindowStore,
|
||||
*,
|
||||
retry_seconds: int = 30,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> None:
|
||||
self._primary = primary
|
||||
self._fallback = fallback
|
||||
self._retry_seconds = max(1, retry_seconds)
|
||||
self._clock = clock
|
||||
self._primary_unavailable_until = 0.0
|
||||
self._state_lock = threading.Lock()
|
||||
|
||||
def read(self, key: str) -> FixedWindowBucket:
|
||||
fallback_result = self._fallback.read(key)
|
||||
primary = self._available_primary()
|
||||
if primary is None:
|
||||
return fallback_result
|
||||
try:
|
||||
return _stricter_bucket(primary.read(key), fallback_result)
|
||||
except (RedisError, OSError, TimeoutError, ConnectionError) as exc:
|
||||
self._mark_primary_unavailable(exc)
|
||||
return fallback_result
|
||||
|
||||
def increment(self, key: str, *, window_seconds: int) -> FixedWindowBucket:
|
||||
fallback_result = self._fallback.increment(
|
||||
key,
|
||||
window_seconds=window_seconds,
|
||||
)
|
||||
primary = self._available_primary()
|
||||
if primary is None:
|
||||
return fallback_result
|
||||
try:
|
||||
primary_result = primary.increment(
|
||||
key,
|
||||
window_seconds=window_seconds,
|
||||
)
|
||||
return _stricter_bucket(primary_result, fallback_result)
|
||||
except (RedisError, OSError, TimeoutError, ConnectionError) as exc:
|
||||
self._mark_primary_unavailable(exc)
|
||||
return fallback_result
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
self._fallback.delete(key)
|
||||
primary = self._available_primary()
|
||||
if primary is None:
|
||||
return
|
||||
try:
|
||||
primary.delete(key)
|
||||
except (RedisError, OSError, TimeoutError, ConnectionError) as exc:
|
||||
self._mark_primary_unavailable(exc)
|
||||
|
||||
def _available_primary(self) -> FixedWindowStore | None:
|
||||
if self._primary is None:
|
||||
return None
|
||||
with self._state_lock:
|
||||
if self._clock() < self._primary_unavailable_until:
|
||||
return None
|
||||
return self._primary
|
||||
|
||||
def _mark_primary_unavailable(self, exc: Exception) -> None:
|
||||
should_log = False
|
||||
with self._state_lock:
|
||||
now = self._clock()
|
||||
if now >= self._primary_unavailable_until:
|
||||
should_log = True
|
||||
self._primary_unavailable_until = now + self._retry_seconds
|
||||
if should_log:
|
||||
logger.warning(
|
||||
"Distributed throttling is unavailable; using the process-local fallback (%s)",
|
||||
type(exc).__name__,
|
||||
)
|
||||
|
||||
|
||||
def _stricter_bucket(
|
||||
first: FixedWindowBucket,
|
||||
second: FixedWindowBucket,
|
||||
) -> FixedWindowBucket:
|
||||
return FixedWindowBucket(
|
||||
count=max(first.count, second.count),
|
||||
retry_after_seconds=max(
|
||||
first.retry_after_seconds,
|
||||
second.retry_after_seconds,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class FixedWindowThrottle:
|
||||
def __init__(
|
||||
self,
|
||||
store: FixedWindowStore,
|
||||
*,
|
||||
window_seconds: int,
|
||||
key_prefix: str = "govoplan:throttle:v1",
|
||||
) -> None:
|
||||
self._store = store
|
||||
self._window_seconds = max(1, window_seconds)
|
||||
self._key_prefix = key_prefix.rstrip(":")
|
||||
|
||||
def check(self, dimensions: Iterable[ThrottleDimension]) -> ThrottleDecision:
|
||||
return self._decision(tuple(dimensions), increment=False)
|
||||
|
||||
def record(self, dimensions: Iterable[ThrottleDimension]) -> ThrottleDecision:
|
||||
return self._decision(tuple(dimensions), increment=True)
|
||||
|
||||
def reset(self, dimensions: Iterable[ThrottleDimension]) -> None:
|
||||
for dimension in tuple(dimensions):
|
||||
self._store.delete(self._key(dimension))
|
||||
|
||||
def _decision(
|
||||
self,
|
||||
dimensions: tuple[ThrottleDimension, ...],
|
||||
*,
|
||||
increment: bool,
|
||||
) -> ThrottleDecision:
|
||||
if not dimensions:
|
||||
raise ValueError("At least one throttle dimension is required")
|
||||
blocked_retry_after = 0
|
||||
for dimension in dimensions:
|
||||
if dimension.limit < 1:
|
||||
raise ValueError("Throttle limits must be positive")
|
||||
key = self._key(dimension)
|
||||
state = (
|
||||
self._store.increment(key, window_seconds=self._window_seconds)
|
||||
if increment
|
||||
else self._store.read(key)
|
||||
)
|
||||
if state.count >= dimension.limit:
|
||||
blocked_retry_after = max(
|
||||
blocked_retry_after,
|
||||
max(1, state.retry_after_seconds),
|
||||
)
|
||||
return ThrottleDecision(
|
||||
allowed=blocked_retry_after == 0,
|
||||
retry_after_seconds=blocked_retry_after,
|
||||
)
|
||||
|
||||
def _key(self, dimension: ThrottleDimension) -> str:
|
||||
namespace = dimension.namespace.strip().casefold()
|
||||
if not _NAMESPACE_PATTERN.fullmatch(namespace):
|
||||
raise ValueError("Throttle namespaces must use lowercase letters, digits, '.', '_' or '-'")
|
||||
subject_digest = hashlib.sha256(dimension.subject.encode("utf-8")).hexdigest()
|
||||
return f"{self._key_prefix}:{namespace}:{subject_digest}"
|
||||
|
||||
|
||||
def build_fixed_window_throttle(
|
||||
*,
|
||||
redis_url: str | None,
|
||||
window_seconds: int,
|
||||
redis_retry_seconds: int = 30,
|
||||
max_local_entries: int = 10_000,
|
||||
key_prefix: str = "govoplan:throttle:v1",
|
||||
) -> FixedWindowThrottle:
|
||||
primary = (
|
||||
RedisFixedWindowStore(redis_url)
|
||||
if redis_url is not None and redis_url.strip()
|
||||
else None
|
||||
)
|
||||
store = ResilientFixedWindowStore(
|
||||
primary,
|
||||
InMemoryFixedWindowStore(max_entries=max_local_entries),
|
||||
retry_seconds=redis_retry_seconds,
|
||||
)
|
||||
return FixedWindowThrottle(
|
||||
store,
|
||||
window_seconds=window_seconds,
|
||||
key_prefix=key_prefix,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"FixedWindowBucket",
|
||||
"FixedWindowStore",
|
||||
"FixedWindowThrottle",
|
||||
"InMemoryFixedWindowStore",
|
||||
"RedisFixedWindowStore",
|
||||
"ResilientFixedWindowStore",
|
||||
"ThrottleDecision",
|
||||
"ThrottleDimension",
|
||||
"build_fixed_window_throttle",
|
||||
]
|
||||
130
tests/test_throttling.py
Normal file
130
tests/test_throttling.py
Normal file
@@ -0,0 +1,130 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
from govoplan_core.core.throttling import (
|
||||
FixedWindowBucket,
|
||||
FixedWindowThrottle,
|
||||
InMemoryFixedWindowStore,
|
||||
ResilientFixedWindowStore,
|
||||
ThrottleDimension,
|
||||
build_fixed_window_throttle,
|
||||
)
|
||||
|
||||
|
||||
class _Clock:
|
||||
def __init__(self) -> None:
|
||||
self.value = 100.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.value
|
||||
|
||||
|
||||
class _RecordingStore:
|
||||
def __init__(self) -> None:
|
||||
self.values: dict[str, int] = {}
|
||||
|
||||
def read(self, key: str) -> FixedWindowBucket:
|
||||
return FixedWindowBucket(self.values.get(key, 0), 60)
|
||||
|
||||
def increment(self, key: str, *, window_seconds: int) -> FixedWindowBucket:
|
||||
del window_seconds
|
||||
self.values[key] = self.values.get(key, 0) + 1
|
||||
return FixedWindowBucket(self.values[key], 60)
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
self.values.pop(key, None)
|
||||
|
||||
|
||||
class _FailingStore(_RecordingStore):
|
||||
def read(self, key: str) -> FixedWindowBucket:
|
||||
del key
|
||||
raise RedisError("offline")
|
||||
|
||||
def increment(self, key: str, *, window_seconds: int) -> FixedWindowBucket:
|
||||
del key, window_seconds
|
||||
raise RedisError("offline")
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
del key
|
||||
raise RedisError("offline")
|
||||
|
||||
|
||||
class FixedWindowThrottleTests(unittest.TestCase):
|
||||
def test_blocks_at_limit_and_expires_without_sleeping(self) -> None:
|
||||
clock = _Clock()
|
||||
store = InMemoryFixedWindowStore(clock=clock)
|
||||
throttle = FixedWindowThrottle(store, window_seconds=60)
|
||||
dimensions = (ThrottleDimension("password", "secret-subject", 3),)
|
||||
|
||||
self.assertTrue(throttle.check(dimensions).allowed)
|
||||
self.assertTrue(throttle.record(dimensions).allowed)
|
||||
self.assertTrue(throttle.record(dimensions).allowed)
|
||||
blocked = throttle.record(dimensions)
|
||||
self.assertFalse(blocked.allowed)
|
||||
self.assertEqual(blocked.retry_after_seconds, 60)
|
||||
self.assertFalse(throttle.check(dimensions).allowed)
|
||||
|
||||
clock.value += 61
|
||||
self.assertTrue(throttle.check(dimensions).allowed)
|
||||
|
||||
def test_subject_is_hashed_and_success_can_reset_it(self) -> None:
|
||||
store = _RecordingStore()
|
||||
throttle = FixedWindowThrottle(store, window_seconds=60)
|
||||
dimensions = (ThrottleDimension("poll-participation-password", "tenant:token-secret", 2),)
|
||||
|
||||
throttle.record(dimensions)
|
||||
|
||||
self.assertEqual(len(store.values), 1)
|
||||
stored_key = next(iter(store.values))
|
||||
self.assertNotIn("tenant:token-secret", stored_key)
|
||||
self.assertIn(":poll-participation-password:", stored_key)
|
||||
throttle.reset(dimensions)
|
||||
self.assertEqual(store.values, {})
|
||||
|
||||
def test_resilient_store_uses_bounded_local_fallback(self) -> None:
|
||||
clock = _Clock()
|
||||
fallback = InMemoryFixedWindowStore(max_entries=2, clock=clock)
|
||||
store = ResilientFixedWindowStore(
|
||||
_FailingStore(),
|
||||
fallback,
|
||||
retry_seconds=30,
|
||||
clock=clock,
|
||||
)
|
||||
|
||||
first = store.increment("first", window_seconds=60)
|
||||
second = store.increment("first", window_seconds=60)
|
||||
store.increment("second", window_seconds=60)
|
||||
store.increment("third", window_seconds=60)
|
||||
|
||||
self.assertEqual(first.count, 1)
|
||||
self.assertEqual(second.count, 2)
|
||||
self.assertEqual(fallback.read("first").count, 0)
|
||||
self.assertEqual(fallback.read("third").count, 1)
|
||||
|
||||
def test_builder_operates_without_redis(self) -> None:
|
||||
throttle = build_fixed_window_throttle(
|
||||
redis_url=None,
|
||||
window_seconds=30,
|
||||
max_local_entries=10,
|
||||
)
|
||||
dimensions = (ThrottleDimension("test", "subject", 1),)
|
||||
|
||||
self.assertFalse(throttle.record(dimensions).allowed)
|
||||
self.assertFalse(throttle.check(dimensions).allowed)
|
||||
|
||||
def test_rejects_invalid_dimensions(self) -> None:
|
||||
throttle = FixedWindowThrottle(_RecordingStore(), window_seconds=60)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "At least one"):
|
||||
throttle.check(())
|
||||
with self.assertRaisesRegex(ValueError, "positive"):
|
||||
throttle.check((ThrottleDimension("test", "subject", 0),))
|
||||
with self.assertRaisesRegex(ValueError, "namespaces"):
|
||||
throttle.check((ThrottleDimension("Invalid namespace", "subject", 1),))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user