Compare commits

..

3 Commits

6 changed files with 489 additions and 84 deletions
+1 -1
View File
@@ -8,7 +8,7 @@ from sqlalchemy.ext.asyncio import create_async_engine
from app.config import settings from app.config import settings
from app.database import Base from app.database import Base
from app.models import user, session, world, asset, wallet # noqa: F401 from app.models import user, session, world, asset # noqa: F401
# this is the Alembic Config object, which provides # this is the Alembic Config object, which provides
# access to the values within the .ini file in use. # access to the values within the .ini file in use.
+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
+1 -10
View File
@@ -2,14 +2,5 @@ from app.models.user import User
from app.models.session import Session from app.models.session import Session
from app.models.world import World from app.models.world import World
from app.models.asset import Asset from app.models.asset import Asset
from app.models.wallet import Transaction, TransactionType, Wallet
__all__ = [ __all__ = ["User", "Session", "World", "Asset"]
"User",
"Session",
"World",
"Asset",
"Wallet",
"Transaction",
"TransactionType",
]
-73
View File
@@ -1,73 +0,0 @@
"""Wallet and transaction ledger models."""
import enum
from sqlalchemy import (
BigInteger,
CheckConstraint,
Column,
DateTime,
ForeignKey,
Integer,
Text,
)
from sqlalchemy.dialects.postgresql import ENUM, JSONB
from sqlalchemy.sql import func
from app.database import Base
class TransactionType(str, enum.Enum):
CREDIT_PURCHASE = "credit_purchase"
MARKETPLACE_PURCHASE = "marketplace_purchase"
MARKETPLACE_SALE = "marketplace_sale"
ESCROW_RELEASE = "escrow_release"
PAYOUT = "payout"
REFUND = "refund"
ADMIN_ADJUSTMENT = "admin_adjustment"
class Wallet(Base):
__tablename__ = "wallets"
__table_args__ = (
CheckConstraint("balance >= 0", name="ck_wallets_balance_nonnegative"),
)
id = Column(BigInteger, primary_key=True, index=True)
user_id = Column(
Integer,
ForeignKey("users.id", ondelete="CASCADE"),
unique=True,
index=True,
nullable=False,
)
balance = Column(BigInteger, nullable=False, default=0, server_default="0")
created_at = Column(DateTime(timezone=True), server_default=func.now())
updated_at = Column(
DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
)
class Transaction(Base):
__tablename__ = "transactions"
id = Column(BigInteger, primary_key=True, index=True)
wallet_id = Column(
BigInteger, ForeignKey("wallets.id"), nullable=False
)
type = Column(
ENUM(
TransactionType,
name="transaction_type",
create_type=False,
values_callable=lambda enum_cls: [member.value for member in enum_cls],
),
nullable=False,
)
amount = Column(BigInteger, nullable=False)
balance_after = Column(BigInteger, nullable=False)
reference_type = Column(Text, nullable=True)
reference_id = Column(Text, nullable=True)
idempotency_key = Column(Text, unique=True, index=True, nullable=True)
metadata_json = Column(JSONB, nullable=True)
created_at = Column(DateTime(timezone=True), server_default=func.now())
+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
+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"}