plm-lottery/app/withdrawals/service.py

from embit import script
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.errors import ApiError
from app.audit.log import write_audit_log
from app.db.models import PendingTransaction, User, UtxoEvent, Withdrawal
from app.electrum.client import ElectrumClient
from app.rounds.config import get_round_config
from app.rounds.events import broadcaster
from app.wallet.address import is_valid_plm_address
from app.wallet.balance import recompute_balance
from app.wallet.hd import derive_user_key
from app.wallet.psbt_builder import (
    BuiltTransaction,
    InsufficientFundsError,
    Utxo,
    build_signed_transaction,
)


class WithdrawalError(ApiError):
    pass


async def request_withdrawal(
    session: AsyncSession, client: ElectrumClient, user: User, external_address: str, amount_sats: int
) -> Withdrawal:
    # Checked before anything else: an address from another chain parses fine as a
    # witness program (see wallet/address.py), so without this the tx would build,
    # broadcast and be irrecoverable rather than fail.
    if not is_valid_plm_address(external_address):
        raise WithdrawalError("invalid_address", "not a valid PLM bech32 address")

    # Withdrawing to your own deposit address is a no-op that costs a network fee,
    # and it breaks two things that assume the recipient and the change are
    # distinguishable by address: the RBF bump would shrink the recipient output
    # instead of the change (tx/broadcast.py:_find_change_output), and
    # compute_pending_balance would count the amount twice (B-17).
    if external_address == user.address:
        raise WithdrawalError(
            "withdrawal_to_own_address",
            "that is your own deposit address — withdraw to an external wallet instead",
        )

    config = await get_round_config(session)
    if amount_sats < config.bet_amount_sats:
        raise WithdrawalError(
            "amount_below_minimum",
            f"amount below the minimum of {config.bet_amount_sats} sats",
            minimum_sats=config.bet_amount_sats,
        )

    unspent = (
        await session.scalars(
            select(UtxoEvent).where(UtxoEvent.user_id == user.id, UtxoEvent.spent_txid.is_(None))
        )
    ).all()
    if sum(u.amount_sats for u in unspent) < amount_sats:
        raise WithdrawalError("insufficient_balance", "insufficient balance", required_sats=amount_sats)

    user_key = derive_user_key(user.derivation_index)
    from_script = script.p2wpkh(user_key.to_public())
    utxos = [Utxo(u.txid, u.vout, u.amount_sats) for u in unspent]

    try:
        built = build_signed_transaction(
            signing_key=user_key,
            from_script=from_script,
            utxos=utxos,
            to_address=external_address,
            amount_sats=amount_sats,
            change_address=user.address,
            fee_rate_sat_vb=config.fee_rate_sat_vb,
        )
    except InsufficientFundsError as exc:
        raise WithdrawalError(exc.code, str(exc)) from exc

    # Persist the intent before broadcasting, and only promote the rows once the
    # network has accepted the tx — same two-phase shape as place_bet (B-08).
    spent_by_key = {(u.txid, u.vout): u for u in unspent}
    for spent in built.spent_utxos:
        spent_by_key[(spent.txid, spent.vout)].spent_txid = built.txid
    await recompute_balance(session, user.id)

    withdrawal = Withdrawal(
        user_id=user.id,
        external_address=external_address,
        amount_requested_sats=amount_sats,
        amount_sent_sats=built.recipient_sats,
        txid=built.txid,
        status="building",
    )
    session.add(withdrawal)
    await session.flush()
    pending = PendingTransaction(
        kind="withdrawal",
        withdrawal_id=withdrawal.id,
        user_id=user.id,
        current_txid=built.txid,
        fee_rate_sat_vb=config.fee_rate_sat_vb,
        raw_tx_hex=built.raw_hex,
        status="building",
    )
    session.add(pending)
    await session.commit()

    try:
        await client.broadcast(built.raw_hex)
    except Exception as exc:
        await _release_failed_withdrawal(session, withdrawal, pending, built, user.id, str(exc))
        raise WithdrawalError(
            "broadcast_failed", f"the network refused the transaction: {exc}"
        ) from exc

    withdrawal.status = "broadcast"
    pending.status = "pending"
    await write_audit_log(
        session,
        "withdrawal_sent",
        {"txid": built.txid, "amount_sent_sats": built.recipient_sats, "external_address": external_address},
        user_id=user.id,
    )

    await session.commit()
    await session.refresh(withdrawal)
    broadcaster.publish()  # balance just went "pending" — nudge the dashboard to refetch
    return withdrawal


async def _release_failed_withdrawal(
    session: AsyncSession,
    withdrawal: Withdrawal,
    pending: PendingTransaction,
    built: BuiltTransaction,
    user_id: int,
    reason: str,
) -> None:
    """Nothing reached the chain, so free the reserved UTXOs and restore the balance.
    The Withdrawal row is kept (marked "failed") rather than deleted: unlike a bet, a
    withdrawal is an instruction the user gave, and they should be able to see that it
    didn't go through."""
    for spent in built.spent_utxos:
        row = await session.scalar(
            select(UtxoEvent).where(UtxoEvent.txid == spent.txid, UtxoEvent.vout == spent.vout)
        )
        if row is not None:
            row.spent_txid = None
    withdrawal.status = "failed"
    withdrawal.txid = None
    pending.status = "failed"
    pending.failure_reason = reason[:128]
    await recompute_balance(session, user_id)
    await write_audit_log(
        session,
        "withdrawal_broadcast_failed",
        {"txid": built.txid, "reason": reason[:200]},
        user_id=user_id,
    )
    await session.commit()