22import time
33import threading
44import math
5+ from datetime import UTC , datetime , timedelta
56from logging import getLogger
6- from typing import Tuple
7+ from typing import Protocol , Tuple , runtime_checkable
78
89from fastapi import Request , Form
910from pydantic import EmailStr
1011from 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
1217logger = getLogger ("uvicorn.error" )
1318load_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
1648class 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
166309def 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