From 048d481fdd1ea48c38f4bf4813ae7ba990bdd8f0 Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Mon, 26 Jun 2023 23:24:58 +1000
Subject: [PATCH] Fix state lookups by root for hot finalized states

---
 beacon_node/beacon_chain/src/beacon_chain.rs |  7 +++--
 beacon_node/http_api/src/state_id.rs         | 29 ++++++++++++++------
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
index ceda7222e69..64dade3b494 100644
--- a/beacon_node/beacon_chain/src/beacon_chain.rs
+++ b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -450,11 +450,11 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     /// Checks if a state is finalized.
     /// The finalization check is done with the slot. The state root is used to verify that
     /// the finalized state is in the canonical chain.
-    pub fn is_finalized_state(
+    pub fn is_canonical_finalized_state(
         &self,
         state_root: &Hash256,
         state_slot: Slot,
-    ) -> Result<bool, Error> {
+    ) -> Result<(bool, bool), Error> {
         let finalized_slot = self
             .canonical_head
             .cached_head()
@@ -464,7 +464,8 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         let is_canonical = self
             .state_root_at_slot(state_slot)?
             .map_or(false, |canonical_root| state_root == &canonical_root);
-        Ok(state_slot <= finalized_slot && is_canonical)
+        let is_finalized_slot = state_slot <= finalized_slot;
+        Ok((is_canonical, is_finalized_slot))
     }
 
     /// Persists the head tracker and fork choice.
diff --git a/beacon_node/http_api/src/state_id.rs b/beacon_node/http_api/src/state_id.rs
index 9e4aadef17e..88c73202d13 100644
--- a/beacon_node/http_api/src/state_id.rs
+++ b/beacon_node/http_api/src/state_id.rs
@@ -70,15 +70,28 @@ impl StateId {
                     .map_err(BeaconChainError::DBError)
                     .map_err(warp_utils::reject::beacon_chain_error)?
                 {
-                    let execution_optimistic = chain
-                        .canonical_head
-                        .fork_choice_read_lock()
-                        .is_optimistic_or_invalid_block_no_fallback(&hot_summary.latest_block_root)
-                        .map_err(BeaconChainError::ForkChoiceError)
-                        .map_err(warp_utils::reject::beacon_chain_error)?;
-                    let finalized = chain
-                        .is_finalized_state(root, hot_summary.slot)
+                    let (canonical, finalized_slot) = chain
+                        .is_canonical_finalized_state(root, hot_summary.slot)
                         .map_err(warp_utils::reject::beacon_chain_error)?;
+                    let finalized = canonical && finalized_slot;
+                    let fork_choice = chain.canonical_head.fork_choice_read_lock();
+                    let execution_optimistic = if finalized_slot && !canonical {
+                        // This block is permanently orphaned and has likely been pruned from fork
+                        // choice. If it isn't found in fork choice, mark it optimistic to be on the
+                        // safe side.
+                        fork_choice
+                            .is_optimistic_or_invalid_block(&hot_summary.latest_block_root)
+                            .unwrap_or(true)
+                    } else {
+                        // This block is either old and finalized, or recent and unfinalized, so
+                        // it's safe to fallback to the optimistic status of the finalized block.
+                        chain
+                            .canonical_head
+                            .fork_choice_read_lock()
+                            .is_optimistic_or_invalid_block(&hot_summary.latest_block_root)
+                            .map_err(BeaconChainError::ForkChoiceError)
+                            .map_err(warp_utils::reject::beacon_chain_error)?
+                    };
                     return Ok((*root, execution_optimistic, finalized));
                 } else if let Some(_cold_state_slot) = chain
                     .store