rostrum 14.0.1

#!/usr/bin/env python3
# Copyright (c) 2019 The Bitcoin Unlimited developers
"""
Tests to check if basic electrum server integration works
"""

import asyncio
from test_framework.util import assert_equal, wait_for
from test_framework.electrumutil import (
    ElectrumTestFramework,
    script_to_scripthash,
)

# node, node_supports, NodeFeature unused - commented out token test
# from test_framework.environment import node, node_supports, NodeFeature
from test_framework.electrumconnection import ElectrumConnection
from test_framework.constants import COIN
from test_framework.script import CScript, OP_TRUE, OP_DROP
from test_framework.utilcompat import lock, unlock
from test_framework.utiltx import pad_tx
from test_framework.serialize import to_hex
from test_framework.environment import on_bch, on_nex

if on_nex():
    from test_framework.nex.script import anyonecanspend_unlocking_script
elif on_bch():
    from test_framework.bch.script import anyonecanspend_unlocking_script
else:
    raise NotImplementedError()

DEFAULT_WAIT = 30


def confirmed_count(history):
    # token.* adds the history list under 'history' key,
    # while address.* responses do not.
    if "history" in history:
        history = history["history"]

    return sum(map(lambda i: i["height"] > 0, history))


class ElectrumReorgTests(ElectrumTestFramework):
    async def run_test(self):
        n = self.nodes[0]

        await self.bootstrap_p2p()
        cli = ElectrumConnection()
        await cli.connect()
        try:
            # Mine blocks using P2P at the beginning to get coinbases for deterministic transactions
            coinbases = await self.mine_blocks(cli, n, 100)
            # for spending wallet in node
            n.generate(100)
            # to mature the spending wallet
            n.generate(110)
            await self.test_reorg(cli, n)
            await self.test_reorg_unspent_index_recreation(cli, n, coinbases)
            await self.test_reorg_scripthash_history(cli, n, coinbases)
        finally:
            cli.disconnect()

    async def test_reorg(self, cli, n):
        address = n.getnewaddress()

        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        # TODO: Token test is broken since it might re-use the funds from 'address' causing an extra tx
        test_tokens = False  # node_supports(node(), NodeFeature.TOKENS)

        # 1 coin transaction
        n.sendtoaddress(address, 10)
        await self.wait_for_mempool_count(cli, count=1)

        if test_tokens:
            # 2 token transactions
            token_id, _ = await self.create_token(
                to_addr=n.getnewaddress(), mint_amount=42
            )
            await self.wait_for_mempool_count(cli, count=3)
        await self.sync_mempool(cli)

        blocks = n.generate(50)

        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        # Check that our transactions got indexed.
        assert (
            confirmed_count(await cli.call("blockchain.address.get_history", address))
            >= 1
        )
        if test_tokens:
            assert_equal(
                2,
                confirmed_count(
                    await cli.call("token.transaction.get_history", token_id)
                ),
            )

        self.info("invalidating %d blocks", len(blocks))
        n.invalidateblock(blocks[0])
        # Clear cached block hashes after reorg to avoid using stale ancestor hashes
        self.hash_at_height.clear()

        # electrum server should trim its chain as well and maybe see our
        # transactions go back into mempool (depending on maturity of inputs)
        await self.sync_all(cli)

        # the transactions should no longer exist on-chain
        async def no_confirmed_txs():
            return 0 == confirmed_count(
                await cli.call("blockchain.address.get_history", address)
            )

        await wait_for(
            10,
            no_confirmed_txs,
            f"got confirmed count {confirmed_count(await cli.call("blockchain.address.get_history", address))}, expecting 0",
        )

        if test_tokens:
            assert_equal(
                0,
                confirmed_count(
                    await cli.call("token.transaction.get_history", token_id)
                ),
            )

        n.generate(50)
        await self.sync_height(cli)

    async def test_reorg_unspent_index_recreation(  # pylint: disable=too-many-locals,too-many-statements
        self, cli, n, coinbases
    ):
        """
        Test that unspent index entries are correctly recreated for outputs
        that were created in earlier blocks but spent in the block being undone.

        This test verifies the fix for the bug where outputs created in block N
        and spent in block N+1 were not being re-added to the unspent index
        when block N+1 was undone during a reorg.
        """
        coinbase = coinbases.pop(0)

        # Use deterministic scriptpubkeys instead of addresses
        scriptpubkey1 = CScript([OP_TRUE, OP_DROP])
        scriptpubkey2 = CScript([OP_TRUE, OP_TRUE, OP_DROP, OP_DROP])
        scripthash1 = script_to_scripthash(lock(scriptpubkey1))

        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        # Step 1: Create funding transaction (creates output in block N)
        self.info("Step 1: Creating output in block N")
        fee = 200
        # Ensure we send at least 10 * COIN for Nexa
        min_value = 10 * COIN
        tx1_value = max(coinbase.vout[0].n_value - fee, min_value)
        tx1 = self.create_transaction(
            coinbase,
            n=0,
            value=tx1_value,
            sig=anyonecanspend_unlocking_script(),
            out=lock(scriptpubkey1),
        )
        pad_tx(tx1)
        # txid1_hex unused - kept for potential debugging
        _ = tx1.get_rpc_hex_id()  # pylint: disable=unused-variable
        await cli.call("blockchain.transaction.broadcast", to_hex(tx1))
        await self.wait_for_mempool_count(cli, count=1)
        await self.sync_mempool(cli)

        # Step 2: Mine block N to confirm the output (use node to avoid P2P issues)
        self.info("Step 2: Mining block N to confirm output")
        n.generate(1)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        # Verify address1 has balance
        balance1 = await cli.call("blockchain.scripthash.get_balance", scripthash1)
        assert_equal(
            tx1_value,
            balance1["confirmed"],
            "scripthash1 should have expected coins confirmed",
        )
        assert_equal(
            0, balance1["unconfirmed"], "scripthash1 should have 0 unconfirmed"
        )

        # Step 3: Create spending transaction (spends output in block N+1)
        self.info("Step 3: Spending output in block N+1")
        # Store the original balance before spending
        original_balance = balance1["confirmed"]

        # Create deterministic spending transaction that spends from tx1's output
        # Send at least 10 * COIN for Nexa
        spend_value = max(tx1.vout[0].n_value - fee, 10 * COIN)
        tx2 = self.create_transaction(
            tx1,
            n=0,
            value=spend_value,
            sig=unlock(scriptpubkey1),
            out=lock(scriptpubkey2),
        )
        pad_tx(tx2)
        # txid2_hex unused - kept for potential debugging
        _ = tx2.get_rpc_hex_id()  # pylint: disable=unused-variable
        await cli.call("blockchain.transaction.broadcast", to_hex(tx2))
        await self.wait_for_mempool_count(cli, count=1)
        await self.sync_mempool(cli)

        # Step 4: Mine block N+1 to confirm the spending (use node to avoid P2P issues)
        self.info("Step 4: Mining block N+1 to confirm spending")
        n.generate(1)
        block_n_plus_1 = n.getbestblockhash()
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        # Verify balance changed (output was spent)
        balance1_after_spend = await cli.call(
            "blockchain.scripthash.get_balance", scripthash1
        )
        # The balance should be different from original (output was spent)
        assert (
            balance1_after_spend["confirmed"] != original_balance
        ), "scripthash1 balance should change after spending transaction"

        # Step 5: Invalidate block N+1 (reorg)
        self.info(
            "Step 5: Invalidating block N+1 to trigger reorg and test unspent index recreation"
        )
        n.invalidateblock(block_n_plus_1)
        # Clear cached block hashes after reorg to avoid using stale ancestor hashes
        self.hash_at_height.clear()
        # Sync height (reorg processing) - don't sync mempool as spending tx may be back in mempool
        await self.sync_height(cli)
        # Give rostrum a moment to process the reorg
        await asyncio.sleep(1)

        # Step 6: Verify balance is restored (unspent index was recreated)
        self.info("Step 6: Verifying unspent index was recreated")
        balance1_after_reorg = await cli.call(
            "blockchain.scripthash.get_balance", scripthash1
        )
        # After reorg, the spending transaction is undone, so confirmed balance should be restored
        # This is the key assertion: if unspent index is not recreated, this will fail
        assert_equal(
            original_balance,
            balance1_after_reorg["confirmed"],
            f"scripthash1 should have {original_balance} confirmed after reorg "
            "(unspent index should be recreated - output from block N should be restored)",
        )
        # Note: unconfirmed balance might be negative if the spending transaction
        # is back in mempool, which is expected behavior after a reorg

        # Cleaning up after test
        n.generate(1)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

    def _has_tx(self, history, txid):
        """Helper to check if a txid is in history"""
        for entry in history:
            if entry["tx_hash"] == txid:
                return True
        return False

    def _confirmed_history(self, history):
        """Helper to get confirmed history (height > 0)"""
        return [h for h in history if h["height"] > 0]

    async def _create_and_confirm_funding_tx(  # pylint: disable=too-many-arguments,too-many-positional-arguments
        self, cli, n, coinbase, scriptpubkey1, scripthash1
    ):
        """Create and confirm a funding transaction, return tx and txid"""
        self.info("Step 1: Creating output in block N")
        fee = 200
        tx1 = self.create_transaction(
            coinbase,
            n=0,
            value=coinbase.vout[0].n_value - fee,
            sig=anyonecanspend_unlocking_script(),
            out=lock(scriptpubkey1),
        )
        pad_tx(tx1)
        txid1_hex = tx1.get_rpc_hex_id()
        await cli.call("blockchain.transaction.broadcast", to_hex(tx1))
        await self.wait_for_mempool_count(cli, count=1)
        await self.sync_mempool(cli)

        self.info(
            "Step 2: Mining block N to confirm output (use node to avoid P2P issues)"
        )
        n.generate(1)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        # Wait for the funding transaction to appear in confirmed history
        confirmed_hist1 = None

        async def wait_for_conf():
            nonlocal confirmed_hist1
            history1 = await cli.call("blockchain.scripthash.get_history", scripthash1)
            confirmed_hist1 = self._confirmed_history(history1)
            # Check if our specific transaction is in the confirmed history
            return self._has_tx(confirmed_hist1, txid1_hex)

        await wait_for(
            10, wait_for_conf, "funding transaction should appear in confirmed history"
        )

        # Verify our transaction is in the confirmed history
        # Note: We don't assert exactly 1 transaction because the same scripthash
        # may be used in other tests, so there could be more transactions
        assert self._has_tx(
            confirmed_hist1, txid1_hex
        ), "history should contain funding tx"
        return tx1, txid1_hex

    async def _create_and_confirm_spending_tx(  # pylint: disable=too-many-arguments,too-many-positional-arguments
        self, cli, n, tx1, scriptpubkey1, scriptpubkey2, scripthash1, txid1_hex
    ):
        """Create and confirm a spending transaction, return txid and block hash"""
        self.info("Step 3: Spending output in block N+1")
        fee = 200
        tx2 = self.create_transaction(
            tx1,
            n=0,
            value=tx1.vout[0].n_value - fee,
            sig=unlock(scriptpubkey1),
            out=lock(scriptpubkey2),
        )
        pad_tx(tx2)
        txid2_hex = tx2.get_rpc_hex_id()
        await cli.call("blockchain.transaction.broadcast", to_hex(tx2))
        await self.wait_for_mempool_count(cli, count=1)
        await self.sync_mempool(cli)

        self.info(
            "Step 4: Mining block N+1 to confirm spending (use node to avoid P2P issues)"
        )
        n.generate(1)
        block_n_plus_1 = n.getbestblockhash()
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        # Wait for the spending transaction to appear in confirmed history
        confirmed_hist2 = None

        async def wait_for_spending():
            nonlocal confirmed_hist2
            history2 = await cli.call("blockchain.scripthash.get_history", scripthash1)
            confirmed_hist2 = self._confirmed_history(history2)
            # Check if both our transactions are in the confirmed history
            return self._has_tx(confirmed_hist2, txid1_hex) and self._has_tx(
                confirmed_hist2, txid2_hex
            )

        await wait_for(
            10,
            wait_for_spending,
            "both funding and spending transactions should appear in confirmed history",
        )

        # Verify both our transactions are in the confirmed history
        # Note: We don't assert exactly 2 transactions because the same scripthash
        # may be used in other tests, so there could be more transactions
        assert self._has_tx(
            confirmed_hist2, txid1_hex
        ), "history should contain funding tx"
        assert self._has_tx(
            confirmed_hist2, txid2_hex
        ), "history should contain spending tx"
        return txid2_hex, block_n_plus_1

    async def test_reorg_scripthash_history(self, cli, n, coinbases):
        """
        Test that scripthash history is correctly updated after reorgs.

        This verifies that when a block is undone:
        - Funding transactions are correctly restored in history
        - Spending transactions are correctly removed from confirmed history
        - The scripthash index is properly maintained
        """
        coinbase = coinbases.pop(0)

        scriptpubkey1 = CScript([OP_TRUE, OP_DROP])
        scriptpubkey2 = CScript([OP_TRUE, OP_TRUE, OP_DROP, OP_DROP])
        scripthash1 = script_to_scripthash(lock(scriptpubkey1))
        _ = script_to_scripthash(lock(scriptpubkey2))

        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        tx1, txid1_hex = await self._create_and_confirm_funding_tx(
            cli, n, coinbase, scriptpubkey1, scripthash1
        )

        txid2_hex, block_n_plus_1 = await self._create_and_confirm_spending_tx(
            cli, n, tx1, scriptpubkey1, scriptpubkey2, scripthash1, txid1_hex
        )

        self.info("Step 5: Invalidating block N+1 to trigger reorg")
        n.invalidateblock(block_n_plus_1)
        # Clear cached block hashes after reorg to avoid using stale ancestor hashes
        self.hash_at_height.clear()
        await self.sync_height(cli)
        await asyncio.sleep(1)

        self.info("Step 6: Verifying history after reorg")
        history3 = await cli.call("blockchain.scripthash.get_history", scripthash1)
        confirmed_hist3 = self._confirmed_history(history3)

        # Verify the funding transaction is still in confirmed history
        # Note: We don't assert exactly 1 transaction because the same scripthash
        # may be used in other tests, so there could be more transactions
        assert self._has_tx(
            confirmed_hist3, txid1_hex
        ), "history should still contain funding tx"
        # Verify the spending transaction is NOT in confirmed history (it was undone by reorg)
        assert not self._has_tx(
            confirmed_hist3, txid2_hex
        ), "history should NOT contain spending tx in confirmed (it was undone)"

        self.info("✓ Test passed: Scripthash history correctly updated after reorg")


if __name__ == "__main__":
    asyncio.run(ElectrumReorgTests().main())