66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
from datetime import datetime, timedelta, timezone
|
|
import hashlib
|
|
import hmac
|
|
from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network, ip_address
|
|
from typing import Protocol
|
|
|
|
from fastapi import HTTPException, Request
|
|
from sqlalchemy import delete, func, select, text
|
|
from sqlalchemy.orm import Session
|
|
|
|
from .models import SubmissionAttempt
|
|
|
|
|
|
class RateLimitConfig(Protocol):
|
|
rate_limit_secret: str
|
|
trusted_proxy_cidrs: list[str]
|
|
|
|
|
|
def is_trusted_proxy(address: str, trusted_cidrs: list[str]) -> bool:
|
|
try:
|
|
addr = IPv4Address(address) if ":" not in address else IPv6Address(address)
|
|
except ValueError:
|
|
return False
|
|
for cidr in trusted_cidrs:
|
|
try:
|
|
network = IPv4Network(cidr) if ":" not in cidr else IPv6Network(cidr)
|
|
if addr in network:
|
|
return True
|
|
except ValueError:
|
|
continue
|
|
return False
|
|
|
|
|
|
def client_address(request: Request, trusted_cidrs: list[str]) -> str:
|
|
client = request.client.host if request.client else "unknown"
|
|
forwarded = request.headers.get("x-forwarded-for")
|
|
if forwarded and request.client and is_trusted_proxy(request.client.host, trusted_cidrs):
|
|
candidate = forwarded.split(",")[0].strip()
|
|
try:
|
|
return str(ip_address(candidate))
|
|
except ValueError:
|
|
return client
|
|
return client
|
|
|
|
|
|
def check_rate_limit(request: Request, db: Session, config: RateLimitConfig) -> None:
|
|
now = datetime.now(timezone.utc)
|
|
cutoff = now - timedelta(minutes=10)
|
|
client = client_address(request, config.trusted_proxy_cidrs)
|
|
client_hash = hmac.new(config.rate_limit_secret.encode(), client.encode(), hashlib.sha256).hexdigest()
|
|
if db.get_bind().dialect.name == "postgresql":
|
|
lock_key = int(client_hash[:16], 16) & 0x7FFF_FFFF_FFFF_FFFF
|
|
db.execute(text("SELECT pg_advisory_xact_lock(:lock_key)"), {"lock_key": lock_key})
|
|
db.execute(delete(SubmissionAttempt).where(SubmissionAttempt.created_at < now - timedelta(days=1)))
|
|
recent = db.scalar(
|
|
select(func.count()).select_from(SubmissionAttempt).where(
|
|
SubmissionAttempt.client_hash == client_hash,
|
|
SubmissionAttempt.created_at >= cutoff,
|
|
)
|
|
) or 0
|
|
if recent >= 5:
|
|
db.commit()
|
|
raise HTTPException(status_code=429, detail="too many submissions")
|
|
db.add(SubmissionAttempt(client_hash=client_hash, created_at=now))
|
|
db.commit()
|