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: Optional[uint32] = None) -> bool:
         return False  # pragma: no cover
 
     def contains_height(self, height: uint32) -> bool:

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

@@ -788,12 +788,16 @@ 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: 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 block_record(self, header_hash: bytes32) -> BlockRecord:
         return self.__block_records[header_hash]
@@ -955,7 +959,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: Optional[uint32] = None) -> 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)
@@ -757,7 +758,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

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,18 @@ 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: 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
+        if not self.contains_height(height):
+            return False
+        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: