plm-lottery/tests/unit/test_confirmation_proof.py

"""B-83 (TX-22, TX-06, ELC-25, ELC-26, DEP-02, DEP-08, AUD-18, ADR-07): confirmations and
reinstatements need the quorum, not the active server's word.

`poll_once` used to mark a transaction confirmed on one server's `height > 0`, and
`reinstate_reappeared_utxos` restored a UTXO on its word — so a hostile active server could
declare a never-relayed payout confirmed (the winner unpaid, the round closed) or put a spent
coin back into a balance. Each test here fails against that.
"""

import hashlib
import struct

import pytest
from sqlalchemy import select

import app.bets.confirmation  # noqa: F401  (registers the "bet" handler)
from app.db.models import AuditLog, PendingTransaction, Round, RoundParticipant, User, UtxoEvent
from app.electrum.client import ElectrumEndpoint
from app.electrum.listener import ElectrumListener
from app.electrum.merkle import header_merkle_root, merkle_root_from_branch, proof_matches_header
from app.rounds.draw import header_meets_its_own_target
from app.tx.confirmation import poll_once
from app.wallet.hd import derive_user_address
from tests.unit.fake_network import FakeNetwork
from tests.unit.test_confirmation import FakeClient, session_factory  # noqa: F401  (fixture and fake)
from tests.unit.test_electrum_listener import (  # noqa: F401  (fixture and helpers)
    _ENDPOINTS,
    _EASY_BITS,
    _ActiveClient,
    _seed_funded_user,
    session_factory as listener_session_factory,
)


def _sha256d(data: bytes) -> bytes:
    return hashlib.sha256(hashlib.sha256(data).digest()).digest()


class Block:
    """A tiny block: `txids` (display order) under a real merkle tree, and a mined header
    that commits to its root."""

    def __init__(self, txids: list[str], prev_hash: str = "00" * 32):
        self.txids = txids
        level = [bytes.fromhex(t)[::-1] for t in txids]
        self._levels = [level]
        while len(level) > 1:
            if len(level) % 2:
                level = level + [level[-1]]
            level = [_sha256d(level[i] + level[i + 1]) for i in range(0, len(level), 2)]
            self._levels.append(level)
        self.root = level[0]
        self.header = self._mine(prev_hash)

    def _mine(self, prev_hash: str) -> str:
        for nonce in range(100_000):
            header = (
                struct.pack("<I", 1)
                + bytes.fromhex(prev_hash)[::-1]
                + self.root
                + struct.pack("<I", 0)
                + struct.pack("<I", _EASY_BITS)
                + struct.pack("<I", nonce)
            ).hex()
            if header_meets_its_own_target(header):
                return header
        raise RuntimeError("could not mine the test header")

    def proof(self, txid: str, height: int) -> dict:
        """What `blockchain.transaction.get_merkle` answers."""
        index = self.txids.index(txid)
        position, branch = index, []
        for level in self._levels[:-1]:
            padded = level + [level[-1]] if len(level) % 2 else level
            branch.append(padded[position ^ 1][::-1].hex())
            position //= 2
        return {"block_height": height, "merkle": branch, "pos": index}


_TXIDS = [f"{i:02x}" * 32 for i in range(1, 8)]  # 7 transactions: an odd level, so a duplicate sibling too
_BLOCK = Block(_TXIDS)


# --- the proof itself -----------------------------------------------------------------------


@pytest.mark.parametrize("txid", _TXIDS)
def test_every_transaction_of_a_block_proves_against_its_header(txid):
    assert proof_matches_header(txid, 100, _BLOCK.proof(txid, 100), _BLOCK.header)


def test_a_single_transaction_block_proves_with_an_empty_branch():
    block = Block(["ab" * 32])
    assert block.proof("ab" * 32, 5)["merkle"] == []
    assert proof_matches_header("ab" * 32, 5, block.proof("ab" * 32, 5), block.header)


def test_header_merkle_root_reads_the_committed_root():
    assert header_merkle_root(_BLOCK.header) == _BLOCK.root


