plm-lottery/app/wallet/psbt_builder.py

from dataclasses import dataclass

from embit import script
from embit.bip32 import HDKey
from embit.finalizer import finalize_psbt
from embit.psbt import PSBT
from embit.transaction import Transaction, TransactionInput, TransactionOutput

# Standard P2WPKH size estimates (vbytes): 10.5-byte overhead (version+counts+locktime+
# segwit marker/flag), ~68 vbytes per input, ~31 vbytes per output. Used to size the fee
# before signing (fee only needs to be "minimized ~1 sat/vB", not maximally precise).
_TX_OVERHEAD_VBYTES = 11
_P2WPKH_INPUT_VBYTES = 68
_P2WPKH_OUTPUT_VBYTES = 31

# BIP125 opt-in RBF: any sequence < 0xfffffffe signals replaceability. Set on every
# input we create so a stuck tx can later be fee-bumped (tx/broadcast.py, stage 9).
RBF_SEQUENCE = 0xFFFFFFFD

# Below this, an output costs more to spend than it's worth and relay policy rejects
# the whole transaction as "dust" — so a small change amount must be left to the fee
# instead of being paid back to ourselves. 294 sat is the standard P2WPKH threshold
# (the output's own 31 vbytes plus the 67-vbyte input needed to spend it, at the
# 3000 sat/kvB dust relay fee). Creating such an output used to make the bet or
# withdrawal fail at broadcast with an opaque error (B-06).
DUST_LIMIT_SATS = 294

# Sanity ceiling on any transaction's fee rate — shared by RoundConfig.fee_rate_sat_vb's
# admin-facing bound (app/api/routes/admin.py, so the two can't drift apart, the same
# reason MIN_PASSWORD_LENGTH is shared in auth/security.py) and tx/broadcast.py's RBF
# bump escalation, which refuses to bump a pending_transaction past this rate (B-32) —
# without a ceiling, a stuck transaction's fee climbed by 1 sat/vB every bump forever,
# eating further and further into the sender's change with no limit.
MAX_FEE_RATE_SAT_VB = 10_000

# Ceiling on how many UTXOs one *user* transaction (bet, withdrawal) may spend (B-48).
# Every extra input costs ~68 vbytes of fee, and that fee comes out of the amount being
# moved — so an address fragmented into hundreds of small deposits would silently erode
# its own bet (shrinking the user's share of the pool) or withdrawal. Failing the build
# with a translatable error is the honest outcome; consolidating the address is the way
# out. This is a *user-protection* limit, which is why the payout gets its own, far
# higher one below.
MAX_TX_INPUTS = 50

# Ceiling on the payout's inputs (B-52). The payout is not a user spending their own
# fragmented balance: it drains the pool, whose UTXO count is simply the number of bets
# in the round, and its fee comes out of a 70% share of that whole pool. So the erosion
# argument behind MAX_TX_INPUTS doesn't apply here — 400 inputs at 1 sat/vB cost ~27_300
# sat, i.e. ~0.00027 PLM out of the winner's share — and reusing that limit was what made
# any round past ~50 participants unpayable: select_utxos raised too_many_inputs, the
# round stayed "paying_out" retrying forever, and since no new round may open while one
# is active, the whole lottery stopped with the pool stuck (B-52).
#
# What actually bounds this is relay policy: a non-standard transaction is refused at
# broadcast past 100 kvB, which at ~68 vbytes per input is ~1470 inputs. 500 stays at
# roughly a third of that budget, and signing that many inputs costs ~0.4s of event loop
# (measured), once per round, inside a background task.
MAX_PAYOUT_TX_INPUTS = 500

# The most participants one round may hold (B-52). Enforced where the money is not yet
# committed — app/bets/service.py refuses the bet — instead of being discovered at payout
# time, when the bets are already in the pool and there is no way back. Deliberately below
# MAX_PAYOUT_TX_INPUTS: the payout also has to be able to spend whatever change UTXOs
# earlier rounds left in the pool, so the gap is the headroom for those.
MAX_PARTICIPANTS_PER_ROUND = 400


class InsufficientFundsError(Exception):
    """`code` is the machine-readable identifier the API layer forwards to the
    client so it can translate the failure (see app/api/errors.py); the message
    itself stays English, and `params` carries the values it interpolates so the
    translation can place them wherever its own grammar needs them."""

    def __init__(
        self,
        message: str,
        code: str = "insufficient_balance",
        **params: int | str,
    ) -> None:
        super().__init__(message)
        self.code = code
        self.params = params


# B-73 (DAT-63): the recipient is always the first output of a single-recipient transaction
# (the change, when there is one, follows it) — which is what a bet records as `bet_vout`.
RECIPIENT_VOUT = 0


