Compare commits

..

3 Commits

5 changed files with 487 additions and 341 deletions
+1
View File
@@ -16,6 +16,7 @@ class Settings(BaseSettings):
# Redis # Redis
REDIS_URL: str = "redis://localhost:6379/0" REDIS_URL: str = "redis://localhost:6379/0"
REDIS_TIMEOUT_SECONDS: float = 2.0
# JWT / Auth # JWT / Auth
# NOTE: production must override SECRET_KEY with a secure random value # NOTE: production must override SECRET_KEY with a secure random value
+293
View File
@@ -0,0 +1,293 @@
"""Redis-backed token-bucket rate limiting.
This module provides a reusable FastAPI dependency that rate-limits requests
per endpoint scope (optionally per-user) using an atomic Lua script on Redis.
If Redis is unavailable the limiter can either allow the request through
(fail-open) or reject it with a ``503`` (fail-closed), depending on the
policy configured for that limit. Chat send uses fail-closed by default
because flooding/abuse is worse than a temporary outage.
Example
-------
from app.rate_limit import rate_limit, named_limit
chat_send_limit = rate_limit(
named_limit("chat_send", capacity=5, refill_rate=0.2, fail_open=False),
user_id_dependency=Depends(get_current_user_id),
)
@router.post("/messages")
async def send_message(
payload: MessageCreate,
current_user: CurrentUser,
db: DbDep,
_: None = Depends(chat_send_limit),
):
...
"""
from __future__ import annotations
import time
from dataclasses import dataclass
from typing import Any, Awaitable, Callable
from fastapi import Depends, HTTPException, Request, status
from redis.asyncio import Redis
from redis.exceptions import RedisError
from app.config import settings
# Atomic token-bucket implementation.
#
# KEYS[1] is the bucket key. ARGV[1] is the bucket capacity, ARGV[2] is the
# refill rate in tokens per second, ARGV[3] is the current Unix time in
# seconds (float), and ARGV[4] is the request cost (tokens to take).
#
# The bucket state is stored as a hash with "tokens" and "ts". The script
# refills the bucket up to capacity, then takes the requested tokens. It
# returns an array:
# [0] allowed (1 or 0)
# [1] tokens remaining after the take
# [2] retry-after seconds (0 when allowed)
_TOKEN_BUCKET_LUA = """
local key = KEYS[1]
local capacity = tonumber(ARGV[1])
local refill_rate = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local cost = tonumber(ARGV[4])
local state = redis.call("HMGET", key, "tokens", "ts")
local tokens = capacity
local ts = now
if state[1] ~= false then
tokens = tonumber(state[1])
ts = tonumber(state[2]) or now
end
local elapsed = math.max(0, now - ts)
tokens = math.min(capacity, tokens + (elapsed * refill_rate))
local allowed = 0
local retry_after = 0
if tokens >= cost then
tokens = tokens - cost
allowed = 1
else
-- Time until enough tokens for a single request have accrued.
retry_after = math.ceil((cost - tokens) / refill_rate)
end
-- Keep the bucket state small and expire-able. Retaining the timestamp as a
-- float lets partial tokens accumulate smoothly between requests.
redis.call("HSET", key, "tokens", tostring(tokens), "ts", tostring(now))
redis.call("PEXPIRE", key, 600000)
return {allowed, tostring(tokens), retry_after}
"""
class RateLimitExceeded(HTTPException):
"""Raised when a rate-limit bucket has no tokens left."""
def __init__(self, retry_after: int) -> None:
super().__init__(
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
detail="Rate limit exceeded",
headers={"Retry-After": str(retry_after)},
)
@dataclass(frozen=True)
class RateLimit:
"""Definition of a named token-bucket rate limit.
``capacity`` is the burst size and ``refill_rate`` is the steady-state
token accumulation rate (tokens/second). For example,
``RateLimit(name="chat_send", capacity=10, refill_rate=1.0)`` allows a
burst of 10 requests and then one request per second on average.
"""
name: str
capacity: int
refill_rate: float
# When Redis is unreachable: ``True`` allows the request (fail-open),
# ``False`` rejects it with 503 (fail-closed).
fail_open: bool = True
def __post_init__(self) -> None:
if not self.name:
raise ValueError("name must be non-empty")
if self.capacity <= 0:
raise ValueError("capacity must be positive")
if self.refill_rate <= 0:
raise ValueError("refill_rate must be positive")
def default_user_scope(request: Request, user_id: int | None) -> str:
"""Scope a bucket per authenticated user, or per client IP if anonymous."""
if user_id is not None:
return f"user:{user_id}"
client = request.client
return f"ip:{client.host if client else 'unknown'}"
@dataclass(frozen=True)
class RateLimitConfig:
"""A named rate limit plus how it is applied to a request."""
limit: RateLimit
# Builds the Redis key fragment for the request. The resolved user id is
# passed as the second argument (or ``None`` for anonymous limits).
scope: Callable[[Request, int | None], str] = default_user_scope
@dataclass(frozen=True)
class RateLimitResult:
"""Outcome of a rate-limit check (used by tests)."""
allowed: bool
remaining: int
retry_after: int
_redis_client: Redis | None = None
def get_redis() -> Redis:
"""Return a shared async Redis client for rate limiting.
The client is created lazily (``Redis.from_url`` does not connect until a
command is issued), so this dependency never blocks or fails just because
Redis is down.
"""
global _redis_client
if _redis_client is None:
_redis_client = Redis.from_url(
settings.REDIS_URL,
encoding="utf-8",
decode_responses=True,
socket_connect_timeout=settings.REDIS_TIMEOUT_SECONDS,
socket_timeout=settings.REDIS_TIMEOUT_SECONDS,
)
return _redis_client
class RateLimiter:
"""Enforce a token bucket in Redis.
Buckets are keyed as ``rate:{limit.name}:{scope_key}`` so multiple
endpoints can share a named limit while remaining independently bucketed.
"""
def __init__(self, config: RateLimitConfig, redis: Redis | None = None) -> None:
self._config = config
self._redis = redis
def _redis_client(self) -> Redis:
return self._redis or get_redis()
def _bucket_key(self, request: Request, user_id: int | None) -> str:
limit = self._config.limit
scope_key = self._config.scope(request, user_id)
return f"rate:{limit.name}:{scope_key}"
async def check(self, request: Request, user_id: int | None = None) -> RateLimitResult:
"""Atomically refill/take a token and return the outcome.
Redis failures are mapped to ``allowed=True`` (fail-open) or
``allowed=False`` (fail-closed, represented with ``retry_after=-1``)
according to the limit's policy, without raising.
"""
limit = self._config.limit
key = self._bucket_key(request, user_id)
try:
result = await self._redis_client().eval(
_TOKEN_BUCKET_LUA,
1,
key,
str(limit.capacity),
str(limit.refill_rate),
str(time.time()),
"1", # each request consumes one token
)
except RedisError:
if limit.fail_open:
return RateLimitResult(allowed=True, remaining=0, retry_after=0)
# Fail-closed: no reliable retry estimate is available.
return RateLimitResult(allowed=False, remaining=0, retry_after=-1)
# Redis may return ints or strings depending on decode_responses.
allowed = bool(int(result[0]))
remaining = int(result[1])
retry_after = int(result[2])
return RateLimitResult(allowed=allowed, remaining=remaining, retry_after=retry_after)
async def __call__(self, request: Request, user_id: int | None = None) -> None:
"""Enforce the rate limit, raising 429 or 503 as appropriate."""
result = await self.check(request, user_id)
if result.allowed:
return
if result.retry_after < 0:
# Redis was unavailable and this limit is fail-closed.
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Rate limiting service unavailable",
)
raise RateLimitExceeded(retry_after=result.retry_after)
def named_limit(
name: str,
*,
capacity: int,
refill_rate: float,
fail_open: bool = True,
) -> RateLimitConfig:
"""Create a named rate-limit config with the default per-user/IP scope.
``capacity`` is the burst size and ``refill_rate`` is tokens/second.
"""
return RateLimitConfig(
limit=RateLimit(
name=name,
capacity=capacity,
refill_rate=refill_rate,
fail_open=fail_open,
)
)
def rate_limit(
config: RateLimitConfig,
*,
user_id_dependency: Any = None,
) -> Callable[[Request, Any], Awaitable[None]]:
"""Build a FastAPI dependency enforcing ``config``.
``user_id_dependency`` optionally injects the authenticated user id, e.g.
``Depends(get_current_user_id)``. If omitted the limiter buckets by client
IP (the ``default_user_scope`` fallback), which is suitable for anonymous
endpoints but should be paired with an auth resolver for per-user limits.
The returned callable is intended for ``Depends(...)``.
"""
async def _anonymous_user_id() -> int | None:
return None
resolver = user_id_dependency if user_id_dependency is not None else Depends(_anonymous_user_id)
async def dependency(
request: Request,
user_id: int | None = resolver,
) -> None:
limiter = RateLimiter(config)
await limiter(request, user_id)
return dependency
-1
View File
@@ -1 +0,0 @@
"""Application services."""
-340
View File
@@ -1,340 +0,0 @@
"""Pluggable server-side text filtering service.
Provides a ``TextFilter`` protocol with a deterministic default
implementation (profanity/term list with normalization) and a hook
for third-party providers.
Usage::
from app.services.text_filter import TextFilter, FilterResult, \
ProfanityTextFilter, get_text_filter
# Default factory reads config from the path below
filt = get_text_filter()
result = filt.check("Hello wor1d!")
# result.action == FilterAction.ALLOW
"""
from __future__ import annotations
import json
import os
import re
import unicodedata
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any
# ---------------------------------------------------------------------------
# Enums / Data classes
# ---------------------------------------------------------------------------
class FilterAction(str, Enum):
"""Action the filter wants the caller to take."""
ALLOW = "allow"
REDACT = "redact"
REJECT = "reject"
@dataclass(frozen=True)
class FilterResult:
"""Immutable result of a single ``TextFilter.check()`` call.
Attributes:
action: What to do with the message.
filtered_text: The text after redaction (may equal *text* if
action is ``ALLOW``).
matched_terms: Sorted list of terms that triggered the filter.
"""
action: FilterAction
filtered_text: str
matched_terms: list[str] = field(default_factory=list)
@property
def is_allowed(self) -> bool:
return self.action == FilterAction.ALLOW
@property
def is_rejected(self) -> bool:
return self.action == FilterAction.REJECT
@property
def is_redacted(self) -> bool:
return self.action == FilterAction.REDACT
# ---------------------------------------------------------------------------
# Protocol / abstract base
# ---------------------------------------------------------------------------
class TextFilter(ABC):
"""Protocol for server-side text filters.
Every concrete filter must implement ``check()``. The default
implementation (``ProfanityTextFilter``) is deterministic and
requires no external service. A third-party provider can be
plugged in by subclassing ``TextFilter`` and overriding ``check``.
"""
@abstractmethod
def check(self, text: str) -> FilterResult:
"""Analyse *text* and return a ``FilterResult``.
Args:
text: Raw message text to filter.
Returns:
A ``FilterResult`` with the decision and any redacted text.
"""
# ---------------------------------------------------------------------------
# Normalisation helpers
# ---------------------------------------------------------------------------
# Common leetspeak substitutions (char -> set of possible originals)
_LEETSPEAK: dict[str, str] = {
"0": "o",
"1": "i",
"3": "e",
"4": "a",
"5": "s",
"7": "t",
"8": "b",
"9": "g",
"@": "a",
}
# Normalise separators that are often used to obfuscate words
_SEPARATORS = re.compile(r"[\s_\-\.]+")
def _normalize(text: str) -> str:
"""Case-fold, strip diacritics, and replace leetspeak / separators.
This produces a canonical form for matching against the term list.
"""
# Case-fold
text = text.casefold()
# Strip diacritics (e.g. é -> e)
text = unicodedata.normalize("NFKD", text)
text = "".join(c for c in text if not unicodedata.combining(c))
# Replace leetspeak
normalized = ""
for ch in text:
normalized += _LEETSPEAK.get(ch, ch)
# Collapse separators
normalized = _SEPARATORS.sub("", normalized)
return normalized
# ---------------------------------------------------------------------------
# Default implementation
# ---------------------------------------------------------------------------
# Default config path (relative to project root)
_DEFAULT_CONFIG_PATH = Path(__file__).resolve().parent.parent.parent / "data" / "text_filter_config.json"
@dataclass
class _FilterPolicy:
"""Filtering policy loaded from config."""
mode: str # "redact" or "reject"
redaction_char: str = "*"
# Minimum term length to consider (avoids matching single chars)
min_term_length: int = 3
# If True, match terms case-insensitively after normalisation
case_insensitive: bool = True
# If True, match whole words only (after normalisation)
whole_word: bool = True
class ProfanityTextFilter(TextFilter):
"""Deterministic text filter backed by a JSON term list.
The config file is a JSON object with:
.. code-block:: json
{
"mode": "redact",
"redaction_char": "*",
"min_term_length": 3,
"case_insensitive": true,
"whole_word": true,
"terms": ["badword", "another", ...]
}
Terms are matched after normalisation (case-fold, leetspeak
substitution, separator removal, diacritic stripping).
"""
def __init__(
self,
config_path: str | Path | None = None,
terms: list[str] | None = None,
policy: _FilterPolicy | None = None,
) -> None:
"""
Args:
config_path: Path to a JSON config file. Ignored if *terms*
is provided.
terms: Inline list of terms (overrides config file).
policy: Filter policy. Defaults to a redact policy.
"""
if terms is not None:
self._terms: list[str] = [t.lower() for t in terms]
elif config_path is not None:
self._terms = self._load_terms(Path(config_path))
else:
# Try default path
if _DEFAULT_CONFIG_PATH.exists():
self._terms = self._load_terms(_DEFAULT_CONFIG_PATH)
else:
self._terms = []
self._policy = policy or _FilterPolicy(mode="redact")
# Pre-compile patterns for each term
self._patterns: list[tuple[re.Pattern[str], str]] = []
for term in self._terms:
if len(term) < self._policy.min_term_length:
continue
if self._policy.whole_word:
pattern = re.compile(
r"(?<!\w)" + re.escape(term) + r"(?!\w)",
re.IGNORECASE if self._policy.case_insensitive else 0,
)
else:
pattern = re.compile(
re.escape(term),
re.IGNORECASE if self._policy.case_insensitive else 0,
)
self._patterns.append((pattern, term))
# -- config loading --------------------------------------------------
@staticmethod
def _load_terms(config_path: Path) -> list[str]:
"""Load terms from a JSON config file."""
with open(config_path, "r", encoding="utf-8") as f:
data: dict[str, Any] = json.load(f)
terms: list[str] = []
for raw in data.get("terms", []):
term = str(raw).strip()
if term:
terms.append(term.casefold())
return terms
# -- public API ------------------------------------------------------
def check(self, text: str) -> FilterResult:
"""Check *text* against the loaded term list.
Returns a ``FilterResult`` with the appropriate action,
filtered text, and matched terms.
"""
if not text:
return FilterResult(
action=FilterAction.ALLOW,
filtered_text=text,
matched_terms=[],
)
normalized = _normalize(text)
matched_terms: list[str] = []
# Collect all matches
for pattern, term in self._patterns:
if pattern.search(normalized):
if term not in matched_terms:
matched_terms.append(term)
if not matched_terms:
return FilterResult(
action=FilterAction.ALLOW,
filtered_text=text,
matched_terms=[],
)
matched_terms.sort()
# Build filtered text by replacing matched terms
filtered = text
for term in matched_terms:
norm_term = _normalize(term)
# Find all occurrences in the normalised text and replace in original
# We need to map positions from normalised text back to original
filtered = self._replace_terms(filtered, term, norm_term)
action = FilterAction.REJECT if self._policy.mode == "reject" else FilterAction.REDACT
return FilterResult(
action=action,
filtered_text=filtered,
matched_terms=matched_terms,
)
def _replace_terms(
self, text: str, term: str, norm_term: str
) -> str:
"""Replace occurrences of *term* in *text* with redaction chars.
Uses normalised text for matching but preserves original casing
in the output.
"""
norm_text = _normalize(text)
norm_term_len = len(norm_term)
result_chars = list(text)
start = 0
while start < len(norm_text):
idx = norm_text.find(norm_term, start)
if idx == -1:
break
# Replace characters in original text
for i in range(idx, min(idx + norm_term_len, len(result_chars))):
result_chars[i] = self._policy.redaction_char
start = idx + norm_term_len
return "".join(result_chars)
@property
def terms(self) -> list[str]:
"""Read-only access to the loaded term list."""
return list(self._terms)
@property
def policy(self) -> _FilterPolicy:
"""Read-only access to the filter policy."""
return self._policy
# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------
def get_text_filter(
config_path: str | Path | None = None,
terms: list[str] | None = None,
policy: _FilterPolicy | None = None,
) -> TextFilter:
"""Return a ``TextFilter`` instance.
By default this creates a ``ProfanityTextFilter``. To use a
third-party provider, subclass ``TextFilter`` and pass an
instance here instead.
"""
return ProfanityTextFilter(
config_path=config_path,
terms=terms,
policy=policy,
)
+193
View File
@@ -0,0 +1,193 @@
"""Tests for the Redis-backed token-bucket rate limiter.
These tests use fakeredis (Lua enabled via lupa) so no live Redis server is
required.
"""
import asyncio
import sys
import pytest
from fastapi import Depends, FastAPI, Request
from fastapi.testclient import TestClient
sys.path.insert(0, "/workspace")
import fakeredis.aioredis
from redis.asyncio import Redis
from redis.exceptions import ConnectionError
from app.rate_limit import (
RateLimit,
RateLimitConfig,
RateLimiter,
default_user_scope,
named_limit,
rate_limit,
)
@pytest.fixture()
def fake_redis():
return fakeredis.aioredis.FakeRedis(decode_responses=True)
def make_request(client_ip: str = "127.0.0.1") -> Request:
scope = {
"type": "http",
"method": "GET",
"path": "/",
"headers": [],
"client": (client_ip, 12345),
}
return Request(scope)
@pytest.mark.asyncio
async def test_token_bucket_allows_burst_then_exhausts(fake_redis):
limiter = RateLimiter(
RateLimitConfig(
limit=RateLimit(name="burst", capacity=3, refill_rate=0.1),
),
redis=fake_redis,
)
request = make_request()
results = []
for _ in range(4):
result = await limiter.check(request, user_id=42)
results.append(result)
assert [r.allowed for r in results] == [True, True, True, False]
assert results[0].remaining == 2
assert results[1].remaining == 1
assert results[2].remaining == 0
assert results[3].retry_after > 0
@pytest.mark.asyncio
async def test_retry_after_reflects_refill_rate(fake_redis):
limiter = RateLimiter(
RateLimitConfig(
limit=RateLimit(name="slow", capacity=1, refill_rate=0.5),
),
redis=fake_redis,
)
request = make_request()
assert (await limiter.check(request, user_id=7)).allowed is True
denied = await limiter.check(request, user_id=7)
assert denied.allowed is False
# 1 token / 0.5 tokens per second = 2 seconds.
assert denied.retry_after == 2
@pytest.mark.asyncio
async def test_separate_scopes_get_separate_buckets(fake_redis):
limiter = RateLimiter(
RateLimitConfig(
limit=RateLimit(name="scoped", capacity=1, refill_rate=0.1),
),
redis=fake_redis,
)
request = make_request()
assert (await limiter.check(request, user_id=1)).allowed is True
assert (await limiter.check(request, user_id=2)).allowed is True
assert (await limiter.check(request, user_id=1)).allowed is False
@pytest.mark.asyncio
async def test_fail_open_returns_allowed_when_redis_unavailable():
class BrokenRedis(Redis):
async def eval(self, *args, **kwargs):
raise ConnectionError("down")
limiter = RateLimiter(
RateLimitConfig(
limit=RateLimit(name="open", capacity=1, refill_rate=0.1, fail_open=True),
),
redis=BrokenRedis(),
)
result = await limiter.check(make_request(), user_id=1)
assert result.allowed is True
@pytest.mark.asyncio
async def test_fail_closed_returns_denied_with_negative_retry_after():
class BrokenRedis(Redis):
async def eval(self, *args, **kwargs):
raise ConnectionError("down")
limiter = RateLimiter(
RateLimitConfig(
limit=RateLimit(name="closed", capacity=1, refill_rate=0.1, fail_open=False),
),
redis=BrokenRedis(),
)
result = await limiter.check(make_request(), user_id=1)
assert result.allowed is False
assert result.retry_after == -1
@pytest.mark.asyncio
async def test_fail_closed_call_raises_503(fake_redis):
class BrokenRedis(Redis):
async def eval(self, *args, **kwargs):
raise ConnectionError("down")
limiter = RateLimiter(
RateLimitConfig(
limit=RateLimit(name="closed", capacity=1, refill_rate=0.1, fail_open=False),
),
redis=BrokenRedis(),
)
with pytest.raises(Exception) as exc_info:
await limiter(make_request(), user_id=1)
assert exc_info.value.status_code == 503
def test_named_limit_builds_expected_config():
config = named_limit(
"chat_send",
capacity=5,
refill_rate=0.2,
fail_open=False,
)
assert config.limit.name == "chat_send"
assert config.limit.capacity == 5
assert config.limit.refill_rate == 0.2
assert config.limit.fail_open is False
assert config.scope is default_user_scope
def test_invalid_limit_definition_rejected():
with pytest.raises(ValueError):
named_limit("bad", capacity=0, refill_rate=1.0)
with pytest.raises(ValueError):
named_limit("bad", capacity=1, refill_rate=0)
def test_dependency_integration_returns_429_with_retry_after():
async def fake_user_id() -> int:
return 99
chat_limit = rate_limit(
named_limit("chat_send_test", capacity=2, refill_rate=0.1, fail_open=False),
user_id_dependency=Depends(fake_user_id),
)
app = FastAPI()
@app.get("/probe")
async def probe(_: None = Depends(chat_limit)):
return {"ok": True}
client = TestClient(app)
assert client.get("/probe").status_code == 200
assert client.get("/probe").status_code == 200
denied = client.get("/probe")
assert denied.status_code == 429
assert denied.headers["retry-after"]
assert denied.json() == {"detail": "Rate limit exceeded"}