def test_a_proof_does_not_prove_anything_it_was_not_made_for():
    txid = _TXIDS[2]
    proof = _BLOCK.proof(txid, 100)
    other_block = Block(["f0" * 32, "f1" * 32])
    assert not proof_matches_header(_TXIDS[3], 100, proof, _BLOCK.header)  # another txid
    assert not proof_matches_header(txid, 101, proof, _BLOCK.header)  # another height
    assert not proof_matches_header(txid, 100, proof, other_block.header)  # another block
    assert not proof_matches_header(txid, 100, {**proof, "pos": proof["pos"] ^ 1}, _BLOCK.header)  # another position
    assert not proof_matches_header(txid, 100, {**proof, "merkle": proof["merkle"][:-1]}, _BLOCK.header)
    tampered = list(proof["merkle"])
    tampered[0] = "00" * 32
    assert not proof_matches_header(txid, 100, {**proof, "merkle": tampered}, _BLOCK.header)


@pytest.mark.parametrize(
    "proof",
    [None, "garbage", {}, {"block_height": 100}, {"block_height": 100, "merkle": "x", "pos": 0},
     {"block_height": 100, "merkle": ["zz"], "pos": 0}, {"block_height": 100, "merkle": ["00" * 32], "pos": 9}],
)
def test_a_malformed_proof_is_simply_a_no(proof):
    assert proof_matches_header(_TXIDS[0], 100, proof, _BLOCK.header) is False


def test_the_root_computation_matches_a_hand_built_two_leaf_tree():
    a, b = "01" * 32, "02" * 32
    expected = _sha256d(bytes.fromhex(a)[::-1] + bytes.fromhex(b)[::-1])
    assert merkle_root_from_branch(a, [b], 0) == expected
    assert merkle_root_from_branch(b, [a], 1) == expected


# --- the listener's verdict: confirmation_proven ---------------------------------------------


class _ProofServer:
    """One other server: `proofs` maps txid -> the get_merkle reply; None = unreachable;
    `headers` maps height -> the header it reports."""

    def __init__(self, proofs=None, headers=None):
        self._proofs = proofs
        self._headers = headers or {}

    async def connect(self):
        if self._proofs is None:
            raise ConnectionRefusedError("down")

    async def get_merkle(self, txid, height):
        return self._proofs[txid]

    async def get_block_header(self, height):
        return self._headers[height]

    async def close(self):
        pass


def _proof_listener(session_factory, others: dict[str, _ProofServer], *, tip_height=100, tip_header=None, **kwargs):
    listener = ElectrumListener(lambda endpoint: others[endpoint.host], session_factory, _ENDPOINTS, **kwargs)
    listener.tip_height = tip_height
    listener.tip_header_hex = tip_header if tip_header is not None else _BLOCK.header
    return listener


def _honest():
    txid = _TXIDS[3]
    return _ProofServer({txid: _BLOCK.proof(txid, 100)})


def _lying():
    """Claims the tx is in the block, with a proof that does not add up."""
    txid = _TXIDS[3]
    proof = _BLOCK.proof(txid, 100)
    return _ProofServer({txid: {**proof, "merkle": ["11" * 32] + proof["merkle"][1:]}})


async def test_a_valid_proof_from_the_quorum_confirms(listener_session_factory):
    listener = _proof_listener(listener_session_factory, {"second.example": _honest(), "third.example": _honest()})
    assert await listener.confirmation_proven(_TXIDS[3], 100) is True


async def test_a_lying_active_server_cannot_confirm_a_transaction(listener_session_factory):
    """The active server claims height 100; the other servers' proofs do not check out.
    The active server's own word is never part of the verdict."""
    listener = _proof_listener(listener_session_factory, {"second.example": _lying(), "third.example": _lying()})
    assert await listener.confirmation_proven(_TXIDS[3], 100) is False


