296 lines
12 KiB
Python
296 lines
12 KiB
Python
"""Process-safe shared rate limiter for auth endpoints.
|
|
|
|
Each action/identity pair is stored as an atomic JSON record. The record store
|
|
uses an OS-level lock, so forked Gunicorn workers cannot independently accept
|
|
attempts from stale in-memory counters or overwrite one another's snapshots.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import time
|
|
|
|
from ..config import Config
|
|
from ..auth.users import AuthError, normalize_identifier
|
|
from ..storage.store import JsonStore, StoreError, StoreNotFoundError
|
|
|
|
# Kept as a compatibility/debug hook for existing tests and callers. It is
|
|
# deliberately not authoritative; the durable JsonStore is read on every
|
|
# check, so clearing this process-local mapping cannot reset the limit.
|
|
_mem: dict[str, list[float]] = {}
|
|
|
|
|
|
def _key(action: str, ident: str) -> str:
|
|
"""Encode the pair without allowing action/identity delimiter collisions."""
|
|
return json.dumps([action, ident], ensure_ascii=False, separators=(",", ":"))
|
|
|
|
|
|
def _record_store() -> JsonStore:
|
|
store = JsonStore(Config.DATA_DIR / "ratelimit")
|
|
_migrate_legacy_state(store)
|
|
return store
|
|
|
|
|
|
def _legacy_paths():
|
|
legacy = Config.DATA_DIR / "ratelimit.json"
|
|
return legacy, legacy.with_name(f"{legacy.name}.migrated")
|
|
|
|
|
|
def _validated_timestamps(raw_timestamps: object) -> list[float]:
|
|
if not isinstance(raw_timestamps, list):
|
|
raise StoreError("legacy rate-limit timestamps are invalid")
|
|
timestamps: list[float] = []
|
|
for timestamp in raw_timestamps:
|
|
if isinstance(timestamp, bool) or not isinstance(timestamp, (int, float)):
|
|
raise StoreError("legacy rate-limit timestamp is invalid")
|
|
try:
|
|
value = float(timestamp)
|
|
except (OverflowError, TypeError, ValueError) as exc:
|
|
raise StoreError("legacy rate-limit timestamp is invalid") from exc
|
|
if not math.isfinite(value):
|
|
raise StoreError("legacy rate-limit timestamp is invalid")
|
|
timestamps.append(value)
|
|
return timestamps
|
|
|
|
|
|
_LEGACY_ACTIONS = (
|
|
"password-change:user",
|
|
"password-change:ip",
|
|
"login:user",
|
|
"login:ip",
|
|
"chat:user",
|
|
)
|
|
|
|
|
|
def _legacy_key_aliases(legacy_key: object) -> set[str]:
|
|
if not isinstance(legacy_key, str) or not legacy_key:
|
|
raise StoreError("legacy rate-limit key is invalid")
|
|
try:
|
|
structured = json.loads(legacy_key)
|
|
except (TypeError, ValueError):
|
|
structured = None
|
|
if (
|
|
isinstance(structured, list)
|
|
and len(structured) == 2
|
|
and all(isinstance(value, str) for value in structured)
|
|
):
|
|
try:
|
|
return {
|
|
_key(
|
|
normalize_identifier(structured[0]),
|
|
normalize_identifier(structured[1]),
|
|
)
|
|
}
|
|
except AuthError as exc:
|
|
raise StoreError("legacy rate-limit key is invalid") from exc
|
|
|
|
# The legacy format was action:identity. Only the known action prefixes
|
|
# are accepted, so an identity containing colons remains part of the
|
|
# identity and cannot create aliases for other action/identity pairs.
|
|
for action in sorted(_LEGACY_ACTIONS, key=len, reverse=True):
|
|
prefix = f"{action}:"
|
|
if not legacy_key.startswith(prefix):
|
|
continue
|
|
raw_ident = legacy_key[len(prefix) :]
|
|
try:
|
|
return {_key(normalize_identifier(action), normalize_identifier(raw_ident))}
|
|
except AuthError as exc:
|
|
raise StoreError("legacy rate-limit key is invalid") from exc
|
|
raise StoreError("legacy rate-limit key is invalid")
|
|
|
|
|
|
def _read_migration_marker(marker_path) -> str:
|
|
try:
|
|
lines = marker_path.read_text(encoding="utf-8").splitlines()
|
|
except (OSError, UnicodeError) as exc:
|
|
raise StoreError("rate-limit migration marker is unreadable") from exc
|
|
if len(lines) != 2 or lines[0] != "version=2" or not lines[1].startswith("sha256="):
|
|
raise StoreError("rate-limit migration marker is invalid")
|
|
digest = lines[1][len("sha256=") :]
|
|
if len(digest) != 64 or any(character not in "0123456789abcdef" for character in digest):
|
|
raise StoreError("rate-limit migration marker is invalid")
|
|
return digest
|
|
|
|
|
|
def _write_migration_marker(marker_path, legacy_digest: str) -> None:
|
|
temporary = marker_path.with_name(f".{marker_path.name}.{os.getpid()}.tmp")
|
|
try:
|
|
with temporary.open("w", encoding="utf-8") as fh:
|
|
fh.write(f"version=2\nsha256={legacy_digest}\n")
|
|
fh.flush()
|
|
os.fsync(fh.fileno())
|
|
os.replace(temporary, marker_path)
|
|
except (OSError, UnicodeError) as exc:
|
|
raise StoreError("could not persist rate-limit migration marker") from exc
|
|
finally:
|
|
if temporary.exists():
|
|
try:
|
|
temporary.unlink()
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _migrate_legacy_state(store: JsonStore) -> None:
|
|
"""Import the pre-S4 single-file limiter state exactly once, safely."""
|
|
legacy_path, marker_path = _legacy_paths()
|
|
if marker_path.exists():
|
|
if not marker_path.is_file():
|
|
raise StoreError("rate-limit migration marker is invalid")
|
|
marker_digest = _read_migration_marker(marker_path)
|
|
if not legacy_path.exists():
|
|
return
|
|
if not legacy_path.is_file():
|
|
raise StoreError("legacy rate-limit state is invalid")
|
|
try:
|
|
legacy_bytes = legacy_path.read_bytes()
|
|
except OSError as exc:
|
|
raise StoreError("legacy rate-limit state is unreadable") from exc
|
|
if hashlib.sha256(legacy_bytes).hexdigest() != marker_digest:
|
|
raise StoreError("rate-limit migration marker does not match legacy state")
|
|
return
|
|
if not legacy_path.exists():
|
|
return
|
|
if not legacy_path.is_file():
|
|
raise StoreError("legacy rate-limit state is invalid")
|
|
|
|
with store.collection_lock():
|
|
if marker_path.exists():
|
|
if not marker_path.is_file():
|
|
raise StoreError("rate-limit migration marker is invalid")
|
|
marker_digest = _read_migration_marker(marker_path)
|
|
if not legacy_path.exists():
|
|
return
|
|
if not legacy_path.is_file():
|
|
raise StoreError("legacy rate-limit state is invalid")
|
|
try:
|
|
legacy_bytes = legacy_path.read_bytes()
|
|
except OSError as exc:
|
|
raise StoreError("legacy rate-limit state is unreadable") from exc
|
|
if hashlib.sha256(legacy_bytes).hexdigest() != marker_digest:
|
|
raise StoreError("rate-limit migration marker does not match legacy state")
|
|
return
|
|
try:
|
|
legacy_bytes = legacy_path.read_bytes()
|
|
legacy_state = json.loads(legacy_bytes.decode("utf-8"))
|
|
except (OSError, UnicodeError, ValueError, TypeError) as exc:
|
|
raise StoreError("legacy rate-limit state is unreadable") from exc
|
|
legacy_digest = hashlib.sha256(legacy_bytes).hexdigest()
|
|
if not isinstance(legacy_state, dict):
|
|
raise StoreError("legacy rate-limit state is invalid")
|
|
|
|
for legacy_key, raw_timestamps in legacy_state.items():
|
|
timestamps = _validated_timestamps(raw_timestamps)
|
|
entry_digest = hashlib.sha256(
|
|
json.dumps(
|
|
[legacy_key, timestamps],
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
).encode("utf-8")
|
|
).hexdigest()
|
|
for key in _legacy_key_aliases(legacy_key):
|
|
record_id = _record_id(key)
|
|
with store.record_lock(record_id):
|
|
current = store.get_or_none(record_id)
|
|
imported = current.get("legacy_imports", []) if current is not None else []
|
|
if not isinstance(imported, list) or not all(
|
|
isinstance(value, str)
|
|
and len(value) == 64
|
|
and all(character in "0123456789abcdef" for character in value)
|
|
for value in imported
|
|
):
|
|
raise StoreError("legacy rate-limit import metadata is invalid")
|
|
if entry_digest in imported:
|
|
continue
|
|
current_timestamps = (
|
|
_validated_timestamps(current.get("timestamps"))
|
|
if current is not None
|
|
else []
|
|
)
|
|
# A timestamp represents one attempt. Keep duplicate
|
|
# values: collapsing them would silently lower a legacy
|
|
# counter when several attempts shared the same clock tick.
|
|
merged = sorted(current_timestamps + timestamps)
|
|
imported = [*imported, entry_digest]
|
|
if current is None:
|
|
store.create(
|
|
{
|
|
"id": record_id,
|
|
"timestamps": merged,
|
|
"legacy_imports": imported,
|
|
},
|
|
key=record_id,
|
|
)
|
|
else:
|
|
store.update(
|
|
record_id,
|
|
timestamps=merged,
|
|
legacy_imports=imported,
|
|
)
|
|
|
|
_write_migration_marker(marker_path, legacy_digest)
|
|
|
|
|
|
def _record_id(key: str) -> str:
|
|
return hashlib.sha256(key.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def check(action: str, ident: str, *, limit: int, window: int) -> bool:
|
|
"""Return True when the attempt is within the configured window."""
|
|
if type(limit) is not int or limit <= 0 or type(window) is not int or window <= 0:
|
|
return False
|
|
try:
|
|
normalized_action = normalize_identifier(action)
|
|
normalized_ident = normalize_identifier(ident)
|
|
key = _key(normalized_action, normalized_ident)
|
|
record_id = _record_id(key)
|
|
except (AuthError, TypeError, UnicodeError, ValueError):
|
|
return False
|
|
now = time.time()
|
|
|
|
try:
|
|
store = _record_store()
|
|
with store.record_lock(record_id):
|
|
try:
|
|
record = store.get(record_id)
|
|
except StoreNotFoundError:
|
|
record = None
|
|
except (OSError, StoreError, TypeError, UnicodeError, ValueError):
|
|
return False
|
|
|
|
if record is not None and not isinstance(record, dict):
|
|
return False
|
|
raw_timestamps = record.get("timestamps") if record is not None else []
|
|
if not isinstance(raw_timestamps, list):
|
|
return False
|
|
|
|
recent = []
|
|
for timestamp in raw_timestamps:
|
|
if isinstance(timestamp, bool) or not isinstance(timestamp, (int, float)):
|
|
return False
|
|
try:
|
|
timestamp = float(timestamp)
|
|
except (OverflowError, TypeError, ValueError):
|
|
return False
|
|
if not math.isfinite(timestamp):
|
|
return False
|
|
if now - timestamp < window:
|
|
recent.append(timestamp)
|
|
|
|
allowed = len(recent) < limit
|
|
if allowed:
|
|
recent.append(now)
|
|
|
|
if record is None:
|
|
store.create({"id": record_id, "timestamps": recent}, key=record_id)
|
|
else:
|
|
store.update(record_id, timestamps=recent)
|
|
|
|
# Diagnostic only; never consult this process-local mapping.
|
|
_mem[key] = recent
|
|
return allowed
|
|
except (OSError, StoreError, TypeError, UnicodeError, ValueError):
|
|
# A durable limiter must fail closed on any storage/serialization fault.
|
|
return False
|