Skip to content

Commit bb982ca

Browse files
fix: add Postgres-backed rate limiter for multi-worker deployments
Store rate-limit attempts in a shared table when RATE_LIMIT_BACKEND=postgres so configured limits apply cluster-wide instead of per process.
1 parent a0b6f8b commit bb982ca

4 files changed

Lines changed: 185 additions & 10 deletions

File tree

.env.example

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@ DB_NAME=
2626
# Domain for production (used by Caddy for TLS)
2727
DOMAIN=
2828

29+
# Rate limit storage: memory (single process) or postgres (multi-worker)
30+
# RATE_LIMIT_BACKEND=memory
31+
2932
# Resend
3033
RESEND_API_KEY=
3134
EMAIL_FROM=

tests/utils/test_rate_limit.py

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from unittest.mock import patch
44

55
import utils.core.rate_limit as rate_limit_module
6-
from utils.core.rate_limit import RateLimitWindow
6+
from utils.core.rate_limit import PostgresRateLimitWindow, RateLimitWindow
77

88

99
# ---------------------------------------------------------------------------
@@ -162,3 +162,21 @@ def test_module_limiters_honor_env_configuration(monkeypatch):
162162
assert rate_limit_module.forgot_password_email_limiter.window_seconds == 91
163163

164164
importlib.reload(rate_limit_module)
165+
166+
167+
def test_postgres_rate_limit_window_enforces_limit(engine, env_vars, monkeypatch):
168+
monkeypatch.setenv("RATE_LIMIT_BACKEND", "postgres")
169+
importlib.reload(rate_limit_module)
170+
171+
limiter = PostgresRateLimitWindow("test_scope", max_attempts=2, window_seconds=60)
172+
limiter.clear()
173+
limiter.record("shared-key")
174+
limiter.record("shared-key")
175+
assert limiter.check("shared-key")[0] is True
176+
177+
second = PostgresRateLimitWindow("test_scope", max_attempts=2, window_seconds=60)
178+
assert second.check("shared-key")[0] is True
179+
180+
limiter.clear()
181+
monkeypatch.delenv("RATE_LIMIT_BACKEND", raising=False)
182+
importlib.reload(rate_limit_module)

utils/core/models.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,17 @@ def is_expired(self) -> bool:
156156
return datetime.now(UTC) > self.expires_at.replace(tzinfo=UTC)
157157

158158

159+
class RateLimitAttempt(SQLModel, table=True):
160+
"""Shared rate-limit counter row for multi-worker deployments."""
161+
162+
__table_args__ = {"schema": "private"}
163+
164+
id: Optional[int] = Field(default=None, primary_key=True)
165+
scope: str = Field(index=True)
166+
key: str = Field(index=True)
167+
attempted_at: datetime = Field(default_factory=utc_now, index=True)
168+
169+
159170
# --- Public database models ---
160171

161172

utils/core/rate_limit.py

Lines changed: 152 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,48 @@
22
import time
33
import threading
44
import math
5+
from datetime import UTC, datetime, timedelta
56
from logging import getLogger
6-
from typing import Tuple
7+
from typing import Protocol, Tuple, runtime_checkable
78

89
from fastapi import Request, Form
910
from pydantic import EmailStr
1011
from dotenv import load_dotenv
12+
from sqlmodel import Session, col, create_engine, delete, select
13+
14+
from utils.core.db import get_connection_url
15+
from utils.core.models import RateLimitAttempt
1116

1217
logger = getLogger("uvicorn.error")
1318
load_dotenv()
1419

20+
_rate_limit_engine = None
21+
22+
23+
def _get_rate_limit_engine():
24+
global _rate_limit_engine
25+
if _rate_limit_engine is None:
26+
_rate_limit_engine = create_engine(get_connection_url())
27+
return _rate_limit_engine
28+
29+
30+
@runtime_checkable
31+
class RateLimiter(Protocol):
32+
max_attempts: int
33+
window_seconds: int
34+
35+
def check(self, key: str) -> Tuple[bool, int]: ...
36+
37+
def record(self, key: str) -> None: ...
38+
39+
def remaining(self, key: str) -> int: ...
40+
41+
def reset(self, key: str) -> None: ...
42+
43+
def prune(self) -> None: ...
44+
45+
def clear(self) -> None: ...
46+
1547

1648
class RateLimitWindow:
1749
"""
@@ -115,6 +147,100 @@ def prune(self) -> None:
115147
self._prune_stale_keys(now)
116148
self._next_prune_at = now + self.prune_interval_seconds
117149

150+
def clear(self) -> None:
151+
"""Clear all recorded attempts."""
152+
with self._lock:
153+
self._attempts.clear()
154+
155+
156+
class PostgresRateLimitWindow:
157+
"""
158+
Sliding-window rate limiter backed by PostgreSQL.
159+
160+
Use when running multiple workers or replicas so configured limits apply
161+
cluster-wide instead of per process.
162+
"""
163+
164+
def __init__(self, scope: str, max_attempts: int, window_seconds: int):
165+
self.scope = scope
166+
self.max_attempts = max_attempts
167+
self.window_seconds = window_seconds
168+
169+
def _cutoff(self, now: datetime) -> datetime:
170+
return now - timedelta(seconds=self.window_seconds)
171+
172+
def _recent_attempts(self, session: Session, key: str, now: datetime):
173+
return session.exec(
174+
select(RateLimitAttempt)
175+
.where(
176+
col(RateLimitAttempt.scope) == self.scope,
177+
col(RateLimitAttempt.key) == key,
178+
col(RateLimitAttempt.attempted_at) > self._cutoff(now),
179+
)
180+
.order_by(col(RateLimitAttempt.attempted_at))
181+
).all()
182+
183+
def check(self, key: str) -> Tuple[bool, int]:
184+
now = datetime.now(UTC)
185+
with Session(_get_rate_limit_engine()) as session:
186+
attempts = self._recent_attempts(session, key, now)
187+
if len(attempts) >= self.max_attempts:
188+
oldest = attempts[0].attempted_at
189+
if oldest.tzinfo is None:
190+
oldest = oldest.replace(tzinfo=UTC)
191+
retry_after = math.ceil(
192+
(
193+
oldest + timedelta(seconds=self.window_seconds) - now
194+
).total_seconds()
195+
)
196+
return True, max(retry_after, 1)
197+
return False, 0
198+
199+
def record(self, key: str) -> None:
200+
with Session(_get_rate_limit_engine()) as session:
201+
session.add(
202+
RateLimitAttempt(
203+
scope=self.scope, key=key, attempted_at=datetime.now(UTC)
204+
)
205+
)
206+
session.commit()
207+
208+
def remaining(self, key: str) -> int:
209+
now = datetime.now(UTC)
210+
with Session(_get_rate_limit_engine()) as session:
211+
attempts = self._recent_attempts(session, key, now)
212+
return max(0, self.max_attempts - len(attempts))
213+
214+
def reset(self, key: str) -> None:
215+
with Session(_get_rate_limit_engine()) as session:
216+
session.exec(
217+
delete(RateLimitAttempt).where(
218+
col(RateLimitAttempt.scope) == self.scope,
219+
col(RateLimitAttempt.key) == key,
220+
)
221+
)
222+
session.commit()
223+
224+
def prune(self) -> None:
225+
cutoff = self._cutoff(datetime.now(UTC))
226+
with Session(_get_rate_limit_engine()) as session:
227+
session.exec(
228+
delete(RateLimitAttempt).where(
229+
col(RateLimitAttempt.scope) == self.scope,
230+
col(RateLimitAttempt.attempted_at) <= cutoff,
231+
)
232+
)
233+
session.commit()
234+
235+
def clear(self) -> None:
236+
with Session(_get_rate_limit_engine()) as session:
237+
session.exec(
238+
delete(RateLimitAttempt).where(
239+
col(RateLimitAttempt.scope) == self.scope
240+
)
241+
)
242+
session.commit()
243+
118244

119245
# --- Configuration helpers ---
120246

@@ -131,25 +257,42 @@ def _int_env(name: str, default: int) -> int:
131257
return default
132258

133259

260+
def _rate_limit_backend() -> str:
261+
return os.environ.get("RATE_LIMIT_BACKEND", "memory").lower()
262+
263+
264+
def _make_rate_limiter(
265+
scope: str, max_attempts: int, window_seconds: int
266+
) -> RateLimiter:
267+
if _rate_limit_backend() == "postgres":
268+
return PostgresRateLimitWindow(scope, max_attempts, window_seconds)
269+
return RateLimitWindow(max_attempts=max_attempts, window_seconds=window_seconds)
270+
271+
134272
# --- Shared limiter instances ---
135273

136-
login_ip_limiter = RateLimitWindow(
274+
login_ip_limiter = _make_rate_limiter(
275+
"login_ip",
137276
max_attempts=_int_env("LOGIN_IP_LIMIT", 10),
138277
window_seconds=_int_env("LOGIN_IP_WINDOW_SECONDS", 60),
139278
)
140-
login_email_limiter = RateLimitWindow(
279+
login_email_limiter = _make_rate_limiter(
280+
"login_email",
141281
max_attempts=_int_env("LOGIN_EMAIL_LIMIT", 5),
142282
window_seconds=_int_env("LOGIN_EMAIL_WINDOW_SECONDS", 60),
143283
)
144-
register_ip_limiter = RateLimitWindow(
284+
register_ip_limiter = _make_rate_limiter(
285+
"register_ip",
145286
max_attempts=_int_env("REGISTER_IP_LIMIT", 5),
146287
window_seconds=_int_env("REGISTER_IP_WINDOW_SECONDS", 60),
147288
)
148-
forgot_password_ip_limiter = RateLimitWindow(
289+
forgot_password_ip_limiter = _make_rate_limiter(
290+
"forgot_password_ip",
149291
max_attempts=_int_env("FORGOT_PASSWORD_IP_LIMIT", 5),
150292
window_seconds=_int_env("FORGOT_PASSWORD_IP_WINDOW_SECONDS", 60),
151293
)
152-
forgot_password_email_limiter = RateLimitWindow(
294+
forgot_password_email_limiter = _make_rate_limiter(
295+
"forgot_password_email",
153296
max_attempts=_int_env("FORGOT_PASSWORD_EMAIL_LIMIT", 3),
154297
window_seconds=_int_env("FORGOT_PASSWORD_EMAIL_WINDOW_SECONDS", 60),
155298
)
@@ -164,9 +307,9 @@ def _int_env(name: str, default: int) -> int:
164307

165308

166309
def clear_all_rate_limiters() -> None:
167-
"""Clear all in-memory rate limiter state."""
310+
"""Clear all rate limiter state for the active backend."""
168311
for limiter in _ALL_LIMITERS:
169-
limiter._attempts.clear()
312+
limiter.clear()
170313

171314

172315
# --- Dependency helpers ---
@@ -184,7 +327,7 @@ def get_client_ip(request: Request) -> str:
184327
return "unknown"
185328

186329

187-
def _enforce_rate_limit(limiter: RateLimitWindow, key: str, scope: str) -> int:
330+
def _enforce_rate_limit(limiter: RateLimiter, key: str, scope: str) -> int:
188331
"""
189332
Check the limiter for the given key. If limited, raise RateLimitError.
190333
Otherwise, record the attempt and return.

0 commit comments

Comments
 (0)