async def test_a_minority_of_liars_does_not_outvote_the_honest_servers(listener_session_factory):
    four = [*_ENDPOINTS, ElectrumEndpoint("fourth.example", 50002, True)]
    listener = ElectrumListener(
        lambda endpoint: {"second.example": _honest(), "third.example": _honest(), "fourth.example": _lying()}[endpoint.host],
        listener_session_factory,
        four,
    )
    listener.tip_height, listener.tip_header_hex = 100, _BLOCK.header
    assert await listener.confirmation_proven(_TXIDS[3], 100) is True

    flipped = ElectrumListener(
        lambda endpoint: {"second.example": _lying(), "third.example": _lying(), "fourth.example": _honest()}[endpoint.host],
        listener_session_factory,
        four,
    )
    flipped.tip_height, flipped.tip_header_hex = 100, _BLOCK.header
    assert await flipped.confirmation_proven(_TXIDS[3], 100) is False


async def test_too_few_answers_do_not_confirm(listener_session_factory):
    listener = _proof_listener(listener_session_factory, {"second.example": _honest(), "third.example": _ProofServer(None)})
    assert await listener.confirmation_proven(_TXIDS[3], 100) is False  # one answer, below the quorum minimum


async def test_a_proof_for_another_height_does_not_confirm(listener_session_factory):
    txid = _TXIDS[3]
    wrong_height = _ProofServer({txid: _BLOCK.proof(txid, 99)})
    listener = _proof_listener(listener_session_factory, {"second.example": wrong_height, "third.example": wrong_height})
    assert await listener.confirmation_proven(txid, 100) is False


async def test_a_height_we_have_not_seen_cannot_be_proven(listener_session_factory):
    listener = _proof_listener(
        listener_session_factory, {"second.example": _honest(), "third.example": _honest()}, tip_height=99
    )
    assert await listener.confirmation_proven(_TXIDS[3], 100) is False  # above our tip: no header to check against


async def test_an_older_block_is_checked_against_the_header_the_quorum_agrees_on(listener_session_factory):
    txid = _TXIDS[3]
    proof = _BLOCK.proof(txid, 90)
    servers = {
        "second.example": _ProofServer({txid: proof}, {90: _BLOCK.header}),
        "third.example": _ProofServer({txid: proof}, {90: _BLOCK.header}),
    }
    listener = _proof_listener(listener_session_factory, servers, tip_height=100, tip_header="00" * 80)
    assert await listener.confirmation_proven(txid, 90) is True

    # The servers disagree about that block's header: nothing trusted to check against.
    other = Block(["f0" * 32, "f1" * 32])
    split = {
        "second.example": _ProofServer({txid: proof}, {90: _BLOCK.header}),
        "third.example": _ProofServer({txid: proof}, {90: other.header}),
    }
    listener = _proof_listener(listener_session_factory, split, tip_height=100, tip_header="00" * 80)
    assert await listener.confirmation_proven(txid, 90) is False

    # A header that was never mined (does not meet its own target) is not accepted either.
    unmined = "00" * 80
    fake = {
        "second.example": _ProofServer({txid: proof}, {90: unmined}),
        "third.example": _ProofServer({txid: proof}, {90: unmined}),
    }
    listener = _proof_listener(listener_session_factory, fake, tip_height=100, tip_header="00" * 80)
    assert await listener.confirmation_proven(txid, 90) is False


async def test_the_reduced_rule_under_allow_no_quorum_accepts_a_single_proof(listener_session_factory):
    listener = _proof_listener(
        listener_session_factory, {"second.example": _honest(), "third.example": _ProofServer(None)}, allow_no_quorum=True
    )
    assert await listener.confirmation_proven(_TXIDS[3], 100) is True


# --- poll_once: the active server only nominates --------------------------------------------------


async def _seed_pending_bet(session_factory, txid: str = "tx1"):
    async with session_factory() as session:
        user = User(username="u0", password_hash="x", derivation_index=0, address=derive_user_address(0))
        session.add(user)
        await session.flush()
        session.add(Round(id=1, status="open"))
        session.add(RoundParticipant(round_id=1, user_id=user.id, bet_amount_sats=1_000, bet_txid=txid, status="broadcast"))
        session.add(
            PendingTransaction(
                kind="bet", round_id=1, user_id=user.id, current_txid=txid, fee_rate_sat_vb=1, raw_tx_hex="00", status="pending"
            )
        )
        await session.commit()


