plm-lottery/app/tx/broadcast.py

import asyncio
import logging
from datetime import datetime, timedelta, timezone

from embit import script
from embit.psbt import PSBT
from embit.transaction import Transaction, TransactionInput, TransactionOutput
from embit.finalizer import finalize_psbt
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from app.db.models import PendingTransaction, Round, RoundParticipant, User, UtxoEvent, Withdrawal
from app.electrum.client import ElectrumClient
from app.rounds.config import get_round_config
from app.wallet.hd import derive_pool_key, derive_user_key
from app.wallet.plm_network import PLM_MAINNET
from app.wallet.psbt_builder import MAX_FEE_RATE_SAT_VB, RBF_SEQUENCE, estimate_vsize

logger = logging.getLogger(__name__)

_POLL_INTERVAL_SECONDS = 30
_FEE_RATE_INCREMENT = 1  # how much pending.fee_rate_sat_vb's *target* rises by per bump

# BIP125 rule 4: a replacement transaction must pay at least this much more, in
# total, per vbyte of its own size, than the transaction it replaces — Bitcoin
# Core's default incremental relay fee. bump_fee's delta must never fall below
# this regardless of what the target-rate arithmetic comes out to (B-32).
_INCREMENTAL_RELAY_FEE_RATE_SAT_VB = 1


class RbfError(Exception):
    pass


def should_bump(pending: PendingTransaction, now: datetime, timeout_seconds: int) -> bool:
    """Pure decision: has this pending tx gone unconfirmed for longer than the
    configured timeout (RoundConfig.rbf_timeout_seconds) *since it was last
    broadcast*? Kept separate from the I/O-heavy bump_fee() so it's trivially
    unit-testable.

    Deliberately measured from last_broadcast_at, not broadcast_at: this decides
    whether *another* bump is due, which should reset after every bump (a tx just
    rebroadcast at a higher fee deserves the same grace period again) — unlike
    reconcile.py's abandon check, which must measure from the *first* broadcast so
    repeated bumping can't indefinitely postpone ever giving up on a tx (B-27)."""
    if pending.status != "pending":
        return False
    return now >= pending.last_broadcast_at.replace(tzinfo=timezone.utc) + timedelta(seconds=timeout_seconds)


async def _signing_context(session: AsyncSession, kind: str, user_id: int | None) -> tuple:
    """Returns (signing_key, own_script, own_address) for the single sender that
    controls every input of this tx — a user for bet/withdrawal, the pool for
    payout. All our builders only ever spend one address's UTXOs per tx."""
    if kind == "payout":
        key = derive_pool_key()
    else:
        user = await session.get(User, user_id)
        key = derive_user_key(user.derivation_index)
    own_script = script.p2wpkh(key.to_public())
    own_address = own_script.address(network=PLM_MAINNET)
    return key, own_script, own_address


async def _prevout_amount(client: ElectrumClient, vin: TransactionInput) -> int:
    """The exact integer satoshi value of the output this input spends.

    Parsed directly from the raw transaction via embit rather than asking the
    server for its own float, whole-coin-denominated "value" field (verbose=True)
    and converting with `* 100_000_000` — embit's TransactionOutput.value is
    already an integer number of satoshis straight from the tx's binary
    encoding, so this never touches floating point in a codebase that is
    otherwise strictly integer-satoshi (B-40).
    """
    txid_hex = vin.txid.hex()
    raw_hex = await client.get_transaction(txid_hex, verbose=False)
    prevout_tx = Transaction.parse(bytes.fromhex(raw_hex))
    return prevout_tx.vout[vin.vout].value


def _find_change_output(tx: Transaction, change_address: str) -> int | None:
    for i, out in enumerate(tx.vout):
        if out.script_pubkey.address(network=PLM_MAINNET) == change_address:
            return i
    return None


async def bump_fee(
    session_factory: async_sessionmaker, client: ElectrumClient, pending_id: int
) -> str | None:
    """Rebuild pending_transaction `pending_id`'s transaction with a higher fee
    (same inputs, same recipient outputs, the extra fee taken from the change
    output) and rebroadcast. Returns the new txid, or None if there was nothing
    to do (the row is gone or already left "pending" — a normal race with
    confirmation, not an error).

    Three phases, so no DB session is held across the network calls this needs
    (one get_transaction per input, then a broadcast) — the same shape used
    elsewhere for exactly this reason (B-18, rounds/scheduler.py:_trigger_payout;
    B-31, electrum/listener.py:refresh_user) and now here too (B-40): read what's
    needed and close the session, do the chain work, then reopen to persist.

    Only handles the common case: exactly one change output paying back to the
    tx's own sender address, large enough to absorb the increase. If there's no
    such output (e.g. an exact-amount bet with no change), this raises RbfError —
    bumping such a tx would require selecting additional inputs, which isn't
    implemented for the MVP; it needs manual operator intervention. Also raises
    RbfError, rather than bumping, once the row is already at MAX_FEE_RATE_SAT_VB
    (B-32) — the reconciler abandons it if it never confirms (B-27), instead of
    this retrying an ever-higher fee forever.
    """
    # --- Phase 1: read what's needed, close the session before any network call ---
    async with session_factory() as session:
        pending = await session.get(PendingTransaction, pending_id)
        if pending is None or pending.status != "pending":
            logger.info("pending_transaction %s no longer pending; skipping bump", pending_id)
            return None
        if pending.fee_rate_sat_vb >= MAX_FEE_RATE_SAT_VB:
            raise RbfError(
                f"pending_transaction {pending_id}: already at the maximum fee rate "
                f"({MAX_FEE_RATE_SAT_VB} sat/vB) — refusing to bump further"
            )

        kind = pending.kind
        current_fee_rate = pending.fee_rate_sat_vb
        raw_tx_hex = pending.raw_tx_hex
        signing_key, own_script, own_address = await _signing_context(session, kind, pending.user_id)

    # --- Phase 2: chain reads, signing, and the broadcast — no DB session open ----
    old_tx = Transaction.parse(bytes.fromhex(raw_tx_hex))
    input_amounts = [await _prevout_amount(client, vin) for vin in old_tx.vin]
    total_in = sum(input_amounts)
    old_fee = total_in - sum(o.value for o in old_tx.vout)
    vsize = estimate_vsize(len(old_tx.vin), len(old_tx.vout))

    target_fee_rate = min(current_fee_rate + _FEE_RATE_INCREMENT, MAX_FEE_RATE_SAT_VB)
    target_fee = vsize * target_fee_rate
    # BIP125 rule 4's minimum, in absolute sats for this tx's size — the floor
    # `fee_delta` must never go below, no matter what `target_fee - old_fee` comes
    # out to. That naive difference used to go to zero or negative whenever
    # old_fee already exceeded target_fee (e.g. a dust change amount folded into
    # the original fee — wallet/psbt_builder.py's DUST_LIMIT_SATS handling), and
    # the previous fallback — a flat 1-satoshi total bump — was nowhere near this
    # relay-mandated minimum, so the node rejected it every time. Because bump_fee
    # raised before touching `pending`, the next tick retried with identical
    # parameters every 30 seconds, forever (B-32).
    min_valid_delta = vsize * _INCREMENTAL_RELAY_FEE_RATE_SAT_VB
    fee_delta = max(target_fee - old_fee, min_valid_delta)

    change_index = _find_change_output(old_tx, own_address)
    if change_index is None or old_tx.vout[change_index].value <= fee_delta:
        raise RbfError(f"pending_transaction {pending_id}: no change output large enough to absorb a fee bump")

    new_vout = list(old_tx.vout)
    bumped_change = new_vout[change_index].value - fee_delta
    new_vout[change_index] = TransactionOutput(bumped_change, new_vout[change_index].script_pubkey)

    new_vin = [TransactionInput(v.txid, v.vout, sequence=RBF_SEQUENCE) for v in old_tx.vin]
    new_tx = Transaction(vin=new_vin, vout=new_vout)
    psbt = PSBT(new_tx)
    for i, amount in enumerate(input_amounts):
        psbt.inputs[i].witness_utxo = TransactionOutput(amount, own_script)

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

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

    raw_hex = final_tx.serialize().hex()
    new_txid = final_tx.txid().hex()
    await client.broadcast(raw_hex)

    # --- Phase 3: persist the outcome ----------------------------------------------
    async with session_factory() as session:
        pending = await session.get(PendingTransaction, pending_id)
        old_txid = pending.current_txid
        pending.replaced_by_txid = old_txid  # points backwards: what current_txid replaced
        pending.current_txid = new_txid
        pending.raw_tx_hex = raw_hex
        # The *actual* resulting rate, not target_fee_rate: when the BIP125-minimum
        # floor above raised fee_delta past the naive target, the tx now pays more
        # than target_fee_rate implied. Recording the true rate keeps the next bump's
        # arithmetic honest instead of drifting from what's really being paid.
        pending.fee_rate_sat_vb = (old_fee + fee_delta) // vsize
        pending.attempt_count += 1
        # last_broadcast_at, not broadcast_at (B-27): broadcast_at must stay the *first*
        # broadcast, since reconcile.py's abandon-after-N-hours grace period is measured
        # from it — overwriting it here used to reset that clock on every bump, so a
        # repeatedly-bumped-but-never-mined tx was never abandoned.
        pending.last_broadcast_at = datetime.now(timezone.utc)
        await _retarget_txid_references(session, pending, old_txid, new_txid)
        await session.commit()

    logger.info("bumped %s pending_transaction %s: %s -> %s", kind, pending_id, old_txid, new_txid)
    return new_txid