@dataclass
class Utxo:
    txid: str
    vout: int
    amount_sats: int


@dataclass
class BuiltTransaction:
    raw_hex: str
    txid: str
    fee_sats: int
    recipient_sats: int
    change_sats: int
    spent_utxos: list[Utxo]


def estimate_vsize(n_inputs: int, n_outputs: int) -> int:
    return _TX_OVERHEAD_VBYTES + n_inputs * _P2WPKH_INPUT_VBYTES + n_outputs * _P2WPKH_OUTPUT_VBYTES


def select_utxos(
    utxos: list[Utxo], target_sats: int, max_inputs: int = MAX_TX_INPUTS
) -> tuple[list[Utxo], int]:
    """Greedily select UTXOs (largest first, to minimize input count) covering
    target_sats — the amount deducted from the sender's balance. The fee is paid
    out of target_sats (see build_signed_transaction), not added on top of it.

    At most `max_inputs` are ever selected (B-48): if the largest `max_inputs`
    UTXOs don't cover the target, the balance is there but too fragmented to spend
    in one transaction, which is a different failure from having no funds at all
    and gets its own code.

    The cap is a parameter, not the constant it used to be, because the two callers
    want different ones (B-52): MAX_TX_INPUTS protects a user from a fee that would
    eat into their own bet/withdrawal, while the payout drains a pool holding one
    UTXO per bet and needs MAX_PAYOUT_TX_INPUTS to be able to pay a full round at
    all."""
    # B-75: a UTXO below the dust limit is never offered, even if one is already on
    # record: it costs more to spend than it is worth.
    ordered = sorted(
        (u for u in utxos if u.amount_sats >= DUST_LIMIT_SATS), key=lambda u: u.amount_sats, reverse=True
    )
    selected: list[Utxo] = []
    total = 0
    for utxo in ordered:
        if len(selected) == max_inputs:
            raise InsufficientFundsError(
                f"balance too fragmented: more than {max_inputs} inputs would be needed",
                code="too_many_inputs",
                max_inputs=max_inputs,
            )
        selected.append(utxo)
        total += utxo.amount_sats
        if total >= target_sats:
            return selected, total
    raise InsufficientFundsError("not enough confirmed balance to cover amount")


def build_signed_transaction(
    *,
    signing_key: HDKey,
    from_script: script.Script,
    utxos: list[Utxo],
    to_address: str,
    amount_sats: int,
    change_address: str,
    fee_rate_sat_vb: int,
    reduce_amount_to_keep_change: bool = False,
) -> BuiltTransaction:
    """Build, sign and finalize a single-recipient P2WPKH transaction with change
    back to change_address.

    `amount_sats` is deducted from the sender's balance in full: the recipient
    receives `amount_sats - fee`, change = total_in - amount_sats. This matches the
    spec everywhere a single-recipient tx is used (bet, withdrawal): "fee deducted
    from the amount being moved", not paid on top by the sender.

    B-62: the transaction always keeps a change output of at least DUST_LIMIT_SATS.
    Change used to be folded into the fee whenever it came out below the dust limit,
    which for an amount equal to the whole input total (the UI's "withdraw
    everything" checkbox, or a bet from a balance exactly equal to the bet amount)
    produced a single-output transaction — and `tx/broadcast.py:bump_fee` has nothing
    to shrink there, so it raised RbfError every 30s until the reconciler abandoned
    the row hours later. Adding inputs instead is no answer for this case in
    particular: the transaction already spends every UTXO the sender has.

    What happens when the change would be too small depends on who's asking, hence
    `reduce_amount_to_keep_change`:

    - withdrawals pass True — the amount moved is reduced just enough to leave a
      dust-limit change output. The fee already comes out of the withdrawn amount by
      design, so this is the same rule applied a little harder, and the caller
      records what was actually sent (`Withdrawal.amount_sent_sats`).
    - bets pass False (the default) and get an InsufficientFundsError instead: the
      bet is a fixed price that cannot be quietly reduced, and "a user's balance must
      never exactly equal the bet" is a documented invariant of the PLAY phase. The
      player needs a little more than the bet amount, which is what the error says.
    """
    selected, total_in = select_utxos(utxos, amount_sats)
    fee = estimate_vsize(len(selected), 2) * fee_rate_sat_vb

    change = total_in - amount_sats
    if change < DUST_LIMIT_SATS:
        if not reduce_amount_to_keep_change:
            raise InsufficientFundsError(
                f"the amount leaves no change output: {DUST_LIMIT_SATS - change} more sats are "
                "needed for the transaction to stay fee-bumpable",
                code="balance_leaves_no_change",
                required_extra_sats=DUST_LIMIT_SATS - change,
            )
        amount_sats -= DUST_LIMIT_SATS - change
        change = DUST_LIMIT_SATS

    recipient_amount = amount_sats - fee
    if recipient_amount <= 0:
        raise InsufficientFundsError(
            "amount too small to cover the network fee", code="amount_below_network_fee"
        )
    if recipient_amount < DUST_LIMIT_SATS:
        raise InsufficientFundsError(
            "amount too small to be sent (dust)", code="amount_below_dust_limit"
        )

    # TransactionInput.txid is natural/display byte order (as in tx_hash from Electrum);
    # embit reverses it internally when serializing to wire format.
    vin = [TransactionInput(bytes.fromhex(u.txid), u.vout, sequence=RBF_SEQUENCE) for u in selected]
    vout = [TransactionOutput(recipient_amount, script.Script.from_address(to_address))]
    if change > 0:
        vout.append(TransactionOutput(change, script.Script.from_address(change_address)))

    tx = Transaction(vin=vin, vout=vout)
    psbt = PSBT(tx)
    for i, utxo in enumerate(selected):
        psbt.inputs[i].witness_utxo = TransactionOutput(utxo.amount_sats, from_script)

    signed_count = psbt.sign_with(signing_key)
    if signed_count != len(selected):
        raise RuntimeError(f"expected {len(selected)} signatures, got {signed_count}")

    final_tx = finalize_psbt(psbt)
    if final_tx is None:
        raise RuntimeError("failed to finalize PSBT")

    raw = final_tx.serialize()
    return BuiltTransaction(
        raw_hex=raw.hex(),
        txid=final_tx.txid().hex(),
        fee_sats=fee,
        recipient_sats=recipient_amount,
        change_sats=change,
        spent_utxos=selected,
    )


@dataclass
class PayoutTransaction:
    raw_hex: str
    txid: str
    fee_sats: int
    winner_sats: int
    commission_sats: int
    change_sats: int
    spent_utxos: list[Utxo]


def build_payout_transaction(
    *,
    signing_key: HDKey,
    from_script: script.Script,
    utxos: list[Utxo],
    winner_address: str,
    winner_share_sats: int,
    fee_address: str,
    commission_sats: int,
    change_address: str,
    fee_rate_sat_vb: int,
) -> PayoutTransaction:
    """Build, sign and finalize the round payout: pool -> winner + fee address,
    with change back to the pool itself. Per spec, only the winner's share
    absorbs the tx fee — the commission (fee_address) output is untouched.

    As in build_signed_transaction, dust-sized pool change is left to the fee
    rather than creating an unrelayable output (B-06).

    Selection uses MAX_PAYOUT_TX_INPUTS, not the much stricter user-facing
    MAX_TX_INPUTS (B-52) — the pool holds one UTXO per bet, so the user-protection
    cap made every round past ~50 participants impossible to pay."""
    target = winner_share_sats + commission_sats
    selected, total_in = select_utxos(utxos, target, max_inputs=MAX_PAYOUT_TX_INPUTS)
    fee = estimate_vsize(len(selected), 3) * fee_rate_sat_vb  # winner + commission + pool change
    winner_amount = winner_share_sats - fee
    if winner_amount < DUST_LIMIT_SATS:
        raise InsufficientFundsError(
            "winner share too small to cover the network fee", code="winner_share_below_network_fee"
        )
    if commission_sats < DUST_LIMIT_SATS:
        raise InsufficientFundsError(
            "commission share too small to be paid out (dust)", code="commission_below_dust_limit"
        )
    change = total_in - target
    if change < DUST_LIMIT_SATS:
        fee += change
        change = 0

    vin = [TransactionInput(bytes.fromhex(u.txid), u.vout, sequence=RBF_SEQUENCE) for u in selected]
    vout = [
        TransactionOutput(winner_amount, script.Script.from_address(winner_address)),
        TransactionOutput(commission_sats, script.Script.from_address(fee_address)),
    ]
    if change > 0:
        vout.append(TransactionOutput(change, script.Script.from_address(change_address)))

    tx = Transaction(vin=vin, vout=vout)
    psbt = PSBT(tx)
    for i, utxo in enumerate(selected):
        psbt.inputs[i].witness_utxo = TransactionOutput(utxo.amount_sats, from_script)

    signed_count = psbt.sign_with(signing_key)
    if signed_count != len(selected):
        raise RuntimeError(f"expected {len(selected)} signatures, got {signed_count}")

    final_tx = finalize_psbt(psbt)
    if final_tx is None:
        raise RuntimeError("failed to finalize PSBT")

    raw = final_tx.serialize()
    return PayoutTransaction(
        raw_hex=raw.hex(),
        txid=final_tx.txid().hex(),
        fee_sats=fee,
        winner_sats=winner_amount,
        commission_sats=commission_sats,
        change_sats=change,
        spent_utxos=selected,
    )