Rate limiting¶

@classmethod
def from_shared_backend(  # noqa: PLR0913
    cls,
    backend: RateLimiterBackend,
    *,
    enabled: Iterable[AuthRateLimitEndpointSlot] | None = None,
    disabled: Iterable[AuthRateLimitEndpointSlot] = (),
    group_backends: Mapping[AuthRateLimitEndpointGroup, RateLimiterBackend] | None = None,
    scope_overrides: Mapping[AuthRateLimitEndpointSlot, RateLimitScope] | None = None,
    namespace_overrides: Mapping[AuthRateLimitEndpointSlot, str] | None = None,
    endpoint_overrides: Mapping[AuthRateLimitEndpointSlot, EndpointRateLimit | None] | None = None,
    trusted_proxy: bool = False,
    identity_fields: tuple[str, ...] = _DEFAULT_IDENTITY_FIELDS,
    trusted_headers: tuple[str, ...] = _DEFAULT_TRUSTED_HEADERS,
) -> Self:
    """Build endpoint-specific limiters from the package-owned shared-backend recipe.

    The builder uses the private endpoint catalog for default scopes and namespace
    tokens, then applies override precedence in this order:

    1. ``backend`` for every enabled slot
    2. ``group_backends`` for slot groups such as ``totp`` or ``verification``
    3. ``scope_overrides`` / ``namespace_overrides`` for slot-specific tweaks
    4. ``endpoint_overrides`` for full slot replacement or explicit ``None`` disablement

    Args:
        backend: Default limiter backend for enabled auth slots.
        enabled: Optional auth slot names to build. Defaults to all supported slots.
        disabled: Auth slot names to leave unset, even when they would otherwise be enabled.
        group_backends: Optional backend overrides keyed by auth slot group:
            ``login``, ``refresh``, ``register``, ``password_reset``, ``totp``, or
            ``verification``.
        scope_overrides: Optional per-slot scope overrides to preserve existing key behavior.
        namespace_overrides: Optional per-slot namespace tokens to preserve existing key names.
        endpoint_overrides: Optional full per-slot replacements. ``None`` disables a slot.
        trusted_proxy: Shared trusted-proxy setting applied to generated limiters.
        identity_fields: Shared request body identity fields applied to generated limiters.
        trusted_headers: Shared trusted proxy header names applied to generated limiters.

    Returns:
        New config populated from the shared-backend builder inputs.
    """
    group_backend_map: dict[AuthRateLimitEndpointGroup, RateLimiterBackend] = dict(group_backends or {})
    scope_override_map: dict[AuthRateLimitEndpointSlot, RateLimitScope] = dict(scope_overrides or {})
    namespace_override_map: dict[AuthRateLimitEndpointSlot, str] = dict(namespace_overrides or {})
    endpoint_override_map: dict[AuthRateLimitEndpointSlot, EndpointRateLimit | None] = dict(
        endpoint_overrides or {},
    )
    enabled_slots: frozenset[AuthRateLimitEndpointSlot] = (
        _AUTH_RATE_LIMIT_ENDPOINT_SLOT_SET if enabled is None else frozenset(enabled)
    )
    disabled_slots: frozenset[AuthRateLimitEndpointSlot] = frozenset(disabled)

    _validate_builder_names(
        enabled_slots,
        allowed=_AUTH_RATE_LIMIT_ENDPOINT_SLOT_SET,
        parameter_name="enabled",
        item_name="auth rate-limit slots",
    )
    _validate_builder_names(
        disabled_slots,
        allowed=_AUTH_RATE_LIMIT_ENDPOINT_SLOT_SET,
        parameter_name="disabled",
        item_name="auth rate-limit slots",
    )
    _validate_builder_names(
        group_backend_map,
        allowed=_AUTH_RATE_LIMIT_ENDPOINT_GROUPS,
        parameter_name="group_backends",
        item_name="auth rate-limit groups",
    )
    _validate_builder_names(
        scope_override_map,
        allowed=_AUTH_RATE_LIMIT_ENDPOINT_SLOT_SET,
        parameter_name="scope_overrides",
        item_name="auth rate-limit slots",
    )
    _validate_builder_names(
        namespace_override_map,
        allowed=_AUTH_RATE_LIMIT_ENDPOINT_SLOT_SET,
        parameter_name="namespace_overrides",
        item_name="auth rate-limit slots",
    )
    _validate_builder_names(
        endpoint_override_map,
        allowed=_AUTH_RATE_LIMIT_ENDPOINT_SLOT_SET,
        parameter_name="endpoint_overrides",
        item_name="auth rate-limit slots",
    )

    config_kwargs: dict[AuthRateLimitEndpointSlot, EndpointRateLimit | None] = {}

    for recipe in _AUTH_RATE_LIMIT_ENDPOINT_RECIPES:
        slot_override = endpoint_override_map.get(recipe.slot, _MISSING_OVERRIDE)
        if slot_override is not _MISSING_OVERRIDE:
            config_kwargs[recipe.slot] = cast("EndpointRateLimit | None", slot_override)
            continue

        if recipe.slot not in enabled_slots or recipe.slot in disabled_slots:
            continue

        config_kwargs[recipe.slot] = EndpointRateLimit(
            backend=group_backend_map.get(recipe.group, backend),
            scope=scope_override_map.get(recipe.slot, recipe.default_scope),
            namespace=namespace_override_map.get(recipe.slot, recipe.default_namespace),
            trusted_proxy=trusted_proxy,
            identity_fields=identity_fields,
            trusted_headers=trusted_headers,
        )

    return cls(**cast("Any", config_kwargs))