async def _retarget_txid_references(
    session: AsyncSession, pending: PendingTransaction, old_txid: str, new_txid: str
) -> None:
    """A bump changes the txid, and everything that recorded the old one has to
    follow — otherwise the bumped tx confirms and nothing recognizes it (B-02).

    The worst case was the bet path: _on_bet_confirmed used to look the participant
    up by bet_txid, so after a bump it found nothing, the participant stayed
    "broadcast" forever, and the scheduler waited on it forever — the round could
    never close and the lottery stopped. The handlers now key off immutable ids
    (round_id/user_id, withdrawal_id) as well, so this update is about keeping the
    stored txids *true* — for the admin UI, for the audit trail, and for
    reconcile.py, which matches UtxoEvent.spent_txid against current_txid.
    """
    if pending.kind == "bet":
        participant = await session.scalar(
            select(RoundParticipant).where(
                RoundParticipant.round_id == pending.round_id,
                RoundParticipant.user_id == pending.user_id,
            )
        )
        if participant is not None:
            participant.bet_txid = new_txid
    elif pending.kind == "withdrawal" and pending.withdrawal_id is not None:
        withdrawal = await session.get(Withdrawal, pending.withdrawal_id)
        if withdrawal is not None:
            withdrawal.txid = new_txid
    elif pending.kind == "payout" and pending.round_id is not None:
        round_ = await session.get(Round, pending.round_id)
        if round_ is not None and round_.payout_txid == old_txid:
            round_.payout_txid = new_txid

    # The UTXOs this tx spends are still the same UTXOs — only the id of the tx
    # spending them changed. Keeping this in step is what lets reconcile.py tell
    # "reserved by this pending tx" from "spent by something else".
    spent = (await session.scalars(select(UtxoEvent).where(UtxoEvent.spent_txid == old_txid))).all()
    for utxo in spent:
        utxo.spent_txid = new_txid


class RbfBumper:
    def __init__(self, session_factory: async_sessionmaker, get_client):
        self._session_factory = session_factory
        self._get_client = get_client

    async def run(self) -> None:
        while True:
            client = self._get_client()
            if client is not None:
                try:
                    await self._tick(client)
                except asyncio.CancelledError:
                    raise
                except Exception:
                    logger.exception("RBF bump tick failed")
            await asyncio.sleep(_POLL_INTERVAL_SECONDS)

    async def _tick(self, client: ElectrumClient) -> None:
        now = datetime.now(timezone.utc)
        async with self._session_factory() as session:
            timeout_seconds = (await get_round_config(session)).rbf_timeout_seconds
            candidates = (
                await session.scalars(select(PendingTransaction).where(PendingTransaction.status == "pending"))
            ).all()
            due_ids = [p.id for p in candidates if should_bump(p, now, timeout_seconds=timeout_seconds)]

        for pending_id in due_ids:
            try:
                await bump_fee(self._session_factory, client, pending_id)
            except RbfError:
                logger.exception("could not bump pending_transaction %s", pending_id)
            except Exception:
                logger.exception("unexpected error bumping pending_transaction %s", pending_id)