async def test_a_height_the_quorum_cannot_prove_leaves_the_row_pending(session_factory):
    """A hostile active server reports a bet (or a payout) as mined at some height. The
    quorum does not prove it, so nothing is confirmed and no handler runs."""
    await _seed_pending_bet(session_factory)
    network = FakeNetwork(proven=False)

    assert await poll_once(session_factory, FakeClient({"tx1": 100}), network) == 0

    assert network.proofs_asked == [("tx1", 100)]  # the claim of the active server was put to the quorum
    async with session_factory() as session:
        assert (await session.scalars(select(PendingTransaction))).one().status == "pending"
        assert (await session.scalars(select(RoundParticipant))).one().status == "broadcast"


async def test_a_proven_height_confirms(session_factory):
    await _seed_pending_bet(session_factory)
    network = FakeNetwork(proven=True)
    assert await poll_once(session_factory, FakeClient({"tx1": 100}), network) == 1
    assert network.proofs_asked == [("tx1", 100)]
    async with session_factory() as session:
        assert (await session.scalars(select(RoundParticipant))).one().status == "confirmed"


async def test_nothing_is_confirmed_without_a_network_to_prove_it(session_factory):
    await _seed_pending_bet(session_factory)
    assert await poll_once(session_factory, FakeClient({"tx1": 100}), None) == 0
    async with session_factory() as session:
        assert (await session.scalars(select(PendingTransaction))).one().status == "pending"


async def test_an_unconfirmed_transaction_is_not_put_to_the_quorum(session_factory):
    await _seed_pending_bet(session_factory)
    network = FakeNetwork()
    assert await poll_once(session_factory, FakeClient({"tx1": 0}), network) == 0
    assert network.proofs_asked == []


async def test_poll_once_confirms_through_a_real_listener_and_a_lying_active_server(session_factory, listener_session_factory):
    """End to end with the real quorum code: the transaction is really in the block, the
    other servers prove it, and the same call refuses when they do not."""
    txid = _TXIDS[3]
    await _seed_pending_bet(session_factory, txid=txid)

    liars = _proof_listener(listener_session_factory, {"second.example": _lying(), "third.example": _lying()})
    assert await poll_once(session_factory, FakeClient({txid: 100}), liars) == 0

    honest = _proof_listener(listener_session_factory, {"second.example": _honest(), "third.example": _honest()})
    assert await poll_once(session_factory, FakeClient({txid: 100}), honest) == 1


# --- reinstatement (ELC-25, DEP-02, DEP-08, AUD-18) ------------------------------------------------------


_FLAGGED_TXID = "dd" * 32
_ADDRESS = "plm1q0xcqpzrky6eff2g52qdye53xkk9jxkvraxkkwd"


class _UnspentServer:
    def __init__(self, entries):
        self._entries = entries

    async def connect(self):
        if self._entries is None:
            raise ConnectionRefusedError("down")

    async def listunspent(self, scripthash):
        return self._entries

    async def close(self):
        pass


_BACK = {"tx_hash": _FLAGGED_TXID, "tx_pos": 0, "height": 100, "value": 20_000_000}
_OTHER = {"tx_hash": "99" * 32, "tx_pos": 0, "height": 100, "value": 1_000_000}  # keeps the reply non-empty


async def _flagged_user(session_factory) -> int:
    user_id = await _seed_funded_user(session_factory, username="carol", address=_ADDRESS)
    async with session_factory() as session:
        utxo = (await session.scalars(select(UtxoEvent).where(UtxoEvent.user_id == user_id))).one()
        utxo.spent_txid = "external-spend"
        user = await session.get(User, user_id)
        user.cached_balance_sats = 0
        await session.commit()
    return user_id


def _reinstate_listener(session_factory, others: dict[str, _UnspentServer], active_entries):
    listener = ElectrumListener(lambda endpoint: others[endpoint.host], session_factory, _ENDPOINTS)
    listener.tip_height = 1000  # B-82: far past every credit threshold, coinbase maturity included
    listener.client = _ActiveClient(active_entries)
    return listener


async def _state(session_factory, user_id):
    async with session_factory() as session:
        utxo = (await session.scalars(select(UtxoEvent).where(UtxoEvent.txid == _FLAGGED_TXID))).one()
        events = [e.event_type for e in (await session.scalars(select(AuditLog))).all()]
        return utxo.spent_txid, (await session.get(User, user_id)).cached_balance_sats, events


async def test_a_reappeared_coin_is_reinstated_once_the_quorum_confirms_it(listener_session_factory):
    user_id = await _flagged_user(listener_session_factory)
    others = {"second.example": _UnspentServer([_BACK, _OTHER]), "third.example": _UnspentServer([_BACK, _OTHER])}
    listener = _reinstate_listener(listener_session_factory, others, [_BACK, _OTHER])

    await listener.refresh_user(user_id, "scripthash")

    spent, balance, events = await _state(listener_session_factory, user_id)
    assert spent is None and balance == 20_000_000 + 1_000_000  # the coin is back (+ the unrelated credit)
    assert "utxo_external_spend_reinstated" in events


@pytest.mark.parametrize(
    "others",
    [
        pytest.param({"second.example": [_OTHER], "third.example": [_OTHER]}, id="others-do-not-list-it"),
        pytest.param({"second.example": [_BACK, _OTHER], "third.example": [_OTHER]}, id="a-minority-lists-it"),
        pytest.param({"second.example": [_BACK, _OTHER], "third.example": None}, id="too-few-answers"),
        pytest.param(
            {"second.example": [{**_BACK, "value": 1}, _OTHER], "third.example": [{**_BACK, "value": 1}, _OTHER]},
            id="other-amount",
        ),
        pytest.param(
            {"second.example": [{**_BACK, "height": 0}, _OTHER], "third.example": [{**_BACK, "height": 0}, _OTHER]},
            id="unconfirmed",
        ),
    ],
)
async def test_a_hostile_active_server_cannot_reinstate_a_coin_alone(listener_session_factory, others):
    """The active server lists the coin again; the quorum does not bear it out. Before the
    fix the reinstatement was written on the active server's word alone."""
    user_id = await _flagged_user(listener_session_factory)
    servers = {host: _UnspentServer(entries) for host, entries in others.items()}
    listener = _reinstate_listener(listener_session_factory, servers, [_BACK, _OTHER])

    await listener.refresh_user(user_id, "scripthash")

    spent, balance, events = await _state(listener_session_factory, user_id)
    assert spent == "external-spend"  # still written off
    assert "utxo_external_spend_reinstated" not in events


async def test_a_coin_listed_with_another_amount_than_recorded_is_not_even_a_candidate(listener_session_factory):
    from app.deposits.service import find_reinstate_candidates

    user_id = await _flagged_user(listener_session_factory)
    async with listener_session_factory() as session:
        assert await find_reinstate_candidates(session, user_id, [_BACK]) == [_BACK]
        assert await find_reinstate_candidates(session, user_id, [{**_BACK, "value": 5}]) == []
        assert await find_reinstate_candidates(session, user_id, [{**_BACK, "height": 0}]) == []
        assert await find_reinstate_candidates(session, user_id, [_OTHER]) == []


async def test_finding_candidates_writes_nothing(listener_session_factory):
    from app.deposits.service import find_reinstate_candidates

    user_id = await _flagged_user(listener_session_factory)
    async with listener_session_factory() as session:
        await find_reinstate_candidates(session, user_id, [_BACK])
    spent, balance, events = await _state(listener_session_factory, user_id)
    assert spent == "external-spend" and balance == 0 and "utxo_external_spend_reinstated" not in events


async def test_the_reinstatement_write_rechecks_that_the_mark_is_still_there(listener_session_factory):
    """DEP-08: a coin whose mark was removed meanwhile (or never carried one) is not touched."""
    from app.deposits.service import reinstate_reappeared_utxos

    user_id = await _seed_funded_user(listener_session_factory, username="dave", address=_ADDRESS)  # not flagged
    async with listener_session_factory() as session:
        assert await reinstate_reappeared_utxos(session, user_id, [_BACK]) == 0