Improve safety of coin store (Chia-Network#2984)

* Improve safety of coin store

* Add test for exception

* Fix issue with fork_point_with_peak

chia/consensus/blockchain.py

@@ -32,7 +32,7 @@
 from chia.types.unfinished_header_block import UnfinishedHeaderBlock
 from chia.types.weight_proof import SubEpochChallengeSegment
 from chia.util.errors import Err
-from chia.util.generator_tools import get_block_header, block_removals_and_additions
+from chia.util.generator_tools import get_block_header, tx_removals_and_additions
 from chia.util.ints import uint16, uint32, uint64, uint128
 from chia.util.streamable import recurse_jsonify
 
@@ -206,10 +206,10 @@ async def receive_block(
                     npc_result = get_name_puzzle_conditions(
                         block_generator, min(self.constants.MAX_BLOCK_COST_CLVM, block.transactions_info.cost), False
                     )
-                    removals, additions = block_removals_and_additions(block, npc_result.npc_list)
+                    removals, tx_additions = tx_removals_and_additions(npc_result.npc_list)
                 else:
-                    removals, additions = [], list(block.get_included_reward_coins())
-                header_block = get_block_header(block, additions, removals)
+                    removals, tx_additions = [], []
+                header_block = get_block_header(block, tx_additions, removals)
             else:
                 npc_result = None
                 header_block = get_block_header(block, [], [])
@@ -314,11 +314,10 @@ async def _reconsider_peak(
                 # Begins a transaction, because we want to ensure that the coin store and block store are only updated
                 # in sync.
                 if npc_result is not None:
-                    removals, additions = block_removals_and_additions(block, npc_result.npc_list)
+                    tx_removals, tx_additions = tx_removals_and_additions(npc_result.npc_list)
                 else:
-                    removals = []
-                    additions = list(block.get_included_reward_coins())
-                await self.coin_store.new_block(block, additions, removals)
+                    tx_removals, tx_additions = [], []
+                await self.coin_store.new_block(block, tx_additions, tx_removals)
                 await self.block_store.set_peak(block.header_hash)
                 return uint32(0), uint32(0), [block_record]
             return None, None, []
@@ -327,12 +326,13 @@ async def _reconsider_peak(
         if block_record.weight > peak.weight:
             # Find the fork. if the block is just being appended, it will return the peak
             # If no blocks in common, returns -1, and reverts all blocks
-            if fork_point_with_peak is not None:
-                fork_height: int = fork_point_with_peak
+            if block_record.prev_hash == peak.header_hash:
+                fork_height: int = peak.height
+            elif fork_point_with_peak is not None:
+                fork_height = fork_point_with_peak
             else:
                 fork_height = find_fork_point_in_chain(self, block_record, peak)
 
-            # Rollback to fork
             if block_record.prev_hash != peak.header_hash:
                 await self.coin_store.rollback_to_block(fork_height)
             # Rollback sub_epoch_summaries
@@ -364,10 +364,12 @@ async def _reconsider_peak(
                 records_to_add.append(fetched_block_record)
                 if fetched_block_record.is_transaction_block:
                     if fetched_block_record.header_hash == block_record.header_hash:
-                        removals, additions = await self.get_removals_and_additions(fetched_full_block, npc_result)
+                        tx_removals, tx_additions = await self.get_tx_removals_and_additions(
+                            fetched_full_block, npc_result
+                        )
                     else:
-                        removals, additions = await self.get_removals_and_additions(fetched_full_block, None)
-                    await self.coin_store.new_block(fetched_full_block, additions, removals)
+                        tx_removals, tx_additions = await self.get_tx_removals_and_additions(fetched_full_block, None)
+                    await self.coin_store.new_block(fetched_full_block, tx_additions, tx_removals)
 
             # Changes the peak to be the new peak
             await self.block_store.set_peak(block_record.header_hash)
@@ -376,7 +378,7 @@ async def _reconsider_peak(
         # This is not a heavier block than the heaviest we have seen, so we don't change the coin set
         return None, None, []
 
-    async def get_removals_and_additions(
+    async def get_tx_removals_and_additions(
         self, block: FullBlock, npc_result: Optional[NPCResult] = None
     ) -> Tuple[List[bytes32], List[Coin]]:
         if block.is_transaction_block():
@@ -385,10 +387,10 @@ async def get_removals_and_additions(
                     block_generator: Optional[BlockGenerator] = await self.get_block_generator(block)
                     assert block_generator is not None
                     npc_result = get_name_puzzle_conditions(block_generator, self.constants.MAX_BLOCK_COST_CLVM, False)
-                removals, additions = block_removals_and_additions(block, npc_result.npc_list)
-                return removals, additions
+                tx_removals, tx_additions = tx_removals_and_additions(npc_result.npc_list)
+                return tx_removals, tx_additions
             else:
-                return [], list(block.get_included_reward_coins())
+                return [], []
         else:
             return [], []
 
@@ -660,10 +662,12 @@ async def get_header_blocks_in_range(self, start: int, stop: int) -> Dict[bytes3
         for block in blocks:
             if self.height_to_hash(block.height) != block.header_hash:
                 raise ValueError(f"Block at {block.header_hash} is no longer in the blockchain (it's in a fork)")
-            additions: List[CoinRecord] = await self.coin_store.get_coins_added_at_height(block.height)
+            tx_additions: List[CoinRecord] = [
+                c for c in (await self.coin_store.get_coins_added_at_height(block.height)) if not c.coinbase
+            ]
             removed: List[CoinRecord] = await self.coin_store.get_coins_removed_at_height(block.height)
             header = get_block_header(
-                block, [record.coin for record in additions], [record.coin.name() for record in removed]
+                block, [record.coin for record in tx_additions], [record.coin.name() for record in removed]
             )
             header_blocks[header.header_hash] = header
 

chia/consensus/multiprocess_validation.py

@@ -22,7 +22,7 @@
 from chia.types.header_block import HeaderBlock
 from chia.util.block_cache import BlockCache
 from chia.util.errors import Err
-from chia.util.generator_tools import get_block_header, block_removals_and_additions
+from chia.util.generator_tools import get_block_header, tx_removals_and_additions
 from chia.util.ints import uint16, uint64, uint32
 from chia.util.streamable import Streamable, dataclass_from_dict, streamable
 
@@ -59,16 +59,16 @@ def batch_pre_validate_blocks(
         for i in range(len(full_blocks_pickled)):
             try:
                 block: FullBlock = FullBlock.from_bytes(full_blocks_pickled[i])
-                additions: List[Coin] = list(block.get_included_reward_coins())
+                tx_additions: List[Coin] = []
                 removals: List[bytes32] = []
                 npc_result: Optional[NPCResult] = None
                 if block.height in npc_results:
                     npc_result = NPCResult.from_bytes(npc_results[block.height])
                     assert npc_result is not None
                     if npc_result.npc_list is not None:
-                        removals, additions = block_removals_and_additions(block, npc_result.npc_list)
+                        removals, tx_additions = tx_removals_and_additions(npc_result.npc_list)
                     else:
-                        removals, additions = block_removals_and_additions(block, [])
+                        removals, tx_additions = [], []
 
                 if block.transactions_generator is not None and npc_result is None:
                     prev_generator_bytes = prev_transaction_generators[i]
@@ -79,9 +79,9 @@ def batch_pre_validate_blocks(
                     npc_result = get_name_puzzle_conditions(
                         block_generator, min(constants.MAX_BLOCK_COST_CLVM, block.transactions_info.cost), True
                     )
-                    removals, additions = block_removals_and_additions(block, npc_result.npc_list)
+                    removals, tx_additions = tx_removals_and_additions(npc_result.npc_list)
 
-                header_block = get_block_header(block, additions, removals)
+                header_block = get_block_header(block, tx_additions, removals)
                 required_iters, error = validate_finished_header_block(
                     constants,
                     BlockCache(blocks),

chia/full_node/coin_store.py

@@ -61,15 +61,15 @@ async def create(cls, db_wrapper: DBWrapper, cache_size: uint32 = uint32(60000))
         self.coin_record_cache = LRUCache(cache_size)
         return self
 
-    async def new_block(self, block: FullBlock, additions: List[Coin], removals: List[bytes32]):
+    async def new_block(self, block: FullBlock, tx_additions: List[Coin], tx_removals: List[bytes32]):
         """
         Only called for blocks which are blocks (and thus have rewards and transactions)
         """
         if block.is_transaction_block() is False:
             return
         assert block.foliage_transaction_block is not None
 
-        for coin in additions:
+        for coin in tx_additions:
             record: CoinRecord = CoinRecord(
                 coin,
                 block.height,
@@ -78,7 +78,7 @@ async def new_block(self, block: FullBlock, additions: List[Coin], removals: Lis
                 False,
                 block.foliage_transaction_block.timestamp,
             )
-            await self._add_coin_record(record)
+            await self._add_coin_record(record, False)
 
         included_reward_coins = block.get_included_reward_coins()
         if block.height == 0:
@@ -95,10 +95,14 @@ async def new_block(self, block: FullBlock, additions: List[Coin], removals: Lis
                 True,
                 block.foliage_transaction_block.timestamp,
             )
-            await self._add_coin_record(reward_coin_r)
+            await self._add_coin_record(reward_coin_r, False)
 
-        for coin_name in removals:
-            await self._set_spent(coin_name, block.height)
+        total_amount_spent: int = 0
+        for coin_name in tx_removals:
+            total_amount_spent += await self._set_spent(coin_name, block.height)
+
+        # Sanity check, already checked in block_body_validation
+        assert sum([a.amount for a in tx_additions]) <= total_amount_spent
 
     # Checks DB and DiffStores for CoinRecord with coin_name and returns it
     async def get_coin_record(self, coin_name: bytes32) -> Optional[CoinRecord]:
@@ -218,9 +222,9 @@ async def rollback_to_block(self, block_index: int):
         await c2.close()
 
     # Store CoinRecord in DB and ram cache
-    async def _add_coin_record(self, record: CoinRecord) -> None:
+    async def _add_coin_record(self, record: CoinRecord, allow_replace: bool) -> None:
         cursor = await self.coin_record_db.execute(
-            "INSERT OR REPLACE INTO coin_record VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)",
+            f"INSERT {'OR REPLACE ' if allow_replace else ''}INTO coin_record VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)",
             (
                 record.coin.name().hex(),
                 record.confirmed_block_index,
@@ -237,10 +241,12 @@ async def _add_coin_record(self, record: CoinRecord) -> None:
         self.coin_record_cache.put(record.coin.name().hex(), record)
 
     # Update coin_record to be spent in DB
-    async def _set_spent(self, coin_name: bytes32, index: uint32):
+    async def _set_spent(self, coin_name: bytes32, index: uint32) -> uint64:
         current: Optional[CoinRecord] = await self.get_coin_record(coin_name)
         if current is None:
             raise ValueError(f"Cannot spend a coin that does not exist in db: {coin_name}")
+
+        assert not current.spent  # Redundant sanity check, already checked in block_body_validation
         spent: CoinRecord = CoinRecord(
             current.coin,
             current.confirmed_block_index,
@@ -249,4 +255,5 @@ async def _set_spent(self, coin_name: bytes32, index: uint32):
             current.coinbase,
             current.timestamp,
         )  # type: ignore # noqa
-        await self._add_coin_record(spent)
+        await self._add_coin_record(spent, True)
+        return current.coin.amount

chia/full_node/full_node_api.py

@@ -1008,8 +1008,8 @@ async def request_block_header(self, request: wallet_protocol.RequestBlockHeader
             return msg
         block: Optional[FullBlock] = await self.full_node.block_store.get_full_block(header_hash)
         if block is not None:
-            removals, additions = await self.full_node.blockchain.get_removals_and_additions(block)
-            header_block = get_block_header(block, additions, removals)
+            tx_removals, tx_additions = await self.full_node.blockchain.get_tx_removals_and_additions(block)
+            header_block = get_block_header(block, tx_additions, tx_removals)
             msg = make_msg(
                 ProtocolMessageTypes.respond_block_header,
                 wallet_protocol.RespondBlockHeader(header_block),
@@ -1211,7 +1211,7 @@ async def request_header_blocks(self, request: wallet_protocol.RequestHeaderBloc
         for block in blocks:
             added_coins_records = await self.full_node.coin_store.get_coins_added_at_height(block.height)
             removed_coins_records = await self.full_node.coin_store.get_coins_removed_at_height(block.height)
-            added_coins = [record.coin for record in added_coins_records]
+            added_coins = [record.coin for record in added_coins_records if not record.coinbase]
             removal_names = [record.coin.name() for record in removed_coins_records]
             header_block = get_block_header(block, added_coins, removal_names)
             header_blocks.append(header_block)

chia/util/generator_tools.py

@@ -1,19 +1,18 @@
 from typing import List, Tuple
 from chiabip158 import PyBIP158
 
-from chia.full_node.mempool_check_conditions import get_name_puzzle_conditions
 from chia.types.blockchain_format.coin import Coin
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.types.full_block import FullBlock
-from chia.types.generator_types import BlockGenerator
 from chia.types.header_block import HeaderBlock
 from chia.types.name_puzzle_condition import NPC
 from chia.util.condition_tools import created_outputs_for_conditions_dict
 
 
-def get_block_header(block: FullBlock, addition_coins: List[Coin], removals_names: List[bytes32]) -> HeaderBlock:
+def get_block_header(block: FullBlock, tx_addition_coins: List[Coin], removals_names: List[bytes32]) -> HeaderBlock:
     # Create filter
     byte_array_tx: List[bytes32] = []
+    addition_coins = tx_addition_coins + list(block.get_included_reward_coins())
     if block.is_transaction_block():
         for coin in addition_coins:
             byte_array_tx.append(bytearray(coin.puzzle_hash))
@@ -65,46 +64,3 @@ def tx_removals_and_additions(npc_list: List[NPC]) -> Tuple[List[bytes32], List[
     additions.extend(additions_for_npc(npc_list))
 
     return removals, additions
-
-
-def block_removals_and_additions(block: FullBlock, npc_list: List[NPC]) -> Tuple[List[bytes32], List[Coin]]:
-    """
-    Returns all coins added and removed in block, including farmer and pool reward.
-    """
-
-    removals: List[bytes32] = []
-    additions: List[Coin] = []
-
-    # build removals list
-    if npc_list is None:
-        return [], []
-    for npc in npc_list:
-        removals.append(npc.coin_name)
-
-    additions.extend(additions_for_npc(npc_list))
-
-    rewards = block.get_included_reward_coins()
-    additions.extend(rewards)
-    return removals, additions
-
-
-def run_and_get_removals_and_additions(
-    block: FullBlock, max_cost: int, safe_mode=False
-) -> Tuple[List[bytes32], List[Coin]]:
-    removals: List[bytes32] = []
-    additions: List[Coin] = []
-
-    assert len(block.transactions_generator_ref_list) == 0
-    if not block.is_transaction_block():
-        return [], []
-
-    if block.transactions_generator is not None:
-        npc_result = get_name_puzzle_conditions(BlockGenerator(block.transactions_generator, []), max_cost, safe_mode)
-        # build removals list
-        for npc in npc_result.npc_list:
-            removals.append(npc.coin_name)
-        additions.extend(additions_for_npc(npc_result.npc_list))
-
-    rewards = block.get_included_reward_coins()
-    additions.extend(rewards)
-    return removals, additions

tests/blockchain/test_blockchain_transactions.py

@@ -11,11 +11,11 @@
 from chia.types.condition_with_args import ConditionWithArgs
 from chia.types.spend_bundle import SpendBundle
 from chia.util.errors import ConsensusError, Err
-from chia.util.generator_tools import run_and_get_removals_and_additions
 from chia.util.ints import uint64
 from chia.util.wallet_tools import WalletTool
 from tests.core.full_node.test_full_node import connect_and_get_peer
 from tests.setup_nodes import bt, setup_two_nodes, test_constants
+from tests.util.generator_tools_testing import run_and_get_removals_and_additions
 
 BURN_PUZZLE_HASH = b"0" * 32