Audit entries written only inside the caller's transaction, database throttle counters that refuse to run inside a transaction that could roll back, and a database session store with the user key in its own column, end_sessions and idle and absolute limits. Ticket 1.3. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Sk72AxLwQnYttF6VFSxQLt
99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
"""Database throttle counters, checked before any outbound call (AD-7)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Callable
|
|
from datetime import datetime, timedelta
|
|
|
|
from django.db import connection, transaction
|
|
from django.db.models import F, Q
|
|
from django.utils import timezone
|
|
|
|
from .models import ThrottleCounter
|
|
|
|
# Defaults, ready to move into the settings registry (AD-13).
|
|
THROTTLE_MAX_FAILURES = 5
|
|
THROTTLE_PAUSE_SECONDS = 15 * 60
|
|
|
|
Clock = Callable[[], datetime]
|
|
|
|
|
|
class Throttled(Exception):
|
|
def __init__(self, retry_after: int) -> None:
|
|
self.retry_after = retry_after
|
|
super().__init__(f"Too many attempts; retry in {retry_after} seconds")
|
|
|
|
|
|
def _subjects(user: str | None, ip: str | None) -> list[str]:
|
|
if not user and not ip:
|
|
raise ValueError("throttle needs a user or an ip; refusing to do nothing")
|
|
subjects = []
|
|
if user:
|
|
subjects.append(f"user:{user.lower()}")
|
|
if ip:
|
|
subjects.append(f"ip:{ip}")
|
|
return subjects
|
|
|
|
|
|
def check(
|
|
endpoint: str,
|
|
*,
|
|
user: str | None = None,
|
|
ip: str | None = None,
|
|
clock: Clock = timezone.now,
|
|
) -> None:
|
|
"""Raise Throttled if the user or the IP is paused at this endpoint. Never writes."""
|
|
now = clock()
|
|
blocked = ThrottleCounter.objects.filter(
|
|
endpoint=endpoint,
|
|
subject__in=_subjects(user, ip),
|
|
blocked_until__gt=now,
|
|
).values_list("blocked_until", flat=True)
|
|
latest = max(blocked, default=None)
|
|
if latest is not None:
|
|
raise Throttled(max(1, math.ceil((latest - now).total_seconds())))
|
|
|
|
|
|
def hit(
|
|
endpoint: str,
|
|
*,
|
|
user: str | None = None,
|
|
ip: str | None = None,
|
|
clock: Clock = timezone.now,
|
|
max_failures: int = THROTTLE_MAX_FAILURES,
|
|
pause_seconds: int = THROTTLE_PAUSE_SECONDS,
|
|
) -> None:
|
|
"""Count one failure for the user and the IP; pause a subject that reaches the limit."""
|
|
subjects = _subjects(user, ip)
|
|
if connection.in_atomic_block:
|
|
raise RuntimeError(
|
|
"throttle.hit() must be called after the failed work was rolled back, "
|
|
"outside any open transaction"
|
|
)
|
|
now = clock()
|
|
window = timedelta(seconds=pause_seconds)
|
|
with transaction.atomic():
|
|
for subject in subjects:
|
|
counter, _ = ThrottleCounter.objects.get_or_create(
|
|
endpoint=endpoint, subject=subject, defaults={"window_started": now}
|
|
)
|
|
qs = ThrottleCounter.objects.filter(pk=counter.pk)
|
|
# Start a fresh window when the old one ran out or the pause is over.
|
|
qs.filter(Q(window_started__lte=now - window) | Q(blocked_until__lte=now)).update(
|
|
count=0, window_started=now, blocked_until=None
|
|
)
|
|
qs.update(count=F("count") + 1)
|
|
qs.filter(count__gte=max_failures, blocked_until__isnull=True).update(
|
|
blocked_until=now + window
|
|
)
|
|
|
|
|
|
def reset(endpoint: str, *, user: str) -> None:
|
|
"""Clear the user's counter (call after a success).
|
|
|
|
IP counters expire on their own and are never reset: a valid sign-in must not let an
|
|
attacker with one account clear the IP counter.
|
|
"""
|
|
ThrottleCounter.objects.filter(endpoint=endpoint, subject__in=_subjects(user, None)).delete()
|