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.config import settings
from app.db.models import PendingTransaction, User
from app.electrum.client import ElectrumClient
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 RBF_SEQUENCE, estimate_vsize

logger = logging.getLogger(__name__)

_POLL_INTERVAL_SECONDS = 30
_FEE_RATE_INCREMENT = 1  # minimum relay-policy-friendly bump per BIP125


class RbfError(Exception):
    pass


def should_bump(pending: PendingTransaction, now: datetime, timeout_seconds: int | None = None) -> bool:
    """Pure decision: has this pending tx been unconfirmed for longer than the
    configured timeout? Kept separate from the I/O-heavy bump_fee() so it's
    trivially unit-testable."""
    timeout = timeout_seconds if timeout_seconds is not None else settings.rbf_timeout_seconds
    if pending.status != "pending":
        return False
    return now >= pending.broadcast_at.replace(tzinfo=timezone.utc) + timedelta(seconds=timeout)


async def _signing_context(session: AsyncSession, pending: PendingTransaction) -> 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 pending.kind == "payout":
        key = derive_pool_key()
    else:
        user = await session.get(User, pending.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:
    txid_hex = vin.txid.hex()
    tx = await client.get_transaction(txid_hex, verbose=True)
    value_coins = tx["vout"][vin.vout]["value"]
    return round(value_coins * 100_000_000)


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: AsyncSession, client: ElectrumClient, pending: PendingTransaction) -> str:
    """Rebuild `pending`'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.

    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.
    """
    old_tx = Transaction.parse(bytes.fromhex(pending.raw_tx_hex))
    signing_key, own_script, own_address = await _signing_context(session, pending)

    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)

    new_fee_rate = pending.fee_rate_sat_vb + _FEE_RATE_INCREMENT
    new_fee = estimate_vsize(len(old_tx.vin), len(old_tx.vout)) * new_fee_rate
    fee_delta = new_fee - old_fee
    if fee_delta <= 0:
        fee_delta = _FEE_RATE_INCREMENT  # already above the new target vsize*rate; bump by a token amount

    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)

    pending.current_txid = new_txid
    pending.raw_tx_hex = raw_hex
    pending.fee_rate_sat_vb = new_fee_rate
    pending.attempt_count += 1
    pending.broadcast_at = datetime.now(timezone.utc)
    await session.commit()

    logger.info("bumped %s pending_transaction %s: %s -> %s", pending.kind, pending.id, pending.current_txid, new_txid)
    return 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:
            candidates = (
                await session.scalars(select(PendingTransaction).where(PendingTransaction.status == "pending"))
            ).all()
            due = [p for p in candidates if should_bump(p, now)]

        for pending in due:
            async with self._session_factory() as session:
                row = await session.get(PendingTransaction, pending.id)
                if row is None or row.status != "pending":
                    continue
                try:
                    await bump_fee(session, client, row)
                except RbfError:
                    logger.exception("could not bump pending_transaction %s", row.id)
                except Exception:
                    logger.exception("unexpected error bumping pending_transaction %s", row.id)