plm-lottery/app/auth/routes.py

from fastapi import APIRouter, Depends, Request, status
from pydantic import BaseModel, Field
from sqlalchemy import func, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.client_ip import client_ip as _client_ip
from app.api.errors import http_error
from app.auth.rate_limit import AuthRateLimiters
from app.auth.security import (
    MIN_PASSWORD_LENGTH,
    create_access_token,
    hash_password_async,
    verify_password_async,
)
from app.db.models import User
from app.db.session import get_session
from app.wallet.hd import derive_user_address

router = APIRouter(prefix="/auth", tags=["auth"])

_MAX_REGISTER_RETRIES = 5


def _rate_limiters(request: Request) -> AuthRateLimiters:
    # B-33: no rate limiting on login was a brute-forceable path to withdrawing
    # someone else's funds. Keyed per-username *and* per-IP so an attacker can't
    # dodge the throttle by spraying one password across many accounts, nor by
    # routing one account's guesses through many IPs alone (the username key
    # still catches that). The IP limiter's threshold is deliberately higher
    # than the username one: a single account should lock out fast, but a
    # shared/NAT IP hosting several genuine users shouldn't be punished for one
    # of them mistyping a password a few times. Registration gets its own,
    # coarser limiter, IP-only — no username exists yet to key on — mainly to
    # bound how many accounts one IP can spin up (B-31), not to protect a
    # secret. Lives on app.state (see AuthRateLimiters) rather than a module
    # global so each app instance gets its own, isolated throttle state.
    if not hasattr(request.app.state, "auth_rate_limiters"):
        request.app.state.auth_rate_limiters = AuthRateLimiters()
    return request.app.state.auth_rate_limiters


def _rate_limited_error(retry_after: float):
    return http_error(
        status.HTTP_429_TOO_MANY_REQUESTS,
        "rate_limited",
        "too many attempts, try again later",
        retry_after_seconds=int(retry_after) + 1,
    )


class RegisterRequest(BaseModel):
    """Registration used to accept an empty username and a one-character password,
    while /users/me/change-password demanded 8 characters — an odd place to be
    lenient on a custodial system holding real funds (B-12). MIN_PASSWORD_LENGTH is
    shared with that endpoint so the two can't drift apart again."""

    username: str = Field(min_length=3, max_length=32, pattern=r"^[A-Za-z0-9_.-]+$")
    password: str = Field(min_length=MIN_PASSWORD_LENGTH, max_length=256)


class TokenResponse(BaseModel):
    access_token: str
    address: str


@router.post("/register", response_model=TokenResponse, status_code=status.HTTP_201_CREATED)
async def register(
    body: RegisterRequest, request: Request, session: AsyncSession = Depends(get_session)
) -> TokenResponse:
    limiters = _rate_limiters(request)
    ip_key = f"ip:{_client_ip(request)}"
    retry_after = limiters.register_ip.retry_after(ip_key)
    if retry_after > 0:
        raise _rate_limited_error(retry_after)
    limiters.register_ip.record_failure(ip_key)

    existing = await session.scalar(select(User).where(User.username == body.username))
    if existing is not None:
        raise http_error(status.HTTP_409_CONFLICT, "username_taken", "username already taken")

    password_hash = await hash_password_async(body.password)

    for _ in range(_MAX_REGISTER_RETRIES):
        max_index = await session.scalar(select(func.max(User.derivation_index)))
        next_index = 0 if max_index is None else max_index + 1
        address = derive_user_address(next_index)
        user = User(
            username=body.username,
            password_hash=password_hash,
            derivation_index=next_index,
            address=address,
        )
        session.add(user)
        try:
            await session.commit()
        except IntegrityError as exc:
            await session.rollback()
            # Only a derivation-index collision is worth retrying. A username
            # collision (someone registered the same name between the check above and
            # this commit) is permanent, and retrying it five times only to report
            # "derivation_index_conflict" told the user the wrong thing entirely (B-12).
            if "username" in str(exc.orig).lower():
                raise http_error(
                    status.HTTP_409_CONFLICT, "username_taken", "username already taken"
                ) from exc
            continue
        await session.refresh(user)
        request.app.state.electrum_listener.address_for_new_user(user.id, user.address)
        return TokenResponse(
            access_token=create_access_token(user.id, user.token_version), address=user.address
        )

    raise http_error(
        status.HTTP_409_CONFLICT,
        "derivation_index_conflict",
        "could not allocate a derivation index, retry",
    )


class LoginRequest(BaseModel):
    username: str
    password: str


@router.post("/login", response_model=TokenResponse)
async def login(
    body: LoginRequest, request: Request, session: AsyncSession = Depends(get_session)
) -> TokenResponse:
    limiters = _rate_limiters(request)
    username_key = f"user:{body.username.lower()}"
    ip_key = f"ip:{_client_ip(request)}"
    retry_after = max(limiters.login.retry_after(username_key), limiters.login_ip.retry_after(ip_key))
    if retry_after > 0:
        raise _rate_limited_error(retry_after)

    user = await session.scalar(select(User).where(User.username == body.username))
    if user is None or not await verify_password_async(body.password, user.password_hash):
        # Same code path (and therefore the same response) whether the username
        # doesn't exist or the password is wrong — no enumeration oracle here.
        limiters.login.record_failure(username_key)
        limiters.login_ip.record_failure(ip_key)
        raise http_error(status.HTTP_401_UNAUTHORIZED, "invalid_credentials", "invalid credentials")

    # Only the username bucket resets on success — the IP bucket is left to decay
    # on its own, so one correct login can't be used to wipe out an IP's failure
    # count while it's mid-attack against other accounts.
    limiters.login.record_success(username_key)
    return TokenResponse(
        access_token=create_access_token(user.id, user.token_version), address=user.address
    )