From 71947641a648a0d3168f078fbf84efb2dfa4af01 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Fri, 17 Jan 2020 16:12:22 +1100
Subject: [PATCH] Fix bug in checkpoint manager

---
 beacon_node/beacon_chain/src/fork_choice.rs   |  2 +-
 .../src/fork_choice/checkpoint_manager.rs     | 20 ++++++++++++++-----
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/beacon_node/beacon_chain/src/fork_choice.rs b/beacon_node/beacon_chain/src/fork_choice.rs
index d7410a8f86..6d06e2c8d9 100644
--- a/beacon_node/beacon_chain/src/fork_choice.rs
+++ b/beacon_node/beacon_chain/src/fork_choice.rs
@@ -134,7 +134,7 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
 
         self.checkpoint_manager
             .write()
-            .process_state(state, chain, &self.backend)?;
+            .process_state(block_root, state, chain, &self.backend)?;
         self.checkpoint_manager.write().update(chain)?;
 
         // Note: we never count the block as a latest message, only attestations.
diff --git a/beacon_node/beacon_chain/src/fork_choice/checkpoint_manager.rs b/beacon_node/beacon_chain/src/fork_choice/checkpoint_manager.rs
index 33211f62c1..fdff2d6b2a 100644
--- a/beacon_node/beacon_chain/src/fork_choice/checkpoint_manager.rs
+++ b/beacon_node/beacon_chain/src/fork_choice/checkpoint_manager.rs
@@ -18,9 +18,17 @@ struct BalancesCache {
 }
 
 impl BalancesCache {
-    pub fn process_state<E: EthSpec>(&mut self, state: &BeaconState<E>) -> Result<(), Error> {
+    pub fn process_state<E: EthSpec>(
+        &mut self,
+        block_root: Hash256,
+        state: &BeaconState<E>,
+    ) -> Result<(), Error> {
         let epoch_boundary_slot = state.current_epoch().start_slot(E::slots_per_epoch());
-        let epoch_boundary_root = *state.get_block_root(epoch_boundary_slot)?;
+        let epoch_boundary_root = if epoch_boundary_slot == state.slot {
+            block_root
+        } else {
+            *state.get_block_root(epoch_boundary_slot)?
+        };
 
         if self.position(epoch_boundary_root).is_none() {
             let item = CacheItem {
@@ -129,12 +137,14 @@ impl CheckpointManager {
         Ok(())
     }
 
-    /// Checks the given `state` to see if it contains a `current_justified_checkpoint` that is
-    /// better than `self.best_justified_checkpoint`. If so, the value is updated.
+    /// Checks the given `state` (must correspond to the given `block_root`) to see if it contains
+    /// a `current_justified_checkpoint` that is better than `self.best_justified_checkpoint`. If
+    /// so, the value is updated.
     ///
     /// Note: this does not update `self.justified_checkpoint`.
     pub fn process_state<T: BeaconChainTypes>(
         &mut self,
+        block_root: Hash256,
         state: &BeaconState<T::EthSpec>,
         chain: &BeaconChain<T>,
         proto_array: &ProtoArrayForkChoice,
@@ -189,7 +199,7 @@ impl CheckpointManager {
             }
 
             // Add the state's balances to the balances cache to avoid a state read later.
-            self.balances_cache.process_state(state)?;
+            self.balances_cache.process_state(block_root, state)?;
         }
 
         Ok(())