plm-lottery/app/tx/broadcast.py

import asyncio
import json
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, update
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from app.audit.log import write_audit_log
from app.db.base import begin_immediate
from app.db.models import (
    AuditLog,
    PendingTransaction,
    PendingTransactionTxid,
    Round,
    RoundParticipant,
    User,
    UtxoEvent,
    Withdrawal,
)
from app.electrum.client import ElectrumClient
from app.rounds.config import get_round_config
from app.tx.outcome import BroadcastOutcome, TransactionNetwork, broadcast_with_verdict
from app.tx.versions import record_version
from app.wallet.hd import derive_pool_address, derive_round_pool_key, derive_user_key
from app.wallet.plm_network import PLM_MAINNET
from app.wallet.psbt_builder import DUST_LIMIT_SATS, 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


# B-87 (TX-24, AUD-40): why a bump could not be made — the `reason` of the `rbf_bump_failed`
# audit row the bumper writes.
REASON_CHANGE_TOO_SMALL = "change_too_small"
REASON_MAX_FEE_RATE = "max_fee_rate"
REASON_REPLACEMENT_REJECTED = "replacement_rejected"


class RbfError(Exception):
    def __init__(self, message: str, reason: str = REASON_CHANGE_TOO_SMALL):
        super().__init__(message)
        self.reason = reason


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, round_id: int | None = 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 (the round's own pool key, B-80). All our builders only ever spend one
    address's UTXOs per tx."""
    if kind == "payout":
        round_ = await session.get(Round, round_id) if round_id is not None else None
        if round_ is None or not round_.pool_address:
            raise RbfError(f"payout of round {round_id}: the round has no pool address")
        key = derive_round_pool_key(round_.id, round_.pool_address)
    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


def _find_reducible_output(tx: Transaction, kind: str, own_address: str) -> int | None:
    """The one output an RBF bump takes the extra fee from (TX-11): the sender's own change
    for a bet or a withdrawal; for a payout the **bump reserve** output paid to the reserve
    address (B-79, RND-38), falling back to a change output back to the pool for a payout
    that was already in flight before the reserve existed. Counterparty outputs — the
    winner's, the fee address's, an external recipient's — are never candidates."""
    if kind == "payout":
        reserve_index = _find_change_output(tx, derive_pool_address())
        if reserve_index is not None:
            return reserve_index
    return _find_change_output(tx, own_address)


async def bump_fee(
    session_factory: async_sessionmaker,
    client: ElectrumClient,
    pending_id: int,
    network: TransactionNetwork | None = None,
) -> 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, or a conflict recorded below).

    Four 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.

    B-77 (TX-12, TX-13, DAT-34): the replacement is **written before it is broadcast**.
    It used to be sent in the chain phase and recorded afterwards, unconditionally, so
    a crash or a `database is locked` in between left a transaction on the network
    that the database did not know, and an abandon or a confirmation that happened in
    the meantime was silently overwritten. Now the signed replacement is committed as
    a *candidate* version (pending_transaction_txids) first; a candidate that is never
    broadcast is harmless, and one that is — and then confirms — is recognised like any
    other version. After the broadcast the row is switched over with a conditional
    UPDATE (status still pending, current_txid still the one replaced); zero rows
    matched is a conflict: the row is left alone, the candidate stays recorded, and
    the reconciler resolves it.

    Only handles one shape: exactly one output to take the increase from — the sender's
    own change, or a payout's bump reserve (B-79) — and when what would be left of it is
    below the dust limit the output is dropped and its whole value goes to the fee, which
    makes that replacement the last possible one (B-87, TX-11). With no such output, or
    when even its whole value is below BIP125's minimum increase, this raises RbfError —
    adding inputs would not help the case that used to hurt (a transaction that already
    spends everything the sender has). Also raises RbfError, rather than bumping, once the
    row is already at MAX_FEE_RATE_SAT_VB (B-32), and when the network explicitly refuses
    the replacement (TX-24). Each carries a `reason` for the `rbf_bump_failed` audit row
    (RbfBumper._record_failure).
    """
    # --- 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",
                REASON_MAX_FEE_RATE,
            )

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

    # --- Phase 2: chain reads and signing — no DB session open, nothing sent yet ---
    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_reducible_output(old_tx, kind, own_address)
    if change_index is None:
        raise RbfError(f"pending_transaction {pending_id}: there is no output to take a fee bump from")
    reducible_value = old_tx.vout[change_index].value
    if reducible_value < min_valid_delta:
        # Not even giving the whole output to the fee reaches BIP125's minimum increase.
        raise RbfError(
            f"pending_transaction {pending_id}: its change ({reducible_value} sat) cannot even cover the "
            f"minimum relay increase ({min_valid_delta} sat)"
        )

    new_vout = list(old_tx.vout)
    bumped_change = reducible_value - fee_delta
    if bumped_change >= DUST_LIMIT_SATS:
        new_vout[change_index] = TransactionOutput(bumped_change, new_vout[change_index].script_pubkey)
    else:
        # B-87 (TX-11): what would be left is below the dust limit, an output no node relays — a
        # 294 sat change shrinks to ~153 after the very first minimum bump, and every replacement
        # kept by it was refused, every 30 seconds, for ever. Drop the output instead and give its
        # whole value to the fee: the replacement is valid, and it is the last one this
        # transaction can have, since nothing is left to reduce.
        del new_vout[change_index]
        fee_delta = reducible_value

    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()

    # --- Phase 3: record the replacement as a candidate, *then* broadcast ----------
    async with session_factory() as session:
        await begin_immediate(session)  # B-88
        pending = await session.get(PendingTransaction, pending_id)
        if pending is None or pending.status != "pending" or pending.current_txid != old_txid:
            logger.info("pending_transaction %s changed while its bump was being built; skipping", pending_id)
            return None
        await record_version(session, pending, new_txid, raw_hex)  # B-77 (TX-12, DAT-71)
        await session.commit()

    verdict = await broadcast_with_verdict(client, network, raw_hex, new_txid)  # B-76: also fans out
    if verdict.outcome is BroadcastOutcome.REFUSED:
        # The network explicitly refused the replacement and no server knows it: the row
        # is unchanged and keeps the version that is on the network (TX-24). The
        # candidate stays recorded — it was never relayed, so nobody will confirm it.
        raise RbfError(
            f"pending_transaction {pending_id}: the network refused the replacement: {verdict.reason}",
            REASON_REPLACEMENT_REJECTED,
        )
    # ACCEPTED, or UNCERTAIN: the replacement may well be on the network, so the row
    # follows it (TX-12); the old version stays recorded and is still recognised should
    # it be the one that confirms.

    # --- Phase 4: switch the row over, conditionally --------------------------------
    async with session_factory() as session:
        await begin_immediate(session)  # B-88
        # 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.
        # 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.
        result = await session.execute(
            update(PendingTransaction)
            .where(
                PendingTransaction.id == pending_id,
                PendingTransaction.status == "pending",
                PendingTransaction.current_txid == old_txid,
            )
            .values(
                replaced_by_txid=old_txid,  # points backwards: what current_txid replaced
                current_txid=new_txid,
                raw_tx_hex=raw_hex,
                fee_rate_sat_vb=(old_fee + fee_delta) // estimate_vsize(len(new_vin), len(new_vout)),
                attempt_count=PendingTransaction.attempt_count + 1,
                last_broadcast_at=datetime.now(timezone.utc),
            )
            .execution_options(synchronize_session=False)
        )
        if result.rowcount != 1:
            # Confirmed, abandoned or already re-bumped while we were broadcasting. The
            # row is not overwritten; the candidate is still recorded, and confirmation
            # and reconciliation look for it like any other version (TX-13).
            await session.rollback()
            logger.warning(
                "bump of pending_transaction %s: the row changed during the broadcast (replaced %s by %s); "
                "left to the reconciler",
                pending_id,
                old_txid,
                new_txid,
            )
            return None
        pending = await session.get(PendingTransaction, pending_id)
        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 adopt_version(session: AsyncSession, pending: PendingTransaction, txid: str) -> None:
    """The network confirmed (or holds) a version of this row other than the current one
    — an earlier version of a bumped transaction, or a candidate whose switch-over never
    happened (B-77, TX-19, DAT-14). `current_txid` and `raw_tx_hex` become that
    version's, and the txid is propagated everywhere it is recorded, exactly as after a
    bump; without it the inputs of a replaced transaction would no longer be
    recognised as the row's (DAT-13)."""
    if txid == pending.current_txid:
        return
    version = await session.scalar(
        select(PendingTransactionTxid).where(
            PendingTransactionTxid.txid == txid, PendingTransactionTxid.pending_transaction_id == pending.id
        )
    )
    old_txid = pending.current_txid
    pending.replaced_by_txid = old_txid
    pending.current_txid = txid
    if version is not None and version.raw_hex:
        pending.raw_tx_hex = version.raw_hex  # empty only for a migration-reconstructed txid
    await _retarget_txid_references(session, pending, old_txid, txid)
    logger.warning("pending_transaction %s: version %s found on the network, replacing %s", pending.id, txid, old_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, network: TransactionNetwork | None = None):
        self._session_factory = session_factory
        self._get_client = get_client
        self._network = network  # B-76/B-77: verifies and fans out a replacement

    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 _record_failure(self, pending_id: int, exc: RbfError) -> None:
        """B-87 (TX-24, AUD-40, AUD-09): a bump that cannot be made is visible from /admin as one
        `rbf_bump_failed` row — in a transaction of its own, the row stays `pending` — and only
        **once per (current txid, reason)**: the same failure on the same transaction every 30
        seconds goes to the log alone, while a new txid (after a successful bump) or another
        reason writes a new row."""
        try:
            async with self._session_factory() as session:
                pending = await session.get(PendingTransaction, pending_id)
                if pending is None:
                    return
                already = (
                    await session.scalars(
                        select(AuditLog.payload_json).where(
                            AuditLog.event_type == "rbf_bump_failed",
                            AuditLog.payload_json.contains(f'"txid": "{pending.current_txid}"'),
                        )
                    )
                ).all()
                for payload_json in already:
                    payload = json.loads(payload_json)
                    if payload.get("txid") == pending.current_txid and payload.get("reason") == exc.reason:
                        return
                await write_audit_log(
                    session,
                    "rbf_bump_failed",
                    {
                        "pending_transaction_id": pending.id,
                        "kind": pending.kind,
                        "txid": pending.current_txid,
                        "reason": exc.reason,
                        "message": str(exc)[:200],
                    },
                    user_id=pending.user_id,
                    round_id=pending.round_id,
                )
                await session.commit()
        except Exception:
            logger.exception("could not record the failed bump of pending_transaction %s", pending_id)

    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, network=self._network)
            except RbfError as exc:
                logger.exception("could not bump pending_transaction %s", pending_id)
                await self._record_failure(pending_id, exc)
            except Exception:
                logger.exception("unexpected error bumping pending_transaction %s", pending_id)