Assert height to hash in contains block

chia/_tests/blockchain/test_augmented_chain.py

@@ -46,7 +46,7 @@ def height_to_block_record(self, height: uint32) -> BlockRecord:
     def height_to_hash(self, height: uint32) -> Optional[bytes32]:
         return self.heights.get(height)
 
-    def contains_block(self, header_hash: bytes32) -> bool:
+    def contains_block(self, header_hash: bytes32, height: uint32) -> bool:
         return False  # pragma: no cover
 
     def contains_height(self, height: uint32) -> bool:
@@ -122,13 +122,13 @@ async def test_augmented_chain(default_10000_blocks: list[FullBlock]) -> None:
         assert abc.try_block_record(blocks[i].header_hash) == block_records[i]
         assert abc.height_to_hash(uint32(i)) == blocks[i].header_hash
         assert await abc.prev_block_hash([blocks[i].header_hash]) == [blocks[i].prev_header_hash]
-        assert abc.contains_block(blocks[i].header_hash) is True
+        assert abc.block_record(blocks[i].header_hash) is not None
         assert await abc.get_block_record_from_db(blocks[i].header_hash) == block_records[i]
         assert abc.contains_height(uint32(i))
 
     for i in range(5, 10):
         assert abc.height_to_hash(uint32(i)) is None
-        assert not abc.contains_block(blocks[i].header_hash)
+        assert abc.block_record(blocks[i].header_hash) is None
         assert not await abc.get_block_record_from_db(blocks[i].header_hash)
         assert not abc.contains_height(uint32(i))
 

chia/_tests/util/blockchain_mock.py

@@ -66,8 +66,16 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
         assert height in self._height_to_hash
         return self._height_to_hash[height]
 
-    def contains_block(self, header_hash: bytes32) -> bool:
-        return header_hash in self._block_records
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
+        if height is None:
+            block_rec = self.try_block_record(header_hash)
+            if block_rec is None:
+                return False
+            height = block_rec.height
+        block_hash_from_hh = self.height_to_hash(height)
+        if block_hash_from_hh is None or block_hash_from_hh != header_hash:
+            return False
+        return True
 
     async def contains_block_from_db(self, header_hash: bytes32) -> bool:
         return header_hash in self._block_records

chia/consensus/blockchain.py

@@ -609,15 +609,13 @@ async def _reconsider_peak(
         )
 
     def get_next_difficulty(self, header_hash: bytes32, new_slot: bool) -> uint64:
-        assert self.contains_block(header_hash)
         curr = self.block_record(header_hash)
         if curr.height <= 2:
             return self.constants.DIFFICULTY_STARTING
 
         return get_next_sub_slot_iters_and_difficulty(self.constants, new_slot, curr, self)[1]
 
     def get_next_slot_iters(self, header_hash: bytes32, new_slot: bool) -> uint64:
-        assert self.contains_block(header_hash)
         curr = self.block_record(header_hash)
         if curr.height <= 2:
             return self.constants.SUB_SLOT_ITERS_STARTING
