Redis contrib¶

def build_rate_limit_config(  # noqa: PLR0913
    self,
    *,
    enabled: typing.Iterable[AuthRateLimitEndpointSlot] | None = None,
    disabled: typing.Iterable[AuthRateLimitEndpointSlot] = (),
    group_backends: typing.Mapping[AuthRateLimitEndpointGroup, RateLimiterBackend] | None = None,
    scope_overrides: typing.Mapping[AuthRateLimitEndpointSlot, RateLimitScope] | None = None,
    namespace_style: AuthRateLimitNamespaceStyle = "route",
    namespace_overrides: typing.Mapping[AuthRateLimitEndpointSlot, str] | None = None,
    endpoint_overrides: typing.Mapping[AuthRateLimitEndpointSlot, EndpointRateLimit | None] | None = None,
    trusted_proxy: bool = False,
    identity_fields: tuple[str, ...] | None = None,
    trusted_headers: tuple[str, ...] | None = None,
) -> AuthRateLimitConfig:
    """Build ``AuthRateLimitConfig`` from the preset's shared Redis client.

    Args:
        enabled: Optional auth slot names to build.
        disabled: Auth slot names to leave unset.
        group_backends: Optional explicit backend overrides keyed by auth
            slot group. These win over ``group_rate_limit_tiers``.
        scope_overrides: Optional per-slot scope overrides.
        namespace_style: Namespace family for generated limiters.
        namespace_overrides: Optional per-slot namespace overrides.
        endpoint_overrides: Optional full per-slot replacements or explicit
            ``None`` disablement.
        trusted_proxy: Shared trusted-proxy setting applied to generated
            limiters.
        identity_fields: Optional shared request body identity fields. When
            omitted, ``AuthRateLimitConfig.from_shared_backend()`` keeps its
            current default.
        trusted_headers: Optional shared trusted proxy header names. When
            omitted, ``AuthRateLimitConfig.from_shared_backend()`` keeps its
            current default.

    Returns:
        The auth rate-limit config built from the preset's shared client and
        tier settings.
    """
    derived_group_backends: dict[AuthRateLimitEndpointGroup, RateLimiterBackend] = {
        group: self._build_rate_limit_backend(tier) for group, tier in self.group_rate_limit_tiers.items()
    }
    if group_backends is not None:
        derived_group_backends.update(group_backends)
    shared_backend: RateLimiterBackend = self._build_rate_limit_backend(self.rate_limit_tier)
    resolved_group_backends = derived_group_backends or None

    if identity_fields is None:
        if trusted_headers is None:
            return AuthRateLimitConfig.from_shared_backend(
                shared_backend,
                enabled=enabled,
                disabled=disabled,
                group_backends=resolved_group_backends,
                scope_overrides=scope_overrides,
                namespace_style=namespace_style,
                namespace_overrides=namespace_overrides,
                endpoint_overrides=endpoint_overrides,
                trusted_proxy=trusted_proxy,
            )
        return AuthRateLimitConfig.from_shared_backend(
            shared_backend,
            enabled=enabled,
            disabled=disabled,
            group_backends=resolved_group_backends,
            scope_overrides=scope_overrides,
            namespace_style=namespace_style,
            namespace_overrides=namespace_overrides,
            endpoint_overrides=endpoint_overrides,
            trusted_proxy=trusted_proxy,
            trusted_headers=trusted_headers,
        )

    if trusted_headers is None:
        return AuthRateLimitConfig.from_shared_backend(
            shared_backend,
            enabled=enabled,
            disabled=disabled,
            group_backends=resolved_group_backends,
            scope_overrides=scope_overrides,
            namespace_style=namespace_style,
            namespace_overrides=namespace_overrides,
            endpoint_overrides=endpoint_overrides,
            trusted_proxy=trusted_proxy,
            identity_fields=identity_fields,
        )

    return AuthRateLimitConfig.from_shared_backend(
        shared_backend,
        enabled=enabled,
        disabled=disabled,
        group_backends=resolved_group_backends,
        scope_overrides=scope_overrides,
        namespace_style=namespace_style,
        namespace_overrides=namespace_overrides,
        endpoint_overrides=endpoint_overrides,
        trusted_proxy=trusted_proxy,
        identity_fields=identity_fields,
        trusted_headers=trusted_headers,
    )

def build_totp_used_tokens_store(self, *, key_prefix: str | None = None) -> RedisUsedTotpCodeStore:
    """Build ``RedisUsedTotpCodeStore`` from the preset's shared Redis client.

    Args:
        key_prefix: Optional per-call Redis key prefix override. When
            omitted, the preset uses ``totp_used_tokens_key_prefix`` and
            finally falls back to the current ``RedisUsedTotpCodeStore``
            default.

    Returns:
        Redis-backed TOTP replay store sharing the preset's client.
    """
    resolved_key_prefix = self.totp_used_tokens_key_prefix if key_prefix is None else key_prefix
    if resolved_key_prefix is None:
        resolved_key_prefix = DEFAULT_TOTP_USED_KEY_PREFIX
    return RedisUsedTotpCodeStore(redis=self.redis, key_prefix=resolved_key_prefix)

@override
async def destroy_token(self, token: str, user: UP) -> None:
    """Delete a persisted Redis token."""
    token_key = self._key(token)
    user_id = str(user.id)
    index_key = self._user_index_key(user_id)
    await self.redis.delete(token_key)
    await self.redis.srem(index_key, token_key)

async def invalidate_all_tokens(self, user: UP) -> None:
    """Delete all Redis-backed tokens associated with the given user.

    This uses a per-user index to delete only the keys associated with the
    user, avoiding keyspace scans under the global prefix.

    Backward compatibility:
        If the per-user index is missing, we fall back to a best-effort scan
        over keys matching the configured prefix and remove those whose
        stored subject matches ``user.id``.
    """
    user_id = str(user.id)
    if await self._invalidate_via_index(user_id):
        return
    await self._invalidate_via_scan(user_id)

@override
async def read_token(
    self,
    token: str | None,
    user_manager: UserManagerProtocol[UP, ID],
) -> UP | None:
    """Resolve a user from a Redis-backed token.

    Returns:
        The resolved user when the token exists and decodes successfully,
        otherwise ``None``.
    """
    if token is None:
        return None

    stored_user_id = await self.redis.get(self._key(token))
    if stored_user_id is None:
        return None

    user_id_text = self._decode_user_id(stored_user_id)

    try:
        user_id = self.subject_decoder(user_id_text) if self.subject_decoder is not None else user_id_text
    except (TypeError, ValueError):
        return None

    return await user_manager.get(cast("ID", user_id))

@override
async def write_token(self, user: UP) -> str:
    """Persist a new opaque token in Redis and return it.

    Returns:
        Newly created opaque token string.
    """
    token = secrets.token_urlsafe(self.token_bytes)
    token_key = self._key(token)
    user_id = str(user.id)
    await self.redis.setex(token_key, self._ttl_seconds, user_id)
    index_key = self._user_index_key(user_id)
    await self.redis.sadd(index_key, token_key)
    await self.redis.expire(index_key, self._ttl_seconds)
    return token

async def mark_used(self, user_id: Hashable, counter: int, ttl_seconds: float) -> bool:
    """Atomically record a used (user_id, counter) pair via SET key 1 NX PX ttl_ms.

    Returns:
        ``True`` when the pair was newly stored, ``False`` for a replay.
    """
    key = self._key(user_id, counter)
    ttl_ms = int(ttl_seconds * 1000)
    result = await self._redis.set(key, "1", nx=True, px=ttl_ms)
    return result is True