async def before_request(self, request: Request[Any, Any, Any]) -> None:
    """Reject the request with 429 when its key is over the configured limit.

    Security:
        Only set ``trusted_proxy=True`` when this service is behind a trusted
        proxy or load balancer that overwrites client IP headers. Otherwise,
        attackers can spoof headers like ``X-Forwarded-For`` and evade or
        poison rate-limiting keys.

    Raises:
        TooManyRequestsException: If the request exceeded the configured limit.
    """
    key = await self.build_key(request)
    if await self.backend.check(key):
        return

    retry_after = await self.backend.retry_after(key)
    logger.warning(
        "Rate limit exceeded",
        extra={
            "event": "rate_limit_triggered",
            "namespace": self.namespace,
            "scope": self.scope,
            "trusted_proxy": self.trusted_proxy,
        },
    )
    msg = "Too many requests."
    raise TooManyRequestsException(
        detail=msg,
        headers={"Retry-After": str(max(retry_after, 1))},
    )

async def build_key(self, request: Request[Any, Any, Any]) -> str:
    """Build the backend key for the given request.

    Returns:
        Namespaced rate-limit key for the request.
    """
    host = _client_host(request, trusted_proxy=self.trusted_proxy, trusted_headers=self.trusted_headers)
    parts = [self.namespace, _safe_key_part(host)]
    if self.scope == "ip_email":
        email = await _extract_email(request, identity_fields=self.identity_fields)
        if email:
            parts.append(_safe_key_part(email.strip().casefold()))

    return ":".join(parts)

async def increment(self, request: Request[Any, Any, Any]) -> None:
    """Record a failed or rate-limited attempt for the current request."""
    await self.backend.increment(await self.build_key(request))

async def reset(self, request: Request[Any, Any, Any]) -> None:
    """Clear stored attempts for the current request key."""
    await self.backend.reset(await self.build_key(request))

async def check(self, key: str) -> bool:
    """Return whether ``key`` can perform another attempt."""
    async with self._lock:
        now = self._clock()
        self._maybe_sweep(now)
        timestamps = self._prune(key, now)
        if timestamps is None:
            return True

        return len(timestamps) < self.max_attempts

async def increment(self, key: str) -> None:
    """Record a new attempt for ``key`` in the current window."""
    async with self._lock:
        now = self._clock()
        self._maybe_sweep(now)
        timestamps = self._prune(key, now)
        if timestamps is None:
            self._evict_oldest_keys()
            timestamps = deque()
            self._windows[key] = timestamps

        timestamps.append(now)

async def reset(self, key: str) -> None:
    """Clear the in-memory counter for ``key``."""
    async with self._lock:
        self._windows.pop(key, None)

async def retry_after(self, key: str) -> int:
    """Return the remaining block duration for ``key`` in whole seconds."""
    async with self._lock:
        now = self._clock()
        timestamps = self._prune(key, now)
        if timestamps is None or len(timestamps) < self.max_attempts:
            return 0

        oldest_timestamp = timestamps[0]
        remaining = self.window_seconds - (now - oldest_timestamp)
        return max(math.ceil(remaining), 1)

async def check(self, key: str) -> bool:
    """Return whether ``key`` can perform another attempt."""
    count = self._decode_integer(
        await self._eval(
            self._CHECK_SCRIPT,
            key,
            self._clock(),
            self.window_seconds,
            self.max_attempts,
        ),
    )
    return count < self.max_attempts

async def increment(self, key: str) -> None:
    """Record a new attempt for ``key`` atomically in Redis."""
    now = self._clock()
    await self._eval(
        self._INCREMENT_SCRIPT,
        key,
        now,
        self.window_seconds,
        f"{now:.9f}:{uuid4().hex}",
        self._ttl_seconds,
    )

async def reset(self, key: str) -> None:
    """Delete the Redis sorted set for ``key``."""
    await self.redis.delete(self._key(key))

async def retry_after(self, key: str) -> int:
    """Return the remaining block duration for ``key`` in whole seconds."""
    return max(
        self._decode_integer(
            await self._eval(
                self._RETRY_AFTER_SCRIPT,
                key,
                self._clock(),
                self.window_seconds,
                self.max_attempts,
            ),
        ),
        0,
    )

async def check(self, key: str) -> bool:
    """Return whether another attempt is allowed for ``key``."""

async def increment(self, key: str) -> None:
    """Record an attempt for ``key``."""

async def reset(self, key: str) -> None:
    """Clear tracked attempts for ``key``."""

async def retry_after(self, key: str) -> int:
    """Return the number of seconds until ``key`` can try again."""

async def before_request(self, endpoint: TotpSensitiveEndpoint, request: Request[Any, Any, Any]) -> None:
    """Run endpoint-specific before-request checks."""
    if limiter := self._limiters.get(endpoint):
        await limiter.before_request(request)

async def on_account_state_failure(self, endpoint: TotpSensitiveEndpoint, request: Request[Any, Any, Any]) -> None:
    """Apply endpoint-specific account-state failure behavior."""
    if endpoint in self._ACCOUNT_STATE_RESET_ENDPOINTS and (limiter := self._limiters.get(endpoint)):
        await limiter.reset(request)

async def on_invalid_attempt(self, endpoint: TotpSensitiveEndpoint, request: Request[Any, Any, Any]) -> None:
    """Record endpoint-specific invalid attempt failures."""
    if limiter := self._limiters.get(endpoint):
        await limiter.increment(request)

async def on_success(self, endpoint: TotpSensitiveEndpoint, request: Request[Any, Any, Any]) -> None:
    """Apply endpoint-specific success behavior."""
    if limiter := self._limiters.get(endpoint):
        await limiter.reset(request)