@@ -704,7 +702,7 @@ async def validate_unfinished_block_header(
             return None, Err.TOO_MANY_GENERATOR_REFS
 
         if (
-            not self.contains_block(block.prev_header_hash)
+            self.block_record(block.prev_header_hash) is None
             and block.prev_header_hash != self.constants.GENESIS_CHALLENGE
         ):
             return None, Err.INVALID_PREV_BLOCK_HASH
@@ -788,12 +786,11 @@ async def validate_unfinished_block(
 
         return PreValidationResult(None, required_iters, conds, uint32(0))
 
-    def contains_block(self, header_hash: bytes32) -> bool:
-        """
-        True if we have already added this block to the chain. This may return false for orphan blocks
-        that we have added but no longer keep in memory.
-        """
-        return header_hash in self.__block_records
+    def contains_block(self, header_hash: bytes32, height: uint32) -> bool:
+        block_hash_from_hh = self.height_to_hash(height)
+        if block_hash_from_hh is None or block_hash_from_hh != header_hash:
+            return False
+        return True
 
     def block_record(self, header_hash: bytes32) -> BlockRecord:
         return self.__block_records[header_hash]
@@ -955,7 +952,7 @@ async def get_block_records_at(self, heights: list[uint32], batch_size: int = 90
         return records
 
     def try_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]:
-        if self.contains_block(header_hash):
+        if header_hash in self.__block_records:
             return self.block_record(header_hash)
         return None
 

chia/consensus/blockchain_interface.py

@@ -14,7 +14,7 @@ class BlockRecordsProtocol(Protocol):
     def try_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]: ...
     def block_record(self, header_hash: bytes32) -> BlockRecord: ...
     def contains_height(self, height: uint32) -> bool: ...
-    def contains_block(self, header_hash: bytes32) -> bool: ...
+    def contains_block(self, header_hash: bytes32, height: uint32) -> bool: ...
     def height_to_hash(self, height: uint32) -> Optional[bytes32]: ...
     def height_to_block_record(self, height: uint32) -> BlockRecord: ...
 

chia/consensus/difficulty_adjustment.py

@@ -224,11 +224,10 @@ def _get_next_sub_slot_iters(
     if next_height < constants.EPOCH_BLOCKS:
         return uint64(constants.SUB_SLOT_ITERS_STARTING)
 
-    if not blocks.contains_block(prev_header_hash):
+    prev_b = blocks.try_block_record(prev_header_hash)
+    if prev_b is None:
         raise ValueError(f"Header hash {prev_header_hash} not in blocks")
 
-    prev_b: BlockRecord = blocks.block_record(prev_header_hash)
-
     # If we are in the same epoch, return same ssi
     if not skip_epoch_check:
         _, can_finish_epoch = can_finish_sub_and_full_epoch(
@@ -304,11 +303,10 @@ def _get_next_difficulty(
         # We are in the first epoch
         return uint64(constants.DIFFICULTY_STARTING)
 
-    if not blocks.contains_block(prev_header_hash):
+    prev_b = blocks.try_block_record(prev_header_hash)
+    if prev_b is None:
         raise ValueError(f"Header hash {prev_header_hash} not in blocks")
 
-    prev_b: BlockRecord = blocks.block_record(prev_header_hash)
-
     # If we are in the same slot as previous block, return same difficulty
     if not skip_epoch_check:
         _, can_finish_epoch = can_finish_sub_and_full_epoch(

chia/full_node/full_node.py

@@ -619,6 +619,7 @@ async def short_sync_batch(self, peer: WSChiaConnection, start_height: uint32, t
                 fork_hash = self.constants.GENESIS_CHALLENGE
             assert fork_hash
             fork_info = ForkInfo(start_height - 1, start_height - 1, fork_hash)
+            blockchain = AugmentedBlockchain(self.blockchain)
             for height in range(start_height, target_height, batch_size):
                 end_height = min(target_height, height + batch_size)
                 request = RequestBlocks(uint32(height), uint32(end_height), True)
@@ -636,7 +637,6 @@ async def short_sync_batch(self, peer: WSChiaConnection, start_height: uint32, t
                         self.constants, new_slot, prev_b, self.blockchain
                     )
                     vs = ValidationState(ssi, diff, None)
-                    blockchain = AugmentedBlockchain(self.blockchain)
                     success, state_change_summary = await self.add_block_batch(
                         response.blocks, peer_info, fork_info, vs, blockchain
                     )
@@ -757,7 +757,7 @@ async def new_peak(self, request: full_node_protocol.NewPeak, peer: WSChiaConnec
         # Store this peak/peer combination in case we want to sync to it, and to keep track of peers
         self.sync_store.peer_has_block(request.header_hash, peer.peer_node_id, request.weight, request.height, True)
 
-        if self.blockchain.contains_block(request.header_hash):
+        if self.blockchain.contains_block(request.header_hash, request.height):
             return None
 
         # Not interested in less heavy peaks
@@ -2018,7 +2018,7 @@ async def add_block(
 
         # Adds the block to seen, and check if it's seen before (which means header is in memory)
         header_hash = block.header_hash
-        if self.blockchain.contains_block(header_hash):
+        if self.blockchain.contains_block(header_hash, block.height):
             if fork_info is not None:
                 await self.blockchain.run_single_block(block, fork_info)
             return None
@@ -2092,7 +2092,7 @@ async def add_block(
             enable_profiler(self.profile_block_validation) as pr,
         ):
             # After acquiring the lock, check again, because another asyncio thread might have added it
-            if self.blockchain.contains_block(header_hash):
+            if self.blockchain.contains_block(header_hash, block.height):
                 if fork_info is not None:
                     await self.blockchain.run_single_block(block, fork_info)
                 return None
@@ -2280,8 +2280,9 @@ async def add_unfinished_block(
         """
         receive_time = time.time()
 
-        if block.prev_header_hash != self.constants.GENESIS_CHALLENGE and not self.blockchain.contains_block(
-            block.prev_header_hash
+        if (
+            block.prev_header_hash != self.constants.GENESIS_CHALLENGE
+            and self.blockchain.block_record(block.prev_header_hash) is None
         ):
             # No need to request the parent, since the peer will send it to us anyway, via NewPeak
             self.log.debug("Received a disconnected unfinished block")

chia/full_node/full_node_api.py

@@ -285,7 +285,7 @@ async def respond_transaction(
     async def request_proof_of_weight(self, request: full_node_protocol.RequestProofOfWeight) -> Optional[Message]:
         if self.full_node.weight_proof_handler is None:
             return None
-        if not self.full_node.blockchain.contains_block(request.tip):
+        if self.full_node.blockchain.block_record(request.tip) is None:
             self.log.error(f"got weight proof request for unknown peak {request.tip}")
             return None
         if request.tip in self.full_node.pow_creation:

chia/simulator/add_blocks_in_batches.py

@@ -35,13 +35,12 @@ async def add_blocks_in_batches(
     fork_info = ForkInfo(fork_height, blocks[0].height - 1, peak_hash)
 
     vs = ValidationState(ssi, diff, None)
-
+    blockchain = AugmentedBlockchain(full_node.blockchain)
     for block_batch in to_batches(blocks, 64):
         b = block_batch.entries[0]
         if (b.height % 128) == 0:
             print(f"main chain: {b.height:4} weight: {b.weight}")
         # vs is updated by the call to add_block_batch()
-        blockchain = AugmentedBlockchain(full_node.blockchain)
         success, state_change_summary = await full_node.add_block_batch(
             block_batch.entries, PeerInfo("0.0.0.0", 0), fork_info, vs, blockchain
         )

chia/util/augmented_chain.py

@@ -47,8 +47,7 @@ def add_extra_block(self, block: FullBlock, block_record: BlockRecord) -> None:
 
     def remove_extra_block(self, hh: bytes32) -> None:
         if hh in self._extra_blocks:
-            block_record = self._extra_blocks.pop(hh)[1]
-            del self._height_to_hash[block_record.height]
+            self._extra_blocks.pop(hh)[1]
 
     # BlocksProtocol
     async def lookup_block_generators(self, header_hash: bytes32, generator_refs: set[uint32]) -> dict[uint32, bytes]:
@@ -82,12 +81,11 @@ async def get_block_record_from_db(self, header_hash: bytes32) -> Optional[Block
 
     def add_block_record(self, block_record: BlockRecord) -> None:
         self._underlying.add_block_record(block_record)
-
+        self._height_to_hash[block_record.height] = block_record.header_hash
         # now that we're adding the block to the underlying blockchain, we don't
         # need to keep the extra block around anymore
         hh = block_record.header_hash
         if hh in self._extra_blocks:
-            del self._height_to_hash[block_record.height]
             del self._extra_blocks[hh]
 
     # BlockRecordsProtocol
@@ -109,6 +107,7 @@ def height_to_block_record(self, height: uint32) -> BlockRecord:
             ret = self._get_block_record(header_hash)
             if ret is not None:
                 return ret
+            return self._underlying.block_record(header_hash)
         return self._underlying.height_to_block_record(height)
 
     def height_to_hash(self, height: uint32) -> Optional[bytes32]:
@@ -117,8 +116,11 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
             return ret
         return self._underlying.height_to_hash(height)
 
-    def contains_block(self, header_hash: bytes32) -> bool:
-        return (header_hash in self._extra_blocks) or self._underlying.contains_block(header_hash)
+    def contains_block(self, header_hash: bytes32, height: uint32) -> bool:
+        block_hash_from_hh = self.height_to_hash(height)
+        if block_hash_from_hh is None or block_hash_from_hh != header_hash:
+            return False
+        return True
 
     def contains_height(self, height: uint32) -> bool:
         return (height in self._height_to_hash) or self._underlying.contains_height(height)

chia/util/block_cache.py

@@ -43,8 +43,16 @@ def height_to_hash(self, height: uint32) -> Optional[bytes32]:
             return None
         return self._height_to_hash[height]
 
-    def contains_block(self, header_hash: bytes32) -> bool:
-        return header_hash in self._block_records
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
+        if height is None:
+            block_rec = self.try_block_record(header_hash)
+            if block_rec is None:
+                return False
+            height = block_rec.height
+        block_hash_from_hh = self.height_to_hash(height)
+        if block_hash_from_hh is None or block_hash_from_hh != header_hash:
+            return False
+        return True
 
     def contains_height(self, height: uint32) -> bool:
         return height in self._height_to_hash

chia/wallet/wallet_blockchain.py

@@ -192,7 +192,11 @@ async def get_finished_sync_up_to(self) -> uint32:
     def get_latest_timestamp(self) -> uint64:
         return self._latest_timestamp
 
-    def contains_block(self, header_hash: bytes32) -> bool:
+    def contains_block(self, header_hash: bytes32, height: Optional[uint32] = None) -> bool:
+        """
+        True if we have already added this block to the chain. This may return false for orphan blocks
+        that we have added but no longer keep in memory.
+        """
         return header_hash in self._block_records
 
     def contains_height(self, height: uint32) -> bool: