180 lines
6.1 KiB
Python
180 lines
6.1 KiB
Python
"""Exponential back-off throttle service for the login endpoint.
|
|
|
|
Design
|
|
------
|
|
Failures are tracked per (scope, key) in the ``auth_login_throttle`` table:
|
|
|
|
* scope ``'ip'`` → key is the client IP address
|
|
* scope ``'user'`` → key is the username
|
|
|
|
On each login attempt the endpoint checks **both** keys and takes the **maximum**
|
|
remaining wait (in seconds). A wait > 0 means the caller is still inside a
|
|
back-off window and should receive 429.
|
|
|
|
Exponential formula
|
|
-------------------
|
|
The first ``N_FREE`` failures are free (no delay). After that::
|
|
|
|
wait = min(BACKOFF_CAP_S, BACKOFF_BASE_S * 2 ** (failures - N_FREE))
|
|
|
|
Constants (module-level so they can be tuned without touching the algorithm):
|
|
|
|
* ``N_FREE`` = 3 — failures before any delay starts
|
|
* ``BACKOFF_BASE_S`` = 1 — base wait in seconds (1st delayed attempt → 2 s)
|
|
* ``BACKOFF_CAP_S`` = 900 — hard cap (15 min)
|
|
|
|
Example delay schedule:
|
|
failures 1, 2, 3 → 0 s (free)
|
|
failures 4 → 2 s
|
|
failures 5 → 4 s
|
|
failures 6 → 8 s
|
|
failures 7 → 16 s
|
|
…
|
|
failures 13+ → 900 s (capped)
|
|
|
|
Timezone notes
|
|
--------------
|
|
The model uses ``DateTime(timezone=True)``. SQLite stores timestamps without
|
|
timezone info, so when rows are read back the column value may be timezone-naive.
|
|
We normalise every stored value to UTC-aware via ``_as_utc`` before comparing
|
|
against ``datetime.now(UTC)``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from sqlalchemy import delete, select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.models.auth_throttle import LoginThrottle
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tunable constants
|
|
# ---------------------------------------------------------------------------
|
|
|
|
N_FREE: int = 3 # failures before any backoff starts
|
|
BACKOFF_BASE_S: int = 1 # seconds — wait after the first delayed failure
|
|
BACKOFF_CAP_S: int = 900 # seconds — maximum wait (15 min)
|
|
|
|
# Internal scope labels (keep in sync with model constraint)
|
|
_SCOPE_IP = "ip"
|
|
_SCOPE_USER = "user"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public API
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def check_and_get_wait(session: Session, *, ip: str, username: str) -> int:
|
|
"""Return the number of seconds the caller must still wait, or 0 if allowed.
|
|
|
|
Checks both the IP row and the username row; returns the *larger* remaining
|
|
wait so that both keys must have their back-off window expire before the
|
|
caller is permitted to try again.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
wait_ip = _remaining_wait(session, scope=_SCOPE_IP, key=ip, now=now)
|
|
wait_user = _remaining_wait(session, scope=_SCOPE_USER, key=username, now=now)
|
|
return max(wait_ip, wait_user)
|
|
|
|
|
|
def register_failure(session: Session, *, ip: str, username: str) -> None:
|
|
"""Record one login failure for both the IP and the username keys.
|
|
|
|
Updates (or creates) the throttle row for each key, increments the failure
|
|
counter, and recomputes ``next_allowed_at`` using the exponential formula.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
_upsert_failure(session, scope=_SCOPE_IP, key=ip, now=now)
|
|
_upsert_failure(session, scope=_SCOPE_USER, key=username, now=now)
|
|
session.commit()
|
|
|
|
|
|
def clear(session: Session, *, ip: str, username: str) -> None:
|
|
"""Delete the throttle rows for the given IP and username (successful login).
|
|
|
|
It is safe to call this even if no rows exist for the given keys.
|
|
"""
|
|
session.execute(
|
|
delete(LoginThrottle).where(
|
|
(LoginThrottle.scope == _SCOPE_IP) & (LoginThrottle.key == ip)
|
|
| (LoginThrottle.scope == _SCOPE_USER) & (LoginThrottle.key == username)
|
|
)
|
|
)
|
|
session.commit()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _remaining_wait(session: Session, *, scope: str, key: str, now: datetime) -> int:
|
|
"""Return the number of seconds that ``key`` in ``scope`` must still wait.
|
|
|
|
Returns 0 if the key has no throttle row or if the back-off window has
|
|
already expired.
|
|
"""
|
|
row = session.scalar(
|
|
select(LoginThrottle)
|
|
.where(LoginThrottle.scope == scope, LoginThrottle.key == key)
|
|
.limit(1)
|
|
)
|
|
if row is None or row.next_allowed_at is None:
|
|
return 0
|
|
|
|
next_allowed = _as_utc(row.next_allowed_at)
|
|
if next_allowed <= now:
|
|
return 0
|
|
|
|
remaining = (next_allowed - now).total_seconds()
|
|
# Round up so callers are never told "0 seconds" while still blocked.
|
|
return math.ceil(remaining)
|
|
|
|
|
|
def _compute_next_allowed_at(failures: int, now: datetime) -> datetime | None:
|
|
"""Compute the earliest time the key should be allowed to try again.
|
|
|
|
Returns ``None`` for the first ``N_FREE`` failures (no delay).
|
|
"""
|
|
if failures <= N_FREE:
|
|
return None
|
|
wait_s = min(BACKOFF_CAP_S, BACKOFF_BASE_S * (2 ** (failures - N_FREE)))
|
|
return now + timedelta(seconds=wait_s)
|
|
|
|
|
|
def _upsert_failure(session: Session, *, scope: str, key: str, now: datetime) -> None:
|
|
"""Create or update the throttle row for (scope, key) with one more failure."""
|
|
row = session.scalar(
|
|
select(LoginThrottle)
|
|
.where(LoginThrottle.scope == scope, LoginThrottle.key == key)
|
|
.limit(1)
|
|
)
|
|
if row is None:
|
|
new_failures = 1
|
|
row = LoginThrottle(
|
|
scope=scope,
|
|
key=key,
|
|
failures=new_failures,
|
|
first_failed_at=now,
|
|
last_failed_at=now,
|
|
next_allowed_at=_compute_next_allowed_at(new_failures, now),
|
|
)
|
|
session.add(row)
|
|
else:
|
|
row.failures += 1
|
|
row.last_failed_at = now
|
|
row.next_allowed_at = _compute_next_allowed_at(row.failures, now)
|
|
|
|
|
|
def _as_utc(value: datetime | None) -> datetime | None:
|
|
"""Normalise a possibly-naive datetime to UTC-aware."""
|
|
if value is None:
|
|
return None
|
|
if value.tzinfo is None:
|
|
return value.replace(tzinfo=UTC)
|
|
return value.astimezone(UTC)
|