From 1ca3fe1aad0ab17745d613f45757fb9259c7d834 Mon Sep 17 00:00:00 2001
From: Amine Khaldi <amine.khaldi@reactos.org>
Date: Wed, 11 Dec 2024 15:26:50 +0100
Subject: [PATCH] Order header blocks by height in get_header_blocks_in_range.

---
 chia/_tests/blockchain/test_blockchain.py | 22 ++++++++++++++++++++++
 chia/consensus/blockchain.py              | 22 ++++++++++++++--------
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/chia/_tests/blockchain/test_blockchain.py b/chia/_tests/blockchain/test_blockchain.py
index 89b015cc524f..c1c101d58af3 100644
--- a/chia/_tests/blockchain/test_blockchain.py
+++ b/chia/_tests/blockchain/test_blockchain.py
@@ -4184,3 +4184,25 @@ async def get_fork_info(blockchain: Blockchain, block: FullBlock, peak: BlockRec
     )
 
     return fork_info
+
+
+@pytest.mark.anyio
+async def test_get_header_blocks_in_range_order(empty_blockchain: Blockchain, bt: BlockTools) -> None:
+    """
+    Tests that we return the headers in the correct order even when we fetch
+    some blocks from the block cache and some from the DB.
+    """
+    bc = empty_blockchain
+    blocks = bt.get_consecutive_blocks(5)
+    blocks_hh = []
+    for block in blocks:
+        await _validate_and_add_block(bc, block)
+        blocks_hh.append(block.header_hash)
+    hh_to_header_block_map = await bc.get_header_blocks_in_range(0, 5)
+    assert list(hh_to_header_block_map) == blocks_hh
+    # Remove 1st and 3rd items from the block cache, so we fetch them from DB
+    for i in [0, 2]:
+        bc.block_store.block_cache.remove(blocks[i].header_hash)
+    # The order should remain the same
+    hh_to_header_block_map = await bc.get_header_blocks_in_range(0, 5)
+    assert list(hh_to_header_block_map) == blocks_hh
diff --git a/chia/consensus/blockchain.py b/chia/consensus/blockchain.py
index 71e5d49c58a9..8b6bd70475bc 100644
--- a/chia/consensus/blockchain.py
+++ b/chia/consensus/blockchain.py
@@ -877,23 +877,29 @@ async def get_block_records_in_range(self, start: int, stop: int) -> dict[bytes3
     async def get_header_blocks_in_range(
         self, start: int, stop: int, tx_filter: bool = True
     ) -> dict[bytes32, HeaderBlock]:
-        hashes = []
+        hashes: list[bytes32] = []
         for height in range(start, stop + 1):
             header_hash: Optional[bytes32] = self.height_to_hash(uint32(height))
             if header_hash is not None:
                 hashes.append(header_hash)
 
-        blocks: list[FullBlock] = []
-        for hash in hashes.copy():
+        header_hash_to_block_map: dict[bytes32, FullBlock] = {}
+        non_cache_hashes: list[bytes32] = []
+        for hash in hashes:
             block = self.block_store.block_cache.get(hash)
             if block is not None:
-                blocks.append(block)
-                hashes.remove(hash)
-        blocks_on_disk: list[FullBlock] = await self.block_store.get_blocks_by_hash(hashes)
-        blocks.extend(blocks_on_disk)
+                header_hash_to_block_map[hash] = block
+            else:
+                non_cache_hashes.append(hash)
+        blocks_on_disk: list[FullBlock] = await self.block_store.get_blocks_by_hash(non_cache_hashes)
+        header_hash_to_block_map.update({block.header_hash: block for block in blocks_on_disk})
         header_blocks: dict[bytes32, HeaderBlock] = {}
 
-        for block in blocks:
+        for hash in hashes:
+            block = header_hash_to_block_map.get(hash)
+            # get_blocks_by_hash throws an exception if the blocks are not
+            # present, so this is just defensive.
+            assert block is not None
             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)")
             if tx_filter is False: